diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b3357579829..27186800e2b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_langfuse_span_scope_value, validate_no_callback_env_reference, ) +from litellm.proxy.auth.access_group_rules import AccessGroupRule, _validate_rules from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled from litellm.types.agents import ( AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export @@ -2507,6 +2508,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + access_group_rules: tuple[AccessGroupRule, ...] | None = None + + @field_validator("access_group_rules", mode="before") + @classmethod + def validate_access_group_rules(cls, value: object) -> tuple[AccessGroupRule, ...] | None: + return _validate_rules(value) if value is not None else None + use_spend_logs_ledger: bool = Field( True, description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.", diff --git a/litellm/proxy/auth/access_group_rules.py b/litellm/proxy/auth/access_group_rules.py new file mode 100644 index 00000000000..b8e3d23f2cc --- /dev/null +++ b/litellm/proxy/auth/access_group_rules.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping, Sequence +from functools import lru_cache, reduce +from importlib import import_module +from typing import TYPE_CHECKING, Final, Literal, Protocol, runtime_checkable + +from pydantic import BaseModel, ConfigDict, TypeAdapter, model_validator + +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.constants import AUTH_OBJECTS_TARGET + +_TRUSTED_ATTRIBUTES: Final = frozenset(("jwt_claims", "user_id", "user_role")) + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.repositories.prisma_protocols import DatabaseClient + + +class AccessGroupRule(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", strict=True) + + attribute: str + contains: str | int | float | bool | None = None + equals: str | int | float | bool | None = None + access_group: str + + @model_validator(mode="after") + def validate_operator(self) -> AccessGroupRule: + if ("contains" in self.model_fields_set) == ("equals" in self.model_fields_set): + raise ValueError("An access group rule requires exactly one of contains or equals") + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + if not self.attribute or any(not part for part in self.attribute.split(".")) or not self.access_group: + raise ValueError("An access group rule requires a nonempty attribute path and access_group") + if self.attribute.split(".")[0] not in _TRUSTED_ATTRIBUTES: + raise ValueError( + "An access group rule attribute must start with an administrator- or IdP-controlled field: " + + ", ".join(sorted(_TRUSTED_ATTRIBUTES)) + ) + if self.attribute.split(".")[0] not in UserAPIKeyAuth.model_fields: + raise ValueError("An access group rule attribute must start with a UserAPIKeyAuth field") + return self + + +class _Enforcer(Protocol): + def enforce(self, subject: str, resource: str) -> bool: ... + + def add_policy(self, subject: str, resource: str) -> bool: ... + + +class _CasbinModel(Protocol): + def load_model_from_text(self, text: str) -> None: ... + + +@runtime_checkable +class _Casbin(Protocol): + Model: Callable[[], _CasbinModel] + Enforcer: Callable[[_CasbinModel], _Enforcer] + + +_MODEL: Final = """ +[request_definition] +r = sub, obj +[policy_definition] +p = sub, obj +[policy_effect] +e = some(where (p.eft == allow)) +[matchers] +m = r.sub == p.sub && r.obj == p.obj +""" +_RULES: Final = TypeAdapter(tuple[AccessGroupRule, ...]) +_ATTRIBUTES: Final = TypeAdapter(dict[str, object]) +_LIST: Final = TypeAdapter(list[object]) +_CASBIN: Final = TypeAdapter(_Casbin, config=ConfigDict(arbitrary_types_allowed=True)) +_MISSING: Final = object() + + +def _subject(operator: Literal["contains", "equals"], claim: str, value: object) -> str: + return json.dumps((operator, claim, value), ensure_ascii=True, separators=(",", ":")) + + +def _attribute_segment(value: object, segment: str) -> object: + if not isinstance(value, Mapping): + return _MISSING + return _ATTRIBUTES.validate_python(value).get(segment, _MISSING) + + +def _claim_value(attributes: Mapping[str, object], path: Sequence[str]) -> object: + return reduce(_attribute_segment, path, attributes) + + +def _subjects(rule: AccessGroupRule, identity: UserAPIKeyAuth) -> tuple[str, ...]: + first, *nested = rule.attribute.split(".") + root: Final[object] = getattr(identity, first, _MISSING) if first in type(identity).model_fields else _MISSING + value: Final = ( + _claim_value(_ATTRIBUTES.validate_python(root), nested) + if nested and isinstance(root, Mapping) + else _MISSING + if nested + else root + ) + if "contains" in rule.model_fields_set: + if not isinstance(value, list): + return () + values: Final = _LIST.validate_python(value) + return tuple(_subject("contains", rule.attribute, item) for item in values if _is_scalar(item)) + if not _is_scalar(value): + return () + return (_subject("equals", rule.attribute, value),) + + +def _is_scalar(value: object) -> bool: + return isinstance(value, (str, int, float, bool)) or value is None + + +@lru_cache(maxsize=1) +def _enforcer(policies: tuple[tuple[str, str], ...]) -> _Enforcer: + casbin: Final = _CASBIN.validate_python(import_module("casbin")) + model: Final = casbin.Model() + model.load_model_from_text(_MODEL) + enforcer: Final = casbin.Enforcer(model) + for subject, resource in policies: + enforcer.add_policy(subject, resource) + return enforcer + + +def _policy(rule: AccessGroupRule) -> tuple[str, str]: + operator: Final = "contains" if "contains" in rule.model_fields_set else "equals" + value: Final = rule.contains if operator == "contains" else rule.equals + return _subject(operator, rule.attribute, value), rule.access_group + + +def _validate_rules(value: object) -> tuple[AccessGroupRule, ...]: + rules: Final = _RULES.validate_python(value) + if rules: + _enforcer(tuple(_policy(rule) for rule in rules)) + return rules + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def _resolve_name(name: str, prisma_client: DatabaseClient | None, cache: UserApiKeyCache) -> str | None: + from litellm.constants import DEFAULT_ACCESS_GROUP_CACHE_TTL + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import get_access_group_tables + from litellm.repositories.table_repositories import AccessGroupRepository + + if prisma_client is None: + verbose_proxy_logger.debug("Access group rule cannot resolve %s without a database", name) + return None + cache_key: Final = f"access_group_rule_name:{name}" + cached: Final = await cache.async_get_cache(key=cache_key, model_type=LiteLLM_AccessGroupTable) + if cached is not None: + groups: Final = await get_access_group_tables( + (cached.access_group_id,), prisma_client=prisma_client, user_api_key_cache=cache + ) + if groups and groups[0].access_group_name == name: + return cached.access_group_id + group: Final = await AccessGroupRepository(prisma_client).find_by_name(name) + if group is None: + verbose_proxy_logger.debug("Access group rule references unknown access group %s", name) + return None + await cache.async_set_cache( + key=cache_key, value=group, model_type=LiteLLM_AccessGroupTable, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL + ) + return group.access_group_id + + +async def access_group_ids_for_identity( + identity: UserAPIKeyAuth, + *, + rules: tuple[AccessGroupRule, ...] | None = None, + prisma_client: DatabaseClient | None = None, + cache: UserApiKeyCache | None = None, +) -> tuple[str, ...]: + if rules is None: + from litellm.proxy import proxy_server + + settings: Final = _ATTRIBUTES.validate_python(proxy_server.general_settings) # pyright: ignore[reportUnknownMemberType] # legacy settings are validated at this boundary + configured: Final = settings.get("access_group_rules") + if not configured: + return () + return await access_group_ids_for_identity( + identity, + rules=_validate_rules(configured), + prisma_client=proxy_server.prisma_client, + cache=proxy_server.user_api_key_cache, + ) + if not rules or cache is None: + return () + enforcer: Final = _enforcer(tuple(_policy(rule) for rule in rules)) + names: Final = frozenset( + rule.access_group + for rule in rules + if any(enforcer.enforce(subject, rule.access_group) for subject in _subjects(rule, identity)) + ) + ids: Final = tuple([await _resolve_name(name, prisma_client, cache) for name in sorted(names)]) + return tuple(group_id for group_id in ids if group_id is not None) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4514882a45f..e546c093473 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5839,7 +5839,7 @@ async def get_authorized_resources_from_key_access_groups( continue team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) - if team_authorized or key_authorized: + if team_authorized or key_authorized or ag_id in valid_token.access_group_rule_grants: authorized_resources.extend(getattr(ag, resource_field, []) or []) return list(set(authorized_resources)) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9e11179c686..4c4e506ce38 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1768,6 +1768,7 @@ class JWTAuthManager: denied_pass_through_team: LiteLLM_TeamTable | None = None route_denied_team_id: str | None = None # rebind-ok: retain route-denied team while checking all claims + fallback_team: tuple[str, LiteLLM_TeamTable] | None = None # rebind-ok: retain first route-allowed candidate if not team_ids: if ( @@ -1796,6 +1797,21 @@ class JWTAuthManager: if team_object and team_object.models is not None: team_models = team_object.models + if ( + fallback_team is None + and allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route=route, + litellm_proxy_roles=jwt_handler.litellm_jwtauth, + ) + and JWTAuthManager._team_has_passthrough_route_access( + team_object=team_object, + route=route, + request_method=request_method, + team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes, + ) + ): + fallback_team = (team_id, team_object) if isinstance(team_models, list) and ( not requested_model or await can_team_access_model( @@ -1828,6 +1844,9 @@ class JWTAuthManager: except Exception: continue + if fallback_team is not None: + return fallback_team + if denied_pass_through_team is not None: JWTAuthManager._raise_team_passthrough_route_denial(route=route, team_object=denied_pass_through_team) @@ -2380,6 +2399,7 @@ class JWTAuthManager: route=route, request_method=request_method, jwt_handler=jwt_handler ) any_team_resolved = False + fallback_team: tuple[str, LiteLLM_TeamTable] | None = None # rebind-ok: retain first route-allowed membership for candidate_team_id in user_team_ids: try: team_object = await get_team_object( @@ -2393,6 +2413,17 @@ class JWTAuthManager: except HTTPException: continue any_team_resolved = True + if ( + team_route_allowed + and fallback_team is None + and JWTAuthManager._team_has_passthrough_route_access( + team_object=team_object, + route=route, + request_method=request_method, + team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes, + ) + ): + fallback_team = (candidate_team_id, team_object) if requested_model: try: await can_team_access_model( @@ -2424,6 +2455,23 @@ class JWTAuthManager: ) return candidate_team_id, team_object, None + if fallback_team is not None: + fallback_id, fallback_object = fallback_team + return ( + fallback_id, + fallback_object, + await get_team_membership( + user_id=user_id, + team_id=fallback_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_id + else None, + ) + if enforce_team_based_model_access: if requested_model and any_team_resolved: raise HTTPException( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0148cafa9e9..9a3b18cf95c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -45,6 +45,7 @@ from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_ from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers +from litellm.proxy.auth.access_group_rules import access_group_ids_for_identity from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports ExperimentalUIJWTToken, TeamNotFoundError, @@ -406,6 +407,17 @@ def _get_model_names_for_budget_checks( return model +async def _normalize_model_for_auth( + request_data: dict[str, object], # mutable-ok: normalization updates dispatch data + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + await _normalize_claude_model(request_data, valid_token, request, route) + await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) + _resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) + + def _budget_fallback_target_model(target: object) -> str | None: if isinstance(target, str): return target @@ -3053,14 +3065,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e param=None, code=status.HTTP_401_UNAUTHORIZED, ) - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) + await _normalize_model_for_auth(request_data, valid_token, request, route) await _prefetch_referenced_auth_objects( valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client ) @@ -3414,6 +3419,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e start_time=start_time, ) virtual_key_auth_obj.via_virtual_key = True + request.state.requires_key_model_check = True # rebind-ok: defer until central grants return virtual_key_auth_obj @@ -4251,7 +4257,33 @@ async def _authorize_authenticated_request( prepare_agent_invocation, ) from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore - from litellm.proxy.proxy_server import general_settings, prisma_client, user_model + from litellm.proxy.proxy_server import ( + general_settings, + llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries + llm_router, + prisma_client, + user_model, + ) + + grants: Final = await access_group_ids_for_identity(user_api_key_auth_obj) + user_api_key_auth_obj.apply_access_group_rule_grants(grants) + key_check_required: Final = ( + getattr(request.state, "requires_key_model_check", False) or force_virtual_key_checks + ) + if key_check_required or grants: + model_identity: Final = ( + user_api_key_auth_obj + if key_check_required + else user_api_key_auth_obj.model_copy(update={"models": list(user_api_key_auth_obj.team_models or ())}) + ) + await enforce_key_and_fallback_model_access( + valid_token=model_identity, + request_data=request_data, + route=route, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) store: Final = AgentIdentityStore.from_client(prisma_client) if prisma_client is not None else None if user_api_key_auth_obj.agent_id is not None: @@ -4368,8 +4400,7 @@ async def _prefetch_referenced_auth_objects( user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient | None, ) -> None: - """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold). - Runs after the key's model access check so a denied request costs no more than it did before.""" + """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold).""" bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None) bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None) await prefetch_auth_objects( @@ -4462,7 +4493,7 @@ async def user_api_key_auth( spend_counter_batch_scope(_spend_counter_store()), ): try: - user_api_key_auth_obj: Final = await user_api_key_auth_builder( + authenticated: Final = await user_api_key_auth_builder( request=request, api_key=api_key, azure_api_key_header=azure_api_key_header, @@ -4479,6 +4510,14 @@ async def user_api_key_auth( if body_parse_exception is not None: raise body_parse_exception raise + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # proxy global is validated below + ) + + settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + user_api_key_auth_obj: Final = ( + authenticated.model_copy() if settings.get("access_group_rules") else authenticated + ) user_api_key_auth_obj.budget_reservation = None user_api_key_auth_obj.agent_caller = agent_caller_from_headers( safe_get_request_headers(request), user_api_key_auth_obj @@ -4817,14 +4856,8 @@ async def _run_post_custom_auth_checks( ) if general_settings.get("custom_auth_run_common_checks", False): - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) + await _normalize_model_for_auth(request_data, valid_token, request, route) + request.state.requires_key_model_check = True # rebind-ok: defer until central grants current_model = _get_model_from_request_context( request_data=request_data, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61f81907ff4..046510d026d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7272,6 +7272,7 @@ class ProxyConfig: "claude_code_marketplace_auth_required", "vertex_ai_pass_through_drop_anthropic_betas", "request_tag_mappings", + "access_group_rules", ) if name in typed_general_settings } diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index aa33da54572..fbf9e5d07f0 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -16,6 +16,7 @@ from pydantic import TypeAdapter from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions @@ -300,6 +301,10 @@ class ModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ModelTab class AccessGroupRepository(PrismaTableRepository["prisma_models.LiteLLM_AccessGroupTable"]): table_name = "litellm_accessgrouptable" + async def find_by_name(self, name: str) -> LiteLLM_AccessGroupTable | None: + row: Final = await self.table.find_unique(where={"access_group_name": name}) + return LiteLLM_AccessGroupTable.model_validate(row, from_attributes=True) if row is not None else None + class SSOConfigRepository(PrismaTableRepository["prisma_models.LiteLLM_SSOConfig"]): table_name = "litellm_ssoconfig" diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index e08459ea14d..8975a4057eb 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -133,6 +133,17 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob """ api_key: str | None = None + _access_group_rule_grants: frozenset[str] = PrivateAttr(default=frozenset()) + + @property + def access_group_rule_grants(self) -> frozenset[str]: + return self._access_group_rule_grants + + def apply_access_group_rule_grants(self, grants: tuple[str, ...]) -> None: + if grants: + self._access_group_rule_grants = frozenset(grants) + self.access_group_ids = list(dict.fromkeys((*(self.access_group_ids or ()), *grants))) + user_role: LitellmUserRoles | None = None allowed_model_region: AllowedModelRegion | None = None parent_otel_span: Span | None = None diff --git a/pyproject.toml b/pyproject.toml index 02834f604d2..de1ea342017 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ Documentation = "https://docs.litellm.ai" [project.optional-dependencies] switchyard = ["nemo-switchyard==0.4.0a1"] proxy = [ + "casbin==1.43.0", "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", "granian>=2.7.4,<3.0; python_version < '3.15'", diff --git a/tests/integration/authorization/test_access_group_rules.py b/tests/integration/authorization/test_access_group_rules.py new file mode 100644 index 00000000000..68decdbf7aa --- /dev/null +++ b/tests/integration/authorization/test_access_group_rules.py @@ -0,0 +1,221 @@ +import json +import time +import uuid +from pathlib import Path +from typing import Final + +import jwt +import pytest +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa + +from litellm.proxy.mcp_server.idp_token_exchange import REJECTED_SUBJECT_TOKEN +from tests.integration._support.client import Gateway, string_value +from tests.integration._support.database import read_rows +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server + +pytestmark: Final = pytest.mark.shares_worker_state + + +@pytest.mark.parametrize( + "enabled,mapping,claimed", ((True, False, True), (False, False, True), (True, True, False), (True, True, True)) +) +def test_rules_add_models_to_jwt_and_scoped_virtual_key_without_opening_other_models( + gateway: Gateway, tmp_path: Path, enabled: bool, mapping: bool, claimed: bool +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_key: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + key_id: Final = "rule-signing-key" + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**public_key, "kid": key_id}]}).encode()) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + explicit: Final = scenario.model() + group_model: Final = scenario.model() + outside: Final = scenario.model() + group_name: Final = f"integration-rule-{uuid.uuid4().hex}" + created: Final = gateway.request( + "POST", "/v1/access_group", {"access_group_name": group_name, "access_model_names": [group_model]} + ) + assert created.status_code == 201, created.text + group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + virtual_user: Final = scenario.user() + virtual_key: Final = scenario.key(models=[explicit], user_id=virtual_user) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "team_ids_jwt_field": "groups", + "user_id_upsert": True, + "enforce_team_based_model_access": True, + **( + {"virtual_key_claim_field": "sub", "unregistered_jwt_client_behavior": "auto_register"} + if mapping + else {} + ), + }, + **( + { + "access_group_rules": [ + {"attribute": "jwt_claims.groups", "contains": "science", "access_group": group_name}, + {"attribute": "user_id", "equals": virtual_user, "access_group": group_name}, + ] + } + if enabled + else {} + ), + } + path: Final = tmp_path / "access_group_rules.yaml" + path.write_text(yaml.safe_dump(config)) + subject: Final = f"integration-rule-user-{uuid.uuid4().hex}" + gateway.post("/user/new", {"user_id": subject, "user_role": "internal_user"}) + scenario.cleanups.callback(scenario.delete_user, subject) + team_id: Final = ( + scenario.team(models=[explicit], members_with_roles=[{"role": "user", "user_id": subject}]) + if claimed + else None + ) + team_claims: Final = [team_id] if team_id is not None else [] + token: Final = jwt.encode( + {"sub": subject, "groups": ["science", *team_claims], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + unmatched: Final = jwt.encode( + {"sub": subject, "groups": ["other", *team_claims], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path, workers=2) as candidate: + fallback_response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": group_model, + "fallbacks": [outside], + "messages": [{"role": "user", "content": "ungranted fallback control"}], + }, + key=token, + ) + assert fallback_response.status_code == 403, fallback_response.text + for key, model, allowed in ( + (token, group_model, enabled and claimed), + *(((token, group_model, enabled and claimed),) if mapping else ()), + (token, outside, False), + (unmatched, group_model, False), + (virtual_key, explicit, True), + (virtual_key, group_model, enabled), + (virtual_key, outside, False), + ): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rule authorization control"}]}, + key=key, + ) + assert response.status_code == (200 if allowed else 403), response.text + if allowed: + assert ( + response.json()["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ), response.text + + if mapping and claimed: + rows: Final = read_rows( + 'SELECT models, access_group_ids FROM "LiteLLM_VerificationToken" ' + 'WHERE user_id=%s AND team_id=%s', + (subject, team_id), + ) + assert len(rows) == 1, rows + assert group_model not in rows[0]["models"], rows + assert not rows[0]["access_group_ids"], rows + + +@pytest.mark.parametrize("consumer", ("auto_register", "token_exchange")) +@pytest.mark.parametrize("matching", (True, False)) +def test_teamless_jwt_cannot_provision_credentials_without_native_team_access( + gateway: Gateway, tmp_path: Path, consumer: str, matching: bool +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_key: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + key_id: Final = "credential-rule-signing-key" + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**public_key, "kid": key_id}]}).encode()) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + model: Final = scenario.model() + subject: Final = f"integration-credential-user-{uuid.uuid4().hex}" + gateway.post("/user/new", {"user_id": subject, "user_role": "internal_user"}) + scenario.cleanups.callback(scenario.delete_user, subject) + group_name: Final = f"credential-rule-{uuid.uuid4().hex}" + created_group: Final = gateway.request( + "POST", "/v1/access_group", {"access_group_name": group_name, "access_model_names": [model]} + ) + assert created_group.status_code == 201, created_group.text + group_id: Final = string_value(created_group.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "enforce_team_based_model_access": True, + **( + {"virtual_key_claim_field": "sub", "unregistered_jwt_client_behavior": "auto_register"} + if consumer == "auto_register" + else {} + ), + }, + "access_group_rules": [ + {"attribute": "jwt_claims.groups", "contains": "science", "access_group": group_name} + ], + } + path: Final = tmp_path / "credential_access_group_rules.yaml" + path.write_text(yaml.safe_dump(config)) + token: Final = jwt.encode( + {"sub": subject, "groups": ["science" if matching else "other"], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path, workers=2) as candidate: + if consumer == "auto_register": + for _ in range(2): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "credential rejection control"}]}, + key=token, + ) + assert response.status_code == 403, response.text + assert "No teams found in token" in response.text, response.text + else: + registered: Final = candidate.client.post( + "/register", json={"redirect_uris": ["http://localhost/callback"], "client_name": "rule-control"} + ) + assert registered.status_code in (200, 201), registered.text + exchanged: Final = candidate.client.post( + "/token", + data={ + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "client_id": registered.json()["client_id"], + "subject_token": token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", + }, + ) + assert exchanged.status_code == 400, exchanged.text + assert exchanged.json() == {"error": "invalid_request", "error_description": REJECTED_SUBJECT_TOKEN}, ( + exchanged.text + ) diff --git a/tests/unit/proxy/auth/test_access_group_rules.py b/tests/unit/proxy/auth/test_access_group_rules.py new file mode 100644 index 00000000000..e73acda7173 --- /dev/null +++ b/tests/unit/proxy/auth/test_access_group_rules.py @@ -0,0 +1,245 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from litellm.proxy._types import ConfigGeneralSettings, LiteLLM_AccessGroupTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.access_group_rules import AccessGroupRule, access_group_ids_for_identity +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +@dataclass(frozen=True) +class _Database: + db: object + + +@pytest.mark.parametrize( + "claims,expected", + ( + ({"groups": ["science"], "tenant": {"id": "production"}}, ("science-id", "standard-id")), + ({"groups": ["other"]}, ()), + ({"groups": "science", "tenant": {"id": ["production"]}}, ()), + ({}, ()), + ({"groups": ["science", "science"]}, ("science-id",)), + ({"groups": [{"name": "science"}]}, ()), + ), +) +@pytest.mark.asyncio +async def test_verified_attributes_grant_union_with_strict_claim_shapes( + claims: Mapping[str, object], expected: tuple[str, ...] +) -> None: + rules: Final = ( + AccessGroupRule(attribute="jwt_claims.groups", contains="science", access_group="science"), + AccessGroupRule(attribute="jwt_claims.tenant.id", equals="production", access_group="standard"), + ) + cache: Final = UserApiKeyCache() + for name in ("science", "standard"): + group: Final = LiteLLM_AccessGroupTable(access_group_id=f"{name}-id", access_group_name=name) + cache.set_cache(f"access_group_rule_name:{name}", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache(f"access_group_id:{name}-id", group, model_type=LiteLLM_AccessGroupTable) + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(jwt_claims=dict(claims)), rules=rules, cache=cache, prisma_client=_Database(object()) + ) + == expected + ) + + +@pytest.mark.parametrize("rule", ({"contains": "science", "equals": "science"}, {})) +def test_invalid_rule_fails_general_settings_validation(rule: Mapping[str, str]) -> None: + with pytest.raises(ValidationError, match="exactly one"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": "jwt_claims.groups", "access_group": "science", **rule}]} + ) + + +@pytest.mark.parametrize("attribute", ("metadata.department", "user_email", "user_email.domain", "team_id", "org_id")) +def test_caller_editable_attributes_are_rejected_as_grant_sources(attribute: str) -> None: + with pytest.raises(ValidationError, match="administrator- or IdP-controlled"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": attribute, "equals": "science", "access_group": "science"}]} + ) + + +@pytest.mark.asyncio +async def test_missing_group_and_no_database_grant_nothing() -> None: + rules: Final = (AccessGroupRule(attribute="user_id", equals="subject", access_group="absent"),) + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value=None) + cache: Final = UserApiKeyCache() + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), + rules=rules, + cache=cache, + prisma_client=_Database(MagicMock(litellm_accessgrouptable=table)), + ) + == () + ) + assert await access_group_ids_for_identity(UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache) == () + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, prisma_client=_Database(object()) + ) == () + + +@pytest.mark.asyncio +async def test_database_failure_does_not_turn_into_an_unrestricted_grant() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + with pytest.raises(RuntimeError, match="database unavailable"): + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), + rules=(AccessGroupRule(attribute="user_id", equals="subject", access_group="science"),), + cache=UserApiKeyCache(), + prisma_client=_Database(MagicMock(litellm_accessgrouptable=table)), + ) + + +@pytest.mark.asyncio +async def test_virtual_key_user_grants_are_additive_and_do_not_mutate_cached_identity() -> None: + identity: Final = UserAPIKeyAuth(models=["explicit"], user_id="science", access_group_ids=["existing"]) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["group-model"] + ) + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + ids: Final = await access_group_ids_for_identity( + identity, + rules=(AccessGroupRule(attribute="user_id", equals="science", access_group="science"),), + cache=cache, + prisma_client=_Database(object()), + ) + granted: Final = identity.model_copy() + granted.apply_access_group_rule_grants(ids) + assert granted.access_group_ids == ["existing", "science-id"] + assert granted.access_group_rule_grants == frozenset({"science-id"}) + assert identity.access_group_ids == ["existing"] + identity.apply_access_group_rule_grants(()) + assert identity.access_group_ids == ["existing"] + + +@pytest.mark.asyncio +async def test_null_equals_does_not_match_an_absent_claim() -> None: + rules: Final = (AccessGroupRule(attribute="jwt_claims.absent", equals=None, access_group="science"),) + assert ( + await access_group_ids_for_identity(UserAPIKeyAuth(jwt_claims={}), rules=rules, cache=UserApiKeyCache()) == () + ) + + +@pytest.mark.parametrize( + "claims,local_email,expected", + ( + ({"email": "science@example.com"}, "other@example.com", ("science-id",)), + ({"email": "other@example.com"}, "science@example.com", ()), + ({}, "science@example.com", ()), + ({"email": ["science@example.com"]}, "science@example.com", ()), + ), +) +@pytest.mark.asyncio +async def test_email_grants_read_verified_claims_not_the_locally_editable_email( + claims: Mapping[str, object], local_email: str, expected: tuple[str, ...] +) -> None: + group: Final = LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="science") + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(jwt_claims=dict(claims), user_email=local_email), + rules=( + AccessGroupRule(attribute="jwt_claims.email", equals="science@example.com", access_group="science"), + ), + cache=cache, + prisma_client=_Database(object()), + ) + == expected + ) + + +@pytest.mark.parametrize("operator", ("contains", "equals")) +@pytest.mark.asyncio +async def test_null_operator_remains_distinct_when_rules_are_reloaded(operator: str) -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + cache: Final = UserApiKeyCache() + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + rule: Final = AccessGroupRule.model_validate( + {"attribute": "jwt_claims.value", operator: None, "access_group": "science"} + ) + principal: Final = UserAPIKeyAuth(jwt_claims={"value": [None] if operator == "contains" else None}) + assert await access_group_ids_for_identity(principal, rules=(rule,), cache=cache, prisma_client=database) == ( + "science-id", + ) + + +@pytest.mark.asyncio +async def test_authenticated_identity_fields_and_claims_are_available_without_metadata_flattening() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + identity: Final = UserAPIKeyAuth( + user_id="verified-user", user_role=LitellmUserRoles.INTERNAL_USER, jwt_claims={"groups": ["science"]} + ) + rules: Final = tuple( + AccessGroupRule.model_validate({"attribute": attribute, operator: value, "access_group": "science"}) + for attribute, operator, value in ( + ("user_id", "equals", "verified-user"), + ("jwt_claims.groups", "contains", "science"), + ("user_role", "equals", "internal_user"), + ) + ) + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + for index, expected in ((0, ("science-id",)), (1, ("science-id",)), (2, ("science-id",))): + assert ( + await access_group_ids_for_identity( + identity, rules=(rules[index],), cache=UserApiKeyCache(), prisma_client=database + ) + == expected + ) + + +@pytest.mark.asyncio +async def test_name_resolution_is_cached_and_renamed_group_does_not_keep_granting_old_name() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + cache: Final = UserApiKeyCache() + rules: Final = (AccessGroupRule(attribute="user_id", equals="subject", access_group="science"),) + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) == ("science-id",) + cache.set_cache( + "access_group_id:science-id", + LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="science"), + model_type=LiteLLM_AccessGroupTable, + ) + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) == ("science-id",) + assert table.find_unique.await_count == 1 + cache.set_cache( + "access_group_id:science-id", + LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="renamed"), + model_type=LiteLLM_AccessGroupTable, + ) + table.find_unique.return_value = None + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) + == () + ) + + +@pytest.mark.parametrize( + "attribute", + ("groups", "department", "metadata.department", "models", "access_group_rule_grants", "jwt_claims..groups"), +) +def test_attribute_requires_a_real_principal_field(attribute: str) -> None: + with pytest.raises(ValidationError, match="attribute"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": attribute, "contains": "science", "access_group": "science"}]} + ) diff --git a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py index 0c3487255f6..e371139dced 100644 --- a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py @@ -1,5 +1,7 @@ import pytest +from typing import Final from unittest.mock import AsyncMock, patch +from starlette.requests import Request import litellm from litellm.proxy.auth.user_api_key_auth import ( _run_post_custom_auth_checks, @@ -19,6 +21,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): # custom_auth_run_common_checks for custom-auth deployments. The # helper itself no longer calls common_checks. valid_token = UserAPIKeyAuth(token="test_token") + http_request: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions"}) # Default: common_checks should NOT be called inside the helper with patch( @@ -27,7 +30,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): mock_common.return_value = True result = await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={}, route="/v1/chat/completions", parent_otel_span=None, @@ -50,7 +53,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): mock_common.return_value = True result = await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={}, route="/v1/chat/completions", parent_otel_span=None, @@ -137,6 +140,7 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( # With custom_auth_run_common_checks set, the wrapper's common_checks() # enforces the end-user budget, so the helper must not double-enforce it. valid_token = UserAPIKeyAuth(token="test_token", end_user_id="customer-1") + http_request: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions"}) end_user_obj = LiteLLM_EndUserTable( user_id="customer-1", blocked=False, @@ -165,7 +169,7 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( ): await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={"model": "gpt-4"}, route="/v1/chat/completions", parent_otel_span=None, diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index ae6bf95c13f..36ba8a734b8 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -1,4 +1,5 @@ import asyncio +import json import re import sys import time @@ -8,26 +9,32 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import jwt import pytest +from cryptography.hazmat.primitives.asymmetric import rsa from fastapi import HTTPException +from starlette.requests import Request +import litellm.proxy.proxy_server as proxy_server from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, - JWTLiteLLMRoleMap, JWTIssuerConfig, - LiteLLMRoutes, + JWTLiteLLMRoleMap, + LiteLLM_AccessGroupTable, LiteLLM_JWTAuth, LiteLLM_ModelTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, + LiteLLMRoutes, LitellmUserRoles, Member, ProxyErrorTypes, ProxyException, RoleBasedPermissions, ScopeMapping, + UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.auth.auth_checks import TeamNotFoundError @@ -44,10 +51,315 @@ from litellm.proxy.auth.handle_jwt import ( NoMatchingJWTPublicKeyError, ) from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException +from litellm.proxy.auth.user_api_key_auth import ( + user_api_key_auth, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.types.agents import AgentResponse +async def _jwt_team_selection_identity( + monkeypatch: pytest.MonkeyPatch, + teams: tuple[LiteLLM_TeamTable, ...], + claimed_teams: tuple[str, ...], + membership_teams: tuple[str, ...], + requested_model: str | None, + *, + authorize: bool = False, + fallback_to_db_teams: bool | None = None, + matching_rule: bool | None = None, + scope_models: tuple[str, ...] | None = None, + role_models: tuple[str, ...] | None = None, +) -> UserAPIKeyAuth: + cache: Final = UserApiKeyCache() + database: Final = MagicMock() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + team_ids_jwt_field="groups", + enforce_team_based_model_access=True, + enforce_rbac=role_models is not None, + fallback_to_db_teams=not claimed_teams if fallback_to_db_teams is None else fallback_to_db_teams, + enforce_scope_based_access=scope_models is not None, + scope_mappings=[ScopeMapping(scope="inference", models=list(scope_models))] + if scope_models is not None + else None, + ), + ) + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_url: Final = "https://identity.invalid/selection/jwks" + cache.set_cache( + f"litellm_jwt_auth_keys_{jwks_url}", + [json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()))], + ) + cache.set_cache( + "selection-user", + LiteLLM_UserTable(user_id="selection-user", user_role="internal_user", teams=list(membership_teams)), + model_type=LiteLLM_UserTable, + ) + for team in teams: + cache.set_cache(f"team_id:{team.team_id}", team, model_type=LiteLLM_TeamTable) + cache.set_cache( + f"team_membership:selection-user:{team.team_id}", + LiteLLM_TeamMembership(user_id="selection-user", team_id=team.team_id), + model_type=LiteLLM_TeamMembership, + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "master_key", "selection-master") + monkeypatch.setattr(proxy_server, "user_custom_auth", None) + monkeypatch.setattr( + proxy_server, + "general_settings", + { + "enable_jwt_auth": True, + **({"role_permissions": [RoleBasedPermissions( + role=LitellmUserRoles.INTERNAL_USER, models=list(role_models), routes=["/chat/completions", "/models"] + )]} if role_models is not None else {}), + **( + { + "access_group_rules": [ + { + "attribute": "jwt_claims.sub", + "equals": "selection-user" if matching_rule else "another-user", + "access_group": "selection-group", + } + ] + } + if matching_rule is not None + else {} + ), + }, + ) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="selection-group-id", + access_group_name="selection-group", + access_model_names=["selection-model"], + ) + cache.set_cache("access_group_rule_name:selection-group", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:selection-group-id", group, model_type=LiteLLM_AccessGroupTable) + token: Final = jwt.encode( + {"sub": "selection-user", "groups": list(claimed_teams), "scope": "inference"}, private_key, algorithm="RS256" + ) + request_data: Final = {"model": requested_model} if requested_model is not None else {} + route: Final = "/chat/completions" if requested_model is not None else "/models" + try: + if authorize: + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": route, + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + } + ) + request._body = json.dumps(request_data).encode() + return await user_api_key_auth(request=request, api_key=f"Bearer {token}") + result: Final = await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=handler, + request_data=request_data, + general_settings=proxy_server.general_settings, + route=route, + prisma_client=database, + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + identity: Final = JWTAuthManager.user_api_key_auth_from_result(result) + return identity + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_unresolved_claim_keeps_model_denial_after_single_membership_is_selected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + membership_team: Final = LiteLLM_TeamTable( + team_id="membership-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (membership_team,), + ("unresolved-claim",), + ("membership-team",), + "selection-model", + authorize=True, + ) + assert int(denied.value.code) == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("claimed_teams", (("team-c", "team-a", "team-b"), ("team-b", "team-c", "team-a"))) +@pytest.mark.parametrize("requested_model,expected_team", (("selection-model", "team-b"), (None, "team-a"))) +async def test_auth_builder_claimed_teams_attribute_to_first_sorted_eligible_team( + monkeypatch: pytest.MonkeyPatch, + claimed_teams: tuple[str, ...], + requested_model: str | None, + expected_team: str, +) -> None: + teams: Final = tuple( + LiteLLM_TeamTable( + team_id=team_id, + models=["other-model"] if team_id == "team-a" else ["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + for team_id in ("team-a", "team-b", "team-c") + ) + identity: Final = await _jwt_team_selection_identity( + monkeypatch, teams, claimed_teams, ("team-c", "team-a", "team-b"), requested_model + ) + assert identity.team_id == expected_team + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("membership_teams", (("team-a", "team-c", "team-b"), ("team-a", "team-b", "team-c"))) +async def test_auth_builder_db_teams_attribute_to_first_model_allowed_membership( + monkeypatch: pytest.MonkeyPatch, membership_teams: tuple[str, ...] +) -> None: + teams: Final = tuple( + LiteLLM_TeamTable( + team_id=team_id, + models=["other-model"] if team_id == "team-a" else ["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + for team_id in membership_teams + ) + identity: Final = await _jwt_team_selection_identity(monkeypatch, teams, (), membership_teams, "selection-model") + assert identity.team_id == membership_teams[1] + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("matching", (True, False)) +async def test_jwt_rule_grants_a_claimed_teams_missing_model_without_losing_attribution( + monkeypatch: pytest.MonkeyPatch, + matching: bool, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="claimed-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + if matching: + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("claimed-team",), + ("claimed-team",), + "selection-model", + authorize=True, + matching_rule=matching, + ) + assert identity.team_id == "claimed-team" + assert identity.user_id == "selection-user" + else: + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("claimed-team",), + ("claimed-team",), + "selection-model", + authorize=True, + matching_rule=matching, + ) + assert denied.value.code == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested_model", (None, "selection-model")) +async def test_jwt_native_team_access_with_rules_configured( + monkeypatch: pytest.MonkeyPatch, + requested_model: str | None, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="native-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("native-team",), + ("native-team",), + requested_model, + authorize=True, + matching_rule=True, + ) + assert identity.team_id == "native-team" + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("granted", (True, False)) +@pytest.mark.parametrize("policy_kind", ("scope", "role")) +async def test_jwt_role_and_scope_model_restrictions_remain_in_builder( + monkeypatch: pytest.MonkeyPatch, + granted: bool, + policy_kind: str, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="scope-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + scope_models: Final = ("selection-model",) if granted else () + if granted: + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("scope-team",), + ("scope-team",), + "selection-model", + authorize=True, + scope_models=scope_models if policy_kind == "scope" else None, + role_models=scope_models if policy_kind == "role" else None, + ) + assert identity.team_id == "scope-team" + else: + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("scope-team",), + ("scope-team",), + "selection-model", + authorize=True, + scope_models=scope_models if policy_kind == "scope" else None, + role_models=scope_models if policy_kind == "role" else None, + ) + assert denied.value.code == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("matching", (True, False)) +async def test_teamless_jwt_is_rejected_before_rule_grants( + monkeypatch: pytest.MonkeyPatch, matching: bool, +) -> None: + with pytest.raises(ProxyException, match="No teams found in token") as denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", authorize=True, + fallback_to_db_teams=False, matching_rule=matching, + ) + assert denied.value.code == "403" + + @pytest.mark.asyncio async def test_map_user_to_teams_user_already_in_team(): """Test that no action is taken when user is already in team""" @@ -213,7 +525,7 @@ async def test_find_team_with_model_access_reports_passthrough_allowlist_denial( assert exc_info.value.status_code == 403 assert "allowed_passthrough_routes" in exc_info.value.detail assert "requested model" not in exc_info.value.detail - mock_is_auth_enforced_pass_through_route.assert_called_once_with( + mock_is_auth_enforced_pass_through_route.assert_called_with( route="/my-pass-through", method="POST" ) @@ -3846,57 +4158,22 @@ async def test_find_team_with_model_access_enforce_team_based_access_still_raise @pytest.mark.asyncio async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403( - monkeypatch, -): - """Regression guard: when the JWT group claim DOES resolve to a real - LiteLLM team but that team does not grant the requested model, keep the - original 403. Only the unresolved-claim case is softened.""" - import sys - import types - - from fastapi import HTTPException - - from litellm.router import Router - - router = Router( - model_list=[ - {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - }, - ] + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="real-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], ) - proxy_server_module = types.ModuleType("proxy_server") - proxy_server_module.llm_router = router - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) - - team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) - - async def mock_get_team_object(*_args, **_kwargs): - return team - - monkeypatch.setattr( - "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + identity: Final = await _jwt_team_selection_identity( + monkeypatch, (team,), ("real-team",), ("real-team",), "selection-model" ) - - jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() - - with pytest.raises(HTTPException) as exc_info: - await JWTAuthManager.find_team_with_model_access( - team_ids={"real-team"}, - requested_model="gpt-4o-mini", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=None, - parent_otel_span=None, - proxy_logging_obj=None, + assert identity.team_id == "real-team" + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, (team,), ("real-team",), ("real-team",), "selection-model", authorize=True ) - - assert exc_info.value.status_code == 403 - assert "No team has access to the requested model" in str(exc_info.value.detail) + assert denied.value.code == "403" @pytest.mark.asyncio @@ -5441,44 +5718,22 @@ async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fal @pytest.mark.asyncio -async def test_find_team_with_model_access_defers_no_team_403_under_db_fallback(): - """find_team_with_model_access raises the early "no teams in token" 403 when - enforcement is on, but defers (returns no team) so auth_builder's DB fallback - can run when fallback_to_db_teams is enabled.""" - jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - enforce_team_based_model_access=True, - fallback_to_db_teams=False, - ) - with pytest.raises(HTTPException) as exc_info: - await JWTAuthManager.find_team_with_model_access( - team_ids=set(), - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), +@pytest.mark.parametrize("fallback", (False, True)) +async def test_auth_builder_preserves_teamless_enforcement_denial( + monkeypatch: pytest.MonkeyPatch, + fallback: bool, +) -> None: + with pytest.raises(HTTPException) as builder_denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", fallback_to_db_teams=fallback ) - assert exc_info.value.status_code == 403 - - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - enforce_team_based_model_access=True, - fallback_to_db_teams=True, - ) - team_id, team_object = await JWTAuthManager.find_team_with_model_access( - team_ids=set(), - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - assert team_id is None - assert team_object is None + assert builder_denied.value.status_code == 403 + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", authorize=True, fallback_to_db_teams=fallback + ) + assert denied.value.code == "403" + assert ("not a member of any team" if fallback else "No teams found in token") in denied.value.message def _db_fallback_handler(litellm_jwtauth: LiteLLM_JWTAuth | None = None) -> JWTHandler: @@ -6265,82 +6520,23 @@ async def test_auth_builder_db_fallback_does_not_validate_rbac_team_against_db_m @pytest.mark.asyncio -async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied(): - """When enforce_team_based_model_access is on, a user with no DB memberships - and a user with memberships that all fail the model-access check must surface - different 403s; collapsing both into the no-membership message hides the real - cause and diverges from find_team_with_model_access's claim-based message.""" - membership_user = LiteLLM_UserTable( - user_id="u_no_model", - user_role=LitellmUserRoles.INTERNAL_USER, - teams=["only_team"], +async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied( + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="only-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], ) - no_membership_user = LiteLLM_UserTable( - user_id="u_empty", - user_role=LitellmUserRoles.INTERNAL_USER, - teams=[], - ) - - async def fake_get_team(team_id, **kwargs): - return LiteLLM_TeamTable(team_id=team_id, models=["other"]) - - async def fake_can_access(model, team_object, llm_router, team_model_aliases=None): - raise ProxyException( - message="team not allowed to access model", - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=403, - ) - - with ( - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", - new_callable=AsyncMock, - side_effect=fake_get_team, - ), - patch( - "litellm.proxy.auth.handle_jwt.can_team_access_model", - new_callable=AsyncMock, - side_effect=fake_can_access, - ), - ): - with pytest.raises(HTTPException) as model_denied: - await JWTAuthManager._resolve_db_team_fallback( - user_object=membership_user, - user_id=None, - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=_db_fallback_handler(), - enforce_team_based_model_access=True, - team_id_upsert=False, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - - with pytest.raises(HTTPException) as no_member: - await JWTAuthManager._resolve_db_team_fallback( - user_object=no_membership_user, - user_id=None, - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=_db_fallback_handler(), - enforce_team_based_model_access=True, - team_id_upsert=False, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - - assert model_denied.value.status_code == 403 - assert "requested model" in model_denied.value.detail - assert "gpt-4" in model_denied.value.detail - assert "only_team" not in model_denied.value.detail - - assert no_member.value.status_code == 403 - assert "not a member of any team" in no_member.value.detail + identity: Final = await _jwt_team_selection_identity(monkeypatch, (team,), (), ("only-team",), "selection-model") + assert identity.team_id == "only-team" + with pytest.raises(ProxyException) as model_denied: + await _jwt_team_selection_identity(monkeypatch, (team,), (), ("only-team",), "selection-model", authorize=True) + assert model_denied.value.code == "403" + assert "not a member of any team" not in model_denied.value.message + with pytest.raises(ProxyException, match="not a member of any team") as no_member: + await _jwt_team_selection_identity(monkeypatch, (), (), (), "selection-model", authorize=True) + assert no_member.value.code == "403" @pytest.mark.asyncio diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 19f9d0dcc73..ce4a844550b 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2,27 +2,29 @@ ## Unit tests for user_api_key_auth helper functions -import litellm.proxy -import litellm.proxy.proxy_server - +import json +import logging from typing import Dict, Final, List, Mapping, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException, WebSocket, status from starlette.datastructures import URL +from starlette.requests import Request as StarletteRequest from starlette.types import Message -from litellm._logging import verbose_proxy_logger -import json -import logging -import litellm -from litellm.proxy.auth.user_api_key_auth import ( - user_api_key_auth, - UserAPIKeyAuth, - get_api_key_from_custom_header, -) -from fastapi import WebSocket, HTTPException, status -from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles +import litellm +import litellm.proxy +import litellm.proxy.proxy_server +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_UserTable, LitellmUserRoles, ProxyException +from litellm.proxy.auth.user_api_key_auth import ( + UserAPIKeyAuth, + _authorize_authenticated_request, + get_api_key_from_custom_header, + user_api_key_auth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from tests._master_key import MASTER_KEY @@ -33,6 +35,146 @@ class Request: self.headers: Dict[str, str] = {} +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,fallbacks,allowed", (("model-a", [], True), ("model-b", [], False), ("model-a", ["model-b"], False)) +) +async def test_virtual_key_model_policy_on_first_lookup_and_cache_hit( + monkeypatch: pytest.MonkeyPatch, model: str, fallbacks: list[str], allowed: bool +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + key: Final = "sk-central-model-check" + stored: Final = UserAPIKeyAuth(token=server.hash_token(key), models=["model-a"]) + database: Final = MagicMock(get_data=AsyncMock(side_effect=(stored, AssertionError("unexpected cache miss")))) + monkeypatch.setattr(server, "master_key", MASTER_KEY) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", database) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + for _ in range(2): + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=AsyncMock( + return_value={ + "type": "http.request", + "body": json.dumps({"model": model, "fallbacks": fallbacks}).encode(), + "more_body": False, + } + ), + ) + if allowed: + identity: Final = await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert identity.models == ["model-a"] + else: + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert denied.value.code == "403" + assert "model-b" in denied.value.message + + +@pytest.mark.asyncio +async def test_scim_deactivation_precedes_central_key_model_denial(monkeypatch: pytest.MonkeyPatch) -> None: + server: Final = litellm.proxy.proxy_server + key: Final = "sk-deactivated-owner" + cache: Final = UserApiKeyCache() + stored: Final = UserAPIKeyAuth(token=server.hash_token(key), models=["model-a"], user_id="inactive-owner") + cache.set_cache( + "inactive-owner", + LiteLLM_UserTable(user_id="inactive-owner", metadata={"scim_active": False}), + model_type=LiteLLM_UserTable, + ) + monkeypatch.setattr(server, "master_key", MASTER_KEY) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=stored))) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=AsyncMock(return_value={"type": "http.request", "body": b'{"model":"model-b"}'}), + ) + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert denied.value.code == "401" + assert "deactivated via SCIM" in denied.value.message + + +@pytest.mark.asyncio +async def test_restricted_key_model_list_denial_names_the_denied_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["model-a"] + ) + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock()) + monkeypatch.setattr( + server, + "general_settings", + {"access_group_rules": [{"attribute": "user_id", "equals": "science", "access_group": "science"}]}, + ) + monkeypatch.setattr(server, "llm_router", None) + identity: Final = UserAPIKeyAuth(user_id="science", models=["model-a"]) + request: Final = StarletteRequest({"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) + with pytest.raises(ProxyException) as denied: + await _authorize_authenticated_request( + identity, request, {"model": ["model-a", "model-b"]}, "/chat/completions", "synthetic", + force_virtual_key_checks=True, + ) + assert denied.value.code == "403" + assert "model-b" in denied.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,enabled,allowed", + (("group-model", True, True), ("outside", True, False), ("group-model", False, False), ("explicit", False, True)), +) +async def test_user_access_group_rules_are_enforced_at_central_authorization( + monkeypatch: pytest.MonkeyPatch, model: str, enabled: bool, allowed: bool +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["group-model"] + ) + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("science", LiteLLM_UserTable(user_id="science"), model_type=LiteLLM_UserTable) + monkeypatch.setattr( + server, + "general_settings", + {"access_group_rules": [{"attribute": "user_id", "equals": "science", "access_group": "science"}]} + if enabled + else {}, + ) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock()) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + principal: Final = UserAPIKeyAuth(models=["explicit"], user_id="science") + request: Final = StarletteRequest({"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) + if allowed: + assert ( + await _authorize_authenticated_request( + principal, request, {"model": model}, "/chat/completions", "synthetic", force_virtual_key_checks=True + ) + is None + ) + assert principal.access_group_ids == (["science-id"] if enabled else None) + else: + with pytest.raises(ProxyException, match="model"): + await _authorize_authenticated_request( + principal, request, {"model": model}, "/chat/completions", "synthetic", force_virtual_key_checks=True + ) + + @pytest.mark.parametrize( "allowed_ips, client_ip, expected_result", [ 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 70208bec0e5..1489afc02af 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 @@ -87,6 +87,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _reserve_budget_after_common_checks, _route_requires_auth_despite_public, _routing_selector_matches_claim, + _authorize_authenticated_request, run_centralized_common_checks, _run_post_custom_auth_checks, user_api_key_auth_builder, @@ -446,29 +447,47 @@ async def test_should_not_reuse_cached_key_object_for_request_state(): assert second_request_key.request_route is None -@pytest.mark.asyncio -async def test_custom_auth_does_not_enforce_key_model_access_by_default(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch( - "litellm.proxy.proxy_server.general_settings", - {}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, +async def _custom_auth_model_request( + monkeypatch: pytest.MonkeyPatch, requested_model: str, *, opt_in: bool, from_header: bool = False +) -> UserAPIKeyAuth: + async def custom_auth(request: Request, api_key: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + models=["allowed-model"], + user_role=LitellmUserRoles.INTERNAL_USER, + request_route=request.url.path, + api_key=api_key, ) - mock_can_key.assert_not_awaited() + + server: Final = litellm.proxy.proxy_server + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + monkeypatch.setattr(server, "user_custom_auth", custom_auth) + monkeypatch.setattr(server, "general_settings", {"custom_auth_run_common_checks": opt_in}) + monkeypatch.setattr(server, "prisma_client", None) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "llm_model_list", None) + monkeypatch.setattr(server, "user_api_key_cache", UserApiKeyCache()) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/files" if from_header else "/v1/chat/completions", + "headers": [(b"x-litellm-model", requested_model.encode())] if from_header else [], + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps({} if from_header else {"model": requested_model}).encode(), + } + ), + ) + return await user_api_key_auth(request=request, api_key="Bearer external-token") + + +@pytest.mark.asyncio +async def test_custom_auth_does_not_enforce_key_model_access_by_default(monkeypatch: pytest.MonkeyPatch): + identity: Final = await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=False) + assert identity.models == ["allowed-model"] @pytest.mark.asyncio @@ -492,99 +511,30 @@ async def test_post_custom_auth_expired_key_returns_unauthorized(): @pytest.mark.asyncio -async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o-mini"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="gpt-4o-mini", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) +async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + identity: Final = await _custom_auth_model_request(monkeypatch, "allowed-model", opt_in=True) + assert identity.models == ["allowed-model"] @pytest.mark.asyncio -async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["allowed-model"]) - request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"}) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=request, - request_data={}, - route="/v1/files", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="restricted-model", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) +async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(ProxyException) as denied: + await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=True, from_header=True) + assert denied.value.code == "403" @pytest.mark.asyncio -async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - mock_can_key.side_effect = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) - with pytest.raises(ProxyException) as exc: - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, - ) - - assert exc.value.type == ProxyErrorTypes.key_model_access_denied +async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(ProxyException) as denied: + await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=True) + assert denied.value.type == ProxyErrorTypes.key_model_access_denied + assert denied.value.code == "403" def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): @@ -1830,8 +1780,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): @pytest.mark.asyncio @pytest.mark.parametrize("model_allowed", [True, False]) -async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_the_model(model_allowed): - """A request denied by the key's model list must not pay for the team/user/org MGET or DB join.""" +async def test_auth_prefetches_referenced_objects_before_central_model_authorization(model_allowed): from fastapi import Request from starlette.datastructures import URL @@ -1841,9 +1790,9 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th api_key = "sk-prefetch-order-test" valid_token = UserAPIKeyAuth(api_key=api_key, token=hash_token(api_key), user_id="u1", team_id="t1") - mock_cache = AsyncMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.delete_cache = MagicMock() + mock_cache = UserApiKeyCache() + mock_cache.set_cache("team_id:t1", LiteLLM_TeamTable(team_id="t1")) + mock_cache.set_cache("team_membership:u1:t1", LiteLLM_TeamMembership(user_id="u1", team_id="t1")) mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) @@ -1873,7 +1822,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) - request = Request(scope={"type": "http"}) + request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions"}) request._url = URL(url="/chat/completions") with ( @@ -1894,7 +1843,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "litellm.proxy.auth.user_api_key_auth.get_user_object", new_callable=AsyncMock, return_value=None ), ): - call = user_api_key_auth_builder( + resolved: Final = await user_api_key_auth_builder( request=request, api_key=f"Bearer {api_key}", azure_api_key_header="", @@ -1903,15 +1852,18 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th azure_apim_header=None, request_data={"model": "gpt-4o"}, ) + assert isinstance(resolved, UserAPIKeyAuth) + mock_prefetch.assert_awaited_once() + assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + call: Final = _authorize_authenticated_request( + resolved, request, {"model": "gpt-4o"}, "/chat/completions", api_key + ) if model_allowed: - assert isinstance(await call, UserAPIKeyAuth) - mock_prefetch.assert_awaited_once() - assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + await call else: with pytest.raises(ProxyException) as exc: await call assert exc.value.type == ProxyErrorTypes.key_model_access_denied - mock_prefetch.assert_not_awaited() finally: for attr, val in _original_values.items(): setattr(_proxy_server_mod, attr, val) @@ -2601,6 +2553,48 @@ async def test_auto_register_map_existing_key_uses_the_issuers_own_user_field_ov assert generate_key.await_count == (0 if expect_reuse else 1) +def _validated_jwt_admission_cache(*, org_id: str | None = None, end_user_id: str | None = None) -> UserApiKeyCache: + cache: Final = UserApiKeyCache() + cache.set_cache( + "validated-user", + LiteLLM_UserTable( + user_id="validated-user", + user_email="validated@example.com", + user_role="internal_user", + teams=["validated-team"], + ), + model_type=LiteLLM_UserTable, + ) + cache.set_cache( + "team_id:validated-team", + LiteLLM_TeamTableCachedObj( + team_id="validated-team", + models=["gpt-4o-mini", "gpt-5.6"], + members_with_roles=[{"user_id": "validated-user", "role": "user"}], + organization_id=org_id, + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + cache.set_cache( + team_membership_reservation_cache_key(user_id="validated-user", team_id="validated-team"), + LiteLLM_TeamMembership(user_id="validated-user", team_id="validated-team"), + model_type=LiteLLM_TeamMembership, + ) + if org_id is not None: + org: Final = LiteLLM_OrganizationTable( + organization_id=org_id, budget_id="org-budget", created_by="admin", updated_by="admin" + ) + cache.set_cache(f"org_id:{org_id}", org, model_type=LiteLLM_OrganizationTable) + cache.set_cache(f"org_id:{org_id}:with_budget", org, model_type=LiteLLM_OrganizationTable) + if end_user_id is not None: + cache.set_cache( + f"end_user_id:{end_user_id}", + LiteLLM_EndUserTable(user_id=end_user_id, blocked=False), + model_type=LiteLLM_EndUserTable, + ) + return cache + + @pytest.mark.asyncio @pytest.mark.parametrize( ("map_existing_key", "master_key", "reused_key_models", "expect_denied"), @@ -2615,7 +2609,7 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( map_existing_key: bool, master_key: str | None, reused_key_models: list[str], expect_denied: bool ) -> None: jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -2648,12 +2642,21 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( "jwt_claims": {"sub": "user1"}, } - mock_request = MagicMock() - mock_request.url.path = "/v1/chat/completions" - mock_request.method = "POST" - mock_request.headers = {"authorization": f"Bearer {jwt_token}"} - mock_request.query_params = {} - mock_request.state = SimpleNamespace() + mock_request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"authorization", f"Bearer {jwt_token}".encode())], + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps({"model": "gpt-4o-mini"}).encode(), + } + ), + ) with ( patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), @@ -2686,14 +2689,13 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( return_value=reused_key, ), ): - call = user_api_key_auth_builder( + call = user_api_key_auth( request=mock_request, api_key=jwt_token, azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, - request_data={"model": "gpt-4o-mini"}, ) if expect_denied: with pytest.raises(ProxyException, match="not available for this API key"): @@ -2719,7 +2721,7 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - """ jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -2952,7 +2954,10 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): id auth_builder resolved must reach the key creation, not be dropped when valid_token is swapped for the freshly registered key.""" jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() + prisma_client: Final = MagicMock() + prisma_client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + prisma_client.writer_db.litellm_retiredagent.find_unique = AsyncMock(return_value=None) jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"}) @@ -3001,7 +3006,7 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): general_settings={"enable_jwt_auth": True}, premium_user=True, master_key="sk-master", - prisma_client=MagicMock(), + prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=MagicMock(), jwt_handler=jwt_handler, @@ -3274,7 +3279,7 @@ class TestJWTOAuth2Coexistence: async def test_auto_register_passes_validated_org_context_to_generated_key(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache(org_id="validated-org", end_user_id="validated-end-user") prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True diff --git a/tests/unit/proxy/mcp_server/test_idp_token_exchange.py b/tests/unit/proxy/mcp_server/test_idp_token_exchange.py index 036ee4a332b..75b0a5cd308 100644 --- a/tests/unit/proxy/mcp_server/test_idp_token_exchange.py +++ b/tests/unit/proxy/mcp_server/test_idp_token_exchange.py @@ -1,20 +1,34 @@ +import json import logging from typing import Final +import jwt import pytest -from fastapi import HTTPException +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import HTTPException, Request from prisma.engine.errors import BinaryNotFoundError from prisma.errors import DataError from litellm.caching.caching import DualCache -from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + JWTIssuerConfig, + LiteLLM_JWTAuth, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + Member, + ProxyException, +) from litellm.proxy.auth.handle_jwt import JWKSUnreachableError, JWTHandler, jwks_unavailable_exception +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, SUBJECT_TOKEN_CHECK_FAULTED, SUBJECT_TOKEN_CHECK_UNAVAILABLE, TokenExchangePrerequisites, + exchange_idp_subject_token, identity_from_subject_token, token_exchange_available, ) @@ -246,3 +260,71 @@ async def test_agent_token_cannot_be_exchanged_for_a_user_identity(user_id: str assert isinstance(result, SubjectTokenRefusal) assert result.error == "invalid_request" assert "direct JWT authentication" in result.description + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enforce_team", "team_id", "expected"), + ( + (True, "exchange-team", SubjectIdentity(user_id="exchange-user", team_id="exchange-team")), + (False, None, SubjectIdentity(user_id="exchange-user", team_id=None)), + (True, None, SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)), + ), + ids=("native-team", "teamless-without-enforcement", "teamless-under-enforcement"), +) +async def test_subject_exchange_applies_team_admission_before_returning_an_identity( + monkeypatch: pytest.MonkeyPatch, + enforce_team: bool, + team_id: str | None, + expected: SubjectIdentity | SubjectTokenRefusal, +) -> None: + cache: Final = UserApiKeyCache() + database: Final = object() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + team_ids_jwt_field="groups", + enforce_team_based_model_access=enforce_team, + team_allowed_routes=["/token"], + ), + ) + signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_url: Final = "https://identity.invalid/exchange/jwks" + cache.set_cache( + f"litellm_jwt_auth_keys_{jwks_url}", + [json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(signing_key.public_key()))], + ) + cache.set_cache( + "exchange-user", + LiteLLM_UserTable(user_id="exchange-user", user_role="internal_user", teams=[team_id] if team_id else []), + model_type=LiteLLM_UserTable, + ) + if team_id is not None: + cache.set_cache( + f"team_id:{team_id}", + LiteLLM_TeamTable(team_id=team_id, members_with_roles=[Member(user_id="exchange-user", role="user")]), + model_type=LiteLLM_TeamTable, + ) + cache.set_cache( + f"team_membership:exchange-user:{team_id}", + LiteLLM_TeamMembership(user_id="exchange-user", team_id=team_id), + model_type=LiteLLM_TeamMembership, + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + token: Final = jwt.encode( + {"sub": "exchange-user", "groups": [team_id] if team_id else []}, signing_key, algorithm="RS256" + ) + request: Final = Request({"type": "http", "method": "POST", "path": "/token", "headers": []}) + assert await exchange_idp_subject_token(token, request) == expected diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py index 1a4751b4598..322752782fd 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -11,6 +11,7 @@ import jwt import pytest import respx from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import Request from fastapi.testclient import TestClient from jwt.algorithms import RSAAlgorithm @@ -704,10 +705,23 @@ def _key_with_alias_to_budgeted_target(monkeypatch: pytest.MonkeyPatch, target_b return key +@pytest.mark.parametrize("custom_auth", (False, True), ids=("virtual-key", "custom-auth")) def test_native_route_refuses_a_saved_alias_when_its_target_is_over_budget( - client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, custom_auth: bool ) -> None: key: Final = _key_with_alias_to_budgeted_target(monkeypatch, target_budget=0.0) + if custom_auth: + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + identity: Final = proxy_server.user_api_key_cache.get_cache(key=hash_token(key), model_type=UserAPIKeyAuth) + assert identity is not None + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return identity + + monkeypatch.setattr(proxy_server, "user_custom_auth", authenticate) + monkeypatch.setattr( + proxy_server, "general_settings", {**proxy_server.general_settings, "custom_auth_run_common_checks": True} + ) _account_catalog(respx_mock) provider: Final = _open_weights_provider(respx_mock) response: Final = client.post( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 24ad068e6df..2c23305d5c9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3691,12 +3691,19 @@ class TestResponsesWSFirstFrameModelAuth: "standard_logging_guardrail_information": [{"guardrail_cost": 0.2}, {"guardrail_cost": 0.4}], } - @pytest.mark.asyncio @pytest.mark.parametrize( "relay_outcome", ["still_relaying", "already_returned", "returns_despite_the_cancel", "crashes_despite_the_cancel"], ) - async def test_session_limit_frees_the_slot_at_once_and_never_cuts_the_failure_row(self, relay_outcome: str): + def test_session_limit_frees_the_slot_at_once_and_never_cuts_the_failure_row(self, relay_outcome: str) -> None: + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(self._assert_session_limit_releases_slot_before_failure_row(relay_outcome)) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + async def _assert_session_limit_releases_slot_before_failure_row(self, relay_outcome: str) -> None: from litellm.exceptions import RateLimitError from litellm.proxy.response_api_endpoints.endpoints import responses_websocket_endpoint diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 00c1b855033..6f8c49cc0f9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26239,6 +26239,17 @@ export interface components { /** Updated By */ updated_by?: string | null; }; + /** AccessGroupRule */ + AccessGroupRule: { + /** Access Group */ + access_group: string; + /** Attribute */ + attribute: string; + /** Contains */ + contains?: string | number | boolean | null; + /** Equals */ + equals?: string | number | boolean | null; + }; /** AccessGroupUpdateRequest */ AccessGroupUpdateRequest: { /** Access Agent Ids */ @@ -30808,6 +30819,8 @@ export interface components { * @description Documents all the fields supported by `general_settings` in config.yaml */ ConfigGeneralSettings: { + /** Access Group Rules */ + access_group_rules?: components["schemas"]["AccessGroupRule"][] | null; /** * Admission Queue Timeout Seconds * @description maximum time a request waits for a worker slot diff --git a/uv.lock b/uv.lock index 8712b78bbd6..62187d596ec 100644 --- a/uv.lock +++ b/uv.lock @@ -1117,6 +1117,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/0f/f897abe4ea0a8c408ae65c8c83bffab4936ad65d6032d4fb4cd35bbdc3ee/cachetools-7.1.1-py3-none-any.whl", hash = "sha256:0335cd7a0952d2b22327441fb0628139e234c565559eeb91a8a4ac7551c5353d", size = 16775, upload-time = "2026-05-03T20:00:27.857Z" }, ] +[[package]] +name = "casbin" +version = "1.43.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "simpleeval" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ad/df/ff2aa55cf0d7c14622ce4f9252cdc34c828c81d4213965d73207ac5434ae/casbin-1.43.0.tar.gz", hash = "sha256:d2e90ce8e72f912877851e94d37999f32c558c6ba7aba0437d483275262e86e0", size = 425727, upload-time = "2025-05-10T06:57:18.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/07/facef6abd81378e6153b757dc1848621675d971fbc88ebb5d182ebc1c37f/casbin-1.43.0-py3-none-any.whl", hash = "sha256:63a3d1228870250e859ccd94133fe478821093f71dd37f05e0baa0c6fea26623", size = 475059, upload-time = "2025-05-10T06:57:16.89Z" }, +] + [[package]] name = "certifi" version = "2026.10.19" @@ -5096,6 +5108,7 @@ proxy = [ { name = "backoff" }, { name = "backports-zstd", marker = "python_full_version < '3.14'" }, { name = "boto3" }, + { name = "casbin" }, { name = "cryptography" }, { name = "expression" }, { name = "fastapi" }, @@ -5312,6 +5325,7 @@ requires-dist = [ { name = "backports-zstd", marker = "python_full_version < '3.14' and extra == 'proxy'", specifier = ">=1.0.0,<2.0" }, { name = "boto3", specifier = ">=1.43.117,<2.0" }, { name = "boto3", marker = "extra == 'proxy'", specifier = ">=1.43.117,<2.0" }, + { name = "casbin", marker = "extra == 'proxy'", specifier = "==1.43.0" }, { name = "certifi", specifier = ">=2026.10.19" }, { name = "click", specifier = ">=8.0.0,<9.0" }, { name = "cryptography", marker = "extra == 'proxy'", specifier = ">=49.0.0,<51.0" }, @@ -10377,6 +10391,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/59/0782e51887ac6b07ffd1570e0364cf901ebc36345fea669969d2084baebb/simple_websocket-1.1.0-py3-none-any.whl", hash = "sha256:4af6069630a38ed6c561010f0e11a5bc0d4ca569b36306eb257cd9a192497c8c", size = 13842, upload-time = "2024-10-10T22:39:29.645Z" }, ] +[[package]] +name = "simpleeval" +version = "1.0.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/89/3e/25326f373871a34e78cac231c0d3c18a0a5c3a80aa4813b50b9f26a5d460/simpleeval-1.0.8.tar.gz", hash = "sha256:b417f869c11f2de034eb919820a825f35d12ce99f63a43971947ce73679eb9ba", size = 31846, upload-time = "2026-09-12T10:16:08.103Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/33/d1ecd9f4f5f272b36ad7b268e307ba0a59c109d8ad134a231cda780baac9/simpleeval-1.0.8-py3-none-any.whl", hash = "sha256:c9c4ff42879ebf0c05f6562cb81031869948916be343bcc950b81fcc081fe9ce", size = 18925, upload-time = "2026-09-12T10:16:06.862Z" }, +] + [[package]] name = "six" version = "1.17.0"