fix(gateway): attach credential_pool to session /model overrides

Per-session /model overrides supplied api_key and provider but omitted
credential_pool, so billing rotation never ran on HTTP 402. Wire the pool
on fast override, rehydrate, and apply paths; backfill from provider for
legacy persisted overrides. Regression tests in tests/gateway/.
This commit is contained in:
Tuna Dev 2026-07-05 01:35:55 +08:00 committed by Teknium
parent 10f7cb043c
commit ff4c8172ce
2 changed files with 101 additions and 1 deletions

View file

@ -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:

View file

@ -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