From a02a8add640403330c0db76dca23e3d1243a91bf 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:35:34 +0000 Subject: [PATCH] Stage registry reads and optional route enablement in native auth --- litellm-rust/crates/python-bridge/src/lib.rs | 2 + .../crates/python-bridge/src/request_auth.rs | 15 ++ .../request-auth-core/src/fallback_access.rs | 29 ++- .../request-auth-core/src/key_model_access.rs | 10 +- .../crates/request-auth-core/src/lib.rs | 2 + .../request-auth-core/src/managed_models.rs | 5 +- .../request-auth-core/src/model_access.rs | 67 ++++- .../request-auth-core/src/model_helpers.rs | 35 +-- .../request-auth-core/src/model_projection.rs | 237 ++++++++++++++++++ .../request-auth-core/src/personal_models.rs | 5 +- .../request-auth-core/src/resolved_models.rs | 5 +- .../request-auth-core/src/route_enable.rs | 139 ++++++++++ .../request-auth-core/src/team_override.rs | 4 +- litellm/_auth/fallback_access.py | 2 + litellm/_auth/model_access.py | 24 ++ litellm/_auth/model_helpers.py | 1 + litellm/_auth/model_projection.py | 26 ++ litellm/_auth/route_enable.py | 42 ++++ litellm/proxy/auth/fallback_access_adapter.py | 52 ++-- .../proxy/auth/fallback_candidate_adapter.py | 4 +- .../proxy/auth/key_model_access_adapter.py | 12 +- litellm/proxy/auth/managed_model_adapter.py | 7 +- litellm/proxy/auth/model_access_adapter.py | 102 +++----- litellm/proxy/auth/model_helper_adapter.py | 13 +- .../proxy/auth/model_projection_adapter.py | 121 +++++++++ litellm/proxy/auth/personal_model_adapter.py | 7 +- litellm/proxy/auth/resolved_model_adapter.py | 6 +- litellm/proxy/auth/route_enable_adapter.py | 122 +++++++++ litellm/proxy/auth/team_alias_adapter.py | 10 +- litellm/proxy/auth/team_override_adapter.py | 6 +- litellm/rust_bridge/_native.pyi | 2 + .../unit/_auth/support/fallback_reference.py | 9 +- tests/unit/_auth/support/model_reference.py | 2 +- .../_auth/support/model_registry_reference.py | 187 ++++++++++++++ .../_auth/support/route_enable_reference.py | 34 +++ tests/unit/_auth/support/route_reference.py | 97 ++++++- tests/unit/_auth/test_fallback_access.py | 10 + tests/unit/_auth/test_model_access.py | 7 +- tests/unit/_auth/test_route_enable.py | 38 +++ tests/unit/_auth/test_routes.py | 3 +- tests/unit/_auth/test_unpriced_models.py | 6 +- .../proxy/auth/test_managed_model_adapter.py | 30 +-- .../auth/test_model_projection_adapter.py | 203 +++++++++++++++ .../proxy/auth/test_route_enable_adapter.py | 88 +++++++ 44 files changed, 1651 insertions(+), 177 deletions(-) create mode 100644 litellm-rust/crates/request-auth-core/src/model_projection.rs create mode 100644 litellm-rust/crates/request-auth-core/src/route_enable.rs create mode 100644 litellm/_auth/model_projection.py create mode 100644 litellm/_auth/route_enable.py create mode 100644 litellm/proxy/auth/model_projection_adapter.py create mode 100644 litellm/proxy/auth/route_enable_adapter.py create mode 100644 tests/unit/_auth/support/model_registry_reference.py create mode 100644 tests/unit/_auth/support/route_enable_reference.py create mode 100644 tests/unit/_auth/test_route_enable.py create mode 100644 tests/unit/proxy/auth/test_model_projection_adapter.py create mode 100644 tests/unit/proxy/auth/test_route_enable_adapter.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 88d61f8dad4..2b08a4afd98 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -63,6 +63,8 @@ mod _native { request_auth_managed_models, request_auth_agent_ceiling_source, request_auth_managed_model_grants, + request_auth_model_projection, + request_auth_route_enable, }; #[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 56dc791c121..1802e852b1d 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -333,3 +333,18 @@ pub(crate) fn request_auth_managed_model_grants(py: Python<'_>, facts: &Bound<'_ let facts: core::managed_models::GrantFacts = from_py(facts)?; to_py(py, &py.detach(move || core::managed_models::grants(&facts))) } + +#[pyfunction] +pub(crate) fn request_auth_model_projection(py: Python<'_>, facts: &Bound<'_, PyAny>, reads: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::model_access::Facts = from_py(facts)?; + let reads: Vec = from_py(reads)?; + let decision = py.detach(move || core::model_projection::decide(&facts, &reads)).map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_enable(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::route_enable::Facts = from_py(facts)?; + let decision = py.detach(move || core::route_enable::decide(&facts)).map_err(map_error)?; + to_py(py, &decision) +} diff --git a/litellm-rust/crates/request-auth-core/src/fallback_access.rs b/litellm-rust/crates/request-auth-core/src/fallback_access.rs index 3f141dd7663..6c080e12ab2 100644 --- a/litellm-rust/crates/request-auth-core/src/fallback_access.rs +++ b/litellm-rust/crates/request-auth-core/src/fallback_access.rs @@ -20,6 +20,7 @@ pub struct Facts { pub team: Option, pub embedded_team: Option, pub assigned_models: Option>, + pub assigned_scope: Option, pub member: Option, pub project: Option, pub team_present: bool, @@ -41,6 +42,7 @@ pub enum Decision { ReadKeyGroups, ReadTeamGroups, ReadAssignedGroups, + ReadAssignedScope, } pub fn grouped(scope: GroupedScope, merge: bool) -> Result { @@ -64,7 +66,10 @@ pub fn decide(facts: Facts) -> Result { return Ok(Decision::ReadKey); }; let key = grouped(key, false)?; - if matches!(key, ModelDecision::ReadGroups) { + if matches!( + key, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadKeyGroups); } if !key.allowed() { @@ -74,9 +79,11 @@ pub fn decide(facts: Facts) -> Result { return Ok(Decision::ReadTeam); } if let Some(team) = facts.team { - let model_facts = team.facts.clone(); let decision = grouped(team, true)?; - if matches!(decision, ModelDecision::ReadGroups) { + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadTeamGroups); } if !decision.allowed() { @@ -88,13 +95,15 @@ pub fn decide(facts: Facts) -> Result { let Some(grants) = facts.assigned_models else { return Ok(Decision::ReadAssignedGroups); }; - if grants.is_empty() - || !model_access::check(&ModelFacts { - grants, - ..model_facts - })? - .allowed() - { + if grants.is_empty() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + let Some(assigned) = facts.assigned_scope else { + return Ok(Decision::ReadAssignedScope); + }; + if !model_access::check(&assigned)?.allowed() { return Ok(Decision::Deny { denial: Some(decision), }); diff --git a/litellm-rust/crates/request-auth-core/src/key_model_access.rs b/litellm-rust/crates/request-auth-core/src/key_model_access.rs index 354e8d5afcd..12fffc55528 100644 --- a/litellm-rust/crates/request-auth-core/src/key_model_access.rs +++ b/litellm-rust/crates/request-auth-core/src/key_model_access.rs @@ -51,7 +51,10 @@ pub fn decide(facts: Facts) -> Result { } if let Some(primary) = facts.primary { let decision = grouped(primary, false)?; - if matches!(decision, ModelDecision::ReadGroups) { + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadPrimaryGroups); } if !decision.allowed() { @@ -66,7 +69,10 @@ pub fn decide(facts: Facts) -> Result { return Ok(Decision::ReadFallbackScope { index }); }; let decision = grouped(scope.clone(), false)?; - if matches!(decision, ModelDecision::ReadGroups) { + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadFallbackGroups { index }); } if !decision.allowed() { diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index 6ac761a1735..8c8f8846378 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -9,6 +9,8 @@ pub mod fallback_candidate; mod key; pub mod key_model_access; pub mod model_access; +pub mod model_projection; +pub mod route_enable; pub mod model_attribution; pub mod model_budget_data; mod model_budgets; diff --git a/litellm-rust/crates/request-auth-core/src/managed_models.rs b/litellm-rust/crates/request-auth-core/src/managed_models.rs index bfd8358ec23..f618b38b371 100644 --- a/litellm-rust/crates/request-auth-core/src/managed_models.rs +++ b/litellm-rust/crates/request-auth-core/src/managed_models.rs @@ -125,7 +125,10 @@ pub fn decide(facts: &Facts) -> Result { return Ok(Decision::ReadTeamScope); }; let result = fallback_access::grouped(team.clone(), true)?; - if matches!(result, ModelDecision::ReadGroups) { + if matches!( + result, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadTeamGroups); } if !result.allowed() { diff --git a/litellm-rust/crates/request-auth-core/src/model_access.rs b/litellm-rust/crates/request-auth-core/src/model_access.rs index 841701b5f44..06211ca5cb9 100644 --- a/litellm-rust/crates/request-auth-core/src/model_access.rs +++ b/litellm-rust/crates/request-auth-core/src/model_access.rs @@ -25,11 +25,14 @@ pub struct Registry { pub aliases: BTreeMap, pub records: BTreeMap, pub compaction_parents: BTreeMap, + pub helper_present: Option, + pub evaluations: Option>, } #[derive(Clone, Default, Deserialize)] pub struct Aliases { pub team: BTreeMap, + pub team_order: Option>, pub key: BTreeMap, #[serde(rename = "global_")] pub global: BTreeMap, @@ -80,6 +83,9 @@ pub enum Decision { }, ReadGroups, ReadScope, + ReadGroupScope { + grants: Vec, + }, } impl Decision { @@ -212,11 +218,19 @@ pub fn candidate_names( .filter(|record| record.deployment_team.as_deref() == Some(team)) .and_then(|record| record.deployment_public_name.clone()) .filter(|name| !name.is_empty()); - let legacy = aliases - .team - .iter() - .filter(|(_, target)| target.as_str() == model) - .map(|(alias, _)| alias.clone()); + let ordered: Vec<&String> = aliases.team_order.as_ref().map_or_else( + || aliases.team.keys().collect(), + |names| names.iter().collect(), + ); + let legacy = ordered + .into_iter() + .filter(|alias| { + aliases + .team + .get(*alias) + .is_some_and(|target| target == model) + }) + .cloned(); (candidates, public.into_iter().chain(legacy).collect()) } @@ -235,6 +249,32 @@ pub fn live_team_target(model: &str, facts: &Facts) -> String { } pub fn check(facts: &Facts) -> Result { + check_at(facts, 0) +} + +fn check_at(facts: &Facts, skip: usize) -> Result { + if let Some(evaluations) = &facts.registry.evaluations { + let Some(evaluation) = evaluations + .iter() + .skip(skip) + .find(|evaluation| evaluation.grants == facts.grants) + else { + return Ok(Decision::ReadGroupScope { + grants: facts.grants.clone(), + }); + }; + let scoped = Facts { + aliases: evaluation + .aliases + .clone() + .unwrap_or_else(|| facts.aliases.clone()), + ..facts.clone() + }; + return match crate::model_projection::decide(&scoped, &evaluation.reads)? { + crate::model_projection::Decision::Complete { decision } => Ok(decision), + _ => Err(Error::InvalidAnswer), + }; + } if facts.fallback_depth >= facts.max_depth { return Ok(Decision::Depth { model: None }); } @@ -372,14 +412,17 @@ pub fn decide(request: Request) -> Result { if models.is_empty() { primary } else { - check(&Facts { - grants: if merge { - unique(facts.grants.iter().cloned().chain(models)) - } else { - models + check_at( + &Facts { + grants: if merge { + unique(facts.grants.iter().cloned().chain(models)) + } else { + models + }, + ..facts }, - ..facts - })? + 1, + )? } } else { Decision::ReadGroups diff --git a/litellm-rust/crates/request-auth-core/src/model_helpers.rs b/litellm-rust/crates/request-auth-core/src/model_helpers.rs index 3e5b1a75c2e..9e4ac468a87 100644 --- a/litellm-rust/crates/request-auth-core/src/model_helpers.rs +++ b/litellm-rust/crates/request-auth-core/src/model_helpers.rs @@ -20,6 +20,7 @@ pub struct Facts { pub names: Option>, pub provider_loaded: bool, pub provider: Option, + pub provider_reads: Option>>, } #[derive(Debug, PartialEq, Eq, Serialize)] @@ -90,21 +91,26 @@ pub fn decide(facts: &Facts) -> Result { return done(true); } } - if !wildcards.is_empty() { - if !facts.provider_loaded { - return Ok(Decision::ReadProvider); - } - if let Some(provider) = &facts.provider - && (provider.stripped_model != facts.model || !facts.model.contains('/')) - { - for pattern in wildcards { - if patterns::model_pattern( - &format!("{}/{}", provider.provider, provider.stripped_model), - pattern, - )? { - return done(true); - } + for (index, pattern) in wildcards.iter().enumerate() { + let provider = if let Some(providers) = &facts.provider_reads { + let Some(provider) = providers.get(index) else { + return Ok(Decision::ReadProvider); + }; + provider + } else { + if !facts.provider_loaded { + return Ok(Decision::ReadProvider); } + &facts.provider + }; + if let Some(provider) = provider + && (provider.stripped_model != facts.model || !facts.model.contains('/')) + && patterns::model_pattern( + &format!("{}/{}", provider.provider, provider.stripped_model), + pattern, + )? + { + return done(true); } } done( @@ -135,6 +141,7 @@ mod tests { names: None, provider_loaded: false, provider: None, + provider_reads: None, }; assert_eq!(decide(&facts).unwrap(), expected); } diff --git a/litellm-rust/crates/request-auth-core/src/model_projection.rs b/litellm-rust/crates/request-auth-core/src/model_projection.rs new file mode 100644 index 00000000000..829b72c703b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_projection.rs @@ -0,0 +1,237 @@ +use crate::{ + Error, + model_access::{self, Decision as ModelDecision, Facts, ModelRecord, Provider, Registry}, + model_helpers, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum ReadKind { + Compaction, + Aliases, + AliasTarget, + Deployment, + Groups, + Names, + Provider, +} + +#[derive(Clone, Deserialize)] +pub struct Reading { + pub kind: ReadKind, + pub model: String, + pub values: Vec, + pub value: Option, + pub team: Option, + pub provider: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Evaluation { + pub grants: Vec, + pub reads: Vec, + pub aliases: Option, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { decision: ModelDecision }, + Read { read: ReadKind, model: String }, +} + +enum Pending { + Read(ReadKind, String), + Invalid(Error), +} +impl From for Pending { + fn from(error: Error) -> Self { + Self::Invalid(error) + } +} + +struct Reader<'a> { + reads: &'a [Reading], + index: usize, +} +impl<'a> Reader<'a> { + fn read(&mut self, kind: ReadKind, model: &str) -> Result<&'a Reading, Pending> { + let Some(value) = self.reads.get(self.index) else { + return Err(Pending::Read(kind, model.to_owned())); + }; + if value.kind != kind || value.model != model { + return Err(Pending::Invalid(Error::InvalidAnswer)); + } + self.index += 1; + Ok(value) + } +} + +fn helper( + model: &str, + facts: &Facts, + aliases: bool, + reader: &mut Reader<'_>, +) -> Result { + let mut helper = model_helpers::Facts { + kind: model_helpers::Kind::Helper, + model: model.to_owned(), + grants: Some(facts.grants.clone()), + team_id: facts.team_id.clone(), + team_aliases: if aliases { + facts.aliases.team.keys().cloned().collect() + } else { + vec![] + }, + router_present: facts + .registry + .helper_present + .unwrap_or(facts.registry.present), + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: Some(vec![]), + }; + loop { + match model_helpers::decide(&helper)? { + model_helpers::Decision::Complete { allowed } => return Ok(allowed), + model_helpers::Decision::ReadGroups => { + helper.groups = Some(reader.read(ReadKind::Groups, model)?.values.clone()) + } + model_helpers::Decision::ReadNames => { + helper.names = Some(reader.read(ReadKind::Names, model)?.values.clone()) + } + model_helpers::Decision::ReadProvider => { + helper + .provider_reads + .as_mut() + .ok_or(Error::InvalidAnswer)? + .push(reader.read(ReadKind::Provider, model)?.provider.clone()); + } + } + } +} + +fn check_one(model: &str, facts: &Facts, reader: &mut Reader<'_>) -> Result { + let compaction = reader.read(ReadKind::Compaction, model)?.value.clone(); + let mut registry = Registry { + present: facts.registry.present, + ..Registry::default() + }; + if let Some(parent) = compaction { + registry.compaction_parents.insert(model.to_owned(), parent); + } + if !facts.aliases.global.contains_key(model) && registry.present { + let aliases = reader.read(ReadKind::Aliases, model)?; + if aliases.values.iter().any(|alias| alias == model) + && let Some(target) = &reader.read(ReadKind::AliasTarget, model)?.value + { + registry.aliases.insert(model.to_owned(), target.clone()); + } + } + let after_team = facts.aliases.team.get(model).map_or(model, String::as_str); + let after_key = facts + .aliases + .key + .get(after_team) + .map_or(after_team, String::as_str); + let after_global = facts + .aliases + .global + .get(after_key) + .map_or(after_key, String::as_str); + let dispatched = facts + .aliases + .key + .get(after_global) + .map_or(after_global, String::as_str); + let key_alias = after_key != after_team || dispatched != after_global; + if !key_alias + && registry.present + && facts + .team_id + .as_ref() + .is_some_and(|team| model.starts_with(&format!("model_name_{team}_"))) + { + let deployment = reader.read(ReadKind::Deployment, model)?; + registry.records.insert( + model.to_owned(), + ModelRecord { + deployment_team: deployment.team.clone(), + deployment_public_name: deployment.value.clone(), + ..ModelRecord::default() + }, + ); + } + let (candidates, owning) = + model_access::candidate_names(model, &facts.aliases, ®istry, facts.team_id.as_deref()); + for candidate in candidates { + if helper(&candidate, facts, true, reader)? { + return Ok(true); + } + } + for candidate in owning { + if helper(&candidate, facts, false, reader)? { + return Ok(true); + } + } + Ok(false) +} + +fn check(facts: &Facts, reader: &mut Reader<'_>) -> Result { + if facts.fallback_depth >= facts.max_depth { + return Ok(ModelDecision::Depth { model: None }); + } + for model in &facts.requested { + if facts.is_list && facts.fallback_depth + 1 >= facts.max_depth { + return Ok(ModelDecision::Depth { + model: Some(model.clone()), + }); + } + if !check_one(model, facts, reader)? { + return Ok(ModelDecision::Deny { + model: model.clone(), + grants: facts.grants.clone(), + object_type: facts.object_type, + }); + } + } + Ok(ModelDecision::Allow) +} + +pub fn decide(facts: &Facts, reads: &[Reading]) -> Result { + let mut reader = Reader { reads, index: 0 }; + match check(facts, &mut reader) { + Ok(decision) if reader.index == reads.len() => Ok(Decision::Complete { decision }), + Ok(_) => Err(Error::InvalidAnswer), + Err(Pending::Read(read, model)) => Ok(Decision::Read { read, model }), + Err(Pending::Invalid(error)) => Err(error), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + #[rstest] + #[case::empty(vec![], true)] + #[case::model(vec!["model".to_owned()], false)] + fn empty_list_never_reads_registry(#[case] requested: Vec, #[case] complete: bool) { + let facts = Facts { + requested, + grants: vec![], + team_id: None, + aliases: Default::default(), + registry: Default::default(), + object_type: model_access::ObjectType::Key, + fallback_depth: 0, + max_depth: 5, + is_list: true, + }; + assert_eq!( + matches!(decide(&facts, &[]).unwrap(), Decision::Complete { .. }), + complete + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_models.rs b/litellm-rust/crates/request-auth-core/src/personal_models.rs index e9911c95ad8..1fce2d6bd88 100644 --- a/litellm-rust/crates/request-auth-core/src/personal_models.rs +++ b/litellm-rust/crates/request-auth-core/src/personal_models.rs @@ -43,7 +43,10 @@ fn scope(facts: Scope) -> Result { return Ok(Decision::ReadTeam); }; let team = fallback_access::grouped(team, true)?; - if matches!(team, model_access::Decision::ReadGroups) { + if matches!( + team, + model_access::Decision::ReadGroups | model_access::Decision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadGroups); } if !team.allowed() { diff --git a/litellm-rust/crates/request-auth-core/src/resolved_models.rs b/litellm-rust/crates/request-auth-core/src/resolved_models.rs index 5718b0966c0..d048a49745d 100644 --- a/litellm-rust/crates/request-auth-core/src/resolved_models.rs +++ b/litellm-rust/crates/request-auth-core/src/resolved_models.rs @@ -82,7 +82,10 @@ pub fn decide(facts: &Facts) -> Result { return Ok(Decision::ReadKey); }; let key = fallback_access::grouped(key.clone(), false)?; - if matches!(key, ModelDecision::ReadGroups) { + if matches!( + key, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { return Ok(Decision::ReadKeyGroups); } if !key.allowed() { diff --git a/litellm-rust/crates/request-auth-core/src/route_enable.rs b/litellm-rust/crates/request-auth-core/src/route_enable.rs new file mode 100644 index 00000000000..2437633fefd --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_enable.rs @@ -0,0 +1,139 @@ +use crate::{ + Error, patterns, + route_access::{self, Classification, RouteContext}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Feature { + Admin, + Llm, +} + +#[derive(Clone, Deserialize)] +pub struct Setting { + pub environment_present: bool, + pub premium: bool, + pub secret: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub context: RouteContext, + pub imported: bool, + pub failed: Option, + pub admin: Option, + pub llm: Option, + pub exemptions: Option>, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum Failure { + Http, + Other, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Reason { + License, + Disabled, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + PropagateHttp, + ReadImport, + ReadExemptions, + ReadSettings { feature: Feature }, + ReadSecret { feature: Feature }, + Deny { feature: Feature, reason: Reason }, +} + +fn setting(setting: Option<&Setting>, feature: Feature) -> Option { + let Some(setting) = setting else { + return Some(Decision::ReadSettings { feature }); + }; + if setting.environment_present && !setting.premium { + return Some(Decision::Deny { + feature, + reason: Reason::License, + }); + } + match setting.secret { + None => Some(Decision::ReadSecret { feature }), + Some(true) => Some(Decision::Deny { + feature, + reason: Reason::Disabled, + }), + Some(false) => None, + } +} + +pub fn decide(facts: &Facts) -> Result { + if let Some(failed) = facts.failed { + return Ok(match failed { + Failure::Http => Decision::PropagateHttp, + Failure::Other => Decision::Allow, + }); + } + if !facts.imported { + return Ok(Decision::ReadImport); + } + let management = facts + .context + .catalog + .groups + .get("management_routes") + .map(Vec::as_slice) + .unwrap_or_default(); + if patterns::route_access(&facts.context.route, management)? + && let Some(decision) = setting(facts.admin.as_ref(), Feature::Admin) + { + return Ok(decision); + } + if !route_access::decide(route_access::Request::Classify { + context: facts.context.clone(), + kind: Classification::Llm, + })? + .allowed + { + return Ok(Decision::Allow); + } + let Some(exemptions) = &facts.exemptions else { + return Ok(Decision::ReadExemptions); + }; + if exemptions.contains(&facts.context.route) { + return Ok(Decision::Allow); + } + Ok(setting(facts.llm.as_ref(), Feature::Llm).unwrap_or(Decision::Allow)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::license(true, false, Decision::Deny { feature: Feature::Admin, reason: Reason::License })] + #[case::absent(false, false, Decision::ReadSecret { feature: Feature::Admin })] + #[case::premium(true, true, Decision::ReadSecret { feature: Feature::Admin })] + fn license_precedes_secret( + #[case] environment_present: bool, + #[case] premium: bool, + #[case] expected: Decision, + ) { + assert_eq!( + setting( + Some(&Setting { + environment_present, + premium, + secret: None + }), + Feature::Admin + ), + Some(expected) + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_override.rs b/litellm-rust/crates/request-auth-core/src/team_override.rs index ca69868c4cc..08e0cc20481 100644 --- a/litellm-rust/crates/request-auth-core/src/team_override.rs +++ b/litellm-rust/crates/request-auth-core/src/team_override.rs @@ -30,7 +30,9 @@ pub fn decide(facts: &Facts) -> Result { let denial = fallback_access::grouped(team.clone(), true)?; match denial { ModelDecision::Allow => return Ok(Decision::Allow), - ModelDecision::ReadGroups => return Ok(Decision::ReadTeamGroups), + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } => { + return Ok(Decision::ReadTeamGroups); + } ModelDecision::Deny { object_type: ObjectType::Team, .. diff --git a/litellm/_auth/fallback_access.py b/litellm/_auth/fallback_access.py index 5ae5bc64c63..057f3a47745 100644 --- a/litellm/_auth/fallback_access.py +++ b/litellm/_auth/fallback_access.py @@ -29,6 +29,7 @@ class FallbackAuthorizationFacts: embedded_team_grants: tuple[str, ...] = () member_grants: tuple[str, ...] = () project_grants: tuple[str, ...] = () + assigned_scope: ModelAccessFacts | None = None @dataclass(frozen=True, slots=True) @@ -44,6 +45,7 @@ class FallbackAuthorizationDecision: "ReadKeyGroups", "ReadTeamGroups", "ReadAssignedGroups", + "ReadAssignedScope", ] denial: ModelAccessDecision | None = None diff --git a/litellm/_auth/model_access.py b/litellm/_auth/model_access.py index fa78779acfc..e3c29ad4c64 100644 --- a/litellm/_auth/model_access.py +++ b/litellm/_auth/model_access.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from collections.abc import Mapping from dataclasses import asdict, dataclass, field from typing import Final, Literal, TypeAlias @@ -8,6 +10,7 @@ from litellm._auth.failure import native_boundary from litellm.rust_bridge import _native ModelObject: TypeAlias = Literal["user", "customer", "team", "key", "org", "project", "agent"] +ModelReadKind: TypeAlias = Literal["Compaction", "Aliases", "AliasTarget", "Deployment", "Groups", "Names", "Provider"] @dataclass(frozen=True, slots=True) @@ -24,6 +27,23 @@ class ModelRecord: deployment_public_name: str | None = None +@dataclass(frozen=True, slots=True) +class ModelReading: + kind: ModelReadKind + model: str + values: tuple[str, ...] = () + value: str | None = None + team: str | None = None + provider: ModelProvider | None = None + + +@dataclass(frozen=True, slots=True) +class ModelEvaluation: + grants: tuple[str, ...] + reads: tuple[ModelReading, ...] + aliases: ModelAliases | None = None + + @dataclass(frozen=True, slots=True) class ModelRegistry: present: bool = False @@ -31,6 +51,8 @@ class ModelRegistry: aliases: Mapping[str, str] = field(default_factory=dict[str, str]) records: Mapping[str, ModelRecord] = field(default_factory=dict[str, ModelRecord]) compaction_parents: Mapping[str, str] = field(default_factory=dict[str, str]) + helper_present: bool | None = None + evaluations: tuple[ModelEvaluation, ...] | None = None @dataclass(frozen=True, slots=True) @@ -38,6 +60,7 @@ class ModelAliases: team: Mapping[str, str] = field(default_factory=dict[str, str]) key: Mapping[str, str] = field(default_factory=dict[str, str]) global_: Mapping[str, str] = field(default_factory=dict[str, str]) + team_order: tuple[str, ...] | None = None @dataclass(frozen=True, slots=True) @@ -65,6 +88,7 @@ class ModelAccessDecision: "PersonalTeam", "ReadGroups", "ReadScope", + "ReadGroupScope", ] model: str | None = None grants: tuple[str, ...] = () diff --git a/litellm/_auth/model_helpers.py b/litellm/_auth/model_helpers.py index 8ba323a1e42..8de363b28ac 100644 --- a/litellm/_auth/model_helpers.py +++ b/litellm/_auth/model_helpers.py @@ -20,6 +20,7 @@ class ModelHelperFacts: names: tuple[str, ...] | None = None provider_loaded: bool = False provider: ModelProvider | None = None + provider_reads: tuple[ModelProvider | None, ...] | None = None @dataclass(frozen=True, slots=True) diff --git a/litellm/_auth/model_projection.py b/litellm/_auth/model_projection.py new file mode 100644 index 00000000000..63022394416 --- /dev/null +++ b/litellm/_auth/model_projection.py @@ -0,0 +1,26 @@ +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, ModelReading, ModelReadKind +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelProjectionDecision: + kind: Literal["Complete", "Read"] + read: ModelReadKind | None = None + model: str | None = None + decision: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(ModelProjectionDecision) + + +@native_boundary +def model_projection(facts: ModelAccessFacts, reads: tuple[ModelReading, ...]) -> ModelProjectionDecision: + return _DECISION.validate_python( + _native.request_auth_model_projection(asdict(facts), tuple(asdict(read) for read in reads)) + ) diff --git a/litellm/_auth/route_enable.py b/litellm/_auth/route_enable.py new file mode 100644 index 00000000000..e528695adfb --- /dev/null +++ b/litellm/_auth/route_enable.py @@ -0,0 +1,42 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.routes import RouteContext +from litellm.rust_bridge import _native + +RouteFeature: TypeAlias = Literal["Admin", "Llm"] + + +@dataclass(frozen=True, slots=True) +class RouteSetting: + environment_present: bool + premium: bool + secret: bool | None = None + + +@dataclass(frozen=True, slots=True) +class RouteEnableFacts: + context: RouteContext + imported: bool = False + failed: Literal["Http", "Other"] | None = None + admin: RouteSetting | None = None + llm: RouteSetting | None = None + exemptions: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class RouteEnableDecision: + kind: Literal["Allow", "PropagateHttp", "ReadImport", "ReadExemptions", "ReadSettings", "ReadSecret", "Deny"] + feature: RouteFeature | None = None + reason: Literal["License", "Disabled"] | None = None + + +_DECISION: Final = TypeAdapter(RouteEnableDecision) + + +@native_boundary +def route_enable(facts: RouteEnableFacts) -> RouteEnableDecision: + return _DECISION.validate_python(_native.request_auth_route_enable(asdict(facts))) diff --git a/litellm/proxy/auth/fallback_access_adapter.py b/litellm/proxy/auth/fallback_access_adapter.py index 7bb20e6ce9a..26e7530c924 100644 --- a/litellm/proxy/auth/fallback_access_adapter.py +++ b/litellm/proxy/auth/fallback_access_adapter.py @@ -3,10 +3,11 @@ from typing import Final from litellm._auth.failure import NativeAuthRuntimeError from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope, fallback_authorization -from litellm._auth.model_access import key_model_grants +from litellm._auth.model_access import ModelAccessFacts, key_model_grants from litellm.proxy._types import LiteLLM_ProjectTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope from litellm.repositories.prisma_protocols import DatabaseClient from litellm.router_api import RegistryReader @@ -45,7 +46,7 @@ async def load_fallback_authorization_facts( async def _read(facts: FallbackAuthorizationFacts, source: _Sources) -> FallbackAuthorizationFacts: - from litellm.proxy.auth.auth_checks import get_models_from_access_groups, key_model_aliases_for_auth_check + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check decision: Final = fallback_authorization(facts) match decision.kind: @@ -106,29 +107,48 @@ async def _read(facts: FallbackAuthorizationFacts, source: _Sources) -> Fallback ) return await _read(replace(facts, member=member), source) case "ReadProject": - if source.project is None: - raise NativeAuthRuntimeError("missing fallback project source") - project: 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", - ) + project: Final = _project_scope(facts, source) return await _read(replace(facts, project=project), source) case "ReadKeyGroups": if facts.key is None: raise NativeAuthRuntimeError("missing fallback key scope") - models: Final = await get_models_from_access_groups(facts.key.group_ids, prisma_client=source.prisma) - return await _read(replace(facts, key=replace(facts.key, group_models=tuple(models))), source) + key_scope: Final = await read_grouped_scope(facts.key, source.router, source.prisma, merge=False) + return await _read(replace(facts, key=key_scope), source) case "ReadTeamGroups": if facts.team is None: raise NativeAuthRuntimeError("missing fallback team scope") - team_models: 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(team_models))), source) + team_scope: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=team_scope), source) case "ReadAssignedGroups": assigned: Final = await get_authorized_resources_from_key_access_groups( source.token, source.team, "access_model_names" ) return await _read(replace(facts, assigned_models=tuple(assigned)), source) + case "ReadAssignedScope": + if facts.assigned_models is None: + raise NativeAuthRuntimeError("missing fallback assigned models") + scope: Final = project_model_facts( + source.model, + source.router, + facts.assigned_models, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ) + return await _read(replace(facts, assigned_scope=scope), source) + + +def _project_scope(facts: FallbackAuthorizationFacts, source: _Sources) -> ModelAccessFacts: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + if source.project is None: + raise NativeAuthRuntimeError("missing fallback project source") + return 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", + ) diff --git a/litellm/proxy/auth/fallback_candidate_adapter.py b/litellm/proxy/auth/fallback_candidate_adapter.py index 586067e0424..eb4440fd382 100644 --- a/litellm/proxy/auth/fallback_candidate_adapter.py +++ b/litellm/proxy/auth/fallback_candidate_adapter.py @@ -1,9 +1,9 @@ from collections.abc import Mapping from dataclasses import dataclass, replace -from typing import ( # noqa: TID251 # preserve opaque legacy budget mappings without coercing their keys +from typing import ( Final, Protocol, - cast, + cast, # noqa: TID251 # preserve opaque legacy budget mappings without coercing their keys ) from pydantic import BaseModel diff --git a/litellm/proxy/auth/key_model_access_adapter.py b/litellm/proxy/auth/key_model_access_adapter.py index 37cb7e337fd..eb79c6d93eb 100644 --- a/litellm/proxy/auth/key_model_access_adapter.py +++ b/litellm/proxy/auth/key_model_access_adapter.py @@ -11,6 +11,7 @@ from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_a from litellm._auth.model_access import key_model_grants from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope from litellm.router import Router @@ -70,11 +71,8 @@ def _scope(model: str | Sequence[str], source: _Sources) -> GroupedModelScope: ) -async def _groups(scope: GroupedModelScope) -> GroupedModelScope: - from litellm.proxy.auth.auth_checks import get_models_from_access_groups - - models: Final = await get_models_from_access_groups(scope.group_ids) - return replace(scope, group_models=tuple(models)) +async def _groups(scope: GroupedModelScope, source: _Sources) -> GroupedModelScope: + return await read_grouped_scope(scope, source.router, merge=False) async def _route(name: str, source: _Sources) -> Exception | None: @@ -124,7 +122,7 @@ async def _read( case "ReadPrimaryGroups": if facts.primary is None: raise NativeAuthRuntimeError("missing primary model scope") - return await _read(replace(facts, primary=await _groups(facts.primary)), source, route_errors) + return await _read(replace(facts, primary=await _groups(facts.primary, source)), source, route_errors) case "ReadFallbackNames": names: Final = tuple( dict.fromkeys( @@ -140,7 +138,7 @@ async def _read( scope: Final = _scope(facts.fallback_names[decision.index], source) return await _read(replace(facts, fallback_scopes=(*facts.fallback_scopes, scope)), source, route_errors) case "ReadFallbackGroups": - grouped: Final = await _groups(facts.fallback_scopes[decision.index]) + grouped: Final = await _groups(facts.fallback_scopes[decision.index], source) scopes: Final = tuple( grouped if index == decision.index else scope for index, scope in enumerate(facts.fallback_scopes) ) diff --git a/litellm/proxy/auth/managed_model_adapter.py b/litellm/proxy/auth/managed_model_adapter.py index f3dfd69b39f..069e739615a 100644 --- a/litellm/proxy/auth/managed_model_adapter.py +++ b/litellm/proxy/auth/managed_model_adapter.py @@ -29,6 +29,7 @@ from litellm.proxy.agent_endpoints.auth.agent_caller import ( 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.model_projection_adapter import read_grouped_scope 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 @@ -252,7 +253,7 @@ def _dispatch(source: _Source) -> str | Sequence[str]: 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 + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check decision: Final = managed_models(facts) match decision.kind: @@ -331,8 +332,8 @@ async def _read(facts: ManagedModelFacts, source: _Source) -> tuple[ManagedModel 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) + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=grouped), 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") diff --git a/litellm/proxy/auth/model_access_adapter.py b/litellm/proxy/auth/model_access_adapter.py index b8df7e88ef5..f5322635870 100644 --- a/litellm/proxy/auth/model_access_adapter.py +++ b/litellm/proxy/auth/model_access_adapter.py @@ -1,24 +1,21 @@ from collections.abc import Mapping, Sequence -from dataclasses import replace from typing import Final, Literal from fastapi import HTTPException from pydantic import BaseModel import litellm -from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope from litellm._auth.model_access import ( ModelAccessDecision, ModelAccessFacts, ModelAliases, ModelObject, - ModelProvider, - ModelRecord, ModelRegistry, grouped_model_access, key_model_grants, model_access, - model_candidates, user_model_precheck, ) from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -36,9 +33,9 @@ from litellm.proxy.auth.model_access_denied import ( customer_model_access_denied_client_message, model_access_denied_client_message, ) +from litellm.proxy.auth.model_projection_adapter import project_scope, read_grouped_scope from litellm.repositories.prisma_protocols import DatabaseClient from litellm.router_api import RegistryReader -from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent class ModelLists(BaseModel, frozen=True): @@ -50,30 +47,8 @@ def object_model_grants(source: object) -> tuple[str, ...]: return ModelLists.model_validate(source, from_attributes=True).models -def _provider(model: str) -> ModelProvider | None: - try: - stripped, provider, _, _ = litellm.get_llm_provider(model=model) - return ModelProvider(stripped, provider) - except NativeAuthFailure: - raise - except Exception: # noqa: BLE001 # provider inference failure leaves the raw provider fact absent - return None - - -def _record(model: str, router: RegistryReader | None, team_id: str | None) -> ModelRecord: - deployment: Final = router.registry().deployments(model, source="first_in_group") if router is not None else None - return ModelRecord( - access_groups=tuple(router.registry().access_groups(model_name=model, team_id=team_id)) - if router is not None - else (), - provider=_provider(model), - deployment_team=deployment.model_info.team_id if deployment is not None else None, - deployment_public_name=deployment.model_info.team_public_model_name if deployment is not None else None, - ) - - @native_boundary -def project_model_facts( +def _initial_model_facts( model: str | Sequence[str], llm_router: RegistryReader | None, models: Sequence[str], @@ -85,22 +60,15 @@ def project_model_facts( ) -> ModelAccessFacts: requested: Final = (model,) if isinstance(model, str) else tuple(model) aliases: Final = ModelAliases( - team=dict(team_model_aliases or {}), key=dict(key_model_aliases or {}), global_=dict(litellm.model_alias_map) + team=dict(team_model_aliases or {}), + key=dict(key_model_aliases or {}), + global_=dict(litellm.model_alias_map), + team_order=tuple(team_model_aliases or ()), ) - names: Final = tuple(dict.fromkeys((*requested, *aliases.team.values()))) registry: Final = ModelRegistry( present=llm_router is not None, - names=tuple(llm_router.registry().names()) if llm_router is not None else (), - aliases={ - name: target - for name in names - if name in llm_router.registry().aliases() - and (target := llm_router.registry().resolve(name, kind="alias")) is not None - } - if llm_router is not None - else {}, - records={name: _record(name, llm_router, team_id) for name in names}, - compaction_parents={name: parent for name in names if (parent := native_compaction_parent(name)) is not None}, + helper_present=bool(llm_router), + evaluations=(), ) initial: Final = ModelAccessFacts( requested=requested, @@ -113,17 +81,23 @@ def project_model_facts( max_depth=DEFAULT_MAX_RECURSE_DEPTH, is_list=not isinstance(model, str), ) - candidates: Final = model_candidates(replace(initial, requested=names)) - return replace( - initial, - registry=replace( - registry, - records={ - **registry.records, - **{name: _record(name, llm_router, team_id) for name in candidates if name not in registry.records}, - }, - ), + return initial + + +def project_model_facts( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> ModelAccessFacts: + initial: Final = _initial_model_facts( + model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type, fallback_depth ) + return project_scope(initial, llm_router) def raise_model_denial( @@ -145,7 +119,7 @@ def raise_model_denial( ) case "NoAgentModels": raise HTTPException(403, "This agent has no model grants") - case "ReadGroups" | "ReadScope": + case "ReadGroups" | "ReadScope" | "ReadGroupScope": raise NativeAuthRuntimeError("model grant lookup was not answered") case "Deny" | "NoDefaultModels" | "EmptyAgentCeiling" | "PersonalTeam": error_model: Final = decision.model if decision.model is not None else model @@ -214,14 +188,12 @@ async def _grouped_model_access( prisma_client: DatabaseClient | None, *, merge: bool, + router: RegistryReader | None = None, ) -> ModelAccessDecision: - from litellm.proxy.auth.auth_checks import get_models_from_access_groups - - decision: Final = grouped_model_access(facts, tuple(group_ids), merge=merge) - if decision.kind != "ReadGroups": - return decision - models: Final = await get_models_from_access_groups(list(group_ids), prisma_client=prisma_client) - return grouped_model_access(facts, tuple(group_ids), tuple(models), merge=merge) + scope: Final = await read_grouped_scope( + GroupedModelScope(facts, tuple(group_ids)), router, prisma_client, merge=merge + ) + return grouped_model_access(scope.facts, scope.group_ids, scope.group_models, merge=merge) async def can_key_call_model( @@ -245,7 +217,9 @@ async def can_key_call_model( key_model_aliases_for_auth_check(valid_token), "key", ) - decision: Final = await _grouped_model_access(facts, valid_token.access_group_ids or (), prisma_client, merge=False) + decision: Final = await _grouped_model_access( + facts, valid_token.access_group_ids or (), prisma_client, merge=False, router=llm_router + ) return raise_model_denial(decision, model) @@ -267,7 +241,11 @@ async def can_team_access_model( "team", ) decision: Final = await _grouped_model_access( - facts, (team_object.access_group_ids or ()) if team_object is not None else (), prisma_client, merge=True + facts, + (team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client, + merge=True, + router=llm_router, ) return raise_model_denial(decision, model) diff --git a/litellm/proxy/auth/model_helper_adapter.py b/litellm/proxy/auth/model_helper_adapter.py index 274efc648d6..b6b02d27281 100644 --- a/litellm/proxy/auth/model_helper_adapter.py +++ b/litellm/proxy/auth/model_helper_adapter.py @@ -4,13 +4,13 @@ from typing import Final from litellm._auth.failure import NativeAuthRuntimeError from litellm._auth.model_helpers import ModelHelperFacts, model_helper -from litellm.proxy.auth.model_access_adapter import ( - _provider, # pyright: ignore[reportPrivateUsage] # shared raw provider projection preserves terminal failures -) +from litellm.proxy.auth.model_projection_adapter import provider_facts from litellm.router_api import RegistryReader def read_model_helper_facts(facts: ModelHelperFacts, router: RegistryReader | None) -> ModelHelperFacts: + if facts.provider_reads is None and not facts.provider_loaded: + return read_model_helper_facts(replace(facts, provider_reads=()), router) decision: Final = model_helper(facts) match decision.kind: case "Complete": @@ -26,7 +26,12 @@ def read_model_helper_facts(facts: ModelHelperFacts, router: RegistryReader | No return read_model_helper_facts(replace(facts, names=tuple(router.registry().names())), router) case "ReadProvider": return read_model_helper_facts( - replace(facts, provider_loaded=True, provider=_provider(facts.model)), router + replace( + facts, + provider_loaded=True, + provider_reads=(*(facts.provider_reads or ()), provider_facts(facts.model)), + ), + router, ) diff --git a/litellm/proxy/auth/model_projection_adapter.py b/litellm/proxy/auth/model_projection_adapter.py new file mode 100644 index 00000000000..50141d7fb3c --- /dev/null +++ b/litellm/proxy/auth/model_projection_adapter.py @@ -0,0 +1,121 @@ +from collections.abc import Iterable +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import Final, assert_never + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ( + ModelAccessFacts, + ModelEvaluation, + ModelProvider, + ModelReading, + ModelReadKind, + grouped_model_access, +) +from litellm._auth.model_projection import ModelProjectionDecision, model_projection +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader +from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent + + +def provider_facts(model: str) -> ModelProvider | None: + try: + stripped, provider, _, _ = litellm.get_llm_provider(model=model) + return ModelProvider(stripped, provider) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # provider inference failure is an absent provider, as in the legacy helper + return None + + +def project_scope( + facts: ModelAccessFacts, router: RegistryReader | None, reads: tuple[ModelReading, ...] = () +) -> ModelAccessFacts: + initial: Final = _Projection(reads, model_projection(facts, reads)) + + def advance(state: _Projection, _tick: None) -> _Projection: + decision: Final = state.decision + if decision.read is None or decision.model is None: + raise NativeAuthRuntimeError("incomplete model registry read") + reading: Final = read_registry(decision.read, decision.model, facts.team_id, router) + updated: Final = (*state.reads, reading) + return _Projection(updated, model_projection(facts, updated)) + + complete: Final = next( + state for state in accumulate(repeat(None), advance, initial=initial) if state.decision.kind == "Complete" + ) + return replace( + facts, + registry=replace( + facts.registry, + evaluations=( + *(facts.registry.evaluations or ()), + ModelEvaluation(facts.grants, complete.reads, facts.aliases), + ), + ), + ) + + +@dataclass(frozen=True, slots=True) +class _Projection: + reads: tuple[ModelReading, ...] + decision: ModelProjectionDecision + + +@native_boundary +def _strings(values: Iterable[str]) -> tuple[str, ...]: + return tuple(values) + + +def read_registry(kind: ModelReadKind, model: str, team_id: str | None, router: RegistryReader | None) -> ModelReading: + if kind == "Compaction": + return ModelReading(kind, model, value=native_compaction_parent(model)) + if kind == "Provider": + return ModelReading(kind, model, provider=provider_facts(model)) + if router is None: + raise NativeAuthRuntimeError("model registry read without router") + match kind: + case "Aliases": + return ModelReading(kind, model, values=_strings(router.registry().aliases())) + case "AliasTarget": + return ModelReading(kind, model, value=router.registry().resolve(model, kind="alias")) + case "Names": + return ModelReading(kind, model, values=_strings(router.registry().names())) + case "Groups": + return ModelReading( + kind, model, values=_strings(router.registry().access_groups(model_name=model, team_id=team_id)) + ) + case "Deployment": + deployment: Final = router.registry().deployments(model, source="first_in_group") + return ModelReading( + kind, + model, + value=deployment.model_info.team_public_model_name if deployment is not None else None, + team=deployment.model_info.team_id if deployment is not None else None, + ) + case _: + assert_never(kind) + + +async def read_grouped_scope( + scope: GroupedModelScope, router: RegistryReader | None, prisma: DatabaseClient | None = None, *, merge: bool +) -> GroupedModelScope: + from litellm.proxy.auth.auth_checks import get_models_from_access_groups + + decision: Final = grouped_model_access(scope.facts, scope.group_ids, scope.group_models, merge=merge) + if decision.kind == "ReadGroups": + models: Final = await get_models_from_access_groups(scope.group_ids, prisma_client=prisma) + return await read_grouped_scope(replace(scope, group_models=tuple(models)), router, prisma, merge=merge) + if decision.kind == "ReadGroupScope": + scoped: Final = project_scope( + replace( + scope.facts, + grants=decision.grants, + aliases=replace(scope.facts.aliases, global_=dict(litellm.model_alias_map)), + ), + router, + ) + return replace(scope, facts=replace(scoped, grants=scope.facts.grants)) + return scope diff --git a/litellm/proxy/auth/personal_model_adapter.py b/litellm/proxy/auth/personal_model_adapter.py index 38ad4fa7592..9fa87d43502 100644 --- a/litellm/proxy/auth/personal_model_adapter.py +++ b/litellm/proxy/auth/personal_model_adapter.py @@ -18,6 +18,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy._types import LiteLLM_TeamTableCachedObj, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router_api import RegistryReader @@ -101,7 +102,7 @@ async def _read_scope(facts: PersonalModelScope, source: _Source) -> PersonalMod async def _scope_effect( decision: PersonalModelDecision, facts: PersonalModelScope, source: _Source ) -> PersonalModelScope: - from litellm.proxy.auth.auth_checks import get_models_from_access_groups, get_team_membership + from litellm.proxy.auth.auth_checks import get_team_membership match decision.kind: case "ReadTeam": @@ -117,8 +118,8 @@ async def _scope_effect( case "ReadGroups": if facts.team is None: raise NativeAuthRuntimeError("personal model group read without team facts") - groups: Final = await get_models_from_access_groups(facts.team.group_ids, prisma_client=source.prisma) - return replace(facts, team=replace(facts.team, group_models=tuple(groups))) + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return replace(facts, team=grouped) case "ReadMembership": if source.token.user_id is None: raise NativeAuthRuntimeError("personal model membership read without identity") diff --git a/litellm/proxy/auth/resolved_model_adapter.py b/litellm/proxy/auth/resolved_model_adapter.py index 66a34a5cd59..6316a43a885 100644 --- a/litellm/proxy/auth/resolved_model_adapter.py +++ b/litellm/proxy/auth/resolved_model_adapter.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( ) 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.model_projection_adapter import read_grouped_scope 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 @@ -108,7 +109,6 @@ async def _read(facts: ResolvedModelFacts, source: _Source) -> ResolvedModelRead 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, @@ -139,8 +139,8 @@ async def _read(facts: ResolvedModelFacts, source: _Source) -> ResolvedModelRead 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) + scope: Final = await read_grouped_scope(facts.key, source.router, merge=False) + return await _read(replace(facts, key=scope), source) case "ReadTeamObject": if source.token.team_id is None: raise NativeAuthRuntimeError("resolved team read without identity") diff --git a/litellm/proxy/auth/route_enable_adapter.py b/litellm/proxy/auth/route_enable_adapter.py new file mode 100644 index 00000000000..abdeb50001f --- /dev/null +++ b/litellm/proxy/auth/route_enable_adapter.py @@ -0,0 +1,122 @@ +import os +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Final, assert_never + +from fastapi import HTTPException +from pydantic import BaseModel + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting, route_enable +from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.auth.route_policy_adapter import route_context + +_KEYS: Final = {"Admin": "DISABLE_ADMIN_ENDPOINTS", "Llm": "DISABLE_LLM_API_ENDPOINTS"} + + +class _Exemptions(BaseModel, frozen=True): + LLM_API_EXEMPT_ROUTES: tuple[str, ...] + + +def read_enterprise() -> object: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + return EnterpriseRouteChecks + + +def read_setting(feature: RouteFeature) -> RouteSetting: + from litellm.proxy.proxy_server import premium_user + from litellm.secret_managers.main import ( + get_secret_bool as _get_secret_bool, # noqa: F401 # pyright: ignore[reportUnusedImport] # preserve the Phase 1 import before environment and license reads + ) + + return RouteSetting(_KEYS[feature] in os.environ, bool(premium_user)) + + +def read_secret(feature: RouteFeature) -> bool: + from litellm.secret_managers.main import get_secret_bool + + return get_secret_bool(_KEYS[feature]) is True + + +@native_boundary +def read_exemptions(enterprise: object) -> tuple[str, ...]: + return _Exemptions.model_validate(enterprise, from_attributes=True).LLM_API_EXEMPT_ROUTES + + +@dataclass(frozen=True, slots=True) +class RouteEnableReaders: + enterprise: Callable[[], object] = read_enterprise + setting: Callable[[RouteFeature], RouteSetting] = read_setting + secret: Callable[[RouteFeature], bool] = read_secret + exemptions: Callable[[object], tuple[str, ...]] = read_exemptions + + +_READERS: Final = RouteEnableReaders() + + +def _answer( + facts: RouteEnableFacts, decision: RouteEnableDecision, readers: RouteEnableReaders, enterprise: object +) -> tuple[RouteEnableFacts, object]: + match decision.kind: + case "ReadImport": + return replace(facts, imported=True), readers.enterprise() + case "ReadExemptions": + return replace(facts, exemptions=readers.exemptions(enterprise)), enterprise + case "ReadSettings" | "ReadSecret": + if decision.feature is None: + raise NativeAuthRuntimeError("route setting read without feature") + previous: Final = facts.admin if decision.feature == "Admin" else facts.llm + if decision.kind == "ReadSecret" and previous is None: + raise NativeAuthRuntimeError("route secret read without settings") + setting: Final = ( + readers.setting(decision.feature) + if decision.kind == "ReadSettings" or previous is None + else replace(previous, secret=readers.secret(decision.feature)) + ) + return ( + replace(facts, admin=setting) if decision.feature == "Admin" else replace(facts, llm=setting) + ), enterprise + case "Allow" | "Deny" | "PropagateHttp": + raise NativeAuthRuntimeError("completed route enable decision requested IO") + case _: + assert_never(decision.kind) + + +def _read( + facts: RouteEnableFacts, + readers: RouteEnableReaders, + enterprise: object = None, + http_error: HTTPException | None = None, +) -> tuple[RouteEnableFacts, HTTPException | None]: + decision: Final = route_enable(facts) + if decision.kind in ("Allow", "Deny", "PropagateHttp"): + return facts, http_error + try: + updated, imported = _answer(facts, decision, readers, enterprise) + except NativeAuthFailure: + raise + except HTTPException as error: + return _read(replace(facts, failed="Http"), readers, enterprise, error) + except Exception: # noqa: BLE001 # optional enterprise import and configuration IO failures were ignored in Phase 1 + return _read(replace(facts, failed="Other"), readers, enterprise) + return _read(updated, readers, imported) + + +def check_route_enabled(route: str, *, readers: RouteEnableReaders = _READERS) -> None: + facts, error = _read(RouteEnableFacts(route_context(route)), readers) + decision: Final = route_enable(facts) + if decision.kind == "Allow": + return + if decision.kind == "PropagateHttp" and error is not None: + raise error + if decision.kind != "Deny" or decision.feature is None or decision.reason is None: + raise NativeAuthRuntimeError("incomplete route enable decision") + if decision.reason == "License": + label: Final = "ADMIN" if decision.feature == "Admin" else "LLM API" + raise HTTPException( + 500, + f"\U0001f6a8\U0001f6a8\U0001f6a8 DISABLING {label} ENDPOINTS is an Enterprise feature\n\U0001f6a8 {CommonProxyErrors.not_premium_user.value}", + ) + disabled_label: Final = "Management" if decision.feature == "Admin" else "LLM API" + raise HTTPException(403, f"{disabled_label} routes are disabled for this instance.") diff --git a/litellm/proxy/auth/team_alias_adapter.py b/litellm/proxy/auth/team_alias_adapter.py index 551cf91ff5e..1c83caf1dba 100644 --- a/litellm/proxy/auth/team_alias_adapter.py +++ b/litellm/proxy/auth/team_alias_adapter.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace from typing import Final @@ -7,7 +7,6 @@ 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, @@ -22,8 +21,13 @@ def resolve_team_alias( 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)) + names: Final = tuple(_names(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) + + +@native_boundary +def _names(names: Iterable[str]) -> tuple[str, ...]: + return tuple(names) diff --git a/litellm/proxy/auth/team_override_adapter.py b/litellm/proxy/auth/team_override_adapter.py index fc471573d1a..ee0ed0c1d3b 100644 --- a/litellm/proxy/auth/team_override_adapter.py +++ b/litellm/proxy/auth/team_override_adapter.py @@ -8,6 +8,7 @@ from litellm._auth.team_override import TeamOverrideFacts, team_override from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope from litellm.repositories.prisma_protocols import DatabaseClient from litellm.router_api import RegistryReader @@ -59,7 +60,6 @@ async def can_team_call_model_with_override( async def _read(facts: TeamOverrideFacts, source: _Source) -> TeamOverrideFacts: - from litellm.proxy.auth.auth_checks import get_models_from_access_groups decision: Final = team_override(facts) match decision.kind: @@ -80,8 +80,8 @@ async def _read(facts: TeamOverrideFacts, source: _Source) -> TeamOverrideFacts: case "ReadTeamGroups": if facts.team is None: raise NativeAuthRuntimeError("missing team override scope") - models: 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(models))), source) + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=grouped), source) case "ReadAssignedGroups": assigned: Final = await get_authorized_resources_from_key_access_groups( source.token, source.team, "access_model_names" diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8330d5bf8b9..2fd1603e146 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -52,6 +52,8 @@ 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: ... +def request_auth_model_projection(facts: object, reads: object) -> object: ... +def request_auth_route_enable(facts: object) -> object: ... @final class NativeClickHouseSpendConfig: diff --git a/tests/unit/_auth/support/fallback_reference.py b/tests/unit/_auth/support/fallback_reference.py index 1459c32801e..d3c47e288a7 100644 --- a/tests/unit/_auth/support/fallback_reference.py +++ b/tests/unit/_auth/support/fallback_reference.py @@ -41,10 +41,11 @@ def authorization(facts: Facts) -> Decision: return Decision("Deny", team) if facts.assigned_models is None: return Decision("ReadAssignedGroups") - if ( - not facts.assigned_models - or model_access(replace(facts.team.facts, grants=facts.assigned_models)).kind != "Allow" - ): + if not facts.assigned_models: + return Decision("Deny", team) + if facts.assigned_scope is None: + return Decision("ReadAssignedScope") + if model_access(facts.assigned_scope).kind != "Allow": return Decision("Deny", team) elif facts.embedded_team is None and facts.embedded_team_grants: return Decision("ReadEmbeddedTeam") diff --git a/tests/unit/_auth/support/model_reference.py b/tests/unit/_auth/support/model_reference.py index 7fb210eed49..e9155f8d256 100644 --- a/tests/unit/_auth/support/model_reference.py +++ b/tests/unit/_auth/support/model_reference.py @@ -6,7 +6,7 @@ from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts def _pattern(model: str, pattern: str) -> bool: - return "*" in pattern and re.match("^" + re.escape(pattern).replace(r"\*", ".*") + "$", model) is not None + return "*" in pattern and re.match("^" + pattern.replace("*", ".*") + "$", model) is not None def helper(model: str, facts: ModelAccessFacts, *, team_aliases: bool = True) -> bool: diff --git a/tests/unit/_auth/support/model_registry_reference.py b/tests/unit/_auth/support/model_registry_reference.py new file mode 100644 index 00000000000..d499bab565e --- /dev/null +++ b/tests/unit/_auth/support/model_registry_reference.py @@ -0,0 +1,187 @@ +"""Model authorization and registry ordering frozen from Phase 1, 8bd7f5089a75.""" + +import re +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy._types import LiteLLM_UserTable, ModelAccessDeniedProxyException, ProxyErrorTypes +from litellm.proxy.auth.model_access_denied import ( + customer_model_access_denied_client_message, + model_access_denied_client_message, +) +from litellm.router_api import RegistryReader +from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent + + +def _pattern(model: str, pattern: str) -> bool: + return "*" in pattern and re.match("^" + pattern.replace("*", ".*") + "$", model) is not None + + +def _provider_pattern(model: str, pattern: str) -> bool: + try: + stripped, provider, _, _ = litellm.get_llm_provider(model=model) + except Exception: # noqa: BLE001 # the Phase 1 provider inference helper swallowed lookup failures + return False + if stripped == model and "/" in model: + return False + return _pattern(f"{provider}/{stripped}", pattern) + + +def _helper( + model: str, + router: RegistryReader | None, + grants: Sequence[str], + team: str | None, + aliases: Mapping[str, str] | None, +) -> bool: + groups: Final = router.registry().access_groups(model_name=model, team_id=team) if router else {} + names: Final = ( + router.registry().names() if "all-team-models" in grants and team is not None and router is not None else None + ) + ordinary: Final = tuple(grant for grant in grants if grant != "all-team-models") + expanded: Final = ( + (tuple(dict.fromkeys((*ordinary, *names))) or tuple(grants)) if names is not None else tuple(grants) + ) + if groups and any(grant in groups for grant in expanded): + return True + filtered: Final = tuple(grant for grant in expanded if grant not in groups) + if aliases and model in aliases: + return True + if any(_pattern(model, grant) for grant in filtered): + return True + if any("*" in grant and _provider_pattern(model, grant) for grant in filtered): + return True + return (not filtered and not expanded) or "*" in filtered or "all-proxy-models" in filtered or model in filtered + + +def _owning( + model: str, router: RegistryReader | None, team: str | None, aliases: Mapping[str, str] | None +) -> tuple[str, ...]: + if router is None or team is None or not model.startswith(f"model_name_{team}_"): + return () + deployment: Final = router.registry().deployments(model, source="first_in_group") + public: Final = ( + deployment.model_info.team_public_model_name + if deployment is not None and deployment.model_info.team_id == team + else None + ) + return (*((public,) if public else ()), *(alias for alias, target in (aliases or {}).items() if target == model)) + + +def can_object_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"] = "user", + fallback_depth: int = 0, +) -> Literal[True]: + if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: + raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") # noqa: TRY002 # frozen Phase 1 exception contract + if not isinstance(model, str): + for item in model: + can_object_call_model( + item, + llm_router, + models, + team_model_aliases, + team_id, + key_model_aliases, + object_type, + fallback_depth + 1, + ) + return True + return _single(model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type) + + +def _single( + model: str, + router: RegistryReader | None, + grants: Sequence[str], + team_aliases: Mapping[str, str] | None, + team: str | None, + key_aliases: Mapping[str, str] | None, + object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"], +) -> Literal[True]: + parent: Final = native_compaction_parent(model) + target: Final = ( + litellm.model_alias_map[model] + if model in litellm.model_alias_map + else router.registry().resolve(model, kind="alias") + if router is not None and model in router.registry().aliases() + else None + ) + after_team: Final = (team_aliases or {}).get(model, model) + after_key: Final = (key_aliases or {}).get(after_team, after_team) + after_global: Final = litellm.model_alias_map.get(after_key, after_key) + dispatched: Final = (key_aliases or {}).get(after_global, after_global) + key_applied: Final = after_key != after_team or dispatched != after_global + candidates: Final = ( + (dispatched,) + if key_applied + else (model, *((parent,) if parent is not None else ()), *((target,) if target else ())) + ) + owning: Final = () if key_applied else _owning(model, router, team, team_aliases) + if any(_helper(candidate, router, grants, team, team_aliases) for candidate in candidates): + return True + if any(_helper(candidate, router, grants, team, None) for candidate in owning): + return True + raise ModelAccessDeniedProxyException( + message=customer_model_access_denied_client_message(model) + if object_type == "customer" + else model_access_denied_client_message(model), + internal_message=f"{object_type} not allowed to access model. This {object_type} can only access models={list(grants)}. Tried to access {model}", + type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type), + param="model", + code=403, + ) + + +class _Grants(BaseModel, frozen=True): + models: tuple[str, ...] = () + + +async def can_user_call_model( + model: str, llm_router: RegistryReader | None, user_object: LiteLLM_UserTable | None +) -> Literal[True]: + if user_object is None: + return True + grants: Final = _Grants.model_validate(user_object, from_attributes=True).models + if "no-default-models" in grants: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model), + internal_message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + if not grants: + return True + return can_object_call_model(model, llm_router, grants) + + +def entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + if "tiered_pricing" in entry and entry["tiered_pricing"] is not None: + return True + return any(_priced(key, value) for key, value in entry.items()) + + +def _priced(key: str, value: object) -> bool: + if "cost_per" not in key: + return False + if isinstance(value, bool): + return False + if isinstance(value, (int, float)): + return value > 0 + if isinstance(value, dict): + return any( + isinstance(number, (int, float)) and not isinstance(number, bool) and number > 0 + for number in TypeAdapter(Mapping[object, object]).validate_python(value).values() + ) + return False diff --git a/tests/unit/_auth/support/route_enable_reference.py b/tests/unit/_auth/support/route_enable_reference.py new file mode 100644 index 00000000000..edbb78d3651 --- /dev/null +++ b/tests/unit/_auth/support/route_enable_reference.py @@ -0,0 +1,34 @@ +from typing import Final + +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting +from tests.unit._auth.support.route_reference import RouteChecks + + +def _setting(facts: RouteSetting | None, feature: RouteFeature) -> RouteEnableDecision | None: + if facts is None: + return RouteEnableDecision("ReadSettings", feature) + if facts.environment_present and not facts.premium: + return RouteEnableDecision("Deny", feature, "License") + if facts.secret is None: + return RouteEnableDecision("ReadSecret", feature) + if facts.secret: + return RouteEnableDecision("Deny", feature, "Disabled") + return None + + +def route_enable(facts: RouteEnableFacts) -> RouteEnableDecision: + if facts.failed is not None: + return RouteEnableDecision("PropagateHttp" if facts.failed == "Http" else "Allow") + if not facts.imported: + return RouteEnableDecision("ReadImport") + if RouteChecks.check_route_access(facts.context.route, facts.context.catalog.groups.get("management_routes", ())): + admin: Final = _setting(facts.admin, "Admin") + if admin is not None: + return admin + if not RouteChecks.is_llm_api_route(facts.context.route): + return RouteEnableDecision("Allow") + if facts.exemptions is None: + return RouteEnableDecision("ReadExemptions") + if facts.context.route in facts.exemptions: + return RouteEnableDecision("Allow") + return _setting(facts.llm, "Llm") or RouteEnableDecision("Allow") diff --git a/tests/unit/_auth/support/route_reference.py b/tests/unit/_auth/support/route_reference.py index 69fc32fc5b0..0361ab8cf23 100644 --- a/tests/unit/_auth/support/route_reference.py +++ b/tests/unit/_auth/support/route_reference.py @@ -4,6 +4,7 @@ from collections.abc import Collection, Iterable, Mapping, Sequence from typing import Final, Protocol, cast from fastapi import HTTPException, Request, status +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -19,7 +20,96 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks_organization import user_is_org_admin # pyright: ignore[reportUnknownVariableType] # preserve the original organization lookup + +def user_is_org_admin(request_data: Mapping[str, object], user_object: LiteLLM_UserTable | None = None) -> bool: + if user_object is None or user_object.organization_memberships is None: + return False + singular: Final = request_data.get("organization_id") + multiple: Final = request_data.get("organizations") + candidates: Final = ( + *((singular,) if singular is not None else ()), + *(TypeAdapter(tuple[object, ...]).validate_python(multiple) if isinstance(multiple, list) else ()), + ) + admins: Final = frozenset( + membership.organization_id + for membership in user_object.organization_memberships + if membership.user_role == LitellmUserRoles.ORG_ADMIN.value and membership.organization_id is not None # pyright: ignore[reportUnnecessaryComparison] # Phase 1 rejected absent legacy membership IDs + ) + return bool(candidates) and all(candidate in admins for candidate in candidates) + + +def managed_agent_route_allowed(route: str, method: str | None) -> bool: + if route in ("/agents", "/v1/agents"): + return method in (None, "GET", "HEAD") + if route in ( + "/realtime", + "/v1/realtime", + "/openai/v1/realtime", + "/realtime/translations", + "/v1/realtime/translations", + "/live/sessions", + "/v1/live/sessions", + "/openai/v1/live/sessions", + ): + return method in (None, "GET") + model_routes: Final = frozenset( + f"{prefix}/{operation}" + for prefix, operation in itertools.product( + ("", "/v1"), + ( + "chat/completions", + "completions", + "embeddings", + "responses", + "messages", + "messages/count_tokens", + "images/generations", + "images/edits", + "audio/transcriptions", + "audio/speech", + "moderations", + "rerank", + "decisions", + "systemone", + "extraction", + "ocr", + ), + ) + ) | frozenset( + ( + "/openai/v1/responses", + "/v2/rerank", + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", + "/cursor/chat/completions", + ) + ) + paths: Final = ( + "/engines/{model:path}/chat/completions", + "/engines/{model:path}/completions", + "/engines/{model:path}/embeddings", + "/openai/deployments/{model:path}/chat/completions", + "/openai/deployments/{model:path}/completions", + "/openai/deployments/{model:path}/embeddings", + "/openai/deployments/{model:path}/images/generations", + "/openai/deployments/{model:path}/images/edits", + "/v1beta/models/{model_name:path}:countTokens", + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/{model_name:path}:streamGenerateContent", + "/models/{model_name:path}:countTokens", + "/models/{model_name:path}:generateContent", + "/models/{model_name:path}:streamGenerateContent", + ) + if route in model_routes or RouteChecks.check_route_access(route, paths): + return method in (None, "POST") + mcp: Final = tuple( + candidate + for candidate in LiteLLMRoutes.mcp_inference_routes.value + if candidate not in ("/token", "/introspect") + ) + return RouteChecks.check_route_access(route, mcp) or RouteChecks.check_route_access( + route, LiteLLMRoutes.agent_inference_routes.value + ) class _ReferenceToken(Protocol): @@ -439,7 +529,10 @@ class RouteChecks: @staticmethod def custom_admin_only_route_check(route: str): - from litellm.proxy.proxy_server import general_settings, premium_user # pyright: ignore[reportUnknownVariableType] # preserve the legacy settings import in this reference + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # preserve the legacy settings import in this reference + premium_user, + ) if "admin_only_routes" in general_settings: if premium_user is not True: diff --git a/tests/unit/_auth/test_fallback_access.py b/tests/unit/_auth/test_fallback_access.py index 3a8f7975da9..ba7c7036c02 100644 --- a/tests/unit/_auth/test_fallback_access.py +++ b/tests/unit/_auth/test_fallback_access.py @@ -46,6 +46,16 @@ _BASE: Final = FallbackAuthorizationFacts(GroupedModelScope(_ALLOW)) ), pytest.param( replace(_BASE, team=GroupedModelScope(_DENY), assigned_models=("model",)), + "ReadAssignedScope", + id="assigned-reads-scope-before-override", + ), + pytest.param( + replace( + _BASE, + team=GroupedModelScope(_DENY), + assigned_models=("model",), + assigned_scope=replace(_DENY, grants=("model",)), + ), "Allow", id="assigned-overrides-team", ), diff --git a/tests/unit/_auth/test_model_access.py b/tests/unit/_auth/test_model_access.py index 8b96492e831..205e10ee2ba 100644 --- a/tests/unit/_auth/test_model_access.py +++ b/tests/unit/_auth/test_model_access.py @@ -23,8 +23,9 @@ from litellm._auth.model_access import ( ModelRegistry as Registry, ) from litellm.proxy._types import LiteLLM_UserTable, ModelAccessDeniedProxyException -from litellm.proxy.auth import auth_checks, model_access_adapter +from litellm.proxy.auth import model_access_adapter from tests.unit._auth.support import model_reference as reference +from tests.unit._auth.support import model_registry_reference _BASE: Final = Facts(("model",), object_type="key") _CASES: Final = ( @@ -220,7 +221,9 @@ def test_user_default_scope(present: bool, grants: tuple[str, ...], kind: str) - assert native.user_model_access(replace(_BASE, grants=grants), present=present).kind == kind -@pytest.mark.parametrize("policy", (auth_checks.can_user_call_model, model_access_adapter.can_user_call_model)) +@pytest.mark.parametrize( + "policy", (model_registry_reference.can_user_call_model, model_access_adapter.can_user_call_model) +) @pytest.mark.asyncio async def test_user_precheck_does_not_read_registry( policy: Callable[[str, MagicMock, LiteLLM_UserTable | None], Awaitable[Literal[True]]], diff --git a/tests/unit/_auth/test_route_enable.py b/tests/unit/_auth/test_route_enable.py new file mode 100644 index 00000000000..fd3e986c3ee --- /dev/null +++ b/tests/unit/_auth/test_route_enable.py @@ -0,0 +1,38 @@ +from collections.abc import Callable +from typing import Final, Literal + +import pytest + +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteSetting, route_enable +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_enable_reference import route_enable as reference + + +@pytest.mark.parametrize("policy", (reference, route_enable)) +@pytest.mark.parametrize( + "route", ("/unknown", "/models", "/v1/models", "/chat/completions", "/key/generate", "/user/info") +) +@pytest.mark.parametrize("imported", (False, True)) +@pytest.mark.parametrize("failed", (None, "Other", "Http")) +@pytest.mark.parametrize("exemptions", (None, (), ("/models", "/v1/models"))) +@pytest.mark.parametrize( + "setting", + ( + None, + RouteSetting(True, False), + RouteSetting(False, False), + RouteSetting(True, True), + RouteSetting(True, True, True), + RouteSetting(True, True, False), + ), +) +def test_optional_route_enable_decision_table( + policy: Callable[[RouteEnableFacts], RouteEnableDecision], + route: str, + imported: bool, + failed: Literal["Other", "Http"] | None, + exemptions: tuple[str, ...] | None, + setting: RouteSetting | None, +) -> None: + facts: Final = RouteEnableFacts(route_context(route), imported, failed, setting, setting, exemptions) + assert policy(facts) == reference(facts) diff --git a/tests/unit/_auth/test_routes.py b/tests/unit/_auth/test_routes.py index b4e12e0e259..31f7d5b6b37 100644 --- a/tests/unit/_auth/test_routes.py +++ b/tests/unit/_auth/test_routes.py @@ -10,6 +10,7 @@ from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.route_policy_adapter import route_context from tests.unit._auth.support.route_reference import RouteChecks as ReferenceRoutes +from tests.unit._auth.support.route_reference import managed_agent_route_allowed @dataclass(frozen=True, slots=True) @@ -169,8 +170,6 @@ def test_route_pattern_modes(route: str, pattern: str) -> None: ) @pytest.mark.parametrize("method", (None, "GET", "HEAD", "POST", "DELETE", "get")) def test_managed_actor_route_methods(route: str, method: str | None) -> None: - from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_route_allowed - assert routes.classify_route(route_context(route, method), "Managed") == managed_agent_route_allowed(route, method) diff --git a/tests/unit/_auth/test_unpriced_models.py b/tests/unit/_auth/test_unpriced_models.py index a25443fb8ad..fac21104611 100644 --- a/tests/unit/_auth/test_unpriced_models.py +++ b/tests/unit/_auth/test_unpriced_models.py @@ -5,10 +5,8 @@ from typing import Final import pytest from litellm._auth.unpriced_models import DeploymentPricing, PriceEntry, UnpricedModelFacts, unpriced_model -from litellm.proxy.auth.auth_checks import ( - _entry_has_priced_metric, # pyright: ignore[reportPrivateUsage] # compare the preserved pricing policy with native facts -) from litellm.proxy.auth.unpriced_model_adapter import price_entry +from tests.unit._auth.support.model_registry_reference import entry_has_priced_metric _BASE: Final = UnpricedModelFacts("model", True) _ROW: Final = DeploymentPricing((), True) @@ -41,7 +39,7 @@ def test_native_price_metrics_match_reference(entry: Mapping[str, object], expec deployments=(replace(_ROW, info_loaded=True, info=price_entry(entry)),), cost_deployments=(), ) - assert _entry_has_priced_metric(entry) == expected + assert entry_has_priced_metric(entry) == expected assert unpriced_model(facts).unpriced == (not expected) diff --git a/tests/unit/proxy/auth/test_managed_model_adapter.py b/tests/unit/proxy/auth/test_managed_model_adapter.py index f278e2d820c..d96e57ae375 100644 --- a/tests/unit/proxy/auth/test_managed_model_adapter.py +++ b/tests/unit/proxy/auth/test_managed_model_adapter.py @@ -5,16 +5,15 @@ 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.auth import 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: +async def test_managed_model_grants_preserve_public_denial(models: object) -> None: token: Final = UserAPIKeyAuth().model_copy( update={ "managed_agent_policy": AgentResponse( @@ -23,7 +22,7 @@ async def test_managed_model_grants_preserve_public_denial(native: bool, models: } ) cache: Final = UserApiKeyCache() - authorize: Final = managed_model_adapter.can_agent_call_model if native else auth_checks.can_agent_call_model + authorize: Final = managed_model_adapter.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)) @@ -38,17 +37,12 @@ async def test_managed_model_grants_preserve_public_denial(native: bool, models: @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 test_agent_empty_ceiling_and_model_denial_envelopes(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 - ) + authorize: Final = managed_model_adapter.check_agent_access_group_model_access with pytest.raises(ModelAccessDeniedProxyException) as denied: await authorize("denied", UserAPIKeyAuth(agent_id="agent"), None, resolve) assert denied.value.internal_message == ( @@ -70,8 +64,16 @@ async def test_managed_alias_dispatch_before_scope_check() -> None: 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: +@pytest.mark.parametrize( + ("models", "expected"), + ( + (None, frozenset[str]()), + ("model", frozenset[str]()), + ((), frozenset[str]()), + (("a", 42, "a", "b"), frozenset(("a", "b"))), + ), +) +def test_managed_listing_filters_non_string_grants(models: object, expected: frozenset[str]) -> None: token: Final = UserAPIKeyAuth().model_copy( update={ "managed_agent_policy": AgentResponse( @@ -79,4 +81,4 @@ def test_managed_listing_filters_non_string_grants(models: object) -> None: ) } ) - 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 + assert managed_model_adapter.managed_agent_model_grants(token) == expected diff --git a/tests/unit/proxy/auth/test_model_projection_adapter.py b/tests/unit/proxy/auth/test_model_projection_adapter.py new file mode 100644 index 00000000000..60817b2fc82 --- /dev/null +++ b/tests/unit/proxy/auth/test_model_projection_adapter.py @@ -0,0 +1,203 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from queue import SimpleQueue +from typing import Final, Literal, cast + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import grouped_model_access +from litellm.proxy._types import ModelAccessDeniedProxyException +from litellm.proxy.auth import model_access_adapter, model_projection_adapter +from litellm.router_api import RegistryReader +from tests.unit._auth.support import model_registry_reference as reference + + +@dataclass(frozen=True, slots=True) +class _Info: + team_id: str | None = "team" + team_public_model_name: str | None = "public" + + +@dataclass(frozen=True, slots=True) +class _Deployment: + model_info: _Info = field(default_factory=_Info) + + +@dataclass(frozen=True, slots=True) +class _Registry: + events: SimpleQueue[tuple[str, str | None]] + entries: tuple[str, ...] = ("registered",) + router_aliases: Mapping[str, str] = field(default_factory=lambda: {"alias": "target"}) + groups: Mapping[str, tuple[str, ...]] = field(default_factory=lambda: {"group": ("m",)}) + deployment: _Deployment | None = field(default_factory=_Deployment) + fail: str | None = None + + def observe(self, kind: str, model: str | None = None) -> None: + self.events.put((kind, model)) + if self.fail == kind: + raise RuntimeError(f"failed {kind}") + + def names(self) -> list[str]: # mutable-ok: the legacy sentinel resolver concatenates the registry's concrete list + self.observe("names") + return list(self.entries) + + def aliases(self) -> Mapping[str, str]: + self.observe("aliases") + return self.router_aliases + + def resolve(self, model: str, *, kind: Literal["alias"]) -> str | None: + self.observe("resolve", model) + return self.router_aliases.get(model) + + def access_groups(self, model_name: str, team_id: str | None) -> Mapping[str, tuple[str, ...]]: + self.observe("groups", model_name) + return self.groups + + def deployments(self, model: str, *, source: Literal["first_in_group"]) -> _Deployment | None: + self.observe("deployment", model) + return self.deployment + + +@dataclass(frozen=True, slots=True) +class _Router: + view: _Registry + + def registry(self) -> _Registry: + self.view.observe("registry") + return self.view + + +def _reader(view: _Registry) -> RegistryReader: + return cast(RegistryReader, _Router(view)) # noqa: TID251 # cast-ok: typed IO double implements only the registry methods consumed by auth + + +def _events(events: SimpleQueue[tuple[str, str | None]]) -> tuple[tuple[str, str | None], ...]: + return tuple(events.get() for _ in range(events.qsize())) + + +def _run(native: bool, model: str | tuple[str, ...], grants: tuple[str, ...], registry: _Registry) -> str: + router: Final = _reader(registry) + aliases: Final = {"z": "model_name_team_deploy", "a": "model_name_team_deploy", "team-alias": "target"} + try: + if native: + model_access_adapter.can_object_call_model(model, router, grants, aliases, "team", {"key": "target"}) + else: + reference.can_object_call_model( + model if isinstance(model, str) else list(model), + router, + list(grants), + aliases, + "team", + {"key": "target"}, + ) + return "allowed" + except ModelAccessDeniedProxyException as error: + return str((error.message, error.internal_message, error.type, error.param, error.code)) + + +@pytest.mark.parametrize( + "model", ("m", "alias", "global", "key", "team-alias", "model_name_team_deploy", ("m", "m"), ("denied", "m"), ()) +) +@pytest.mark.parametrize( + "grants", + ( + (), + ("m",), + ("target",), + ("group",), + ("all-team-models",), + ("all-team-models", "other"), + ("*",), + ("other/*", "provider/*"), + ("a",), + ("public",), + ), +) +def test_registry_observation_order_and_denials_match_legacy( + monkeypatch: pytest.MonkeyPatch, model: str | tuple[str, ...], grants: tuple[str, ...] +) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + registry: Final = _Registry(events) + + def provider(model: str) -> tuple[str, str, None, None]: + events.put(("provider", model)) + return "leaf", "provider", None, None + + monkeypatch.setattr(litellm, "get_llm_provider", provider) + monkeypatch.setattr(litellm, "model_alias_map", {"global": "target"}) + expected: Final = _run(False, model, grants, registry) + expected_events: Final = _events(events) + assert _run(True, model, grants, registry) == expected + assert _events(events) == expected_events + + +@pytest.mark.parametrize("native", (False, True)) +@pytest.mark.parametrize("failed", ("aliases", "groups", "names", "deployment")) +def test_registry_failure_stops_before_later_reads(native: bool, failed: str) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + registry: Final = _Registry(events, fail=failed) + with pytest.raises(RuntimeError, match=f"failed {failed}") as error: + _run(native, "model_name_team_deploy", ("all-team-models",), registry) + assert not isinstance(error.value, NativeAuthFailure) + assert _events(events)[-1][0] == failed + + +@pytest.mark.parametrize("native", (False, True)) +def test_depth_denial_never_reads_registry(native: bool) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + router: Final = _reader(_Registry(events, fail="registry")) + policy: Final = model_access_adapter.can_object_call_model if native else reference.can_object_call_model + with pytest.raises(Exception, match="max fallback depth exceeded"): + policy("m", router, (), fallback_depth=100) + assert _events(events) == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group_models", (("other",), ("other/*",), ("m",))) +async def test_group_retry_gets_fresh_registry_reads_even_for_duplicate_grants(group_models: tuple[str, ...]) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + router: Final = _reader(_Registry(events)) + facts: Final = model_access_adapter.project_model_facts("m", router, ("other",)) + assert _events(events) == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) + scope: Final = await model_projection_adapter.read_grouped_scope( + GroupedModelScope(facts, ("group",), group_models), router, merge=False + ) + assert _events(events)[:4] == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) + assert grouped_model_access(scope.facts, scope.group_ids, scope.group_models).kind == ( + "Allow" if group_models == ("m",) else "Deny" + ) + + +def test_native_provider_failure_is_not_a_missing_provider(monkeypatch: pytest.MonkeyPatch) -> None: + def provider(model: str) -> tuple[str, str, None, None]: + raise NativeAuthRuntimeError("terminal provider failure") + + monkeypatch.setattr(litellm, "get_llm_provider", provider) + with pytest.raises(NativeAuthRuntimeError, match="terminal provider failure"): + model_access_adapter.can_object_call_model("model", None, ("provider/*",)) + + +@pytest.mark.parametrize("native", (False, True)) +def test_compaction_parent_is_read_only_after_child_denial(native: bool) -> None: + from litellm.router_strategy.complexity_router.context_compaction import native_compaction_call + + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + with native_compaction_call("parent", "child"): + assert _run(native, "child", ("parent",), _Registry(events)) == "allowed" + assert _events(events) == ( + ("registry", None), + ("aliases", None), + ("registry", None), + ("groups", "child"), + ("registry", None), + ("groups", "parent"), + ) + + +def test_large_model_lists_do_not_consume_python_recursion_depth() -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + assert _run(True, ("m",) * 180, ("m",), _Registry(events)) == "allowed" + assert _events(events) == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) * 180 diff --git a/tests/unit/proxy/auth/test_route_enable_adapter.py b/tests/unit/proxy/auth/test_route_enable_adapter.py new file mode 100644 index 00000000000..56f1bc96ee5 --- /dev/null +++ b/tests/unit/proxy/auth/test_route_enable_adapter.py @@ -0,0 +1,88 @@ +from queue import SimpleQueue +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_enable import RouteFeature, RouteSetting +from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.auth.route_enable_adapter import RouteEnableReaders, check_route_enabled + + +@pytest.mark.parametrize( + ("route", "expected"), + ( + ("/unknown", ("import",)), + ("/models", ("import", "exemptions")), + ("/chat/completions", ("import", "exemptions", "Llm", "secret-Llm")), + ("/key/generate", ("import", "Admin", "secret-Admin")), + ), +) +def test_optional_route_reads_are_lazy(route: str, expected: tuple[str, ...]) -> None: + events: Final[SimpleQueue[str]] = SimpleQueue() + + def enterprise() -> object: + events.put("import") + return object() + + def setting(feature: RouteFeature) -> RouteSetting: + events.put(feature) + return RouteSetting(False, False) + + def secret(feature: RouteFeature) -> bool: + events.put(f"secret-{feature}") + return False + + def exemptions(_enterprise: object) -> tuple[str, ...]: + events.put("exemptions") + return ("/models",) + + check_route_enabled(route, readers=RouteEnableReaders(enterprise, setting, secret, exemptions)) + assert tuple(events.get() for _ in range(events.qsize())) == expected + + +@pytest.mark.parametrize(("route", "label"), (("/key/generate", "ADMIN"), ("/chat/completions", "LLM API"))) +@pytest.mark.parametrize("premium", (False, True)) +def test_optional_route_denial_envelopes(route: str, label: str, premium: bool) -> None: + readers: Final = RouteEnableReaders( + enterprise=lambda: object(), + setting=lambda _: RouteSetting(True, premium), + secret=lambda _: True, + exemptions=lambda _: (), + ) + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + if premium: + disabled: Final = "Management" if label == "ADMIN" else label + assert (error.value.status_code, error.value.detail) == ( + 403, + f"{disabled} routes are disabled for this instance.", + ) + else: + assert (error.value.status_code, error.value.detail) == ( + 500, + f"\U0001f6a8\U0001f6a8\U0001f6a8 DISABLING {label} ENDPOINTS is an Enterprise feature\n\U0001f6a8 {CommonProxyErrors.not_premium_user.value}", + ) + + +@pytest.mark.parametrize( + "error", + ( + ImportError("optional package"), + RuntimeError("settings unavailable"), + HTTPException(418, "read failure"), + NativeAuthRuntimeError("terminal native failure"), + ), +) +def test_optional_import_failure_policy(error: Exception) -> None: + def enterprise() -> object: + raise error + + readers: Final = RouteEnableReaders(enterprise=enterprise) + if isinstance(error, (HTTPException, NativeAuthRuntimeError)): + with pytest.raises(type(error)) as caught: + check_route_enabled("/key/generate", readers=readers) + assert caught.value is error + else: + assert check_route_enabled("/key/generate", readers=readers) is None