perf: take per-API-call token accounting off the turn thread

[Soju06]

Problem

Every API call inside the tool loop persists its token/cost delta by calling SessionDB.update_token_counts() synchronously on the turn thread — the hot site in agent/conversation_loop.py plus both codex app-server sites in agent/codex_runtime.py. Each call is a BEGIN IMMEDIATE sessions UPDATE plus a session_model_usage upsert.

In a production deployment we measured this write at 3.3 ms p50 / 70.4 ms p95 per API call, with spikes up to 299 ms against a cold multi-GB state.db. The tool loop stalls for that long between calls; a multi-tool turn with N API calls pays it N times. None of it needs to block the turn: the in-memory per-turn counters (agent.session_estimated_cost_usd etc.) that live displays read are updated separately and stay synchronous.

Change

SessionDB gains a single-writer background queue for token deltas:

• queue_token_counts(session_id, **kwargs) — same signature and semantics as update_token_counts, but the critical path is a deque append + condvar notify. A dedicated daemon thread (session-db-token-writer, started lazily on first enqueue) applies deltas in enqueue order through the existing update_token_counts → _execute_write path, so SQLite access keeps the established self._lock / BEGIN IMMEDIATE / jitter-retry discipline unchanged.
• Coalescing — when a backlog forms, adjacent deltas for the same (session, model, cost_status, cost_source, pricing_version, billing_provider/base_url/mode) route merge into one UPDATE: token/api-call fields sum, cost fields sum None-preservingly (an all-None run stays None so COALESCE keeps the stored value). absolute=True deltas (cumulative overwrites) never merge and act as ordering barriers.
• flush_token_counts(timeout=5) — blocks until the queue is drained; a plain attribute check when nothing is queued. Readers that surface token/cost totals call it first: get_session, list_sessions_rich, _get_session_rich_row, list_gateway_sessions, and InsightsEngine.generate.
• Call sites — the conversation_loop.py per-call site and both codex_runtime.py sites now enqueue instead of writing inline. Their existing row-existence retry and try/except guards are unchanged.

Correctness notes

• Ordering — one writer thread applies deltas strictly in enqueue order; only adjacent same-route incremental deltas merge, so ordering across sessions and across a mid-session /model switch is preserved exactly. Route equality is required for a merge precisely because those fields feed the COALESCE backfill, the last-non-None-wins status fields, and the per-model usage attribution key — a merged apply is equivalent to sequential applies (covered by a row-equivalence test comparing full sessions + session_model_usage rows).
• Flush barriers before route switches — update_session_model, update_session_billing_route, and update_session_meta write the sessions row synchronously, bypassing the queue. A first-of-session delta still queued at switch time carries the pre-switch route; applying it after the switch UPDATE would trip the first-accounted-route branch in update_token_counts (row sees api_call_count == 0 plus a route mismatch) and resurrect the old model/provider. All three flush the queue first, restoring the pre-queue happens-before.
• Lock-free fast path vs in-flight batches — the writer sets its busy flag before popping the queue, so flush_token_counts's unlocked queue-then-busy check can never observe "empty and idle" while a popped batch is still unapplied.
• Flush vs a stop-flagged writer — a live writer is authoritative even when stop-flagged (close in progress): its loop drains the queue before exiting. Flush only drains on the caller's thread when the writer is dead or never started, claiming the same busy flag so concurrent flushes wait instead of racing.
• Durability — AIAgent._persist_session (turn finalize and every error-exit persist point) flushes; close() stops and drains the writer before the WAL checkpoint; an atexit hook (registered on first enqueue, unregistered on close() so closed instances aren't pinned until interpreter exit) drains at shutdown. Worst-case crash loss is the in-flight call's delta — the same window as the old inline write.
• Enqueue after close — once the writer is permanently stopped, queue_token_counts applies the delta inline instead of parking it on a queue nothing will drain; a closed-connection failure then raises at the call site, which already guards for it, exactly like the old synchronous path.
• Failure isolation — writer apply errors are logged (logger.warning) and never raise into a turn; the writer thread survives and keeps applying subsequent deltas.

Tests

• New: tests/agent/test_async_token_accounting.py (19 tests) — strict enqueue-order across sessions, absolute-as-barrier semantics, backlog coalescing with exact sums including session_model_usage, coalesced-vs-sequential row equivalence, unit merge rules, None-cost preservation, get_session read-your-writes under a slow writer, empty-flush fast path, flush waiting out a stop-flagged live writer and a concurrent caller-drain, inline apply after writer stop, enqueue-after-close raising at the call site, close() drain surviving reopen, atexit idempotence and unregistration (weakref + gc), _persist_session drain via a real AIAgent, and writer failure logging + survival.
• Updated: tests/run_agent/test_token_persistence_non_cli.py to the queue_token_counts contract.
• Suites for every touched file pass: tests/test_hermes_state.py, tests/hermes_state/, tests/test_sql_injection.py, tests/agent/test_insights.py, tests/agent/test_codex_app_server_persist.py, and the full tests/run_agent tree (1974 passed, 3 skipped). The full tests/agent tree was run before and after the change on the same machine: the failure set is byte-for-byte identical (144 environment-dependent tests — LSP tooling, model-metadata network cache, codex transport, and similar — that fail the same way on a clean checkout of main; the after run passes the same 5468 tests plus the 19 new ones).

How to test:

pytest tests/agent/test_async_token_accounting.py tests/test_hermes_state.py tests/run_agent -q

Tested on Linux (Ubuntu, Python 3.12); the change touches no platform-specific I/O — it is pure threading + the existing SQLite write path.

Measured impact

From a production deployment of this change:

• Enqueue critical path: p50 0.0015 ms / p95 0.0036 ms, vs 2.6 ms / 7.7 ms for the synchronous UPDATE on a warm, small DB — and vs 3.3 ms p50 / 70.4 ms p95 (up to 299 ms cold) observed in production before the change.
• Removes that entire cost from the turn thread after every API call; multi-tool turns with N calls save N× it.
• Total DB work is unchanged or lower: the same deltas are applied off-thread, and backlogs coalesce into fewer transactions.

🤖 Generated with Claude Code

[Soju06]

The failing Python tests slices here are test_model_validation.py::TestProbeApiModelsUserAgent — broken on main itself since b8eb89f (the tests mock urllib.request.urlopen, but the probe goes through open_credentialed_url's OpenerDirector). Test-only fix: #64200. Unrelated to this PR's diff.

[teknium1]

Thanks for addressing a real turn-thread latency source. Current main still invokes SessionDB.update_token_counts() inline in agent/conversation_loop.py:2268 and in both Codex app-server branches at agent/codex_runtime.py:129 and agent/codex_runtime.py:209.

Static review found no specific blocking defect in the queue, coalescing, flush, and route-switch-barrier design proposed in bcce3ea076ea1237898710f7203c20ef1ac9351a.

Automated hermes-sweeper review.

[kshitijk4poor]

Merged via #73359 (rebase-merge) — your commit perf(state): apply per-call token accounting on a background single-writer queue landed on main with your authorship preserved, plus three small hardening follow-ups from our review (dead-writer respawn + coalesce-failure fallback, busy-before-clear ordering in the stop-drain, and a signature guard test for the coalescing field lists).

We validated it end-to-end with real API calls before and after: identical persisted accounting, turn-thread cost down from ~2.7-3.3 ms to ~0.25 ms per call, and totals byte-exact under a 200-call backlog and a multi-thread hammer. Excellent, carefully-reasoned work — thank you!