Agent-System/internal/squad/step_retry_test.go

package squad_test

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

	"chat-agent-system/internal/agent"
	"chat-agent-system/internal/squad"
)

type countingBackend struct {
	mu       sync.Mutex
	calls    int
	mode     string
	failText string
}

func (b *countingBackend) Name() string { return "counting" }

func (b *countingBackend) Execute(ctx context.Context, req agent.Request) (*agent.Result, error) {
	b.mu.Lock()
	b.calls++
	n := b.calls
	mode := b.mode
	failText := b.failText
	b.mu.Unlock()

	switch mode {
	case "transient-then-ok":
		if n == 1 {
			_ = os.WriteFile(filepath.Join(req.WorkDir, "execution.log"), []byte(
				`{"type":"auto_retry_end","success":false,"finalError":"Connection error."}`+"\n"+
					`{"type":"agent_settled"}`+"\n",
			), 0o644)
			return &agent.Result{Success: false, ErrorReason: "Connection error.", Retryable: true}, nil
		}
		body := []byte("# recovered draft\n\nfresh step invocation\n")
		_ = os.WriteFile(filepath.Join(req.WorkDir, "output.md"), body, 0o644)
		return &agent.Result{Success: true, Patches: map[string][]byte{"output.md": body}}, nil

	case "old-runner-lie-then-ok":
		if n == 1 {
			_ = os.WriteFile(filepath.Join(req.WorkDir, "execution.log"), []byte(
				"Outcome: SUCCESS\nError: None\n"+
					`{"type":"message_end","message":{"stopReason":"error","errorMessage":"terminated"}}`+"\n"+
					`{"type":"auto_retry_end","success":false,"attempt":3,"finalError":"Connection error."}`+"\n"+
					`{"type":"agent_settled"}`+"\n",
			), 0o644)
			// Old runner: process 0 + settled + empty patches.
			return &agent.Result{Success: true, Patches: map[string][]byte{
				"execution.log": []byte("Outcome: SUCCESS\n"),
			}}, nil
		}
		body := []byte("# recovered after old-runner lie\n")
		_ = os.WriteFile(filepath.Join(req.WorkDir, "output.md"), body, 0o644)
		return &agent.Result{Success: true, Patches: map[string][]byte{"output.md": body}}, nil

	case "permanent-missing":
		return &agent.Result{Success: false, ErrorReason: "missing target deliverable output.md"}, nil

	case "settled-no-deliverable-then-ok":
		if n == 1 {
			_ = os.WriteFile(filepath.Join(req.WorkDir, "execution.log"), []byte(
				`{"type":"agent_settled"}`+"\n",
			), 0o644)
			return &agent.Result{Success: true, Patches: map[string][]byte{
				"execution.log": []byte(`{"type":"agent_settled"}` + "\n"),
			}}, nil
		}
		if !strings.Contains(req.Prompt, "纠偏重试") || !strings.Contains(req.Prompt, "write") {
			return &agent.Result{Success: false, ErrorReason: "corrective prompt missing on retry"}, nil
		}
		body := []byte("# recovered after settled-no-deliverable\n")
		_ = os.WriteFile(filepath.Join(req.WorkDir, "output.md"), body, 0o644)
		return &agent.Result{Success: true, Patches: map[string][]byte{"output.md": body}}, nil

	case "always-transient":
		return &agent.Result{Success: false, ErrorReason: failText, Retryable: true}, nil

	case "fail-after-writing-target":
		body := []byte("# already on disk\n\ndo not overwrite\n")
		_ = os.WriteFile(filepath.Join(req.WorkDir, "output.md"), body, 0o644)
		return &agent.Result{Success: false, ErrorReason: "Connection error.", Retryable: true}, nil

	case "success-without-patch":
		return &agent.Result{Success: true, Patches: map[string][]byte{
			"execution.log": []byte("Outcome: SUCCESS\n"),
		}}, nil

	default:
		return &agent.Result{Success: false, ErrorReason: "unknown mode"}, nil
	}
}

func prepareSingleWriterWS(t *testing.T) (string, *squad.WorkflowDriver) {
	t.Helper()
	tmpDir := t.TempDir()
	engine := squad.NewWorkspaceEngine(nil)
	ws, err := engine.PrepareWorkspace(squad.WorkspacePrepareRequest{
		WorkDir:     tmpDir,
		TaskKey:     "prelim_design",
		StageTitle:  "系统初步设计方案",
		ProjectName: "retry-audit",
		Backend:     "pi-agent",
		WorkflowSteps: []squad.StepSpec{
			{
				ID:           "step_01_draft",
				Role:         "writer",
				Title:        "撰写初稿",
				SubWorkspace: "tasks/01_draft",
				Mode:         "execute",
			},
		},
	})
	if err != nil {
		t.Fatalf("PrepareWorkspace: %v", err)
	}
	return ws.WorkDir, squad.NewWorkflowDriver(engine)
}

