diff --git a/gateway/run.py b/gateway/run.py index 3c2abd834..ca96e5aec 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -1434,10 +1434,11 @@ async def _handle_message(self, event: MessageEvent) -> Optional[str]: skip_db=agent_persisted, ) - # Update session with actual prompt token count from the agent + # Update session with token counts and model from this turn self.session_store.update_session( session_entry.session_key, last_prompt_tokens=agent_result.get("last_prompt_tokens", 0), + model=agent_result.get("model"), ) return response @@ -3234,6 +3235,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: tags that need to be delivered @@ -3278,6 +3280,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 diff --git a/gateway/session.py b/gateway/session.py index 17ca8e4d5..de5d1306b 100644 --- a/gateway/session.py +++ b/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) diff --git a/hermes_state.py b/hermes_state.py index 84c3bf44a..2f273b325 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -227,15 +227,17 @@ 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.""" + """Increment token counters on a session. Optionally backfill model.""" self._conn.execute( """UPDATE sessions SET input_tokens = input_tokens + ?, - output_tokens = output_tokens + ? + output_tokens = output_tokens + ?, + model = COALESCE(model, ?) WHERE id = ?""", - (input_tokens, output_tokens, session_id), + (input_tokens, output_tokens, model, session_id), ) self._conn.commit()