Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions gateway/slash_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
78 changes: 59 additions & 19 deletions hermes_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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, ?)``
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down
87 changes: 87 additions & 0 deletions tests/cli/test_resume_model_restore.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
2 changes: 1 addition & 1 deletion tui_gateway/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading