diff --git a/internal/runtime/persist/persist.go b/internal/runtime/persist/persist.go new file mode 100644 index 0000000..0fb4ef7 --- /dev/null +++ b/internal/runtime/persist/persist.go @@ -0,0 +1,290 @@ +// 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 = ?` + + loadLatestSQL = `SELECT id, agent_id, step, state, created_at, updated_at + FROM checkpoints WHERE agent_id = ? + ORDER BY updated_at DESC, step 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. +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) + } + 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 +} + +// formatTime renders t as a UTC RFC3339 timestamp with nanosecond precision. +func formatTime(t time.Time) string { + return t.UTC().Format(time.RFC3339Nano) +} + +// parseTime parses a timestamp written by formatTime. +func parseTime(s string) (time.Time, error) { + return time.Parse(time.RFC3339Nano, s) +} + +// 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)) +} diff --git a/internal/runtime/persist/persist_test.go b/internal/runtime/persist/persist_test.go new file mode 100644 index 0000000..d9755c4 --- /dev/null +++ b/internal/runtime/persist/persist_test.go @@ -0,0 +1,362 @@ +package persist + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "sync" + "testing" + "time" + + dbpkg "git.yeij.top/AskaEth/Cyrene_v2/internal/db" +) + +// testLogger returns a logger that writes nowhere, keeping test output clean. +func testLogger() *slog.Logger { + return slog.New(slog.NewTextHandler(io.Discard, nil)) +} + +// newTestStore opens an isolated in-memory SQLite database and returns a Store +// backed by it. Each call owns a private database, so parallel tests do not +// interfere with one another. +func newTestStore(t *testing.T) *Store { + t.Helper() + ctx := context.Background() + + db, err := dbpkg.Open(ctx, ":memory:", testLogger()) + if err != nil { + t.Fatalf("open in-memory db: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + s, err := New(ctx, db, testLogger()) + if err != nil { + t.Fatalf("New(store) error: %v", err) + } + return s +} + +func TestStore_SaveLoad_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + cp := &Checkpoint{ + ID: "cp-1", + AgentID: "agent-a", + Step: 3, + State: []byte(`{"turn":3,"msg":"hello"}`), + } + if err := s.Save(ctx, cp); err != nil { + t.Fatalf("Save error: %v", err) + } + + got, err := s.Load(ctx, "cp-1") + if err != nil { + t.Fatalf("Load error: %v", err) + } + if got.ID != cp.ID || got.AgentID != cp.AgentID || got.Step != cp.Step { + t.Fatalf("round-trip mismatch: got %+v", got) + } + if string(got.State) != string(cp.State) { + t.Fatalf("state mismatch: got %q want %q", got.State, cp.State) + } + if got.CreatedAt.IsZero() || got.UpdatedAt.IsZero() { + t.Fatalf("timestamps not set: created=%v updated=%v", got.CreatedAt, got.UpdatedAt) + } +} + +func TestStore_Load_Missing_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + if _, err := s.Load(ctx, "does-not-exist"); !errors.Is(err, ErrNotFound) { + t.Fatalf("Load(missing) err = %v, want wrap of ErrNotFound", err) + } +} + +func TestStore_LoadLatest_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + // Save steps in ascending order; the latest save must be the newest step. + for step := 1; step <= 3; step++ { + cp := &Checkpoint{ + ID: fmt.Sprintf("cp-%d", step), + AgentID: "agent-a", + Step: step, + State: []byte(fmt.Sprintf(`{"step":%d}`, step)), + } + if err := s.Save(ctx, cp); err != nil { + t.Fatalf("Save(step %d) error: %v", step, err) + } + } + + got, err := s.LoadLatest(ctx, "agent-a") + if err != nil { + t.Fatalf("LoadLatest error: %v", err) + } + if got.Step != 3 { + t.Fatalf("LoadLatest step = %d, want 3", got.Step) + } + if got.ID != "cp-3" { + t.Fatalf("LoadLatest id = %q, want cp-3", got.ID) + } +} + +func TestStore_LoadLatest_Missing_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + if _, err := s.LoadLatest(ctx, "no-such-agent"); !errors.Is(err, ErrNotFound) { + t.Fatalf("LoadLatest(missing) err = %v, want wrap of ErrNotFound", err) + } +} + +func TestStore_List_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + // agent-a gets steps 1 and 2; agent-b gets a single step. + for _, cp := range []*Checkpoint{ + {ID: "a-1", AgentID: "agent-a", Step: 1, State: []byte(`{"step":1}`)}, + {ID: "a-2", AgentID: "agent-a", Step: 2, State: []byte(`{"step":2}`)}, + {ID: "b-1", AgentID: "agent-b", Step: 1, State: []byte(`{"step":1}`)}, + } { + if err := s.Save(ctx, cp); err != nil { + t.Fatalf("Save(%s) error: %v", cp.ID, err) + } + } + + listA, err := s.List(ctx, "agent-a") + if err != nil { + t.Fatalf("List(agent-a) error: %v", err) + } + if len(listA) != 2 { + t.Fatalf("List(agent-a) length = %d, want 2", len(listA)) + } + if listA[0].Step != 1 || listA[1].Step != 2 { + t.Fatalf("List(agent-a) not ascending by step: %+v", listA) + } + + listB, err := s.List(ctx, "agent-b") + if err != nil { + t.Fatalf("List(agent-b) error: %v", err) + } + if len(listB) != 1 { + t.Fatalf("List(agent-b) length = %d, want 1", len(listB)) + } + + // An agent with no checkpoints returns an empty, non-nil slice. + none, err := s.List(ctx, "no-such-agent") + if err != nil { + t.Fatalf("List(none) error: %v", err) + } + if none == nil { + t.Fatalf("List(none) returned nil slice, want empty non-nil") + } + if len(none) != 0 { + t.Fatalf("List(none) length = %d, want 0", len(none)) + } +} + +func TestStore_Save_Upsert_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + // A clearly past timestamp: the store must preserve it across the upsert + // while refreshing UpdatedAt to the current time. + originalCreated := time.Date(2024, time.January, 1, 0, 0, 0, 0, time.UTC) + first := &Checkpoint{ + ID: "cp-1", + AgentID: "agent-a", + Step: 1, + State: []byte(`{"step":1}`), + CreatedAt: originalCreated, + } + if err := s.Save(ctx, first); err != nil { + t.Fatalf("Save(first) error: %v", err) + } + + // A second save with the same ID must keep the original CreatedAt while + // refreshing step, state and UpdatedAt. + second := &Checkpoint{ + ID: "cp-1", + AgentID: "agent-a", + Step: 2, + State: []byte(`{"step":2}`), + } + if err := s.Save(ctx, second); err != nil { + t.Fatalf("Save(second) error: %v", err) + } + + got, err := s.Load(ctx, "cp-1") + if err != nil { + t.Fatalf("Load error: %v", err) + } + if got.Step != 2 { + t.Fatalf("upsert step = %d, want 2", got.Step) + } + if string(got.State) != `{"step":2}` { + t.Fatalf("upsert state = %q, want step 2", got.State) + } + if !got.CreatedAt.Equal(originalCreated) { + t.Fatalf("upsert created_at = %v, want preserved %v", got.CreatedAt, originalCreated) + } + if got.UpdatedAt.Before(got.CreatedAt) { + t.Fatalf("upsert updated_at %v is before created_at %v", got.UpdatedAt, got.CreatedAt) + } + + // Only one row must remain for the ID after the upsert. + all, err := s.List(ctx, "agent-a") + if err != nil { + t.Fatalf("List(after upsert) error: %v", err) + } + if len(all) != 1 { + t.Fatalf("upsert left %d rows, want 1", len(all)) + } +} + +func TestStore_Delete_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + cp := &Checkpoint{ID: "cp-1", AgentID: "agent-a", Step: 1, State: []byte(`{}`)} + if err := s.Save(ctx, cp); err != nil { + t.Fatalf("Save error: %v", err) + } + if err := s.Delete(ctx, "cp-1"); err != nil { + t.Fatalf("Delete error: %v", err) + } + if _, err := s.Load(ctx, "cp-1"); !errors.Is(err, ErrNotFound) { + t.Fatalf("Load after delete err = %v, want ErrNotFound", err) + } +} + +func TestStore_Delete_Missing_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + if err := s.Delete(ctx, "never-existed"); !errors.Is(err, ErrNotFound) { + t.Fatalf("Delete(missing) err = %v, want wrap of ErrNotFound", err) + } +} + +func TestStore_Save_SQLInjection_安全(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + // Malicious strings must be treated as opaque data, never executed as SQL. + cp := &Checkpoint{ + ID: "1'; DROP TABLE checkpoints;--", + AgentID: "agent'; DELETE FROM checkpoints;--", + Step: 1, + State: []byte(`{"x":"'); DROP TABLE checkpoints;--"}`), + } + if err := s.Save(ctx, cp); err != nil { + t.Fatalf("Save(injection) error: %v", err) + } + + got, err := s.Load(ctx, cp.ID) + if err != nil { + t.Fatalf("Load(injection id) error: %v", err) + } + if got.ID != cp.ID || got.AgentID != cp.AgentID { + t.Fatalf("injection round-trip mismatch: got %+v", got) + } + if string(got.State) != string(cp.State) { + t.Fatalf("injection state mismatch: got %q", got.State) + } + + // The table must still exist and hold exactly the injected row. + var tableCount int + if err := s.db.QueryRowContext(ctx, + "SELECT count(*) FROM sqlite_master WHERE type='table' AND name='checkpoints'").Scan(&tableCount); err != nil { + t.Fatalf("query sqlite_master: %v", err) + } + if tableCount != 1 { + t.Fatalf("checkpoints table missing after injection, count=%d", tableCount) + } + var rowCount int + if err := s.db.QueryRowContext(ctx, "SELECT count(*) FROM checkpoints").Scan(&rowCount); err != nil { + t.Fatalf("count checkpoints: %v", err) + } + if rowCount != 1 { + t.Fatalf("checkpoints row count = %d, want 1", rowCount) + } +} + +func TestStore_SaveConcurrent_并发(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + const agents = 8 + const perAgent = 5 + + var wg sync.WaitGroup + errCh := make(chan error, agents*perAgent) + for a := 0; a < agents; a++ { + agentID := fmt.Sprintf("agent-%d", a) + for step := 0; step < perAgent; step++ { + wg.Add(1) + go func(id string, st int) { + defer wg.Done() + cp := &Checkpoint{ + ID: fmt.Sprintf("cp-%s-%d", id, st), + AgentID: id, + Step: st, + State: []byte(fmt.Sprintf(`{"agent":%q,"step":%d}`, id, st)), + } + if err := s.Save(ctx, cp); err != nil { + errCh <- err + } + }(agentID, step) + } + } + wg.Wait() + close(errCh) + for err := range errCh { + t.Fatalf("concurrent save failed: %v", err) + } + + for a := 0; a < agents; a++ { + agentID := fmt.Sprintf("agent-%d", a) + all, err := s.List(ctx, agentID) + if err != nil { + t.Fatalf("List(%s) error: %v", agentID, err) + } + if len(all) != perAgent { + t.Fatalf("List(%s) length = %d, want %d", agentID, len(all), perAgent) + } + } +} + +func TestStore_New_NilDB_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + + if _, err := New(ctx, nil, testLogger()); !errors.Is(err, ErrNilDB) { + t.Fatalf("New(nil db) err = %v, want wrap of ErrNilDB", err) + } +} + +func TestStore_Save_MissingID_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newTestStore(t) + + if err := s.Save(ctx, &Checkpoint{AgentID: "a", Step: 1, State: []byte(`{}`)}); !errors.Is(err, ErrMissingID) { + t.Fatalf("Save(blank id) err = %v, want wrap of ErrMissingID", err) + } +}