mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(auth): grant access groups through principal attribute rules (#50455)
* feat(auth): grant access groups through principal attribute rules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): satisfy cache tracing and recursion CI checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): close access group rule authorization bypasses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): reject self-editable identity rule attributes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Centralize JWT model authorization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Cover centralized key checks and name denied JWT models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Test centralized authorization for custom auth and JWT mappings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Preserve model normalization before budget validation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Use a controlled clock for WebSocket session deadline tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Unify JWT credential admission and simplify team selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Simplify request-time access group authorization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Use real requests in custom auth budget fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia <krrishdholakia@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Krrish Dholakia
GitHub
parent
74cad3a2c3
commit
70f399b8c5
20 changed files with 1599 additions and 339 deletions
|
|
@ -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.",
|
||||
|
|
|
|||
201
litellm/proxy/auth/access_group_rules.py
Normal file
201
litellm/proxy/auth/access_group_rules.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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'",
|
||||
|
|
|
|||
221
tests/integration/authorization/test_access_group_rules.py
Normal file
221
tests/integration/authorization/test_access_group_rules.py
Normal file
|
|
@ -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
|
||||
)
|
||||
245
tests/unit/proxy/auth/test_access_group_rules.py
Normal file
245
tests/unit/proxy/auth/test_access_group_rules.py
Normal file
|
|
@ -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"}]}
|
||||
)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
13
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
13
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
23
uv.lock
generated
23
uv.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue