Agent-System/internal/squad/quality_repair_contract_test.go

package squad_test

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

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

const repairTemplate = "# 用户初始需求收集表\n\n## 1. 结论\n\n## 2. 详细方案\n"

const repairBadDraft = "# 用户初始需求收集表\n\n## 1. 结论\n由于篇幅有限，不做展开。\n"

const repairGoodDraft = `# 用户初始需求收集表

## 1. 结论
本项目面向专业医疗机构，采集与供电方案已按材料展开。

## 2. 详细方案
覆盖信号采集、电磁兼容与供电冗余，并标注出处〔材料: 需求表 §2〕。
`

func prepareRepairWS(t *testing.T) string {
	t.Helper()
	engine := squad.NewWorkspaceEngine()
	ws, err := engine.PrepareWorkspace(squad.WorkspacePrepareRequest{
		WorkDir:         t.TempDir(),
		TaskKey:         "intake",
		StageTitle:      "用户初始需求收集表",
		TemplateContent: repairTemplate,
	})
	if err != nil {
		t.Fatalf("PrepareWorkspace: %v", err)
	}
	if err := os.WriteFile(filepath.Join(ws.WorkDir, "output.md"), []byte(repairBadDraft), 0o644); err != nil {
		t.Fatal(err)
	}
	return ws.WorkDir
}

type localRepairBackend struct{ body string }

func (b *localRepairBackend) Name() string { return "local-repair" }
func (b *localRepairBackend) Execute(_ context.Context, req agent.Request) (*agent.Result, error) {
	return &agent.Result{Success: true, Patches: map[string][]byte{"output.md": []byte(b.body)}}, nil
}

type parentEscapeBackend struct{ body string }

func (b *parentEscapeBackend) Name() string { return "parent-escape" }
func (b *parentEscapeBackend) Execute(_ context.Context, req agent.Request) (*agent.Result, error) {
	parent := filepath.Clean(filepath.Join(req.WorkDir, "..", "..", "output.md"))
	_ = os.WriteFile(parent, []byte(b.body), 0o644)
	return &agent.Result{
		Success: true,
		Patches: map[string][]byte{"../../output.md": []byte(b.body)},
	}, nil
}

type bothWriteBackend struct{ local, parent string }

func (b *bothWriteBackend) Name() string { return "both-write" }
func (b *bothWriteBackend) Execute(_ context.Context, req agent.Request) (*agent.Result, error) {
	parent := filepath.Clean(filepath.Join(req.WorkDir, "..", "..", "output.md"))
	_ = os.WriteFile(parent, []byte(b.parent), 0o644)
	return &agent.Result{
		Success: true,
		Patches: map[string][]byte{"output.md": []byte(b.local)},
	}, nil
}

func TestEnsureQualityOrRepair_LocalTargetPromotesAndRevalidates(t *testing.T) {
	workDir := prepareRepairWS(t)
	driver := squad.NewWorkflowDriver(squad.NewWorkspaceEngine())
	gate := squad.NewQualityGate()
	report, err := gate.EnsureQualityOrRepair(context.Background(), workDir, "intake", driver, &localRepairBackend{body: repairGoodDraft}, 1, squad.ExecuteOptions{
		Timeout: 5 * time.Second,
	})
	if err != nil {
		t.Fatalf("EnsureQualityOrRepair: %v", err)
	}
	if report == nil || !report.Passed {
		t.Fatalf("expected revalidate PASS after local repair, got %+v", report)
	}
	root, err := os.ReadFile(filepath.Join(workDir, "output.md"))
	if err != nil || !strings.Contains(string(root), "详细方案") {
		t.Fatalf("root output.md must be overwritten from local repair, got %q", root)
	}
	task, err := os.ReadFile(filepath.Join(workDir, "tasks", "repair_1", "TASK.md"))
	if err != nil {
		t.Fatal(err)
	}
	if squad.TaskMarkdownEscapesParent(string(task)) {
		t.Fatalf("repair TASK.md must not contain ../ :\n%s", task)
	}
	if !strings.Contains(string(task), "source_output.md") {
		t.Fatalf("repair TASK.md must cite local source_output.md:\n%s", task)
	}
	src, err := os.ReadFile(filepath.Join(workDir, "tasks", "repair_1", squad.RepairSourceOutputName))
	if err != nil || !strings.Contains(string(src), "不做展开") {
		t.Fatalf("source_output.md should keep the failing draft, got %q", src)
	}
}

