Agent-Squad/squad/runstore.go

package squad

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"time"
)

// Store 是运行状态持久化端口（Q1 WP4）：run/gate/checkpoint 的读写以
// 运行工作区为键。现状文件态收编为 FileStore（行为与此前逐文件直写完全一致）；
// SQLite 实现进 adapters/sqlite（WP5）。事件流仍归 store/eventsink（子包），
// 不在本端口内重复。
type Store interface {
	SaveRun(st *RunState) error
	LoadRun(wsDir string) (*RunState, error)
	SaveGate(wsDir string, gate *HitlGate) error
	LoadGate(wsDir string) (*HitlGate, error)
	SaveCheckpoint(wsDir string, cp Checkpoint) error
	LoadLatestCheckpoint(wsDir string) (*Checkpoint, error)
	// SaveLeader/LoadLeader 持久化 leader 循环状态（raw JSON bytes；
	// 文件态即 leader.json）。LoadLeader 不存在时返回 (nil, nil)。
	SaveLeader(wsDir string, data []byte) error
	LoadLeader(wsDir string) ([]byte, error)
}

// FileStore 把 runs/<id>/ 目录布局呈现为 Store（现状文件态）。
type FileStore struct{}

func (FileStore) SaveRun(st *RunState) error {
	return st.persist()
}

func (FileStore) LoadRun(wsDir string) (*RunState, error) {
	return loadRunState(wsDir)
}

func (FileStore) SaveGate(wsDir string, gate *HitlGate) error {
	st := &RunState{}
	st.Meta.WorkDir = wsDir
	return st.saveGate(gate)
}

func (FileStore) LoadGate(wsDir string) (*HitlGate, error) {
	return loadGate(wsDir)
}

func (FileStore) SaveLeader(wsDir string, data []byte) error {
	return os.WriteFile(filepath.Join(wsDir, "leader.json"), data, 0o644)
}

func (FileStore) LoadLeader(wsDir string) ([]byte, error) {
	data, err := os.ReadFile(filepath.Join(wsDir, "leader.json"))
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return data, nil
}

func (FileStore) SaveCheckpoint(wsDir string, cp Checkpoint) error {
	dir := filepath.Join(wsDir, "checkpoints")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	path := filepath.Join(dir, fmt.Sprintf("%03d.json", cp.Superstep))
	if err := writeJSON(path, cp); err != nil {
		return err
	}
	return writeJSON(filepath.Join(dir, "latest.json"), cp)
}

func (FileStore) LoadLatestCheckpoint(wsDir string) (*Checkpoint, error) {
	data, err := os.ReadFile(filepath.Join(wsDir, "checkpoints", "latest.json"))
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	var cp Checkpoint
	if err := json.Unmarshal(data, &cp); err != nil {
		return nil, fmt.Errorf("parse latest checkpoint: %w", err)
	}
	return &cp, nil
}

// Checkpoint 是一个 superstep 边界的状态快照（Q1 WP4）。
// GraphHash 是图结构指纹：恢复时与当前 Steps 重编译比对，不一致说明崩溃
// 期间图形态被改过（动态调整或外力），按 WP1 裁定拒绝恢复并提示重新规划。
type Checkpoint struct {
	Superstep  int            `json:"superstep"`
	At         time.Time      `json:"at"`
	GraphHash  string         `json:"graph_hash"`
	StepStates []StepSnapshot `json:"steps"`
	wsDir      string         `json:"-"`
}

// write 立即同步落盘（经当前 Store 端口——SQLite 后端时进数据库）。
func (cp Checkpoint) write(store Store) error {
	return store.SaveCheckpoint(cp.wsDir, cp)
}

// StepSnapshot 是检查点里的单步快照。
type StepSnapshot struct {
	ID        string `json:"id"`
	Status    string `json:"status"`
	AttemptID string `json:"attempt_id,omitempty"`
}

// CheckpointMode 快照写盘档位（Q1 WP4 三档）。
type CheckpointMode string

const (
	// CheckpointSync 每个 superstep 边界同步写（默认，与现状「逐步全量落盘」一致）。
	CheckpointSync CheckpointMode = "sync"
	// CheckpointAsync 边界快照进队列由单 goroutine 串行写（边界不阻塞；恢复粒度同 sync）。
	CheckpointAsync CheckpointMode = "async"
	// CheckpointOnExit 只在推进退出（终态/挂起/错误）时写最终快照。
	CheckpointOnExit CheckpointMode = "on_exit"
)

// checkpointWriter 按档位执行边界快照；async 档单 goroutine 串行落盘。
type checkpointWriter struct {
	mode  CheckpointMode
	store Store
	ch    chan Checkpoint
	done  chan error
}

func newCheckpointWriter(mode CheckpointMode, store Store) *checkpointWriter {
	if mode == "" {
		mode = CheckpointSync
	}
	if store == nil {
		store = FileStore{}
	}
	w := &checkpointWriter{mode: mode, store: store}
	if mode == CheckpointAsync {
		w.ch = make(chan Checkpoint, 16)
		w.done = make(chan error, 1)
		go func() {
			var first error
			for cp := range w.ch {
				if first != nil {
					continue
				}
				if err := cp.write(w.store); err != nil {
					first = err
				}
			}
			w.done <- first
		}()
	}
	return w
}

// boundary 由调用方（flowDriver）在 superstep 边界触发：
// sync 立即写；async 入队（满即错，不静默丢）；on_exit 跳过。
func (w *checkpointWriter) boundary(wsDir string, cp Checkpoint) error {
	cp.wsDir = wsDir
	switch w.mode {
	case CheckpointAsync:
		select {
		case w.ch <- cp:
		default:
			return fmt.Errorf("checkpoint queue full at superstep %d", cp.Superstep)
		}
		return nil
	case CheckpointOnExit:
		return nil
	default:
		return cp.write(w.store)
	}
}

// flush 在推进退出前调用：async 落完并报告首个错误；on_exit 写最终快照。
func (w *checkpointWriter) flush(wsDir string, cp Checkpoint) error {
	cp.wsDir = wsDir
	switch w.mode {
	case CheckpointAsync:
		close(w.ch)
		return <-w.done
	case CheckpointOnExit:
		return cp.write(w.store)
	default:
		return nil
	}
}

// validateGraphContinuity 恢复连续性校验（Q1 WP4；WP1 的图版本裁定）。
// 最新检查点存在且其图指纹与当前 Steps 重编译不一致时拒绝恢复——
// 崩溃期间图形态被改过，继续推进会把已完成的结果接到错误的图上。
// 检查点不存在（旧运行）或指纹一致则放行。
func validateGraphContinuity(st *RunState) error {
	latest, err := st.storeOr().LoadLatestCheckpoint(st.Meta.WorkDir)
	if err != nil || latest == nil {
		return err
	}
	g, gerr := compileGraph(st.Plan.Workflow.Steps)
	if gerr != nil {
		return fmt.Errorf("graph continuity: 当前 Steps 无法编译：%w（请人工介入或重新规划）", gerr)
	}
	if latest.GraphHash != g.hash() {
		return fmt.Errorf("graph continuity: 崩溃前后图形态不一致（checkpoint superstep=%d hash=%s，当前=%s）——可能在动态调整或外部修改后崩溃，请重新规划或人工介入",
			latest.Superstep, latest.GraphHash, g.hash())
	}
	return nil
}