AetherForge/server/internal/ai/commands.go

package ai

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
)

const (
	CmdBulkCommand      = "bulk_command"
	CmdAgentCommand     = "agent_command"
	CmdDiscoverAndJoin  = "discover_and_join"
	CmdRestartMining    = "restart_mining"
	CmdReorderTiers     = "reorder_tiers"
	CmdSpreadNow        = "spread_now"
	CmdSpreadRetryLane  = "spread_retry_lane"
	CmdSkipTier         = "skip_tier"
	CmdStageFetch       = "stage_fetch"
	CmdSetAgentVersion  = "set_agent_version"
	CmdSpreadGraft      = "spread_graft"
	CmdPersonaTweak     = "persona_tweak"
	CmdEnableErasure    = "enable_erasure"
	CmdStrainHospice    = "strain_hospice"
	CmdNoop             = "noop"
)

var knownCommands = map[string]bool{
	CmdBulkCommand:     true,
	CmdAgentCommand:    true,
	CmdDiscoverAndJoin: true,
	CmdRestartMining:   true,
	CmdReorderTiers:    true,
	CmdSpreadNow:       true,
	CmdSpreadRetryLane: true,
	CmdSkipTier:        true,
	CmdStageFetch:      true,
	CmdSetAgentVersion: true,
	CmdSpreadGraft:     true,
	CmdPersonaTweak:    true,
	CmdEnableErasure:   true,
	CmdStrainHospice:   true,
	CmdNoop:            true,
}

var jsonBlockRe = regexp.MustCompile(`(?s)\{[\s\n]*"commands"\s*:\s*\[[\s\S]*?\]\s*\}`)

// ParseCommands extracts fleet commands from LLM text (JSON block, tool-call, or COMMAND: lines).
func ParseCommands(raw string) []Command {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return nil
	}
	if cmds := parseCommandsJSON(raw); len(cmds) > 0 {
		return cmds
	}
	if block := jsonBlockRe.FindString(raw); block != "" && block != raw {
		if cmds := parseCommandsJSON(block); len(cmds) > 0 {
			return cmds
		}
	}
	if cmd := parseToolCall(raw); cmd != nil {
		return []Command{*cmd}
	}
	return parseCommandLines(raw)
}

func parseCommandsJSON(raw string) []Command {
	var envelope struct {
		Commands []Command `json:"commands"`
	}
	if err := json.Unmarshal([]byte(raw), &envelope); err == nil && len(envelope.Commands) > 0 {
		return normalizeCommands(envelope.Commands)
	}
	// Bare array
	var arr []Command
	if err := json.Unmarshal([]byte(raw), &arr); err == nil && len(arr) > 0 {
		return normalizeCommands(arr)
	}
	return nil
}

func parseToolCall(raw string) *Command {
	var tool struct {
		Tool string                 `json:"tool"`
		Args map[string]interface{} `json:"args"`
		Type string                 `json:"type"`
	}
	start := strings.Index(raw, "{")
	end := strings.LastIndex(raw, "}")
	if start < 0 || end <= start {
		return nil
	}
	if err := json.Unmarshal([]byte(raw[start:end+1]), &tool); err != nil {
		return nil
	}
	name := strings.TrimSpace(tool.Tool)
	if name == "" {
		name = strings.TrimSpace(tool.Type)
	}
	if name == "" {
		return nil
	}
	name = normalizeCommandType(name)
	if !knownCommands[name] && name != "restart_agent" {
		return &Command{Type: CmdAgentCommand, Args: map[string]interface{}{"action": name, "args": tool.Args}}
	}
	if name == "restart_agent" {
		name = CmdRestartMining
	}
	return &Command{Type: name, Args: tool.Args}
}

func parseCommandLines(raw string) []Command {
	var out []Command
	for _, line := range strings.Split(raw, "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(strings.ToUpper(line), "COMMAND:") {
			continue
		}
		rest := strings.TrimSpace(line[len("COMMAND:"):])
		if rest == "" {
			continue
		}
		parts := strings.Fields(rest)
		cmdType := normalizeCommandType(parts[0])
		args := map[string]interface{}{}
		for _, p := range parts[1:] {
			kv := strings.SplitN(p, "=", 2)
			if len(kv) == 2 {
				args[kv[0]] = kv[1]
			}
		}
		out = append(out, Command{Type: cmdType, Args: args})
	}
	return normalizeCommands(out)
}

func normalizeCommands(cmds []Command) []Command {
	out := make([]Command, 0, len(cmds))
	for _, c := range cmds {
		typ := normalizeCommandType(c.Type)
		if typ == "" {
			continue
		}
		if typ == "restart_agent" {
			typ = CmdRestartMining
		}
		args := c.Args
		if args == nil {
			args = map[string]interface{}{}
		}
		out = append(out, Command{Type: typ, Args: args})
	}
	return out
}

func normalizeCommandType(s string) string {
	s = strings.TrimSpace(strings.ToLower(s))
	s = strings.ReplaceAll(s, "-", "_")
	if s == "restart" {
		return CmdRestartMining
	}
	if knownCommands[s] {
		return s
	}
	return s
}

// FormatExecuted summarizes commands for audit log storage.
func FormatExecuted(cmds []Command, results []string) string {
	parts := make([]string, 0, len(cmds))
	for i, c := range cmds {
		msg := c.Type
		if i < len(results) && results[i] != "" {
			msg = fmt.Sprintf("%s:%s", c.Type, results[i])
		}
		parts = append(parts, msg)
	}
	return strings.Join(parts, "; ")
}