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}
}

// runSchemaVersion 是 squad run 存储的 schema 版本（squad_run_schema_version 记档）。
// v1 起步，v2 补齐协作态列（mailbox/handoffs/board）。迁移只做增量：
// 新库直接建到最新形态，旧库经 ALTER TABLE ADD COLUMN 补列（既有行读出为空默认，不丢数据）。
const runSchemaVersion = 2

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,
  mailbox TEXT, handoffs TEXT, board TEXT
);
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
);
CREATE TABLE IF NOT EXISTS squad_run_schema_version (
  id INTEGER PRIMARY KEY CHECK (id = 1), version INTEGER NOT NULL
);
`

// runStateV2Columns 是 v2 在 squad_run_state 上新增的协作态列（列名为代码内常量，非外部输入）。
var runStateV2Columns = []string{"mailbox", "handoffs", "board"}

// migrate 建表并按版本步进：老库补列后再记版本号，逐次读写都是幂等的
// （每次 Save/Load 都调；补列只在版本落后时发生，且既有行不被改写）。
func (s *SQLiteStore) migrate() error {
	if _, err := s.db.Exec(runSchema); err != nil {
		return err
	}
	cur, err := s.runSchemaVersion()
	if err != nil {
		return err
	}
	if cur >= runSchemaVersion {
		return nil
	}
	if err := s.addRunStateColumns(runStateV2Columns); err != nil {
		return err
	}
	_, err = s.db.Exec(`INSERT INTO squad_run_schema_version (id, version) VALUES (1, ?)
	  ON CONFLICT(id) DO UPDATE SET version = excluded.version`, runSchemaVersion)
	return err
}

// runSchemaVersion 读取当前 schema 版本；无记录 = 0（v1 及更早的库）。
func (s *SQLiteStore) runSchemaVersion() (int, error) {
	var v int
	err := s.db.QueryRow(`SELECT version FROM squad_run_schema_version WHERE id = 1`).Scan(&v)
	if err != nil {
		if err == sql.ErrNoRows {
			return 0, nil
		}
		return 0, err
	}
	return v, nil
}

// runStateColumns 读出 squad_run_state 现存列名（迁移判存在性用）。
func (s *SQLiteStore) runStateColumns() (map[string]bool, error) {
	rows, err := s.db.Query(`PRAGMA table_info(squad_run_state)`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[string]bool{}
	for rows.Next() {
		var cid, notnull, pk int
		var name, ctype string
		var dflt sql.NullString
		if err := rows.Scan(&cid, &name, &ctype, &notnull, &dflt, &pk); err != nil {
			return nil, err
		}
		out[name] = true
	}
	return out, rows.Err()
}

// addRunStateColumns 为缺失的列补 ALTER TABLE ADD COLUMN（既有行保持原值，
// 新列读出为 NULL → 空默认）。
func (s *SQLiteStore) addRunStateColumns(cols []string) error {
	have, err := s.runStateColumns()
	if err != nil {
		return err
	}
	for _, col := range cols {
		if have[col] {
			continue
		}
		if _, err := s.db.Exec(`ALTER TABLE squad_run_state ADD COLUMN ` + col + ` TEXT`); err != nil {
			// 并发迁移：另一进程可能刚补上同一列，复核后不算失败。
			if again, err2 := s.runStateColumns(); err2 == nil && again[col] {
				continue
			}
			return fmt.Errorf("migrate squad_run_state: add column %s: %w", col, err)
		}
	}
	return nil
}

// 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)
}

// runStateColumn 是 squad_run_state 的一列：列名 + 待编码值（编码失败即报错，
// 不用空串顶替——空串读回来是「空默认」，等于静默丢数据）。
type runStateColumn struct {
	name string
	val  any
}

// runStateColumnsOf 列出 RunState 的全部持久化字段（含协作态 mailbox/handoffs/board）。
func runStateColumnsOf(st *squad.RunState) []runStateColumn {
	return []runStateColumn{
		{"meta", st.Meta},
		{"workflow", st.Plan.Workflow},
		{"squad", st.Plan.Squad},
		{"bb", st.BB},
		{"approved", st.Approved},
		{"clarifications", st.Clarifications},
		{"steering", st.Steering},
		{"usage", st.Usage},
		{"mailbox", st.Mailbox},
		{"handoffs", st.Handoffs},
		{"board", st.Board},
	}
}

func (s *SQLiteStore) SaveRun(st *squad.RunState) error {
	if err := s.migrate(); err != nil {
		return err
	}
	encoded := map[string]string{}
	for _, col := range runStateColumnsOf(st) {
		data, err := json.Marshal(col.val)
		if err != nil {
			return fmt.Errorf("save run %s: encode column %s: %w", st.Meta.WorkDir, col.name, err)
		}
		encoded[col.name] = string(data)
	}
	_, 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, mailbox, handoffs, board, 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, mailbox=excluded.mailbox, handoffs=excluded.handoffs,
	    board=excluded.board, updated_at=excluded.updated_at`,
		st.Meta.WorkDir, encoded["meta"], encoded["workflow"], encoded["squad"], encoded["bb"],
		st.Template, encoded["approved"], st.RepairBudgetRemaining, st.RepairBudgetGranted,
		encoded["clarifications"], encoded["steering"], encoded["usage"],
		encoded["mailbox"], encoded["handoffs"], encoded["board"], now())
	return err
}

