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
+2 -1
View File
@@ -77,7 +77,8 @@ def _queue_token_counts(agent, fail_msg: str, *fail_extra: Any, counts: Callable
try:
if not agent._session_db_created:
agent._ensure_db_session()
agent._session_db.queue_token_counts(agent.session_id, **counts())
from agent.turn_usage import _agent_session_source
agent._session_db.queue_token_counts(agent.session_id, source=_agent_session_source(agent), **counts())
except Exception as exc:
logger.debug(fail_msg, agent.session_id, *fail_extra, exc)
+8
View File
@@ -22,6 +22,13 @@ from agent.usage_pricing import estimate_usage_cost, normalize_usage
logger = logging.getLogger("agent.conversation_loop")
def _agent_session_source(agent: Any) -> str:
"""The surface the agent's own row create would stamp (``_ensure_db_session``), so an
accounting guard that wins the row-creation race never mints an anonymous session."""
from run_agent import _session_source_for_agent # late: run_agent imports this module
return _session_source_for_agent(getattr(agent, "platform", None))
@dataclass
class ResponseUsageOutcome:
"""``compression_attempts`` is the (possibly rearmed-to-zero) budget counter;
@@ -248,6 +255,7 @@ def record_response_usage(
agent._ensure_db_session()
agent._session_db.queue_token_counts(
agent.session_id,
source=_agent_session_source(agent),
input_tokens=canonical_usage.input_tokens,
output_tokens=canonical_usage.output_tokens,
cache_read_tokens=canonical_usage.cache_read_tokens,