Agent-Squad/adapters/sqlite/runstore_test.go

package platform

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

	"agent-squad/squad"
)

func newTestRunState(wsDir string) *squad.RunState {
	return &squad.RunState{
		Meta: squad.RunMetadata{
			ID: "run_test", Status: "executing", WorkDir: wsDir,
			CreatedAt: time.Now(), UpdatedAt: time.Now(),
			Metadata: map[string]string{"mode": ""},
		},
		Plan: squad.SquadPlan{
			Workflow: squad.WorkflowPlan{
				WorkflowID: "wf1", Status: "in_progress",
				Steps: []squad.StepSpec{
					{ID: "s1", Role: "researcher", Status: "completed", Attempts: 1, AttemptID: "run_test/s1/1"},
					{ID: "s2", Role: "writer", Status: "pending"},
				},
			},
			Squad: squad.SquadConfig{SquadName: "demo"},
		},
		BB:                    squad.BlackboardState{Facts: []squad.FactItem{{ID: "f1", Claim: "x"}}},
		Template:              "tpl",
		Approved:              map[string]bool{"step:s1": true},
		RepairBudgetRemaining: 1,
		RepairBudgetGranted:   0,
		Clarifications:        []string{"c1"},
		Steering:              []string{"st1"},
		Usage:                 map[string]*squad.Usage{"writer": {PromptTokens: 10, CompletionTokens: 5, Steps: 1}},
	}
}

