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
This commit is contained in:
Devin AI 2026-11-15 00:34:27 +00:00
parent cb622a8542
commit 3d4b99f543
8 changed files with 181 additions and 0 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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