diff --git a/litellm/proxy/auth/key_disallowed_models.py b/litellm/proxy/auth/key_disallowed_models.py index aec2e1cd50e..9842fa4a4e6 100644 --- a/litellm/proxy/auth/key_disallowed_models.py +++ b/litellm/proxy/auth/key_disallowed_models.py @@ -37,16 +37,17 @@ def _names_model_resolves_through( after_key_alias: Final = key_aliases.get(after_team_alias, after_team_alias) after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias) dispatched: Final = key_aliases.get(after_global_alias, after_global_alias) - router_alias_target: Final = ( - llm_router.registry().resolve(model, kind="alias") - if llm_router is not None and model in llm_router.registry().aliases() - else None - ) - chain: Final = frozenset( - name - for name in (model, after_team_alias, after_key_alias, after_global_alias, dispatched, router_alias_target) - if name is not None + aliased: Final = frozenset((model, after_team_alias, after_key_alias, after_global_alias, dispatched)) + router_alias_targets: Final = ( + () + if llm_router is None + else ( + llm_router.registry().resolve(name, kind="alias") + for name in aliased + if name in llm_router.registry().aliases() + ) ) + chain: Final = aliased | frozenset(target for target in router_alias_targets if target is not None) return chain | frozenset( resolved for name in chain diff --git a/tests/integration/authorization/test_key_disallowed_models.py b/tests/integration/authorization/test_key_disallowed_models.py index fed158ed68b..b7a02898e62 100644 --- a/tests/integration/authorization/test_key_disallowed_models.py +++ b/tests/integration/authorization/test_key_disallowed_models.py @@ -21,6 +21,7 @@ from tests.integration._support.client import ( ) from tests.integration._support.database import read_rows from tests.integration._support.upstream import delete_scenario, register_scenario +from tests.integration.authorization._hidden_alias_budget import install_aliases, remove_aliases from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse @@ -206,6 +207,48 @@ def test_a_key_alias_to_a_deployment_id_cannot_reach_a_disallowed_model(gateway: _assert_denied(gateway, key, {"model": alias}, alias) +def test_a_key_alias_to_a_router_alias_cannot_reach_a_disallowed_model(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + disallowed: Final = scenario.model() + router_alias: Final = f"router-alias-{uuid.uuid4().hex[:8]}" + key_alias: Final = f"alias-{uuid.uuid4().hex[:8]}" + install_aliases(gateway, {router_alias: disallowed}) + scenario.cleanups.callback(remove_aliases, gateway, frozenset({router_alias})) + unrestricted: Final = scenario.key(aliases={key_alias: router_alias}) + served: Final = eventually( + lambda: tuple(_chat(gateway, unrestricted, {"model": key_alias}).status_code for _ in range(6)), + lambda statuses: frozenset(statuses) == {200}, + seconds=60, + ) + assert frozenset(served) == {200}, served + key: Final = scenario.key(aliases={key_alias: router_alias}, disallowed_models=[disallowed]) + _assert_denied(gateway, key, {"model": key_alias}, key_alias) + + +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', + ([sha256(key.encode()).hexdigest() for key in keys],), + ) + + +@pytest.mark.parametrize("count", [1, 2], ids=["one-key", "two-keys"]) +def test_deleting_keys_keeps_their_disallowed_models_on_the_deleted_records(gateway: Gateway, count: int) -> None: + prefix: Final = f"deleted-disallowed-{uuid.uuid4().hex[:8]}" + keys: Final = tuple( + string_value( + gateway.post("/key/generate", {"key_alias": f"{prefix}-{index}", "disallowed_models": ["a", f"b-{index}"]})[ + "key" + ] + ) + for index in range(count) + ) + deleted: Final = gateway.request("POST", "/key/delete", {"keys": list(keys)}) + assert deleted.status_code == 200, f"{deleted.status_code} {deleted.text}" + assert _deleted_disallowed_models(keys) == [{"disallowed_models": ["a", f"b-{index}"]} for index in range(count)] + assert [row for key in keys for row in _stored_disallowed_models(key)] == [] + + def test_a_team_deployments_routing_name_cannot_reach_a_disallowed_model(gateway: Gateway) -> None: with gateway.scenario() as scenario: team: Final = scenario.team()