Agent-Squad/squad/attempt_test.go

package squad

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

// perStepRecordingBackend 按子工作区归因计数（steps 的 role→tasks/NN_role 布局）。
type perStepRecordingBackend struct {
	mu    sync.Mutex
	calls map[string]int
}

func (b *perStepRecordingBackend) Name() string { return "recording" }

func (b *perStepRecordingBackend) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	b.mu.Lock()
	if b.calls == nil {
		b.calls = map[string]int{}
	}
	b.calls[filepath.Base(req.WorkDir)]++
	b.mu.Unlock()
	out := map[string][]byte{}
	for _, t := range req.Targets {
		out[t] = []byte(goodReport)
	}
	if _, err := MaterializeTargets(req.WorkDir, req.Targets, out); err != nil {
		return &ExecResult{Success: false, ErrorReason: err.Error()}, nil
	}
	return &ExecResult{Success: true, Logs: "ok"}, nil
}

func (b *perStepRecordingBackend) countOf(sub string) int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.calls[sub]
}

func newAttemptTestEngine(t *testing.T, backend *perStepRecordingBackend, dir string) *SquadEngine {
	t.Helper()
	reg := NewRoleRegistry()
	for _, role := range []string{"researcher", "writer", "reviewer"} {
		if err := reg.Register(RoleSpec{Role: role, Backend: "recording"}); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("recording", backend)
	eng, err := NewEngine(Options{
		DataDir: dir, Registry: reg, Planner: mockScriptPlanner(baseSteps()...),
		Backends: backends, Sink: NopSink{},
	})
	if err != nil {
		t.Fatal(err)
	}
	return eng
}

func flipToInterrupted(t *testing.T, workDir string, stepIdx int) {
	t.Helper()
	st, err := loadRunState(workDir)
	if err != nil {
		t.Fatal(err)
	}
	st.Meta.Status = string(RunExecuting)
	st.Meta.CompletedAt = nil
	wf := st.Workflow()
	wf.Status = "running"
	wf.Steps[stepIdx].Status = "running"
	wf.Steps[stepIdx].Attempts = 1
	st.setWorkflow(wf)
	if err := st.save(); err != nil {
		t.Fatal(err)
	}
}

// TestAttemptIDFormat：attempt id 稳定落盘且形如 runID/stepID/N（Q1 WP3）。
func TestAttemptIDFormat(t *testing.T) {
	backend := &perStepRecordingBackend{}
	dir := t.TempDir()
	eng := newAttemptTestEngine(t, backend, dir)
	res, err := eng.Run(context.Background(), Goal{Text: "格式"}, RunOptions{})
	if err != nil || res.Status != RunCompleted {
		t.Fatalf("run: %v %+v", err, res)
	}
	if res.Plan == nil {
		t.Fatal("no plan")
	}
	for _, s := range res.Plan.Steps {
		want := res.RunID + "/" + s.ID + "/1"
		if s.AttemptID != want {
			t.Fatalf("step %s attempt id = %q, want %q", s.ID, s.AttemptID, want)
		}
	}
}

// TestKilledAtSuperstepBoundarySideEffectsOnce：杀进程于 superstep 边界后恢复，
// 已完成步骤不重新执行（执行次数=1）；running 且产物齐全的步骤按 attempt 判定
// 完成、不 rerun——Q1 验收第一条（WP3）。
func TestKilledAtSuperstepBoundarySideEffectsOnce(t *testing.T) {
	backend := &perStepRecordingBackend{}
	dir := t.TempDir()
	eng := newAttemptTestEngine(t, backend, dir)
	res, err := eng.Run(context.Background(), Goal{Text: "先跑完再中断", HitlMode: HitlAuto}, RunOptions{})
	if err != nil || res.Status != RunCompleted {
		t.Fatalf("setup run: %v %+v", err, res)
	}
	before := map[string]int{}
	st, err := loadRunState(res.WorkDir)
	if err != nil {
		t.Fatal(err)
	}
	for _, s := range st.Plan.Workflow.Steps {
		before[s.ID] = backend.countOf(filepath.Base(filepath.FromSlash(s.SubWorkspace)))
	}

	// 第二步停在 running（产物其实在盘上——首次运行真实写出的）。
	flipToInterrupted(t, res.WorkDir, 1)

	eng2 := newAttemptTestEngine(t, backend, dir)
	res, err = eng2.Resume(context.Background(), res.RunID)
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("resume: %+v", res)
	}
	for _, s := range res.Plan.Steps {
		if s.Status != "completed" && s.Status != "skipped" {
			t.Fatalf("step %s not finished: %s", s.ID, s.Status)
		}
		got := backend.countOf(filepath.Base(filepath.FromSlash(s.SubWorkspace)))
		if got != before[s.ID] {
			t.Fatalf("step %s executed %d times after kill (before=%d)——已完成/running 步骤被重复执行", s.ID, got, before[s.ID])
		}
	}
}

// TestInterruptedMissingProductsRetries：running 崩溃且产物缺失时按 auto 一次重试（WP3）。
func TestInterruptedMissingProductsRetries(t *testing.T) {
	backend := &perStepRecordingBackend{}
	dir := t.TempDir()
	eng := newAttemptTestEngine(t, backend, dir)
	res, err := eng.Run(context.Background(), Goal{Text: "先跑完再中断", HitlMode: HitlAuto}, RunOptions{})
	if err != nil || res.Status != RunCompleted {
		t.Fatalf("setup run: %v %+v", err, res)
	}
	// 第二步停 running 并删掉其交付物（副作用未发生的形态）。
	flipToInterrupted(t, res.WorkDir, 1)
	st, err := loadRunState(res.WorkDir)
	if err != nil {
		t.Fatal(err)
	}
	sub := filepath.FromSlash(st.Plan.Workflow.Steps[1].SubWorkspace)
	if err := os.Remove(filepath.Join(res.WorkDir, sub, "output.md")); err != nil {
		t.Fatal(err)
	}
	eng2 := newAttemptTestEngine(t, backend, dir)
	res, err = eng2.Resume(context.Background(), res.RunID)
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("resume: %+v", res)
	}
	got := backend.countOf(filepath.Base(sub))
	if got != 2 {
		t.Fatalf("missing-product step executed %d times, want 2（auto 一次重试）", got)
	}
}