fix: persist token counts from gateway to SessionEntry and SQLite (#316)

Closes #316

All sessions had input_tokens=0, output_tokens=0 because nothing
persisted them. The agent accumulates session_prompt_tokens and
session_completion_tokens internally, and the gateway extracts them
into the result dict — but update_session() only accepted
last_prompt_tokens, silently dropping the rest.

Changes:

- gateway/session.py: Extend update_session() to accept and persist
  input_tokens and output_tokens to SessionEntry

- gateway/run.py: After every conversation turn, extract token counts
  from the agent instance and pass them to update_session() AND to
  _session_db.set_token_counts() for SQLite persistence

- tests/test_token_tracking_persistence.py: Regression tests for
  SessionEntry token field roundtrip, update_session persistence,
  and disk save/load survival

Token flow: API response → agent.session_input_tokens →
gateway extracts → update_session(input_tokens=X) →
SessionEntry.input_tokens → sessions.json + state.db

gateway/run.py

@@ -3075,14 +3075,30 @@ class GatewayRunner:
                             skip_db=agent_persisted,
                         )
             
-            # Token counts and model are now persisted by the agent directly.
-            # Keep only last_prompt_tokens here for context-window tracking and
-            # compression decisions.
+            # Token counts — persist to SessionEntry and SQLite (Issue #316).
+            # The agent instance accumulates session_prompt_tokens and
+            # session_completion_tokens across API calls within a turn.
+            _input_toks = getattr(agent, "session_prompt_tokens", 0) if agent else 0
+            _output_toks = getattr(agent, "session_completion_tokens", 0) if agent else 0
+
             self.session_store.update_session(
                 session_entry.session_key,
                 last_prompt_tokens=agent_result.get("last_prompt_tokens", 0),
+                input_tokens=_input_toks,
+                output_tokens=_output_toks,
             )
 
+            # Persist to SQLite if session DB is available
+            if self._session_db and session_entry.session_id:
+                try:
+                    self._session_db.set_token_counts(
+                        session_entry.session_id,
+                        input_tokens=_input_toks,
+                        output_tokens=_output_toks,
+                    )
+                except Exception as e:
+                    logger.debug("Failed to persist token counts to SQLite: %s", e)
+
             # Auto voice reply: send TTS audio before the text response
             _already_sent = bool(agent_result.get("already_sent"))
             if self._should_send_voice_reply(event, response, agent_messages, already_sent=_already_sent):

gateway/session.py

@@ -810,6 +810,8 @@ class SessionStore:
         self,
         session_key: str,
         last_prompt_tokens: int = None,
+        input_tokens: int = None,
+        output_tokens: int = None,
     ) -> None:
         """Update lightweight session metadata after an interaction."""
         with self._lock:
@@ -820,6 +822,10 @@ class SessionStore:
                 entry.updated_at = _now()
                 if last_prompt_tokens is not None:
                     entry.last_prompt_tokens = last_prompt_tokens
+                if input_tokens is not None:
+                    entry.input_tokens = input_tokens
+                if output_tokens is not None:
+                    entry.output_tokens = output_tokens
                 self._save()
 
     def reset_session(self, session_key: str) -> Optional[SessionEntry]:

tests/test_token_tracking_persistence.py

@@ -0,0 +1,100 @@
+"""Test token tracking persistence in SessionEntry and SessionStore.
+
+Refs: #316 — Token tracking - all token counts are zero
+"""
+
+from __future__ import annotations
+
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from gateway.session import SessionEntry, SessionStore
+
+
+class TestSessionEntryTokenFields:
+    """Verify SessionEntry has and persists token fields."""
+
+    def test_has_token_fields(self):
+        entry = SessionEntry(session_key="test", session_id="s1")
+        assert hasattr(entry, "input_tokens")
+        assert hasattr(entry, "output_tokens")
+        assert entry.input_tokens == 0
+        assert entry.output_tokens == 0
+
+    def test_token_fields_roundtrip(self):
+        """Tokens survive serialization to/from dict."""
+        entry = SessionEntry(
+            session_key="test",
+            session_id="s1",
+            input_tokens=1234,
+            output_tokens=567,
+        )
+        data = entry.to_dict()
+        restored = SessionEntry.from_dict(data)
+        assert restored.input_tokens == 1234
+        assert restored.output_tokens == 567
+
+    def test_token_fields_in_json(self):
+        """Tokens appear in serialized JSON."""
+        entry = SessionEntry(
+            session_key="test",
+            session_id="s1",
+            input_tokens=9999,
+            output_tokens=8888,
+        )
+        data = json.loads(json.dumps(entry.to_dict()))
+        assert data["input_tokens"] == 9999
+        assert data["output_tokens"] == 8888
+
+
+class TestUpdateSessionTokenPersistence:
+    """Verify update_session() persists token counts."""
+
+    def test_update_session_sets_tokens(self, tmp_path: Path):
+        """update_session() with input/output tokens persists to entry."""
+        store = SessionStore(sessions_dir=tmp_path)
+        entry = store.get_or_create_session_key("test-session")
+
+        store.update_session(
+            entry.session_key,
+            last_prompt_tokens=100,
+            input_tokens=5000,
+            output_tokens=2000,
+        )
+
+        # Re-read the entry
+        reloaded = store.get_session(entry.session_key)
+        assert reloaded is not None
+        assert reloaded.input_tokens == 5000
+        assert reloaded.output_tokens == 2000
+        assert reloaded.last_prompt_tokens == 100
+
+    def test_update_session_partial(self, tmp_path: Path):
+        """update_session() with only input_tokens preserves output_tokens."""
+        store = SessionStore(sessions_dir=tmp_path)
+        entry = store.get_or_create_session_key("test-session")
+
+        # Set both first
+        store.update_session(entry.session_key, input_tokens=100, output_tokens=200)
+        # Update only input
+        store.update_session(entry.session_key, input_tokens=300)
+
+        reloaded = store.get_session(entry.session_key)
+        assert reloaded.input_tokens == 300
+        assert reloaded.output_tokens == 200  # Preserved
+
+    def test_tokens_persist_to_disk(self, tmp_path: Path):
+        """Tokens survive save/load cycle."""
+        store = SessionStore(sessions_dir=tmp_path)
+        entry = store.get_or_create_session_key("test-session")
+        store.update_session(entry.session_key, input_tokens=7777, output_tokens=8888)
+
+        # Create new store from same dir — simulates restart
+        store2 = SessionStore(sessions_dir=tmp_path)
+        reloaded = store2.get_session(entry.session_key)
+        assert reloaded is not None
+        assert reloaded.input_tokens == 7777
+        assert reloaded.output_tokens == 8888