fix(sessions): token-accounting guard stamps the agent's real source; trim salvage

When every row create of a turn loses to the SQLite lock, the queued token delta's
"ensure the row exists" guard becomes the session's first writer and minted the row as
source='unknown'. That placeholder was permanent on the real path even with the upsert
repair from #112045: the turn lease (turn_facade_lease.admit_durable_turn) treats an existing
row as proof the create already happened and sets _session_db_created, so the creator never
returns to repair it. Live probe: a platform="desktop" AIAgent whose create_session raised
"database is locked" for the whole first turn ended with a source='unknown' row on base AND
on the contributor head; with this change the row is minted 'desktop' by the guard itself.

Producer fix: update_token_counts gains an optional source= that the two agent call sites
(agent/turn_usage.py, agent/codex_runtime.py) fill from _session_source_for_agent(platform),
the same value _ensure_db_session would stamp. record_auxiliary_usage has no surface and
keeps the placeholder, which the creator's upsert now repairs.

Salvage trims: the contributor's SimpleNamespace dispatch test is replaced by a real-AIAgent
invariant test under tests/agent/ (the dispatch hunk in _run_prompt_submit is kept; the
INSERT-OR-IGNORE is idempotent under prompt.submit's own persist); narration comments cut
to the WHY; docs list 'unknown' among the startup-sweep sources.

Refs #111999
This commit is contained in:
teknium1
2026-09-15 12:17:38 -07:00
committed by Teknium
parent a7dde8a57d
commit cfd752e6f7
8 changed files with 95 additions and 113 deletions
+10 -9
View File
@@ -278,18 +278,19 @@ class SessionUsageMixin:
actual_cost_usd: Optional[float]=None, cost_status: Optional[str]=None, cost_source: Optional[str]=None,
pricing_version: Optional[str]=None, billing_provider: Optional[str]=None, billing_base_url: Optional[str]=None,
billing_mode: Optional[str]=None, api_call_count: int=0, absolute: bool=False,
source: Optional[str]=None,
) -> None:
"""Update token counters and backfill model if unset. *absolute*=False increments
(per-API-call deltas, CLI path); *absolute*=True sets directly (gateway path,
where the cached agent holds cumulative totals)."""
where the cached agent holds cumulative totals). ``source`` is the session's real surface
for the row-existence guard; callers that don't know it leave the placeholder."""
usage = {k: v for k, v in locals().items() if k in _MODEL_USAGE_FIELDS}
# Ensure the row exists: under concurrent load create_session() may have failed on
# locking, and the UPDATE would silently affect 0 rows. The minted row carries the
# placeholder ``unknown`` source; a later writer's real surface replaces it in
# _insert_session_row's upsert, so the placeholder cannot outlive the session's creator
# (#111999). Until then the token guard is the only thing holding the row — never a
# session the user is shown as theirs.
self._insert_session_row(session_id, "unknown", model=model)
# locking, and the UPDATE would silently affect 0 rows. When this guard is the first
# writer it must carry the agent's real source: the turn lease treats an existing row as
# proof the create already happened, so the creator never returns to repair an anonymous
# ``unknown`` placeholder and the session stays a phantom for life (#111999).
self._insert_session_row(session_id, source or "unknown", model=model)
sql = _TOKEN_UPDATE_ABSOLUTE_SQL if absolute else _TOKEN_UPDATE_DELTA_SQL
has_usage = bool(input_tokens or output_tokens or cache_read_tokens or cache_write_tokens or reasoning_tokens
or api_call_count or estimated_cost_usd)
@@ -384,8 +385,8 @@ class SessionUsageMixin:
if not session_id or not task:
return
usage["api_call_count"] = 1 if api_call_count is None else int(api_call_count)
# FK to sessions.id: same INSERT OR IGNORE guard as update_token_counts (its placeholder
# source is repairable by the session's real creator — see _insert_session_row).
# FK to sessions.id: same guard as update_token_counts; the aux path carries no surface, so
# the placeholder stays repairable by the creator's upsert (_insert_session_row).
self._insert_session_row(session_id, "unknown")
self._execute_write(lambda conn: self._record_model_usage(conn, session_id, task=task, **usage))