Agent-Squad/squad/evaluator_test.go

package squad

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

// scriptedEvaluator 第一步后建议：追加一个补漏步骤 + 跳过冗余步骤。
type scriptedEvaluator struct{ fired bool }

func (e *scriptedEvaluator) Evaluate(ctx context.Context, req EvaluateRequest) (*EvaluationResult, error) {
	if req.Step.ID != "s1" || e.fired {
		return &EvaluationResult{StepID: req.Step.ID, Action: ActionContinue}, nil
	}
	e.fired = true
	return &EvaluationResult{
		StepID: req.Step.ID,
		Action: ActionAddStep,
		Reason: "首步产出暴露竞品价格缺口，需补一步价格核对；s3 与 s1 内容重复可跳过",
		NewAssumptions: []AssumptionItem{
			{ID: "asm_1", Topic: "price", Statement: "价格取公开电商口径", Source: "inferred"},
		},
		DynamicStepsToAdd: []StepSpec{
			{ID: "s_price", Role: "writer", Title: "竞品价格核对补充", Targets: []string{"price.md"}},
		},
		StepsToSkip:   []string{"s3"},
		NeedsApproval: true,
	}, nil
}

func dynamicSteps() []StepSpec {
	return []StepSpec{
		{ID: "s1", Role: "researcher", Title: "调研", AllowDynamicAdjustment: true, Targets: []string{"notes.md"}},
		{ID: "s2", Role: "writer", Title: "成文", Targets: []string{"output.md"}},
		{ID: "s3", Role: "reviewer", Title: "冗余复核", Targets: []string{"review.md"}},
	}
}

const priceNote = "价格核对补充：A 品牌售价 99 元，B 品牌 149 元，C 品牌 199 元，均取公开电商旗舰店口径，含税包邮。"

func newDynamicEngine(t *testing.T, script map[string]map[string]string, ev StepEvaluator) *SquadEngine {
	t.Helper()
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, r := range []string{"researcher", "writer", "reviewer"} {
		_ = reg.Register(RoleSpec{Role: r, Backend: "script"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: script})
	engine, err := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner(dynamicSteps()...), Backends: backends,
		Evaluator: ev, Sink: NopSink{}})
	if err != nil {
		t.Fatal(err)
	}
	return engine
}

// TestDynamicAdjustmentWithApproval 验证：步骤产出 → 评估器建议 →
// HITL 审批挂起 → 批准后新步骤插入执行、冗余步骤跳过、假设入黑板。
func TestDynamicAdjustmentWithApproval(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "# 笔记\n内容"},
		"02_writer":     {"output.md": goodReport},
		// 动态加步发生在子工作区构建之后，runOneStep 按索引兜底目录名。
		"03_writer":   {"price.md": priceNote},
		"03_reviewer": {"review.md": goodReport},
	}
	engine := newDynamicEngine(t, script, &scriptedEvaluator{})

	res, err := engine.Run(context.Background(), Goal{Text: "调研并成文"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunWaitingHuman || res.Pending == nil || res.Pending.Kind != HitlAdjustment {
		t.Fatalf("want adjustment hold, got %s pending=%v", res.Status, res.Pending)
	}
	if !strings.Contains(res.Pending.Context, "竞品价格核对补充") || !strings.Contains(res.Pending.Context, "s3") {
		t.Fatalf("adjustment context missing suggestions: %s", res.Pending.Context)
	}

	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: res.Pending.ID, Choice: "approve"})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted && res2.Status != RunDegraded {
		t.Fatalf("want completed/degraded after approve, got %s err=%s", res2.Status, res2.Error)
	}
	var added, skipped bool
	for _, s := range res2.Plan.Steps {
		if s.ID == "s_price" && s.Status == "completed" {
			added = true
		}
		if s.ID == "s3" && s.Status == "skipped" {
			skipped = true
		}
	}
	if !added || !skipped {
		t.Fatalf("adjustment not applied: added=%v skipped=%v steps=%+v", added, skipped, res2.Plan.Steps)
	}
	bb, _ := engine.Blackboard(res.RunID)
	if len(bb.Assumptions) == 0 || bb.Assumptions[0].Topic != "price" {
		t.Fatalf("assumption not recorded: %+v", bb.Assumptions)
	}
}

// TestDynamicAdjustmentRejected 验证：拒绝调整 → 忽略建议按原计划继续。
func TestDynamicAdjustmentRejected(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "# 笔记\n内容"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"review.md": goodReport},
	}
	engine := newDynamicEngine(t, script, &scriptedEvaluator{})

	res, _ := engine.Run(context.Background(), Goal{Text: "调研并成文"}, RunOptions{})
	if res.Status != RunWaitingHuman || res.Pending.Kind != HitlAdjustment {
		t.Fatalf("want adjustment hold, got %s", res.Status)
	}
	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: res.Pending.ID, Choice: "reject"})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted && res2.Status != RunDegraded {
		t.Fatalf("want completed/degraded after reject, got %s err=%s", res2.Status, res2.Error)
	}
	for _, s := range res2.Plan.Steps {
		if s.ID == "s_price" {
			t.Fatal("rejected adjustment must not insert step")
		}
	}
}

// mutableScriptBackend 可变脚本后端：第 N 次执行前可改写脚本。
type mutableScriptBackend struct {
	script   map[string]map[string]string
	mutateAt int
	mutate   func(map[string]map[string]string)
	calls    int
}

func (b *mutableScriptBackend) Name() string { return "mut" }

func (b *mutableScriptBackend) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	b.calls++
	if b.mutate != nil && b.calls == b.mutateAt {
		b.mutate(b.script)
	}
	key := filepath.Base(filepath.Clean(req.WorkDir))
	files := b.script[key]
	patches := map[string][]byte{}
	for _, tg := range req.Targets {
		content := files[tg]
		if content == "" {
			content = "\n"
		}
		out := filepath.Join(req.WorkDir, tg)
		_ = os.MkdirAll(filepath.Dir(out), 0o755)
		_ = os.WriteFile(out, []byte(content), 0o644)
		patches[tg] = []byte(content)
	}
	return &ExecResult{Success: true, Patches: patches}, nil
}

// TestQualityRepairAutoRerun 验证：质检失败 → 自动重跑交付步骤 → 通过。
func TestQualityRepairAutoRerun(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n"},
		"02_writer":     {"output.md": "# 报告\n\n{{待填充}}\n"},
	}
	mut := &mutableScriptBackend{script: script, mutateAt: 2,
		mutate: func(s map[string]map[string]string) {
			s["02_writer"] = map[string]string{"output.md": goodReport}
		}}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, r := range []string{"researcher", "writer"} {
		_ = reg.Register(RoleSpec{Role: r, Backend: "mut"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("mut", mut)
	engine, _ := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner(
			StepSpec{ID: "s1", Role: "researcher", Title: "调研", Targets: []string{"notes.md"}},
			StepSpec{ID: "s2", Role: "writer", Title: "成文", Targets: []string{"output.md"}},
		), Backends: backends, Sink: NopSink{}})

	res, err := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("auto repair should complete, got %s err=%s", res.Status, res.Error)
	}
	if res.Quality == nil || !res.Quality.Passed || len(res.Quality.Placeholders) != 0 {
		t.Fatalf("final quality should pass without placeholders: %+v", res.Quality)
	}
	data, _ := os.ReadFile(filepath.Join(res.WorkDir, "output.md"))
	if !strings.Contains(string(data), "测试交付物") {
		t.Fatalf("repaired deliverable wrong: %q", string(data)[:min(len(data), 40)])
	}
}

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