Provider/route persistence after mid-session /model switches

Fixes on record

Merged to main. A stamp counts the issues a fix is linked to.

• #52666 fix(hermes_state): persist billing route after mid-session /model switch (#48248)

  merged to main

  files: agent/agent_runtime_helpers.py, hermes_state.py, tests/test_hermes_state.py

  Show the diff

  • 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,

  • hermes_state.py

    @@ -1584,6 +1584,36 @@ def _do(conn):
                 )
             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,

  • tests/test_hermes_state.py

    @@ -210,6 +210,54 @@ def test_update_session_model_overwrites_existing(self, db):
                                    model="xiaomi/mimo-v2.5-pro")
             assert db.get_session("s1")["model"] == "xiaomi/mimo-v2.5"
     
    +    def test_update_session_billing_route_overwrites_after_switch(self, db):
    +        """A mid-session provider switch must overwrite the billing route.
    +
    +        update_token_counts writes billing fields with
    +        COALESCE(billing_provider, ?) (first-writer-wins), so after a
    +        provider switch the dashboard kept attributing cost to the original
    +        provider (#48248). update_session_billing_route sets them
    +        unconditionally and nulls system_prompt so the next turn rebuilds
    +        the Model:/Provider: header (#48173).
    +        """
    +        db.create_session(session_id="s1", source="telegram")
    +        # First token update seeds the billing route.
    +        db.update_token_counts("s1", input_tokens=10, output_tokens=5,
    +                               billing_provider="openrouter",
    +                               billing_base_url="https://openrouter.ai/api/v1",
    +                               billing_mode="api_key")
    +        sess = db.get_session("s1")
    +        assert sess["billing_provider"] == "openrouter"
    +        # A later token update never changes it (COALESCE first-writer-wins).
    +        db.update_token_counts("s1", input_tokens=10, output_tokens=5,
    +                               billing_provider="ollama",
    +                               billing_base_url="http://localhost:11434/v1",
    +                               billing_mode="local")
    +        assert db.get_session("s1")["billing_provider"] == "openrouter"
    +
    +        # Seed a stale prompt snapshot, then switch the billing route.
    +        db.update_system_prompt("s1", "Model: x/old\nProvider: openrouter")
    +        assert db.get_session("s1")["system_prompt"] is not None
    +        db.update_session_billing_route(
    +            "s1", provider="ollama",
    +            base_url="http://localhost:11434/v1", billing_mode="local",
    +        )
    +        sess = db.get_session("s1")
    +        assert sess["billing_provider"] == "ollama"
    +        assert sess["billing_base_url"] == "http://localhost:11434/v1"
    +        assert sess["billing_mode"] == "local"
    +        assert sess["system_prompt"] is None, \
    +            "system_prompt must be nulled so the next turn rebuilds Model:/Provider:"
    +
    +        # billing_mode defaults to COALESCE — omitting it preserves the value.
    +        db.update_session_billing_route(
    +            "s1", provider="openai",
    +            base_url="https://api.openai.com/v1",
    +        )
    +        sess = db.get_session("s1")
    +        assert sess["billing_provider"] == "openai"
    +        assert sess["billing_mode"] == "local"  # preserved (COALESCE on None)
    +
         def test_parent_session(self, db):
             db.create_session(session_id="parent", source="cli")
             db.create_session(session_id="child", source="cli", parent_session_id="parent")