From 3d4b99f5433df0f8aac8e00859791a777b68b628 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 15 Nov 2026 00:34:27 +0000 Subject: [PATCH] fix(proxy): keep auto routers from sending a key to a tier in its disallowed_models The deny check only saw the requested route name, so a complexity or auto router could pick a tier the key disallows and dispatch it. The pre-routing step now asks the router host about the selected model, and the proxy host runs the key's deny list on it. Requests with no key in metadata are unchanged. Also adds unit tests for the deny list permission check, the null update and the batch row check --- litellm/proxy/_router_host.py | 12 ++++ litellm/proxy/auth/auto_router_checks.py | 10 +++ litellm/router_host.py | 8 +++ litellm/router_rust/steps/prerouting.py | 6 ++ .../test_key_disallowed_models.py | 21 +++++++ .../proxy/hooks/test_batch_file_validation.py | 2 + .../test_key_management_endpoints.py | 60 ++++++++++++++++++ tests/unit/test_router/test_router_host.py | 62 +++++++++++++++++++ 8 files changed, 181 insertions(+) diff --git a/litellm/proxy/_router_host.py b/litellm/proxy/_router_host.py index 98735d382e5..b9e391c5f9b 100644 --- a/litellm/proxy/_router_host.py +++ b/litellm/proxy/_router_host.py @@ -45,6 +45,18 @@ class ProxyRouterHost(RouterHost): llm_router=llm_router, ) + def authorize_routed_model( + self, + model: str, + request_kwargs: Mapping[str, object], + llm_router: "Router", + ) -> None: + from litellm.proxy.auth import auto_router_checks + + auto_router_checks.ensure_key_may_call_routed_model( + model=model, request_kwargs=request_kwargs, llm_router=llm_router + ) + def compression_policy( self, llm_router: "Router | None", diff --git a/litellm/proxy/auth/auto_router_checks.py b/litellm/proxy/auth/auto_router_checks.py index b83e1f3fffe..30711fba4f5 100644 --- a/litellm/proxy/auth/auto_router_checks.py +++ b/litellm/proxy/auth/auto_router_checks.py @@ -20,6 +20,16 @@ def _mapping(value: object) -> Mapping[str, object] | None: return None +def ensure_key_may_call_routed_model(*, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> None: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.key_disallowed_models import ensure_key_may_call_model + + metadata: Final = _mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))) + actor: Final = metadata.get("user_api_key_auth") if metadata is not None else None + if isinstance(actor, UserAPIKeyAuth): + ensure_key_may_call_model(model, actor, llm_router) + + async def authorize_member_auto_router_inference( *, deployment: Mapping[str, object] | None, diff --git a/litellm/router_host.py b/litellm/router_host.py index 050f5ad6a05..8fe5e76460d 100644 --- a/litellm/router_host.py +++ b/litellm/router_host.py @@ -63,6 +63,14 @@ class RouterHost: response=httpx.Response(status_code=403, request=httpx.Request("POST", "https://litellm.invalid")), ) + def authorize_routed_model( + self, + model: str, + request_kwargs: Mapping[str, object], + llm_router: "Router", + ) -> None: + """Raise when the caller may not be sent to ``model``, the one a routing strategy selected.""" + def compression_policy( self, llm_router: "Router | None", diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index 9e5d5ede967..4b9a1a5f97b 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -921,6 +921,12 @@ class PreRouter: call_type=call_type, ) hooked: Final[PreRoutingHookResponse | None] = routed + if routed is not None: + router_host().authorize_routed_model( + model=routed.model, + request_kwargs=request_kwargs, + llm_router=public_router, # pyright: ignore[reportArgumentType] # the litellm.Router this core serves + ) traced: Final = ( hooked.model_copy(update={"routing_decision": {**provenance, **hooked.routing_decision}}) if hooked is not None and hooked.routing_decision is not None diff --git a/tests/integration/authorization/test_key_disallowed_models.py b/tests/integration/authorization/test_key_disallowed_models.py index b7a02898e62..1a4d3dec192 100644 --- a/tests/integration/authorization/test_key_disallowed_models.py +++ b/tests/integration/authorization/test_key_disallowed_models.py @@ -225,6 +225,27 @@ def test_a_key_alias_to_a_router_alias_cannot_reach_a_disallowed_model(gateway: _assert_denied(gateway, key, {"model": key_alias}, key_alias) +def test_a_complexity_router_cannot_send_a_key_to_a_disallowed_tier(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + keyword: Final = f"simple-{uuid.uuid4().hex}" + tier_upstream: Final = f"tier-{uuid.uuid4().hex[:8]}" + disallowed: Final = scenario.model(model=f"openai/{tier_upstream}") + other_tier: Final = scenario.model() + router: Final = scenario.model( + model="auto_router/complexity_router", + complexity_router_config={ + "classifier_type": "heuristic", + "keyword_tier_rules": [{"keywords": [keyword], "tier": "SIMPLE"}], + "tiers": {"SIMPLE": disallowed, "MEDIUM": other_tier, "COMPLEX": other_tier, "REASONING": other_tier}, + }, + ) + body: Final[dict[str, JsonValue]] = {"model": router, "messages": [{"role": "user", "content": keyword}]} + served: Final = _chat(gateway, scenario.key(), body) + assert served.status_code == 200, f"{served.status_code} {served.text}" + assert tier_upstream in _observed_models(gateway) + _assert_denied(gateway, scenario.key(disallowed_models=[disallowed]), body, disallowed) + + def _deleted_disallowed_models(keys: tuple[str, ...]) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT disallowed_models FROM "LiteLLM_DeletedVerificationToken" WHERE token = ANY(%s) ORDER BY key_alias', diff --git a/tests/unit/proxy/hooks/test_batch_file_validation.py b/tests/unit/proxy/hooks/test_batch_file_validation.py index df554eacf23..01ea6590e8a 100644 --- a/tests/unit/proxy/hooks/test_batch_file_validation.py +++ b/tests/unit/proxy/hooks/test_batch_file_validation.py @@ -980,6 +980,8 @@ def test_key_requires_batch_model_access_check_branches(): check(UserAPIKeyAuth(api_key="sk", models=[], access_group_ids=["grp"])) is True ) assert check(UserAPIKeyAuth(api_key="sk", models=[])) is False + assert check(UserAPIKeyAuth(api_key="sk", models=[], disallowed_models=["gpt-denied"])) is True + assert check(UserAPIKeyAuth(api_key="sk", models=["*"], disallowed_models=["gpt-denied"])) is True assert check(UserAPIKeyAuth(api_key="sk", models=["gpt-4o-mini"])) is True # Wildcard / all-proxy-models grant access to every model, so # can_key_call_model passes any model regardless of access groups (which diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 5a51a3ad013..d4a65392fa7 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -2143,6 +2143,66 @@ async def test_update_key_keeps_a_due_budget_reset_when_the_duration_is_resent( assert ("budget_reset_at" in updated) is moves_reset +@pytest.mark.asyncio +async def test_update_key_with_null_disallowed_models_clears_the_stored_deny_list() -> None: + existing_key: Final = LiteLLM_VerificationToken(token="hashed", disallowed_models=("claude-*",)) + + updated: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key="sk-1", disallowed_models=None), existing_key_row=existing_key + ) + + assert updated["disallowed_models"] == [] + + +_DENY_LIST_TEAM: Final = LiteLLM_TeamTableCachedObj( + team_id="team-1", + members_with_roles=[Member(user_id="team-admin", role="admin"), Member(user_id="member", role="user")], +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller_id", "key_team_id", "requested", "allowed"), + [ + ("team-admin", "team-1", ["claude-*"], True), + ("member", "team-1", ["claude-*"], False), + ("member", None, ["claude-*"], False), + ("member", "team-1", ["claude-*", "gpt-denied", "another"], True), + ], + ids=["team-admin-removes", "team-member-removes", "no-team-removes", "team-member-only-adds"], +) +async def test_only_admins_of_the_keys_team_may_remove_disallowed_models( + monkeypatch: pytest.MonkeyPatch, caller_id: str, key_team_id: str | None, requested: list[str], allowed: bool +) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import check_disallowed_models_caller_permission + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(return_value=_DENY_LIST_TEAM), + ) + + async def check() -> None: + await check_disallowed_models_caller_permission( + data=UpdateKeyRequest(key="sk-1", disallowed_models=requested), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-caller", user_id=caller_id + ), + existing_key_row=LiteLLM_VerificationToken( + token="hashed", team_id=key_team_id, disallowed_models=("claude-*", "gpt-denied") + ), + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + + if allowed: + await check() + return + with pytest.raises(HTTPException) as refused: + await check() + assert refused.value.status_code == 403 + assert "gpt-denied" in str(refused.value.detail) + + @pytest.mark.asyncio async def test_update_key_clears_end_user_budget_id_with_empty_string(): """Sending an empty end_user_budget_id detaches the key default without touching any budget row, diff --git a/tests/unit/test_router/test_router_host.py b/tests/unit/test_router/test_router_host.py index 1da96e71d7b..aab3281417e 100644 --- a/tests/unit/test_router/test_router_host.py +++ b/tests/unit/test_router/test_router_host.py @@ -75,3 +75,65 @@ def test_default_host_keeps_nothing_and_compresses_nothing(): assert host.compression_policy(None, "auto", {}, ()) is None assert host.model_hop_compression_armed() is False assert asyncio.run(host.messages_for_routing(None, [{"role": "user", "content": "hi"}], {})) is None + + +_TIER_KEYWORD = "zzsimpletier" + + +def _complexity_router_reply(metadata: dict[str, object] | None) -> str: + import asyncio + + from litellm import Router + + router = Router( + model_list=[ + {"model_name": tier, "litellm_params": {"model": f"openai/{tier}", "api_key": "k", "mock_response": tier}} + for tier in ("cheap", "smart") + ] + + [ + { + "model_name": "auto", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "heuristic", + "keyword_tier_rules": [{"keywords": [_TIER_KEYWORD], "tier": "SIMPLE"}], + "tiers": {"SIMPLE": "cheap", "MEDIUM": "smart", "COMPLEX": "smart", "REASONING": "smart"}, + }, + }, + } + ] + ) + extra = {} if metadata is None else {"metadata": metadata} + response = asyncio.run( + router.acompletion(model="auto", messages=[{"role": "user", "content": _TIER_KEYWORD}], **extra) + ) + return response.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming reply + + +@pytest.mark.parametrize("disallowed_models", [None, [], ["smart"]], ids=["no-key", "no-deny-list", "other-tier-denied"]) +def test_a_complexity_router_serves_its_selected_tier_when_the_key_does_not_disallow_it( + disallowed_models: list[str] | None, +): + from litellm.proxy._types import UserAPIKeyAuth + + metadata = ( + None + if disallowed_models is None + else {"user_api_key_auth": UserAPIKeyAuth(api_key="sk", disallowed_models=disallowed_models)} + ) + + assert _complexity_router_reply(metadata) == "cheap" + + +@pytest.mark.parametrize("deny_entry", ["cheap", "che*"]) +def test_a_complexity_router_cannot_send_a_key_to_a_tier_it_disallows(deny_entry: str): + from litellm.proxy._types import ModelAccessDeniedProxyException, ProxyErrorTypes, UserAPIKeyAuth + + key = UserAPIKeyAuth(api_key="sk", disallowed_models=[deny_entry]) + + with pytest.raises(ModelAccessDeniedProxyException) as denied: + _complexity_router_reply({"user_api_key_auth": key}) + + assert denied.value.type == ProxyErrorTypes.key_model_access_denied + assert denied.value.code == "403"