diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index ccf15307b07..3473eb0b54c 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -1697,6 +1697,27 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo old_model, old_provider, new_model, new_provider, ) + # ── Persist billing route to session DB ── + # The agent's _session_db / session_id may not be set in all contexts + # (tests, bare agents without a session DB, etc.). This ensures the + # dashboard Model cards show the actual provider after a mid-session + # /model switch instead of the stale session-creation provider. + # See #48248 for the full bug description. + _session_db = getattr(agent, "_session_db", None) + _session_id = getattr(agent, "session_id", None) + if _session_db is not None and _session_id: + try: + _session_db.update_session_billing_route( + _session_id, + provider=agent.provider, + base_url=agent.base_url, + billing_mode=getattr(agent, "api_mode", None), + ) + except Exception: + logger.warning( + "Failed to persist billing route after model switch", + exc_info=True, + ) def invoke_tool(agent, function_name: str, function_args: dict, effective_task_id: str, diff --git a/hermes_state.py b/hermes_state.py index 0c77d9c21dd..56efd47dabc 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -1584,6 +1584,36 @@ class SessionDB: ) self._execute_write(_do) + def update_session_billing_route( + self, + session_id: str, + *, + provider: str, + base_url: str, + billing_mode: Optional[str] = None, + ) -> None: + """Unconditionally update the billing provider/base_url for a session. + + Unlike ``update_token_counts`` which uses ``COALESCE(billing_provider, ?)`` + (only filling in NULL), this unconditionally sets the billing fields so + that the dashboard reflects the user's latest /model switch. + + Also nulls ``system_prompt`` so the cached snapshot (which embeds a + stale ``Model:`` / ``Provider:`` header) is rebuilt — matching the + behavior of ``update_session_model`` (see #48173, #48248). + """ + def _do(conn): + conn.execute( + """UPDATE sessions SET + billing_provider = ?, + billing_base_url = ?, + billing_mode = COALESCE(?, billing_mode), + system_prompt = NULL + WHERE id = ?""", + (provider, base_url, billing_mode, session_id), + ) + self._execute_write(_do) + def update_token_counts( self, session_id: str,