The state store: one table, conditional writes
Everything the platform remembers about a running job (its events, its budget, which worker holds which unit) has to live somewhere durable, concurrent, and fast. Part 1 built that store as a JSON-lines file on one laptop; a fleet of workers on many machines needs the same design with real durability and real concurrency control. This chapter is that upgrade. The event-sourcing design does not change at all; only the storage engine grows up, and the engine this book reaches for is DynamoDB, for reasons that are about shape, not brand loyalty.
Why a key-value store, not SQL
An agent state store has an unusual load profile, and naming it explains the choice. The writes are almost all appends (new events, never edited). The reads are almost all "give me everything for run X," a single-partition scan. The concurrency need is narrow but sharp: two workers must never both claim the same sequence number or the same work unit. And the scale is wide and spiky: a quiet afternoon and a five-hundred-worker night hit the same table.
That profile wants predictable single-digit-millisecond operations at any scale far more than it wants SQL's flexible joins and ad-hoc queries, which it will essentially never use. DynamoDB is built for exactly this trade, and its lineage is the reason: it descends from Amazon's 2007 Dynamo paper, the internal store born from shopping-cart outages whose whole thesis was giving up relational flexibility to buy predictable performance and availability at scale. When your access pattern is "look up by key, append by key, scan one partition," you are the customer that paper was written for. (If you already run Postgres and your scale is modest, Postgres does this job fine; the design below ports directly. The choice is workload, not dogma.)
Single-table design
The move that surprises people coming from SQL: one table holds
every entity type. Not a runs table, an events table, a
budgets table, and a leases table, but one table whose rows are
told apart by their keys. Each row has a partition key (which
logical group it belongs to) and a sort key (its position within
that group), and different entity types use different key patterns in
the same table:
partition key sort key the row is...
run#42 meta run metadata (task, status)
run#42 budget the budget ledger (cap, spent)
run#42 evt#000001 an event
run#42 evt#000002 an event
...
unit#7 lease who holds work-unit 7, until when
The payoff is the access pattern falling out for free: "everything
about run 42" is a single query on partition key run#42, and because
the sort keys sort lexically, the events come back in order and sit
contiguously after the metadata. One query, one partition, the whole
story of a run, which is exactly what resume, audit, and the ops
dashboards all need. Designing the keys is designing the queries, and
in a single-table world that is most of the schema work.
The two conditional writes that make it a control plane
A dumb store would let any write land. A control plane needs writes that can refuse, and DynamoDB's conditional writes (put or update only if the item's current state satisfies a condition) are the primitive the whole platform's correctness rests on. Two uses carry almost everything.
Sequence integrity. Append event N only if no item already exists
at that sort key. Two workers racing to write evt#000005 cannot both
win: the first succeeds, the second's condition ("this key is absent")
now fails, and it retries at the next free sequence. No lock server, no
coordination round trip, just a conditional put that the losing writer
handles as a normal retry.
Leases. Chapter 25 promised leases as the mechanism that distinguishes "worker is slow" from "worker is gone" without asking the worker. Here is the mechanism: a lease is a row with an owner and an expiry, acquired by a conditional write whose condition is "no lease exists, or the existing one has expired." One worker acquires it; a second is refused while it is live; and if the holder crashes, its lease simply ages out and the next worker's conditional acquire succeeds. Crash recovery with no crash detector, which is the only kind that works when the thing that crashed is the thing you would have asked.
Lab 7.1: the store, simulated
DynamoDB semantics are small enough to model faithfully in a local
table (a dict keyed by partition-and-sort, with put_if_absent and
put_if conditional operations), so the design is visible without an
AWS account. The lab seeds a run, then exercises the three behaviors:
python3 agent_state.py
--- one table, one query: the run's whole story ---
run#42 budget {'cap': 500000, 'spent': 0}
run#42 evt#000001 {'type': 'run.created'}
run#42 evt#000002 {'type': 'model.responded'}
run#42 evt#000003 {'type': 'tool.returned'}
run#42 evt#000004 {'type': 'run.completed'}
run#42 meta {'task': 'audit checkout', 'status': 'running'}
--- conditional append: two workers race for seq 5 ---
worker A: append seq 5 -> ok
worker B: append seq 5 -> ConditionalCheckFailed (retries at seq 6)
worker B: append seq 6 -> ok
--- leases: acquire, refuse, reclaim after a crash ---
t=0 worker A acquires unit#7 : True
t=10 worker B tries unit#7 : False (A's lease is live; B is refused)
... worker A crashes; its lease keeps ticking toward expiry
t=70 worker B reclaims unit#7 : True (lease expired at t=60; B wins)
unit#7 lease now held by B until t=130
Read the three blocks as the chapter's three claims, verified. The
first shows single-table design working: metadata, budget, and an
ordered event stream, all one entity, all returned by one query
(note they come back sorted by sort key, which is why budget and
meta bracket the events). The second shows sequence integrity: the
colliding writer is refused and recovers by retrying, no lock in
sight. The third shows the lease lifecycle end to end: acquired,
refused while live, reclaimed after the holder's crash lets it expire,
which is Chapter 25's worker-death story
reduced to three conditional writes and a clock.
What the real service adds
The simulation is faithful to the semantics; production adds
durability and scale the dict cannot. TTL auto-expires old rows
(finished runs, dead leases) so the table self-cleans. Atomic
counters let the budget ledger's spent field increment
conditionally in one operation ("add this cost only if it keeps us
under cap"), which is admission control as a single write. Streams
emit a change feed other services can react to. And the partition-key
design carries a scaling caveat worth knowing early: a single
partition key is a throughput unit, so a run so hot that all its
writes hammer one run#id can throttle, which is why very high-write
entities sometimes shard their key (run#42#shard3). None of this
changes the model; it is the model at fleet scale.
Don't be confused: the state store vs memory. This chapter's store holds run state: events, budgets, leases, the machinery of jobs in flight, wanted with perfect fidelity and discarded when the run's relevance ends. The next two chapters build memory: distilled knowledge that outlives any run. They use different engines for different reasons (a key-value store for exact keyed access here; a vector index for similarity search there), and the state-vs-memory distinction from Part 3 is why a platform runs both rather than forcing one to do the other's job.
Full source
"""Lab 7.1: the agent state store, single-table, with conditional writes.
DynamoDB semantics, simulated locally so the design is visible without
an AWS account: one table holds every entity type for a run (metadata,
events, budget, unit leases), keyed by a partition key and a sort key.
The two operations that make it a control plane are conditional writes
(succeed only if the item is absent, or matches an expected value):
1. event append that cannot collide on a sequence number,
2. a work-unit lease that one worker wins, another is refused, and a
third reclaims after the holder's crash lets it expire.
Standard library only. Simulated clock. Deterministic.
"""
from __future__ import annotations
class ConditionalCheckFailed(Exception):
"""Raised when a conditional write's precondition is not met."""
class Table:
"""A single table: {(pk, sk): item}. Mirrors DynamoDB's core ops."""
def __init__(self):
self._items: dict[tuple[str, str], dict] = {}
def put_if_absent(self, pk: str, sk: str, item: dict) -> None:
if (pk, sk) in self._items:
raise ConditionalCheckFailed(f"{pk}/{sk} already exists")
self._items[(pk, sk)] = {"pk": pk, "sk": sk, **item}
def put_if(self, pk: str, sk: str, item: dict,
expect) -> None:
"""Write only if the current item satisfies expect(current|None)."""
current = self._items.get((pk, sk))
if not expect(current):
raise ConditionalCheckFailed(f"{pk}/{sk} precondition failed")
self._items[(pk, sk)] = {"pk": pk, "sk": sk, **item}
def query(self, pk: str, sk_prefix: str = "") -> list[dict]:
"""All items under one partition key, sorted by sort key."""
rows = [v for (p, s), v in self._items.items()
if p == pk and s.startswith(sk_prefix)]
return sorted(rows, key=lambda r: r["sk"])
# ---------------------------------------------------------------------------
# 1. One table, many entity types, one query returns a run's whole story
# ---------------------------------------------------------------------------
def seed_run(t: Table, run: str) -> None:
pk = f"run#{run}"
t.put_if_absent(pk, "meta", {"task": "audit checkout", "status": "running"})
t.put_if_absent(pk, "budget", {"cap": 500_000, "spent": 0})
def append_event(t: Table, run: str, seq: int, etype: str) -> None:
"""seq is unique per run: two writers cannot both claim it."""
t.put_if_absent(f"run#{run}", f"evt#{seq:06d}", {"type": etype})
# ---------------------------------------------------------------------------
# 2. Leases: one owner at a time, reclaimable after expiry
# ---------------------------------------------------------------------------
def acquire_lease(t: Table, unit: str, owner: str, now: int,
ttl: int = 60) -> bool:
"""Acquire only if no lease exists, or the existing one has expired."""
def free_or_expired(cur):
return cur is None or cur["expires"] <= now
try:
t.put_if(f"unit#{unit}", "lease",
{"owner": owner, "expires": now + ttl},
expect=free_or_expired)
return True
except ConditionalCheckFailed:
return False
if __name__ == "__main__":
t = Table()
seed_run(t, "42")
print("--- one table, one query: the run's whole story ---")
for seq, etype in [(1, "run.created"), (2, "model.responded"),
(3, "tool.returned"), (4, "run.completed")]:
append_event(t, "42", seq, etype)
for row in t.query("run#42"):
detail = {k: v for k, v in row.items() if k not in ("pk", "sk")}
print(f" {row['pk']:<8} {row['sk']:<12} {detail}")
print("\n--- conditional append: two workers race for seq 5 ---")
append_event(t, "42", 5, "tool.called") # worker A wins
print(" worker A: append seq 5 -> ok")
try:
append_event(t, "42", 5, "tool.called") # worker B collides
except ConditionalCheckFailed:
print(" worker B: append seq 5 -> ConditionalCheckFailed "
"(retries at seq 6)")
append_event(t, "42", 6, "tool.called")
print(" worker B: append seq 6 -> ok")
print("\n--- leases: acquire, refuse, reclaim after a crash ---")
print(f" t=0 worker A acquires unit#7 : {acquire_lease(t, '7', 'A', 0)}")
print(f" t=10 worker B tries unit#7 : {acquire_lease(t, '7', 'B', 10)}"
" (A's lease is live; B is refused)")
print(f" ... worker A crashes; its lease keeps ticking toward expiry")
print(f" t=70 worker B reclaims unit#7 : {acquire_lease(t, '7', 'B', 70)}"
" (lease expired at t=60; B wins)")
lease = t.query("unit#7", "lease")[0]
print(f" unit#7 lease now held by {lease['owner']} until "
f"t={lease['expires']}")
👉 Next: vector memory, where the store stops answering "give me run 42" and starts answering "what do we know about X," a different question that needs a different index, built from scratch.