From eb4f0348e1f6505f58cb4e30e62273332df4650f Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Mon, 16 Mar 2026 00:23:13 -0700
Subject: [PATCH] fix: persist CLI token counts to session DB for /insights

Token usage was tracked in-memory during CLI sessions (session_prompt_tokens,
session_completion_tokens) but never written to the SQLite session DB. The
gateway persisted tokens via session_store.update_session(), but CLI sessions
always showed 0 tokens in /insights.

Now run_agent.py persists token deltas to the DB after each API call for CLI
sessions. Gateway sessions continue to use their existing persist path to
avoid double-counting.
---
 run_agent.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/run_agent.py b/run_agent.py
index 7ee22876..431c0e6a 100644
--- a/run_agent.py
+++ b/run_agent.py
@@ -5069,6 +5069,22 @@ class AIAgent:
                         self.session_completion_tokens += completion_tokens
                         self.session_total_tokens += total_tokens
                         self.session_api_calls += 1
+
+                        # Persist token counts to session DB for /insights.
+                        # Gateway sessions persist via session_store.update_session()
+                        # after run_conversation returns, so only persist here for
+                        # CLI (and other non-gateway) platforms to avoid double-counting.
+                        if (self._session_db and self.session_id
+                                and getattr(self, 'platform', None) == 'cli'):
+                            try:
+                                self._session_db.update_token_counts(
+                                    self.session_id,
+                                    input_tokens=prompt_tokens,
+                                    output_tokens=completion_tokens,
+                                    model=self.model,
+                                )
+                            except Exception:
+                                pass  # never block the agent loop
                         
                         if self.verbose_logging:
                             logging.debug(f"Token usage: prompt={usage_dict['prompt_tokens']:,}, completion={usage_dict['completion_tokens']:,}, total={usage_dict['total_tokens']:,}")