Durable Agent Runtime
Build a crash-safe agent runner on PostgreSQL where every model call and tool call is checkpointed, runs resume after a worker dies, budgets and timeouts are enforced, and tokens, latency and cost are observable.
Revise first: Multithreading & ConcurrencyTransactions & PropagationTransactions & Isolation LevelsMicroservices & ResilienceLogging, Monitoring, Tracing, Prometheus & Grafana
An agent is a loop: call the model, run the tools it asks for, feed the results back, repeat until it stops. In a demo that loop lives in one method. In production a run lasts minutes, costs real money, calls flaky APIs, and the process hosting it will be redeployed or killed halfway. This project makes the loop durable. A run is a database row, every step is written down the moment it finishes, and any worker can pick up a half-finished run and continue from the last checkpoint instead of starting again and paying twice.
Target behaviour
$ curl -s -X POST localhost:8080/api/runs -H 'Content-Type: application/json' \
-d '{"input":"Which acme tickets are overdue?","tokenBudget":50000,"timeoutSeconds":300}'
{"id":"6f1c0a52","status":"QUEUED"}
worker-A run=6f1c0a52 seq=1 MODEL stop=tool_use in=812 out=64 latency=1830ms
worker-A run=6f1c0a52 seq=2 TOOL list_tickets id=toolu_01Ab started
worker-A killed with kill -9 # lease expires 30s later
worker-B run=6f1c0a52 resumed at seq=2; seq=2 TOOL list_tickets id=toolu_01Ab done 41ms
worker-B run=6f1c0a52 seq=3 MODEL stop=end_turn in=1210 out=140 latency=2210ms
worker-B run=6f1c0a52 SUCCEEDED tokens=2226
$ curl -s localhost:8080/api/runs/6f1c0a52/timeline | jq -c '.[] | [.seq,.kind,.status,.latencyMs]'
[1,"MODEL","DONE",1830] [2,"TOOL","DONE",41] [3,"MODEL","DONE",2210]Architecture
POST /api/runs ──▶ ┌──────────── PostgreSQL ────────────┐ ◀── GET /runs/{id}/timeline
POST /cancel ──▶ │ run (status, lease, budgets) │
│ run_step (seq, kind, tokens, body) │
└──▲────────────────────▲────────────┘
claim (SKIP LOCKED) checkpoint (fenced, same txn as token count)
┌──┴────────────────────┴───┐
│ RunWorker (virtual thread) │──▶ LlmClient = Retrying(TimeLimited(Anthropic))
│ guard ▶ replay ▶ next step │──▶ ToolExecutor (idempotent by tool_use id)
└────────────────────────────┘──▶ Micrometer + MDC trace loggingNothing important lives in memory. The conversation is rebuilt each iteration from the ordered run_step rows, so the log is the single source of truth and resuming is the same code path as continuing. The lease on the run row replaces a separate reaper: an expired RUNNING row is simply claimable again.
The core of it
Two tables. run_step is append-only apart from a tool step moving from STARTED to a final state. The unique index on tool steps is the idempotency key from the model's tool_use id.
CREATE TABLE run (
id UUID PRIMARY KEY, status TEXT NOT NULL DEFAULT 'QUEUED', input TEXT NOT NULL, model TEXT NOT NULL,
token_budget INT NOT NULL, tokens_used INT NOT NULL DEFAULT 0, deadline_at TIMESTAMPTZ NOT NULL,
run_at TIMESTAMPTZ NOT NULL DEFAULT now(), locked_by TEXT, locked_until TIMESTAMPTZ,
cancel_requested BOOLEAN NOT NULL DEFAULT false, error TEXT);
CREATE INDEX run_claim_idx ON run (run_at) WHERE status IN ('QUEUED', 'RUNNING');
CREATE TABLE run_step (
run_id UUID NOT NULL REFERENCES run(id), seq INT NOT NULL, kind TEXT NOT NULL,
tool_name TEXT, tool_use_id TEXT, status TEXT NOT NULL, request TEXT, response TEXT,
stop_reason TEXT, input_tokens INT NOT NULL DEFAULT 0, output_tokens INT NOT NULL DEFAULT 0,
latency_ms INT, started_at TIMESTAMPTZ NOT NULL DEFAULT now(), PRIMARY KEY (run_id, seq));
CREATE UNIQUE INDEX step_tool_uk ON run_step (run_id, tool_use_id) WHERE kind = 'TOOL';Claiming is the pattern from project 10 with one extension: a RUNNING row whose lease has lapsed is eligible too. One statement, no explicit transaction; locked_by becomes the fence for every later write.
public Optional<Run> claim(String worker, Duration lease) {
return jdbc.sql("""
WITH next AS (
SELECT id FROM run WHERE run_at <= now() AND (status = 'QUEUED'
OR (status = 'RUNNING' AND locked_until < now()))
ORDER BY run_at FOR UPDATE SKIP LOCKED LIMIT 1)
UPDATE run r SET status = 'RUNNING', locked_by = :w,
locked_until = now() + make_interval(secs => :lease)
FROM next WHERE r.id = next.id
RETURNING r.id, r.input, r.model, r.token_budget, r.deadline_at""")
.param("w", worker).param("lease", (double) lease.toSeconds()).query(RUN_MAPPER).optional();
}The checkpoint is where durability is decided. The step row, the token counter and the lease renewal commit together or not at all, and the fenced UPDATE makes a worker that lost its lease fail loudly instead of writing into someone else's run. Wrap it with TransactionTemplate; the thrown runtime exception rolls everything back.
public void checkpointModel(Run run, String worker, int seq, String request, LlmResponse r, long ms) {
tx.executeWithoutResult(s -> {
int owned = jdbc.sql("""
UPDATE run SET tokens_used = tokens_used + :t, locked_until = now() + make_interval(secs => :lease)
WHERE id = :id AND locked_by = :w AND status = 'RUNNING'""")
.param("t", r.inputTokens() + r.outputTokens()).param("lease", lease.toSeconds())
.param("id", run.id()).param("w", worker).update();
if (owned == 0) throw new LeaseLostException(run.id());
jdbc.sql("""
INSERT INTO run_step(run_id, seq, kind, status, request, response, stop_reason, input_tokens, output_tokens, latency_ms)
VALUES (:id, :seq, 'MODEL', 'DONE', :req, :resp, :stop, :in, :out, :ms)""")
.param("id", run.id()).param("seq", seq).param("req", request).param("stop", r.stopReason())
.param("resp", json(r.content())).param("in", r.inputTokens()).param("out", r.outputTokens())
.param("ms", ms).update();
});
}json(Object) is a helper that wraps Jackson's checked exceptions in UncheckedIOException. Replay rebuilds the message list. A model step adds an assistant message holding its content blocks; the tool_result blocks are gathered into one user message, flushed only once every tool_use in that turn has a result, as the API requires. Whatever is still pending is the next work.
static Conversation replay(Run run, List<Step> steps, ObjectMapper m) throws JsonProcessingException {
List<Map<String, Object>> msgs = new ArrayList<>();
msgs.add(Map.of("role", "user", "content", run.input()));
Map<String, JsonNode> pending = new LinkedHashMap<>(); String stop = null;
List<Map<String, Object>> results = new ArrayList<>();
for (Step s : steps) {
if (s.kind().equals("MODEL")) {
if (!results.isEmpty()) { msgs.add(Map.of("role", "user", "content", List.copyOf(results))); results.clear(); }
JsonNode content = m.readTree(s.response());
msgs.add(Map.of("role", "assistant", "content", content));
stop = s.stopReason(); pending.clear(); content.forEach(b -> { if ("tool_use".equals(b.path("type").asText())) pending.put(b.get("id").asText(), b); });
} else if (!s.status().equals("STARTED")) {
pending.remove(s.toolUseId());
results.add(Map.of("type", "tool_result", "tool_use_id", s.toolUseId(), "content", s.response(),
"is_error", s.status().equals("ERROR")));
}
}
if (pending.isEmpty() && !results.isEmpty()) msgs.add(Map.of("role", "user", "content", results));
return new Conversation(msgs, List.copyOf(pending.values()), stop != null && !stop.equals("tool_use"));
}The driver is a small state machine over that replay. Every iteration starts with guard, which checks cancellation, the deadline, the token budget and a step cap (a model that loops on a tool must not run forever). Budgets are checked before a model step, so overshoot is bounded by one step.
private void drive(Run run) throws Exception {
while (true) {
List<Step> steps = runs.steps(run.id());
String stop = guard(run, runs.state(run.id()), steps.size()); // CANCELLED, TIMED_OUT, FAILED or null
if (stop != null) { runs.finish(run, worker, stop, "stopped by guard"); return; }
Conversation c = Conversation.replay(run, steps, mapper);
if (!c.pending().isEmpty()) toolStep(run, c.pending().get(0));
else if (c.finished()) { runs.finish(run, worker, "SUCCEEDED", null); return; }
else modelStep(run, steps.size() + 1, c);
}
}The model step builds a request from the replayed messages. Mark the stable prefix for prompt caching: systemBlocks is List.of(Map.of("type", "text", "text", systemPrompt, "cache_control", Map.of("type", "ephemeral"))), so later iterations of the same run read it from cache. Cached tokens show up as cache_creation_input_tokens and cache_read_input_tokens in usage, and they are billed at different rates, so extend LlmResponse and the cost function with them.
private void modelStep(Run run, int seq, Conversation c) throws Exception {
LlmRequest req = new LlmRequest(run.model(), 4096, systemBlocks, c.messages(), tools.definitions(), null);
long t0 = System.nanoTime();
LlmResponse r = llm.send(req); // Retrying(TimeLimited(Anthropic))
long ms = (System.nanoTime() - t0) / 1_000_000;
runs.checkpointModel(run, worker, seq, mapper.writeValueAsString(req), r, ms);
metrics.model(run.model(), r, ms);
}A crash between response and checkpoint repeats that one model call: the unavoidable at-least-once window, costing one step of tokens, never a restart. Tools have side effects, so they differ. Write the STARTED row first, run the tool with the tool_use id as its idempotency key, then record the result. If the worker dies in between, the row stays STARTED, replay still reports the tool as pending, and the resumed worker re-executes it with the same key. The tool must honour that key, for example INSERT ... ON CONFLICT (idempotency_key) DO NOTHING or an Idempotency-Key header to the downstream API.
private void toolStep(Run run, JsonNode use) {
String id = use.get("id").asText(), name = use.get("name").asText();
runs.startTool(run.id(), id, name); // INSERT at seq max+1 ... ON CONFLICT (run_id, tool_use_id) DO NOTHING
long t0 = System.nanoTime(); String out, status = "DONE";
try { out = tools.execute(new ToolCall(id, name, use.get("input"))); }
catch (Exception e) { out = "error: " + e.getMessage(); status = "ERROR"; } // model sees is_error
runs.finishTool(run, worker, id, status, out, (System.nanoTime() - t0) / 1_000_000);
}finishTool is fenced exactly like checkpointModel. The executor enforces its own per-call timeout. Tools that throw become is_error: true results, because a failing tool is information for the model. Errors and lease loss are not caught; they abort the run attempt so another worker can take over.
Retries live in a decorator. The API signals back-pressure with HTTP 429 (with a retry-after header) and 529 (overloaded); 5xx and I/O failures are also transient, while 400, 401, 403 and 404 never are. Use full jitter, and take the larger of the jitter and the server's retry-after.
public LlmResponse send(LlmRequest req) {
for (int attempt = 0; ; attempt++) {
try {
return delegate.send(req);
} catch (LlmException e) {
if (!e.retryable() || attempt + 1 >= maxAttempts) throw e;
long wait = ThreadLocalRandom.current().nextLong(Math.min(capMs, baseMs << Math.min(attempt, 10)) + 1);
if (e.retryAfter() != null) wait = Math.max(wait, e.retryAfter().toMillis());
sleeper.sleep(Duration.ofMillis(wait)); // Sleeper wraps Thread.sleep, no checked exception
}
}
}retryable() is true for 429, 500, 502, 503, 504, 529 and -1 (I/O or timeout). The step timeout wraps a single attempt, inside the retry loop: TimeLimitedLlmClient submits delegate.send to a virtual-thread executor (Executors.newVirtualThreadPerTaskExecutor()), waits with Future.get(timeout, MILLISECONDS), and on TimeoutException calls cancel(true) and throws LlmException(-1, null, "step timeout"). Set the HttpRequest timeout slightly above it so interruption does the work.
Cancellation is cooperative. The endpoint runs UPDATE run SET cancel_requested = true, status = CASE WHEN status = 'QUEUED' THEN 'CANCELLED' ELSE status END WHERE id = :id AND status IN ('QUEUED','RUNNING'): a queued run ends immediately, a running one notices at the next guard.
Observability has three parts. Metrics: counters agent.tokens (tags model, direction) and agent.cost.usd, a timer agent.step.latency (tags kind, outcome), and a gauge of runs by status, all via Counter.builder(...).register(registry) and Timer.builder(...).register(registry).record(duration). Price per million tokens comes from properties, never constants. Tracing: put run_id and step_seq in the SLF4J MDC (MDC.putCloseable in try-with-resources) and switch the console to JSON with logging.structured.format.console=ecs on Spring Boot 3.4+, or a Logstash encoder before that. Timeline: GET /api/runs/{id}/timeline returns the run_step rows without request bodies.
Testing
Reuse the scripted LlmClient from project 16, with fluent helpers thenToolUse(id, name, inputJson, in, out), thenText(text, in, out) and thenFail(LlmException). Test the retry decorator with a fake Sleeper recording durations: two 429s with retry-after: 7 then a success gives two sleeps of at least 7000 ms and three calls; a 400 gives exactly one call.
The test that proves the design simulates a crash with an Error subclass, which the tool step deliberately does not catch, so the worker thread dies holding its lease.
@SpringBootTest @Testcontainers
class ResumeIT {
@Container @ServiceConnection static PostgreSQLContainer<?> pg = new PostgreSQLContainer<>("postgres:16");
@Autowired RunStore runs; @Autowired JdbcClient jdbc; ScriptedLlm llm = new ScriptedLlm(); List<String> executed = new ArrayList<>();
@Test
void resumedRunDoesNotRepeatModelCalls() {
llm.thenToolUse("toolu_1", "lookup", "{\"q\":\"x\"}", 100, 20).thenText("done", 150, 5);
UUID id = runs.create("find x", "claude-sonnet-5-5", 10_000, Duration.ofMinutes(5));
ToolExecutor crashOnce = call -> { // SimulatedCrash extends Error
executed.add(call.id()); if (executed.size() == 1) throw new SimulatedCrash(); return "found";
};
assertThrows(SimulatedCrash.class, () -> worker("A", crashOnce).runOnce());
jdbc.sql("UPDATE run SET locked_until = now() - interval '1 second' WHERE id = :id").param("id", id).update();
worker("B", crashOnce).runOnce();
assertThat(runs.get(id).status()).isEqualTo("SUCCEEDED");
assertThat(llm.seen).hasSize(2); // first model call was not repeated
assertThat(executed).containsExactly("toolu_1", "toolu_1");
}
}Bugs you will hit
- Lease shorter than one step. A 90 second model call under a 30 second lease makes a second worker take over a live run and both append steps. Size the lease above the worst single step including its retries, renew it at every checkpoint, and rely on the fence to make the loser stop.
- Rebuilding messages wrongly. Sending
tool_resultblocks split across messages, or atool_resultwithout its precedingtool_use, gets a 400. Replay must emit the assistant content blocks exactly as returned and group all results for one turn into a single user message. - Non-idempotent tools. Re-execution after a crash charges the card twice. The
tool_useid is a stable key only if the tool or its downstream honours it; if it cannot, mark the stepSTARTEDand fail the run for a human rather than guessing. - Retrying everything. Retrying a 400 or 401 burns time and quota, and retrying without honouring
retry-afterwhile rate limited extends the outage. Classify errors, cap attempts, and add jitter so recovered workers do not stampede.
Interview angle
- How is this different from retrying the whole run? Whole-run retry repeats completed model calls (money) and completed tool effects (correctness). Checkpointing makes the unit of redo one step, and idempotency keys make that redo safe.
- Why is the lease on the row better than an in-memory lock? It survives the worker dying, needs no coordinator, and the fenced updates give you a correctness guarantee at the data layer: a stale worker cannot corrupt state even if it wakes up late.
- Where can duplicates still happen? Between a model response and its checkpoint (one repeated call) and between a tool's side effect and its recording (one repeated execution). Delivery is at-least-once; correctness comes from idempotent tools, not from preventing the repeat.
- How do you stop a runaway agent? Defence in depth: a token budget, a wall-clock deadline, a step cap, per-step timeouts and a cooperative cancel flag, all checked in the loop, plus alerts on cost per run.
Milestones
Done when
- A worker killed mid-run leaves a run that another worker finishes, without repeating any completed model call.
- Token budget, deadline, step cap and cancellation each stop a run with a recorded reason; 429 and 529 are retried with jitter honouring
retry-after, client errors are not. - A stale worker cannot write after losing its lease, proven by a test; tokens, latency and cost are in Prometheus and every log line carries
run_id.
Stretch goals
- Execute independent tool calls from one turn in parallel on virtual threads, writing all results before the next model step.
- Add a heartbeat thread that renews the lease during long steps instead of only at checkpoints.
- Stream the model response (server-sent events) so the timeline updates live.
- Add human approval steps: a tool marked sensitive pauses the run in
WAITINGuntil an API call approves or rejects it.