From 85d98f0adf97243a2dc67da6bc996da56c7a41ac Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 8 Nov 2026 23:39:58 +0000 Subject: [PATCH] feat(auth): stage resolved model scopes and member grant reads --- litellm-rust/crates/python-bridge/src/lib.rs | 18 +- .../crates/python-bridge/src/request_auth.rs | 24 ++ .../crates/request-auth-core/src/lib.rs | 2 + .../request-auth-core/src/member_models.rs | 77 +++++ .../request-auth-core/src/resolved_models.rs | 216 +++++++++++++ litellm/_auth/member_models.py | 32 ++ litellm/_auth/resolved_models.py | 78 +++++ litellm/proxy/auth/member_model_adapter.py | 132 ++++++++ litellm/proxy/auth/resolved_model_adapter.py | 283 ++++++++++++++++++ litellm/proxy/auth/team_override_adapter.py | 10 +- litellm/rust_bridge/_native.pyi | 2 + tests/unit/_auth/test_member_models.py | 56 ++++ tests/unit/_auth/test_resolved_models.py | 156 ++++++++++ 13 files changed, 1072 insertions(+), 14 deletions(-) create mode 100644 litellm-rust/crates/request-auth-core/src/member_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/resolved_models.rs create mode 100644 litellm/_auth/member_models.py create mode 100644 litellm/_auth/resolved_models.py create mode 100644 litellm/proxy/auth/member_model_adapter.py create mode 100644 litellm/proxy/auth/resolved_model_adapter.py create mode 100644 tests/unit/_auth/test_member_models.py create mode 100644 tests/unit/_auth/test_resolved_models.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index a00057e0013..5acadfb1fb1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -50,15 +50,15 @@ mod _native { request_auth_custom_flow, request_auth_expired, request_auth_fallback_access, request_auth_fallback_candidate, request_auth_key_model_access, request_auth_litellm_header_token, request_auth_managed, request_auth_member_limit, - request_auth_model_access, request_auth_model_attribution, request_auth_model_budget_check, - request_auth_model_budget_query, request_auth_model_budget_resolve, - request_auth_model_groups, request_auth_model_helpers, request_auth_org_admin, - request_auth_org_route, request_auth_owner_teams, request_auth_personal_budget, - request_auth_personal_model_request, request_auth_personal_models, - request_auth_personal_owner, request_auth_reservation, request_auth_route_access, - request_auth_scheme_credential, request_auth_select_api_key, - request_auth_select_websocket_api_key, request_auth_start, request_auth_team_override, - request_auth_unpriced_model, + request_auth_member_models, request_auth_model_access, request_auth_model_attribution, + request_auth_model_budget_check, request_auth_model_budget_query, + request_auth_model_budget_resolve, request_auth_model_groups, request_auth_model_helpers, + request_auth_org_admin, request_auth_org_route, request_auth_owner_teams, + request_auth_personal_budget, request_auth_personal_model_request, + request_auth_personal_models, request_auth_personal_owner, request_auth_reservation, + request_auth_resolved_models, request_auth_route_access, request_auth_scheme_credential, + request_auth_select_api_key, request_auth_select_websocket_api_key, request_auth_start, + request_auth_team_override, request_auth_unpriced_model, }; #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index 9fe009e647b..3e69ffb1126 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -622,3 +622,27 @@ pub(crate) fn request_auth_personal_model_request( .map_err(map_error)?; to_py(py, &decision) } + +#[pyfunction] +pub(crate) fn request_auth_member_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::member_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::member_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_resolved_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::resolved_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::resolved_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index c3320a1a1d8..25a60e0806d 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -17,6 +17,7 @@ pub mod fallback_candidate; mod key; pub mod key_model_access; pub mod managed; +pub mod member_models; pub mod model_access; pub mod model_attribution; pub mod model_budget_data; @@ -32,6 +33,7 @@ pub mod personal_models; pub mod personal_owner; pub mod protocol; pub mod reservation; +pub mod resolved_models; pub mod route_access; pub mod team_override; pub mod unpriced_models; diff --git a/litellm-rust/crates/request-auth-core/src/member_models.rs b/litellm-rust/crates/request-auth-core/src/member_models.rs new file mode 100644 index 00000000000..a4cacc8959b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/member_models.rs @@ -0,0 +1,77 @@ +use crate::{Error, model_access}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub team_present: bool, + pub membership_loaded: bool, + pub grants: Option>, + pub scope: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: model_access::Decision, + member_models: Option>, + }, + ReadMembership, + ReadGrants, + ReadScope, +} + +pub fn decide(facts: &Facts) -> Result { + if !facts.user_present || !facts.team_present { + return Ok(Decision::Allow); + } + if !facts.membership_loaded { + return Ok(Decision::ReadMembership); + } + let Some(grants) = &facts.grants else { + return Ok(Decision::ReadGrants); + }; + if grants.is_empty() { + return Ok(Decision::Allow); + } + let Some(scope) = &facts.scope else { + return Ok(Decision::ReadScope); + }; + let denial = model_access::check(scope)?; + if denial.allowed() { + return Ok(Decision::Allow); + } + let member_models = + matches!(denial, model_access::Decision::Deny { .. }).then(|| grants.clone()); + Ok(Decision::Deny { + denial, + member_models, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(false, false, Decision::Allow)] + #[case::lookup(true, false, Decision::ReadMembership)] + #[case::preloaded(true, true, Decision::ReadGrants)] + fn membership_read( + #[case] user_present: bool, + #[case] membership_loaded: bool, + #[case] expected: Decision, + ) { + let facts = Facts { + user_present, + team_present: true, + membership_loaded, + grants: None, + scope: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/resolved_models.rs b/litellm-rust/crates/request-auth-core/src/resolved_models.rs new file mode 100644 index 00000000000..5718b0966c0 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/resolved_models.rs @@ -0,0 +1,216 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + member_models, + model_access::{self, Decision as ModelDecision}, + personal_model_request, team_override, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct TeamRead { + pub found: bool, + pub failed: bool, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub model: String, + pub config_nonempty: bool, + pub key_models: Vec, + pub team_id: Option, + pub user_id: Option, + pub project_present: bool, + pub customer_present: bool, + pub database_present: bool, + pub personal_enabled: bool, + pub master: bool, + pub role: Option, + pub key: Option, + pub team_read: Option, + pub embedded_ready: bool, + pub team: Option, + pub member: Option, + pub owner_loaded: bool, + pub owner_present: bool, + pub owner_failed: bool, + pub personal: Option, + pub project_loaded: bool, + pub project_grants: Vec, + pub project: Option, + pub customer_loaded: bool, + pub customer_grants: Vec, + pub customer: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: ModelDecision, + member_models: Option>, + }, + ReadKey, + ReadKeyGroups, + ReadTeamObject, + ReadEmbeddedTeam, + ReadTeamScope, + ReadMember, + ReadOwner, + ReadPersonal, + ReadProjectObject, + ReadProjectScope, + ReadCustomerObject, + ReadCustomerScope, +} + +pub fn decide(facts: &Facts) -> Result { + let deny = |denial| { + Ok(Decision::Deny { + denial, + member_models: None, + }) + }; + if !facts.config_nonempty + && !facts + .key_models + .iter() + .any(|model| model == "all-team-models") + { + let Some(key) = &facts.key else { + return Ok(Decision::ReadKey); + }; + let key = fallback_access::grouped(key.clone(), false)?; + if matches!(key, ModelDecision::ReadGroups) { + return Ok(Decision::ReadKeyGroups); + } + if !key.allowed() { + return deny(key); + } + } + if facts.team_id.is_some() { + let Some(read) = &facts.team_read else { + return Ok(Decision::ReadTeamObject); + }; + if read.failed && !facts.embedded_ready { + return Ok(Decision::ReadEmbeddedTeam); + } + if read.found || read.failed { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeamScope); + }; + match team_override::decide(team)? { + team_override::Decision::Allow => (), + team_override::Decision::Deny { denial } => return deny(denial), + _ => return Err(Error::InvalidAnswer), + } + if read.found && facts.user_id.is_some() { + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember); + }; + match member_models::decide(member)? { + member_models::Decision::Allow => (), + member_models::Decision::Deny { + denial, + member_models, + } => { + return Ok(Decision::Deny { + denial, + member_models, + }); + } + _ => return Err(Error::InvalidAnswer), + } + } + } + } + let model_access::Response::Predicate(personal) = + model_access::decide(model_access::Request::PersonalApplies { + enabled: facts.personal_enabled, + token_present: true, + team_id: facts.team_id.clone(), + user_id: facts.user_id.clone(), + master: facts.master, + role: facts.role.clone(), + api_key_present: true, + inheritance: false, + })? + else { + return Err(Error::InvalidAnswer); + }; + if personal { + if !facts.owner_loaded { + return Ok(Decision::ReadOwner); + } + if facts.owner_failed || !facts.owner_present { + return deny(ModelDecision::PersonalTeam { + model: facts.model.clone(), + }); + } + let Some(personal) = &facts.personal else { + return Ok(Decision::ReadPersonal); + }; + match personal_model_request::decide(personal)? { + personal_model_request::Decision::Allow => (), + personal_model_request::Decision::Deny { model } => { + return deny(ModelDecision::PersonalTeam { model }); + } + _ => return Err(Error::InvalidAnswer), + } + } + if facts.project_present { + if !facts.project_loaded { + return Ok(Decision::ReadProjectObject); + } + if !facts.project_grants.is_empty() { + let Some(project) = &facts.project else { + return Ok(Decision::ReadProjectScope); + }; + let decision = model_access::check(project)?; + if !decision.allowed() { + return deny(decision); + } + } + } + if facts.customer_present && facts.database_present { + if !facts.customer_loaded { + return Ok(Decision::ReadCustomerObject); + } + if !facts.customer_grants.is_empty() { + let Some(customer) = &facts.customer else { + return Ok(Decision::ReadCustomerScope); + }; + let model_access::Response::Decision(decision) = + model_access::decide(model_access::Request::Customer(customer.clone()))? + else { + return Err(Error::InvalidAnswer); + }; + if !decision.allowed() { + return deny(decision); + } + } + } + Ok(Decision::Allow) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::key(false, Decision::ReadKey)] + #[case::team(true, Decision::ReadTeamObject)] + fn key_before_team(#[case] config_nonempty: bool, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + config_nonempty, + team_id: Some("team".into()), + ..Facts::default() + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm/_auth/member_models.py b/litellm/_auth/member_models.py new file mode 100644 index 00000000000..a2e845b69f4 --- /dev/null +++ b/litellm/_auth/member_models.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class MemberModelFacts: + user_present: bool + team_present: bool + membership_loaded: bool = False + grants: tuple[str, ...] | None = None + scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class MemberModelDecision: + kind: Literal["Allow", "Deny", "ReadMembership", "ReadGrants", "ReadScope"] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + + +_DECISION: Final = TypeAdapter(MemberModelDecision) + + +@native_boundary +def member_models(facts: MemberModelFacts) -> MemberModelDecision: + return _DECISION.validate_python(_native.request_auth_member_models(asdict(facts))) diff --git a/litellm/_auth/resolved_models.py b/litellm/_auth/resolved_models.py new file mode 100644 index 00000000000..5e784b035d8 --- /dev/null +++ b/litellm/_auth/resolved_models.py @@ -0,0 +1,78 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm._auth.personal_model_request import PersonalModelRequestFacts +from litellm._auth.team_override import TeamOverrideFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ResolvedTeamRead: + found: bool = False + failed: bool = False + + +@dataclass(frozen=True, slots=True) +class ResolvedModelFacts: + model: str = "" + config_nonempty: bool = False + key_models: tuple[str, ...] = () + team_id: str | None = None + user_id: str | None = None + project_present: bool = False + customer_present: bool = False + database_present: bool = False + personal_enabled: bool = False + master: bool = False + role: str | None = None + key: GroupedModelScope | None = None + team_read: ResolvedTeamRead | None = None + embedded_ready: bool = False + team: TeamOverrideFacts | None = None + member: MemberModelFacts | None = None + owner_loaded: bool = False + owner_present: bool = False + owner_failed: bool = False + personal: PersonalModelRequestFacts | None = None + project_loaded: bool = False + project_grants: tuple[str, ...] = () + project: ModelAccessFacts | None = None + customer_loaded: bool = False + customer_grants: tuple[str, ...] = () + customer: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ResolvedModelDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadKey", + "ReadKeyGroups", + "ReadTeamObject", + "ReadEmbeddedTeam", + "ReadTeamScope", + "ReadMember", + "ReadOwner", + "ReadPersonal", + "ReadProjectObject", + "ReadProjectScope", + "ReadCustomerObject", + "ReadCustomerScope", + ] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + + +_DECISION: Final = TypeAdapter(ResolvedModelDecision) + + +@native_boundary +def resolved_models(facts: ResolvedModelFacts) -> ResolvedModelDecision: + return _DECISION.validate_python(_native.request_auth_resolved_models(asdict(facts))) diff --git a/litellm/proxy/auth/member_model_adapter.py b/litellm/proxy/auth/member_model_adapter.py new file mode 100644 index 00000000000..08c2af532d6 --- /dev/null +++ b/litellm/proxy/auth/member_model_adapter.py @@ -0,0 +1,132 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.member_models import MemberModelFacts, member_models +from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import project_model_facts, raise_model_denial +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + team: LiteLLM_TeamTable + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + membership: LiteLLM_TeamMembership | None + aliases: Mapping[str, str] | None + + +async def read_member_model_facts( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, + key_model_aliases: Mapping[str, str] | None = None, +) -> MemberModelFacts: + facts: Final = MemberModelFacts( + user_present=valid_token.user_id is not None, + team_present=team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # preserve unvalidated cache row identity guard + membership_loaded=team_membership_loaded, + ) + return await _read( + facts, + _Source( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + key_model_aliases, + ), + ) + + +async def check_team_member_model_access( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, + key_model_aliases: Mapping[str, str] | None = None, +) -> None: + facts: Final = await read_member_model_facts( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + key_model_aliases, + ) + decision: Final = member_models(facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete member model access facts") + raise_model_denial( + decision.denial, + model, + user_id=valid_token.user_id, + team_id=team_object.team_id, + member_models=decision.member_models, + ) + + +async def _read(facts: MemberModelFacts, source: _Source) -> MemberModelFacts: + from litellm.proxy.auth.auth_checks import get_team_membership + + decision: Final = member_models(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("member model lookup without user identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + return await _read(replace(facts, membership_loaded=True), replace(source, membership=membership)) + case "ReadGrants": + row: Final = source.membership.litellm_budget_table if source.membership is not None else None + return await _read( + replace(facts, grants=tuple(row.allowed_models or ()) if row is not None else ()), source + ) + case "ReadScope": + if facts.grants is None: + raise NativeAuthRuntimeError("member model scope read without grants") + scope: Final = project_model_facts( + source.model, + source.router, + facts.grants, + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return await _read(replace(facts, scope=scope), source) diff --git a/litellm/proxy/auth/resolved_model_adapter.py b/litellm/proxy/auth/resolved_model_adapter.py new file mode 100644 index 00000000000..66a34a5cd59 --- /dev/null +++ b/litellm/proxy/auth/resolved_model_adapter.py @@ -0,0 +1,283 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from pydantic import BaseModel, Field, TypeAdapter + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import key_model_grants +from litellm._auth.resolved_models import ResolvedModelFacts, ResolvedTeamRead, resolved_models +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.member_model_adapter import read_member_model_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.personal_model_request_adapter import read_personal_model_request_facts +from litellm.proxy.auth.team_override_adapter import read_team_override_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +class _Config(BaseModel, frozen=True): + config: Mapping[str, object] = Field(default_factory=dict[str, object]) + + +_METADATA: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + +@dataclass(frozen=True, slots=True) +class ResolvedModelReading: + facts: ResolvedModelFacts + team: LiteLLM_TeamTableCachedObj | None = None + owner: LiteLLM_UserTable | None = None + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + team: LiteLLM_TeamTableCachedObj | None = None + owner: LiteLLM_UserTable | None = None + project: LiteLLM_ProjectTable | None = None + customer: LiteLLM_EndUserTable | None = None + + +async def read_resolved_model_facts( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> ResolvedModelReading: + facts: Final = ResolvedModelFacts( + model=model, + config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config), + key_models=object_model_grants(valid_token), + team_id=valid_token.team_id, + user_id=valid_token.user_id, + project_present=valid_token.project_id is not None, + customer_present=valid_token.end_user_id is not None, + database_present=prisma_client is not None, + personal_enabled=bool(litellm.personal_key_model_access_from_teams), + master=valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS, + role=valid_token.user_role, + ) + return await _read( + facts, _Source(model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj) + ) + + +async def can_key_call_resolved_model( + model: str, llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None +) -> None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + reading: Final = await read_resolved_model_facts( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = resolved_models(reading.facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete resolved model access facts") + raise_model_denial( + decision.denial, + model, + user_id=valid_token.user_id, + team_id=reading.team.team_id if reading.team is not None else None, + teams=reading.owner.teams if reading.owner is not None else None, + member_models=decision.member_models, + ) + + +async def _read(facts: ResolvedModelFacts, source: _Source) -> ResolvedModelReading: + from litellm.proxy.auth.auth_checks import ( + get_end_user_object, + get_key_end_user_budget_id, + get_models_from_access_groups, + get_project_object, + get_team_object, + get_user_object, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, + ) + + decision: Final = resolved_models(facts) + match decision.kind: + case "Allow" | "Deny": + return ResolvedModelReading(facts, source.team, source.owner) + case "ReadKey": + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + key: Final = project_model_facts( + source.model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ) + return await _read( + replace(facts, key=GroupedModelScope(key, tuple(source.token.access_group_ids or ()))), source + ) + case "ReadKeyGroups": + if facts.key is None: + raise NativeAuthRuntimeError("resolved key group read without scope") + models: Final = await get_models_from_access_groups(facts.key.group_ids) + return await _read(replace(facts, key=replace(facts.key, group_models=tuple(models))), source) + case "ReadTeamObject": + if source.token.team_id is None: + raise NativeAuthRuntimeError("resolved team read without identity") + try: + team: Final = await get_team_object( + team_id=source.token.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy resolved-model team lookup falls back to the embedded token row + return await _read(replace(facts, team_read=ResolvedTeamRead(failed=True)), source) + return await _read( + replace(facts, team_read=ResolvedTeamRead(found=team is not None)), # pyright: ignore[reportUnnecessaryComparison] # preserve the legacy missing-row result + replace(source, team=team), + ) + case "ReadEmbeddedTeam": + embedded: Final = LiteLLM_TeamTableCachedObj( + team_id=source.token.team_id, # pyright: ignore[reportArgumentType] # native selects this read only for a present team ID + models=source.token.team_models, + blocked=source.token.team_blocked, + team_alias=source.token.team_alias, + metadata=source.token.team_metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy typed row validates embedded metadata + object_permission_id=source.token.team_object_permission_id, + object_permission=source.token.team_object_permission, + ) + return await _read(replace(facts, embedded_ready=True), replace(source, team=embedded)) + case "ReadTeamScope": + team_scope: Final = await read_team_override_facts( + source.model, + source.token, + source.team, + source.router, + source.token.team_model_aliases, + key_model_aliases_for_auth_check(source.token), + ) + return await _read(replace(facts, team=team_scope), source) + case "ReadMember": + if source.team is None: + raise NativeAuthRuntimeError("resolved member model read without team") + member: Final = await read_member_model_facts( + source.model, + source.team, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + ) + return await _read(replace(facts, member=member), source) + case "ReadOwner": + if source.token.user_id is None: + raise NativeAuthRuntimeError("resolved personal owner read without identity") + try: + owner: Final = await get_user_object( + user_id=source.token.user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + user_id_upsert=False, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve raw owner lookup failure and native fail-closed model selection + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", error + ) + return await _read(replace(facts, owner_loaded=True, owner_failed=True), source) + return await _read( + replace(facts, owner_loaded=True, owner_present=owner is not None), replace(source, owner=owner) + ) + case "ReadPersonal": + personal: Final = await read_personal_model_request_facts( + source.model, source.owner, source.token, source.router, source.prisma, source.cache, source.logging + ) + return await _read(replace(facts, personal=personal), source) + case "ReadProjectObject": + if source.token.project_id is None: + raise NativeAuthRuntimeError("resolved project read without identity") + project: Final = await get_project_object( + project_id=source.token.project_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + return await _read( + replace( + facts, + project_loaded=True, + project_grants=object_model_grants(project) if project is not None else (), + ), + replace(source, project=project), + ) + case "ReadProjectScope": + if source.project is None: + raise NativeAuthRuntimeError("resolved project scope read without row") + project_scope: Final = project_model_facts( + source.model, + source.router, + facts.project_grants, + team_id=source.project.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="project", + ) + return await _read(replace(facts, project=project_scope), source) + case "ReadCustomerObject": + if source.token.end_user_id is None or source.prisma is None: + raise NativeAuthRuntimeError("resolved customer read without identity or database") + customer: Final = await get_end_user_object( + end_user_id=source.token.end_user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + token_end_user_max_budget=source.token.end_user_max_budget, + key_end_user_budget_id=get_key_end_user_budget_id(_METADATA.validate_python(source.token.metadata)), # pyright: ignore[reportUnknownMemberType] # validate the legacy metadata field at the IO boundary + ) + return await _read( + replace( + facts, + customer_loaded=True, + customer_grants=tuple(customer.models or ()) if customer is not None else (), + ), + replace(source, customer=customer), + ) + case "ReadCustomerScope": + customer_scope: Final = project_model_facts( + source.model, + source.router, + facts.customer_grants, + team_model_aliases_for_auth_check(source.token), + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "customer", + ) + return await _read(replace(facts, customer=customer_scope), source) diff --git a/litellm/proxy/auth/team_override_adapter.py b/litellm/proxy/auth/team_override_adapter.py index bc3fc6ffed6..fc471573d1a 100644 --- a/litellm/proxy/auth/team_override_adapter.py +++ b/litellm/proxy/auth/team_override_adapter.py @@ -15,7 +15,7 @@ from litellm.router_api import RegistryReader @dataclass(frozen=True, slots=True) class _Source: model: str | Sequence[str] - token: UserAPIKeyAuth + token: UserAPIKeyAuth | None team: LiteLLM_TeamTable | None router: RegistryReader | None team_aliases: Mapping[str, str] | None @@ -25,7 +25,7 @@ class _Source: async def read_team_override_facts( model: str | Sequence[str], - valid_token: UserAPIKeyAuth, + valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, llm_router: RegistryReader | None, team_model_aliases: Mapping[str, str] | None = None, @@ -40,7 +40,7 @@ async def read_team_override_facts( async def can_team_call_model_with_override( model: str | Sequence[str], - valid_token: UserAPIKeyAuth, + valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, llm_router: RegistryReader | None, team_model_aliases: Mapping[str, str] | None = None, @@ -94,8 +94,8 @@ async def _read(facts: TeamOverrideFacts, source: _Source) -> TeamOverrideFacts: source.model, source.router, facts.assigned_models, - source.token.team_model_aliases, - source.token.team_id, + source.token.team_model_aliases if source.token is not None else None, + source.token.team_id if source.token is not None else None, source.key_aliases, "key", ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index fd3cc7ef39f..b3dae3408c0 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -64,6 +64,8 @@ def request_auth_team_override(facts: object) -> object: ... def request_auth_model_helpers(facts: object) -> object: ... def request_auth_personal_owner(facts: object) -> object: ... def request_auth_personal_model_request(facts: object) -> object: ... +def request_auth_member_models(facts: object) -> object: ... +def request_auth_resolved_models(facts: object) -> object: ... @final class NativeClickHouseSpendConfig: diff --git a/tests/unit/_auth/test_member_models.py b/tests/unit/_auth/test_member_models.py new file mode 100644 index 00000000000..b9901fd8706 --- /dev/null +++ b/tests/unit/_auth/test_member_models.py @@ -0,0 +1,56 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.member_models import MemberModelDecision as Decision +from litellm._auth.member_models import MemberModelFacts as Facts +from litellm._auth.member_models import member_models +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Facts) -> Decision: + if not facts.user_present or not facts.team_present: + return Decision("Allow") + if not facts.membership_loaded: + return Decision("ReadMembership") + if facts.grants is None: + return Decision("ReadGrants") + if not facts.grants: + return Decision("Allow") + if facts.scope is None: + return Decision("ReadScope") + denied: Final = model_access(facts.scope) + if denied.kind == "Allow": + return Decision("Allow") + return Decision("Deny", denied, facts.grants if denied.kind == "Deny" else None) + + +_BASE: Final = Facts(True, True) +_GRANTS: Final = replace(_BASE, membership_loaded=True, grants=("allowed",)) +_SCOPE: Final = ModelAccessFacts(("requested",), ("allowed",), object_type="team") + + +@pytest.mark.parametrize("policy", (reference, member_models)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadMembership")), + (replace(_BASE, user_present=False), Decision("Allow")), + (replace(_BASE, team_present=False), Decision("Allow")), + (replace(_BASE, membership_loaded=True), Decision("ReadGrants")), + (replace(_GRANTS, grants=()), Decision("Allow")), + (_GRANTS, Decision("ReadScope")), + (replace(_GRANTS, scope=replace(_SCOPE, grants=())), Decision("Allow")), + ( + replace(_GRANTS, scope=_SCOPE), + Decision("Deny", ModelAccessDecision("Deny", "requested", ("allowed",), "team"), ("allowed",)), + ), + ), +) +def test_member_scope_load_order_and_denial_envelope( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_resolved_models.py b/tests/unit/_auth/test_resolved_models.py new file mode 100644 index 00000000000..7b301e9b3be --- /dev/null +++ b/tests/unit/_auth/test_resolved_models.py @@ -0,0 +1,156 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm._auth.resolved_models import ResolvedModelDecision as Decision +from litellm._auth.resolved_models import ResolvedModelFacts as Facts +from litellm._auth.resolved_models import ResolvedTeamRead, resolved_models +from litellm._auth.team_override import TeamOverrideFacts +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_member_models import reference as member_reference +from tests.unit._auth.test_personal_model_request import reference as personal_reference +from tests.unit._auth.test_team_override import reference as team_reference + + +def reference(facts: Facts) -> Decision: + if not facts.config_nonempty and "all-team-models" not in facts.key_models: + if facts.key is None: + return Decision("ReadKey") + key: Final = grouped(facts.key, merge=False) + if key.kind == "ReadGroups": + return Decision("ReadKeyGroups") + if key.kind != "Allow": + return Decision("Deny", key) + if facts.team_id is not None: + if facts.team_read is None: + return Decision("ReadTeamObject") + if facts.team_read.failed and not facts.embedded_ready: + return Decision("ReadEmbeddedTeam") + if facts.team_read.found or facts.team_read.failed: + if facts.team is None: + return Decision("ReadTeamScope") + team: Final = team_reference(facts.team) + if team.kind == "Deny": + return Decision("Deny", team.denial) + if facts.team_read.found and facts.user_id is not None: + if facts.member is None: + return Decision("ReadMember") + member: Final = member_reference(facts.member) + if member.kind == "Deny": + return Decision("Deny", member.denial, member.member_models) + if ( + facts.personal_enabled + and facts.team_id is None + and facts.user_id is not None + and not facts.master + and facts.role != "proxy_admin" + ): + if not facts.owner_loaded: + return Decision("ReadOwner") + if facts.owner_failed or not facts.owner_present: + return Decision("Deny", ModelAccessDecision("PersonalTeam", model=facts.model)) + if facts.personal is None: + return Decision("ReadPersonal") + personal: Final = personal_reference(facts.personal) + if personal.kind == "Deny": + return Decision("Deny", ModelAccessDecision("PersonalTeam", model=personal.model)) + if facts.project_present: + if not facts.project_loaded: + return Decision("ReadProjectObject") + if facts.project_grants: + if facts.project is None: + return Decision("ReadProjectScope") + project: Final = model_access(facts.project) + if project.kind != "Allow": + return Decision("Deny", project) + if facts.customer_present and facts.database_present: + if not facts.customer_loaded: + return Decision("ReadCustomerObject") + if facts.customer_grants: + if facts.customer is None: + return Decision("ReadCustomerScope") + customer: Final = model_access(facts.customer) + if customer.kind != "Allow": + return Decision("Deny", customer) + return Decision("Allow") + + +_MODEL: Final = ModelAccessFacts(("model",), ("model",)) +_PASS: Final = Facts(model="model", config_nonempty=True) +_TEAM: Final = replace( + _PASS, + team_id="team", + team_read=ResolvedTeamRead(found=True), + team=TeamOverrideFacts(GroupedModelScope(replace(_MODEL, object_type="team"))), +) +_PERSONAL: Final = replace(_PASS, user_id="owner", personal_enabled=True) +_PROJECT: Final = replace(_PASS, project_present=True, project_loaded=True, project_grants=("model",)) +_CUSTOMER: Final = replace( + _PASS, customer_present=True, database_present=True, customer_loaded=True, customer_grants=("model",) +) + + +@pytest.mark.parametrize("policy", (reference, resolved_models)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (Facts(), Decision("ReadKey")), + (_PASS, Decision("Allow")), + (Facts(key_models=("all-team-models",)), Decision("Allow")), + (Facts(key=GroupedModelScope(_MODEL)), Decision("Allow")), + ( + Facts(key=GroupedModelScope(replace(_MODEL, grants=("other",)), ("group",)), project_present=True), + Decision("ReadKeyGroups"), + ), + (replace(_TEAM, team_read=None), Decision("ReadTeamObject")), + (replace(_TEAM, team_read=ResolvedTeamRead(failed=True)), Decision("ReadEmbeddedTeam")), + (replace(_TEAM, team=None), Decision("ReadTeamScope")), + (_TEAM, Decision("Allow")), + (replace(_TEAM, user_id="member", project_present=True), Decision("ReadMember")), + ( + replace(_TEAM, user_id="member", member=MemberModelFacts(True, False), project_present=True), + Decision("ReadProjectObject"), + ), + ( + replace( + _TEAM, + user_id="member", + team_read=ResolvedTeamRead(failed=True), + embedded_ready=True, + project_present=True, + ), + Decision("ReadProjectObject"), + ), + (_PERSONAL, Decision("ReadOwner")), + (replace(_PERSONAL, master=True), Decision("Allow")), + (replace(_PERSONAL, role="proxy_admin"), Decision("Allow")), + (replace(_PERSONAL, owner_loaded=True), Decision("Deny", ModelAccessDecision("PersonalTeam", model="model"))), + ( + replace(_PERSONAL, owner_loaded=True, owner_present=True, owner_failed=True), + Decision("Deny", ModelAccessDecision("PersonalTeam", model="model")), + ), + (replace(_PERSONAL, owner_loaded=True, owner_present=True), Decision("ReadPersonal")), + (replace(_PROJECT, project_loaded=False), Decision("ReadProjectObject")), + (_PROJECT, Decision("ReadProjectScope")), + (replace(_PROJECT, project_grants=()), Decision("Allow")), + ( + replace( + _PROJECT, project=replace(_MODEL, object_type="project"), customer_present=True, database_present=True + ), + Decision("ReadCustomerObject"), + ), + (replace(_CUSTOMER, database_present=False), Decision("Allow")), + (replace(_CUSTOMER, customer_loaded=False), Decision("ReadCustomerObject")), + (_CUSTOMER, Decision("ReadCustomerScope")), + (replace(_CUSTOMER, customer_grants=()), Decision("Allow")), + (replace(_CUSTOMER, customer=replace(_MODEL, object_type="customer")), Decision("Allow")), + ), +) +def test_resolved_scope_and_lookup_order(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected