feat(persist): durable checkpoint(SQLite upsert/LoadLatest/List/Delete,ErrNotFound)
This commit is contained in:
1 parent
ecd3cdf28c
commit
011d25a53d
2 files changed
+652
No files matched your search
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user