func TestWorkflowDriver_RetriesTransientSettledFailure(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	backend := &countingBackend{mode: "transient-then-ok"}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err != nil {
		t.Fatalf("transient then recover must succeed: %v", err)
	}
	if !res.Success {
		t.Fatalf("want success, got %s", res.ErrorReason)
	}
	if backend.calls != 2 {
		t.Fatalf("want one fresh re-invocation, calls=%d", backend.calls)
	}
	body, err := os.ReadFile(filepath.Join(wsDir, "tasks", "01_draft", "output.md"))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(body), "fresh step invocation") {
		t.Fatalf("recovered body missing: %s", body)
	}
	entries := squad.LoadTrace(wsDir)
	if len(entries) < 2 {
		t.Fatalf("each invocation must be traced, got %d", len(entries))
	}
	if entries[0].Success {
		t.Fatal("first settled-failure invocation must not be recorded as success")
	}
	if !entries[1].Success {
		t.Fatal("recovered invocation must be success")
	}
}

func TestWorkflowDriver_RetriesOldRunnerLieWhenTraceShowsSettledFailure(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	backend := &countingBackend{mode: "old-runner-lie-then-ok"}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err != nil {
		t.Fatalf("old-runner lie + transient trace must retry: %v", err)
	}
	if !res.Success {
		t.Fatalf("want success after fresh invocation, got %s", res.ErrorReason)
	}
	if backend.calls != 2 {
		t.Fatalf("want fresh invocation, calls=%d", backend.calls)
	}
	if _, err := os.Stat(filepath.Join(wsDir, "tasks", "01_draft", "output.md")); err != nil {
		t.Fatal("must not invent output from truncated write args; recovered file should come from second invocation")
	}
}

func TestWorkflowDriver_DoesNotRetryPermanentMissing(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	backend := &countingBackend{mode: "permanent-missing"}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err == nil || res.Success {
		t.Fatal("permanent missing must fail-closed")
	}
	if backend.calls != 1 {
		t.Fatalf("permanent missing must not be retried, calls=%d", backend.calls)
	}
	if !strings.Contains(err.Error(), "missing target deliverable output.md") {
		t.Fatalf("reason: %v", err)
	}
	if _, serr := os.Stat(filepath.Join(wsDir, "tasks", "01_draft", "output.md")); serr == nil {
		t.Fatal("must not forge output.md")
	}
}

func TestWorkflowDriver_RetriesSettledNoDeliverableWithCorrectivePrompt(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	backend := &countingBackend{mode: "settled-no-deliverable-then-ok"}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err != nil {
		t.Fatalf("settled-no-deliverable must retry with corrective prompt: %v", err)
	}
	if !res.Success {
		t.Fatalf("want success after corrective retry, got %s", res.ErrorReason)
	}
	if backend.calls != 2 {
		t.Fatalf("want one corrective re-invocation, calls=%d", backend.calls)
	}
	body, err := os.ReadFile(filepath.Join(wsDir, "tasks", "01_draft", "output.md"))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(body), "settled-no-deliverable") {
		t.Fatalf("recovered body missing: %s", body)
	}
}

func TestWorkflowDriver_DoesNotOverwriteExistingTarget(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	backend := &countingBackend{mode: "fail-after-writing-target"}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err == nil || res.Success {
		t.Fatal("backend explicit failure must fail-closed; a possibly truncated file is not success")
	}
	if backend.calls != 1 {
		t.Fatalf("must not re-invoke over this-invocation target, calls=%d", backend.calls)
	}
	if !strings.Contains(err.Error(), "Connection error.") {
		t.Fatalf("want backend failure preserved, got %v", err)
	}
	body, err := os.ReadFile(filepath.Join(wsDir, "tasks", "01_draft", "output.md"))
	if err != nil {
		t.Fatal(err)
	}
	if string(body) != "# already on disk\n\ndo not overwrite\n" {
		t.Fatalf("target was rewritten: %s", body)
	}
}

func TestWorkflowDriver_PriorResidueIsNotSuccess(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	leftover := []byte("# leftover from previous round\n")
	if err := os.WriteFile(filepath.Join(wsDir, "tasks", "01_draft", "output.md"), leftover, 0o644); err != nil {
		t.Fatal(err)
	}
	backend := &countingBackend{mode: "success-without-patch"}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err == nil || res.Success {
		t.Fatal("prior-round residue must not count as this-invocation success")
	}
	if backend.calls != 1 {
		t.Fatalf("residue without trusted patch is a permanent contract miss, calls=%d", backend.calls)
	}
	got, err := os.ReadFile(filepath.Join(wsDir, "tasks", "01_draft", "output.md"))
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != string(leftover) {
		t.Fatalf("residue must stay on disk unchanged: %s", got)
	}
}

func TestWorkflowDriver_BoundedTransientRetryThenFailClosed(t *testing.T) {
	wsDir, driver := prepareSingleWriterWS(t)
	backend := &countingBackend{mode: "always-transient", failText: "Connection error."}
	res, err := driver.Execute(context.Background(), wsDir, backend, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err == nil || res.Success {
		t.Fatal("exhausted transient retries must fail-closed")
	}
	if backend.calls != agent.MaxTransientStepInvocations {
		t.Fatalf("calls=%d want bound %d", backend.calls, agent.MaxTransientStepInvocations)
	}
	if _, serr := os.Stat(filepath.Join(wsDir, "tasks", "01_draft", "output.md")); serr == nil {
		t.Fatal("must not forge output.md after exhausted retries")
	}
}