diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 49d23d82aa9..ad186e5d2a8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2778,6 +2778,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): default=False, description="When True, a search tool must be explicitly listed in object_permission.search_tools: a virtual key needs its own grant plus its team's, a keyless team member needs the team's, and a user with neither needs their own. A missing permission record, an empty list, or an unresolved team grants nothing, and the unregistered search fallback is denied. The master key and dashboard sessions are exempt", ) + model_access_group_deny_by_default: bool = Field( + default=False, + description="When True, a model tagged with model_info.access_groups is only reachable by a caller whose grant names one of its groups or the model itself. The grant that counts is the team's for a team key, the user's for a personal key, and the key's own for a key with neither. An empty models list, '*', all-proxy-models and wildcard patterns no longer reach these models, by any name, alias, deployment id, routing group or fallback. A model name with both tagged and untagged deployments serves only the untagged ones to such callers. Untagged models are unaffected, and the master key and proxy admins are exempt", + ) missing_session_id: Literal["generate", "reject", "omit"] | None = Field( None, description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 448643ef59b..82505c83a0f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -161,6 +161,7 @@ from litellm.repositories.table_repositories import ( from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.router_utils.model_access_group_grants import explicit_model_grants from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict @@ -5941,6 +5942,140 @@ def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) - ) +_MODEL_NAMES: Final[TypeAdapter[tuple[str, ...]]] = TypeAdapter(tuple[str, ...]) + + +def model_access_group_deny_by_default(general_settings: Mapping[str, object]) -> bool: + """ + Startup rejects a non-boolean value from the config file. A non-boolean value that reaches + general_settings another way enables the policy, so only access-group models are denied. + """ + try: + return ConfigGeneralSettings.model_validate( + MappingProxyType( + { + "model_access_group_deny_by_default": general_settings.get( + "model_access_group_deny_by_default", False + ) + } + ) + ).model_access_group_deny_by_default + except ValidationError: + return True + + +async def model_access_group_grants( + deny_by_default: bool, + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> frozenset[str] | None: + """ + The names that reach a model access group's deployments under general_settings.model_access_group_deny_by_default: + the team's grant for a team key, the owner teams' for a personal key that takes its models from them, the user's + for any other personal key, and the key's own otherwise, plus the key's unified access groups. None when the + setting is off or the caller is the master key, a dashboard session or a proxy admin + """ + if ( + not deny_by_default + or not _is_strict_grant_identity(valid_token) + or valid_token.user_role == LitellmUserRoles.PROXY_ADMIN + ): + return None + owner_models: Final = await _model_access_group_owner_models( + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + key_group_models: Final = await _access_group_model_names( + access_group_ids=valid_token.access_group_ids, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return explicit_model_grants((*owner_models, *key_group_models)) + + +async def _model_access_group_owner_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + if valid_token.team_id is not None: + if team_object is None: + return _MODEL_NAMES.validate_python(valid_token.team_models) + return await _team_model_grants(team_object, prisma_client, user_api_key_cache, proxy_logging_obj) + if personal_key_team_model_access_applies(valid_token): + teams: Final = ( + await _load_personal_key_teams( + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if user_object is not None + else () + ) + team_grants: Final = await asyncio.gather( + *( + _team_model_grants(team, prisma_client, user_api_key_cache, proxy_logging_obj) + for team in teams + if team is not None + ) + ) + return tuple(chain.from_iterable(team_grants)) + if valid_token.user_id is not None and user_object is not None: + from litellm.proxy.auth.resolvers.grants import user_models + + return user_models(user_object) + return tuple(_resolve_key_models_for_auth_check(valid_token=valid_token)) + + +async def _team_model_grants( + team_object: LiteLLM_TeamTable, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + return ( + *_MODEL_NAMES.validate_python(team_object.models), + *await _access_group_model_names( + access_group_ids=team_object.access_group_ids, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + ) + + +async def _access_group_model_names( + access_group_ids: Sequence[str] | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + if not access_group_ids: + return () + return tuple( + await get_models_from_access_groups( + access_group_ids=access_group_ids, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + + def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool: """MCP server, MCP tool and agent inheritance is for a virtual key only. A keyless auth with no team (one grant source of an admitted subject, a JWT user) resolves as the plain caller it is.""" diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b1e8bdbe787..33d0649ae4f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -85,6 +85,8 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo jwt_key_mapping_cache_key, key_access_group_grants_model, key_model_aliases_for_auth_check, + model_access_group_deny_by_default, + model_access_group_grants, resolve_and_validate_end_user_id, resolve_default_end_user_budget, route_skips_budget_checks, @@ -3794,6 +3796,16 @@ async def run_centralized_common_checks( ) bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id) + grants: Final = await model_access_group_grants( + deny_by_default=model_access_group_deny_by_default(_GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)), + valid_token=user_api_key_auth_obj, + team_object=typed_team_object, + user_object=user_object if isinstance(user_object, LiteLLM_UserTable) else None, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + user_api_key_auth_obj.model_access_group_grants = grants # rebind-ok: carries server-only grants to the router try: _ = await common_checks( request=request, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ca0f8892b9d..962120603af 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6959,6 +6959,7 @@ class ProxyConfig: for name in ( "vector_store_deny_by_default", "search_tool_deny_by_default", + "model_access_group_deny_by_default", "vertex_ai_pass_through_drop_anthropic_betas", "request_tag_mappings", ) diff --git a/litellm/router.py b/litellm/router.py index 8d95e3e91d4..20e865cdb16 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -301,6 +301,11 @@ from litellm.router_utils.kubernetes_pod_discovery import ( resolve_pods_after, resolve_pods_after_bound, ) +from litellm.router_utils.model_access_group_grants import ( + deployment_is_granted, + granted_deployments, + model_access_denied, +) from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_MARKER, ContinuationOutput, @@ -14710,6 +14715,27 @@ class Router: specific_deployment: bool | None = False, request_kwargs: Mapping[str, object] | None = None, ) -> tuple[str, list | dict]: + """The candidates `_resolve_available_deployments` picks, less any the caller's model access grants miss""" + resolved_model, deployments = self._resolve_available_deployments( + model=model, specific_deployment=specific_deployment, request_kwargs=request_kwargs + ) + if not isinstance(deployments, Sequence): + if not deployment_is_granted(deployments, model, request_kwargs): + raise model_access_denied(model) + return resolved_model, cast( # cast-ok: contract requires a plain dict for a single deployment + dict[str, object], deployments + ) + granted: Final = granted_deployments(deployments, model, request_kwargs) + if deployments and not granted: + raise model_access_denied(model) + return resolved_model, list(granted) + + def _resolve_available_deployments( + self, + model: str, + specific_deployment: bool | None, + request_kwargs: Mapping[str, object] | None, + ) -> tuple[str, Sequence[DeploymentTypedDict] | Mapping[str, object]]: """ Common checks for 'get_available_deployment' across sync + async call. @@ -15933,6 +15959,9 @@ class Router: model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs ) definition: Final = selected_strategy.deployment + strategy_source: Final = marker if marker is not None else definition.model_dump() if definition else None + if strategy_source is not None and not deployment_is_granted(strategy_source, model, request_kwargs): + raise model_access_denied(model) config_id: Final = definition.model_info.id if definition else None updated_at: Final = definition.model_info.updated_at if definition else None router_kind: Final = classify_strategy_router_model(definition.litellm_params.model) if definition else None diff --git a/litellm/router_utils/model_access_group_grants.py b/litellm/router_utils/model_access_group_grants.py new file mode 100644 index 00000000000..55534a4508a --- /dev/null +++ b/litellm/router_utils/model_access_group_grants.py @@ -0,0 +1,79 @@ +""" +Under general_settings.model_access_group_deny_by_default, auth stamps the caller's explicit model grants on its +auth object, and the router drops every candidate deployment in a model access group those grants do not name. +Filtering the deployments the router is about to pick, rather than predicting them at auth time, covers aliases, +routing groups, deployment ids, default fallbacks and models loaded after auth alike +""" + +from collections.abc import Iterable, Mapping, Sequence +from typing import Final, TypeVar + +import httpx +from pydantic import ConfigDict + +import litellm +from litellm.router_utils.common_utils import request_user_api_key_auth +from litellm.types.llms.base_model import LiteLLMBaseModel +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + +_BLANKET_GRANTS: Final = frozenset({"*", "all-proxy-models", "all-team-models"}) + +_Deployment: Final = TypeVar("_Deployment", bound=Mapping[str, object]) + + +class _DeploymentModelInfo(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + access_groups: tuple[str, ...] | None = None + + +class _DeploymentRow(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + model_name: str | None = None + model_info: _DeploymentModelInfo | None = None + + +def explicit_model_grants(models: Iterable[str]) -> frozenset[str]: + """The names in `models` that grant one thing by name: blanket entries and wildcard patterns name nothing""" + return frozenset(model for model in models if model not in _BLANKET_GRANTS and "*" not in model) + + +def request_model_access_group_grants(request_kwargs: Mapping[str, object] | None) -> frozenset[str] | None: + if request_kwargs is None: + return None + user_api_key_auth: Final = request_user_api_key_auth(request_kwargs) + return user_api_key_auth.model_access_group_grants if isinstance(user_api_key_auth, UserAPIKeyAuth) else None + + +def _is_granted(deployment: Mapping[str, object], requested_model: str, grants: frozenset[str]) -> bool: + row: Final = _DeploymentRow.model_validate(deployment) + groups: Final = (row.model_info.access_groups if row.model_info is not None else None) or () + return not groups or requested_model in grants or row.model_name in grants or not grants.isdisjoint(groups) + + +def model_access_denied(model: str) -> litellm.PermissionDeniedError: + from litellm.proxy.auth.model_access_denied import model_access_denied_client_message + + return litellm.PermissionDeniedError( + message=model_access_denied_client_message(model=model), + llm_provider="", + model=model, + response=httpx.Response(status_code=403, request=httpx.Request(method="POST", url="https://litellm.ai")), + ) + + +def deployment_is_granted( + deployment: Mapping[str, object], requested_model: str, request_kwargs: Mapping[str, object] | None +) -> bool: + grants: Final = request_model_access_group_grants(request_kwargs) + return grants is None or _is_granted(deployment, requested_model, grants) + + +def granted_deployments( + deployments: Sequence[_Deployment], requested_model: str, request_kwargs: Mapping[str, object] | None +) -> tuple[_Deployment, ...]: + grants: Final = request_model_access_group_grants(request_kwargs) + if grants is None: + return tuple(deployments) + return tuple(deployment for deployment in deployments if _is_granted(deployment, requested_model, grants)) diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 2c1c75fbe04..7d5fbfb155b 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -203,6 +203,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) + model_access_group_grants: frozenset[str] | None = Field( + default=None, + exclude=True, + description=( + "Server-only. Set by auth under general_settings.model_access_group_deny_by_default to the model and " + "access group names the caller was explicitly granted; the router then drops tagged deployments these " + "names miss. None means no such restriction" + ), + ) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: object | None = None # Expanded user object when expand=user is used created_by_user: object | None = None # Expanded created_by user when expand=user is used @@ -240,6 +249,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("invoked_agent_policy", None) values.pop("agent_invocation_cost", None) values.pop("billing_agent_policy", None) + values.pop("model_access_group_grants", None) api_key: Final = values.get("api_key") if api_key is not None: values.update({"token": cls._safe_hash_litellm_api_key(api_key)}) # pyright: ignore[reportArgumentType] # non-str fails in the hasher diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..e40e375faa5 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -31,10 +31,13 @@ from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_UserTable, LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, Litellm_EntityType, LitellmUserRoles, ProxyException, + UI_TEAM_ID, ) +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.utils import PrismaClient import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module from litellm.caching.in_memory_cache import InMemoryCache @@ -160,12 +163,15 @@ async def test_team_require_trace_id_accepts_supported_trace_carriers( @pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"]) @pytest.mark.asyncio async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None: - assert await _run_team_trace_id_check( - route=route, - method="POST", - request_body={}, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route=route, + method="POST", + request_body={}, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.asyncio @@ -197,11 +203,14 @@ async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat() async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route( request_body: dict[str, object], ) -> None: - assert await _run_team_trace_id_check( - route="/v1/responses", - request_body=request_body, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route="/v1/responses", + request_body=request_body, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.parametrize( @@ -240,12 +249,15 @@ async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trac @pytest.mark.asyncio async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None: - assert await _run_team_trace_id_check( - route="/mcp/", - request_body={}, - headers={"x-litellm-trace-id": "trace-mcp-123"}, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route="/mcp/", + request_body={}, + headers={"x-litellm-trace-id": "trace-mcp-123"}, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.parametrize( @@ -278,12 +290,15 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_ @pytest.mark.asyncio async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None: - assert await _run_team_trace_id_check( - route="/anthropic/v1/messages", - request_body={}, - headers={"x-litellm-trace-id": "trace-provider-123"}, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route="/anthropic/v1/messages", + request_body={}, + headers={"x-litellm-trace-id": "trace-provider-123"}, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.parametrize( @@ -329,7 +344,9 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pa "jsonrpc": "2.0", "id": "1", "method": "message/send", - "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}}, + "params": { + "message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]} + }, }, ), ], @@ -370,7 +387,9 @@ async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes "jsonrpc": "2.0", "id": "1", "method": "message/send", - "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}}, + "params": { + "message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]} + }, "metadata": {"trace_id": "trace-body-123"}, }, ), @@ -463,7 +482,7 @@ def test_lens_management_does_not_require_an_inference_user_parameter(route: str async def test_get_end_user_object(customer_spend, customer_budget): """ Scenario 1: normal - get_end_user_object returns the cached user - Scenario 2: user over budget - NOTE: budget enforcement now happens in + Scenario 2: user over budget - NOTE: budget enforcement now happens in common_checks() via _check_end_user_budget(), not in get_end_user_object() This test verifies that get_end_user_object correctly retrieves the end user @@ -591,7 +610,7 @@ async def test_can_key_call_model(model, expect_to_work): if expect_to_work: await can_key_call_model(**args) else: - with pytest.raises(Exception, match='is not available for this API key') as e: + with pytest.raises(Exception, match="is not available for this API key") as e: await can_key_call_model(**args) print(e) @@ -863,16 +882,12 @@ async def test_is_valid_fallback_model(): ) try: - await is_valid_fallback_model( - model="gpt-3.5-turbo", llm_router=router, user_model=None - ) + await is_valid_fallback_model(model="gpt-3.5-turbo", llm_router=router, user_model=None) except Exception as e: pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}") with pytest.raises(Exception, match="Invalid") as exc_info: - await is_valid_fallback_model( - model="gpt-4o", llm_router=router, user_model=None - ) + await is_valid_fallback_model(model="gpt-4o", llm_router=router, user_model=None) e = exc_info.value assert "Invalid" in str(e) @@ -886,9 +901,7 @@ async def test_is_valid_fallback_model(): ], ) @pytest.mark.asyncio -async def test_virtual_key_max_budget_check( - token_spend, max_budget, expect_budget_error -): +async def test_virtual_key_max_budget_check(token_spend, max_budget, expect_budget_error): """ Test if virtual key budget checks work as expected: 1. Triggers budget alert for all cases @@ -1003,9 +1016,7 @@ async def test_can_team_access_model(model, team_models, expect_to_work): team_model_aliases=None, ) if not expect_to_work: - pytest.fail( - f"Expected model access check to fail for model={model}, team_models={team_models}" - ) + pytest.fail(f"Expected model access check to fail for model={model}, team_models={team_models}") except Exception as e: if expect_to_work: pytest.fail( @@ -1058,9 +1069,9 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert): await asyncio.sleep(0.1) # Allow time for the alert task to complete - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + ) @pytest.mark.parametrize( @@ -1113,9 +1124,7 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert): ], ) @pytest.mark.asyncio -async def test_team_soft_budget_check( - spend, soft_budget, expect_alert, metadata, expected_alert_emails -): +async def test_team_soft_budget_check(spend, soft_budget, expect_alert, metadata, expected_alert_emails): """ Test cases for _team_soft_budget_check: 1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured) @@ -1163,9 +1172,9 @@ async def test_team_soft_budget_check( await asyncio.sleep(0.1) # Allow time for the alert task to complete - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + ) if expect_alert: assert captured_call_info is not None @@ -1177,10 +1186,7 @@ async def test_team_soft_budget_check( if expected_alert_emails is not None: assert captured_call_info.alert_emails == expected_alert_emails else: - assert ( - captured_call_info.alert_emails is None - or captured_call_info.alert_emails == [] - ) + assert captured_call_info.alert_emails is None or captured_call_info.alert_emails == [] @pytest.mark.asyncio @@ -1311,9 +1317,7 @@ async def test_get_fuzzy_user_object(): # Test 5: Only email provided (no SSO ID) mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user) - result = await _get_fuzzy_user_object( - prisma_client=mock_prisma, user_email="test@example.com" - ) + result = await _get_fuzzy_user_object(prisma_client=mock_prisma, user_email="test@example.com") assert result == test_user mock_prisma.db.litellm_usertable.find_first.assert_called_with( where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}}, @@ -1322,9 +1326,7 @@ async def test_get_fuzzy_user_object(): # Test 6: Only SSO ID provided (no email) mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user) - result = await _get_fuzzy_user_object( - prisma_client=mock_prisma, sso_user_id="sso_123" - ) + result = await _get_fuzzy_user_object(prisma_client=mock_prisma, sso_user_id="sso_123") assert result == test_user mock_prisma.db.litellm_usertable.find_unique.assert_called_with( where={"sso_user_id": "sso_123"}, include={"organization_memberships": True} @@ -1371,7 +1373,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work) llm_router=router, ) else: - with pytest.raises(Exception, match='is not available for this API key') as e: + with pytest.raises(Exception, match="is not available for this API key") as e: await can_key_call_model( model=model, llm_model_list=llm_model_list, @@ -1457,9 +1459,7 @@ async def test_delete_cache_access_object(): ], ) @pytest.mark.asyncio -async def test_get_resources_from_access_groups( - resource_field, access_group_data, expected -): +async def test_get_resources_from_access_groups(resource_field, access_group_data, expected): """Test _get_resources_from_access_groups returns correct resource list from access groups.""" from unittest.mock import AsyncMock, MagicMock, patch @@ -2049,7 +2049,9 @@ def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_ro model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"}, ) expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False} - order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry") + order: Final = ( + ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry") + ) verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order} @@ -2292,10 +2294,7 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) client = _WedgedPublishRedisClient() cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) - hashed_tokens: Final = tuple( - hashlib.sha256(f"token-{index}".encode()).hexdigest() - for index in range(600) - ) + hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600)) for hashed_token in hashed_tokens: cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token}) seeded: Final = tuple( @@ -2355,10 +2354,7 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) client = _WedgedPublishRedisClient() - hashed_tokens: Final = tuple( - hashlib.sha256(f"token-{index}".encode()).hexdigest() - for index in range(600) - ) + hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600)) redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens}) cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer @@ -2380,3 +2376,136 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ assert refilled == [] assert redis.entries == {} + + +_GRANT_ACCESS_GROUPS: Final = {"ag-team": ["team-group-model"], "ag-key": ["key-group-model"]} + + +async def _model_access_group_grants( + token: UserAPIKeyAuth, + *, + team: LiteLLM_TeamTable | None = None, + user: LiteLLM_UserTable | None = None, + deny_by_default: bool = True, + cached_teams: Sequence[LiteLLM_TeamTableCachedObj] = (), +) -> frozenset[str] | None: + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import cache_access_object, model_access_group_grants + + cache: Final = UserApiKeyCache() + for group_id, models in _GRANT_ACCESS_GROUPS.items(): + await cache_access_object( + access_group_id=group_id, + access_group_table=LiteLLM_AccessGroupTable( + access_group_id=group_id, access_group_name=group_id, access_model_names=models + ), + user_api_key_cache=cache, + ) + for cached_team in cached_teams: + await cache.async_set_cache( + key=f"team_id:{cached_team.team_id}", value=cached_team, model_type=LiteLLM_TeamTableCachedObj + ) + return await model_access_group_grants( + deny_by_default=deny_by_default, + valid_token=token, + team_object=team, + user_object=user, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token, deny_by_default", + [ + (UserAPIKeyAuth(token="k1", models=["group-x"]), False), + (UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, models=["group-x"]), True), + (UserAPIKeyAuth(token="k1", team_id=UI_TEAM_ID, team_models=["group-x"]), True), + (UserAPIKeyAuth(token="k1", models=["group-x"], user_role=LitellmUserRoles.PROXY_ADMIN), True), + ], + ids=["setting-off", "master-key", "dashboard-session", "proxy-admin"], +) +async def test_model_access_group_grants_is_unrestricted_when_off_or_for_admin_identities( + token: UserAPIKeyAuth, deny_by_default: bool +) -> None: + assert await _model_access_group_grants(token, deny_by_default=deny_by_default) is None + + +@pytest.mark.asyncio +async def test_team_key_grants_are_the_team_models_and_team_access_group_models_without_blankets() -> None: + token: Final = UserAPIKeyAuth(token="k1", team_id="t1", models=["key-model"]) + team: Final = LiteLLM_TeamTable( + team_id="t1", models=["*", "openai/*", "all-proxy-models", "group-x"], access_group_ids=["ag-team"] + ) + + assert await _model_access_group_grants(token, team=team) == frozenset({"group-x", "team-group-model"}) + + +@pytest.mark.asyncio +async def test_team_key_without_a_loaded_team_falls_back_to_the_token_team_models() -> None: + token: Final = UserAPIKeyAuth(token="k1", team_id="t1", models=["key-model"], team_models=["group-x", "*"]) + + assert await _model_access_group_grants(token) == frozenset({"group-x"}) + + +@pytest.mark.asyncio +async def test_key_access_groups_add_their_models_to_the_owner_grants() -> None: + token: Final = UserAPIKeyAuth(token="k1", team_id="t1", access_group_ids=["ag-key"]) + team: Final = LiteLLM_TeamTable(team_id="t1", models=["group-x"]) + + assert await _model_access_group_grants(token, team=team) == frozenset({"group-x", "key-group-model"}) + + +@pytest.mark.asyncio +async def test_personal_key_grants_are_the_owner_user_models_not_the_key_models() -> None: + token: Final = UserAPIKeyAuth(token="k1", user_id="u1", models=["key-model"]) + user: Final = LiteLLM_UserTable(user_id="u1", models=["user-model", "all-proxy-models"]) + + assert await _model_access_group_grants(token, user=user) == frozenset({"user-model"}) + + +@pytest.mark.asyncio +async def test_personal_key_taking_models_from_teams_gets_the_union_of_its_owner_teams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1", models=["key-model"]) + user: Final = LiteLLM_UserTable(user_id="u1", models=["user-model"], teams=["team-a", "team-b"]) + owner_teams: Final = ( + LiteLLM_TeamTableCachedObj(team_id="team-a", models=["a-model", "*"]), + LiteLLM_TeamTableCachedObj(team_id="team-b", models=["b-model"], access_group_ids=["ag-team"]), + ) + + assert await _model_access_group_grants(token, user=user, cached_teams=owner_teams) == frozenset( + {"a-model", "b-model", "team-group-model"} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_models, expected", + [ + (["key-model", "all-proxy-models", "gpt-*"], frozenset({"key-model"})), + (["all-team-models"], frozenset()), + ], +) +async def test_teamless_key_without_a_loaded_user_is_granted_its_own_models( + key_models: list[str], expected: frozenset[str] +) -> None: + token: Final = UserAPIKeyAuth(token="k1", user_id="u1", models=key_models, team_models=["team-model"]) + + assert await _model_access_group_grants(token) == expected + + +@pytest.mark.parametrize( + "build", + [ + lambda: UserAPIKeyAuth(token="k1", model_access_group_grants=frozenset({"group-x"})), + lambda: UserAPIKeyAuth.model_validate({"token": "k1", "model_access_group_grants": ["group-x"]}), + ], + ids=["constructor", "model_validate"], +) +def test_caller_input_cannot_set_model_access_group_grants(build: Callable[[], UserAPIKeyAuth]) -> None: + assert build().model_access_group_grants is None diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 2d9c781b0d1..20d16fc5c7c 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5202,6 +5202,54 @@ async def test_centralized_common_checks_enforces_team_model_max_budget_from_the assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings, expected_grants", + [ + ({"model_access_group_deny_by_default": True}, frozenset({"group-x"})), + ({"model_access_group_deny_by_default": "true"}, frozenset({"group-x"})), + ({"model_access_group_deny_by_default": "not-a-bool"}, frozenset({"group-x"})), + ({"model_access_group_deny_by_default": False}, None), + ({}, None), + ], + ids=["deny-by-default", "quoted-true", "invalid-value-fails-closed", "setting-false", "setting-absent"], +) +async def test_centralized_common_checks_stamps_the_team_model_access_group_grants_on_the_token( + monkeypatch: pytest.MonkeyPatch, + general_settings: Mapping[str, object], + expected_grants: frozenset[str] | None, +) -> None: + from starlette.datastructures import URL + + from litellm.proxy import proxy_server + + token: Final = UserAPIKeyAuth(token="k1", team_id="t1", models=["key-model"]) + request: Final = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", models=["group-x", "*"]), + model_type=LiteLLM_TeamTableCachedObj, + ) + for name, value in { + **_proxy_attrs_for_centralized_checks(), + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "general_settings": general_settings, + }.items(): + monkeypatch.setattr(proxy_server, name, value) + + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "group-x"}, + route="/chat/completions", + ) + + assert token.model_access_group_grants == expected_grants + + @pytest.mark.asyncio async def test_centralized_common_checks_loads_member_budget_and_exempts_discovery_routes(): from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index e59dd19271c..2a49cd11151 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -2546,7 +2546,9 @@ async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monke @pytest.mark.asyncio -@pytest.mark.parametrize("setting", ["vector_store_deny_by_default", "search_tool_deny_by_default"]) +@pytest.mark.parametrize( + "setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "model_access_group_deny_by_default"] +) @pytest.mark.parametrize(("yaml_value", "expected"), [("true", True), ("false", False)]) async def test_load_config_yaml_deny_by_default_is_boolean( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: str, yaml_value: str, expected: bool @@ -2564,7 +2566,9 @@ async def test_load_config_yaml_deny_by_default_is_boolean( @pytest.mark.asyncio -@pytest.mark.parametrize("setting", ["vector_store_deny_by_default", "search_tool_deny_by_default"]) +@pytest.mark.parametrize( + "setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "model_access_group_deny_by_default"] +) @pytest.mark.parametrize("yaml_value", ["", "enabled"], ids=["null", "string"]) async def test_load_config_rejects_non_boolean_deny_by_default( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: str, yaml_value: str diff --git a/tests/unit/router_utils/test_model_access_group_grants.py b/tests/unit/router_utils/test_model_access_group_grants.py new file mode 100644 index 00000000000..c4a7ac7609f --- /dev/null +++ b/tests/unit/router_utils/test_model_access_group_grants.py @@ -0,0 +1,183 @@ +from typing import Final + +import pytest + +import litellm +from litellm.router_utils.model_access_group_grants import explicit_model_grants +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + +def _deployment(model_name: str, reply: str, model_id: str, access_groups: tuple[str, ...] = ()) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": f"openai/{model_id}-upstream", "api_key": "sk-fake", "mock_response": reply}, + "model_info": {"id": model_id, "access_groups": list(access_groups)}, + } + + +def _router() -> litellm.Router: + return litellm.Router( + model_list=[ + _deployment("private-model", "private", "private-id", ("group-x",)), + _deployment("public-model", "public", "public-id"), + _deployment("shared-model", "shared-x", "shared-x-id", ("group-x",)), + _deployment("shared-model", "shared-y", "shared-y-id", ("group-y",)), + ], + model_group_alias={"alias-to-private": "private-model"}, + routing_groups=[ + { + "group_name": "mixed-group", + "models": ["private-model", "public-model"], + "routing_strategy": "simple-shuffle", + } + ], + ) + + +def _metadata(grants: frozenset[str] | None) -> dict[str, object]: + caller: Final = UserAPIKeyAuth(api_key="sk-caller") + caller.model_access_group_grants = grants + return {"user_api_key_auth": caller} + + +async def _reply(router: litellm.Router, model: str, grants: frozenset[str] | None) -> str | None: + response: Final = await router.acompletion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_metadata(grants) + ) + assert isinstance(response, litellm.ModelResponse) + return response.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming choices + + +@pytest.mark.parametrize( + ("model", "grants", "reply"), + [ + pytest.param("private-model", frozenset({"group-x"}), "private", id="group-grant"), + pytest.param("private-model", frozenset({"private-model"}), "private", id="model-name-grant"), + pytest.param("alias-to-private", frozenset({"private-model"}), "private", id="alias-of-granted-model"), + pytest.param("private-model", None, "private", id="setting-off-or-exempt"), + pytest.param("public-model", frozenset(), "public", id="untagged-model"), + ], +) +async def test_granted_callers_reach_the_deployment(model: str, grants: frozenset[str] | None, reply: str) -> None: + assert await _reply(_router(), model, grants) == reply + + +@pytest.mark.parametrize( + "model", + [ + pytest.param("private-model", id="model-name"), + pytest.param("alias-to-private", id="alias"), + pytest.param("private-id", id="deployment-id"), + pytest.param("openai/private-id-upstream", id="provider-model-name"), + ], +) +async def test_ungranted_callers_get_403_on_every_name_for_a_tagged_deployment(model: str) -> None: + with pytest.raises(litellm.PermissionDeniedError) as denied: + await _reply(_router(), model, explicit_model_grants(["*", "openai/*", "all-proxy-models", "group-y"])) + assert denied.value.status_code == 403 + assert f"The requested model '{model}' is not available for this API key" in denied.value.message + + +@pytest.mark.parametrize( + ("grants", "replies"), + [ + pytest.param(frozenset(), {"public"}, id="ungranted-caller-gets-only-untagged-members"), + pytest.param(frozenset({"mixed-group"}), {"private", "public"}, id="group-named-explicitly-reaches-all"), + ], +) +async def test_a_routing_group_serves_the_members_the_caller_is_granted( + grants: frozenset[str], replies: set[str] +) -> None: + router: Final = _router() + assert {await _reply(router, "mixed-group", grants) for _ in range(40)} == replies + + +async def test_a_model_split_across_groups_serves_the_copy_in_the_callers_group() -> None: + router: Final = _router() + replies: Final = {await _reply(router, "shared-model", frozenset({"group-x"})) for _ in range(20)} + assert replies == {"shared-x"} + + +async def test_a_catch_all_deployment_in_a_group_is_not_granted_by_a_catch_all_grant() -> None: + router: Final = litellm.Router(model_list=[_deployment("*", "catch-all", "catch-all-id", ("group-x",))]) + with pytest.raises(litellm.PermissionDeniedError): + await _reply(router, "*", explicit_model_grants(["*"])) + with pytest.raises(litellm.PermissionDeniedError): + await _reply(router, "openai/gpt-4o", explicit_model_grants(["*"])) + assert await _reply(router, "openai/gpt-4o", explicit_model_grants(["group-x"])) == "catch-all" + + +async def test_the_default_fallback_does_not_reach_a_tagged_deployment() -> None: + router: Final = litellm.Router( + model_list=[_deployment("private-model", "private", "private-id", ("group-x",))], + default_fallbacks=["private-model"], + ) + with pytest.raises(litellm.PermissionDeniedError): + await _reply(router, "unknown-model", frozenset()) + assert await _reply(router, "unknown-model", frozenset({"group-x"})) == "private" + + +@pytest.mark.parametrize( + ("models", "grants"), + [ + pytest.param(["*", "all-proxy-models", "all-team-models"], frozenset(), id="blanket"), + pytest.param(["openai/*", "group-x", "private-model"], frozenset({"group-x", "private-model"}), id="mixed"), + pytest.param([], frozenset(), id="empty"), + ], +) +def test_only_names_that_grant_one_thing_count_as_explicit(models: list[str], grants: frozenset[str]) -> None: + assert explicit_model_grants(models) == grants + + +@pytest.mark.parametrize( + ("model", "resolved_name", "deployment_ids"), + [ + pytest.param("alias-to-private", "private-model", ("private-id",), id="alias"), + pytest.param("mixed-group", "mixed-group", ("private-id", "public-id"), id="routing-group"), + ], +) +def test_resolution_finds_tagged_deployments_for_any_caller( + model: str, resolved_name: str, deployment_ids: tuple[str, ...] +) -> None: + name, deployments = _router()._resolve_available_deployments( + model=model, specific_deployment=False, request_kwargs={"metadata": _metadata(frozenset())} + ) + assert isinstance(deployments, list) + assert (name, tuple(d["model_info"]["id"] for d in deployments)) == (resolved_name, deployment_ids) + + +def test_resolution_finds_a_tagged_deployment_by_id_for_any_caller() -> None: + name, deployment = _router()._resolve_available_deployments( + model="private-id", specific_deployment=False, request_kwargs={"metadata": _metadata(frozenset())} + ) + assert isinstance(deployment, dict) + assert (name, deployment["model_info"]["id"]) == ("openai/private-id-upstream", "private-id") + + +def _router_with_tagged_complexity_router() -> litellm.Router: + return litellm.Router( + model_list=[ + _deployment("public-model", "public", "public-id"), + { + "model_name": "private-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "public-model", "MEDIUM": "public-model", "COMPLEX": "public-model"} + }, + "complexity_router_default_model": "public-model", + }, + "model_info": {"id": "private-router-id", "access_groups": ["group-x"]}, + }, + ], + ) + + +async def test_a_tagged_auto_router_refuses_ungranted_callers_even_when_it_picks_a_public_target() -> None: + with pytest.raises(litellm.PermissionDeniedError) as denied: + await _reply(_router_with_tagged_complexity_router(), "private-router", frozenset({"public-model"})) + assert denied.value.status_code == 403 + + +async def test_a_tagged_auto_router_serves_callers_granted_its_group() -> None: + assert await _reply(_router_with_tagged_complexity_router(), "private-router", frozenset({"group-x"})) == "public" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0fc0d86b479..4411c1f9357 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30225,6 +30225,12 @@ export interface components { * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. */ missing_session_id?: ("generate" | "reject" | "omit") | null; + /** + * Model Access Group Deny By Default + * @description When True, a model tagged with model_info.access_groups is only reachable by a caller whose grant names one of its groups or the model itself. The grant that counts is the team's for a team key, the user's for a personal key, and the key's own for a key with neither. An empty models list, '*', all-proxy-models and wildcard patterns no longer reach these models, by any name, alias, deployment id, routing group or fallback. A model name with both tagged and untagged deployments serves only the untagged ones to such callers. Untagged models are unaffected, and the master key and proxy admins are exempt + * @default false + */ + model_access_group_deny_by_default: boolean; /** * Model List Healthy Only * @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called.