Follow-up to PR #64171. The _TOKEN_DELTA_* classification must exactly
cover update_token_counts' keyword surface: an unclassified kwarg is
silently kept only from the first delta of a merged run. Introspect the
live signature and fail with a pointed message when a future kwarg is
added without classification (or a classified field is removed).
Follow-up to PR #64171. The writer loop and flush_token_counts'
caller-drain both set _token_writer_busy BEFORE popping the queue —
that ordering is what makes flush's lock-free fast path (reads queue
then busy, no cond held) sound. _stop_token_writer's leftover drain did
it backwards (clear queue, then set busy), leaving a few-bytecode window
where a concurrent flush could observe 'empty and idle' and return True
with the popped batch still unapplied. Shutdown-only staleness, no data
loss — but the protocol now matches at all three drain sites.
Test: concurrent flush during a stop-drain mid-apply must time out
(False), never report drained.
Follow-up to PR #64171. Two writer-death hardening gaps:
- queue_token_counts only spawned the writer when the thread object was
None, so a writer that died from an unexpected exception could never be
replaced: deltas piled up on the uncapped deque until a reader's flush
drained them synchronously. Respawn on 'not thread.is_alive()' instead.
(atexit re-registration on respawn is safe: unregister removes all equal
bound-method registrations and the drain hook is idempotent.)
- _coalesce_token_deltas ran outside the per-delta try/except in
_apply_token_batch, so a merge bug (e.g. an unclassified future kwarg
summing None + 0) would escape and kill the writer thread. Wrap it and
fall back to applying the raw batch — coalescing is an optimization,
never load-bearing.
Tests: coalesce-failure fallback + dead-writer respawn.
Every API call in the tool loop persisted its token/cost delta by
calling SessionDB.update_token_counts() synchronously on the turn
thread — a BEGIN IMMEDIATE sessions UPDATE plus a session_model_usage
upsert, measured in production at p50 3.3ms / p95 70.4ms per call and
up to 299ms against a cold multi-GB state.db. The tool loop stalls for
that long between calls, N times per multi-tool turn.
SessionDB gains queue_token_counts(): same signature and semantics as
update_token_counts(), but the critical path is a deque append plus a
condvar notify. A lazily started daemon thread applies deltas in
enqueue order through the existing update_token_counts ->
_execute_write path, so the established self._lock / BEGIN IMMEDIATE /
jitter-retry discipline is unchanged. When a backlog forms, adjacent
same-route incremental deltas coalesce into one UPDATE: token and
api-call fields sum, cost fields sum None-preservingly (an all-None
run stays None so COALESCE keeps the stored value), and absolute=True
deltas never merge and act as ordering barriers. Route equality is
required for a merge because those fields feed COALESCE backfill, the
last-non-None-wins status fields, and the per-model usage attribution
key — a merged apply is row-equivalent to sequential applies.
Correctness and durability:
- flush_token_counts() gives read-your-writes to token/cost readers
(get_session, list_sessions_rich, _get_session_rich_row,
list_gateway_sessions, InsightsEngine.generate) — a plain attribute
check when nothing is queued. The writer sets its busy flag before
popping the queue so the lock-free fast path can never miss an
in-flight batch.
- update_session_model / update_session_billing_route /
update_session_meta write the sessions row synchronously, bypassing
the queue, so they flush it first: a still-queued first-of-session
delta carries the pre-switch route, and applying it after the switch
UPDATE would trip the first_accounted_route branch (api_call_count
== 0 plus a route mismatch) and resurrect the old model/provider.
- AIAgent._persist_session flushes at turn finalize and every
error-exit persist point; close() stops and drains the writer before
the WAL checkpoint; an atexit hook (registered on first enqueue,
unregistered on close so closed instances are not 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.
- A flush trusts a live stop-flagged writer (its loop drains before
exiting) and 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 an in-flight batch.
- After close() has stopped the writer, 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.
- Writer apply failures are logged and never raise into a turn; the
writer thread survives and keeps applying.
Call sites switched to the queue: the per-call site in
agent/conversation_loop.py and both codex app-server sites in
agent/codex_runtime.py. In-memory per-turn counters
(agent.session_estimated_cost_usd etc.) stay synchronous, so live turn
displays never see the queue.
Tests: tests/agent/test_async_token_accounting.py (19 tests: enqueue
ordering, absolute-as-barrier, backlog coalescing with exact sums,
coalesced-vs-sequential row equivalence, merge unit rules, None-cost
preservation, read-your-writes, flush vs stop-flagged/concurrent
drains, inline apply after writer stop, close/atexit durability,
_persist_session drain, writer failure isolation);
tests/run_agent/test_token_persistence_non_cli.py updated to the
queue_token_counts contract.
