fix: backfill model on gateway sessions after agent runs

gateway/run.py

@@ -1437,7 +1437,10 @@ async def _handle_message(self, event: MessageEvent) -> Optional[str]:
             # Update session with actual prompt token count from the agent
             self.session_store.update_session(
                 session_entry.session_key,
+                input_tokens=agent_result.get("input_tokens", 0),
+                output_tokens=agent_result.get("output_tokens", 0),
                 last_prompt_tokens=agent_result.get("last_prompt_tokens", 0),
+                model=agent_result.get("model"),
             )
 
             return response
@@ -3234,6 +3237,7 @@ def run_sync():
                     "tools": tools_holder[0] or [],
                     "history_offset": len(agent_history),
                     "last_prompt_tokens": _last_prompt_toks,
+                    "model": agent_holder[0].model if agent_holder[0] else None,
                 }
 
             # Scan tool results for MEDIA:<path> tags that need to be delivered
@@ -3278,6 +3282,7 @@ def run_sync():
                 "tools": tools_holder[0] or [],
                 "history_offset": len(agent_history),
                 "last_prompt_tokens": _last_prompt_toks,
+                "model": agent_holder[0].model if agent_holder[0] else None,
             }
 
         # Start progress message sender if enabled

gateway/session.py

@@ -554,15 +554,16 @@ def get_or_create_session(
         return entry
 
     def update_session(
-        self, 
+        self,
         session_key: str,
         input_tokens: int = 0,
         output_tokens: int = 0,
         last_prompt_tokens: int = None,
+        model: str = None,
     ) -> None:
         """Update a session's metadata after an interaction."""
         self._ensure_loaded()
-        
+
         if session_key in self._entries:
             entry = self._entries[session_key]
             entry.updated_at = datetime.now()
@@ -572,11 +573,12 @@ def update_session(
                 entry.last_prompt_tokens = last_prompt_tokens
             entry.total_tokens = entry.input_tokens + entry.output_tokens
             self._save()
-            
+
             if self._db:
                 try:
                     self._db.update_token_counts(
-                        entry.session_id, input_tokens, output_tokens
+                        entry.session_id, input_tokens, output_tokens,
+                        model=model,
                     )
                 except Exception as e:
                     logger.debug("Session DB operation failed: %s", e)

hermes_state.py

@@ -227,16 +227,27 @@ def update_system_prompt(self, session_id: str, system_prompt: str) -> None:
         self._conn.commit()
 
     def update_token_counts(
-        self, session_id: str, input_tokens: int = 0, output_tokens: int = 0
+        self, session_id: str, input_tokens: int = 0, output_tokens: int = 0,
+        model: str = None,
     ) -> None:
-        """Increment token counters on a session."""
-        self._conn.execute(
-            """UPDATE sessions SET
-               input_tokens = input_tokens + ?,
-               output_tokens = output_tokens + ?
-               WHERE id = ?""",
-            (input_tokens, output_tokens, session_id),
-        )
+        """Increment token counters on a session. Optionally backfill model."""
+        if model:
+            self._conn.execute(
+                """UPDATE sessions SET
+                   input_tokens = input_tokens + ?,
+                   output_tokens = output_tokens + ?,
+                   model = COALESCE(model, ?)
+                   WHERE id = ?""",
+                (input_tokens, output_tokens, model, session_id),
+            )
+        else:
+            self._conn.execute(
+                """UPDATE sessions SET
+                   input_tokens = input_tokens + ?,
+                   output_tokens = output_tokens + ?
+                   WHERE id = ?""",
+                (input_tokens, output_tokens, session_id),
+            )
         self._conn.commit()
 
     def get_session(self, session_id: str) -> Optional[Dict[str, Any]]: