From 342848ec81461038621c751b69a0eb7c5ce2b02e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 00:11:02 +0000 Subject: [PATCH] Stage managed model scopes and team alias resolution in native auth --- litellm-rust/crates/python-bridge/src/lib.rs | 42 +- .../crates/python-bridge/src/request_auth.rs | 26 ++ .../crates/request-auth-core/src/lib.rs | 2 + .../request-auth-core/src/managed_models.rs | 250 ++++++++++++ .../request-auth-core/src/team_aliases.rs | 77 ++++ litellm/_auth/managed_models.py | 95 +++++ litellm/_auth/team_aliases.py | 34 ++ litellm/proxy/auth/managed_model_adapter.py | 365 ++++++++++++++++++ litellm/proxy/auth/team_alias_adapter.py | 29 ++ litellm/rust_bridge/_native.pyi | 4 + tests/unit/_auth/test_managed_models.py | 116 ++++++ tests/unit/_auth/test_team_aliases.py | 49 +++ .../proxy/auth/test_managed_model_adapter.py | 82 ++++ 13 files changed, 1151 insertions(+), 20 deletions(-) create mode 100644 litellm-rust/crates/request-auth-core/src/managed_models.rs create mode 100644 litellm-rust/crates/request-auth-core/src/team_aliases.rs create mode 100644 litellm/_auth/managed_models.py create mode 100644 litellm/_auth/team_aliases.py create mode 100644 litellm/proxy/auth/managed_model_adapter.py create mode 100644 litellm/proxy/auth/team_alias_adapter.py create mode 100644 tests/unit/_auth/test_managed_models.py create mode 100644 tests/unit/_auth/test_team_aliases.py create mode 100644 tests/unit/proxy/auth/test_managed_model_adapter.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index d4b402f57e8..3168734c3a4 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -40,30 +40,32 @@ mod _native { #[pymodule_export] use crate::request_auth::{ NativeAuthContinuation, request_auth_access, request_auth_admission, - request_auth_agent_access, request_auth_agent_ceiling, request_auth_agent_grants, - request_auth_agent_group_load, request_auth_agent_invocation, request_auth_budget_alert, - request_auth_budget_alert_plan, request_auth_budget_exemption, request_auth_budget_price, - request_auth_budget_route, request_auth_budget_window_start, request_auth_budgets, - request_auth_common_budget, request_auth_common_budget_scopes, request_auth_common_flow, - request_auth_common_read, request_auth_context, request_auth_context_end_user, - request_auth_context_member, request_auth_context_models, request_auth_context_percentage, - request_auth_context_temporary, request_auth_context_throttle, request_auth_context_token, - request_auth_custom_flow, request_auth_end_user_apply, request_auth_end_user_assignments, + request_auth_agent_access, request_auth_agent_ceiling, request_auth_agent_ceiling_source, + request_auth_agent_grants, request_auth_agent_group_load, request_auth_agent_invocation, + request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption, + request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start, + request_auth_budgets, request_auth_common_budget, request_auth_common_budget_scopes, + request_auth_common_flow, request_auth_common_read, request_auth_context, + request_auth_context_end_user, request_auth_context_member, request_auth_context_models, + request_auth_context_percentage, request_auth_context_temporary, + request_auth_context_throttle, request_auth_context_token, request_auth_custom_flow, + request_auth_end_user_apply, request_auth_end_user_assignments, request_auth_end_user_default_id, request_auth_end_user_lookup, request_auth_end_user_normalize_id, request_auth_end_user_registry, request_auth_end_user_validation, 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_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_listing, - 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_spend_counter_policy, - request_auth_start, request_auth_team_override, request_auth_unpriced_model, + request_auth_litellm_header_token, request_auth_managed, request_auth_managed_model_grants, + request_auth_managed_models, request_auth_member_limit, 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_listing, 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_spend_counter_policy, request_auth_start, request_auth_team_aliases, + 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 f14a5ea5ab7..95b8a7b65a2 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -784,3 +784,29 @@ pub(crate) fn request_auth_agent_group_load( let facts: core::agent_ceilings::LoadFacts = from_py(facts)?; to_py(py, &py.detach(move || core::agent_ceilings::load(&facts))) } + +#[pyfunction] +pub(crate) fn request_auth_team_aliases(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::team_aliases::Facts = from_py(facts)?; + let decision = py.detach(move || core::team_aliases::decide(&facts)).map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_managed_models(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::managed_models::Facts = from_py(facts)?; + let decision = py.detach(move || core::managed_models::decide(&facts)).map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_ceiling_source(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::managed_models::CeilingSource = from_py(facts)?; + to_py(py, &py.detach(move || core::managed_models::ceiling_source(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_managed_model_grants(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::managed_models::GrantFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::managed_models::grants(&facts))) +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index f8bdd9bc834..4382273bff0 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -19,6 +19,7 @@ pub mod fallback_candidate; mod key; pub mod key_model_access; pub mod managed; +pub mod managed_models; pub mod member_models; pub mod model_access; pub mod model_attribution; @@ -39,6 +40,7 @@ pub mod reservation; pub mod resolved_models; pub mod route_access; pub mod spend_counter_policy; +pub mod team_aliases; pub mod team_override; pub mod unpriced_models; diff --git a/litellm-rust/crates/request-auth-core/src/managed_models.rs b/litellm-rust/crates/request-auth-core/src/managed_models.rs new file mode 100644 index 00000000000..bfd8358ec23 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/managed_models.rs @@ -0,0 +1,250 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + member_models, + model_access::{self, Decision as ModelDecision}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Default, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Purpose { + #[default] + All, + Ceilings, + Caller, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub purpose: Purpose, + pub model_present: bool, + pub token_present: bool, + pub managed_loaded: bool, + pub managed_present: bool, + pub managed_models_valid: bool, + pub managed_models_count: usize, + pub managed: Option, + pub ceilings: Option>>, + pub aliases_loaded: bool, + pub scopes: Vec, + pub caller_loaded: bool, + pub caller_present: bool, + pub team_loaded: bool, + pub team_present: bool, + pub team: Option, + pub member: Option, + pub user_loaded: bool, + pub user_present: bool, + pub user_grants: Vec, + pub user: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: ModelDecision, + member_models: Option>, + }, + ReadManaged, + ReadManagedScope, + ReadCeilings, + ReadAliases, + ReadCeilingScope { + index: usize, + }, + ReadCaller, + ReadTeam, + ReadTeamScope, + ReadTeamGroups, + ReadMember, + ReadUser, + ReadUserScope, +} + +pub fn decide(facts: &Facts) -> Result { + let deny = |denial| { + Ok(Decision::Deny { + denial, + member_models: None, + }) + }; + if facts.purpose == Purpose::All && !facts.managed_loaded { + return Ok(Decision::ReadManaged); + } + if !facts.model_present || !facts.token_present { + return Ok(Decision::Allow); + } + if facts.purpose == Purpose::All && facts.managed_present { + if !facts.managed_models_valid || facts.managed_models_count == 0 { + return deny(ModelDecision::NoAgentModels); + } + let Some(scope) = &facts.managed else { + return Ok(Decision::ReadManagedScope); + }; + let result = model_access::check(scope)?; + if !result.allowed() { + return deny(result); + } + } + if facts.purpose != Purpose::Caller { + let Some(ceilings) = &facts.ceilings else { + return Ok(Decision::ReadCeilings); + }; + if !facts.aliases_loaded { + return Ok(Decision::ReadAliases); + } + for (index, models) in ceilings.iter().enumerate() { + if models.is_empty() { + return deny(ModelDecision::EmptyAgentCeiling); + } + let Some(scope) = facts.scopes.get(index) else { + return Ok(Decision::ReadCeilingScope { index }); + }; + let result = model_access::check(scope)?; + if !result.allowed() { + return deny(result); + } + } + if facts.purpose == Purpose::Ceilings { + return Ok(Decision::Allow); + } + } + if !facts.caller_loaded { + return Ok(Decision::ReadCaller); + } + if !facts.caller_present { + return Ok(Decision::Allow); + } + if !facts.team_loaded { + return Ok(Decision::ReadTeam); + } + if facts.team_present { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeamScope); + }; + let result = fallback_access::grouped(team.clone(), true)?; + if matches!(result, ModelDecision::ReadGroups) { + return Ok(Decision::ReadTeamGroups); + } + if !result.allowed() { + return deny(result); + } + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember); + }; + return match member_models::decide(member)? { + member_models::Decision::Allow => Ok(Decision::Allow), + member_models::Decision::Deny { + denial, + member_models, + } => Ok(Decision::Deny { + denial, + member_models, + }), + _ => Err(Error::InvalidAnswer), + }; + } + if !facts.user_loaded { + return Ok(Decision::ReadUser); + } + let precheck = model_access::decide(model_access::Request::UserPrecheck { + present: facts.user_present, + grants: facts.user_grants.clone(), + })?; + let model_access::Response::Decision(precheck) = precheck else { + return Err(Error::InvalidAnswer); + }; + if !matches!(precheck, ModelDecision::ReadScope) { + return if precheck.allowed() { + Ok(Decision::Allow) + } else { + deny(precheck) + }; + } + let Some(user) = &facts.user else { + return Ok(Decision::ReadUserScope); + }; + let result = model_access::check(user)?; + if result.allowed() { + Ok(Decision::Allow) + } else { + deny(result) + } +} + +#[derive(Clone, Default, Deserialize)] +pub struct CeilingSource { + pub agent_present: bool, + pub managed_loaded: bool, + pub managed_present: bool, + pub loaded: bool, + pub count: usize, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum SourceDecision { + ReadManaged, + ReadManagedCeilings, + ReadUnmanagedCeiling, + Complete { indices: Vec }, +} + +pub fn ceiling_source(facts: &CeilingSource) -> SourceDecision { + if !facts.agent_present { + return SourceDecision::Complete { indices: vec![] }; + } + if !facts.managed_loaded { + return SourceDecision::ReadManaged; + } + if facts.loaded { + return SourceDecision::Complete { + indices: (0..facts.count).collect(), + }; + } + if facts.managed_present { + SourceDecision::ReadManagedCeilings + } else { + SourceDecision::ReadUnmanagedCeiling + } +} + +#[derive(Clone, Deserialize)] +pub struct GrantFacts { + pub models: Option>>, + pub ceilings: Vec>, +} + +pub fn grants(facts: &GrantFacts) -> std::collections::BTreeSet { + facts + .models + .iter() + .flatten() + .flatten() + .chain(facts.ceilings.iter().flatten()) + .cloned() + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::main(Purpose::All, Decision::ReadManaged)] + #[case::groups(Purpose::Ceilings, Decision::Allow)] + #[case::caller(Purpose::Caller, Decision::Allow)] + fn empty_model_read_order(#[case] purpose: Purpose, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + purpose, + ..Facts::default() + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_aliases.rs b/litellm-rust/crates/request-auth-core/src/team_aliases.rs new file mode 100644 index 00000000000..6abd1e477b8 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/team_aliases.rs @@ -0,0 +1,77 @@ +use crate::Error; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub aliases: BTreeMap, + pub team_id: Option, + pub router_present: bool, + pub names: Option>>, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadNames { count: usize }, + Complete { models: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let prefix = format!("model_name_{}_", facts.team_id.as_deref().unwrap_or("None")); + let needs_names = |target: &str| facts.router_present && target.starts_with(&prefix); + let count = facts + .requested + .iter() + .filter_map(|model| facts.aliases.get(model)) + .filter(|target| needs_names(target)) + .count(); + let names = match &facts.names { + None if count != 0 => return Ok(Decision::ReadNames { count }), + Some(names) if names.len() != count => return Err(Error::InvalidAnswer), + names => names.as_deref().unwrap_or_default(), + }; + let mut reads = names.iter(); + let models = facts + .requested + .iter() + .map(|model| { + let Some(target) = facts.aliases.get(model) else { + return model.clone(); + }; + if needs_names(target) && !reads.next().is_some_and(|names| names.contains(target)) { + model.clone() + } else { + target.clone() + } + }) + .collect(); + Ok(Decision::Complete { models }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(false, "alias")] + #[case::live(true, "model_name_team_target")] + fn deleted_target(#[case] live: bool, #[case] expected: &str) { + let target = "model_name_team_target".to_owned(); + let facts = Facts { + requested: vec!["alias".into()], + aliases: BTreeMap::from([("alias".into(), target.clone())]), + team_id: Some("team".into()), + router_present: true, + names: Some(vec![if live { vec![target] } else { vec![] }]), + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Complete { + models: vec![expected.into()] + } + ); + } +} diff --git a/litellm/_auth/managed_models.py b/litellm/_auth/managed_models.py new file mode 100644 index 00000000000..2a41347d329 --- /dev/null +++ b/litellm/_auth/managed_models.py @@ -0,0 +1,95 @@ +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.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ManagedModelFacts: + purpose: Literal["All", "Ceilings", "Caller"] + model_present: bool + token_present: bool + managed_loaded: bool = False + managed_present: bool = False + managed_models_valid: bool = False + managed_models_count: int = 0 + managed: ModelAccessFacts | None = None + ceilings: tuple[tuple[str, ...], ...] | None = None + aliases_loaded: bool = False + scopes: tuple[ModelAccessFacts, ...] = () + caller_loaded: bool = False + caller_present: bool = False + team_loaded: bool = False + team_present: bool = False + team: GroupedModelScope | None = None + member: MemberModelFacts | None = None + user_loaded: bool = False + user_present: bool = False + user_grants: tuple[str, ...] = () + user: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ManagedModelDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadManaged", + "ReadManagedScope", + "ReadCeilings", + "ReadAliases", + "ReadCeilingScope", + "ReadCaller", + "ReadTeam", + "ReadTeamScope", + "ReadTeamGroups", + "ReadMember", + "ReadUser", + "ReadUserScope", + ] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + index: int = 0 + + +@dataclass(frozen=True, slots=True) +class AgentCeilingSourceFacts: + agent_present: bool + managed_loaded: bool = False + managed_present: bool = False + loaded: bool = False + count: int = 0 + + +@dataclass(frozen=True, slots=True) +class AgentCeilingSourceDecision: + kind: Literal["ReadManaged", "ReadManagedCeilings", "ReadUnmanagedCeiling", "Complete"] + indices: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(ManagedModelDecision) +_CEILINGS: Final = TypeAdapter(AgentCeilingSourceDecision) +_GRANTS: Final = TypeAdapter(frozenset[str]) + + +@native_boundary +def managed_models(facts: ManagedModelFacts) -> ManagedModelDecision: + return _DECISION.validate_python(_native.request_auth_managed_models(asdict(facts))) + + +@native_boundary +def agent_ceiling_source(facts: AgentCeilingSourceFacts) -> AgentCeilingSourceDecision: + return _CEILINGS.validate_python(_native.request_auth_agent_ceiling_source(asdict(facts))) + + +@native_boundary +def managed_model_grants( + models: tuple[str | None, ...] | None, ceilings: tuple[tuple[str, ...], ...] +) -> frozenset[str]: + return _GRANTS.validate_python(_native.request_auth_managed_model_grants({"models": models, "ceilings": ceilings})) diff --git a/litellm/_auth/team_aliases.py b/litellm/_auth/team_aliases.py new file mode 100644 index 00000000000..c72ff652111 --- /dev/null +++ b/litellm/_auth/team_aliases.py @@ -0,0 +1,34 @@ +from collections.abc import Mapping +from dataclasses import asdict, dataclass, replace +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class TeamAliasFacts: + requested: tuple[str, ...] + aliases: Mapping[str, str] + team_id: str | None + router_present: bool + names: tuple[tuple[str, ...], ...] | None = None + + +@dataclass(frozen=True, slots=True) +class TeamAliasDecision: + kind: Literal["ReadNames", "Complete"] + count: int = 0 + models: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(TeamAliasDecision) + + +@native_boundary +def team_aliases(facts: TeamAliasFacts) -> TeamAliasDecision: + return _DECISION.validate_python( + _native.request_auth_team_aliases(asdict(replace(facts, aliases=dict(facts.aliases)))) + ) diff --git a/litellm/proxy/auth/managed_model_adapter.py b/litellm/proxy/auth/managed_model_adapter.py new file mode 100644 index 00000000000..f3dfd69b39f --- /dev/null +++ b/litellm/proxy/auth/managed_model_adapter.py @@ -0,0 +1,365 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.managed_models import ( + AgentCeilingSourceFacts, + ManagedModelFacts, + agent_ceiling_source, + managed_model_grants, + managed_models, +) +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, + resolve_agent_access_group_ceiling, + resolve_managed_agent_ceilings, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import ( + agent_caller_auth, + load_agent_caller_team, + load_agent_caller_user, +) +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy +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.team_alias_adapter import resolve_team_alias +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader +from litellm.types.agents import AgentResponse + +CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] +CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] +CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_UserTable | None]] +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] | None + token: UserAPIKeyAuth | None + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + ceiling: CeilingResolver + team_loader: CallerTeamLoader + user_loader: CallerUserLoader + managed_grants: tuple[object, ...] = () + dispatched: str | Sequence[str] = () + caller: UserAPIKeyAuth | None = None + team: LiteLLM_TeamTable | None = None + user: LiteLLM_UserTable | None = None + key_aliases: Mapping[str, str] | None = None + + +async def agent_access_group_ceilings( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> tuple[AgentAccessGroupCeiling, ...]: + return await _ceilings(AgentCeilingSourceFacts(bool(valid_token.agent_id)), valid_token, resolve_ceiling) + + +async def _ceilings( + facts: AgentCeilingSourceFacts, + token: UserAPIKeyAuth, + resolver: CeilingResolver, + managed: AgentResponse | None = None, + rows: tuple[AgentAccessGroupCeiling, ...] = (), +) -> tuple[AgentAccessGroupCeiling, ...]: + decision: Final = agent_ceiling_source(facts) + match decision.kind: + case "Complete": + return tuple(rows[index] for index in decision.indices) + case "ReadManaged": + policy: Final = managed_agent_policy(token) + return await _ceilings( + replace(facts, managed_loaded=True, managed_present=policy is not None), token, resolver, policy + ) + case "ReadManagedCeilings": + if managed is None: + raise NativeAuthRuntimeError("missing managed agent ceiling source") + loaded: Final = await resolve_managed_agent_ceilings(managed) + return await _ceilings(replace(facts, loaded=True, count=len(loaded)), token, resolver, managed, loaded) + case "ReadUnmanagedCeiling": + if token.agent_id is None: + raise NativeAuthRuntimeError("missing agent ceiling identity") + row: Final = await resolver(token.agent_id) + projected: Final = (row,) if row is not None else () + return await _ceilings( + replace(facts, loaded=True, count=len(projected)), token, resolver, managed, projected + ) + + +async def read_managed_model_facts( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + purpose: Literal["All", "Ceilings", "Caller"] = "All", + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> ManagedModelFacts: + facts, _source = await _read( + ManagedModelFacts(purpose, bool(model), valid_token is not None), + _Source( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + resolve_ceiling, + load_team, + load_user, + ), + ) + return facts + + +async def can_agent_call_model( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + purpose: Literal["All", "Ceilings", "Caller"] = "All", + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + facts, source = await _read( + ManagedModelFacts(purpose, bool(model), valid_token is not None), + _Source( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + resolve_ceiling, + load_team, + load_user, + ), + ) + decision: Final = managed_models(facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None or model is None: + raise NativeAuthRuntimeError("incomplete managed model authorization") + raise_model_denial( + decision.denial, + model, + user_id=source.caller.user_id + if decision.member_models is not None and source.caller is not None + else valid_token.agent_id + if valid_token is not None + else None, + team_id=source.caller.team_id if source.caller is not None else None, + member_models=decision.member_models, + ) + + +@native_boundary +def _managed_grants(raw: tuple[object, ...]) -> tuple[str, ...]: + return _NAMES.validate_python(raw) + + +def _declared_grants(token: UserAPIKeyAuth) -> tuple[str | None, ...] | None: + managed: Final = managed_agent_policy(token) + raw: Final = (managed.object_permission or {}).get("models") if managed is not None else None + if not isinstance(raw, (list, tuple)): + return None + return tuple(value if isinstance(value, str) else None for value in _ITEMS.validate_python(raw)) + + +def managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: + return managed_model_grants(_declared_grants(valid_token), ()) + + +async def agent_model_grants( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> frozenset[str]: + ceilings: Final = await agent_access_group_ceilings(valid_token, resolve_ceiling) + return managed_model_grants(_declared_grants(valid_token), tuple(tuple(ceiling.models) for ceiling in ceilings)) + + +async def check_agent_access_group_model_access( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> Literal[True]: + await can_agent_call_model( + model, valid_token, llm_router, None, None, None, purpose="Ceilings", resolve_ceiling=resolve_ceiling + ) + return True + + +async def check_agent_caller_model_access( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + await can_agent_call_model( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + purpose="Caller", + load_team=load_team, + load_user=load_user, + ) + + +def _model(source: _Source) -> str | Sequence[str]: + if source.model is None: + raise NativeAuthRuntimeError("managed model read without a request model") + return source.model + + +def _token(source: _Source) -> UserAPIKeyAuth: + if source.token is None: + raise NativeAuthRuntimeError("managed model read without a token") + return source.token + + +def _dispatch(source: _Source) -> str | Sequence[str]: + from litellm.proxy.auth.auth_checks import team_model_aliases_for_auth_check + + token: Final = _token(source) + return resolve_team_alias(_model(source), team_model_aliases_for_auth_check(token), token.team_id, source.router) + + +async def _read(facts: ManagedModelFacts, source: _Source) -> tuple[ManagedModelFacts, _Source]: + from litellm.proxy.auth.auth_checks import get_models_from_access_groups, key_model_aliases_for_auth_check + + decision: Final = managed_models(facts) + match decision.kind: + case "Allow" | "Deny": + return facts, source + case "ReadManaged": + policy: Final = managed_agent_policy(source.token) + raw: Final = (policy.object_permission or {}).get("models", ()) if policy is not None else () + items: Final = _ITEMS.validate_python(raw) if isinstance(raw, (list, tuple)) else () + return await _read( + replace( + facts, + managed_loaded=True, + managed_present=policy is not None, + managed_models_valid=isinstance(raw, (list, tuple)), + managed_models_count=len(items), + ), + replace(source, managed_grants=items), + ) + case "ReadManagedScope": + dispatched: Final = _dispatch(source) + managed: Final = project_model_facts( + dispatched, + source.router, + _managed_grants(source.managed_grants), + team_id=_token(source).team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="agent", + ) + return await _read(replace(facts, managed=managed), source) + case "ReadCeilings": + ceilings: Final = await agent_access_group_ceilings(_token(source), source.ceiling) + return await _read( + replace(facts, ceilings=tuple(tuple(sorted(ceiling.models)) for ceiling in ceilings)), source + ) + case "ReadAliases": + return await _read(replace(facts, aliases_loaded=True), replace(source, dispatched=_dispatch(source))) + case "ReadCeilingScope": + if facts.ceilings is None: + raise NativeAuthRuntimeError("managed model scope without ceilings") + ceiling_scope: Final = project_model_facts( + source.dispatched, + source.router, + facts.ceilings[decision.index], + team_id=_token(source).team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="agent", + ) + return await _read(replace(facts, scopes=(*facts.scopes, ceiling_scope)), source) + case "ReadCaller": + caller: Final = agent_caller_auth(_token(source)) + return await _read( + replace(facts, caller_loaded=True, caller_present=caller is not None), replace(source, caller=caller) + ) + case "ReadTeam": + team: Final = await source.team_loader(_token(source)) + aliases: Final = key_model_aliases_for_auth_check(source.token) + return await _read( + replace(facts, team_loaded=True, team_present=team is not None), + replace(source, team=team, key_aliases=aliases), + ) + case "ReadTeamScope": + if source.team is None: + raise NativeAuthRuntimeError("missing caller team for model scope") + scope: Final = project_model_facts( + _model(source), + source.router, + object_model_grants(source.team), + team_id=source.team.team_id, + key_model_aliases=source.key_aliases, + object_type="team", + ) + return await _read( + replace(facts, team=GroupedModelScope(scope, tuple(source.team.access_group_ids or ()))), source + ) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing caller team group scope") + grants: Final = await get_models_from_access_groups(facts.team.group_ids, prisma_client=source.prisma) + return await _read(replace(facts, team=replace(facts.team, group_models=tuple(grants))), source) + case "ReadMember": + if source.team is None or source.caller is None or source.cache is None or source.logging is None: + raise NativeAuthRuntimeError("missing caller member model context") + member: Final = await read_member_model_facts( + _model(source), + source.team, + source.caller, + source.router, + source.prisma, + source.cache, + source.logging, + key_model_aliases=source.key_aliases, + ) + return await _read(replace(facts, member=member), source) + case "ReadUser": + user: Final = await source.user_loader(_token(source)) + return await _read( + replace( + facts, + user_loaded=True, + user_present=user is not None, + user_grants=object_model_grants(user) if user is not None else (), + ), + replace(source, user=user), + ) + case "ReadUserScope": + user_scope: Final = project_model_facts( + _model(source), source.router, facts.user_grants, key_model_aliases=source.key_aliases + ) + return await _read(replace(facts, user=user_scope), source) diff --git a/litellm/proxy/auth/team_alias_adapter.py b/litellm/proxy/auth/team_alias_adapter.py new file mode 100644 index 00000000000..551cf91ff5e --- /dev/null +++ b/litellm/proxy/auth/team_alias_adapter.py @@ -0,0 +1,29 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.team_aliases import TeamAliasFacts, team_aliases +from litellm.router_api import RegistryReader + + +@native_boundary +def resolve_team_alias( + model: str | Sequence[str], + team_model_aliases: Mapping[str, str] | None, + team_id: str | None, + llm_router: RegistryReader | None, +) -> str | list[str]: # mutable-ok: the existing public wrapper preserves the input scalar/list shape + initial: Final = TeamAliasFacts( + (model,) if isinstance(model, str) else tuple(model), team_model_aliases or {}, team_id, llm_router is not None + ) + decision: Final = team_aliases(initial) + if decision.kind == "Complete": + return decision.models[0] if isinstance(model, str) else list(decision.models) + if llm_router is None: + raise NativeAuthRuntimeError("team aliases requested names without a router") + names: Final = tuple(tuple(llm_router.registry().index("model_name")) for _ in range(decision.count)) + completed: Final = team_aliases(replace(initial, names=names)) + if completed.kind != "Complete": + raise NativeAuthRuntimeError("incomplete team alias resolution") + return completed.models[0] if isinstance(model, str) else list(completed.models) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 9293ab7d5bd..94f0b2b0837 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -71,6 +71,10 @@ def request_auth_resolved_models(facts: object) -> object: ... def request_auth_personal_model_listing(facts: object) -> object: ... def request_auth_agent_ceiling(facts: object) -> object: ... def request_auth_agent_group_load(facts: object) -> object: ... +def request_auth_team_aliases(facts: object) -> object: ... +def request_auth_managed_models(facts: object) -> object: ... +def request_auth_agent_ceiling_source(facts: object) -> object: ... +def request_auth_managed_model_grants(facts: object) -> object: ... @final class NativeClickHouseSpendConfig: diff --git a/tests/unit/_auth/test_managed_models.py b/tests/unit/_auth/test_managed_models.py new file mode 100644 index 00000000000..167ab0da702 --- /dev/null +++ b/tests/unit/_auth/test_managed_models.py @@ -0,0 +1,116 @@ +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.managed_models import ManagedModelDecision as Decision +from litellm._auth.managed_models import ManagedModelFacts as Facts +from litellm._auth.managed_models import managed_models +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +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 + + +def reference(facts: Facts) -> Decision: + if facts.purpose == "All" and not facts.managed_loaded: + return Decision("ReadManaged") + if not facts.model_present or not facts.token_present: + return Decision("Allow") + if facts.purpose == "All" and facts.managed_present: + if not facts.managed_models_valid or not facts.managed_models_count: + return Decision("Deny", ModelAccessDecision("NoAgentModels")) + if facts.managed is None: + return Decision("ReadManagedScope") + managed: Final = model_access(facts.managed) + if managed.kind != "Allow": + return Decision("Deny", managed) + if facts.purpose != "Caller": + if facts.ceilings is None: + return Decision("ReadCeilings") + if not facts.aliases_loaded: + return Decision("ReadAliases") + for index, ceiling in enumerate(facts.ceilings): + if not ceiling: + return Decision("Deny", ModelAccessDecision("EmptyAgentCeiling")) + if index >= len(facts.scopes): + return Decision("ReadCeilingScope", index=index) + if (scope := model_access(facts.scopes[index])).kind != "Allow": + return Decision("Deny", scope) + if facts.purpose == "Ceilings": + return Decision("Allow") + if not facts.caller_loaded: + return Decision("ReadCaller") + if not facts.caller_present: + return Decision("Allow") + if not facts.team_loaded: + return Decision("ReadTeam") + if facts.team_present: + if facts.team is None: + return Decision("ReadTeamScope") + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if team.kind != "Allow": + return Decision("Deny", team) + if facts.member is None: + return Decision("ReadMember") + member: Final = member_reference(facts.member) + return Decision("Allow") if member.kind == "Allow" else Decision("Deny", member.denial, member.member_models) + if not facts.user_loaded: + return Decision("ReadUser") + if not facts.user_present or not facts.user_grants: + return Decision("Allow") + if "no-default-models" in facts.user_grants: + return Decision("Deny", ModelAccessDecision("NoDefaultModels")) + if facts.user is None: + return Decision("ReadUserScope") + user: Final = model_access(facts.user) + return Decision("Allow") if user.kind == "Allow" else Decision("Deny", user) + + +_ALL: Final = Facts("All", True, True) +_GROUPS: Final = Facts("Ceilings", True, True) +_CALLER: Final = Facts("Caller", True, True, caller_loaded=True, caller_present=True, team_loaded=True) +_ALLOW: Final = ModelAccessFacts(("model",), ("model",), object_type="agent") +_DENY: Final = replace(_ALLOW, grants=("other",)) +_TEAM: Final = GroupedModelScope(replace(_ALLOW, object_type="team")) + + +@pytest.mark.parametrize("policy", (reference, managed_models)) +@pytest.mark.parametrize( + "facts", + ( + _ALL, + replace(_ALL, model_present=False), + replace(_ALL, managed_loaded=True, model_present=False), + replace(_ALL, managed_loaded=True, managed_present=True), + replace(_ALL, managed_loaded=True, managed_present=True, managed_models_valid=True, managed_models_count=1), + replace( + _ALL, + managed_loaded=True, + managed_present=True, + managed_models_valid=True, + managed_models_count=1, + managed=_DENY, + ), + _GROUPS, + replace(_GROUPS, ceilings=()), + replace(_GROUPS, ceilings=(), aliases_loaded=True), + replace(_GROUPS, ceilings=((),), aliases_loaded=True), + replace(_GROUPS, ceilings=(("model",), ("other",)), aliases_loaded=True, scopes=(_ALLOW,)), + replace(_GROUPS, ceilings=(("model",), ("other",)), aliases_loaded=True, scopes=(_ALLOW, _DENY)), + _CALLER, + replace(_CALLER, team_present=True), + replace(_CALLER, team_present=True, team=replace(_TEAM, group_ids=("g",))), + replace(_CALLER, team_present=True, team=_TEAM), + replace(_CALLER, team_present=True, team=_TEAM, member=MemberModelFacts(True, True, True, ())), + replace(_CALLER, user_loaded=True, user_present=True, user_grants=("no-default-models",)), + replace(_CALLER, user_loaded=True, user_present=True, user_grants=("model",)), + ), +) +def test_managed_caller_and_ceiling_read_denial_order(policy: Callable[[Facts], Decision], facts: Facts) -> None: + assert policy(facts) == reference(facts) diff --git a/tests/unit/_auth/test_team_aliases.py b/tests/unit/_auth/test_team_aliases.py new file mode 100644 index 00000000000..b61a34204dc --- /dev/null +++ b/tests/unit/_auth/test_team_aliases.py @@ -0,0 +1,49 @@ +from collections.abc import Callable +from typing import Final + +import pytest + +from litellm._auth.team_aliases import TeamAliasDecision as Decision +from litellm._auth.team_aliases import TeamAliasFacts as Facts +from litellm._auth.team_aliases import team_aliases + + +def reference(facts: Facts) -> Decision: + targets: Final = tuple(facts.aliases.get(name) for name in facts.requested) + count: Final = sum( + bool(target is not None and facts.router_present and target.startswith(f"model_name_{facts.team_id}_")) + for target in targets + ) + if count and facts.names is None: + return Decision("ReadNames", count=count) + reads: Final = iter(facts.names or ()) + return Decision( + "Complete", + models=tuple( + name + if target is None + or ( + facts.router_present and target.startswith(f"model_name_{facts.team_id}_") and target not in next(reads) + ) + else target + for name, target in zip(facts.requested, targets, strict=True) + ), + ) + + +@pytest.mark.parametrize("policy", (reference, team_aliases)) +@pytest.mark.parametrize( + "facts", + ( + Facts(("alias",), {}, "team", True), + Facts(("alias",), {"alias": "target"}, "team", True), + Facts(("alias",), {"alias": "model_name_team_target"}, "team", False), + Facts(("alias",), {"alias": "model_name_team_target"}, "team", True), + Facts(("alias", "alias"), {"alias": "model_name_team_target"}, "team", True, (("model_name_team_target",), ())), + Facts(("alias",), {"alias": "model_name_None_target"}, None, True, ((),)), + ), +) +def test_aliases_preserve_deleted_targets_and_repeated_registry_reads( + policy: Callable[[Facts], Decision], facts: Facts +) -> None: + assert policy(facts) == reference(facts) diff --git a/tests/unit/proxy/auth/test_managed_model_adapter.py b/tests/unit/proxy/auth/test_managed_model_adapter.py new file mode 100644 index 00000000000..f278e2d820c --- /dev/null +++ b/tests/unit/proxy/auth/test_managed_model_adapter.py @@ -0,0 +1,82 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ModelAccessDeniedProxyException, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling +from litellm.proxy.auth import auth_checks, managed_model_adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from litellm.types.agents import AgentResponse + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True)) +@pytest.mark.parametrize("models", (None, "model", (), ("allowed",))) +async def test_managed_model_grants_preserve_public_denial(native: bool, models: object) -> None: + token: Final = UserAPIKeyAuth().model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": models} + ) + } + ) + cache: Final = UserApiKeyCache() + authorize: Final = managed_model_adapter.can_agent_call_model if native else auth_checks.can_agent_call_model + if models == ("allowed",): + with pytest.raises(ModelAccessDeniedProxyException) as denied: + await authorize("denied", token, None, None, cache, ProxyLogging(user_api_key_cache=cache)) + assert ( + denied.value.internal_message + == "agent not allowed to access model. This agent can only access models=['allowed']. Tried to access denied" + ) + else: + with pytest.raises(HTTPException) as denied: + await authorize("denied", token, None, None, cache, ProxyLogging(user_api_key_cache=cache)) + assert (denied.value.status_code, denied.value.detail) == (403, "This agent has no model grants") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True)) +@pytest.mark.parametrize("grants", ((), ("allowed",))) +async def test_agent_empty_ceiling_and_model_denial_envelopes(native: bool, grants: tuple[str, ...]) -> None: + async def resolve(_agent: str) -> AgentAccessGroupCeiling: + return AgentAccessGroupCeiling(("group",), frozenset(grants), frozenset(), frozenset()) + + authorize: Final = ( + managed_model_adapter.check_agent_access_group_model_access + if native + else auth_checks._check_agent_access_group_model_access # pyright: ignore[reportPrivateUsage] # compare the legacy callable contract + ) + with pytest.raises(ModelAccessDeniedProxyException) as denied: + await authorize("denied", UserAPIKeyAuth(agent_id="agent"), None, resolve) + assert denied.value.internal_message == ( + "agent agent access groups grant no models" + if not grants + else "agent not allowed to access model. This agent can only access models=['allowed']. Tried to access denied" + ) + + +@pytest.mark.asyncio +async def test_managed_alias_dispatch_before_scope_check() -> None: + token: Final = UserAPIKeyAuth(team_model_aliases={"alias": "allowed"}).model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": ("allowed",)} + ), + } + ) + assert await managed_model_adapter.can_agent_call_model("alias", token, None, None, None, None) is None + + +@pytest.mark.parametrize("models", (None, "model", (), ("a", 42, "a", "b"))) +def test_managed_listing_filters_non_string_grants(models: object) -> None: + token: Final = UserAPIKeyAuth().model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": models} + ) + } + ) + assert managed_model_adapter.managed_agent_model_grants(token) == auth_checks._managed_agent_model_grants(token) # pyright: ignore[reportPrivateUsage] # compare legacy filtering rather than mocking its policy