Agent-System/internal/squad/gov_test.go

package squad_test

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"chat-agent-system/internal/agent"
	"chat-agent-system/internal/squad"
)

// ============================================================================
// StepFSM 状态机
// ============================================================================

func TestStepFSM_ValidAndInvalidTransitions(t *testing.T) {
	cases := []struct {
		from, to string
		ok       bool
	}{
		{squad.StepStatePending, squad.StepStateExecuting, true},
		{squad.StepStateExecuting, squad.StepStateCompleted, true},
		{squad.StepStateExecuting, squad.StepStateVerifying, true},
		{squad.StepStateVerifying, squad.StepStateCompleted, true},
		{squad.StepStateVerifying, squad.StepStateRepairing, true},
		{squad.StepStateRepairing, squad.StepStateVerifying, true},
		{squad.StepStateFailed, squad.StepStateExecuting, true},      // 断点续跑：失败可重试
		{squad.StepStateAborted, squad.StepStateExecuting, true},     // 中止可恢复
		{squad.StepStateCompleted, squad.StepStateExecuting, false},  // 完成不可直接重跑
		{squad.StepStateCompleted, squad.StepStateRepairing, true},   // 仅可经显式重开进入修复
		{squad.StepStateSkipped, squad.StepStateExecuting, false},    // 跳过为终态
		{"", squad.StepStateExecuting, true},                         // 空态视为 pending
		{squad.StepStatePending, squad.StepStateVerifying, false},    // 不可跳过执行态
	}
	for _, c := range cases {
		if got := squad.ValidStepTransition(c.from, c.to); got != c.ok {
			t.Errorf("ValidStepTransition(%q→%q) = %v, want %v", c.from, c.to, got, c.ok)
		}
	}
}

func TestStepFSM_TransitionStep(t *testing.T) {
	wf := &squad.WorkflowPlan{Steps: []squad.StepSpec{{ID: "s1"}}}
	if err := squad.TransitionStep(wf, "s404", squad.StepStateExecuting); err == nil {
		t.Errorf("expected error for unknown step id")
	}
	if err := squad.TransitionStep(wf, "s1", squad.StepStateExecuting); err != nil {
		t.Fatalf("valid transition failed: %v", err)
	}
	if wf.Steps[0].Status != squad.StepStateExecuting {
		t.Errorf("status not updated: %s", wf.Steps[0].Status)
	}
	if err := squad.TransitionStep(wf, "s1", squad.StepStateExecuting); err == nil {
		t.Errorf("running→running should be rejected")
	}
}

// ============================================================================
// BudgetLedger 预算账本
// ============================================================================

func TestBudgetLedger_HardStopAndPersist(t *testing.T) {
	dir := t.TempDir()
	policy := squad.BudgetPolicy{MaxSteps: 2, MaxRepairRounds: 1, MaxAgentCalls: 3}
	l := squad.LoadOrCreateLedger(dir, policy)

	if err := l.ReserveStep(); err != nil {
		t.Fatalf("reserve 1: %v", err)
	}
	if err := l.ReserveStep(); err != nil {
		t.Fatalf("reserve 2: %v", err)
	}
	if err := l.ReserveStep(); err == nil {
		t.Errorf("expected hard stop at step 3")
	}
	if !l.Aborted {
		t.Errorf("expected aborted flag set")
	}
	if err := l.ReserveRepair(); err == nil {
		t.Errorf("aborted ledger must refuse further reserves")
	}

	if err := l.Save(dir); err != nil {
		t.Fatalf("save: %v", err)
	}
	loaded := squad.LoadOrCreateLedger(dir, policy)
	if !loaded.Aborted || loaded.StepExecutions != 2 {
		t.Errorf("ledger roundtrip failed: %+v", loaded)
	}
}

func TestBudgetLedger_SoftLimitIdempotent(t *testing.T) {
	policy := squad.BudgetPolicy{MaxSteps: 10, MaxRepairRounds: 2, MaxAgentCalls: 10}
	l := squad.LoadOrCreateLedger(t.TempDir(), policy)
	for i := 0; i < 8; i++ {
		if err := l.ReserveStep(); err != nil {
			t.Fatalf("reserve %d: %v", i, err)
		}
	}
	if msg := l.CheckSoftLimit(); msg == "" {
		t.Errorf("expected soft warning at 80%%")
	}
	if msg := l.CheckSoftLimit(); msg != "" {
		t.Errorf("soft warning should fire exactly once, got: %s", msg)
	}
	if !strings.Contains(l.UsageSummary(), "步骤 8/10") {
		t.Errorf("usage summary wrong: %s", l.UsageSummary())
	}
}

// ============================================================================
// Trace 轨迹 + 相似度
// ============================================================================

func TestTrace_AppendAndLoad(t *testing.T) {
	dir := t.TempDir()
	_ = squad.AppendTrace(dir, squad.TraceEntry{Seq: 1, StepID: "s1", Role: "organizer", Timestamp: time.Now(), PromptDigest: "p1", OutputDigest: "o1"})
	_ = squad.AppendTrace(dir, squad.TraceEntry{Seq: 2, StepID: "s1", Role: "organizer", Timestamp: time.Now(), PromptDigest: "p1", OutputDigest: "o1", SimilarityToPrev: 1.0})
	entries := squad.LoadTrace(dir)
	if len(entries) != 2 {
		t.Fatalf("expected 2 entries, got %d", len(entries))
	}
	if entries[1].SimilarityToPrev != 1.0 {
		t.Errorf("similarity not persisted: %v", entries[1].SimilarityToPrev)
	}
}

func TestOutputSimilarity(t *testing.T) {
	a := "掌上超声设备采用双频探头设计，功耗控制在5W以内，支持无线传输与电池供电，整机重量不超过800克。"
	if s := squad.OutputSimilarity(a, a); s != 1.0 {
		t.Errorf("identical text similarity = %v, want 1.0", s)
	}
	c := "完全无关的另一段文字，讲述洁净车间的压差梯度设计与悬浮粒子沉降菌监测规程的要求。"
	if s := squad.OutputSimilarity(a, c); s > 0.3 {
		t.Errorf("disjoint text similarity too high: %v", s)
	}
	d := a + "另外补充一句关于散热设计的说明。"
	if s := squad.OutputSimilarity(a, d); s < 0.5 || s >= 1.0 {
		t.Errorf("similar text similarity unexpected: %v", s)
	}
	if s := squad.OutputSimilarity("", ""); s != 1.0 {
		t.Errorf("both empty should be 1.0, got %v", s)
	}
}

// ============================================================================
// LoopDetector 重复识别
// ============================================================================

func mkTraceEntry(stepID, prompt, output string, success bool) squad.TraceEntry {
	return squad.TraceEntry{StepID: stepID, PromptDigest: prompt, OutputDigest: output, Success: success}
}

func TestDetectRepetition_RepeatStallOscillation(t *testing.T) {
	// 原地打转：连续两次完全一致
	repeat := []squad.TraceEntry{mkTraceEntry("s1", "p", "o", true), mkTraceEntry("s1", "p", "o", true)}
	if v := squad.DetectRepetition(repeat, "s1"); !v.Detected || v.Kind != "repeat" {
		t.Errorf("expected repeat verdict, got %+v", v)
	} else if v.Hint == "" {
		t.Errorf("repeat verdict should carry anti-repetition hint")
	}

	// 收敛停滞：相似度 0.99
	stall := []squad.TraceEntry{
		mkTraceEntry("s1", "p", "oA", true),
		{StepID: "s1", PromptDigest: "p", OutputDigest: "oA2", SimilarityToPrev: 0.99, Success: true},
	}
	if v := squad.DetectRepetition(stall, "s1"); !v.Detected || v.Kind != "stall" {
		t.Errorf("expected stall verdict, got %+v", v)
	}

	// 振荡：A→B→A
	osc := []squad.TraceEntry{mkTraceEntry("s1", "p", "A", true), mkTraceEntry("s1", "p", "B", true), mkTraceEntry("s1", "p", "A", true)}
	if v := squad.DetectRepetition(osc, "s1"); !v.Detected || v.Kind != "oscillation" {
		t.Errorf("expected oscillation verdict, got %+v", v)
	}
}

