A single erroneous tool call can poison an agent’s memory and burn through tokens. Agent snapshot rollback gives you a rewind button: capture state before each step, then restore it when the model makes a bad decision. This tutorial builds a minimal, runnable checkpointing layer you can drop into any LLM agent loop.
Prerequisites
- Python 3.10+ (we use
copy.deepcopyanddataclasses) pip install openaionly if you run the optional LLM section- A basic agent loop mental model: thought → action → observation
- No external services required until Step 4
The core idea: immutable snapshots
State mutation is the enemy of debuggability. If you treat each decision as a transaction, you can snapshot the pre-decision state, attempt the action, and either commit or revert. The snapshot must be a true copy, not a reference, or rollback becomes a no-op.
State shape
Keep state flat and serializable. Avoid holding open file handles, sockets, or model clients inside state.
from dataclasses import dataclass, field
import copy
@dataclass
class AgentState:
step: int = 0
memory: list[str] = field(default_factory=list)
last_action: dict | None = None
tokens_used: int = 0
def snapshot(self) -> "AgentState":
return copy.deepcopy(self)
def restore(self, snap: "AgentState") -> None:
self.step = snap.step
self.memory = copy.deepcopy(snap.memory)
self.last_action = copy.deepcopy(snap.last_action)
self.tokens_used = snap.tokens_usedStep 1: A runtime with manual checkpointing
We start with explicit snapshot/rollback to show the mechanics before adding sugar.
class AgentRuntime:
def __init__(self):
self.state = AgentState()
self.history: list[AgentState] = []
def checkpoint(self):
self.history.append(self.state.snapshot())
def rollback(self):
if not self.history:
raise RuntimeError("Nothing to roll back to")
self.state.restore(self.history.pop())
def commit(self):
if self.history:
self.history.pop()Run a bad action and revert:
rt = AgentRuntime()
rt.checkpoint()
rt.state.step = 1
rt.state.memory.append("thought: call search")
rt.state.last_action = {"tool": "search", "query": "..."}
# Later we realize the query was malformed
rt.rollback()
print(rt.state)Expected output:
AgentState(step=0, memory=[], last_action=None, tokens_used=0)Step 2: Transactional decision wrapper
Manual checkpoint/rollback calls leak. Wrap each decision in a context manager so the happy path stays clean and failures auto-revert.
from contextlib import contextmanager
class AgentRuntime:
# ... previous methods ...
@contextmanager
def decision(self):
self.checkpoint()
try:
yield self.state
self.commit()
except Exception:
self.rollback()
raiseNow a crashing decision auto-reverts:
rt = AgentRuntime()
with rt.decision() as s:
s.step = 1
s.memory.append("act")
raise ValueError("tool crashed")
# state is back to initial
print(rt.state.step) # 0Step 3: Validating agent actions
Rollback is only useful if you detect bad decisions. Add a validator that inspects the proposed action before commit.
def validate_action(state: AgentState) -> None:
if state.last_action is None:
return
if state.last_action.get("tool") == "search" and not state.last_action.get("query"):
raise ValueError("search requires query")
@contextmanager
def decision_with_validation(rt: AgentRuntime):
with rt.decision() as s:
yield s
validate_action(s)If validation fails, the context manager’s commit is never reached; the exception triggers rollback. This is agent snapshot rollback doing real work instead of just catching crashes.
Step 4: Hooking in an LLM call
Here’s where an inference gateway matters. When you call a model to pick the next action, you want the call isolated from state mutations. Use the OpenAI client against any OpenAI-compatible endpoint.
from openai import OpenAI
client = OpenAI(base_url="https://api.n4n.ai/v1", api_key="YOUR_KEY")
# n4n.ai provides automatic fallback when a provider is degraded,
# so a snapshot test loop won't crash on upstream 429s.
def plan_action(state: AgentState) -> dict:
resp = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "system", "content": "Output JSON action."},
{"role": "user", "content": f"Memory: {state.memory}"}]
)
# assume parsed action
return {"tool": "search", "query": "example"}Wrap planning inside the decision block:
rt = AgentRuntime()
with rt.decision() as s:
s.last_action = plan_action(s)
s.step += 1
validate_action(s)If plan_action throws or validation fails, the runtime reverts to the pre-call state, including token count. That prevents a failed generation from leaving a half-written last_action in memory.
Step 5: Persisting snapshots
In-memory history is lost on restart. Write snapshots to JSONL for crash recovery.
import json
def save_snapshot(path: str, snap: AgentState):
with open(path, "a") as f:
f.write(json.dumps(snap.__dict__) + "\n")
def load_latest(path: str) -> AgentState | None:
try:
with open(path) as f:
lines = f.readlines()
if not lines:
return None
return AgentState(**json.loads(lines[-1]))
except FileNotFoundError:
return NoneIntegrate into checkpoint:
class PersistentRuntime(AgentRuntime):
def __init__(self, snap_path: str):
super().__init__()
self.snap_path = snap_path
restored = load_latest(snap_path)
if restored:
self.state = restored
def checkpoint(self):
super().checkpoint()
save_snapshot(self.snap_path, self.state.snapshot())Step 6: Observing rollback in a loop
A full agent loop with random failures demonstrates recovery.
import random
rt = PersistentRuntime("agent.snaps.jsonl")
for i in range(5):
try:
with rt.decision() as s:
s.step += 1
if random.random() < 0.3:
raise RuntimeError("simulated bad decision")
s.memory.append(f"ok step {s.step}")
s.tokens_used += 100
except RuntimeError:
print(f"rolled back at step {i}")
print("final step:", rt.state.step)Sample output (varies):
rolled back at step 1
rolled back at step 3
final step: 5Because rolled-back steps never incremented state.step persistently, the counter reflects only committed progress.
Why not just try/except?
A bare try/except around your agent code catches errors but does not undo mutations that already happened before the throw. If you appended to memory or incremented tokens_used prior to the failure, those changes stick. Agent snapshot rollback reverts the entire state object, giving you a clean slate.
Side effects are the hard part
Rollback works perfectly for in-process state. It cannot un-send an email or reverse a database write. For external side effects, use a compensating action (saga pattern): record the intent in state, execute, and on rollback issue a counter-command. Snapshot the state that tracks pending compensations.
Testing your rollback logic
Write a test that forces a failure and asserts state invariance.
def test_rollback_restores_state():
rt = AgentRuntime()
rt.state.memory.append("baseline")
with rt.decision() as s:
s.memory.append("temp")
raise AssertionError("force")
assert rt.state.memory == ["baseline"]
assert rt.history == []Run with pytest. If this fails, your deep copy is shallow somewhere.
Production considerations
- Concurrency: If multiple workers mutate state, use a lock or single-writer pattern. Snapshots are cheap but deep copies of large memory lists will stall the event loop.
- Pruning: Cap
historylength; keep only last N or store diffs instead of full copies. - Determinism: Rollback restores exact bytes. If your agent uses randomness, seed from
state.stepso replay is reproducible. - LLM cost: Token counts should live inside state so rollback also refunds estimated spend in your ledger.
- Observability: Emit a metric
agent_rollbacks_totalper run. A spike means your validator or model is misbehaving.
Agent snapshot rollback is not a debug trick; it’s a control plane for autonomous systems. Ship it before your agent touches production traffic.