Agent-Squad/squad/engine_test.go

package squad

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// scriptBackend 测试用脚本后端：按子目录名回放产出。
type scriptBackend struct{ script map[string]map[string]string }

func (b *scriptBackend) Name() string { return "script" }

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

const goodReport = "# 测试交付物\n\n## 概述\n\n这是由脚本后端生成的合格交付物，长度超过五十个字符以满足最小质检要求。\n\n## 内容\n\n正文内容。\n"

func newTestEngine(t *testing.T, planner Planner, script map[string]map[string]string, sink ProgressSink) (*SquadEngine, string) {
	t.Helper()
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, spec := range []RoleSpec{
		{Role: "researcher", Name: "调研", Backend: "script"},
		{Role: "writer", Name: "撰写", Backend: "script"},
		{Role: "reviewer", Name: "审核", Backend: "script"},
	} {
		if err := reg.Register(spec); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: script})
	if sink == nil {
		sink = NopSink{}
	}
	engine, err := NewEngine(Options{
		DataDir: dir, Registry: reg, Planner: planner, Backends: backends, Sink: sink,
	})
	if err != nil {
		t.Fatal(err)
	}
	return engine, dir
}

func mockScriptPlanner(steps ...StepSpec) *MockPlanner { return &MockPlanner{Script: steps} }

func baseSteps() []StepSpec {
	return []StepSpec{
		{ID: "s1", Role: "researcher", Title: "调研", Targets: []string{"notes.md", "fields.json"}},
		{ID: "s2", Role: "writer", Title: "成文", Targets: []string{"output.md"}},
		{ID: "s3", Role: "reviewer", Title: "终审", Targets: []string{"output.md"}},
	}
}

func gapSteps() []StepSpec {
	steps := baseSteps()
	steps[0].Targets = []string{"notes.md", "fields.json", "gaps.json"}
	return steps
}

func TestEngineHappyPath(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "# 笔记\n内容", "fields.json": `{"k":"v"}`},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	engine, dir := newTestEngine(t, mockScriptPlanner(baseSteps()...), script, nil)
	res, err := engine.Run(context.Background(), Goal{Text: "生成测试报告"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("want completed, got %s (err=%s)", res.Status, res.Error)
	}
	if res.Deliverable != "output.md" {
		t.Fatalf("want output.md, got %s", res.Deliverable)
	}
	// 工作区布局。
	for _, f := range []string{"workspace.json", "squad.json", "workflow.json", "blackboard.json",
		"manifest.json", "output.md", "context/goal.md", "context/index.json", "tasks/01_researcher/TASK.md", "tasks/01_researcher/AGENTS.md"} {
		if _, err := os.Stat(filepath.Join(res.WorkDir, f)); err != nil {
			t.Errorf("workspace missing %s: %v", f, err)
		}
	}
	// 黑板字段合并。
	bb, _ := engine.Blackboard(res.RunID)
	if bb.Fields["k"] != "v" {
		t.Errorf("blackboard field k=%q, want v", bb.Fields["k"])
	}
	// 清单包含质检报告。
	data, _ := os.ReadFile(filepath.Join(res.WorkDir, "manifest.json"))
	var m Manifest
	if err := json.Unmarshal(data, &m); err != nil {
		t.Fatal(err)
	}
	if !m.QualityReport.Passed || len(m.Files) == 0 {
		t.Errorf("manifest quality/files invalid: %+v files=%d", m.QualityReport, len(m.Files))
	}
	_ = dir
}

func TestEngineGapHitlSuspendRespondAndRestartRecovery(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {
			"notes.md":    "# 笔记\n内容",
			"fields.json": `{"k":"v"}`,
			"gaps.json":   `[{"item":"缺少产量","blocking":true}]`,
		},
		"02_writer":   {"output.md": goodReport},
		"03_reviewer": {"output.md": goodReport},
	}
	engine, dir := newTestEngine(t, mockScriptPlanner(gapSteps()...), script, nil)
	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 != HitlGap {
		t.Fatalf("want waiting_human/gap, got %s pending=%v", res.Status, res.Pending)
	}
	if res.Plan.Steps[0].Status != "completed" || res.Plan.Steps[1].Status != "pending" {
		t.Fatalf("step states wrong: %+v", res.Plan.Steps)
	}

	// 进程重启恢复：用同一 DataDir 新建引擎实例续跑。
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: script})
	engine2, err := NewEngine(Options{DataDir: dir, Registry: NewRoleRegistry(), Backends: backends, Sink: NopSink{}})
	if err != nil {
		t.Fatal(err)
	}
	pend, err := engine2.Pending(res.RunID)
	if err != nil || pend == nil {
		t.Fatalf("pending not recovered: %v %v", pend, err)
	}
	res2, err := engine2.Respond(context.Background(), res.RunID, RespondDecision{
		RequestID: res.Pending.ID, Choice: "resolved", Answer: "年产量 10 万套",
	})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted {
		t.Fatalf("want completed after respond, got %s (err=%s)", res2.Status, res2.Error)
	}
	bb, _ := engine2.Blackboard(res.RunID)
	found := false
	for _, g := range bb.Gaps {
		if g.Item == "缺少产量" && g.Resolved {
			found = true
		}
	}
	if !found {
		t.Errorf("blocking gap not resolved: %+v", bb.Gaps)
	}
	if len(bb.Facts) == 0 || !strings.Contains(bb.Facts[0].Claim, "10 万套") {
		t.Errorf("human answer not recorded as fact: %+v", bb.Facts)
	}
}

func TestEngineClarificationFlow(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	// 先澄清后规划成功：用可变 mock。
	mp := &MockPlanner{ClarifyReason: "目标缺少范围"}
	engine, _ := newTestEngine(t, mp, script, nil)
	res, err := engine.Run(context.Background(), Goal{Text: "做一份计划"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunAwaitingClarification || res.Pending.Kind != HitlClarification {
		t.Fatalf("want awaiting_clarification, got %s", res.Status)
	}
	// 人工澄清后 mock 切换为脚本计划。
	mp.ClarifyReason = ""
	mp.Script = baseSteps()
	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: res.Pending.ID, Answer: "范围限定为华东市场"})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted {
		t.Fatalf("want completed after clarification, got %s err=%s", res2.Status, res2.Error)
	}
	// 澄清答案并入目标（工作区 goal.md 可追溯）。
	data, _ := os.ReadFile(filepath.Join(res2.WorkDir, "context", "goal.md"))
	if !strings.Contains(string(data), "华东市场") {
		t.Errorf("clarification answer not recorded in goal.md:\n%s", data)
	}
}

func TestEnginePermissionAsk(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	if err := reg.Register(RoleSpec{Role: "researcher", Name: "调研", Backend: "script", BashPermission: "ask"}); err != nil {
		t.Fatal(err)
	}
	for _, r := range []string{"writer", "reviewer"} {
		_ = reg.Register(RoleSpec{Role: r, Backend: "script"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: script})
	engine, _ := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner(baseSteps()...), Backends: backends, Sink: NopSink{}})

	res, err := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunWaitingHuman || res.Pending.Kind != HitlPermission {
		t.Fatalf("want waiting_human/permission, got %s", res.Status)
	}
	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 {
		t.Fatalf("want completed after approve, got %s err=%s", res2.Status, res2.Error)
	}
	// 拒绝路径：关键步骤失败。
	engine2, _ := NewEngine(Options{DataDir: t.TempDir(), Registry: reg,
		Planner: mockScriptPlanner(baseSteps()...), Backends: backends, Sink: NopSink{}})
	res3, _ := engine2.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	res4, err := engine2.Respond(context.Background(), res3.RunID, RespondDecision{RequestID: res3.Pending.ID, Choice: "reject"})
	if err != nil {
		t.Fatal(err)
	}
	if res4.Status != RunFailed {
		t.Fatalf("want failed after reject, got %s", res4.Status)
	}
}

type failChecker struct{}

func (failChecker) Check(req CheckRequest) *QualityReport {
	return &QualityReport{Passed: false, Score: 10, Violations: []string{"内容不达标"}, CheckedAt: time.Now()}
}

func TestEngineQualityReviewManualPass(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	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, _ := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner(baseSteps()...), Backends: backends,
		Checker: failChecker{}, Sink: NopSink{}})

	res, _ := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	if res.Status != RunQualityReview || res.Pending.Kind != HitlQualityReview {
		t.Fatalf("want quality_review, got %s", res.Status)
	}
	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: res.Pending.ID, Choice: "approve", Answer: "内容已人工核对"})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted || res2.Quality == nil || !res2.Quality.ManualPass {
		t.Fatalf("want completed+manual_pass, got %s %+v", res2.Status, res2.Quality)
	}
}

func TestEngineAutoModeGapFallbackAssumption(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {
			"notes.md":    "# 笔记\n内容",
			"fields.json": `{"k":"v"}`,
			"gaps.json":   `[{"item":"缺少产量","blocking":true}]`,
		},
		"02_writer":   {"output.md": goodReport},
		"03_reviewer": {"output.md": goodReport},
	}
	engine, _ := newTestEngine(t, mockScriptPlanner(gapSteps()...), script, nil)
	res, err := engine.Run(context.Background(), Goal{Text: "g", HitlMode: HitlAuto}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("auto mode should complete with fallback assumption, got %s err=%s", res.Status, res.Error)
	}
	bb, _ := engine.Blackboard(res.RunID)
	if len(bb.Assumptions) == 0 || bb.Assumptions[0].Source != "fallback" {
		t.Fatalf("fallback assumption missing: %+v", bb.Assumptions)
	}
}

