Agent-Squad/adapters/sqlite/runstore.go

package platform

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"time"

	"agent-squad/squad"
)

// SQLiteStore 把 squad 的 Store 端口落到本库（Q1 WP5）：run/gate/checkpoint/leader
// 与平台其余表同一个 SQLite 文件——运行状态与平台状态不再各记一份。
// 表名带 squad_run_ 前缀，与平台既有 23 张表并存不冲突。
type SQLiteStore struct {
	db *sql.DB
}

// OpenRunStore 在本 Store 的底层连接上提供 squad.Store 端口（schema 随 migrate 建）。
func (s *Store) RunStore() *SQLiteStore {
	return &SQLiteStore{db: s.db}
}

const runSchema = `
CREATE TABLE IF NOT EXISTS squad_run_state (
  ws_dir TEXT PRIMARY KEY,
  meta TEXT NOT NULL, workflow TEXT NOT NULL, squad TEXT NOT NULL, bb TEXT NOT NULL,
  template TEXT, approved TEXT, repair_budget_remaining INTEGER NOT NULL DEFAULT 0,
  repair_budget_granted INTEGER NOT NULL DEFAULT 0,
  clarifications TEXT, steering TEXT, usage TEXT, updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS squad_run_gate (
  ws_dir TEXT PRIMARY KEY, gate TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS squad_run_checkpoint (
  ws_dir TEXT NOT NULL, superstep INTEGER NOT NULL, at TEXT NOT NULL,
  graph_hash TEXT NOT NULL, steps TEXT NOT NULL,
  PRIMARY KEY (ws_dir, superstep)
);
CREATE TABLE IF NOT EXISTS squad_run_leader (
  ws_dir TEXT PRIMARY KEY, data TEXT NOT NULL
);
`

func (s *SQLiteStore) migrate() error {
	_, err := s.db.Exec(runSchema)
	return err
}

// OpenSQLiteStore 独立打开一个 squad.Store（不经 platform.Store 的场景，如测试）。
func OpenSQLiteStore(path string) (*SQLiteStore, error) {
	db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)")
	if err != nil {
		return nil, err
	}
	s := &SQLiteStore{db: db}
	if err := s.migrate(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return s, nil
}

func j(v any) string {
	b, err := json.Marshal(v)
	if err != nil {
		return ""
	}
	return string(b)
}

func (s *SQLiteStore) SaveRun(st *squad.RunState) error {
	if err := s.migrate(); err != nil {
		return err
	}
	_, err := s.db.Exec(`INSERT INTO squad_run_state
	  (ws_dir, meta, workflow, squad, bb, template, approved, repair_budget_remaining, repair_budget_granted, clarifications, steering, usage, updated_at)
	  VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)
	  ON CONFLICT(ws_dir) DO UPDATE SET
	    meta=excluded.meta, workflow=excluded.workflow, squad=excluded.squad, bb=excluded.bb,
	    template=excluded.template, approved=excluded.approved,
	    repair_budget_remaining=excluded.repair_budget_remaining,
	    repair_budget_granted=excluded.repair_budget_granted,
	    clarifications=excluded.clarifications, steering=excluded.steering,
	    usage=excluded.usage, updated_at=excluded.updated_at`,
		st.Meta.WorkDir, j(st.Meta), j(st.Plan.Workflow), j(st.Plan.Squad), j(st.BB),
		st.Template, j(st.Approved), st.RepairBudgetRemaining, st.RepairBudgetGranted,
		j(st.Clarifications), j(st.Steering), j(st.Usage), now())
	return err
}

func (s *SQLiteStore) LoadRun(wsDir string) (*squad.RunState, error) {
	if err := s.migrate(); err != nil {
		return nil, err
	}
	st := &squad.RunState{}
	var meta, wf, sq, bb, template, approved, clar, steering, usage string
	row := s.db.QueryRow(`SELECT meta, workflow, squad, bb, template, approved,
	  repair_budget_remaining, repair_budget_granted, clarifications, steering, usage
	  FROM squad_run_state WHERE ws_dir = ?`, wsDir)
	var granted int
	err := row.Scan(&meta, &wf, &sq, &bb, &template, &approved,
		&st.RepairBudgetRemaining, &granted, &clar, &steering, &usage)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("load run %s: %w", wsDir, os.ErrNotExist)
		}
		return nil, err
	}
	st.RepairBudgetGranted = granted
	for pair, dst := range map[string]any{
		meta: &st.Meta, wf: &st.Plan.Workflow, sq: &st.Plan.Squad, bb: &st.BB,
		approved: &st.Approved, clar: &st.Clarifications, steering: &st.Steering, usage: &st.Usage,
	} {
		if err := json.Unmarshal([]byte(pair), dst); err != nil {
			return nil, fmt.Errorf("load run %s: parse: %w", wsDir, err)
		}
	}
	st.Template = template
	if st.Approved == nil {
		st.Approved = map[string]bool{}
	}
	return st, nil
}

