Agent-Squad/squad/recovery_test.go

package squad

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

	"agent-native/exec"
)

func approvalScript() map[string]map[string]string {
	return map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "fields.json": `{}`},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
}

func TestStepApprovalApproveExecutes(t *testing.T) {
	steps := baseSteps()
	steps[0].RequiresApproval = true
	eng, _ := newTestEngine(t, mockScriptPlanner(steps...), approvalScript(), nil)
	res, err := eng.Run(context.Background(), Goal{Text: "需要审批"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunAwaitingConfirmation || res.Pending == nil || res.Pending.Reason != "step_approval" {
		t.Fatalf("want step approval hold, got %s %+v", res.Status, res.Pending)
	}
	res, err = eng.Respond(context.Background(), res.RunID, RespondDecision{Choice: "approve"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("approve should execute, got %s %s", res.Status, res.Error)
	}
}

func TestStepApprovalSkipDoesNotCancel(t *testing.T) {
	steps := baseSteps()
	steps[0].RequiresApproval = true
	eng, _ := newTestEngine(t, mockScriptPlanner(steps...), approvalScript(), nil)
	res, err := eng.Run(context.Background(), Goal{Text: "跳过一步"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	res, err = eng.Respond(context.Background(), res.RunID, RespondDecision{Choice: "skip"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status == RunCancelled || res.Status == RunFailed {
		t.Fatalf("skip must not cancel the run, got %s %s", res.Status, res.Error)
	}
	if res.Plan == nil || len(res.Plan.Steps) == 0 || res.Plan.Steps[0].Status != "skipped" {
		t.Fatalf("first step should be skipped: %+v", res.Plan)
	}
}

func TestApprovalSurvivesRestart(t *testing.T) {
	steps := baseSteps()
	steps[0].RequiresApproval = true
	eng, dir := newTestEngine(t, mockScriptPlanner(steps...), approvalScript(), nil)
	res, err := eng.Run(context.Background(), Goal{Text: "授权要留下"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := eng.Respond(context.Background(), res.RunID, RespondDecision{Choice: "approve"}); err != nil {
		t.Fatal(err)
	}
	st, err := loadRunState(res.WorkDir)
	if err != nil {
		t.Fatal(err)
	}
	if !st.Approved["step:s1"] {
		t.Fatalf("approval not stored: %+v", st.Approved)
	}
	st.Meta.Status = string(RunExecuting)
	st.Meta.CompletedAt = nil
	wf := st.Workflow()
	wf.Status = "running"
	wf.Steps[0].Status = "pending"
	st.setWorkflow(wf)
	if err := st.save(); err != nil {
		t.Fatal(err)
	}
	eng2, err := NewEngine(Options{
		DataDir: dir, Registry: eng.Registry(), Planner: mockScriptPlanner(steps...),
		Backends: eng.opts.Backends, Sink: NopSink{},
	})
	if err != nil {
		t.Fatal(err)
	}
	res, err = eng2.Resume(context.Background(), res.RunID)
	if err != nil {
		t.Fatal(err)
	}
	if res.Status == RunAwaitingConfirmation {
		t.Fatal("restart asked for approval again")
	}
	if res.Status != RunCompleted {
		t.Fatalf("resume after restart: %s %s", res.Status, res.Error)
	}
}

type countBackend struct {
	mu sync.Mutex
	n  int
}

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

func (b *countBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	b.mu.Lock()
	b.n++
	b.mu.Unlock()
	return (&scriptBackend{script: approvalScript()}).Execute(ctx, req)
}

func (b *countBackend) calls() int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.n
}

func TestInterruptedRunningStepStaysHeld(t *testing.T) {
	backend := &countBackend{}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, role := range []string{"researcher", "writer", "reviewer"} {
		if err := reg.Register(RoleSpec{Role: role, Backend: "script"}); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", backend)
	eng, err := NewEngine(Options{
		DataDir: dir, Registry: reg, Planner: mockScriptPlanner(baseSteps()...),
		Backends: backends, Sink: NopSink{},
	})
	if err != nil {
		t.Fatal(err)
	}
	res, err := eng.Run(context.Background(), Goal{Text: "先跑完再假装中断"}, RunOptions{})
	if err != nil || res.Status != RunCompleted {
		t.Fatalf("setup run: %v %+v", err, res)
	}
	before := backend.calls()
	st, err := loadRunState(res.WorkDir)
	if err != nil {
		t.Fatal(err)
	}
	st.Meta.Status = string(RunExecuting)
	st.Meta.CompletedAt = nil
	wf := st.Workflow()
	wf.Status = "running"
	wf.Steps[0].Status = "running"
	wf.Steps[0].Attempts = 1
	st.setWorkflow(wf)
	if err := st.save(); err != nil {
		t.Fatal(err)
	}
	eng2, err := NewEngine(Options{DataDir: dir, Registry: reg, Planner: eng.opts.Planner, Backends: backends, Sink: NopSink{}})
	if err != nil {
		t.Fatal(err)
	}
	res, err = eng2.Resume(context.Background(), res.RunID)
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunAwaitingConfirmation || res.Pending == nil || res.Pending.Reason != "interrupted" {
		t.Fatalf("want interrupted hold, got %s %+v", res.Status, res.Pending)
	}
	if backend.calls() != before {
		t.Fatalf("interrupted step executed again: calls %d -> %d", before, backend.calls())
	}
}

type blockBackend struct {
	started chan struct{}
	release chan struct{}
	once    sync.Once
}

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

func (b *blockBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	b.once.Do(func() { close(b.started) })
	select {
	case <-b.release:
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	return (&scriptBackend{script: approvalScript()}).Execute(ctx, req)
}

func TestTwoEnginesOneRunIsBusy(t *testing.T) {
	backend := &blockBackend{started: make(chan struct{}), release: make(chan struct{})}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, role := range []string{"researcher", "writer", "reviewer"} {
		_ = reg.Register(RoleSpec{Role: role, Backend: "script"})
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", backend)
	mk := func() *SquadEngine {
		t.Helper()
		eng, err := NewEngine(Options{
			DataDir: dir, Registry: reg, Planner: mockScriptPlanner(baseSteps()...),
			Backends: backends, Sink: NopSink{},
		})
		if err != nil {
			t.Fatal(err)
		}
		return eng
	}
	engA := mk()
	done := make(chan struct{})
	go func() {
		defer close(done)
		_, _ = engA.Run(context.Background(), Goal{Text: "占住锁"}, RunOptions{})
	}()
	select {
	case <-backend.started:
	case <-time.After(3 * time.Second):
		t.Fatal("run did not start")
	}
	runs, err := engA.ListRuns()
	if err != nil || len(runs) != 1 {
		t.Fatalf("runs: %v %+v", err, runs)
	}
	engB := mk()
	ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
	defer cancel()
	_, err = engB.Resume(ctx, runs[0].ID)
	if !errors.Is(err, ErrRunBusy) {
		t.Fatalf("want busy, got %v", err)
	}
	close(backend.release)
	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("holder did not finish")
	}
}

func TestWakeupDoesNotApprovePermission(t *testing.T) {
	var kinds []string
	sink := ProgressFunc(func(ev Event) { kinds = append(kinds, ev.Type) })
	steps := gapSteps()
	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},
	}
	eng, dir := newTestEngine(t, mockScriptPlanner(steps...), script, sink)
	res, err := eng.Run(context.Background(), Goal{Text: "权限不能被唤醒放行"}, RunOptions{})
	if err != nil || res.Status != RunWaitingHuman {
		t.Fatalf("setup: %v %+v", err, res)
	}
	gate, err := loadGate(filepath.Join(dir, "runs", res.RunID))
	if err != nil || gate == nil {
		t.Fatal(err)
	}
	gate.Request.Kind = HitlPermission
	gate.Request.WaitCondition = &WaitCondition{Type: "field_set", Field: "ext"}
	if err := writeJSON(filepath.Join(dir, "runs", res.RunID, "pending.json"), gate); err != nil {
		t.Fatal(err)
	}
	if err := eng.SetBlackboardField(res.RunID, "ext", "ready"); err != nil {
		t.Fatal(err)
	}
	if n := eng.SweepWakeups(context.Background()); n != 0 {
		t.Fatalf("permission wakeup must not count as resolved, got %d", n)
	}
	insp, err := eng.Inspect(res.RunID)
	if err != nil {
		t.Fatal(err)
	}
	if insp.Status != RunWaitingHuman || insp.Pending == nil || insp.Pending.Kind != HitlPermission {
		t.Fatalf("permission hold changed: %s %+v", insp.Status, insp.Pending)
	}
	saw := false
	for _, k := range kinds {
		if k == "wakeup_held" {
			saw = true
		}
	}
	if !saw {
		t.Fatalf("expected wakeup_held event, got %v", kinds)
	}
}

func TestParallelGroupSecondStepApprovalHolds(t *testing.T) {
	steps := baseSteps()
	steps[0].Parallel = true
	steps[1].Parallel = true
	steps[1].RequiresApproval = true
	eng, _ := newTestEngine(t, mockScriptPlanner(steps...), approvalScript(), nil)
	res, err := eng.Run(context.Background(), Goal{Text: "并行组第二步要审批"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunAwaitingConfirmation || res.Pending == nil || res.Pending.StepID != "s2" {
		t.Fatalf("group should hold on the second step, got %s %+v", res.Status, res.Pending)
	}
	if _, err := os.Stat(filepath.Join(res.WorkDir, "tasks", "01_researcher", "notes.md")); !os.IsNotExist(err) {
		t.Fatalf("parallel group ran before approval: %v", err)
	}
}

func TestParallelGroupSecondStepPermissionHolds(t *testing.T) {
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, spec := range []RoleSpec{
		{Role: "researcher", Name: "调研", Backend: "script"},
		{Role: "writer", Name: "撰写", Backend: "script", BashPermission: BashPermissionAsk},
		{Role: "reviewer", Name: "审核", Backend: "script"},
	} {
		if err := reg.Register(spec); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: approvalScript()})
	steps := baseSteps()
	steps[0].Parallel = true
	steps[1].Parallel = true
	eng, err := NewEngine(Options{
		DataDir: dir, Registry: reg, Planner: mockScriptPlanner(steps...), Backends: backends, Sink: NopSink{},
	})
	if err != nil {
		t.Fatal(err)
	}
	var res *RunResult
	res, err = eng.Run(context.Background(), Goal{Text: "并行组第二步要授权"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunWaitingHuman || res.Pending == nil || res.Pending.Kind != HitlPermission || res.Pending.StepID != "s2" {
		t.Fatalf("group should hold on writer permission, got %s %+v", res.Status, res.Pending)
	}
	if _, err := os.Stat(filepath.Join(res.WorkDir, "tasks", "01_researcher", "notes.md")); !os.IsNotExist(err) {
		t.Fatalf("parallel group ran before permission: %v", err)
	}
}

func TestFreeTextKeepsChoiceHold(t *testing.T) {
	for _, kind := range []HumanRequestKind{HitlConfirmation, HitlPermission, HitlQualityReview} {
		t.Run(string(kind), func(t *testing.T) {
			steps := baseSteps()
			steps[0].RequiresApproval = true
			eng, dir := newTestEngine(t, mockScriptPlanner(steps...), approvalScript(), nil)
			res, err := eng.Run(context.Background(), Goal{Text: "自由文本不能取消"}, RunOptions{})
			if err != nil || res.Pending == nil {
				t.Fatalf("setup: %v %+v", err, res)
			}
			work := filepath.Join(dir, "runs", res.RunID)
			gate, err := loadGate(work)
			if err != nil || gate == nil {
				t.Fatal(err)
			}
			id := gate.Request.ID
			gate.Request.Kind = kind
			if err := writeJSON(filepath.Join(work, "pending.json"), gate); err != nil {
				t.Fatal(err)
			}
			res, err = eng.Respond(context.Background(), res.RunID, RespondDecision{Answer: "补充一句"})
			if err != nil {
				t.Fatal(err)
			}
			if res.Pending == nil || res.Pending.ID != id || res.Pending.Status != "pending" {
				t.Fatalf("hold was cleared: %+v", res.Pending)
			}
			data, err := os.ReadFile(filepath.Join(work, "steering.json"))
			if err != nil || !strings.Contains(string(data), "补充一句") {
				t.Fatalf("steering not recorded: %v %s", err, data)
			}
		})
	}
}

func TestFailedDecisionRestoresPending(t *testing.T) {
	steps := gapSteps()
	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},
	}
	eng, _ := newTestEngine(t, mockScriptPlanner(steps...), script, nil)
	res, err := eng.Run(context.Background(), Goal{Text: "错误选择后仍可回应"}, RunOptions{})
	if err != nil || res.Pending == nil {
		t.Fatalf("setup: %v %+v", err, res)
	}
	if _, err := eng.Respond(context.Background(), res.RunID, RespondDecision{Choice: "nope"}); err == nil {
		t.Fatal("unknown choice should fail")
	}
	res, err = eng.Respond(context.Background(), res.RunID, RespondDecision{Choice: "resolved"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("restored hold should still be answerable, got %s %s", res.Status, res.Error)
	}
}