diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index eb19f61bd0c52..52d8fc78b3e1d 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -1937,7 +1937,8 @@ async def _on_model_selected_scoped( event.source ) await _sess_db.update_session_model( - _sess_entry.session_id, result.new_model + _sess_entry.session_id, result.new_model, + provider=result.target_provider, ) except Exception as exc: logger.debug( @@ -2247,7 +2248,8 @@ async def _finish_switch() -> str: if getattr(_sess_entry, "was_auto_reset", False): _sess_entry.was_auto_reset = False await _sess_db.update_session_model( - _sess_entry.session_id, result.new_model + _sess_entry.session_id, result.new_model, + provider=result.target_provider, ) except Exception as exc: logger.debug( diff --git a/hermes_state.py b/hermes_state.py index 0188e5d59bd12..1f44cfc184f5d 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -255,6 +255,15 @@ def _delegate_from_json(col: str = "model_config") -> str: # None result ("merged config is empty → store NULL"). _MODEL_CONFIG_ROW_MISSING = object() +# Billing-bucket classes that aren't a routable provider identity on their +# own — used by session_gateway_runtime's billing_provider fallback and by +# tui_gateway.server._stored_session_runtime_overrides. A session that +# persisted only one of these (never ran /model) must fall back to the +# ambient config default rather than restore a bare bucket. Shared here so +# both consumers stay in sync (previously duplicated as a set in +# tui_gateway/server.py). +_BARE_BILLING_PROVIDERS = frozenset({"auto", "custom"}) + def _cwd_prefix_clause(cwd_prefix: str) -> Tuple[str, List[str]]: prefix = cwd_prefix.rstrip("/\\") or cwd_prefix @@ -5893,7 +5902,9 @@ def _do(conn): self._delete_unreferenced_system_prompts(conn) self._execute_write(_do) - def update_session_model(self, session_id: str, model: str) -> None: + def update_session_model( + self, session_id: str, model: str, provider: Optional[str] = None + ) -> None: """Update the model for a session after a mid-session switch. Unlike ``update_token_counts`` which uses ``COALESCE(model, ?)`` @@ -5903,6 +5914,13 @@ def update_session_model(self, session_id: str, model: str) -> None: footer metadata is rebuilt on the next turn. A successful /model switch explicitly replaces any confirmed Browser runtime lock while preserving unrelated lineage markers in ``model_config``. + + When *provider* is given, it is merged into ``model_config`` + alongside the model (``$.model`` / ``$.provider``) so a later + resume recombines the persisted model with the provider that + actually serves it instead of the config.yaml primary provider + (#79536). Callers without provider knowledge leave any stored + provider untouched. """ # This write bypasses the token queue, so deltas enqueued before the # switch must land first: a still-queued first delta carries the @@ -5913,19 +5931,24 @@ def update_session_model(self, session_id: str, model: str) -> None: self.flush_token_counts() def _do(conn): + # Use the shared merge discipline so lineage markers like + # _branched_from / _delegate_from survive. browser_model_lock + # is deleted via a None patch value (same semantics as the + # old json_remove). + patch: Dict[str, Any] = {"browser_model_lock": None} + if model: + patch["model"] = model + if provider: + patch["provider"] = provider + merged = self._merge_model_config_json(conn, session_id, patch) + if merged is _MODEL_CONFIG_ROW_MISSING: + return conn.execute( - """UPDATE sessions SET - model = ?, - model_config = CASE - WHEN model_config IS NULL THEN NULL - WHEN json_valid(model_config) - THEN json_remove(model_config, '$.browser_model_lock') - ELSE model_config - END, - system_prompt = NULL, - system_prompt_hash = NULL - WHERE id = ?""", - (model, session_id), + "UPDATE sessions SET " + "model = ?, model_config = ?, " + "system_prompt = NULL, system_prompt_hash = NULL " + "WHERE id = ?", + (model, merged, session_id), ) self._delete_unreferenced_system_prompts(conn) self._execute_write(_do) @@ -6118,18 +6141,21 @@ def session_gateway_runtime(session_meta: Optional[Dict[str, Any]]) -> Dict[str, ``gateway_runtime`` key (written by the gateway's ``_sync_session_model_from_agent`` and the CLI ``/model`` persist), falling back to the top-level ``provider``/``base_url``/``api_mode`` - keys the TUI gateway's ``_runtime_model_config`` writes. Returns an - empty dict on any parse failure — resume falls back to ambient - config resolution. + keys the TUI gateway's ``_runtime_model_config`` writes. As a last + resort, falls back to the ``billing_provider`` column (written on + every session's first accounted API call) so sessions that never ran + ``/model`` still restore the provider that actually served them. + Returns an empty dict on any parse failure — resume falls back to + ambient config resolution. """ raw = (session_meta or {}).get("model_config") if isinstance(raw, str): try: raw = json.loads(raw) except Exception: - return {} + raw = {} if not isinstance(raw, dict): - return {} + raw = {} runtime = raw.get("gateway_runtime") if isinstance(runtime, dict) and runtime.get("provider"): # Filter None values: the persist path writes or-None to trigger @@ -6143,7 +6169,21 @@ def session_gateway_runtime(session_meta: Optional[Dict[str, Any]]) -> Dict[str, } if top_level: return top_level - return dict(runtime) if isinstance(runtime, dict) else {} + # Last resort: billing_provider column. Written via COALESCE on every + # session's first accounted API call — the only durable record for + # sessions that never ran /model. Mirrors the TUI gateway's + # _stored_session_runtime_overrides fallback. Bare billing buckets + # ("auto"/"custom") are not routable identities — filter them out so + # resume falls back to the ambient config default instead. + billing_provider = str( + (session_meta or {}).get("billing_provider") or "" + ).strip() + if ( + billing_provider + and billing_provider.lower() not in _BARE_BILLING_PROVIDERS + ): + return {"provider": billing_provider} + return {k: v for k, v in (runtime or {}).items() if v is not None} if isinstance(runtime, dict) else {} def update_session_billing_route( self, diff --git a/tests/cli/test_resume_model_restore.py b/tests/cli/test_resume_model_restore.py index 766490c54d88e..aa01462527e7d 100644 --- a/tests/cli/test_resume_model_restore.py +++ b/tests/cli/test_resume_model_restore.py @@ -278,3 +278,90 @@ def test_round_trip_persist_then_restore(tmp_path, monkeypatch): assert restored.model == "deepseek-v4-flash-free" assert restored.provider == "custom:opencode-zen" assert restored.base_url == "https://oz/v1" + + +# ── update_session_model provider persistence (#79536) ────────────── + + +def test_update_session_model_persists_provider(tmp_path, monkeypatch): + """update_session_model writes $.model + $.provider into model_config.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + db = SessionDB(db_path=tmp_path / "state.db") + db.create_session(session_id="s1", source="cli", model="m0") + db.update_session_model("s1", "claude-x", provider="custom:feather") + meta = db.get_session("s1") + assert meta["model"] == "claude-x" + config = json.loads(meta["model_config"]) + assert config["model"] == "claude-x" + assert config["provider"] == "custom:feather" + + +def test_update_session_model_without_provider_preserves_existing(tmp_path, monkeypatch): + """Without provider, existing $.provider is left untouched.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + db = SessionDB(db_path=tmp_path / "state.db") + db.create_session(session_id="s2", source="cli", model="m0") + db.update_session_model("s2", "claude-x", provider="custom:feather") + db.update_session_model("s2", "gpt-5.4") # no provider + meta = db.get_session("s2") + config = json.loads(meta["model_config"]) + assert config["model"] == "gpt-5.4" + assert config["provider"] == "custom:feather" # preserved + + +def test_update_session_model_null_model_config_with_provider(tmp_path, monkeypatch): + """Provider persistence works when model_config starts as NULL.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + db = SessionDB(db_path=tmp_path / "state.db") + db.create_session(session_id="s3", source="cli", model="m0") + # model_config is NULL at creation — update_session_model must create it + db.update_session_model("s3", "claude-x", provider="minimax") + meta = db.get_session("s3") + config = json.loads(meta["model_config"]) + assert config["model"] == "claude-x" + assert config["provider"] == "minimax" + + +# ── session_gateway_runtime billing_provider fallback (#85721) ───── + + +def test_session_gateway_runtime_falls_back_to_billing_provider(): + """Sessions that never ran /model have only billing_provider.""" + meta = { + "model": "glm-4.7", + "model_config": None, + "billing_provider": "minimax", + } + runtime = SessionDB.session_gateway_runtime(meta) + assert runtime == {"provider": "minimax"} + + +def test_session_gateway_runtime_billing_provider_bare_bucket_ignored(): + """Bare billing buckets (auto/custom) are not routable — skip them.""" + for bare in ("auto", "custom"): + meta = { + "model": "m", + "model_config": None, + "billing_provider": bare, + } + assert SessionDB.session_gateway_runtime(meta) == {} + + +def test_session_gateway_runtime_explicit_provider_wins_over_billing(): + """Explicit model_config provider takes precedence over billing_provider.""" + meta = _row(model_config={"provider": "nous"}) + meta["billing_provider"] = "minimax" + runtime = SessionDB.session_gateway_runtime(meta) + assert runtime == {"provider": "nous"} + + +def test_restore_session_model_restores_billing_provider_fallback(): + """End-to-end: _restore_session_model uses billing_provider fallback.""" + stub = _make_stub() + stub._restore_session_model({ + "model": "glm-4.7", + "model_config": None, + "billing_provider": "minimax", + }) + assert stub.model == "glm-4.7" + assert stub.provider == "minimax" diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 04b3c956e090f..d1b58dd185500 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -3775,7 +3775,7 @@ def _resolve_startup_runtime() -> tuple[str, str | None]: # ``billing_provider="openrouter"``; dropping it forces resume to the current # global model (e.g. a custom endpoint), which is the wrong provider for the # stored model. See #57588. -_BARE_BILLING_PROVIDERS = {"auto", "custom"} +from hermes_state import _BARE_BILLING_PROVIDERS def _stored_session_runtime_overrides(row: dict | None) -> dict: