From 33761972678617dd6735ef3655784864f08f32f9 Mon Sep 17 00:00:00 2001 From: Jesus Velarde Date: Thu, 5 Nov 2026 18:24:23 -0800 Subject: [PATCH] feat(proxy): add key-level disallowed_models column and enforce it at auth and listing --- .../migration.sql | 7 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/verification_token.py | 1 + litellm/proxy/auth/auth_checks.py | 6 +- litellm/proxy/auth/key_disallowed_models.py | 80 ++++++++ litellm/proxy/auth/user_api_key_auth.py | 78 ++++---- .../key_management_endpoints.py | 45 ++++- litellm/proxy/proxy_server.py | 29 +-- litellm/proxy/schema.prisma | 2 + litellm/proxy/utils.py | 7 +- .../management_endpoints/key_requests.py | 1 + schema.prisma | 2 + .../test_key_disallowed_models.py | 172 ++++++++++++++++++ .../proxy/auth/test_key_disallowed_models.py | 118 ++++++++++++ 14 files changed, 494 insertions(+), 56 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_add_key_disallowed_models/migration.sql create mode 100644 litellm/proxy/auth/key_disallowed_models.py create mode 100644 tests/integration/authorization/test_key_disallowed_models.py create mode 100644 tests/unit/proxy/auth/test_key_disallowed_models.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_add_key_disallowed_models/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_add_key_disallowed_models/migration.sql new file mode 100644 index 00000000000..e9bbc3bb6a8 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261106000000_add_key_disallowed_models/migration.sql @@ -0,0 +1,7 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "disallowed_models" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "disallowed_models" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 818f2dc4858..cbf5e883794 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -530,6 +530,7 @@ model LiteLLM_VerificationToken { key_type String? policies String[] @default([]) access_group_ids String[] @default([]) + disallowed_models String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_fallbacks Json @default("{}") @@ -632,6 +633,7 @@ model LiteLLM_DeletedVerificationToken { key_type String? policies String[] @default([]) access_group_ids String[] @default([]) + disallowed_models String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_fallbacks Json @default("{}") diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 89048c56a9f..73fa654e093 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -56,6 +56,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): object_permission_id: str | None = None object_permission: LiteLLM_ObjectPermissionTable | None = None access_group_ids: list[str] | None = None + disallowed_models: tuple[str, ...] | None = None rotation_count: int | None = 0 auto_rotate: bool | None = False rotation_interval: str | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af60a315c60..b8982459440 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -100,6 +100,7 @@ from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, should_throttle_budget_exceeded, ) +from litellm.proxy.auth.key_disallowed_models import ensure_key_may_call_model from litellm.proxy.auth.model_access_denied import ( customer_model_access_denied_client_message, model_access_denied_client_message, @@ -5521,6 +5522,7 @@ async def can_key_call_model( Raises: - Exception: If token not allowed to call model """ + ensure_key_may_call_model(model, valid_token, llm_router) key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) try: return can_object_call_model( @@ -5596,7 +5598,9 @@ async def can_key_call_resolved_model( skip_key_model_check: Final = valid_token.config or ( isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models ) - if not skip_key_model_check: + if skip_key_model_check: + ensure_key_may_call_model(model, valid_token, llm_router) + else: await can_key_call_model( model=model, llm_model_list=llm_model_list, diff --git a/litellm/proxy/auth/key_disallowed_models.py b/litellm/proxy/auth/key_disallowed_models.py new file mode 100644 index 00000000000..ab10fc77f97 --- /dev/null +++ b/litellm/proxy/auth/key_disallowed_models.py @@ -0,0 +1,80 @@ +from collections.abc import Collection, Iterable +from typing import Final + +from fastapi import status +from pydantic import BaseModel + +import litellm +from litellm.proxy._types import ModelAccessDeniedProxyException, ProxyErrorTypes, UserAPIKeyAuth +from litellm.proxy.auth.model_access_denied import model_access_denied_client_message +from litellm.proxy.common_utils.model_listing_utils import alias_map +from litellm.router_api import RegistryReader + + +def key_disallowed_models(valid_token: UserAPIKeyAuth) -> frozenset[str]: + return frozenset(valid_token.disallowed_models or ()) + + +def _names_model_resolves_through( + model: str, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None +) -> frozenset[str]: + team_aliases: Final = alias_map(valid_token.team_model_aliases) if valid_token.team_model_aliases else {} + key_aliases: Final = alias_map(valid_token.aliases) if valid_token.aliases else {} + after_team_alias: Final = team_aliases.get(model, model) + 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 + ) + return frozenset( + name + for name in (model, after_team_alias, after_key_alias, after_global_alias, dispatched, router_alias_target) + if name is not None + ) + + +def first_disallowed_model( + model: str | Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None +) -> str | None: + disallowed: Final = key_disallowed_models(valid_token) + if not disallowed: + return None + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return next( + (m for m in requested if _names_model_resolves_through(m, valid_token, llm_router) & disallowed), + None, + ) + + +def ensure_key_may_call_model( + model: str | Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None +) -> None: + denied: Final = first_disallowed_model(model, valid_token, llm_router) + if denied is None: + return + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=denied), + internal_message=f"Key is not allowed to call model={denied}: it is listed in the key's disallowed_models.", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + + +def without_key_disallowed_models( + models: Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None +) -> list[str]: + disallowed: Final = key_disallowed_models(valid_token) + if not disallowed: + return list(models) + return [m for m in models if not _names_model_resolves_through(m, valid_token, llm_router) & disallowed] + + +def disallowed_models_removed_by(data: BaseModel | None, stored: Collection[str] | None) -> frozenset[str]: + if data is None or "disallowed_models" not in data.model_fields_set: + return frozenset() + requested: Final[Collection[str]] = getattr(data, "disallowed_models", None) or () + return frozenset(stored or ()) - frozenset(requested) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82135cb1b74..23a038fa039 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -123,6 +123,7 @@ from litellm.proxy.auth.auth_utils import ( route_in_additonal_public_routes, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.auth.key_disallowed_models import ensure_key_may_call_model, key_disallowed_models from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request @@ -4628,49 +4629,50 @@ async def enforce_key_and_fallback_model_access( """ await _normalize_claude_model(request_data, valid_token, request, route) await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) - config: Final = valid_token.config + skips_key_allowlist: Final = valid_token.config != {} or ( + isinstance(valid_token.models, list) and "all-team-models" in valid_token.models + ) + if skips_key_allowlist and not key_disallowed_models(valid_token): + return - if config != {}: - model_list: Final = config.get("model_list", []) - new_model_list: Final = model_list - verbose_proxy_logger.debug("\n new llm router model list %s", new_model_list) - elif isinstance(valid_token.models, list) and "all-team-models" in valid_token.models: - pass - else: - model: Final = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, + model: Final = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + fallback_names: Final = tuple( + name + for target in iter_request_fallback_targets(request_data) + if (name := fallback_target_model_name(target)) is not None + ) + + if skips_key_allowlist: + requested_models: Final = () if model is None else (model,) if isinstance(model, str) else tuple(model) + ensure_key_may_call_model((*requested_models, *fallback_names), valid_token, llm_router) + return + + if model is not None: + await can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=valid_token, llm_router=llm_router, - team_id=valid_token.team_id, ) - if model is not None: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - - fallback_names: Final = tuple( - name - for target in iter_request_fallback_targets(request_data) - if (name := fallback_target_model_name(target)) is not None + for _name in dict.fromkeys(fallback_names): # dedupe, preserve order + await can_key_call_model( + model=_name, + llm_model_list=llm_model_list, + valid_token=valid_token, + llm_router=llm_router, + ) + await is_valid_fallback_model( + model=_name, + llm_router=llm_router, + user_model=None, ) - - for _name in dict.fromkeys(fallback_names): # dedupe, preserve order - await can_key_call_model( - model=_name, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - await is_valid_fallback_model( - model=_name, - llm_router=llm_router, - user_model=None, - ) _enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d28401524ed..a9fdf9220d7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, c import fastapi import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, enforce_output_token_estimates_are_admin_only, ) +from litellm.proxy.auth.key_disallowed_models import disallowed_models_removed_by from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, @@ -3274,6 +3275,37 @@ async def _check_single_key_update_team_permissions( ) +async def check_disallowed_models_caller_permission( + data: BaseModel | None, + user_api_key_dict: UserAPIKeyAuth, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + removed: Final = disallowed_models_removed_by(data, existing_key_row.disallowed_models) + if not removed or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + team_table: Final = ( + await get_team_object( + team_id=existing_key_row.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + if existing_key_row.team_id is not None and prisma_client is not None + else None + ) + if team_table is not None and is_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_table): + return + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Only proxy admins and admins of the key's team can remove models from a key's " + f"`disallowed_models`. Removed: {sorted(removed)}" + }, + ) + + async def _process_single_key_update( update_key_request: UpdateKeyRequest, user_api_key_dict: UserAPIKeyAuth, @@ -3348,6 +3380,9 @@ async def _process_single_key_update( user_api_key_dict, existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) + await check_disallowed_models_caller_permission( + update_key_request, user_api_key_dict, existing_key_row, prisma_client, user_api_key_cache + ) # Custom key update hook if user_custom_key_update is not None: @@ -3894,6 +3929,9 @@ async def _validate_update_key_data( user_api_key_dict, existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) + await check_disallowed_models_caller_permission( + data, user_api_key_dict, existing_key_row, checked_prisma_client, user_api_key_cache + ) # Check team limits if key has a team_id (from request or existing key) team_obj: LiteLLM_TeamTableCachedObj | None = None @@ -5298,6 +5336,7 @@ async def generate_key_helper_fn( rotation_interval: str | None = None, router_settings: dict[str, object] | None = None, access_group_ids: list[str] | None = None, + disallowed_models: Sequence[str] | None = None, budget_limits: list | None = None, # multiple concurrent budget windows *, llm_router: Router | None = None, @@ -5434,6 +5473,7 @@ async def generate_key_helper_fn( "object_permission_id": object_permission_id, "router_settings": router_settings_json, "access_group_ids": access_group_ids or [], + "disallowed_models": list(disallowed_models or ()), } # Add rotation fields if auto_rotate is enabled @@ -6677,6 +6717,9 @@ async def regenerate_key_fn( user_api_key_dict, existing_metadata=_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) + await check_disallowed_models_caller_permission( + data, user_api_key_dict, _key_in_db, prisma_client, user_api_key_cache + ) if data is not None and (data.access_group_ids or data.object_permission is not None): regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f837e75039e..9a43ced03e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -452,6 +452,7 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck, router_fallback_budget_check from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.key_disallowed_models import key_disallowed_models, without_key_disallowed_models from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.login_throttle import ( LoginThrottle, @@ -18311,7 +18312,7 @@ async def _v1_model_info_lists_every_deployment( model_access_groups=model_access_groups, ): return False - if personal_key_team_model_access_applies(user_api_key_dict): + if personal_key_team_model_access_applies(user_api_key_dict) or key_disallowed_models(user_api_key_dict): return False return not user_api_key_dict.agent_id or await resolve_agent_ceiling(user_api_key_dict.agent_id) is None @@ -18915,17 +18916,21 @@ async def model_group_info( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) all_models_str: Final = ( - await run_model_listing( - partial( - get_complete_model_list, - key_models=(), - team_models=(), - proxy_model_list=llm_router.registry().names(), - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=False, - llm_router=llm_router, - ) + without_key_disallowed_models( + await run_model_listing( + partial( + get_complete_model_list, + key_models=(), + team_models=(), + proxy_model_list=llm_router.registry().names(), + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=False, + llm_router=llm_router, + ) + ), + user_api_key_dict, + llm_router, ) if is_proxy_admin else await get_available_models_for_user( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 818f2dc4858..cbf5e883794 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -530,6 +530,7 @@ model LiteLLM_VerificationToken { key_type String? policies String[] @default([]) access_group_ids String[] @default([]) + disallowed_models String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_fallbacks Json @default("{}") @@ -632,6 +633,7 @@ model LiteLLM_DeletedVerificationToken { key_type String? policies String[] @default([]) access_group_ids String[] @default([]) + disallowed_models String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_fallbacks Json @default("{}") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3102674376a..d4f7da3bef8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -10080,6 +10080,7 @@ async def get_available_models_for_user( personal_key_team_model_access_applies, personal_key_team_visible_models, ) + from litellm.proxy.auth.key_disallowed_models import without_key_disallowed_models from litellm.proxy.auth.model_checks import ( get_complete_model_list, get_key_models, @@ -10204,10 +10205,8 @@ async def get_available_models_for_user( team_id=effective_team_id, resolve_agent_ceiling=resolve_agent_ceiling, ) - if agent_visible is None: - return all_models - capped: Final = [m for m in all_models if m in agent_visible] - return capped + visible_models: Final = all_models if agent_visible is None else [m for m in all_models if m in agent_visible] + return without_key_disallowed_models(visible_models, user_api_key_dict, llm_router) def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None: diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index f42a412cf2d..0fdd37a7220 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -38,6 +38,7 @@ class KeyRequestBase(GenerateRequestBase): allowed_routes: list[Any] | None = [] allowed_passthrough_routes: list[object] | None = None denied_passthrough_routes: list[str] | None = None + disallowed_models: list[str] | None = None allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = ( None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm diff --git a/schema.prisma b/schema.prisma index 818f2dc4858..cbf5e883794 100644 --- a/schema.prisma +++ b/schema.prisma @@ -530,6 +530,7 @@ model LiteLLM_VerificationToken { key_type String? policies String[] @default([]) access_group_ids String[] @default([]) + disallowed_models String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_fallbacks Json @default("{}") @@ -632,6 +633,7 @@ model LiteLLM_DeletedVerificationToken { key_type String? policies String[] @default([]) access_group_ids String[] @default([]) + disallowed_models String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_fallbacks Json @default("{}") diff --git a/tests/integration/authorization/test_key_disallowed_models.py b/tests/integration/authorization/test_key_disallowed_models.py new file mode 100644 index 00000000000..9f7f5739b7e --- /dev/null +++ b/tests/integration/authorization/test_key_disallowed_models.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from hashlib import sha256 +from typing import Final + +import httpx +from pydantic import JsonValue + +from tests.integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + list_value, + object_value, + string_value, +) +from tests.integration._support.database import read_rows + + +@dataclass(frozen=True, slots=True) +class _TeamGrant: + team_id: str + allowed: str + disallowed: str + + +def _team_grant(scenario: Scenario) -> _TeamGrant: + group: Final = f"integration-group-{uuid.uuid4().hex[:10]}" + allowed: Final = scenario.model(model_info={"access_groups": [group]}) + disallowed: Final = scenario.model(model_info={"access_groups": [group]}) + return _TeamGrant(team_id=scenario.team(models=[group]), allowed=allowed, disallowed=disallowed) + + +def _observed_models(gateway: Gateway) -> tuple[str, ...]: + response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=10) + response.raise_for_status() + requests: Final = list_value(JSON_OBJECT.validate_json(response.content)["requests"]) + return tuple(str(object_value(object_value(entry).get("body") or {}).get("model")) for entry in requests) + + +def _chat(gateway: Gateway, key: str, body: dict[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"messages": [{"role": "user", "content": f"disallowed {uuid.uuid4().hex}"}], **body}, + key=key, + ) + + +def _assert_served(gateway: Gateway, key: str, model: str) -> None: + response: Final = _chat(gateway, key, {"model": model}) + assert response.status_code == 200, f"{model}: {response.status_code} {response.text}" + + +def _assert_denied(gateway: Gateway, key: str, body: dict[str, JsonValue], denied: str) -> None: + before: Final = _observed_models(gateway) + response: Final = _chat(gateway, key, body) + assert response.status_code == 403, f"{response.status_code} {response.text}" + error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"]) + assert error["type"] == "key_model_access_denied", response.text + assert error["param"] == "model", response.text + assert denied in string_value(error["message"]), response.text + assert _observed_models(gateway) == before + + +def _listed_models(gateway: Gateway, key: str) -> frozenset[str]: + response: Final = gateway.request("GET", "/v1/models", key=key) + assert response.status_code == 200, response.text + entries: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"]) + return frozenset(string_value(object_value(entry)["id"]) for entry in entries) + + +def _stored_disallowed_models(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT disallowed_models FROM "LiteLLM_VerificationToken" WHERE token=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +def test_a_key_inheriting_a_team_access_group_cannot_call_or_list_its_disallowed_models(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + grant: Final = _team_grant(scenario) + key: Final = scenario.key(team_id=grant.team_id, disallowed_models=[grant.disallowed]) + assert _stored_disallowed_models(key) == [{"disallowed_models": [grant.disallowed]}] + _assert_served(gateway, key, grant.allowed) + _assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed) + listed: Final = eventually( + lambda: _listed_models(gateway, key), + lambda models: grant.allowed in models, + seconds=30, + ) + assert grant.allowed in listed + assert grant.disallowed not in listed + + +def test_an_all_team_models_key_cannot_call_its_disallowed_models(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + grant: Final = _team_grant(scenario) + key: Final = scenario.key( + team_id=grant.team_id, models=["all-team-models"], disallowed_models=[grant.disallowed] + ) + _assert_served(gateway, key, grant.allowed) + _assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed) + + +def test_disallowed_models_win_over_the_keys_own_allow_list(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + allowed: Final = scenario.model() + disallowed: Final = scenario.model() + key: Final = scenario.key(models=[allowed, disallowed], disallowed_models=[disallowed]) + _assert_served(gateway, key, allowed) + _assert_denied(gateway, key, {"model": disallowed}, disallowed) + assert disallowed not in _listed_models(gateway, key) + + +def test_an_alias_or_a_fallback_cannot_reach_a_disallowed_model(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + grant: Final = _team_grant(scenario) + alias: Final = f"alias-{uuid.uuid4().hex[:8]}" + key: Final = scenario.key( + team_id=grant.team_id, aliases={alias: grant.disallowed}, disallowed_models=[grant.disallowed] + ) + _assert_denied(gateway, key, {"model": alias}, alias) + _assert_denied(gateway, key, {"model": grant.allowed, "fallbacks": [grant.disallowed]}, grant.disallowed) + + +def test_adding_and_clearing_disallowed_models_takes_effect_on_the_next_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + grant: Final = _team_grant(scenario) + key: Final = scenario.key(team_id=grant.team_id) + _assert_served(gateway, key, grant.disallowed) + gateway.post("/key/update", {"key": key, "disallowed_models": [grant.disallowed]}) + _assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed) + info: Final = object_value(gateway.get("/key/info", {"key": key})["info"]) + assert info["disallowed_models"] == [grant.disallowed], info + gateway.post("/key/update", {"key": key, "disallowed_models": []}) + _assert_served(gateway, key, grant.disallowed) + + +def test_a_key_owner_cannot_remove_disallowed_models_but_can_add_them(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + grant: Final = _team_grant(scenario) + owner: Final = scenario.member(grant.team_id) + key: Final = scenario.key(team_id=grant.team_id, user_id=owner, disallowed_models=[grant.disallowed]) + cleared: Final = gateway.request("POST", "/key/update", {"key": key, "disallowed_models": []}, key=key) + assert cleared.status_code == 403, f"{cleared.status_code} {cleared.text}" + assert _stored_disallowed_models(key) == [{"disallowed_models": [grant.disallowed]}] + _assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed) + widened: Final = gateway.request( + "POST", + "/key/update", + {"key": key, "disallowed_models": [grant.disallowed, grant.allowed]}, + key=key, + ) + assert widened.status_code == 200, f"{widened.status_code} {widened.text}" + _assert_denied(gateway, key, {"model": grant.allowed}, grant.allowed) + + +def test_an_admin_of_the_keys_team_can_remove_disallowed_models(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + grant: Final = _team_grant(scenario) + admin: Final = scenario.member(grant.team_id, role="admin") + admin_key: Final = scenario.key(team_id=grant.team_id, user_id=admin) + key: Final = scenario.key(team_id=grant.team_id, disallowed_models=[grant.disallowed]) + _assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed) + cleared: Final = gateway.request("POST", "/key/update", {"key": key, "disallowed_models": []}, key=admin_key) + assert cleared.status_code == 200, f"{cleared.status_code} {cleared.text}" + assert _stored_disallowed_models(key) == [{"disallowed_models": []}] + _assert_served(gateway, key, grant.disallowed) diff --git a/tests/unit/proxy/auth/test_key_disallowed_models.py b/tests/unit/proxy/auth/test_key_disallowed_models.py new file mode 100644 index 00000000000..04a51be0788 --- /dev/null +++ b/tests/unit/proxy/auth/test_key_disallowed_models.py @@ -0,0 +1,118 @@ +from typing import Final + +import pytest + +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import can_key_call_model +from litellm.proxy.auth.key_disallowed_models import ( + disallowed_models_removed_by, + ensure_key_may_call_model, + without_key_disallowed_models, +) +from litellm.proxy.auth.user_api_key_auth import enforce_key_and_fallback_model_access +from litellm.types.proxy.management_endpoints.key_requests import UpdateKeyRequest + + +def _token(**fields: object) -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="hashed-token", **fields) + + +def _denied_model(error: pytest.ExceptionInfo[ProxyException]) -> str: + assert error.value.type == ProxyErrorTypes.key_model_access_denied + assert str(error.value.code) == "403" + return error.value.message + + +@pytest.mark.parametrize( + "fields,requested", + [ + pytest.param({}, "gpt-denied", id="exact-name"), + pytest.param({"aliases": {"short": "gpt-denied"}}, "short", id="key-alias"), + pytest.param({"team_model_aliases": {"team-short": "gpt-denied"}}, "team-short", id="team-alias"), + pytest.param({}, ["gpt-allowed", "gpt-denied"], id="denied-not-first-in-list"), + ], +) +def test_a_disallowed_model_is_refused_however_it_is_named(fields: dict[str, object], requested: object) -> None: + token: Final = _token(disallowed_models=("gpt-denied",), **fields) + with pytest.raises(ProxyException) as error: + ensure_key_may_call_model(requested, token, None) # pyright: ignore[reportArgumentType] # parametrized str or list + assert "is not available for this API key" in _denied_model(error) + + +@pytest.mark.parametrize("disallowed", [None, (), ("gpt-other",)]) +def test_a_model_outside_disallowed_models_is_not_refused(disallowed: tuple[str, ...] | None) -> None: + ensure_key_may_call_model("gpt-allowed", _token(disallowed_models=disallowed), None) + + +def test_listing_drops_every_disallowed_model_and_every_alias_pointing_at_one() -> None: + token: Final = _token(disallowed_models=("b", "d"), aliases={"alias-b": "b"}) + assert without_key_disallowed_models(["a", "b", "alias-b", "c", "d"], token, None) == ["a", "c"] + + +@pytest.mark.asyncio +async def test_disallowed_models_win_over_the_keys_own_allow_list() -> None: + token: Final = _token(models=["gpt-allowed", "gpt-denied"], disallowed_models=("gpt-denied",)) + assert await can_key_call_model("gpt-allowed", None, token, None) is True + with pytest.raises(ProxyException) as error: + await can_key_call_model("gpt-denied", None, token, None) + _denied_model(error) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fields", + [ + pytest.param({"models": ["all-team-models"], "team_id": "team-1"}, id="all-team-models"), + pytest.param({"config": {"model_list": []}}, id="config-key"), + ], +) +@pytest.mark.parametrize( + "request_data", + [ + pytest.param({"model": "gpt-denied"}, id="model"), + pytest.param({"model": "gpt-allowed", "fallbacks": ["gpt-denied"]}, id="fallback"), + ], +) +async def test_keys_that_skip_the_allow_list_still_honor_disallowed_models( + fields: dict[str, object], request_data: dict[str, object] +) -> None: + token: Final = _token(disallowed_models=("gpt-denied",), **fields) + with pytest.raises(ProxyException) as error: + await enforce_key_and_fallback_model_access( + valid_token=token, + request_data=request_data, + route="/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + _denied_model(error) + + +@pytest.mark.asyncio +async def test_keys_that_skip_the_allow_list_are_served_models_outside_disallowed_models() -> None: + await enforce_key_and_fallback_model_access( + valid_token=_token(models=["all-team-models"], team_id="team-1", disallowed_models=("gpt-denied",)), + request_data={"model": "gpt-allowed"}, + route="/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + +@pytest.mark.parametrize( + "update,stored,removed", + [ + pytest.param({}, ["a", "b"], frozenset(), id="field-not-sent"), + pytest.param({"disallowed_models": []}, ["a", "b"], frozenset({"a", "b"}), id="cleared"), + pytest.param({"disallowed_models": None}, ["a"], frozenset({"a"}), id="null"), + pytest.param({"disallowed_models": ["b"]}, ["a", "b"], frozenset({"a"}), id="narrowed"), + pytest.param({"disallowed_models": ["a", "b", "c"]}, ["a", "b"], frozenset(), id="added"), + pytest.param({"disallowed_models": []}, None, frozenset(), id="nothing-stored"), + ], +) +def test_removed_disallowed_models_are_the_stored_entries_the_update_drops( + update: dict[str, object], stored: list[str] | None, removed: frozenset[str] +) -> None: + assert disallowed_models_removed_by(UpdateKeyRequest(key="sk-1", **update), stored) == removed