// TestSQLiteStoreRoundTrip：run/gate/checkpoint/leader 全链路（Q1 WP5 端口实现）。
func TestSQLiteStoreRoundTrip(t *testing.T) {
	st, err := OpenSQLiteStore(filepath.Join(t.TempDir(), "runs.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer st.Close()
	ws := filepath.Join(t.TempDir(), "runs", "run_test")
	orig := newTestRunState(ws)
	if err := st.SaveRun(orig); err != nil {
		t.Fatal(err)
	}
	got, err := st.LoadRun(ws)
	if err != nil {
		t.Fatal(err)
	}
	if got.Meta.ID != "run_test" || got.Plan.Workflow.WorkflowID != "wf1" || got.Template != "tpl" {
		t.Fatalf("loaded = %+v", got.Meta)
	}
	if !got.Approved["step:s1"] || got.RepairBudgetRemaining != 1 || len(got.Clarifications) != 1 || len(got.Steering) != 1 {
		t.Fatalf("scalars = %+v", got)
	}
	if got.Usage["writer"].PromptTokens != 10 {
		t.Fatalf("usage = %+v", got.Usage)
	}
	if _, err := st.LoadRun(filepath.Join(t.TempDir(), "nope")); err == nil {
		t.Fatal("missing run should error")
	}

	// gate：写→读→删。
	gate := &squad.HitlGate{PrevStatus: squad.RunExecuting}
	if err := st.SaveGate(ws, gate); err != nil {
		t.Fatal(err)
	}
	if g2, err := st.LoadGate(ws); err != nil || g2 == nil {
		t.Fatalf("load gate: %v %+v", err, g2)
	}
	if err := st.SaveGate(ws, nil); err != nil {
		t.Fatal(err)
	}
	if g2, err := st.LoadGate(ws); err != nil || g2 != nil {
		t.Fatalf("gate after delete: %v %+v", err, g2)
	}

	// checkpoint：两层，latest 取最大 superstep。
	if err := st.SaveCheckpoint(ws, squad.Checkpoint{Superstep: 1, At: time.Now(), GraphHash: "h1", StepStates: []squad.StepSnapshot{{ID: "s1", Status: "completed", AttemptID: "run_test/s1/1"}}}); err != nil {
		t.Fatal(err)
	}
	if err := st.SaveCheckpoint(ws, squad.Checkpoint{Superstep: 2, At: time.Now(), GraphHash: "h2"}); err != nil {
		t.Fatal(err)
	}
	latest, err := st.LoadLatestCheckpoint(ws)
	if err != nil || latest == nil {
		t.Fatalf("latest: %v", err)
	}
	if latest.Superstep != 2 || latest.GraphHash != "h2" {
		t.Fatalf("latest = %+v", latest)
	}
	if l2, err := st.LoadLatestCheckpoint(filepath.Join(t.TempDir(), "nope")); err != nil || l2 != nil {
		t.Fatalf("latest for unknown: %v %+v", err, l2)
	}

	// leader bytes。
	if err := st.SaveLeader(ws, []byte(`{"turns":2}`)); err != nil {
		t.Fatal(err)
	}
	data, err := st.LoadLeader(ws)
	if err != nil || string(data) != `{"turns":2}` {
		t.Fatalf("leader = %q %v", data, err)
	}
}

// TestEngineRunsOnSQLiteStore：引擎注入 SQLiteStore 后运行状态进 SQLite、
// 不再写 run 级 JSON 文件；第二引擎可从中恢复（验收第三条的实证形态）。
func TestEngineRunsOnSQLiteStore(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "runs.db")
	runStore, err := OpenSQLiteStore(dbPath)
	if err != nil {
		t.Fatal(err)
	}
	defer runStore.Close()
	dir := t.TempDir()
	reg := squad.NewRoleRegistry()
	for _, spec := range []squad.RoleSpec{{Role: "researcher", Backend: "script"}, {Role: "writer", Backend: "script"}, {Role: "reviewer", Backend: "script"}} {
		if err := reg.Register(spec); err != nil {
			t.Fatal(err)
		}
	}
	backends := squad.NewBackendResolver()
	backends.RegisterInstance("script", stubExecutor{result: &squad.ExecResult{Success: false, Logs: "ok", Patches: map[string][]byte{"output.md": []byte(goodReportLite), "notes.md": []byte(goodReportLite), "fields.json": []byte("{}")}}})
	blueprints := squad.NewBlueprintRegistry()
	if err := blueprints.Register(squad.Blueprint{
		Key: "demo", Title: "入库演示", Match: []string{"入库演示"},
		Steps: []squad.StepSpec{
			{ID: "s1", Role: "researcher", Title: "调研", Targets: []string{"notes.md", "fields.json"}},
			{ID: "s2", Role: "writer", Title: "成文", Targets: []string{"output.md"}},
			{ID: "s3", Role: "reviewer", Title: "终审", Targets: []string{"output.md"}},
		},
	}); err != nil {
		t.Fatal(err)
	}
	eng, err := squad.NewEngine(squad.Options{
		DataDir: dir, Registry: reg, Blueprints: blueprints,
		Planner:  squad.NewRulePlanner(),
		Backends: backends, Sink: squad.NopSink{},
		Store: runStore,
	})
	if err != nil {
		t.Fatal(err)
	}
	res, err := eng.Run(context.Background(), squad.Goal{Text: "做一个三步入库演示", HitlMode: squad.HitlAuto}, squad.RunOptions{})
	if err != nil || res.Status != squad.RunCompleted {
		t.Fatalf("run: %v %+v", err, res)
	}
	// run 级 JSON 不在盘上（状态全进了 SQLite）。
	for _, f := range []string{"workspace.json", "workflow.json", "squad.json", "blackboard.json", "runstate.json"} {
		if _, err := os.Stat(filepath.Join(res.WorkDir, f)); !os.IsNotExist(err) {
			t.Fatalf("%s should not be on disk with SQLite store (err=%v)", f, err)
		}
	}
	// 状态可经端口读回。
	st, err := runStore.LoadRun(res.WorkDir)
	if err != nil {
		t.Fatal(err)
	}
	if st.Plan.Workflow.Status != "completed" {
		t.Fatalf("workflow status = %s", st.Plan.Workflow.Status)
	}
	latest, err := runStore.LoadLatestCheckpoint(res.WorkDir)
	if err != nil || latest == nil {
		t.Fatalf("latest checkpoint: %v", err)
	}
	// 第二引擎（同库）恢复可用。
	eng2, err := squad.NewEngine(squad.Options{
		DataDir: dir, Registry: reg, Blueprints: blueprints,
		Planner:  squad.NewRulePlanner(),
		Backends: backends, Sink: squad.NopSink{},
		Store: runStore,
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := eng2.Inspect(res.RunID); err != nil {
		t.Fatalf("inspect from second engine: %v", err)
	}
}

const goodReportLite = "# 测试交付物\n\n## 概述\n\n这是由脚本后端生成的合格交付物，长度超过五十个字符以满足最小质检要求。\n"