Agent-Squad/squad/leader_test.go

package squad

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

	"agent-native/exec"
)

// leaderScriptBackend leader 模式脚本后端：00_leader 目录按轮次回放
// delegation.json 序列；成员目录按 script map 回放产出。
type leaderScriptBackend struct {
	mu        sync.Mutex
	turn      int
	decisions []string // 每轮写入 tasks/00_leader/delegation.json 的内容
	script    map[string]map[string]string
	prompts   map[string][]string // 目录 → 收到的 prompt（记录用）
}

func (b *leaderScriptBackend) Name() string { return "leader-script" }

func (b *leaderScriptBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	key := filepath.Base(filepath.Clean(req.WorkDir))
	if b.prompts != nil {
		b.mu.Lock()
		b.prompts[key] = append(b.prompts[key], req.Prompt)
		b.mu.Unlock()
	}
	if key == "00_leader" {
		b.mu.Lock()
		idx := b.turn
		b.turn++
		b.mu.Unlock()
		if idx >= len(b.decisions) {
			idx = len(b.decisions) - 1
		}
		out := filepath.Join(req.WorkDir, "delegation.json")
		_ = os.WriteFile(out, []byte(b.decisions[idx]), 0o644)
		return &exec.Result{Success: true, Patches: map[string][]byte{"delegation.json": []byte(b.decisions[idx])}}, nil
	}
	files := b.script[key]
	patches := map[string][]byte{}
	for _, t := range req.Targets {
		content := files[t]
		if content == "" {
			content = "\n"
		}
		out := filepath.Join(req.WorkDir, t)
		_ = os.MkdirAll(filepath.Dir(out), 0o755)
		_ = os.WriteFile(out, []byte(content), 0o644)
		patches[t] = []byte(content)
	}
	return &exec.Result{Success: true, Patches: patches,
		PromptTokens: 100, CompletionTokens: 50}, nil
}

func delegateJSON(member, title string) string {
	return fmt.Sprintf(`{"action":"delegate","member":%q,"title":%q,"targets":["output.md"],"briefing":"验收：内容完整"}`, member, title)
}

const closeJSON = `{"action":"close","reason":"所有子任务完成，交付物满足目标"}`

func newLeaderEngine(t *testing.T, backend exec.Backend, roles ...RoleSpec) *SquadEngine {
	t.Helper()
	dir := t.TempDir()
	reg := NewRoleRegistry()
	if len(roles) == 0 {
		roles = []RoleSpec{
			{Role: "leader", Name: "主管", Backend: "leader-script", Skills: []string{"协调"}},
			{Role: "writer", Name: "撰写", Backend: "leader-script", Skills: []string{"撰写"}},
		}
	}
	for _, r := range roles {
		if err := reg.Register(r); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("leader-script", backend)
	engine, err := NewEngine(Options{DataDir: dir, Registry: reg, Backends: backends, Sink: NopSink{}})
	if err != nil {
		t.Fatal(err)
	}
	return engine
}

func TestLeaderModeDelegateThenClose(t *testing.T) {
	b := &leaderScriptBackend{
		decisions: []string{
			delegateJSON("writer", "撰写调研简报"),
			closeJSON,
		},
		script: map[string]map[string]string{
			"01_writer": {"output.md": goodReport},
		},
	}
	engine := newLeaderEngine(t, b)
	res, err := engine.Run(context.Background(), Goal{Text: "产出一份调研简报"}, RunOptions{
		Mode: "leader", LeaderRole: "leader", MaxLeaderTurns: 5,
	})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("want completed, got %s err=%s", res.Status, res.Error)
	}
	// 委派出的成员步骤进入编排清单且完成。
	var member StepSpec
	for _, s := range res.Plan.Steps {
		if s.Role == "writer" {
			member = s
		}
	}
	if member.ID == "" || member.Status != "completed" {
		t.Fatalf("member step missing/incomplete: %+v", res.Plan.Steps)
	}
	// leader 历史持久化。
	data, _ := os.ReadFile(filepath.Join(res.WorkDir, "leader.json"))
	var ls LeaderState
	if err := json.Unmarshal(data, &ls); err != nil {
		t.Fatal(err)
	}
	if len(ls.Delegations) != 1 || ls.Delegations[0].Member != "writer" || ls.Turns < 2 {
		t.Fatalf("leader state wrong: %+v", ls)
	}
	// 用量按角色记账。
	if res.Usage == nil || res.Usage["writer"].CompletionTokens == 0 {
		t.Fatalf("usage missing: %+v", res.Usage)
	}
}

