Cyrene_v2/internal/runtime/persist/persist.go

// Package persist provides durable, SQLite-backed checkpoint storage for the
// runtime agent loop.
//
// It implements the checkpoint (durable execution) requirement described in
// DevDocs/design/04-运行时.md §1: the agent loop persists a Checkpoint at every
// step so a long-running task can be interrupted and resumed later without
// losing context.
//
// The store reuses the *sql.DB opened by internal/db, which already limits the
// pool to a single connection (SetMaxOpenConns(1)) so all writes are serialized
// and never contend on SQLite's single-writer lock. The package therefore owns
// the checkpoints table only; it does not open or configure connections.
//
// Checkpoints are keyed by ID and grouped by AgentID. The full state of a step
// is carried as JSON bytes in State, so a checkpoint record is self-contained
// and can be restored by simply setting it back as the loop's working state.
package persist

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"time"
)

// Package-level sentinel errors. Each is wrapped with %w at the call site so
// callers can use errors.Is.
var (
	// ErrNotFound is returned when a checkpoint for the requested ID or agent
	// does not exist.
	ErrNotFound = errors.New("persist: checkpoint not found")

	// ErrNilDB is returned when New receives a nil *sql.DB.
	ErrNilDB = errors.New("persist: nil database handle")

	// ErrMissingID is returned by Save when the checkpoint ID is blank.
	ErrMissingID = errors.New("persist: checkpoint id is empty")
)

// Checkpoint is a single durable snapshot of the runtime agent loop state at one
// step.
type Checkpoint struct {
	// ID uniquely identifies the checkpoint.
	ID string
	// AgentID groups checkpoints by the agent loop that produced them.
	AgentID string
	// Step is the loop step index this checkpoint represents.
	Step int
	// State is the JSON-encoded loop state at this step.
	State []byte
	// CreatedAt is when the checkpoint was first written. On upsert it is
	// preserved from the original insert.
	CreatedAt time.Time
	// UpdatedAt is when the checkpoint was last written. Save always refreshes
	// it to the current time.
	UpdatedAt time.Time
}

// Store persists and queries checkpoints against a shared *sql.DB.
//
// A Store is safe for concurrent use; the underlying single-connection pool
// serializes writes. The caller supplies the *sql.DB (typically from
// internal/db.Open) and remains responsible for closing it.
type Store struct {
	db  *sql.DB
	log *slog.Logger
}

// DDL for the checkpoints table. Idempotent: both statements are safe to run
// repeatedly.
const (
	createTableSQL = `CREATE TABLE IF NOT EXISTS checkpoints (
		id         TEXT PRIMARY KEY,
		agent_id   TEXT NOT NULL,
		step       INTEGER NOT NULL,
		state      TEXT NOT NULL,
		created_at TEXT NOT NULL,
		updated_at TEXT NOT NULL
	)`

	createIndexSQL = `CREATE INDEX IF NOT EXISTS idx_checkpoints_agent_id ON checkpoints(agent_id)`

	upsertSQL = `INSERT INTO checkpoints (id, agent_id, step, state, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT(id) DO UPDATE SET
			agent_id = excluded.agent_id,
			step     = excluded.step,
			state    = excluded.state,
			updated_at = excluded.updated_at`

	loadSQL = `SELECT id, agent_id, step, state, created_at, updated_at
		FROM checkpoints WHERE id = ?`

	// updated_at is stored as fixed-width UTC text (see formatTime), so its text
	// order is its time order. step breaks ties between checkpoints written in
	// the same tick, and rowid makes that tiebreak deterministic (checkpoints is
	// an ordinary rowid table, not WITHOUT ROWID).
	loadLatestSQL = `SELECT id, agent_id, step, state, created_at, updated_at
		FROM checkpoints WHERE agent_id = ?
		ORDER BY updated_at DESC, step DESC, rowid DESC LIMIT 1`

	listSQL = `SELECT id, agent_id, step, state, created_at, updated_at
		FROM checkpoints WHERE agent_id = ?
		ORDER BY step ASC`

	deleteSQL = `DELETE FROM checkpoints WHERE id = ?`
)

// New creates a Store on db and ensures the checkpoints table and its index
// exist. It is idempotent and may be called more than once.
//
// db must be non-nil; it is typically the *sql.DB returned by internal/db.Open,
// which already constrains the pool to a single connection for write
// serialization. The caller owns db and must Close it when done. log may be nil
// and is replaced with a discard logger so the Store can always log safely.
// Legacy variable-width timestamps are normalized in place (see
// normalizeTimestamps) so text ordering matches time ordering.
func New(ctx context.Context, db *sql.DB, log *slog.Logger) (*Store, error) {
	if db == nil {
		return nil, fmt.Errorf("persist: new: %w", ErrNilDB)
	}
	log = discardLogger(log)

	if _, err := db.ExecContext(ctx, createTableSQL); err != nil {
		return nil, fmt.Errorf("persist: create table: %w", err)
	}
	if _, err := db.ExecContext(ctx, createIndexSQL); err != nil {
		return nil, fmt.Errorf("persist: create index: %w", err)
	}
	if err := normalizeTimestamps(ctx, db); err != nil {
		return nil, err
	}
	return &Store{db: db, log: log}, nil
}

