Skip to content

feat: agent state checkpointing (rollback, restore, session resume)#193

Open
EvilFreelancer wants to merge 9 commits into
mainfrom
feat/agent-checkpointing
Open

feat: agent state checkpointing (rollback, restore, session resume)#193
EvilFreelancer wants to merge 9 commits into
mainfrom
feat/agent-checkpointing

Conversation

@EvilFreelancer

Copy link
Copy Markdown
Member

Summary

Adds agent state checkpointing so an agent can snapshot its execution state as it runs, then roll back to an earlier step, restore after a process restart, and resume a session (including over ACP). Checkpointing is opt-in — disabled by default, so existing behavior is unchanged.

Built bottom-up (base → services → config → agent → factory → server → acp), TDD, one layer per commit.

What's included

  • AgentContext.to_snapshot()/from_snapshot() + AgentCheckpoint (models.py) — JSON-serializable snapshot; excludes the live asyncio clarification event and re-resolvable skills; serialize_as_any preserves a Pydantic custom_context.
  • Pluggable CheckpointStore (services/checkpoint_store.py) — InMemoryCheckpointStore (default) and FileCheckpointStore (JSON on disk, survives restarts), each with optional max_history ring buffer; build_checkpoint_store() helper.
  • Config (execution.checkpoint) — enabled / backend (memory|file) / dir / max_history; off by default.
  • BaseAgentcheckpoint(), list_checkpoints(), rollback(step); automatic snapshot at the start of every iteration (crash-safe) and at the clarification pause. Auto-checkpoint failures are logged, never abort a run.
  • AgentFactorycreate(checkpoint_store=…, session_id=…) and restore(checkpoint) (rebinds id + streaming generator, re-applies conversation/context).
  • REST APIGET /agents/{id}/checkpoints, POST /agents/{id}/rollback (cancels a running task first), POST /agents/{id}/restore; store wired from config at startup.
  • ACP — advertises load_session=true, tags per-prompt checkpoints with the session id, and load_session rebuilds the session from its latest checkpoint; the next prompt resumes the restored agent instead of discarding it.
  • Docs & examplesframework/checkpoints.md (en+ru) + nav, three endpoints in the API reference (en+ru), examples/checkpoint_and_rollback.py, and a commented execution.checkpoint block in config.yaml.example.

Testing

  • 654 passed, 6 deselected (.venv/bin/pytest); ~66 new tests across models, store, config, agent, factory, server, ACP.
  • pre-commit run -a clean (ruff, ruff-format, docformatter, mdformat).

Review

An adversarial review pass ran over the full diff; its findings were folded in as a dedicated fix commit:

  • custom_context Pydantic model was silently reduced to {} on snapshot → fixed with serialize_as_any.
  • A checkpoint-store failure could fail an otherwise-healthy run → auto-checkpoint isolated.
  • ACP resume discarded restored state on the next prompt → prompt() now continues the restored agent; a checkpoint is captured at the clarification pause.
  • find_by_session ordered by step across multi-agent sessions → ordered by created_at.
  • Server rollback raced the running execute task → task is cancelled first.
  • FileCheckpointStore.delete() broke on stray non-JSON files → uses rmtree.

Notes / scope

  • Fully backward compatible: with checkpointing disabled the agent takes no snapshots and behaves as before.
  • Restoring an agent requires its definition (def_name) to be present in the current configuration.

🤖 Generated with Claude Code

EvilFreelancer and others added 9 commits July 16, 2026 18:19
Add AgentContext.to_snapshot/from_snapshot (JSON-serializable, excludes the
live clarification event and re-resolvable skills) and the AgentCheckpoint
snapshot model. Foundation for agent state checkpointing.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add BaseCheckpointStore with save/list/get/latest/delete/agent_ids/
find_by_session, an in-memory backend with optional ring-buffer history, and a
JSON file backend that survives process restarts.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add execution.checkpoint with enabled/backend/dir/max_history; defaults keep
checkpointing off so existing behavior is unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
BaseAgent gains checkpoint(), list_checkpoints(), rollback(step) and an
optional checkpoint_store/session_id. When a store is attached the execution
loop snapshots state at the start of every iteration; rollback restores the
conversation and context (with a fresh clarification event). No store = prior
behavior.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
create() now accepts checkpoint_store/session_id. restore() rebuilds a live
agent from a checkpoint: resolves the definition (arg or GlobalConfig by
def_name), rebinds the saved id (and streaming generator id), and re-applies
conversation + context.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add GET /agents/{id}/checkpoints, POST /agents/{id}/rollback and
POST /agents/{id}/restore, a shared checkpoint_store built from
execution.checkpoint at startup, and thread it into agent creation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Advertise load_session=True, accept a checkpoint_store, tag per-prompt
checkpoints with the ACP session id, and implement load_session to rebuild the
session from its latest checkpoint. Wire the store from config in the ACP entry
point.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add framework/checkpoints.md (en+ru) and nav entry, document the three REST
endpoints in SGR-Description-API (en+ru), add examples/checkpoint_and_rollback.py,
and a commented checkpoint block in config.yaml.example. Apply ruff/docformatter.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- models: serialize custom_context with serialize_as_any so a Pydantic
  custom_context is not silently reduced to {} on snapshot
- base_agent: isolate auto-checkpoint failures (never abort a run) and capture
  a checkpoint at the clarification pause so a restore observes the waiting state
- store: order find_by_session by created_at (correct 'latest turn' across
  multi-agent sessions); delete via rmtree (tolerate stray files)
- server: cancel a running execute task before rollback to avoid the race
- acp: prompt() now resumes a restored agent instead of discarding it

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant