Checkpointing and Recovery
347 Steps Inā
It is 11:47 PM. Your agent has been running for two hours, systematically processing 500 research papers - extracting key findings, cross-referencing citations, and building a structured knowledge graph. It is on paper 347 when the OpenAI API returns a 429 rate limit error. Your retry logic exhausts its attempts. The agent crashes.
Without checkpointing: you start over tomorrow morning. Two hours of API calls, thousands of tokens, real money - gone.
With checkpointing: you restart the agent, it reads its last saved state, and resumes from paper 347. The first 346 papers are untouched.
This is not a hypothetical. Any task that runs for more than a few minutes will, in production, encounter failures: rate limits, network timeouts, OOM errors, power outages, deployment restarts, or simply bugs in your own code. Checkpointing is not optional for long-horizon agents - it is the difference between a viable system and a toy.
:::tip š® Interactive Playground Visualize this concept: Try the Agent Checkpointing & Recovery demo on the EngineersOfAI Playground - no code required. :::
What to Checkpointā
The agent's full state at any point in execution includes:
Not all of this must be checkpointed at the same frequency:
| State Component | Checkpoint Frequency | Rationale |
|---|---|---|
| Completed task results | After every task | Losing this requires rerunning the task |
| Task graph status | After every task | Determines what to run next |
| Tool call history | After every tool call | Fine-grained audit + idempotency |
| Memory state | Every N tasks | Memory is often reconstructible |
| Configuration | At start only | Does not change during execution |
Checkpoint Granularityā
After every tool call (finest granularity):
- Pros: maximum recovery resolution; never redo more than one tool call
- Cons: highest storage and I/O overhead; can slow down fast tool calls significantly
After every completed step (medium granularity):
- Pros: good balance of safety and performance
- Cons: if a multi-tool-call step fails partway through, you redo the whole step
At milestone boundaries (coarsest granularity):
- Pros: minimal overhead; makes sense for naturally phased tasks
- Cons: a failure within a long phase means redoing the entire phase
Recommendation: default to after-every-step. Use after-every-tool-call only for expensive operations (large LLM calls, slow API requests). Use milestone checkpointing for tightly coupled step groups where partial completion is meaningless.
Recovery Strategiesā
When an agent restarts and finds a checkpoint:
Exact resume: load the full state, jump to the first incomplete task, continue from there. Works when the environment has not changed between crash and restart (e.g., local file system tasks).
Fuzzy resume: reprocess from the last verified good state, re-validating results along the way. Use when the environment might have changed (e.g., data sources updated, APIs changed).
Restart with context: start a new execution but inject the checkpoint as rich context: "Here is what we completed, here are the key results, continue from this point." The new execution treats the prior results as given facts rather than re-executing them. Most flexible, useful when the checkpoint data is rich enough to inform a fresh plan.
Idempotency: The Checkpoint-Recovery Contractā
A checkpoint system is only useful if the actions you replay are idempotent - calling the same action twice with the same inputs produces the same result without side effects.
Common violations:
- Writing to a file without checking if it exists (creates duplicate content)
- Inserting into a database without checking for existing records (duplicate rows)
- Sending an API request that triggers an email or payment (sends it twice)
- Creating cloud resources without checking if they already exist
Idempotency patterns:
# BAD: not idempotent
def write_result(path: str, content: str):
with open(path, "a") as f:
f.write(content) # Will append duplicates on replay
# GOOD: idempotent
def write_result_idempotent(path: str, content: str):
if not os.path.exists(path):
with open(path, "w") as f:
f.write(content)
# If file exists, assume it was written correctly
# BAD: not idempotent
def insert_record(db, record: dict):
db.execute("INSERT INTO results VALUES (?)", record)
# GOOD: idempotent (upsert)
def insert_record_idempotent(db, record: dict):
db.execute(
"INSERT OR REPLACE INTO results (id, data) VALUES (?, ?)",
(record["id"], record["data"])
)
The rule: every agent action should be safe to replay. If an action cannot be made idempotent (e.g., sending an email), checkpoint before it and wrap it in a "has this been done?" guard that checks external state.
Full Implementation: Checkpoint Systemā
"""
checkpointing.py
Production-grade checkpoint and recovery system for long-horizon agents.
Uses SQLite for structured state and can be extended to S3 for large artifacts.
Requirements:
pip install pydantic aiosqlite
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import sqlite3
import time
import uuid
from contextlib import contextmanager
from enum import Enum
from pathlib import Path
from typing import Any, Generator, Optional
from pydantic import BaseModel, Field
logger = logging.getLogger(__name__)
# āāā State Models āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā
class StepStatus(str, Enum):
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
SKIPPED = "skipped"
class ToolCall(BaseModel):
"""Record of a single tool invocation."""
id: str = Field(default_factory=lambda: str(uuid.uuid4())[:12])
tool_name: str
arguments: dict[str, Any]
result: Optional[Any] = None
error: Optional[str] = None
timestamp: float = Field(default_factory=time.time)
duration_ms: Optional[float] = None
class Step(BaseModel):
"""A unit of work within the agent run."""
id: str
title: str
description: str
status: StepStatus = StepStatus.PENDING
result: Optional[str] = None
error: Optional[str] = None
tool_calls: list[ToolCall] = Field(default_factory=list)
dependencies: list[str] = Field(default_factory=list)
started_at: Optional[float] = None
completed_at: Optional[float] = None
attempt_count: int = 0
metadata: dict[str, Any] = Field(default_factory=dict)
def is_ready(self, completed_ids: set[str]) -> bool:
return all(dep in completed_ids for dep in self.dependencies)
@property
def duration_seconds(self) -> Optional[float]:
if self.started_at and self.completed_at:
return self.completed_at - self.started_at
return None
class AgentRun(BaseModel):
"""Complete state of an agent execution run."""
run_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
goal: str
steps: dict[str, Step] = Field(default_factory=dict)
created_at: float = Field(default_factory=time.time)
updated_at: float = Field(default_factory=time.time)
status: str = "running" # running, completed, failed, paused
total_cost_usd: float = 0.0
total_tokens: int = 0
replan_count: int = 0
metadata: dict[str, Any] = Field(default_factory=dict)
def add_step(self, step: Step) -> None:
self.steps[step.id] = step
self.updated_at = time.time()
def get_ready_steps(self) -> list[Step]:
completed = {sid for sid, s in self.steps.items() if s.status == StepStatus.COMPLETED}
return [
s for s in self.steps.values()
if s.status == StepStatus.PENDING and s.is_ready(completed)
]
def is_complete(self) -> bool:
return all(
s.status in (StepStatus.COMPLETED, StepStatus.FAILED, StepStatus.SKIPPED)
for s in self.steps.values()
)
@property
def completed_count(self) -> int:
return sum(1 for s in self.steps.values() if s.status == StepStatus.COMPLETED)
@property
def failed_count(self) -> int:
return sum(1 for s in self.steps.values() if s.status == StepStatus.FAILED)
# āāā SQLite Checkpoint Store āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā
class CheckpointStore:
"""
SQLite-backed checkpoint storage.
Stores full agent run state with fine-grained step-level checkpoints.
"""
def __init__(self, db_path: str = "agent_checkpoints.db"):
self.db_path = db_path
self._init_db()
def _init_db(self) -> None:
"""Initialize database schema."""
with self._conn() as conn:
conn.executescript("""
CREATE TABLE IF NOT EXISTS runs (
run_id TEXT PRIMARY KEY,
goal TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'running',
state_json TEXT NOT NULL,
created_at REAL NOT NULL,
updated_at REAL NOT NULL
);
CREATE TABLE IF NOT EXISTS step_checkpoints (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL,
step_id TEXT NOT NULL,
status TEXT NOT NULL,
result TEXT,
error TEXT,
checkpoint_at REAL NOT NULL,
FOREIGN KEY (run_id) REFERENCES runs(run_id),
UNIQUE(run_id, step_id, status)
);
CREATE TABLE IF NOT EXISTS tool_call_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL,
step_id TEXT NOT NULL,
tool_call_json TEXT NOT NULL,
logged_at REAL NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_runs_status ON runs(status);
CREATE INDEX IF NOT EXISTS idx_step_checkpoints_run ON step_checkpoints(run_id);
""")
@contextmanager
def _conn(self) -> Generator[sqlite3.Connection, None, None]:
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.close()
def save_run(self, run: AgentRun) -> None:
"""Save or update the full run state."""
run.updated_at = time.time()
with self._conn() as conn:
conn.execute("""
INSERT OR REPLACE INTO runs (run_id, goal, status, state_json, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?)
""", (
run.run_id,
run.goal,
run.status,
run.model_dump_json(),
run.created_at,
run.updated_at,
))
def save_step_checkpoint(self, run_id: str, step: Step) -> None:
"""Checkpoint an individual step's state."""
with self._conn() as conn:
conn.execute("""
INSERT OR REPLACE INTO step_checkpoints
(run_id, step_id, status, result, error, checkpoint_at)
VALUES (?, ?, ?, ?, ?, ?)
""", (
run_id,
step.id,
step.status,
step.result,
step.error,
time.time(),
))
def log_tool_call(self, run_id: str, step_id: str, tool_call: ToolCall) -> None:
"""Log a tool call for audit and idempotency checking."""
with self._conn() as conn:
conn.execute("""
INSERT INTO tool_call_log (run_id, step_id, tool_call_json, logged_at)
VALUES (?, ?, ?, ?)
""", (run_id, step_id, tool_call.model_dump_json(), time.time()))
def load_run(self, run_id: str) -> Optional[AgentRun]:
"""Load a full run state by ID."""
with self._conn() as conn:
row = conn.execute(
"SELECT state_json FROM runs WHERE run_id = ?", (run_id,)
).fetchone()
if row:
return AgentRun.model_validate_json(row["state_json"])
return None
def find_resumable_runs(self) -> list[dict]:
"""List all runs that can be resumed (status = running or paused)."""
with self._conn() as conn:
rows = conn.execute("""
SELECT run_id, goal, status, updated_at,
(SELECT COUNT(*) FROM step_checkpoints WHERE run_id = r.run_id AND status = 'completed') as completed_steps
FROM runs r
WHERE status IN ('running', 'paused')
ORDER BY updated_at DESC
""").fetchall()
return [dict(row) for row in rows]
def get_tool_call_history(self, run_id: str, step_id: str) -> list[ToolCall]:
"""Retrieve all tool calls for a step (for idempotency checking)."""
with self._conn() as conn:
rows = conn.execute("""
SELECT tool_call_json FROM tool_call_log
WHERE run_id = ? AND step_id = ?
ORDER BY logged_at
""", (run_id, step_id)).fetchall()
return [ToolCall.model_validate_json(row["tool_call_json"]) for row in rows]
def mark_run_complete(self, run_id: str, status: str = "completed") -> None:
"""Mark a run as finished."""
with self._conn() as conn:
conn.execute(
"UPDATE runs SET status = ?, updated_at = ? WHERE run_id = ?",
(status, time.time(), run_id)
)
# āāā Checkpoint-Aware Agent āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā
class CheckpointedAgent:
"""
An agent that automatically checkpoints its state after each step
and can resume from any saved checkpoint.
"""
def __init__(
self,
checkpoint_store: CheckpointStore,
checkpoint_every_n_steps: int = 1, # checkpoint after every N steps
):
self.store = checkpoint_store
self.checkpoint_every_n = checkpoint_every_n_steps
self._steps_since_checkpoint = 0
def start_run(self, goal: str, steps: list[Step]) -> AgentRun:
"""Initialize a new run and save the initial checkpoint."""
run = AgentRun(goal=goal)
for step in steps:
run.add_step(step)
self.store.save_run(run)
logger.info(f"Started run {run.run_id} with {len(steps)} steps")
print(f"\nš Started run {run.run_id}")
print(f" Goal: {goal[:60]}")
print(f" Steps: {len(steps)}")
return run
def resume_run(self, run_id: str) -> Optional[AgentRun]:
"""
Resume a previously interrupted run.
Returns the loaded run if found, None otherwise.
"""
run = self.store.load_run(run_id)
if run is None:
logger.warning(f"No run found with ID {run_id}")
return None
completed = run.completed_count
total = len(run.steps)
print(f"\nā»ļø Resuming run {run_id}")
print(f" Goal: {run.goal[:60]}")
print(f" Progress: {completed}/{total} steps completed")
return run
def execute_run(self, run: AgentRun, executor) -> AgentRun:
"""
Execute all pending steps in the run, checkpointing after each.
The executor is a callable: executor(step, tool_call_history) -> str
"""
print(f"\n{'ā'*50}")
print(f"Executing run: {run.run_id[:12]}...")
while not run.is_complete():
ready = run.get_ready_steps()
if not ready:
pending = [s for s in run.steps.values() if s.status == StepStatus.PENDING]
if pending:
logger.error("Deadlock: pending steps with all dependencies failed")
run.status = "failed"
break
for step in ready:
self._execute_step(run, step, executor)
# Checkpoint after step
self._steps_since_checkpoint += 1
if self._steps_since_checkpoint >= self.checkpoint_every_n:
self.store.save_run(run)
self._steps_since_checkpoint = 0
logger.debug(f"Checkpoint saved at step {step.id}")
run.status = "completed" if run.failed_count == 0 else "completed_with_errors"
self.store.save_run(run)
self.store.mark_run_complete(run.run_id, run.status)
print(f"\n{'ā'*50}")
print(f"Run complete: {run.run_id[:12]}")
print(f" Completed: {run.completed_count} / {len(run.steps)}")
print(f" Failed: {run.failed_count}")
return run
def _execute_step(self, run: AgentRun, step: Step, executor) -> None:
"""Execute a single step with checkpointing and error handling."""
step.status = StepStatus.RUNNING
step.started_at = time.time()
step.attempt_count += 1
print(f"\n [{step.id}] {step.title}")
self.store.save_step_checkpoint(run.run_id, step)
# Check tool call history for idempotency
prior_calls = self.store.get_tool_call_history(run.run_id, step.id)
if prior_calls:
print(f" ā” Found {len(prior_calls)} prior tool calls for this step (attempt {step.attempt_count})")
try:
result = executor(step, prior_calls)
step.status = StepStatus.COMPLETED
step.result = result
step.completed_at = time.time()
print(f" ā
Done in {step.duration_seconds:.1f}s")
except Exception as e:
step.status = StepStatus.FAILED
step.error = str(e)
step.completed_at = time.time()
print(f" ā Failed: {str(e)[:80]}")
logger.exception(f"Step {step.id} failed")
self.store.save_step_checkpoint(run.run_id, step)
def record_tool_call(self, run_id: str, step_id: str, tool_call: ToolCall) -> None:
"""Record a tool call for audit and idempotency checking."""
self.store.log_tool_call(run_id, step_id, tool_call)
def pause_run(self, run: AgentRun) -> None:
"""Pause a run - saves state and marks as paused."""
run.status = "paused"
self.store.save_run(run)
self.store.mark_run_complete(run.run_id, "paused")
print(f"\nāøļø Run {run.run_id[:12]} paused. Resume with: agent.resume_run('{run.run_id}')")
# āāā Idempotency Helper āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā
class IdempotentToolRunner:
"""
Wraps tool calls with idempotency checking.
If the same tool + arguments were called successfully before, return the cached result.
"""
def __init__(self, agent: CheckpointedAgent, run_id: str, step_id: str):
self.agent = agent
self.run_id = run_id
self.step_id = step_id
self._history: Optional[list[ToolCall]] = None
def _get_history(self) -> list[ToolCall]:
if self._history is None:
self._history = self.agent.store.get_tool_call_history(self.run_id, self.step_id)
return self._history
def run(self, tool_name: str, arguments: dict, actual_tool_fn) -> Any:
"""
Execute a tool call, returning cached result if the same call was made before.
"""
# Check if we've already called this tool with these args
for prior in self._get_history():
if prior.tool_name == tool_name and prior.arguments == arguments and prior.result is not None:
print(f" š Idempotent hit: {tool_name}({arguments}) ā using cached result")
return prior.result
# First time - actually call the tool
start = time.time()
call = ToolCall(tool_name=tool_name, arguments=arguments)
try:
result = actual_tool_fn(**arguments)
call.result = result
call.duration_ms = (time.time() - start) * 1000
except Exception as e:
call.error = str(e)
self.agent.record_tool_call(self.run_id, self.step_id, call)
raise
self.agent.record_tool_call(self.run_id, self.step_id, call)
return result
# āāā Demo āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā
def simulate_document_processing():
"""
Simulates processing 10 documents with an intentional failure at step 5
to demonstrate checkpoint-and-resume.
"""
store = CheckpointStore(db_path=":memory:") # Use file path for persistence
agent = CheckpointedAgent(store, checkpoint_every_n_steps=1)
# Build steps
steps = [
Step(
id=f"doc_{i:03d}",
title=f"Process document {i:03d}",
description=f"Extract entities and relationships from document {i:03d}",
dependencies=[f"doc_{i-1:03d}"] if i > 0 else [],
)
for i in range(10)
]
run = agent.start_run(
goal="Process 10 research documents and build knowledge graph",
steps=steps,
)
run_id = run.run_id
# Executor that fails at step 5
call_count = 0
def mock_executor(step: Step, prior_calls: list[ToolCall]) -> str:
nonlocal call_count
call_count += 1
time.sleep(0.01) # Simulate work
doc_num = int(step.id.split("_")[1])
if doc_num == 5 and step.attempt_count == 1:
raise ConnectionError("API rate limit exceeded (simulated failure)")
return f"Extracted 12 entities, 8 relationships from doc_{doc_num:03d}"
# First run - will fail at step 5
print("\n=== FIRST RUN (will fail at step 5) ===")
try:
agent.execute_run(run, mock_executor)
except Exception as e:
print(f"\nā ļø Run interrupted: {e}")
# Show resumable runs
print("\n=== RESUMABLE RUNS ===")
resumable = store.find_resumable_runs()
for r in resumable:
print(f" Run {r['run_id'][:12]}: {r['completed_steps']} steps done - {r['goal'][:40]}")
# Resume from checkpoint
print("\n=== RESUMING FROM CHECKPOINT ===")
resumed_run = agent.resume_run(run_id)
if resumed_run:
agent.execute_run(resumed_run, mock_executor)
print("\nFinal state:")
print(f" Completed: {resumed_run.completed_count} / {len(resumed_run.steps)}")
print(f" Status: {resumed_run.status}")
def demonstrate_idempotency():
"""Shows how IdempotentToolRunner prevents duplicate API calls on replay."""
store = CheckpointStore(db_path=":memory:")
agent = CheckpointedAgent(store)
run = AgentRun(goal="demo")
store.save_run(run)
step = Step(id="s1", title="Fetch data from API")
call_count = 0
def expensive_api_call(query: str) -> str:
nonlocal call_count
call_count += 1
return f"Result for: {query}"
runner = IdempotentToolRunner(agent, run.run_id, step.id)
# First call - actually calls the tool
result1 = runner.run("search", {"query": "Python performance"}, expensive_api_call)
print(f"\nFirst call result: {result1} (total real calls: {call_count})")
# Second call with same args - returns cached result, does NOT call the tool
result2 = runner.run("search", {"query": "Python performance"}, expensive_api_call)
print(f"Second call result: {result2} (total real calls: {call_count}) ā no new API call!")
if __name__ == "__main__":
print("=== CHECKPOINT AND RECOVERY DEMO ===\n")
simulate_document_processing()
demonstrate_idempotency()
Checkpoint-Recovery State Machineā
Distributed Agents: Checkpointing Across Workersā
When multiple workers execute parts of the same task graph:
- Each worker checkpoints its own steps to a shared store (PostgreSQL, not SQLite)
- A coordinator process reads the shared state to determine what to schedule next
- Distributed locking ensures two workers do not claim the same step simultaneously
- The coordinator's own state must be checkpointed separately
For most use cases, SQLite with WAL mode is sufficient for single-process agents. Switch to PostgreSQL when you need true multi-process concurrent access.
Production Notesā
:::warning Checkpoint Store as Single Point of Failure If your checkpoint store fails, you lose the ability to recover. Use at minimum SQLite with WAL mode for durability, or PostgreSQL for production. For critical runs, mirror checkpoints to S3 every N steps as a backup. :::
:::danger Partial Checkpoint Corruption
Always write checkpoints atomically (use a transaction). A partial write during a crash is worse than no checkpoint - the agent thinks state X was saved when it was not. SQLite transactions are atomic by default. Verify with PRAGMA integrity_check after writing.
:::
Memory vs. disk: for short runs (< 5 minutes), in-memory checkpoints (a simple dict) are fine. For anything longer, use SQLite. For runs measured in hours or days, use PostgreSQL with a connection pool.
Checkpoint size: for long-running agents that process large documents, store only summaries in the checkpoint, not full content. Keep a reference (S3 path, file path) to the full artifact instead.
Interview Questions and Answersā
Q: Why is checkpointing essential for long-horizon agents, and what are the key design decisions?
A: Long-horizon agents run for extended periods during which failures (network errors, rate limits, crashes, OOM) are near-certain. Without checkpointing, any failure requires restarting from scratch, wasting time and money. Key design decisions are: granularity (after every tool call vs every step vs every milestone - I default to every step), storage backend (SQLite for single-agent, PostgreSQL for distributed), what state to persist (task graph, results, tool call history), and recovery strategy (exact resume vs fuzzy resume). The most important decision is granularity - too fine increases overhead, too coarse wastes work on recovery.
Q: What is idempotency and why does it matter for checkpointing?
A: Idempotency means calling the same action twice with the same inputs produces the same result with no extra side effects. It matters for checkpointing because recovery involves re-executing the step that was interrupted - if the step is not idempotent, you might write duplicate data, send duplicate messages, or charge customers twice. I enforce idempotency by logging every tool call before execution and checking the log before re-calling. For external APIs, I use the tool call ID as an idempotency key if the API supports it.
Q: How do you handle checkpoint corruption or invalid state on recovery?
A: First, use atomic writes - SQLite transactions ensure partial writes cannot happen. Second, validate the loaded state before using it: check that the task graph is internally consistent (no cycles, referenced dependencies exist), check that step statuses are coherent (no COMPLETED step missing a result). If validation fails, fall back to the previous checkpoint (I keep N=3 checkpoint snapshots). If all checkpoints are corrupt, restart from scratch with the completed step summaries as context, treating them as given facts rather than re-executing.
Q: How would you extend this system to support distributed execution across multiple workers?
A: Replace SQLite with PostgreSQL. Add a worker_id column to steps with a CLAIMED_AT timestamp. Workers use SELECT FOR UPDATE SKIP LOCKED to claim unclaimed ready steps - this provides distributed locking without a separate lock manager. A coordinator process monitors the PostgreSQL state and handles worker failures by releasing steps that have been claimed but not completed within a timeout. The coordinator's own state is checkpointed to PostgreSQL. I would use connection pooling (asyncpg + SQLAlchemy async) to handle many concurrent workers efficiently.
Q: What is the difference between exact resume and restart with context?
A: Exact resume loads the full serialized agent state and continues execution from the interrupted point - the agent does not know it was ever stopped. This works perfectly when the environment is stable (same files, same API state). Restart with context starts a fresh agent but injects the checkpoint as rich context: "here are the 347 documents we already processed with these results." The new agent generates a fresh continuation plan informed by those results. Use exact resume for automation pipelines where the environment is controlled. Use restart with context when the environment may have changed during the interruption (data updated, APIs changed) or when the checkpoint is too complex to deserialize cleanly across agent versions.