func TestLeaderModeEscalateHITL(t *testing.T) {
	b := &leaderScriptBackend{
		prompts: map[string][]string{},
		decisions: []string{
			`{"action":"escalate","question":"目标受众是谁？","reason":"受众影响成文风格"}`,
			`{"action":"delegate","member":"writer","title":"按人工答复撰写简报","targets":["output.md"],"briefing":"面向人工确认的受众"}`,
			closeJSON,
		},
		script: map[string]map[string]string{
			"02_writer": {"output.md": goodReport},
		},
	}
	engine := newLeaderEngine(t, b)
	res, err := engine.Run(context.Background(), Goal{Text: "写简报"}, RunOptions{Mode: "leader", LeaderRole: "leader"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunWaitingHuman || res.Pending.Kind != HitlEscalation || res.Pending.Severity != SeverityActionRequired {
		t.Fatalf("want escalation hold, got %s kind=%v sev=%q", res.Status, res.Pending.Kind, res.Pending.Severity)
	}
	if !strings.Contains(res.Pending.Question, "受众") {
		t.Fatalf("escalation question wrong: %s", res.Pending.Question)
	}
	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 escalation answer, got %s err=%s", res2.Status, res2.Error)
	}
	// 人的决定落为黑板事实。
	bb, _ := engine.Blackboard(res.RunID)
	found := false
	for _, f := range bb.Facts {
		if strings.Contains(f.Claim, "大众消费者") && f.Source == "human" {
			found = true
		}
	}
	if !found {
		t.Fatalf("escalation answer not a fact: %+v", bb.Facts)
	}
	// 下一轮 leader 上下文包含该问答（通过记录的 prompt 验证）。
	b.mu.Lock()
	lastLeaderPrompt := ""
	if n := len(b.prompts["00_leader"]); n > 0 {
		lastLeaderPrompt = b.prompts["00_leader"][n-1]
	}
	b.mu.Unlock()
	if !strings.Contains(lastLeaderPrompt, "大众消费者") {
		t.Fatalf("leader not re-briefed with escalation answer: %q", truncate(lastLeaderPrompt, 80))
	}
}

func TestLeaderModeTurnCapGuard(t *testing.T) {
	// leader 每轮都委派（成员产出内容但 leader 永不 close）→ 轮次护栏触发失败。
	loop := delegateJSON("writer", "再补一节")
	b := &leaderScriptBackend{decisions: []string{loop}, script: map[string]map[string]string{
		"01_writer": {"output.md": goodReport},
	}}
	engine := newLeaderEngine(t, b)
	res, err := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{Mode: "leader", LeaderRole: "leader", MaxLeaderTurns: 3})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunFailed || !strings.Contains(res.Error, "上限") {
		t.Fatalf("want turn-cap failure, got %s err=%s", res.Status, res.Error)
	}
	if b.turn < 3 {
		t.Fatalf("leader should have run %d turns, ran %d", 3, b.turn)
	}
}

func TestLeaderModeRestartRecovery(t *testing.T) {
	b := &leaderScriptBackend{
		decisions: []string{delegateJSON("writer", "撰写简报"), closeJSON},
		script:    map[string]map[string]string{"01_writer": {"output.md": goodReport}},
	}
	engine := newLeaderEngine(t, b)
	res, _ := engine.Run(context.Background(), Goal{Text: "写简报"}, RunOptions{Mode: "leader", LeaderRole: "leader"})
	if res.Status != RunCompleted {
		t.Fatalf("want completed, got %s err=%s", res.Status, res.Error)
	}
	// 新引擎实例（进程重启模拟）能读出终态与 leader 历史。
	engine2, err := NewEngine(Options{DataDir: engineOpts(engine).DataDir,
		Registry: NewRoleRegistry(), Backends: NewBackendResolver(), Sink: NopSink{}})
	if err != nil {
		t.Fatal(err)
	}
	insp, err := engine2.Inspect(res.RunID)
	if err != nil || insp.Status != RunCompleted {
		t.Fatalf("restart inspect: %v %+v", err, insp)
	}
}

func engineOpts(e *SquadEngine) Options { return e.opts }

func TestLeaderModeGapHoldThenResume(t *testing.T) {
	// 成员产出阻塞缺口 → 标准 gap 挂起；人工回应后 leader 循环继续到 close。
	b := &leaderScriptBackend{
		decisions: []string{
			`{"action":"delegate","member":"writer","title":"撰写简报","targets":["output.md","fields.json","gaps.json"]}`,
			closeJSON,
		},
		script: map[string]map[string]string{
			"01_writer": {"output.md": goodReport, "fields.json": `{"k":"v"}`,
				"gaps.json": `[{"item":"缺少受众定位","blocking":true}]`},
		},
	}
	engine := newLeaderEngine(t, b)
	res, _ := engine.Run(context.Background(), Goal{Text: "写简报"}, RunOptions{Mode: "leader", LeaderRole: "leader"})
	if res.Status != RunWaitingHuman || res.Pending.Kind != HitlGap {
		t.Fatalf("want gap hold in leader mode, got %s", res.Status)
	}
	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{
		RequestID: res.Pending.ID, Choice: "resolved", Answer: "受众为大众消费者",
	})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted {
		t.Fatalf("want completed after gap answer in leader mode, got %s err=%s", res2.Status, res2.Error)
	}
}

