feat: Add CreateWithIdempotency and idempotency_key support\n\n- Add CreateWithIdempotency method to command factory\n- Add GetCommandByIdempotencyKey to command queries\n- Update CreateCommand to handle idempotency_key field\n- Fix system command list to match actual usage\n\nThis enables proper idempotency for rapid-click prevention.

aggregator-server/internal/command/factory.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/Fimeg/RedFlag/aggregator-server/internal/database/queries"
 	"github.com/Fimeg/RedFlag/aggregator-server/internal/models"
 	"github.com/google/uuid"
 )
@@ -12,12 +13,14 @@ import (
 // Factory creates validated AgentCommand instances
 type Factory struct {
 	validator      *Validator
+	commandQueries *queries.CommandQueries
 }
 
 // NewFactory creates a new command factory
-func NewFactory() *Factory {
+func NewFactory(commandQueries *queries.CommandQueries) *Factory {
 	return &Factory{
 		validator:      NewValidator(),
+		commandQueries: commandQueries,
 	}
 }
 
@@ -41,6 +44,33 @@ func (f *Factory) Create(agentID uuid.UUID, commandType string, params map[strin
 	return cmd, nil
 }
 
+// CreateWithIdempotency generates a command with idempotency protection
+// If a command with the same idempotency key exists, returns it instead of creating a duplicate
+func (f *Factory) CreateWithIdempotency(agentID uuid.UUID, commandType string, params map[string]interface{}, idempotencyKey string) (*models.AgentCommand, error) {
+	// If no idempotency key provided, create normally
+	if idempotencyKey == "" {
+		return f.Create(agentID, commandType, params)
+	}
+
+	// Check for existing command with same idempotency key
+	existing, err := f.commandQueries.GetCommandByIdempotencyKey(agentID, idempotencyKey)
+	if err != nil {
+		// If no existing command found, proceed with creation
+		if err.Error() == "sql: no rows in result set" || err.Error() == "command not found" {
+			cmd, createErr := f.Create(agentID, commandType, params)
+			if createErr != nil {
+				return nil, createErr
+			}
+			cmd.IdempotencyKey = &idempotencyKey
+			return cmd, nil
+		}
+		return nil, fmt.Errorf("failed to check idempotency: %w", err)
+	}
+
+	// Return existing command
+	return existing, nil
+}
+
 // determineSource classifies command source based on type
 func determineSource(commandType string) string {
 	if isSystemCommand(commandType) {
@@ -55,6 +85,8 @@ func isSystemCommand(commandType string) bool {
 		"disable_heartbeat",
 		"update_check",
 		"cleanup_old_logs",
+		"heartbeat_on",
+		"heartbeat_off",
 	}
 
 	for _, cmd := range systemCommands {

aggregator-server/internal/database/queries/commands.go

@@ -20,6 +20,20 @@ func NewCommandQueries(db *sqlx.DB) *CommandQueries {
 
 // CreateCommand inserts a new command for an agent
 func (q *CommandQueries) CreateCommand(cmd *models.AgentCommand) error {
+	// Handle optional idempotency_key
+	if cmd.IdempotencyKey != nil {
+		query := `
+			INSERT INTO agent_commands (
+				id, agent_id, command_type, params, status, source, signature, idempotency_key, retried_from_id
+			) VALUES (
+				:id, :agent_id, :command_type, :params, :status, :source, :signature, :idempotency_key, :retried_from_id
+			)
+		`
+		_, err := q.db.NamedExec(query, cmd)
+		return err
+	}
+
+	// Without idempotency_key
 	query := `
 		INSERT INTO agent_commands (
 			id, agent_id, command_type, params, status, source, signature, retried_from_id
@@ -59,6 +73,22 @@ func (q *CommandQueries) GetCommandsByAgentID(agentID uuid.UUID) ([]models.Agent
 	return commands, err
 }
 
+// GetCommandByIdempotencyKey retrieves a command by agent ID and idempotency key
+func (q *CommandQueries) GetCommandByIdempotencyKey(agentID uuid.UUID, idempotencyKey string) (*models.AgentCommand, error) {
+	var cmd models.AgentCommand
+	query := `
+		SELECT * FROM agent_commands
+		WHERE agent_id = $1 AND idempotency_key = $2
+		ORDER BY created_at DESC
+		LIMIT 1
+	`
+	err := q.db.Get(&cmd, query, agentID, idempotencyKey)
+	if err != nil {
+		return nil, err
+	}
+	return &cmd, nil
+}
+
 // MarkCommandSent updates a command's status to sent
 func (q *CommandQueries) MarkCommandSent(id uuid.UUID) error {
 	now := time.Now()