diff --git a/api/routes.py b/api/routes.py index 9236db6b6..8db3c0598 100644 --- a/api/routes.py +++ b/api/routes.py @@ -13138,7 +13138,8 @@ def _handle_live_models(handler, parsed): # Fallback: try credential pool for base_url + api_key if (not _base_url or not _api_key) and provider.startswith("custom:"): try: - from api.config import _has_explicit_pool_credentials, _resolve_provider_alias + from api.config import _has_explicit_pool_credentials + if _has_explicit_pool_credentials(provider): from agent.credential_pool import load_pool as _lpool _resolved = _resolve_provider_alias(provider) @@ -13218,11 +13219,19 @@ def _handle_live_models(handler, parsed): import urllib.request _providers_cfg = cfg.get("providers", {}) _prov = _providers_cfg.get(provider, {}) if isinstance(_providers_cfg, dict) else {} - # Only use provider-scoped key — never fall back to a top-level - # api_key which may belong to a different provider. + # Only use a provider-scoped key. A top-level model.api_key + # is safe here only when it belongs to the requested provider; + # otherwise /api/models/live?provider= could forward + # the active provider's credential to the wrong third party. _key = _prov.get("api_key") if isinstance(_prov, dict) else None if not _key: - _key = cfg.get("model", {}).get("api_key") + _model_cfg = cfg.get("model", {}) + if isinstance(_model_cfg, dict): + _active_provider = _resolve_provider_alias( + (_model_cfg.get("provider") or "").strip().lower() + ) + if _active_provider == provider: + _key = _model_cfg.get("api_key") if _key: _req = urllib.request.Request( f"{_ep}/models", diff --git a/tests/test_byok_model_dropdown.py b/tests/test_byok_model_dropdown.py index 451a3c089..8b9bdcd51 100644 --- a/tests/test_byok_model_dropdown.py +++ b/tests/test_byok_model_dropdown.py @@ -247,7 +247,6 @@ class TestLiveModelsCustomProviderFallback: return True monkeypatch.setattr(r, "j", fake_j) - from urllib.parse import urlparse parsed = mock.MagicMock() parsed.query = "provider=custom" @@ -373,6 +372,140 @@ class TestLiveModelsCustomProviderFallback: ] assert [m["id"] for m in resp["models"]] == ["right-live-model"] + def test_standard_provider_live_fetch_does_not_reuse_active_provider_key(self, monkeypatch): + """A requested provider must not receive another provider's top-level key.""" + import urllib.request + + requests = [] + + def fake_urlopen(req, timeout=None): + requests.append( + { + "url": req.full_url, + "authorization": req.headers.get("Authorization"), + "timeout": timeout, + } + ) + raise AssertionError("unexpected cross-provider live fetch") + + cfg = { + "model": { + "provider": "openai", + "api_key": "active-provider-canary", + }, + "providers": {"mistralai": {}}, + } + monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) + + resp = self._call_live_models(monkeypatch, cfg, "mistralai") + + assert requests == [] + assert resp["provider"] == "mistralai" + assert resp["models"], "static fallback models should still be returned" + + def test_standard_provider_live_fetch_can_use_matching_top_level_key(self, monkeypatch): + """The active provider's top-level key remains valid for that same provider.""" + import json + import urllib.request + + requests = [] + + class Response: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def read(self): + return json.dumps({"data": [{"id": "mistral-live-model"}]}).encode("utf-8") + + def fake_urlopen(req, timeout=None): + requests.append( + { + "url": req.full_url, + "authorization": req.headers.get("Authorization"), + "timeout": timeout, + } + ) + return Response() + + cfg = { + "model": { + "provider": "mistralai", + "api_key": "active-provider-canary", + }, + "providers": {"mistralai": {}}, + } + monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) + + resp = self._call_live_models(monkeypatch, cfg, "mistralai") + + assert requests == [ + { + "url": "https://api.mistral.ai/v1/models", + "authorization": "Bearer active-provider-canary", + "timeout": 8, + } + ] + assert resp["provider"] == "mistralai" + + def test_standard_provider_live_fetch_allows_matching_active_provider_alias(self, monkeypatch): + """Alias-equivalent active providers should still count as the same provider.""" + import json + import urllib.request + + requests = [] + + class Response: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def read(self): + return json.dumps({"data": [{"id": "zai-live-model"}]}).encode("utf-8") + + def fake_urlopen(req, timeout=None): + requests.append( + { + "url": req.full_url, + "authorization": req.headers.get("Authorization"), + "timeout": timeout, + } + ) + return Response() + + cfg = { + "model": { + "provider": "z.ai", + "api_key": "active-provider-canary", + }, + "providers": {"zai": {}}, + } + import api.config as c + import api.routes as r + + monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) + monkeypatch.setattr(c, "get_config", lambda: cfg) + monkeypatch.setattr(r, "j", lambda _handler, payload, **_kw: payload) + self._install_provider_model_ids(monkeypatch, lambda _p: []) + + parsed = mock.MagicMock() + parsed.query = "provider=zai" + r._clear_live_models_cache() + resp = r._handle_live_models(object(), parsed) + + assert requests == [ + { + "url": "https://api.z.ai/v1/models", + "authorization": "Bearer active-provider-canary", + "timeout": 8, + } + ] + assert resp["provider"] == "zai" + # ── Regression: known-good providers still work ───────────────────────────────