From 6e6ad053d4555125016bd0637f4b38cbbd953c23 Mon Sep 17 00:00:00 2001
From: Fimeg <casey.tunturi@gmail.com>
Date: Sat, 20 Dec 2025 15:59:56 -0500
Subject: [PATCH] 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 | 38 +++++++++++++++++--
 .../internal/database/queries/commands.go     | 30 +++++++++++++++
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/aggregator-server/internal/command/factory.go b/aggregator-server/internal/command/factory.go
index 7dd25d4..13b1b3a 100644
--- a/aggregator-server/internal/command/factory.go
+++ b/aggregator-server/internal/command/factory.go
@@ -5,19 +5,22 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/Fimeg/RedFlag/aggregator-server/internal/database/queries"
 	"github.com/Fimeg/RedFlag/aggregator-server/internal/models"
 	"github.com/google/uuid"
 )
 
 // Factory creates validated AgentCommand instances
 type Factory struct {
-	validator *Validator
+	validator      *Validator
+	commandQueries *queries.CommandQueries
 }
 
 // NewFactory creates a new command factory
-func NewFactory() *Factory {
+func NewFactory(commandQueries *queries.CommandQueries) *Factory {
 	return &Factory{
-		validator: NewValidator(),
+		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 {
diff --git a/aggregator-server/internal/database/queries/commands.go b/aggregator-server/internal/database/queries/commands.go
index c82d35c..eb9c1da 100644
--- a/aggregator-server/internal/database/queries/commands.go
+++ b/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()