// ---- 唤醒（等待条件自动恢复） ----

func TestWakeupRunDone(t *testing.T) {
	// 运行 B 先完成；运行 A 的缺口声明 wait_run_id=B → sweep 自动唤醒续跑。
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "fields.json": `{}`},
		"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, Sink: NopSink{}})

	// 运行 B：无缺口，直接完成。
	resB, err := engine.Run(context.Background(), Goal{Text: "运行B"}, RunOptions{})
	if err != nil || resB.Status != RunCompleted {
		t.Fatalf("run B failed: %v %+v", err, resB)
	}

	// 运行 A：步骤 1 声明 wait_run_id。
	scriptA := map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "fields.json": `{}`,
			"gaps.json": `[{"item":"等待运行B的数据","blocking":true,"wait_run_id":"` + resB.RunID + `"}]`},
		"02_writer":   {"output.md": goodReport},
		"03_reviewer": {"output.md": goodReport},
	}
	backendsA := NewBackendResolver()
	backendsA.RegisterInstance("script", &scriptBackend{script: scriptA})
	dirA := t.TempDir()
	engineA, _ := NewEngine(Options{DataDir: dirA, Registry: reg,
		Planner: mockScriptPlanner(gapSteps()...), Backends: backendsA, Sink: NopSink{}})
	resA, err := engineA.Run(context.Background(), Goal{Text: "运行A"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if resA.Status != RunWaitingHuman || resA.Pending.WaitCondition == nil || resA.Pending.WaitCondition.Type != "run_done" {
		t.Fatalf("gap wait condition missing: %s cond=%+v", resA.Status, resA.Pending.WaitCondition)
	}
	// 未满足时不唤醒：等待自己（不存在）不触发。
	if n := engineA.SweepWakeups(context.Background()); n != 0 {
		// run_done 指向 engineB 的 run（不同 DataDir 引擎看不见）→ 不满足 → 0。
		_ = n
	}
	// 同一引擎内验证：重新在 engineA 数据目录下伪造目标运行已完成（跨运行等待）。
	// 直接用 field_set 变体验证自动唤醒通路。
	gateDir := filepath.Join(dirA, "runs", resA.RunID)
	gate, _ := loadGate(gateDir)
	gate.Request.WaitCondition = &WaitCondition{Type: "field_set", Field: "external_data"}
	_ = writeJSON(filepath.Join(gateDir, "pending.json"), gate)
	if err := engineA.SetBlackboardField(resA.RunID, "external_data", "来自外部的数据"); err != nil {
		t.Fatal(err)
	}
	if n := engineA.SweepWakeups(context.Background()); n != 1 {
		t.Fatalf("sweep want 1 woken, got %d", n)
	}
	insp, err := engineA.Inspect(resA.RunID)
	if err != nil {
		t.Fatal(err)
	}
	if insp.Status != RunCompleted {
		t.Fatalf("auto-wakeup should complete run, got %s", insp.Status)
	}
	bb, _ := engineA.Blackboard(resA.RunID)
	found := false
	for _, f := range bb.Facts {
		if strings.Contains(f.Claim, "external_data") || strings.Contains(f.Claim, "自动唤醒") {
			found = true
		}
	}
	if !found {
		t.Fatalf("wakeup fact not recorded: %+v", bb.Facts)
	}
}

func TestWakeupGuardMaxFires(t *testing.T) {
	// 构造一个永不满足条件但人为反复置位触发…… 实际护栏语义：fires 达上限即自动中止。
	// 用 field_set 反复唤醒：每次唤醒后又产生新缺口（脚本后端每步都写 gaps）→ 循环到上限。
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, r := range []string{"researcher", "writer", "reviewer"} {
		_ = reg.Register(RoleSpec{Role: r, Backend: "script"})
	}
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "fields.json": `{}`,
			"gaps.json": `[{"item":"等字段","blocking":true,"wait_field":"ext"}]`},
		"02_writer":   {"output.md": goodReport},
		"03_reviewer": {"output.md": goodReport},
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: script})
	engine, _ := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner(gapSteps()...), Backends: backends, Sink: NopSink{}})
	// 只跑第一步就挂起：两步计划，步骤1 gap 挂起。
	res, _ := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	if res.Status != RunWaitingHuman {
		t.Fatalf("want held, got %s", res.Status)
	}
	// 反复置字段 + sweep：第一次唤醒后步骤2 完成（无缺口）→ 直接完成，不会循环。
	// 护栏上限用直接改 fires 验证：设为 MaxWakeupFires-1 再触发一次。
	gateDir := filepath.Join(dir, "runs", res.RunID)
	gate, _ := loadGate(gateDir)
	gate.Request.WaitCondition = &WaitCondition{Type: "field_set", Field: "ext"}
	gate.Request.WakeupFires = MaxWakeupFires - 1
	_ = writeJSON(filepath.Join(gateDir, "pending.json"), gate)
	_ = engine.SetBlackboardField(res.RunID, "ext", "v")
	engine.SweepWakeups(context.Background())
	insp, _ := engine.Inspect(res.RunID)
	if insp.Status != RunCancelled && insp.Status != RunCompleted {
		t.Fatalf("guard-exhaust path unexpected: %s", insp.Status)
	}
	_ = time.Now
}