// Save upserts cp inside a transaction. For a new ID the CreatedAt is set to
// now when cp.CreatedAt is zero; for an existing ID the original CreatedAt is
// preserved and only the step, state and UpdatedAt are refreshed.
func (s *Store) Save(ctx context.Context, cp *Checkpoint) error {
	if cp == nil {
		return fmt.Errorf("persist: save: %w", ErrNotFound)
	}
	if cp.ID == "" {
		return fmt.Errorf("persist: save: %w", ErrMissingID)
	}

	now := time.Now().UTC()
	createdAt := cp.CreatedAt
	if createdAt.IsZero() {
		createdAt = now
	}

	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("persist: save: begin: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	if _, err := tx.ExecContext(ctx, upsertSQL,
		cp.ID, cp.AgentID, cp.Step, string(cp.State),
		formatTime(createdAt), formatTime(now),
	); err != nil {
		return fmt.Errorf("persist: save: upsert: %w", err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("persist: save: commit: %w", err)
	}

	s.log.Debug("checkpoint saved", "id", cp.ID, "agent_id", cp.AgentID, "step", cp.Step)
	return nil
}

// Load returns the checkpoint with the given id, or ErrNotFound when it does
// not exist.
func (s *Store) Load(ctx context.Context, id string) (*Checkpoint, error) {
	cp, err := scanCheckpoint(s.db.QueryRowContext(ctx, loadSQL, id))
	if errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("persist: load(%s): %w", id, ErrNotFound)
	}
	if err != nil {
		return nil, fmt.Errorf("persist: load(%s): %w", id, err)
	}
	return cp, nil
}

// LoadLatest returns the most recently updated checkpoint for agentID, or
// ErrNotFound when the agent has no checkpoints.
func (s *Store) LoadLatest(ctx context.Context, agentID string) (*Checkpoint, error) {
	cp, err := scanCheckpoint(s.db.QueryRowContext(ctx, loadLatestSQL, agentID))
	if errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("persist: load_latest(%s): %w", agentID, ErrNotFound)
	}
	if err != nil {
		return nil, fmt.Errorf("persist: load_latest(%s): %w", agentID, err)
	}
	return cp, nil
}

// List returns every checkpoint for agentID in ascending step order. It returns
// an empty (non-nil) slice when the agent has no checkpoints.
func (s *Store) List(ctx context.Context, agentID string) ([]Checkpoint, error) {
	rows, err := s.db.QueryContext(ctx, listSQL, agentID)
	if err != nil {
		return nil, fmt.Errorf("persist: list(%s): %w", agentID, err)
	}
	defer func() { _ = rows.Close() }()

	checkpoints := make([]Checkpoint, 0)
	for rows.Next() {
		cp, err := scanCheckpoint(rows)
		if err != nil {
			return nil, fmt.Errorf("persist: list(%s): scan: %w", agentID, err)
		}
		checkpoints = append(checkpoints, *cp)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("persist: list(%s): %w", agentID, err)
	}
	return checkpoints, nil
}

// Delete removes the checkpoint with the given id. It returns ErrNotFound when
// no checkpoint with that id existed.
func (s *Store) Delete(ctx context.Context, id string) error {
	res, err := s.db.ExecContext(ctx, deleteSQL, id)
	if err != nil {
		return fmt.Errorf("persist: delete(%s): %w", id, err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return fmt.Errorf("persist: delete(%s): rows affected: %w", id, err)
	}
	if n == 0 {
		return fmt.Errorf("persist: delete(%s): %w", id, ErrNotFound)
	}
	s.log.Debug("checkpoint deleted", "id", id)
	return nil
}

// rowScanner is the subset of *sql.Row/*sql.Rows that scanCheckpoint needs,
// allowing it to serve both a single-row and a multi-row query.
type rowScanner interface {
	Scan(dest ...any) error
}

// scanCheckpoint scans one checkpoint row following the fixed column order of
// the SELECT statements. It returns sql.ErrNoRows untouched so callers can
// translate it to ErrNotFound.
func scanCheckpoint(sc rowScanner) (*Checkpoint, error) {
	var (
		id              string
		agentID         string
		step            int
		state           string
		createdAt, upAt string
	)
	if err := sc.Scan(&id, &agentID, &step, &state, &createdAt, &upAt); err != nil {
		return nil, err
	}

	createdTime, err := parseTime(createdAt)
	if err != nil {
		return nil, fmt.Errorf("parse created_at: %w", err)
	}
	updatedTime, err := parseTime(upAt)
	if err != nil {
		return nil, fmt.Errorf("parse updated_at: %w", err)
	}
	return &Checkpoint{
		ID:        id,
		AgentID:   agentID,
		Step:      step,
		State:     []byte(state),
		CreatedAt: createdTime,
		UpdatedAt: updatedTime,
	}, nil
}

// discardLogger returns log if non-nil, otherwise a logger that writes nowhere,
// so the Store can always log safely.
func discardLogger(log *slog.Logger) *slog.Logger {
	if log != nil {
		return log
	}
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}