diff --git a/gateway/run.py b/gateway/run.py index b08650d08ac..021e5f1ba80 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -1870,6 +1870,23 @@ def _resolve_runtime_agent_kwargs_for_provider(provider: str) -> dict: } +def _credential_pool_for_provider(provider: Optional[str]): + """Return the live credential pool for a provider id (e.g. ``custom:hyper``).""" + if not provider or not str(provider).strip(): + return None + try: + return _resolve_runtime_agent_kwargs_for_provider(str(provider).strip()).get( + "credential_pool" + ) + except Exception: + logger.debug( + "Failed to resolve credential pool for provider=%s", + provider, + exc_info=True, + ) + return None + + def _try_resolve_fallback_provider() -> dict | None: """Attempt to resolve credentials from the fallback_model/fallback_providers config.""" from hermes_cli.runtime_provider import resolve_runtime_provider @@ -3666,8 +3683,13 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew "base_url": override.get("base_url"), "api_mode": override.get("api_mode"), "max_tokens": override.get("max_tokens"), + "credential_pool": override.get("credential_pool"), } if override_runtime.get("api_key"): + if override_runtime.get("credential_pool") is None: + override_runtime["credential_pool"] = _credential_pool_for_provider( + override.get("provider") + ) logger.debug( "Session model override (fast): session=%s config_model=%s -> override_model=%s provider=%s", resolved_session_key or "", model, override_model, @@ -15275,6 +15297,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew runtime = _resolve_runtime_agent_kwargs_for_provider(provider) override["api_key"] = runtime.get("api_key") override["api_mode"] = runtime.get("api_mode") + override["credential_pool"] = runtime.get("credential_pool") if not override.get("base_url"): override["base_url"] = runtime.get("base_url") except Exception: @@ -15304,10 +15327,18 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew if not override: return model, runtime_kwargs model = override.get("model", model) - for key in ("provider", "api_key", "base_url", "api_mode"): + for key in ("provider", "api_key", "base_url", "api_mode", "credential_pool"): val = override.get(key) if val is not None: runtime_kwargs[key] = val + if ( + runtime_kwargs.get("api_key") + and runtime_kwargs.get("credential_pool") is None + and override.get("provider") + ): + runtime_kwargs["credential_pool"] = _credential_pool_for_provider( + override.get("provider") + ) return model, runtime_kwargs def _is_intentional_model_switch(self, session_key: str, agent_model: str) -> bool: diff --git a/tests/gateway/test_session_model_override_credential_pool.py b/tests/gateway/test_session_model_override_credential_pool.py new file mode 100644 index 00000000000..8720240db82 --- /dev/null +++ b/tests/gateway/test_session_model_override_credential_pool.py @@ -0,0 +1,69 @@ +"""Session /model overrides must attach credential_pool for 402 rotation.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from gateway.run import GatewayRunner, _credential_pool_for_provider + + +def test_fast_session_override_includes_credential_pool(monkeypatch): + runner = object.__new__(GatewayRunner) + runner._session_model_overrides = { + "sess-1": { + "model": "kimi-k2.7", + "provider": "custom:hyper", + "api_key": "sk-test", + "base_url": "https://hyper.charm.land/v1", + "api_mode": "chat_completions", + }, + } + fake_pool = object() + + monkeypatch.setattr( + "gateway.run._resolve_gateway_model", + lambda _uc=None: "default-model", + ) + monkeypatch.setattr( + "gateway.run._credential_pool_for_provider", + lambda provider: fake_pool if provider == "custom:hyper" else None, + ) + + model, runtime = runner._resolve_session_agent_runtime(session_key="sess-1") + + assert model == "kimi-k2.7" + assert runtime.get("credential_pool") is fake_pool + + +def test_apply_session_override_backfills_credential_pool(monkeypatch): + runner = object.__new__(GatewayRunner) + fake_pool = MagicMock(name="pool") + runner._session_model_overrides = { + "sess-2": { + "model": "kimi-k2.7", + "provider": "custom:hyper", + "api_key": "sk-test", + }, + } + monkeypatch.setattr( + "gateway.run._credential_pool_for_provider", + lambda provider: fake_pool, + ) + + model, runtime = runner._apply_session_model_override( + "sess-2", + "default-model", + {"api_key": "old", "provider": "x"}, + ) + + assert model == "kimi-k2.7" + assert runtime["credential_pool"] is fake_pool + + +def test_credential_pool_for_provider_delegates(monkeypatch): + sentinel = object() + monkeypatch.setattr( + "gateway.run._resolve_runtime_agent_kwargs_for_provider", + lambda p: {"credential_pool": sentinel, "provider": p}, + ) + assert _credential_pool_for_provider("custom:hyper") is sentinel \ No newline at end of file