fix: integration hardening for gateway token tracking

Follow-up to 58dbd81 — ensures smooth transition for existing users:

- Backward compat: old session files without last_prompt_tokens
  default to 0 via data.get('last_prompt_tokens', 0)
- /compress, /undo, /retry: reset last_prompt_tokens to 0 after
  rewriting transcripts (stale token counts would under-report)
- Auto-compression hygiene: reset last_prompt_tokens after rewriting
- update_session: use None sentinel (not 0) as default so callers
  can explicitly reset to 0 while normal calls don't clobber
- 6 new tests covering: default value, serialization roundtrip,
  old-format migration, set/reset/no-change semantics
- /reset: new SessionEntry naturally gets last_prompt_tokens=0

2942 tests pass.
This commit is contained in:
teknium1 2026-03-10 23:40:24 -07:00
parent 5eb62ef423
commit 909e048ad4
3 changed files with 128 additions and 2 deletions

View file

@ -556,7 +556,7 @@ class SessionStore:
session_key: str,
input_tokens: int = 0,
output_tokens: int = 0,
last_prompt_tokens: int = 0,
last_prompt_tokens: int = None,
) -> None:
"""Update a session's metadata after an interaction."""
self._ensure_loaded()
@ -566,7 +566,7 @@ class SessionStore:
entry.updated_at = datetime.now()
entry.input_tokens += input_tokens
entry.output_tokens += output_tokens
if last_prompt_tokens > 0:
if last_prompt_tokens is not None:
entry.last_prompt_tokens = last_prompt_tokens
entry.total_tokens = entry.input_tokens + entry.output_tokens
self._save()