func TestEnsureQualityOrRepair_ParentEscapeRestoresAndFailsClosed(t *testing.T) {
	workDir := prepareRepairWS(t)
	orig, _ := os.ReadFile(filepath.Join(workDir, "output.md"))
	driver := squad.NewWorkflowDriver(squad.NewWorkspaceEngine())
	gate := squad.NewQualityGate()
	report, err := gate.EnsureQualityOrRepair(context.Background(), workDir, "intake", driver, &parentEscapeBackend{body: repairGoodDraft}, 1, squad.ExecuteOptions{
		Timeout: 5 * time.Second,
	})
	if err != nil {
		t.Fatalf("EnsureQualityOrRepair: %v", err)
	}
	if report == nil || report.Passed {
		t.Fatalf("escaped parent write must not count as success: %+v", report)
	}
	if len(report.MissingSections) == 0 {
		t.Fatalf("must keep original quality_failed missing sections, got %+v", report)
	}
	got, _ := os.ReadFile(filepath.Join(workDir, "output.md"))
	if string(got) != string(orig) {
		t.Fatalf("parent output.md must be restored to pre-repair bytes, got %q", got)
	}
	if _, err := os.Stat(filepath.Join(workDir, "tasks", "repair_1", "output.md")); err == nil {
		t.Fatal("must not invent a local repair target")
	}
	audit, _ := os.ReadFile(filepath.Join(workDir, "gov_audit.jsonl"))
	if !strings.Contains(string(audit), squad.GovAuditContainment) {
		t.Fatalf("expected containment audit, got %s", audit)
	}
}

func TestEnsureQualityOrRepair_BothChangedPromotesLocalAfterRestore(t *testing.T) {
	workDir := prepareRepairWS(t)
	driver := squad.NewWorkflowDriver(squad.NewWorkspaceEngine())
	gate := squad.NewQualityGate()
	parentDecoy := repairGoodDraft + "\n\n## 越界父级不应被采用\n"
	report, err := gate.EnsureQualityOrRepair(context.Background(), workDir, "intake", driver, &bothWriteBackend{
		local:  repairGoodDraft,
		parent: parentDecoy,
	}, 1, squad.ExecuteOptions{Timeout: 5 * time.Second})
	if err != nil {
		t.Fatalf("EnsureQualityOrRepair: %v", err)
	}
	if report == nil || !report.Passed {
		t.Fatalf("trusted local target should still pass after restoring parent: %+v", report)
	}
	root, _ := os.ReadFile(filepath.Join(workDir, "output.md"))
	if strings.Contains(string(root), "越界父级不应被采用") {
		t.Fatalf("must not keep escaped parent bytes: %s", root)
	}
	if !strings.Contains(string(root), "详细方案") {
		t.Fatalf("must promote local repair body, got %s", root)
	}
	audit, _ := os.ReadFile(filepath.Join(workDir, "gov_audit.jsonl"))
	if !strings.Contains(string(audit), squad.GovAuditContainment) {
		t.Fatalf("both-changed must still record containment, got %s", audit)
	}
}

func TestEnsureQualityOrRepair_RepairBudgetNotRegressed(t *testing.T) {
	workDir := prepareRepairWS(t)
	driver := squad.NewWorkflowDriver(squad.NewWorkspaceEngine())
	gate := squad.NewQualityGate()
	_, err := gate.EnsureQualityOrRepair(context.Background(), workDir, "intake", driver, &parentEscapeBackend{body: repairGoodDraft}, 1, squad.ExecuteOptions{
		Timeout: 5 * time.Second,
	})
	if err != nil {
		t.Fatal(err)
	}
	ledger := squad.LoadOrCreateLedger(workDir, squad.DefaultBudgetPolicy())
	if ledger.RepairRounds != 1 {
		t.Fatalf("containment fail still consumes one repair round, got %d", ledger.RepairRounds)
	}
}

func TestAddStep_SnapshotLayoutCopiesNeedsLocally(t *testing.T) {
	engine := squad.NewWorkspaceEngine()
	ws, err := engine.PrepareWorkspace(squad.WorkspacePrepareRequest{
		WorkDir:         t.TempDir(),
		TaskKey:         "intake",
		StageTitle:      "用户初始需求收集表",
		TemplateContent: repairTemplate,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(ws.WorkDir, "output.md"), []byte(repairBadDraft), 0o644); err != nil {
		t.Fatal(err)
	}
	driver := squad.NewWorkflowDriver(engine)
	err = driver.AddStep(ws.WorkDir, squad.StepSpec{
		ID:           "step_repair_1",
		Role:         "writer",
		Title:        "repair",
		SubWorkspace: "tasks/repair_1",
		Mode:         "execute",
		Needs:        []string{"output.md"},
		InputLayout:  squad.InputLayoutSnapshot,
	})
	if err != nil {
		t.Fatal(err)
	}
	sub := filepath.Join(ws.WorkDir, "tasks", "repair_1")
	task, _ := os.ReadFile(filepath.Join(sub, "TASK.md"))
	if squad.TaskMarkdownEscapesParent(string(task)) {
		t.Fatalf("AddStep snapshot TASK.md has ../:\n%s", task)
	}
	if _, err := os.Lstat(filepath.Join(sub, "template.md")); err != nil {
		t.Fatal(err)
	}
	info, err := os.Lstat(filepath.Join(sub, "template.md"))
	if err != nil || !info.Mode().IsRegular() {
		t.Fatalf("snapshot template.md must be a regular file: %v", err)
	}
	if _, err := os.Stat(filepath.Join(sub, squad.RepairSourceOutputName)); err != nil {
		t.Fatal(err)
	}
}