Agent-Squad/squad/checkpoint_test.go

package squad

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

// TestCheckpointFilesAtBoundaries：sync 默认档——每个 superstep 边界有快照，
// latest 指向最后一层且图指纹与重编译一致（Q1 WP4）。
func TestCheckpointFilesAtBoundaries(t *testing.T) {
	eng, _ := newTestEngine(t, mockScriptPlanner(baseSteps()...), map[string]map[string]string{
		"01_researcher": {"notes.md": goodReport, "fields.json": "{}"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}, NopSink{})
	res, err := eng.Run(context.Background(), Goal{Text: "快照"}, RunOptions{})
	if err != nil || res.Status != RunCompleted {
		t.Fatalf("run: %v %+v", err, res)
	}
	latest, err := FileStore{}.LoadLatestCheckpoint(res.WorkDir)
	if err != nil || latest == nil {
		t.Fatalf("latest checkpoint: %v", err)
	}
	if latest.Superstep < 3 {
		t.Fatalf("superstep = %d, want >= 3（三步各占一层）", latest.Superstep)
	}
	g, err := compileGraph(res.Plan.Steps)
	if err != nil {
		t.Fatal(err)
	}
	if latest.GraphHash != g.hash() {
		t.Fatalf("hash = %s, want %s", latest.GraphHash, g.hash())
	}
	if len(latest.StepStates) != 3 {
		t.Fatalf("step states = %d", len(latest.StepStates))
	}
	if _, err := os.Stat(filepath.Join(res.WorkDir, "checkpoints", "001.json")); err != nil {
		t.Fatalf("first boundary snapshot missing: %v", err)
	}
}

// TestCheckpointAsyncAndOnExit：async 在退出前落完；on_exit 只有最终快照。
func TestCheckpointAsyncAndOnExit(t *testing.T) {
	for _, mode := range []string{"async", "on_exit"} {
		eng, _ := newTestEngine(t, mockScriptPlanner(baseSteps()...), map[string]map[string]string{
			"01_researcher": {"notes.md": goodReport, "fields.json": "{}"},
			"02_writer":     {"output.md": goodReport},
			"03_reviewer":   {"output.md": goodReport},
		}, NopSink{})
		res, err := eng.Run(context.Background(), Goal{Text: "快照"}, RunOptions{CheckpointMode: mode})
		if err != nil || res.Status != RunCompleted {
			t.Fatalf("%s run: %v %+v", mode, err, res)
		}
		latest, err := FileStore{}.LoadLatestCheckpoint(res.WorkDir)
		if err != nil || latest == nil {
			t.Fatalf("%s latest: %v", mode, err)
		}
		if latest.Superstep < 3 {
			t.Fatalf("%s superstep = %d", mode, latest.Superstep)
		}
	}
}

// TestGraphContinuityRejectsStructuralChange：崩溃期间图形态被改则拒绝恢复（WP1 裁定）。
func TestGraphContinuityRejectsStructuralChange(t *testing.T) {
	eng, _ := newTestEngine(t, mockScriptPlanner(baseSteps()...), map[string]map[string]string{
		"01_researcher": {"notes.md": goodReport, "fields.json": "{}"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}, NopSink{})
	dir := eng.opts.DataDir
	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（仅状态变化，图结构不变——应先过连续性校验）。
	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[1].Status = "running"
	wf.Steps[1].Attempts = 1
	// 再改结构：把 s3 的 needs 改成另一个合法 id——编译能过但指纹必变。
	wf.Steps[2].Needs = []string{"s1"}
	st.setWorkflow(wf)
	if err := st.save(); err != nil {
		t.Fatal(err)
	}
	// 用一个新的引擎在同一目录上恢复（图编译本身合法：needs 引用存在的 id）。
	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", &scriptBackend{script: map[string]map[string]string{}})
	eng2, err := NewEngine(Options{DataDir: dir, Registry: reg, Planner: eng.opts.Planner, Backends: backends, Sink: NopSink{}})
	if err != nil {
		t.Fatal(err)
	}
	_, err = eng2.Resume(context.Background(), res.RunID)
	if err == nil || !strings.Contains(err.Error(), "图形态不一致") && !strings.Contains(err.Error(), "graph continuity") {
		t.Fatalf("expected continuity rejection, got %v", err)
	}
}