// LoadRun 按列位置读回 RunState。解码目标逐列具名指定——绝不用 JSON 文本当键：
// 文本相同的两列（例如两条一模一样的澄清/补充指令）会互相覆盖，把一列丢掉。
func (s *SQLiteStore) LoadRun(wsDir string) (*squad.RunState, error) {
	if err := s.migrate(); err != nil {
		return nil, err
	}
	st := &squad.RunState{}
	// 列 → 扫描目标（顺序与 SQL 一一对应）。
	var tpl sql.NullString
	var granted int
	var meta, wf, sq, bb, approved, clar, steering, usage, mailbox, handoffs, board sql.NullString
	row := s.db.QueryRow(`SELECT meta, workflow, squad, bb, template, approved,
	  repair_budget_remaining, repair_budget_granted, clarifications, steering, usage,
	  mailbox, handoffs, board
	  FROM squad_run_state WHERE ws_dir = ?`, wsDir)
	if err := row.Scan(&meta, &wf, &sq, &bb, &tpl, &approved,
		&st.RepairBudgetRemaining, &granted, &clar, &steering, &usage,
		&mailbox, &handoffs, &board); err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("load run %s: %w", wsDir, os.ErrNotExist)
		}
		return nil, err
	}
	st.RepairBudgetGranted = granted
	for _, dec := range []struct {
		col string
		raw sql.NullString
		dst any
	}{
		{"meta", meta, &st.Meta},
		{"workflow", wf, &st.Plan.Workflow},
		{"squad", sq, &st.Plan.Squad},
		{"bb", bb, &st.BB},
		{"approved", approved, &st.Approved},
		{"clarifications", clar, &st.Clarifications},
		{"steering", steering, &st.Steering},
		{"usage", usage, &st.Usage},
		{"mailbox", mailbox, &st.Mailbox},
		{"handoffs", handoffs, &st.Handoffs},
		{"board", board, &st.Board},
	} {
		// 旧行缺列（NULL）或空串 = 该字段的空默认，不报错、不影响其他列。
		if !dec.raw.Valid || dec.raw.String == "" {
			continue
		}
		if err := json.Unmarshal([]byte(dec.raw.String), dec.dst); err != nil {
			return nil, fmt.Errorf("load run %s: parse column %s: %w", wsDir, dec.col, err)
		}
	}
	st.Template = tpl.String
	if st.Meta.WorkDir == "" {
		// 行键（ws_dir）即权威工作区路径：meta 旧数据缺该字段时由行键兜底。
		st.Meta.WorkDir = wsDir
	}
	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)
}