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"