Agent-Squad/squad/steering_test.go

package squad

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

	"agent-native/exec"
)

// recordingBackend 记录每次收到的 prompt（按目录），产出按 script。
type recordingBackend struct {
	mu      sync.Mutex
	prompts map[string][]string
	script  map[string]map[string]string
	tokens  bool
}

func (b *recordingBackend) Name() string { return "rec" }

func (b *recordingBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	key := filepath.Base(filepath.Clean(req.WorkDir))
	b.mu.Lock()
	b.prompts[key] = append(b.prompts[key], req.Prompt)
	b.mu.Unlock()
	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)
	}
	pt, ct := 0, 0
	if b.tokens {
		pt, ct = 120, 80
	}
	return &exec.Result{Success: true, Patches: patches, PromptTokens: pt, CompletionTokens: ct}, nil
}

func TestSteeringInjectedAtNextStep(t *testing.T) {
	// 步骤1 产生非阻塞挂起点：用 gap 挂起，steer 后 respond → 步骤2 收到注入。
	rec := &recordingBackend{prompts: map[string][]string{}, script: map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "fields.json": `{}`,
			"gaps.json": `[{"item":"缺信息","blocking":true}]`},
		"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: "rec"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("rec", rec)
	engine, _ := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner(gapSteps()...), Backends: backends, Sink: NopSink{}})

	res, _ := engine.Run(context.Background(), Goal{Text: "g"}, RunOptions{})
	if res.Status != RunWaitingHuman {
		t.Fatalf("want held for steering window, got %s", res.Status)
	}
	if err := engine.Steer(res.RunID, "补充：结论部分要加一节成本分析"); err != nil {
		t.Fatal(err)
	}
	// 终态运行拒绝 steer。
	if err := engine.Steer("run_nope", "x"); err == nil {
		t.Fatal("steer unknown run must fail")
	}
	res2, err := engine.Respond(context.Background(), res.RunID, RespondDecision{RequestID: res.Pending.ID, Choice: "resolved", Answer: "ok"})
	if err != nil {
		t.Fatal(err)
	}
	if res2.Status != RunCompleted {
		t.Fatalf("want completed, got %s err=%s", res2.Status, res2.Error)
	}
	rec.mu.Lock()
	writerPrompts := rec.prompts["02_writer"]
	reviewerPrompts := rec.prompts["03_reviewer"]
	rec.mu.Unlock()
	if len(writerPrompts) == 0 || !strings.Contains(writerPrompts[0], "成本分析") {
		t.Fatalf("steering not injected into next step: %q", truncate(writerPrompts[0], 100))
	}
	if len(reviewerPrompts) > 0 && strings.Contains(reviewerPrompts[0], "成本分析") {
		t.Fatalf("steering must be consumed once, leaked to later step")
	}
}

func TestUsageAggregatedPerRole(t *testing.T) {
	rec := &recordingBackend{prompts: map[string][]string{}, tokens: true,
		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: "rec"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("rec", rec)
	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 != RunCompleted {
		t.Fatalf("want completed, got %s", res.Status)
	}
	if res.Usage == nil || res.Usage["writer"].PromptTokens != 120 || res.Usage["writer"].Steps != 1 {
		t.Fatalf("usage wrong: %+v", res.Usage)
	}
	// usage 持久化可恢复。
	data, err := os.ReadFile(filepath.Join(res.WorkDir, "usage.json"))
	if err != nil || !strings.Contains(string(data), "researcher") {
		t.Fatalf("usage.json missing: %v", err)
	}
}

// slowCountBackend 统计并发峰值（验证角色级并发上限）。
type slowCountBackend struct {
	mu    sync.Mutex
	cur   int
	peak  int
	delay time.Duration
}

func (b *slowCountBackend) Name() string { return "slow" }

func (b *slowCountBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	b.mu.Lock()
	b.cur++
	if b.cur > b.peak {
		b.peak = b.cur
	}
	b.mu.Unlock()
	time.Sleep(b.delay)
	b.mu.Lock()
	b.cur--
	b.mu.Unlock()
	for _, t := range req.Targets {
		out := filepath.Join(req.WorkDir, t)
		_ = os.MkdirAll(filepath.Dir(out), 0o755)
		_ = os.WriteFile(out, []byte(goodReport), 0o644)
	}
	return &exec.Result{Success: true, Patches: map[string][]byte{}}, nil
}

func TestRoleConcurrencyLimit(t *testing.T) {
	// 两个并行步骤同角色，MaxConcurrentTasks=1 → 串行（峰值 1）。
	slow := &slowCountBackend{delay: 80 * time.Millisecond}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	_ = reg.Register(RoleSpec{Role: "writer", Backend: "slow", MaxConcurrentTasks: 1})
	backends := NewBackendResolver()
	backends.RegisterInstance("slow", slow)
	engine, _ := NewEngine(Options{DataDir: dir, Registry: reg,
		Planner: mockScriptPlanner([]StepSpec{
			{ID: "s1", Role: "writer", Title: "A", Parallel: true, Targets: []string{"a.md"}},
			{ID: "s2", Role: "writer", Title: "B", Parallel: true, Targets: []string{"b.md"}},
		}...), Backends: backends, Sink: ProgressFunc(func(ev Event) {
			t.Logf("EVENT %s step=%s text=%q", ev.Type, ev.StepID, ev.Text)
		})})
	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)
	}
	slow.mu.Lock()
	peak := slow.peak
	slow.mu.Unlock()
	if peak != 1 {
		t.Fatalf("role concurrency peak want 1, got %d", peak)
	}
}