func TestDetectRepetition_NoFalsePositive(t *testing.T) {
	// 正常演进（产出不同）
	normal := []squad.TraceEntry{mkTraceEntry("s1", "p", "o1", true), mkTraceEntry("s1", "p", "o2", true)}
	if v := squad.DetectRepetition(normal, "s1"); v.Detected {
		t.Errorf("normal evolution should not trigger: %+v", v)
	}
	// 失败后重试（前次失败）不应误报
	failRetry := []squad.TraceEntry{mkTraceEntry("s1", "p", "o", false), mkTraceEntry("s1", "p", "o", true)}
	if v := squad.DetectRepetition(failRetry, "s1"); v.Detected {
		t.Errorf("failed retry should not trigger: %+v", v)
	}
	// 其他步骤的轨迹不应影响判定
	if v := squad.DetectRepetition(normal, "s9"); v.Detected {
		t.Errorf("unknown step should not trigger: %+v", v)
	}
	// 单条轨迹不判定
	if v := squad.DetectRepetition(normal[:1], "s1"); v.Detected {
		t.Errorf("single entry should not trigger: %+v", v)
	}
}

// ============================================================================
// StepManifest 归档锁定 + Driver 治理集成
// ============================================================================

func prepareGovWorkspace(t *testing.T, steps []squad.StepSpec) string {
	t.Helper()
	dir := t.TempDir()
	engine := squad.NewWorkspaceEngine()
	_, err := engine.PrepareWorkspace(squad.WorkspacePrepareRequest{
		WorkDir:       dir,
		TaskKey:       "intake",
		StageTitle:    "用户初始需求收集表",
		ProjectName:   "治理测试项目",
		CustomerName:  "测试客户",
		Backend:       "stub",
		WorkflowSteps: steps,
	})
	if err != nil {
		t.Fatalf("PrepareWorkspace failed: %v", err)
	}
	return dir
}

func TestFreezeStep_RefuseEmptyAndReopen(t *testing.T) {
	dir := prepareGovWorkspace(t, []squad.StepSpec{{ID: "s1", Role: "organizer", Title: "整理", SubWorkspace: "tasks/01_organize"}})
	step := squad.StepSpec{ID: "s1", Role: "organizer", SubWorkspace: "tasks/01_organize"}

	// 空子工作区（不存在）→ 拒绝归档
	emptyStep := squad.StepSpec{ID: "s_empty", Role: "organizer", SubWorkspace: "tasks/99_missing"}
	if _, err := squad.FreezeStep(dir, emptyStep); err == nil {
		t.Errorf("expected empty-archive refusal for missing sub workspace")
	}

	// 缺少 sub_workspace → 拒绝归档
	if _, err := squad.FreezeStep(dir, squad.StepSpec{ID: "s_nosub"}); err == nil {
		t.Errorf("expected refusal for missing sub_workspace")
	}

	// 有产物的步骤可正常归档
	_ = os.WriteFile(filepath.Join(dir, "tasks", "01_organize", "output.md"), []byte("# 产出"), 0o644)
	m, err := squad.FreezeStep(dir, step)
	if err != nil {
		t.Fatalf("freeze: %v", err)
	}
	if len(m.Files) == 0 {
		t.Errorf("expected files in step manifest")
	}
	if _, err := os.Stat(filepath.Join(dir, "tasks", "01_organize", "step_manifest.json")); err != nil {
		t.Errorf("step manifest not written: %v", err)
	}

	// 显式重开：completed+Frozen → repairing，审计入账
	wf := &squad.WorkflowPlan{Steps: []squad.StepSpec{{ID: "s1", Status: squad.StepStateCompleted, Frozen: true}}}
	if err := squad.ReopenStep(dir, wf, "s1", "需要修订"); err != nil {
		t.Fatalf("reopen: %v", err)
	}
	if wf.Steps[0].Status != squad.StepStateRepairing || wf.Steps[0].Frozen || wf.Steps[0].ReopenCount != 1 {
		t.Errorf("reopen state wrong: %+v", wf.Steps[0])
	}
	// 未冻结的步骤不可重开
	wf2 := &squad.WorkflowPlan{Steps: []squad.StepSpec{{ID: "s1", Status: squad.StepStateCompleted}}}
	if err := squad.ReopenStep(dir, wf2, "s1", "x"); err == nil {
		t.Errorf("reopen on non-frozen step should fail")
	}
}