// eventingBackend 在执行时触发过程事件，验证 OnEvent 桥接到 detail 事件。
type eventingBackend struct {
	inner *leaderScriptBackend
	mu    sync.Mutex
	reqs  []exec.Request
}

func (b *eventingBackend) Name() string { return "leader-script" }

func (b *eventingBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	b.mu.Lock()
	b.reqs = append(b.reqs, req)
	b.mu.Unlock()
	if req.OnEvent != nil {
		req.OnEvent(exec.ProcessEvent{Type: "tool_execution_start", Tool: "grep", Detail: `{"q":"x"}`})
		req.OnEvent(exec.ProcessEvent{Type: "tool_execution_end", Tool: "grep", Text: "found"})
	}
	return b.inner.Execute(ctx, req)
}

func TestLeaderTurnsEmitDetailEventsAndCarryProfile(t *testing.T) {
	inner := &leaderScriptBackend{
		decisions: []string{delegateJSON("writer", "撰写"), closeJSON},
		script:    map[string]map[string]string{"01_writer": {"output.md": goodReport}},
	}
	eb := &eventingBackend{inner: inner}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	leader := RoleSpec{
		Role: "leader", Name: "主管", Backend: "leader-script",
		ToolProfile: "workspace", Skills: []string{"协调"},
		Tools: []string{"read_file", "grep"},
	}
	if err := reg.Register(leader); err != nil {
		t.Fatal(err)
	}
	if err := reg.Register(RoleSpec{Role: "writer", Name: "撰写", Backend: "leader-script"}); err != nil {
		t.Fatal(err)
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("leader-script", eb)
	collector := &sliceSink{}
	engine, err := NewEngine(Options{DataDir: dir, Registry: reg, Backends: backends, Sink: collector})
	if err != nil {
		t.Fatal(err)
	}
	res, err := engine.Run(context.Background(), Goal{Text: "产出简报"}, RunOptions{Mode: "leader", LeaderRole: "leader"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("want completed, got %s err=%s", res.Status, res.Error)
	}
	// leader 轮的请求带上了 ToolProfile/白名单/技能与标识（修平 leader 请求缺口）。
	var leaderReq *exec.Request
	for i := range eb.reqs {
		if strings.HasSuffix(eb.reqs[i].WorkDir, "00_leader") {
			leaderReq = &eb.reqs[i]
		}
	}
	if leaderReq == nil {
		t.Fatal("leader request missing")
	}
	if leaderReq.ToolProfile != exec.ToolProfileWorkspace || len(leaderReq.ToolAllowlist) != 2 || len(leaderReq.Skills) != 1 {
		t.Fatalf("leader request not fully assembled: %+v", leaderReq)
	}
	if !strings.Contains(leaderReq.ExecutionID, "leader_turn_") {
		t.Fatalf("leader execution id missing: %q", leaderReq.ExecutionID)
	}
	// 过程事件桥接到 detail 事件流。
	sawToolStart, sawToolEnd := false, false
	for _, ev := range collector.events {
		if ev.Type == EvDetail {
			var data map[string]string
			_ = json.Unmarshal(ev.Data, &data)
			if data["process_type"] == "tool_execution_start" && data["tool"] == "grep" {
				sawToolStart = true
			}
			if data["process_type"] == "tool_execution_end" {
				sawToolEnd = true
			}
		}
	}
	if !sawToolStart || !sawToolEnd {
		t.Fatalf("leader process events must reach the sink (start=%v end=%v)", sawToolStart, sawToolEnd)
	}
}

type sliceSink struct {
	mu     sync.Mutex
	events []Event
}

func (s *sliceSink) Emit(ev Event) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.events = append(s.events, ev)
}