func (s *SQLiteStore) SaveGate(wsDir string, gate *squad.HitlGate) error {
	if err := s.migrate(); err != nil {
		return err
	}
	if gate == nil {
		_, err := s.db.Exec(`DELETE FROM squad_run_gate WHERE ws_dir = ?`, wsDir)
		return err
	}
	_, err := s.db.Exec(`INSERT INTO squad_run_gate (ws_dir, gate) VALUES (?,?)
	  ON CONFLICT(ws_dir) DO UPDATE SET gate=excluded.gate`, wsDir, j(gate))
	return err
}

func (s *SQLiteStore) LoadGate(wsDir string) (*squad.HitlGate, error) {
	if err := s.migrate(); err != nil {
		return nil, err
	}
	var raw string
	err := s.db.QueryRow(`SELECT gate FROM squad_run_gate WHERE ws_dir = ?`, wsDir).Scan(&raw)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	gate := &squad.HitlGate{}
	if err := json.Unmarshal([]byte(raw), gate); err != nil {
		return nil, fmt.Errorf("parse gate: %w", err)
	}
	return gate, nil
}

func (s *SQLiteStore) SaveCheckpoint(wsDir string, cp squad.Checkpoint) error {
	if err := s.migrate(); err != nil {
		return err
	}
	_, err := s.db.Exec(`INSERT INTO squad_run_checkpoint (ws_dir, superstep, at, graph_hash, steps)
	  VALUES (?,?,?,?,?)
	  ON CONFLICT(ws_dir, superstep) DO UPDATE SET at=excluded.at, graph_hash=excluded.graph_hash, steps=excluded.steps`,
		wsDir, cp.Superstep, cp.At.UTC().Format("2006-01-02T15:04:05Z07:00"), cp.GraphHash, j(cp.StepStates))
	return err
}

func (s *SQLiteStore) LoadLatestCheckpoint(wsDir string) (*squad.Checkpoint, error) {
	if err := s.migrate(); err != nil {
		return nil, err
	}
	var cp squad.Checkpoint
	var at, steps string
	err := s.db.QueryRow(`SELECT superstep, at, graph_hash, steps FROM squad_run_checkpoint
	  WHERE ws_dir = ? ORDER BY superstep DESC LIMIT 1`, wsDir).
		Scan(&cp.Superstep, &at, &cp.GraphHash, &steps)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	t, _ := parseTime(at)
	cp.At = t
	if err := json.Unmarshal([]byte(steps), &cp.StepStates); err != nil {
		return nil, fmt.Errorf("parse checkpoint steps: %w", err)
	}
	return &cp, nil
}

func (s *SQLiteStore) SaveLeader(wsDir string, data []byte) error {
	if err := s.migrate(); err != nil {
		return err
	}
	_, err := s.db.Exec(`INSERT INTO squad_run_leader (ws_dir, data) VALUES (?,?)
	  ON CONFLICT(ws_dir) DO UPDATE SET data=excluded.data`, wsDir, string(data))
	return err
}

func (s *SQLiteStore) LoadLeader(wsDir string) ([]byte, error) {
	if err := s.migrate(); err != nil {
		return nil, err
	}
	var raw string
	err := s.db.QueryRow(`SELECT data FROM squad_run_leader WHERE ws_dir = ?`, wsDir).Scan(&raw)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	return []byte(raw), nil
}

// Close 只在独立打开（OpenSQLiteStore）时调用；经 platform.Store 打开时由平台统一关闭。
func (s *SQLiteStore) Close() error { return s.db.Close() }

var _ squad.Store = (*SQLiteStore)(nil)

// parseTime 解析检查点时间（RFC3339Nano 与 RFC3339 两态）。
func parseTime(s string) (time.Time, error) {
	for _, layout := range []string{time.RFC3339Nano, time.RFC3339} {
		if t, err := time.Parse(layout, s); err == nil {
			return t, nil
		}
	}
	return time.Time{}, fmt.Errorf("bad time %q", s)
}