diff --git a/gateway/run.py b/gateway/run.py index 221f8f916..bc16b224d 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -1578,10 +1578,11 @@ class GatewayRunner: skip_db=agent_persisted, ) - # Update session with actual prompt token count from the agent + # Update session with actual prompt token count and model from the agent 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 @@ -3586,6 +3587,7 @@ class GatewayRunner: "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 @@ -3648,6 +3650,7 @@ class GatewayRunner: "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, "session_id": effective_session_id, } diff --git a/gateway/session.py b/gateway/session.py index 3e42db4fe..965f60793 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -590,6 +590,7 @@ class SessionStore: 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() @@ -607,7 +608,8 @@ class SessionStore: 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 5e29321ec..8945e195d 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -227,15 +227,17 @@ class SessionDB: 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 and backfill model if not already set.""" 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()