func TestEngineRulePlannerAutonomousAndEmptyRegistryClarify(t *testing.T) {
	// 空注册表 → 澄清。
	dir := t.TempDir()
	backends := NewBackendResolver()
	engine, _ := NewEngine(Options{DataDir: dir, Registry: NewRoleRegistry(),
		Planner: NewRulePlanner(), Backends: backends, Sink: NopSink{}})
	res, err := engine.Run(context.Background(), Goal{Text: "做点事"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunAwaitingClarification {
		t.Fatalf("empty registry should clarify, got %s", res.Status)
	}

	// 选"最小编制"：引导注册 4 个通用角色 → 角色齐备无派生 → 直接完成。
	script := map[string]map[string]string{
		"01_organizer":  {"notes.md": "n", "fields.json": `{}`},
		"02_specialist": {"notes.md": "n"},
		"03_writer":     {"output.md": goodReport},
		"04_leader":     {"output.md": goodReport},
	}
	backends.RegisterInstance("stub", &scriptBackend{script: script})
	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: res.Pending.ID, Choice: "minimal"})
	if err != nil {
		t.Fatal(err)
	}
	if roles := engine.Registry().List(); len(roles) != 4 {
		t.Fatalf("minimal squad want 4 roles, got %d: %+v", len(roles), roles)
	}
	if res2.Status != RunCompleted {
		t.Fatalf("minimal bootstrap should complete directly, got %s err=%s", res2.Status, res2.Error)
	}
}

func TestEngineAutonomousWithRolesNoConfirmWhenDirect(t *testing.T) {
	// 角色齐备（与四阶段键一致）时自主规划直接执行（无派生→无需确认）。
	script := map[string]map[string]string{
		"01_organizer":  {"notes.md": "n", "fields.json": `{}`},
		"02_specialist": {"notes.md": "n"},
		"03_writer":     {"output.md": goodReport},
		"04_leader":     {"output.md": goodReport},
	}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, r := range []string{"organizer", "specialist", "writer", "leader"} {
		_ = reg.Register(RoleSpec{Role: r, Backend: "stub"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("stub", &scriptBackend{script: script})
	engine, _ := NewEngine(Options{DataDir: dir, Registry: reg, Planner: NewRulePlanner(), Backends: backends, Sink: NopSink{}})
	res, err := engine.Run(context.Background(), Goal{Text: "整理并撰写报告"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("want completed without confirmation, got %s err=%s", res.Status, res.Error)
	}
}

func TestEngineBlueprintRegistrationFlow(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n"},
		"02_writer":     {"output.md": goodReport},
	}
	engine, _ := newTestEngine(t, NewRulePlanner(), script, nil)
	bp := Blueprint{
		Key: "report", Title: "标准报告", Match: []string{"报告"},
		Steps: []StepSpec{
			{ID: "b1", Role: "researcher", Title: "调研", Targets: []string{"notes.md"}},
			{ID: "b2", Role: "writer", Title: "成文", Targets: []string{"output.md"}},
		},
	}
	if err := engine.Blueprints().Register(bp); err != nil {
		t.Fatal(err)
	}
	res, err := engine.Run(context.Background(), Goal{Text: "帮我出一份竞品报告"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	// writer 是最后完成步骤，无 reviewer；应 completed。
	if res.Status != RunCompleted {
		t.Fatalf("blueprint run want completed, got %s err=%s", res.Status, res.Error)
	}
	if res.Plan.TaskKey != "report" || res.Plan.PlanningMode != PlanModeBlueprint {
		t.Fatalf("plan not from blueprint: key=%s mode=%s", res.Plan.TaskKey, res.Plan.PlanningMode)
	}
}

func TestEngineParallelStepsAndEvents(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	steps := []StepSpec{
		{ID: "s1", Role: "researcher", Title: "调研", Targets: []string{"notes.md"}},
		{ID: "s2", Role: "writer", Title: "成文A", Parallel: true, Targets: []string{"draft_a.md"}},
		{ID: "s3", Role: "reviewer", Title: "成文B", Parallel: true, Targets: []string{"output.md"}},
	}
	script["02_writer"] = map[string]string{"draft_a.md": goodReport}
	sink := NewFanoutSink(1024)
	engine, _ := newTestEngine(t, mockScriptPlanner(steps...), script, sink)
	res, err := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("want completed, got %s err=%s", res.Status, res.Error)
	}
	events, _ := sink.Snapshot()
	seen := map[string]int{}
	for _, ev := range events {
		seen[ev.Type]++
	}
	for _, want := range []string{EvRunStarted, EvStepStarted, EvStepCompleted, EvDeliverable, EvRunCompleted} {
		if seen[want] == 0 {
			t.Errorf("missing event %s (seen=%v)", want, seen)
		}
	}
}

func TestRespondIdempotencyAndMismatch(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "gaps.json": `[{"item":"x","blocking":true}]`},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	engine, _ := newTestEngine(t, mockScriptPlanner(gapSteps()...), script, nil)
	res, _ := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	// 错误请求 ID → 拒绝。
	if _, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: "hr_bogus", Choice: "resolved"}); err == nil {
		t.Fatal("bogus request id should fail")
	}
	// 无挂起运行 → 报错。
	if _, err := engine.Respond(context.Background(), "run_nope", RespondDecision{}); err == nil {
		t.Fatal("unknown run should fail")
	}
}