func TestWorkflowDriver_GovernanceOnExecute(t *testing.T) {
	steps := []squad.StepSpec{
		{ID: "step_01", Role: "organizer", Title: "整理", SubWorkspace: "tasks/01_organize"},
		{ID: "step_02", Role: "writer", Title: "成文", SubWorkspace: "tasks/02_draft"},
	}
	dir := prepareGovWorkspace(t, steps)
	driver := squad.NewWorkflowDriver(squad.NewWorkspaceEngine())

	res, err := driver.Execute(context.Background(), dir, &agent.StubBackend{}, squad.ExecuteOptions{Timeout: 10 * time.Second})
	if err != nil || !res.Success {
		t.Fatalf("execute failed: %v, reason=%s", err, res.ErrorReason)
	}

	// 1. 每步归档：manifest 文件 + workflow frozen 标记
	wf, _ := driver.LoadWorkflow(dir)
	for _, s := range wf.Steps {
		if !s.Frozen {
			t.Errorf("step %s not frozen after completion", s.ID)
		}
		if s.Status != squad.StepStateCompleted {
			t.Errorf("step %s status = %s", s.ID, s.Status)
		}
		if _, err := os.Stat(filepath.Join(dir, filepath.FromSlash(s.SubWorkspace), "step_manifest.json")); err != nil {
			t.Errorf("step manifest missing for %s: %v", s.ID, err)
		}
	}

	// 2. 预算账本入账
	ledger := squad.LoadOrCreateLedger(dir, squad.DefaultBudgetPolicy())
	if ledger.StepExecutions != 2 || ledger.AgentCalls != 2 {
		t.Errorf("ledger counts wrong: %s", ledger.UsageSummary())
	}

	// 3. 轨迹入账
	if entries := squad.LoadTrace(dir); len(entries) != 2 {
		t.Errorf("expected 2 trace entries, got %d", len(entries))
	}

	// 4. 二次 Execute：冻结步骤被跳过（不重跑、不重复入账）
	res2, err := driver.Execute(context.Background(), dir, &agent.StubBackend{}, squad.ExecuteOptions{Timeout: 10 * time.Second})
	if err != nil || !res2.Success {
		t.Fatalf("re-execute failed: %v, reason=%s", err, res2.ErrorReason)
	}
	if res2.StepsRun != 0 {
		t.Errorf("frozen steps should be skipped on re-run, ran %d", res2.StepsRun)
	}
	if entries := squad.LoadTrace(dir); len(entries) != 2 {
		t.Errorf("trace should not grow on frozen skip, got %d", len(entries))
	}
}

func TestWorkflowDriver_BudgetHardStopKeepsPartialArchive(t *testing.T) {
	steps := []squad.StepSpec{
		{ID: "step_01", Role: "organizer", Title: "整理", SubWorkspace: "tasks/01_organize"},
		{ID: "step_02", Role: "writer", Title: "成文", SubWorkspace: "tasks/02_draft"},
		{ID: "step_03", Role: "leader", Title: "终审", SubWorkspace: "tasks/03_review"},
	}
	dir := prepareGovWorkspace(t, steps)
	driver := squad.NewWorkflowDriver(squad.NewWorkspaceEngine())

	// 预算收紧为 2 步 → 第 3 步触发硬停
	tight := squad.BudgetPolicy{MaxSteps: 2, MaxRepairRounds: 1, MaxAgentCalls: 10}
	res, err := driver.Execute(context.Background(), dir, &agent.StubBackend{}, squad.ExecuteOptions{Timeout: 10 * time.Second, BudgetPolicy: &tight})
	if err == nil || res.Success {
		t.Fatalf("expected budget hard stop, got success")
	}
	if !strings.Contains(res.ErrorReason, "预算硬停") {
		t.Errorf("unexpected abort reason: %s", res.ErrorReason)
	}

	// 已完成的 2 步保留归档（partial output 有价值）
	wf, _ := driver.LoadWorkflow(dir)
	frozen := 0
	for _, s := range wf.Steps {
		if s.Frozen {
			frozen++
		}
	}
	if frozen != 2 {
		t.Errorf("expected 2 frozen partial steps, got %d", frozen)
	}
	// 账本持久化中止状态
	ledger := squad.LoadOrCreateLedger(dir, tight)
	if !ledger.Aborted {
		t.Errorf("ledger should record abort")
	}
}