diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 9bc7b42a37b..81d4454b58a 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -68,7 +68,8 @@ mod _native { request_auth_agent_grants, request_auth_agent_group_load, request_auth_agent_invocation, request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption, request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start, - request_auth_budgets, request_auth_common_budget, request_auth_common_budget_scopes, + request_auth_budgets, request_auth_codex_model, request_auth_codex_normalization, + request_auth_codex_turn, request_auth_common_budget, request_auth_common_budget_scopes, request_auth_common_flow, request_auth_common_read, request_auth_context, request_auth_context_end_user, request_auth_context_member, request_auth_context_models, request_auth_context_percentage, request_auth_context_temporary, @@ -90,7 +91,7 @@ mod _native { request_auth_resolved_models, request_auth_route_access, request_auth_route_enable, request_auth_route_role, request_auth_route_role_api, request_auth_route_role_expand, request_auth_route_role_inside, request_auth_route_role_jwt, - request_auth_route_role_listed, request_auth_scheme_credential, + request_auth_route_role_listed, request_auth_route_scope, request_auth_scheme_credential, request_auth_select_api_key, request_auth_select_websocket_api_key, request_auth_spend_counter_policy, request_auth_start, request_auth_team_aliases, request_auth_team_override, request_auth_unpriced_model, diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index 9d07c977f29..3f8378fc66a 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -279,6 +279,51 @@ fn map_error(error: core::Error) -> PyErr { } } +#[pyfunction] +pub(crate) fn request_auth_codex_normalization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::codex_normalization::Facts = from_py(facts)?; + let decision = py + .detach(move || core::codex_normalization::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_codex_turn( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::codex_turn::Request = from_py(facts)?; + let decision = py + .detach(move || core::codex_turn::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_scope( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_scope::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::route_scope::decide(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_codex_model( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::codex_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::codex_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + fn project(py: Python<'_>, step: Result) -> PyResult { let step = step.map_err(map_error)?; Ok(( diff --git a/litellm-rust/crates/request-auth-core/src/codex_models.rs b/litellm-rust/crates/request-auth-core/src/codex_models.rs new file mode 100644 index 00000000000..a1e8daed5dd --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/codex_models.rs @@ -0,0 +1,348 @@ +use crate::{Error, model_helpers}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Routes, + NativeRoutes, + Chatgpt, + Grants, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Candidate { + pub provider_loaded: bool, + pub provider: Option, + pub model: Option, + pub name: Option, + pub public_loaded: bool, + pub public_name: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Grants { + pub key: Vec, + pub team: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub model: Option, + pub router_present: bool, + pub team_id: Option, + pub bundled: Option>, + pub priced: Option, + pub candidates: Option>, + pub alias_loaded: bool, + pub alias: Option, + pub grants: Option, + pub prefixed_grants: Option, + pub bare: Vec, + pub prefixed: Vec, + pub route: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Field { + Provider, + Model, + Name, + PublicName, +} + +#[derive(Debug, Serialize, PartialEq, Eq)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { + allowed: bool, + target: Option, + }, + ReadBundled, + ReadCost { + model: String, + }, + ReadCandidates { + model: String, + }, + ReadAlias { + model: String, + }, + ReadCandidate { + index: usize, + field: Field, + }, + ReadGrants { + prefixed: bool, + }, + ReadHelper { + index: usize, + model: String, + prefixed: bool, + }, +} + +fn complete(allowed: bool) -> Decision { + Decision::Complete { + allowed, + target: None, + } +} + +fn grants(facts: &Facts, model: &str, prefixed: bool) -> Result, Error> { + let selected = if prefixed { + &facts.prefixed_grants + } else { + &facts.grants + }; + let Some(grants) = selected else { + return Ok(Err(Decision::ReadGrants { prefixed })); + }; + let scopes = if prefixed { + &facts.prefixed + } else { + &facts.bare + }; + for (index, expected) in [&grants.key, &grants.team].into_iter().enumerate() { + let Some(scope) = scopes.get(index) else { + return Ok(Err(Decision::ReadHelper { + index, + model: model.to_owned(), + prefixed, + })); + }; + if scope.model != model + || scope.grants.as_ref() != Some(expected) + || scope.kind != model_helpers::Kind::Helper + { + return Err(Error::InvalidAnswer); + } + match model_helpers::decide(scope)? { + model_helpers::Decision::Complete { allowed: false } => return Ok(Ok(false)), + model_helpers::Decision::Complete { allowed: true } => (), + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Ok(true)) +} + +fn candidate( + candidate: &Candidate, + index: usize, + facts: &Facts, + model: &str, +) -> Result { + let read = |field| Decision::ReadCandidate { index, field }; + if !candidate.provider_loaded { + return Err(read(Field::Provider)); + } + let chatgpt = if candidate.provider.as_deref() == Some("chatgpt") { + true + } else { + let backend = candidate + .model + .as_deref() + .ok_or_else(|| read(Field::Model))?; + backend.split('/').next() == Some("chatgpt") + }; + if facts.kind == Kind::Chatgpt { + return Ok(chatgpt); + } + if chatgpt { + return Ok(false); + } + let alias = facts + .alias + .as_deref() + .filter(|alias| !alias.is_empty()) + .unwrap_or(model); + let name = candidate.name.as_deref().ok_or_else(|| read(Field::Name))?; + if name == model || name == alias { + return Ok(true); + } + if !candidate.public_loaded { + return Err(read(Field::PublicName)); + } + Ok(candidate + .public_name + .as_deref() + .is_some_and(|name| name == model || name == alias)) +} + +pub fn decide(facts: &Facts) -> Result { + if facts.kind == Kind::NativeRoutes + && facts.route.as_deref() != Some("/backend-api/codex/responses") + { + return Ok(complete(false)); + } + let Some(model) = facts.model.as_deref().filter(|_| facts.router_present) else { + return Ok(complete(false)); + }; + if facts.kind == Kind::Grants { + return Ok(match grants(facts, model, false)? { + Ok(allowed) => complete(allowed), + Err(read) => read, + }); + } + if facts.kind == Kind::Chatgpt { + if model.starts_with("chatgpt/") { + return Ok(complete(false)); + } + match grants(facts, model, false)? { + Ok(true) => return Ok(complete(false)), + Err(read) => return Ok(read), + Ok(false) => (), + } + } else { + let Some(bundled) = &facts.bundled else { + return Ok(Decision::ReadBundled); + }; + if bundled.iter().any(|name| name == model) { + return Ok(complete(false)); + } + match facts.priced { + None => { + return Ok(Decision::ReadCost { + model: format!("chatgpt/{model}"), + }); + } + Some(true) => return Ok(complete(false)), + Some(false) => (), + } + } + let target = if facts.kind == Kind::Chatgpt { + format!("chatgpt/{model}") + } else { + model.to_owned() + }; + let Some(candidates) = &facts.candidates else { + return Ok(Decision::ReadCandidates { model: target }); + }; + if matches!(facts.kind, Kind::Routes | Kind::NativeRoutes) && !facts.alias_loaded { + return Ok(Decision::ReadAlias { + model: model.to_owned(), + }); + } + if candidates.is_empty() { + return Ok(complete(false)); + } + for (index, value) in candidates.iter().enumerate() { + match candidate(value, index, facts, model) { + Ok(false) => return Ok(complete(false)), + Err(read) => return Ok(read), + Ok(true) => (), + } + } + if matches!(facts.kind, Kind::Routes | Kind::NativeRoutes) { + return Ok(complete(true)); + } + Ok(match grants(facts, &target, true)? { + Ok(true) => Decision::Complete { + allowed: true, + target: Some(target), + }, + Ok(false) => complete(false), + Err(read) => read, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn base() -> Facts { + Facts { + kind: Kind::Routes, + model: Some("model".into()), + router_present: true, + team_id: None, + bundled: None, + priced: None, + candidates: None, + alias_loaded: false, + alias: None, + grants: None, + prefixed_grants: None, + bare: vec![], + prefixed: vec![], + route: None, + } + } + + #[rstest] + #[case::no_model(None, true, false)] + #[case::no_router(Some("model"), false, false)] + #[case::needs_bundled_catalog(Some("model"), true, true)] + fn routing_starts_lazily( + base: Facts, + #[case] model: Option<&str>, + #[case] router: bool, + #[case] read: bool, + ) { + let facts = Facts { + model: model.map(str::to_owned), + router_present: router, + ..base + }; + let expected = if read { + Decision::ReadBundled + } else { + Decision::Complete { + allowed: false, + target: None, + } + }; + assert_eq!(decide(&facts).unwrap(), expected); + } + + #[rstest] + #[case::bundled("model", false)] + #[case::other("other", true)] + fn bundled_route_stops_before_cost(base: Facts, #[case] bundled: &str, #[case] read: bool) { + let facts = Facts { + bundled: Some(vec![bundled.into()]), + ..base + }; + let expected = if read { + Decision::ReadCost { + model: "chatgpt/model".into(), + } + } else { + Decision::Complete { + allowed: false, + target: None, + } + }; + assert_eq!(decide(&facts).unwrap(), expected); + } + + #[rstest] + #[case::alias_pending(false, true)] + #[case::empty_candidates(true, false)] + fn empty_candidates_still_read_alias( + base: Facts, + #[case] alias_loaded: bool, + #[case] read: bool, + ) { + let facts = Facts { + bundled: Some(vec![]), + priced: Some(false), + candidates: Some(vec![]), + alias_loaded, + ..base + }; + let expected = if read { + Decision::ReadAlias { + model: "model".into(), + } + } else { + Decision::Complete { + allowed: false, + target: None, + } + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/codex_normalization.rs b/litellm-rust/crates/request-auth-core/src/codex_normalization.rs new file mode 100644 index 00000000000..041c0cca72d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/codex_normalization.rs @@ -0,0 +1,201 @@ +use crate::{Error, codex_models, codex_turn, route_access::RouteContext, route_enable}; +use serde::{Deserialize, Serialize}; + +const NATIVE_ROUTE: &str = "/backend-api/codex/responses"; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Alias, + Prefix, +} + +#[derive(Deserialize)] +pub struct Facts { + pub kind: Kind, + pub route: String, + pub request_present: bool, + pub router_present: Option, + pub context: Option, + pub eu_jev: Option, + pub resolved: Option, + pub marked: bool, + pub model_loaded: bool, + pub model: Option, + pub explicit: Option, + pub provider: Option, + pub routing: Option, + pub body_present: Option, + pub settings_loaded: bool, + pub target: Option, + pub provider_after: Option, + pub turn: Option, + pub prefix: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete, + ReadRouter, + ReadContext, + ReadEuJev, + ReadResolved, + MarkResolved, + ReadModel, + ReadExplicit, + ReadProvider, + ReadRouting, + ReadBody, + ReadSettings, + ReadProviderAfter, + ReadTurn, + ReadPrefix, + Rewrite { target: String }, +} + +fn routed(facts: &codex_models::Facts) -> Result { + match codex_models::decide(facts)? { + codex_models::Decision::Complete { allowed, .. } => Ok(allowed), + _ => Err(Error::InvalidAnswer), + } +} + +pub fn decide(facts: &Facts) -> Result { + if !facts.request_present { + return Ok(Decision::Complete); + } + if facts.kind == Kind::Prefix && (facts.route != NATIVE_ROUTE || facts.model.is_none()) { + return Ok(Decision::Complete); + } + match facts.router_present { + None => return Ok(Decision::ReadRouter), + Some(false) => return Ok(Decision::Complete), + Some(true) => (), + } + if facts.kind == Kind::Prefix { + let Some(prefix) = &facts.prefix else { + return Ok(Decision::ReadPrefix); + }; + return match codex_models::decide(prefix)? { + codex_models::Decision::Complete { + target: Some(target), + .. + } => Ok(Decision::Rewrite { target }), + codex_models::Decision::Complete { .. } => Ok(Decision::Complete), + _ => Err(Error::InvalidAnswer), + }; + } + let Some(context) = &facts.context else { + return Ok(Decision::ReadContext); + }; + match route_enable::llm_route(context, facts.eu_jev)? { + route_enable::LlmDecision::Deny => return Ok(Decision::Complete), + route_enable::LlmDecision::ReadEuJev => return Ok(Decision::ReadEuJev), + route_enable::LlmDecision::Allow => (), + } + match facts.resolved { + None => return Ok(Decision::ReadResolved), + Some(true) => return Ok(Decision::Complete), + Some(false) => (), + } + if !facts.marked { + return Ok(Decision::MarkResolved); + } + if !facts.model_loaded { + return Ok(Decision::ReadModel); + } + match facts.explicit { + None => return Ok(Decision::ReadExplicit), + Some(true) => return Ok(Decision::Complete), + Some(false) => (), + } + let Some(provider) = facts.provider else { + return Ok(Decision::ReadProvider); + }; + if provider { + if facts.route != NATIVE_ROUTE { + return Ok(Decision::Complete); + } + let Some(routing) = &facts.routing else { + return Ok(Decision::ReadRouting); + }; + if !routed(routing)? { + return Ok(Decision::Complete); + } + } + let Some(model) = &facts.model else { + return Ok(Decision::Complete); + }; + match facts.body_present { + None => return Ok(Decision::ReadBody), + Some(false) => return Ok(Decision::Complete), + Some(true) => (), + } + if !facts.settings_loaded { + return Ok(Decision::ReadSettings); + } + let Some(target) = facts + .target + .as_ref() + .filter(|target| !target.is_empty() && *target != model) + else { + return Ok(Decision::Complete); + }; + let Some(provider_after) = facts.provider_after else { + return Ok(Decision::ReadProviderAfter); + }; + if provider_after { + let Some(turn) = &facts.turn else { + return Ok(Decision::ReadTurn); + }; + match codex_turn::turn(turn)? { + codex_turn::Decision::Complete { allowed: true } => (), + codex_turn::Decision::Complete { allowed: false } => return Ok(Decision::Complete), + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Rewrite { + target: target.clone(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing_request(false, NATIVE_ROUTE, Some("model"), Decision::Complete)] + #[case::wrong_route(true, "/v1/responses", Some("model"), Decision::Complete)] + #[case::non_string(true, NATIVE_ROUTE, None, Decision::Complete)] + #[case::native(true, NATIVE_ROUTE, Some("model"), Decision::ReadRouter)] + fn prefix_does_not_read_router_on_early_exit( + #[case] request: bool, + #[case] route: &str, + #[case] model: Option<&str>, + #[case] expected: Decision, + ) { + let facts = Facts { + kind: Kind::Prefix, + route: route.into(), + request_present: request, + router_present: None, + context: None, + eu_jev: None, + resolved: None, + marked: false, + model_loaded: true, + model: model.map(str::to_owned), + explicit: None, + provider: None, + routing: None, + body_present: None, + settings_loaded: false, + target: None, + provider_after: None, + turn: None, + prefix: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/codex_turn.rs b/litellm-rust/crates/request-auth-core/src/codex_turn.rs new file mode 100644 index 00000000000..4215b762193 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/codex_turn.rs @@ -0,0 +1,230 @@ +use crate::{Error, codex_models, route_access, route_enable, route_scope}; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +pub struct KeyFacts { + pub proxy: Option, + pub enabled: Option, + pub virtual_route: Option, + pub role_loaded: bool, + pub role: Option, + pub admin: Option, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum JwtRead { + Returned, + Exception, +} + +#[derive(Deserialize)] +pub struct AccessFacts { + pub key: KeyFacts, + pub jwt_token_present: bool, + pub jwt: Option, +} + +#[derive(Deserialize)] +pub struct Credentials { + pub bearer_present: bool, + pub gateway_present: bool, + pub equal: bool, +} + +#[derive(Deserialize)] +pub struct Facts { + pub model: Option, + pub native: bool, + pub credentials: Option, + pub access: Option, + pub routing: Option, + pub catalog_loaded: bool, + pub slugs: Option>, +} + +#[derive(Deserialize)] +pub enum Request { + Key(Box), + Access(Box), + Turn(Box), + Jwt(JwtRead), + JwtToken { + claims_present: bool, + token_present: bool, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + Failure { message: String }, + ReadProxy, + ReadEnabled, + ReadVirtual, + ReadRole, + ReadAdmin, + ReadCredentials, + ReadAccess, + ReadRouting, + ReadJwt, + ReadCatalog, +} + +fn complete(allowed: bool) -> Decision { + Decision::Complete { allowed } +} + +fn scope(facts: &route_scope::Facts) -> Result, Error> { + match route_scope::decide(facts) { + route_scope::Decision::Allow => Ok(Ok(true)), + route_scope::Decision::Deny => Ok(Ok(false)), + route_scope::Decision::Failure { message } => Ok(Err(message)), + _ => Err(Error::InvalidAnswer), + } +} + +pub fn key(facts: &KeyFacts) -> Result { + let Some(proxy) = &facts.proxy else { + return Ok(Decision::ReadProxy); + }; + match scope(proxy)? { + Ok(true) => (), + Ok(false) => return Ok(complete(false)), + Err(message) => return Ok(Decision::Failure { message }), + } + let Some(enabled) = &facts.enabled else { + return Ok(Decision::ReadEnabled); + }; + match route_enable::decide(enabled)? { + route_enable::Decision::Allow => (), + route_enable::Decision::Deny { .. } | route_enable::Decision::PropagateHttp => { + return Ok(complete(false)); + } + _ => return Err(Error::InvalidAnswer), + } + let Some(virtual_route) = &facts.virtual_route else { + return Ok(Decision::ReadVirtual); + }; + if !route_access::virtual_route(virtual_route)?.allowed { + return Ok(complete(false)); + } + if !facts.role_loaded { + return Ok(Decision::ReadRole); + } + if facts.role.as_deref() == Some("proxy_admin") { + return Ok(complete(true)); + } + let Some(admin) = &facts.admin else { + return Ok(Decision::ReadAdmin); + }; + Ok(match scope(admin)? { + Ok(allowed) => complete(allowed), + Err(message) => Decision::Failure { message }, + }) +} + +pub fn access(facts: &AccessFacts) -> Result { + match key(&facts.key)? { + Decision::Complete { allowed: true } => (), + Decision::Complete { allowed: false } => return Ok(complete(false)), + _ => return Err(Error::InvalidAnswer), + } + if !facts.jwt_token_present { + return Ok(complete(true)); + } + Ok(match facts.jwt { + None => Decision::ReadJwt, + Some(read) => jwt(read), + }) +} + +fn jwt(read: JwtRead) -> Decision { + complete(matches!(read, JwtRead::Returned)) +} + +pub fn turn(facts: &Facts) -> Result { + if facts.native { + let Some(credentials) = &facts.credentials else { + return Ok(Decision::ReadCredentials); + }; + if !credentials.bearer_present || !credentials.gateway_present || credentials.equal { + return Ok(complete(false)); + } + } + let Some(access_facts) = &facts.access else { + return Ok(Decision::ReadAccess); + }; + let Some(model) = &facts.model else { + return Ok(complete(false)); + }; + let Some(routing) = &facts.routing else { + return Ok(Decision::ReadRouting); + }; + match codex_models::decide(routing)? { + codex_models::Decision::Complete { allowed: true, .. } => (), + codex_models::Decision::Complete { allowed: false, .. } => return Ok(complete(false)), + _ => return Err(Error::InvalidAnswer), + } + match access(access_facts)? { + Decision::Complete { allowed: true } => (), + other => return Ok(other), + } + if !facts.catalog_loaded { + return Ok(Decision::ReadCatalog); + } + Ok(complete( + facts + .slugs + .as_ref() + .is_some_and(|slugs| !slugs.contains(model)), + )) +} + +pub fn decide(request: Request) -> Result { + match request { + Request::Key(facts) => key(&facts), + Request::Access(facts) => access(&facts), + Request::Turn(facts) => turn(&facts), + Request::Jwt(read) => Ok(jwt(read)), + Request::JwtToken { + claims_present, + token_present, + } => Ok(complete(claims_present && !token_present)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, Decision::ReadCredentials)] + #[case::no_bearer(Some(Credentials { bearer_present: false, gateway_present: true, equal: false }), complete(false))] + #[case::no_gateway(Some(Credentials { bearer_present: true, gateway_present: false, equal: false }), complete(false))] + #[case::equal(Some(Credentials { bearer_present: true, gateway_present: true, equal: true }), complete(false))] + #[case::separate(Some(Credentials { bearer_present: true, gateway_present: true, equal: false }), Decision::ReadAccess)] + fn credentials_precede_route_reads( + #[case] credentials: Option, + #[case] expected: Decision, + ) { + let facts = Facts { + model: Some("model".into()), + native: true, + credentials, + access: None, + routing: None, + catalog_loaded: false, + slugs: None, + }; + assert_eq!(turn(&facts).unwrap(), expected); + } + + #[rstest] + #[case::returned(JwtRead::Returned, true)] + #[case::exception(JwtRead::Exception, false)] + fn jwt_wrapper_uses_only_validator_outcome(#[case] read: JwtRead, #[case] allowed: bool) { + assert_eq!(jwt(read), complete(allowed)); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index 592e7268c52..6363f426296 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -10,6 +10,9 @@ pub mod budget_alerts; pub mod budget_context; pub mod budget_lookup; pub mod budget_pricing; +pub mod codex_models; +pub mod codex_normalization; +pub mod codex_turn; pub mod common_budgets; pub mod common_flow; pub mod credential; @@ -43,6 +46,7 @@ pub mod resolved_models; pub mod route_access; pub mod route_enable; pub mod route_roles; +pub mod route_scope; pub mod spend_counter_policy; pub mod team_aliases; pub mod team_override; 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 06211ca5cb9..8c9e7cb1b87 100644 --- a/litellm-rust/crates/request-auth-core/src/model_access.rs +++ b/litellm-rust/crates/request-auth-core/src/model_access.rs @@ -213,10 +213,16 @@ pub fn candidate_names( if !registry.present || !model.starts_with(&format!("model_name_{team}_")) { return (candidates, vec![]); } - let record = registry.records.get(model); - let public = record + let Some(record) = registry + .records + .get(model) .filter(|record| record.deployment_team.as_deref() == Some(team)) - .and_then(|record| record.deployment_public_name.clone()) + else { + return (candidates, vec![]); + }; + let public = record + .deployment_public_name + .clone() .filter(|name| !name.is_empty()); let ordered: Vec<&String> = aliases.team_order.as_ref().map_or_else( || aliases.team.keys().collect(), @@ -454,6 +460,9 @@ pub fn decide(request: Request) -> Result { } } Request::Customer(facts) => { + if facts.registry.evaluations.is_some() { + return Ok(Response::Decision(check(&facts)?)); + } for name in &facts.requested { let target = live_team_target(name, &facts); if target != *name && facts.grants.contains(name) { @@ -634,6 +643,49 @@ mod tests { use super::*; use rstest::rstest; + #[rstest] + #[case::missing(None, false)] + #[case::foreign(Some("other"), false)] + #[case::owned(Some("team"), true)] + fn ownership_aliases_require_matching_deployment( + #[case] owner: Option<&str>, + #[case] owned: bool, + ) { + let model = "model_name_team_private"; + let aliases = Aliases { + team: BTreeMap::from([ + ("z-last".into(), model.into()), + ("a-first".into(), model.into()), + ]), + team_order: Some(vec!["z-last".into(), "a-first".into()]), + ..Default::default() + }; + let records = owner + .map(|team| { + BTreeMap::from([( + model.into(), + ModelRecord { + deployment_team: Some(team.into()), + deployment_public_name: Some("public".into()), + ..Default::default() + }, + )]) + }) + .unwrap_or_default(); + let registry = Registry { + present: true, + records, + ..Default::default() + }; + let (_, names) = candidate_names(model, &aliases, ®istry, Some("team")); + let expected: Vec = if owned { + vec!["public".into(), "z-last".into(), "a-first".into()] + } else { + vec![] + }; + assert_eq!(names, expected); + } + #[rstest] #[case::absent(false, vec!["no-default-models"], Decision::Allow)] #[case::restricted(true, vec!["no-default-models"], Decision::NoDefaultModels)] diff --git a/litellm-rust/crates/request-auth-core/src/model_projection.rs b/litellm-rust/crates/request-auth-core/src/model_projection.rs index 829b72c703b..77a5e286715 100644 --- a/litellm-rust/crates/request-auth-core/src/model_projection.rs +++ b/litellm-rust/crates/request-auth-core/src/model_projection.rs @@ -36,12 +36,18 @@ pub struct Evaluation { #[derive(Serialize)] #[serde(tag = "kind")] pub enum Decision { - Complete { decision: ModelDecision }, - Read { read: ReadKind, model: String }, + Complete { + decision: ModelDecision, + }, + Read { + read: ReadKind, + model: String, + team_id: Option, + }, } enum Pending { - Read(ReadKind, String), + Read(ReadKind, String, Option), Invalid(Error), } impl From for Pending { @@ -53,11 +59,12 @@ impl From for Pending { struct Reader<'a> { reads: &'a [Reading], index: usize, + team_id: Option, } 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())); + return Err(Pending::Read(kind, model.to_owned(), self.team_id.clone())); }; if value.kind != kind || value.model != model { return Err(Pending::Invalid(Error::InvalidAnswer)); @@ -200,12 +207,67 @@ fn check(facts: &Facts, reader: &mut Reader<'_>) -> Result) -> Result { + for model in &facts.requested { + let target = facts.aliases.team.get(model).unwrap_or(model); + let target = if facts.registry.present + && target.starts_with(&format!( + "model_name_{}_", + facts.team_id.as_deref().unwrap_or("None") + )) + && facts.aliases.team.contains_key(model) + && !reader.read(ReadKind::Names, model)?.values.contains(target) + { + model + } else { + target + }; + if target != model && facts.grants.contains(model) { + continue; + } + let scope = Facts { + requested: vec![target.clone()], + team_id: None, + is_list: false, + aliases: model_access::Aliases { + key: facts.aliases.key.clone(), + global: facts.aliases.global.clone(), + ..Default::default() + }, + ..facts.clone() + }; + let decision = check(&scope, reader)?; + if !decision.allowed() { + return Ok(decision); + } + } + Ok(ModelDecision::Allow) +} + pub fn decide(facts: &Facts, reads: &[Reading]) -> Result { - let mut reader = Reader { reads, index: 0 }; - match check(facts, &mut reader) { + let is_customer = facts.object_type == model_access::ObjectType::Customer; + let mut reader = Reader { + reads, + index: 0, + team_id: if is_customer { + None + } else { + facts.team_id.clone() + }, + }; + let decision = if is_customer { + customer(facts, &mut reader) + } else { + check(facts, &mut reader) + }; + match decision { 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::Read(read, model, team_id)) => Ok(Decision::Read { + read, + model, + team_id, + }), Err(Pending::Invalid(error)) => Err(error), } } diff --git a/litellm-rust/crates/request-auth-core/src/route_access.rs b/litellm-rust/crates/request-auth-core/src/route_access.rs index b1eb51177bf..6f090c91ad6 100644 --- a/litellm-rust/crates/request-auth-core/src/route_access.rs +++ b/litellm-rust/crates/request-auth-core/src/route_access.rs @@ -311,7 +311,7 @@ fn managed(context: &RouteContext) -> Result { || catalog.access(route, "agent_inference_routes")?) } -fn virtual_route(facts: &VirtualRouteFacts) -> Result { +pub(crate) fn virtual_route(facts: &VirtualRouteFacts) -> Result { let Some(allowed) = &facts.allowed_routes else { return Ok(Decision::allow()); }; diff --git a/litellm-rust/crates/request-auth-core/src/route_scope.rs b/litellm-rust/crates/request-auth-core/src/route_scope.rs new file mode 100644 index 00000000000..2275dbc4881 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_scope.rs @@ -0,0 +1,97 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Proxy, + Admin, +} + +#[derive(Clone, Deserialize)] +pub struct Values { + pub text: Option, + pub names: Vec, + pub failure: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub route: String, + pub present: bool, + pub premium: bool, + pub warned: bool, + pub values: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny, + Warn, + ReadValues, + Failure { message: String }, +} + +pub fn decide(facts: &Facts) -> Decision { + if !facts.present { + return Decision::Allow; + } + if !facts.premium { + if !facts.warned { + return Decision::Warn; + } + if facts.kind == Kind::Admin { + return Decision::Allow; + } + } + let Some(values) = &facts.values else { + return Decision::ReadValues; + }; + if let Some(message) = &values.failure { + return Decision::Failure { + message: message.clone(), + }; + } + let member = values.text.as_ref().map_or_else( + || values.names.contains(&facts.route), + |text| text.contains(&facts.route), + ); + if member == (facts.kind == Kind::Proxy) { + Decision::Allow + } else { + Decision::Deny + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::proxy_warn(Kind::Proxy, false, false, Decision::Warn)] + #[case::admin_warn(Kind::Admin, false, false, Decision::Warn)] + #[case::proxy_after_warn(Kind::Proxy, false, true, Decision::ReadValues)] + #[case::admin_after_warn(Kind::Admin, false, true, Decision::Allow)] + #[case::licensed_proxy(Kind::Proxy, true, false, Decision::ReadValues)] + #[case::licensed_admin(Kind::Admin, true, false, Decision::ReadValues)] + fn licensing_retains_distinct_warning_behavior( + #[case] kind: Kind, + #[case] premium: bool, + #[case] warned: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(&Facts { + kind, + route: "/responses".into(), + present: true, + premium, + warned, + values: None + }), + expected + ); + } +} diff --git a/litellm/_auth/codex_models.py b/litellm/_auth/codex_models.py new file mode 100644 index 00000000000..f4b2e13613d --- /dev/null +++ b/litellm/_auth/codex_models.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_helpers import ModelHelperFacts +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + +CodexModelKind: TypeAlias = Literal["Routes", "NativeRoutes", "Chatgpt", "Grants"] +CandidateField: TypeAlias = Literal["Provider", "Model", "Name", "PublicName"] + + +@dataclass(frozen=True, slots=True) +class CodexCandidate: + provider_loaded: bool = False + provider: str | None = None + model: str | None = None + name: str | None = None + public_loaded: bool = False + public_name: str | None = None + + +@dataclass(frozen=True, slots=True) +class CodexGrants: + key: tuple[str, ...] + team: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class CodexModelFacts: + kind: CodexModelKind + model: str | None + router_present: bool + team_id: str | None = None + bundled: tuple[str, ...] | None = None + priced: bool | None = None + candidates: tuple[CodexCandidate, ...] | None = None + alias_loaded: bool = False + alias: str | None = None + grants: CodexGrants | None = None + prefixed_grants: CodexGrants | None = None + bare: tuple[ModelHelperFacts, ...] = () + prefixed: tuple[ModelHelperFacts, ...] = () + route: str | None = None + + +@dataclass(frozen=True, slots=True) +class CodexModelDecision: + kind: Literal[ + "Complete", + "ReadBundled", + "ReadCost", + "ReadCandidates", + "ReadAlias", + "ReadCandidate", + "ReadGrants", + "ReadHelper", + ] + allowed: bool = False + target: str | None = None + model: str | None = None + index: int = 0 + field: CandidateField | None = None + prefixed: bool = False + + +_DECISION: Final = TypeAdapter(CodexModelDecision) + + +@native_boundary +def codex_model(facts: CodexModelFacts) -> CodexModelDecision: + return _DECISION.validate_python(_native.request_auth_codex_model(encode_facts(facts))) diff --git a/litellm/_auth/codex_normalization.py b/litellm/_auth/codex_normalization.py new file mode 100644 index 00000000000..adf3ecedf6d --- /dev/null +++ b/litellm/_auth/codex_normalization.py @@ -0,0 +1,64 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.codex_models import CodexModelFacts +from litellm._auth.codex_turn import CodexTurnFacts +from litellm._auth.failure import native_boundary +from litellm._auth.routes import RouteContext +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class CodexNormalizationFacts: + kind: Literal["Alias", "Prefix"] + route: str + request_present: bool + router_present: bool | None = None + context: RouteContext | None = None + eu_jev: bool | None = None + resolved: bool | None = None + marked: bool = False + model_loaded: bool = False + model: str | None = None + explicit: bool | None = None + provider: bool | None = None + routing: CodexModelFacts | None = None + body_present: bool | None = None + settings_loaded: bool = False + target: str | None = None + provider_after: bool | None = None + turn: CodexTurnFacts | None = None + prefix: CodexModelFacts | None = None + + +@dataclass(frozen=True, slots=True) +class CodexNormalizationDecision: + kind: Literal[ + "Complete", + "ReadRouter", + "ReadContext", + "ReadEuJev", + "ReadResolved", + "MarkResolved", + "ReadModel", + "ReadExplicit", + "ReadProvider", + "ReadRouting", + "ReadBody", + "ReadSettings", + "ReadProviderAfter", + "ReadTurn", + "ReadPrefix", + "Rewrite", + ] + target: str | None = None + + +_DECISION: Final = TypeAdapter(CodexNormalizationDecision) + + +@native_boundary +def codex_normalization(facts: CodexNormalizationFacts) -> CodexNormalizationDecision: + return _DECISION.validate_python(_native.request_auth_codex_normalization(asdict(facts))) diff --git a/litellm/_auth/codex_turn.py b/litellm/_auth/codex_turn.py new file mode 100644 index 00000000000..dac1bbc1dda --- /dev/null +++ b/litellm/_auth/codex_turn.py @@ -0,0 +1,100 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.codex_models import CodexModelFacts +from litellm._auth.failure import native_boundary +from litellm._auth.route_enable import RouteEnableFacts +from litellm._auth.route_scope import RouteScopeFacts +from litellm._auth.routes import VirtualRouteFacts +from litellm.rust_bridge import _native + +JwtRead: TypeAlias = Literal["Returned", "Exception"] + + +@dataclass(frozen=True, slots=True) +class CodexKeyFacts: + proxy: RouteScopeFacts | None = None + enabled: RouteEnableFacts | None = None + virtual_route: VirtualRouteFacts | None = None + role_loaded: bool = False + role: str | None = None + admin: RouteScopeFacts | None = None + + +@dataclass(frozen=True, slots=True) +class CodexAccessFacts: + key: CodexKeyFacts + jwt_token_present: bool + jwt: JwtRead | None = None + + +@dataclass(frozen=True, slots=True) +class CodexCredentials: + bearer_present: bool + gateway_present: bool + equal: bool + + +@dataclass(frozen=True, slots=True) +class CodexTurnFacts: + model: str | None + native: bool = False + credentials: CodexCredentials | None = None + access: CodexAccessFacts | None = None + routing: CodexModelFacts | None = None + catalog_loaded: bool = False + slugs: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class CodexTurnDecision: + kind: Literal[ + "Complete", + "Failure", + "ReadProxy", + "ReadEnabled", + "ReadVirtual", + "ReadRole", + "ReadAdmin", + "ReadCredentials", + "ReadAccess", + "ReadRouting", + "ReadJwt", + "ReadCatalog", + ] + allowed: bool = False + message: str | None = None + + +_DECISION: Final = TypeAdapter(CodexTurnDecision) + + +@native_boundary +def codex_key(facts: CodexKeyFacts) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Key": asdict(facts)})) + + +@native_boundary +def codex_access(facts: CodexAccessFacts) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Access": asdict(facts)})) + + +@native_boundary +def codex_turn(facts: CodexTurnFacts) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Turn": asdict(facts)})) + + +@native_boundary +def codex_jwt(read: JwtRead) -> CodexTurnDecision: + return _DECISION.validate_python(_native.request_auth_codex_turn({"Jwt": read})) + + +@native_boundary +def codex_jwt_token(*, claims_present: bool, token_present: bool) -> CodexTurnDecision: + return _DECISION.validate_python( + _native.request_auth_codex_turn( + {"JwtToken": {"claims_present": claims_present, "token_present": token_present}} + ) + ) diff --git a/litellm/_auth/model_projection.py b/litellm/_auth/model_projection.py index 63022394416..832b1cf13b9 100644 --- a/litellm/_auth/model_projection.py +++ b/litellm/_auth/model_projection.py @@ -13,6 +13,7 @@ class ModelProjectionDecision: kind: Literal["Complete", "Read"] read: ModelReadKind | None = None model: str | None = None + team_id: str | None = None decision: ModelAccessDecision | None = None diff --git a/litellm/_auth/route_scope.py b/litellm/_auth/route_scope.py new file mode 100644 index 00000000000..a2f0616f942 --- /dev/null +++ b/litellm/_auth/route_scope.py @@ -0,0 +1,40 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +RouteScopeKind: TypeAlias = Literal["Proxy", "Admin"] + + +@dataclass(frozen=True, slots=True) +class RouteValues: + text: str | None = None + names: tuple[str, ...] = () + failure: str | None = None + + +@dataclass(frozen=True, slots=True) +class RouteScopeFacts: + kind: RouteScopeKind + route: str + present: bool + premium: bool + warned: bool = False + values: RouteValues | None = None + + +@dataclass(frozen=True, slots=True) +class RouteScopeDecision: + kind: Literal["Allow", "Deny", "Warn", "ReadValues", "Failure"] + message: str | None = None + + +_DECISION: Final = TypeAdapter(RouteScopeDecision) + + +@native_boundary +def route_scope(facts: RouteScopeFacts) -> RouteScopeDecision: + return _DECISION.validate_python(_native.request_auth_route_scope(asdict(facts))) diff --git a/litellm/proxy/auth/codex_model_adapter.py b/litellm/proxy/auth/codex_model_adapter.py new file mode 100644 index 00000000000..8f3ea7f5207 --- /dev/null +++ b/litellm/proxy/auth/codex_model_adapter.py @@ -0,0 +1,166 @@ +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import Final, assert_never + +from pydantic import BaseModel + +import litellm +from litellm._auth.codex_models import ( + CandidateField, + CodexCandidate, + CodexGrants, + CodexModelFacts, + CodexModelKind, + codex_model, +) +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.model_helpers import ModelHelperFacts +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.model_helper_adapter import read_model_helper_facts +from litellm.proxy.common_utils.codex_model_catalog import bundled_codex_models +from litellm.router_api import RegistryReader +from litellm.types.router import Deployment + + +class _ModelGrants(BaseModel, frozen=True): + models: tuple[str, ...] + team_models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _State: + facts: CodexModelFacts + deployments: tuple[Deployment, ...] = () + + +@native_boundary +def _grants(token: UserAPIKeyAuth) -> CodexGrants: + fields: Final = _ModelGrants.model_validate(token, from_attributes=True) + return CodexGrants(fields.models, fields.team_models) + + +def _candidate(field: CandidateField, previous: CodexCandidate, deployment: Deployment) -> CodexCandidate: + match field: + case "Provider": + return replace(previous, provider_loaded=True, provider=deployment.litellm_params.custom_llm_provider) + case "Model": + return replace(previous, model=deployment.litellm_params.model) + case "Name": + return replace(previous, name=deployment.model_name) + case "PublicName": + return replace(previous, public_loaded=True, public_name=deployment.model_info.team_public_model_name) + case _: + assert_never(field) + + +def read_codex_model_facts( + model: object, + router: RegistryReader | None, + team_id: str | None, + *, + kind: CodexModelKind = "Routes", + token: UserAPIKeyAuth | None = None, + route: str | None = None, +) -> CodexModelFacts: + initial: Final = _State( + CodexModelFacts(kind, model if isinstance(model, str) else None, router is not None, team_id, route=route) + ) + + def advance(state: _State, _tick: None) -> _State: + facts: Final = state.facts + decision: Final = codex_model(facts) + match decision.kind: + case "ReadBundled": + return replace(state, facts=replace(facts, bundled=tuple(bundled_codex_models()))) + case "ReadCost": + if decision.model is None: + raise NativeAuthRuntimeError("Codex price lookup without model") + return replace(state, facts=replace(facts, priced=decision.model in litellm.model_cost)) # pyright: ignore[reportUnknownMemberType] # dynamically loaded registry, only raw key presence is read + case "ReadCandidates": + if router is None or decision.model is None: + raise NativeAuthRuntimeError("Codex deployment lookup without router or model") + registry: Final = router.registry() + loaded: Final = tuple( + registry.deployment(model_id) + for model_id in registry.names(decision.model, kind="route_candidates", team_id=team_id) + ) + deployments: Final = tuple(deployment for deployment in loaded if deployment is not None) + return _State(replace(facts, candidates=tuple(CodexCandidate() for _ in deployments)), deployments) + case "ReadAlias": + if router is None or decision.model is None: + raise NativeAuthRuntimeError("Codex alias lookup without router or model") + return replace( + state, + facts=replace( + facts, alias_loaded=True, alias=router.registry().resolve(decision.model, kind="alias") + ), + ) + case "ReadCandidate": + if facts.candidates is None or decision.field is None: + raise NativeAuthRuntimeError("Codex candidate projection without deployment") + updated: Final = _candidate( + decision.field, facts.candidates[decision.index], state.deployments[decision.index] + ) + return replace( + state, + facts=replace( + facts, + candidates=tuple( + updated if index == decision.index else value + for index, value in enumerate(facts.candidates) + ), + ), + ) + case "ReadGrants": + if token is None: + raise NativeAuthRuntimeError("Codex model grants without token") + grant_read: Final = _grants(token) + return replace( + state, + facts=( + replace(facts, prefixed_grants=grant_read) + if decision.prefixed + else replace(facts, grants=grant_read) + ), + ) + case "ReadHelper": + selected: Final = facts.prefixed_grants if decision.prefixed else facts.grants + if selected is None or decision.model is None: + raise NativeAuthRuntimeError("Codex helper lookup without model grants") + grants: Final = (selected.key, selected.team)[decision.index] + scope: Final = read_model_helper_facts( + ModelHelperFacts("Helper", decision.model, grants, team_id=team_id, router_present=bool(router)), + router, + ) + updated_facts: Final = ( + replace(facts, prefixed=(*facts.prefixed, scope)) + if decision.prefixed + else replace(facts, bare=(*facts.bare, scope)) + ) + return replace(state, facts=updated_facts) + case "Complete": + raise NativeAuthRuntimeError("completed Codex model decision requested IO") + case _: + assert_never(decision.kind) + + return next( + state.facts + for state in accumulate(repeat(None), advance, initial=initial) + if codex_model(state.facts).kind == "Complete" + ) + + +def routes_to_litellm(model: object, llm_router: RegistryReader | None, team_id: str | None) -> bool: + return codex_model(read_codex_model_facts(model, llm_router, team_id)).allowed + + +def is_codex_native_litellm_turn(route: str, model: object, llm_router: RegistryReader, team_id: str | None) -> bool: + return codex_model(read_codex_model_facts(model, llm_router, team_id, kind="NativeRoutes", route=route)).allowed + + +def key_and_team_allow(model: str, llm_router: RegistryReader, token: UserAPIKeyAuth) -> bool: + return codex_model(read_codex_model_facts(model, llm_router, token.team_id, kind="Grants", token=token)).allowed + + +def native_chatgpt_model(model: str, llm_router: RegistryReader, token: UserAPIKeyAuth) -> str | None: + return codex_model(read_codex_model_facts(model, llm_router, token.team_id, kind="Chatgpt", token=token)).target diff --git a/litellm/proxy/auth/codex_normalization_adapter.py b/litellm/proxy/auth/codex_normalization_adapter.py new file mode 100644 index 00000000000..12fe38af6d5 --- /dev/null +++ b/litellm/proxy/auth/codex_normalization_adapter.py @@ -0,0 +1,166 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Final, cast + +from fastapi import Request + +from litellm._auth.codex_normalization import CodexNormalizationFacts, codex_normalization +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._logging import verbose_proxy_logger +from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_model_adapter import read_codex_model_facts +from litellm.proxy.auth.codex_turn_adapter import read_native_turn_facts +from litellm.proxy.auth.route_policy_adapter import read_eu_jev_configuration, route_context +from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body, rewrite_request_model +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + data: dict[str, object] # mutable-ok: request-local rewrite is a selected effect + token: UserAPIKeyAuth + request: Request | None + router: RegistryReader | None + + +def _value(value: object, key: str) -> object: + if not isinstance(value, Mapping): + return None + return cast(Mapping[object, object], value).get( + key + ) # cast-ok: Mapping shape was checked without coercing unrelated keys + + +def _rewrite(facts: CodexNormalizationFacts, source: _Source, target: str | None) -> None: + if source.request is None or facts.model is None or target is None: + raise NativeAuthRuntimeError("Codex rewrite without request model") + if facts.kind == "Alias": + verbose_proxy_logger.debug( + "router_settings.model_group_alias resolved %s -> %s before auth", + facts.model.replace("\r", "").replace("\n", ""), + target.replace("\r", "").replace("\n", ""), + ) + source.request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, facts.model) + rewrite_request_model(source.data, source.request, target) + + +async def _alias(facts: CodexNormalizationFacts, source: _Source) -> None: + from litellm.proxy.auth.auth_utils import ( + request_dispatched_to_pass_through_endpoint, + request_dispatched_to_provider_pass_through, + ) + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + decision: Final = codex_normalization(facts) + if decision.kind == "Complete": + return + if source.request is None: + raise NativeAuthRuntimeError("alias normalization read without request") + match decision.kind: + case "ReadContext": + return await _alias(replace(facts, context=route_context(facts.route)), source) + case "ReadEuJev": + return await _alias(replace(facts, eu_jev=read_eu_jev_configuration()), source) + case "ReadResolved": + return await _alias( + replace(facts, resolved=source.request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True), source + ) + case "MarkResolved": + source.request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True + return await _alias(replace(facts, marked=True), source) + case "ReadModel": + model: Final = source.data.get("model") + return await _alias( + replace(facts, model_loaded=True, model=model if isinstance(model, str) else None), source + ) + case "ReadExplicit": + return await _alias( + replace(facts, explicit=request_dispatched_to_pass_through_endpoint(source.request)), source + ) + case "ReadProvider": + return await _alias( + replace(facts, provider=request_dispatched_to_provider_pass_through(source.request)), source + ) + case "ReadRouting": + routing: Final = read_codex_model_facts(facts.model, source.router, source.token.team_id) + return await _alias(replace(facts, routing=routing), source) + case "ReadBody": + return await _alias( + replace(facts, body_present=await read_raw_json_body(request=source.request) is not None), source + ) + case "ReadSettings": + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=source.token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + if facts.model is None: + raise NativeAuthRuntimeError("alias settings read without requested model") + entry: Final = _value(_value(settings, "model_group_alias"), facts.model) + target: Final = ( + _value(cast(object, entry), "model") if isinstance(entry, Mapping) else entry + ) # cast-ok: retain the object boundary after Mapping narrowing + return await _alias( + replace(facts, settings_loaded=True, target=target if isinstance(target, str) else None), source + ) + case "ReadProviderAfter": + return await _alias( + replace(facts, provider_after=request_dispatched_to_provider_pass_through(source.request)), source + ) + case "ReadTurn": + if facts.model is None or source.router is None: + raise NativeAuthRuntimeError("alias turn read without router model") + turn: Final = await read_native_turn_facts(facts.model, source.router, source.request, source.token) + return await _alias(replace(facts, turn=turn), source) + case "Rewrite": + _rewrite(facts, source, decision.target) + case _: + raise NativeAuthRuntimeError("invalid alias normalization stage") + + +async def resolve_router_settings_model_group_alias( + request_data: dict[str, object], # mutable-ok: preserve request-local model mutation + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + from litellm.proxy.proxy_server import llm_router + + await _alias( + CodexNormalizationFacts("Alias", route, request is not None, router_present=llm_router is not None), + _Source(request_data, valid_token, request, llm_router), + ) + + +def _prefix(facts: CodexNormalizationFacts, source: _Source) -> None: + decision: Final = codex_normalization(facts) + match decision.kind: + case "Complete": + return + case "ReadRouter": + from litellm.proxy.proxy_server import llm_router + + return _prefix(replace(facts, router_present=llm_router is not None), replace(source, router=llm_router)) + case "ReadPrefix": + prefix: Final = read_codex_model_facts( + facts.model, source.router, source.token.team_id, kind="Chatgpt", token=source.token + ) + return _prefix(replace(facts, prefix=prefix), source) + case "Rewrite": + _rewrite(facts, source, decision.target) + case _: + raise NativeAuthRuntimeError("invalid Codex prefix normalization stage") + + +def resolve_codex_native_chatgpt_model( + request_data: dict[str, object], # mutable-ok: preserve request-local model mutation + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + model: Final = request_data.get("model") + _prefix( + CodexNormalizationFacts( + "Prefix", route, request is not None, model_loaded=True, model=model if isinstance(model, str) else None + ), + _Source(request_data, valid_token, request, None), + ) diff --git a/litellm/proxy/auth/codex_turn_adapter.py b/litellm/proxy/auth/codex_turn_adapter.py new file mode 100644 index 00000000000..c73d1d7c694 --- /dev/null +++ b/litellm/proxy/auth/codex_turn_adapter.py @@ -0,0 +1,218 @@ +from __future__ import annotations + +import hmac +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Final + +from fastapi import Request + +from litellm._auth.codex_turn import ( + CodexAccessFacts, + CodexCredentials, + CodexKeyFacts, + CodexTurnFacts, + JwtRead, + codex_access, + codex_jwt_token, + codex_key, + codex_turn, +) +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_model_adapter import read_codex_model_facts +from litellm.proxy.auth.route_enable_adapter import read_route_enable_facts +from litellm.proxy.auth.route_policy_adapter import read_virtual_route_facts +from litellm.proxy.auth.route_scope_adapter import read_route_scope_facts +from litellm.router_api import RegistryReader + +if TYPE_CHECKING: + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _SharedCache, # pyright: ignore[reportPrivateUsage] # existing account-catalog IO protocol + ) + +RESPONSES_PATH: Final = "/v1/responses" + + +def read_codex_key_facts(token: UserAPIKeyAuth, facts: CodexKeyFacts = CodexKeyFacts()) -> CodexKeyFacts: + decision: Final = codex_key(facts) + match decision.kind: + case "Complete": + return facts + case "Failure": + raise TypeError(decision.message) + case "ReadProxy": + return read_codex_key_facts(token, replace(facts, proxy=read_route_scope_facts(RESPONSES_PATH, "Proxy"))) + case "ReadEnabled": + return read_codex_key_facts(token, replace(facts, enabled=read_route_enable_facts(RESPONSES_PATH))) + case "ReadVirtual": + return read_codex_key_facts( + token, replace(facts, virtual_route=read_virtual_route_facts(RESPONSES_PATH, token, None)) + ) + case "ReadRole": + return read_codex_key_facts(token, replace(facts, role_loaded=True, role=token.user_role)) + case "ReadAdmin": + return read_codex_key_facts(token, replace(facts, admin=read_route_scope_facts(RESPONSES_PATH, "Admin"))) + case _: + raise NativeAuthRuntimeError("invalid Codex key read stage") + + +async def read_jwt_result(gateway_token: str, model: str, headers: Mapping[str, str]) -> JwtRead: + from litellm.proxy.auth.handle_jwt import JWTAuthManager + from litellm.proxy.proxy_server import ( + general_settings_view, + jwt_handler, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + try: + _ = await JWTAuthManager.auth_builder( # pyright: ignore[reportUnknownMemberType] # unchanged validator accepts legacy dictionaries; its return value is not interpreted here + api_key=gateway_token, + jwt_handler=jwt_handler, + request_data={"model": model}, + general_settings=dict(general_settings_view()), + route=RESPONSES_PATH, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=headers, + request_method="POST", + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # the unchanged JWT validator reports refusals by raising + return "Exception" + return "Returned" + + +@dataclass(frozen=True, slots=True) +class CodexAccess: + key: CodexKeyFacts + jwt_token: str | None = field(repr=False) + request_headers: Mapping[str, str] = field(repr=False) + + @property + def key_admits(self) -> bool: + return codex_key(self.key).allowed + + async def admits(self, model: str) -> bool: + return codex_access(await read_codex_access_facts(self, model)).allowed + + +def read_codex_access(token: UserAPIKeyAuth, request: Request) -> CodexAccess: + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _gateway_token, # pyright: ignore[reportPrivateUsage] # raw credential extraction, not auth policy + ) + + selection: Final = codex_jwt_token(claims_present=bool(token.jwt_claims), token_present=bool(token.token)) + key: Final = read_codex_key_facts(token) + return CodexAccess(key, _gateway_token(request) if selection.allowed else None, request.headers) + + +async def read_codex_access_facts(source: CodexAccess, model: str) -> CodexAccessFacts: + facts: Final = CodexAccessFacts(source.key, source.jwt_token is not None) + decision: Final = codex_access(facts) + match decision.kind: + case "Complete": + return facts + case "ReadJwt": + if source.jwt_token is None: + raise NativeAuthRuntimeError("Codex JWT read without credential") + return replace(facts, jwt=await read_jwt_result(source.jwt_token, model, source.request_headers)) + case _: + raise NativeAuthRuntimeError("invalid Codex access read stage") + + +@dataclass(frozen=True, slots=True) +class _Source: + router: RegistryReader | None + token: UserAPIKeyAuth + request: Request | None = None + access: CodexAccess | None = None + caller: str | None = field(default=None, repr=False) + headers: Mapping[str, str] | None = field(default=None, repr=False) + cache: _SharedCache | None = None + gateway_token: str | None = field(default=None, repr=False) + + +async def _read(facts: CodexTurnFacts, source: _Source) -> CodexTurnFacts: + from litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints import ( + _account_slugs, # pyright: ignore[reportPrivateUsage] # unchanged catalog IO + _chatgpt_bearer, # pyright: ignore[reportPrivateUsage] # raw credential parsing + _forwarded_headers, # pyright: ignore[reportPrivateUsage] # transport header projection + _gateway_token, # pyright: ignore[reportPrivateUsage] # raw credential parsing + _models_etag_caller, # pyright: ignore[reportPrivateUsage] # cache identity projection + codex_shared_cache, + ) + + decision: Final = codex_turn(facts) + match decision.kind: + case "Complete": + return facts + case "ReadCredentials": + if source.request is None: + raise NativeAuthRuntimeError("native Codex turn without request") + bearer: Final = _chatgpt_bearer(source.request) + gateway: Final = _gateway_token(source.request) + credentials: Final = CodexCredentials( + bool(bearer), bool(gateway), hmac.compare_digest(bearer.encode(), gateway.encode()) + ) + return await _read(replace(facts, credentials=credentials), replace(source, gateway_token=gateway)) + case "ReadAccess": + if source.request is None or source.gateway_token is None: + raise NativeAuthRuntimeError("native Codex access without request credential") + access: Final = read_codex_access(source.token, source.request) + caller: Final = _models_etag_caller(source.token, source.gateway_token) + headers: Final = _forwarded_headers(source.request) + cache: Final = codex_shared_cache() + return await _read( + replace(facts, access=CodexAccessFacts(access.key, access.jwt_token is not None)), + replace(source, access=access, caller=caller, headers=headers, cache=cache), + ) + case "ReadRouting": + routing: Final = read_codex_model_facts(facts.model, source.router, source.token.team_id) + return await _read(replace(facts, routing=routing), source) + case "ReadJwt": + if source.access is None or facts.model is None: + raise NativeAuthRuntimeError("Codex JWT read without access context") + return await _read(replace(facts, access=await read_codex_access_facts(source.access, facts.model)), source) + case "ReadCatalog": + if source.headers is None or source.caller is None or source.cache is None: + raise NativeAuthRuntimeError("Codex catalog read without IO context") + slugs: Final = await _account_slugs(source.headers, source.caller, source.cache) + return await _read( + replace(facts, catalog_loaded=True, slugs=tuple(slugs) if slugs is not None else None), source + ) + case _: + raise NativeAuthRuntimeError("invalid Codex turn read stage") + + +async def read_native_turn_facts( + model: str, router: RegistryReader, request: Request, token: UserAPIKeyAuth +) -> CodexTurnFacts: + return await _read(CodexTurnFacts(model, native=True), _Source(router, token, request=request)) + + +async def native_turn_serves_through_litellm( + model: str, llm_router: RegistryReader, request: Request, user_api_key_dict: UserAPIKeyAuth +) -> bool: + return codex_turn(await read_native_turn_facts(model, llm_router, request, user_api_key_dict)).allowed + + +async def serves_through_litellm( + model: object, + llm_router: RegistryReader | None, + user_api_key_dict: UserAPIKeyAuth, + access: CodexAccess, + caller: str, + headers: Mapping[str, str], + cache: _SharedCache, +) -> bool: + facts: Final = CodexTurnFacts( + model if isinstance(model, str) else None, access=CodexAccessFacts(access.key, access.jwt_token is not None) + ) + source: Final = _Source(llm_router, user_api_key_dict, access=access, caller=caller, headers=headers, cache=cache) + return codex_turn(await _read(facts, source)).allowed diff --git a/litellm/proxy/auth/key_model_access_adapter.py b/litellm/proxy/auth/key_model_access_adapter.py index eb79c6d93eb..fdb32b58a0a 100644 --- a/litellm/proxy/auth/key_model_access_adapter.py +++ b/litellm/proxy/auth/key_model_access_adapter.py @@ -10,6 +10,10 @@ from litellm._auth.fallback_access import GroupedModelScope from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_and_fallback_model_access from litellm._auth.model_access import key_model_grants from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_normalization_adapter import ( + resolve_codex_native_chatgpt_model, + resolve_router_settings_model_group_alias, +) 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 @@ -39,11 +43,11 @@ async def read_key_and_fallback_model_access_facts( ) -> KeyAndFallbackModelAccessFacts: from litellm.proxy.auth.user_api_key_auth import ( _normalize_claude_model, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # legacy helper takes the mutable request body - _resolve_router_settings_model_group_alias, # pyright: ignore[reportPrivateUsage] # preserve request-local router alias normalization ) await _normalize_claude_model(request_data, valid_token, request, route) - await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) + await resolve_router_settings_model_group_alias(request_data, valid_token, request, route) + resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) facts: Final = KeyAndFallbackModelAccessFacts( config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config), key_models=object_model_grants(valid_token), diff --git a/litellm/proxy/auth/model_projection_adapter.py b/litellm/proxy/auth/model_projection_adapter.py index 50141d7fb3c..b95bf682867 100644 --- a/litellm/proxy/auth/model_projection_adapter.py +++ b/litellm/proxy/auth/model_projection_adapter.py @@ -39,7 +39,7 @@ def project_scope( 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) + reading: Final = read_registry(decision.read, decision.model, decision.team_id, router) updated: Final = (*state.reads, reading) return _Projection(updated, model_projection(facts, updated)) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 4f1e5df0505..4d822f22300 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -4,11 +4,9 @@ from typing import Final, cast from fastapi import HTTPException, Request, status from litellm._auth import routes -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( MODEL_LIST_ROUTES, UI_TEAM_ID, - CommonProxyErrors, KeyManagementRoutes, LiteLLM_UserTable, LiteLLMRoutes, @@ -208,33 +206,15 @@ class RouteChecks: @staticmethod def custom_admin_only_route_check(route: str): - from litellm.proxy.proxy_server import general_settings, premium_user + from litellm.proxy.auth.route_scope_adapter import custom_admin_only_route_check - if "admin_only_routes" in general_settings: - if premium_user is not True: - verbose_proxy_logger.error( - "Trying to use 'admin_only_routes' this is an Enterprise only feature. %s", - CommonProxyErrors.not_premium_user.value, - ) - return - if route in general_settings["admin_only_routes"]: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route. Route={route} is an admin only route", - ) + return custom_admin_only_route_check(route) @staticmethod def is_allowed_on_proxy(route: str) -> bool: - from litellm.proxy.proxy_server import general_settings, premium_user + from litellm.proxy.auth.route_scope_adapter import is_allowed_on_proxy - if "allowed_routes" not in general_settings: - return True - if premium_user is not True: - verbose_proxy_logger.error( - "Trying to set allowed_routes. This is an Enterprise feature. %s", - CommonProxyErrors.not_premium_user.value, - ) - return route in general_settings["allowed_routes"] + return is_allowed_on_proxy(route) @staticmethod def is_unbilled_route(route: str) -> bool: diff --git a/litellm/proxy/auth/route_enable_adapter.py b/litellm/proxy/auth/route_enable_adapter.py index 4cb119c0381..91358dbb73b 100644 --- a/litellm/proxy/auth/route_enable_adapter.py +++ b/litellm/proxy/auth/route_enable_adapter.py @@ -56,6 +56,11 @@ class RouteEnableReaders: _READERS: Final = RouteEnableReaders() +def read_route_enable_facts(route: str) -> RouteEnableFacts: + facts, _error = _read(RouteEnableFacts(route_context(route)), _READERS) + return facts + + def _answer( facts: RouteEnableFacts, decision: RouteEnableDecision, readers: RouteEnableReaders, enterprise: object ) -> tuple[RouteEnableFacts, object]: diff --git a/litellm/proxy/auth/route_policy_adapter.py b/litellm/proxy/auth/route_policy_adapter.py index 459798b5568..de7309c7ad4 100644 --- a/litellm/proxy/auth/route_policy_adapter.py +++ b/litellm/proxy/auth/route_policy_adapter.py @@ -143,17 +143,20 @@ def raise_route_denial( ) -def check_virtual_route(route: str, token: UserAPIKeyAuth, request: Request | None) -> bool: +def read_virtual_route_facts(route: str, token: UserAPIKeyAuth, request: Request | None) -> routes.VirtualRouteFacts: from litellm.proxy.auth.route_checks import RouteChecks fields: Final = _TokenFields.model_validate(token, from_attributes=True) - facts: Final = routes.VirtualRouteFacts( + return routes.VirtualRouteFacts( context=route_context(route, RouteChecks.get_request_method(request), registration=True), allowed_routes=_allowed_routes(fields.allowed_routes), passthrough=passthrough_facts(route, token), password_reset_required=(fields.metadata or {}).get("password_reset_required") is True, ) - raise_route_denial(routes.virtual_route(facts), route, token=token) + + +def check_virtual_route(route: str, token: UserAPIKeyAuth, request: Request | None) -> bool: + raise_route_denial(routes.virtual_route(read_virtual_route_facts(route, token, request)), route, token=token) return True diff --git a/litellm/proxy/auth/route_scope_adapter.py b/litellm/proxy/auth/route_scope_adapter.py new file mode 100644 index 00000000000..a4c5de07178 --- /dev/null +++ b/litellm/proxy/auth/route_scope_adapter.py @@ -0,0 +1,73 @@ +from collections.abc import Iterable, Mapping +from dataclasses import replace +from typing import Final, assert_never + +from fastapi import HTTPException +from pydantic import BaseModel, TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_scope import RouteScopeFacts, RouteScopeKind, RouteValues, route_scope +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors + +_OBJECTS: Final = TypeAdapter(Iterable[object]) + + +class _Settings(BaseModel, frozen=True): + general_settings: Mapping[str, object] + premium_user: object + + +def route_values(value: object) -> RouteValues: + if isinstance(value, str): + return RouteValues(text=value) + if isinstance(value, (bytes, bytearray)): + return RouteValues(failure="a bytes-like object is required, not 'str'") + if isinstance(value, Iterable): + return RouteValues(names=tuple(item for item in _OBJECTS.validate_python(value) if isinstance(item, str))) + return RouteValues(failure=f"argument of type '{type(value).__name__}' is not iterable") + + +def read_route_scope_facts(route: str, kind: RouteScopeKind) -> RouteScopeFacts: + from litellm.proxy import proxy_server + + settings: Final = _Settings.model_validate(vars(proxy_server)) + setting: Final = "allowed_routes" if kind == "Proxy" else "admin_only_routes" + initial: Final = RouteScopeFacts(kind, route, setting in settings.general_settings, settings.premium_user is True) + + def advance(facts: RouteScopeFacts) -> RouteScopeFacts: + decision: Final = route_scope(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "Warn": + verbose_proxy_logger.error( + "Trying to set allowed_routes. This is an Enterprise feature. %s" + if kind == "Proxy" + else "Trying to use 'admin_only_routes' this is an Enterprise only feature. %s", + CommonProxyErrors.not_premium_user.value, + ) + return advance(replace(facts, warned=True)) + case "ReadValues": + return advance(replace(facts, values=route_values(settings.general_settings[setting]))) + case "Failure": + raise TypeError(decision.message) + case _: + assert_never(decision.kind) + + return advance(initial) + + +def is_allowed_on_proxy(route: str) -> bool: + return route_scope(read_route_scope_facts(route, "Proxy")).kind == "Allow" + + +def custom_admin_only_route_check(route: str) -> None: + decision: Final = route_scope(read_route_scope_facts(route, "Admin")) + match decision.kind: + case "Allow": + return + case "Deny": + raise HTTPException(403, f"user not allowed to access this route. Route={route} is an admin only route") + case _: + raise NativeAuthRuntimeError("incomplete custom admin route decision") diff --git a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py index d321f5b5861..74739e67ea6 100644 --- a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py @@ -4,7 +4,6 @@ import hmac import json import re from collections.abc import Iterator, Mapping -from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol @@ -14,7 +13,7 @@ from fastapi.responses import StreamingResponse from pydantic import TypeAdapter, ValidationError from starlette.types import Message -import litellm +from litellm._auth.failure import NativeAuthFailure from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.exceptions import BadRequestError @@ -24,14 +23,12 @@ from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared client factory has untyped params ) -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import check_model_access_helper, check_tools_allowlist -from litellm.proxy.auth.handle_jwt import JWTAuthManager -from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import check_tools_allowlist +from litellm.proxy.auth.codex_turn_adapter import CodexAccess as _LitellmResponsesAccess from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.codex_model_catalog import ( CodexCatalogRow, - bundled_codex_models, codex_catalog_etag, codex_entry_names_default_reasoning_level, merged_codex_catalog, @@ -69,7 +66,6 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _INPUT_ITEMS: Final = TypeAdapter(tuple[object, ...]) _CACHED_SLUGS: Final = TypeAdapter(frozenset[str]) _CACHED_MODELS_ETAG: Final = TypeAdapter(dict[Literal["etag"], str]) -_MODEL_NAMES: Final = TypeAdapter(list[str]) _HEADER_DERIVED_STATE: Final = frozenset({"_cached_headers"}) _ACCOUNT_CATALOG_HEADERS: Final = frozenset({"authorization", "chatgpt-account-id", "originator", "user-agent"}) _CODEX_USER_AGENT_PRODUCT: Final = re.compile(r"([^/]+)/(\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?)") @@ -108,19 +104,6 @@ _HEADERS: Final = frozenset( ) -def _is_chatgpt(deployment: Deployment) -> bool: - return ( - deployment.litellm_params.custom_llm_provider == _CHATGPT_PROVIDER - or deployment.litellm_params.model.partition("/")[0] == _CHATGPT_PROVIDER - ) - - -def _serves_by_name(deployment: Deployment, names: frozenset[str]) -> bool: - return not _is_chatgpt(deployment) and ( - deployment.model_name in names or deployment.model_info.team_public_model_name in names - ) - - def _candidate_deployments(model: str, llm_router: "Router", team_id: str | None) -> tuple[Deployment, ...]: registry: Final = llm_router.registry() candidates: Final = tuple( @@ -130,37 +113,15 @@ def _candidate_deployments(model: str, llm_router: "Router", team_id: str | None def routes_to_litellm(model: object, llm_router: "Router | None", team_id: str | None) -> bool: - if ( - not isinstance(model, str) - or llm_router is None - or model in bundled_codex_models() - or f"{_CHATGPT_PROVIDER}/{model}" in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed - ): - return False - deployments: Final = _candidate_deployments(model, llm_router, team_id) - names: Final = frozenset({model, llm_router.registry().resolve(model, kind="alias") or model}) - return bool(deployments) and all(_serves_by_name(deployment, names) for deployment in deployments) + from litellm.proxy.auth.codex_model_adapter import routes_to_litellm as native - -def _key_and_team_allow(model: str, llm_router: "Router", user_api_key_dict: UserAPIKeyAuth) -> bool: - allowed_lists: Final = ( - _MODEL_NAMES.validate_python(user_api_key_dict.models), - _MODEL_NAMES.validate_python(user_api_key_dict.team_models), - ) - return all( - check_model_access_helper(model, llm_router, allowed, team_id=user_api_key_dict.team_id) - for allowed in allowed_lists - ) + return native(model, llm_router, team_id) def native_chatgpt_model(model: str, llm_router: "Router", user_api_key_dict: UserAPIKeyAuth) -> str | None: - if model.startswith(f"{_CHATGPT_PROVIDER}/") or _key_and_team_allow(model, llm_router, user_api_key_dict): - return None - chatgpt_model: Final = f"{_CHATGPT_PROVIDER}/{model}" - deployments: Final = _candidate_deployments(chatgpt_model, llm_router, user_api_key_dict.team_id) - if not deployments or not all(_is_chatgpt(deployment) for deployment in deployments): - return None - return chatgpt_model if _key_and_team_allow(chatgpt_model, llm_router, user_api_key_dict) else None + from litellm.proxy.auth.codex_model_adapter import native_chatgpt_model as native + + return native(model, llm_router, user_api_key_dict) class _SharedCache(Protocol): @@ -177,69 +138,10 @@ def codex_shared_cache() -> _SharedCache: return proxy_logging_obj.internal_usage_cache.dual_cache -def _key_may_call_litellm_responses(user_api_key_dict: UserAPIKeyAuth) -> bool: - if not RouteChecks.is_allowed_on_proxy(_LITELLM_RESPONSES_PATH): - return False - try: - _ = RouteChecks.should_call_route(route=_LITELLM_RESPONSES_PATH, valid_token=user_api_key_dict) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - RouteChecks.custom_admin_only_route_check(route=_LITELLM_RESPONSES_PATH) - except HTTPException: - return False - return True - - -async def _jwt_policy_admits_litellm_responses( - gateway_token: str, model: str, request_headers: Mapping[str, str] -) -> bool: - from litellm.proxy.proxy_server import ( - general_settings_view, - jwt_handler, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - try: - _ = await JWTAuthManager.auth_builder( - api_key=gateway_token, - jwt_handler=jwt_handler, - request_data={"model": model}, - general_settings=dict(general_settings_view()), - route=_LITELLM_RESPONSES_PATH, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=proxy_logging_obj, - request_headers=request_headers, - request_method="POST", - ) - except Exception: # noqa: BLE001 # any JWT policy refusal keeps the caller on ChatGPT's models - return False - return True - - -@dataclass(frozen=True, slots=True) -class _LitellmResponsesAccess: - key_admits: bool - jwt_token: str | None - request_headers: Mapping[str, str] - - async def admits(self, model: str) -> bool: - if not self.key_admits: - return False - if self.jwt_token is None: - return True - return await _jwt_policy_admits_litellm_responses(self.jwt_token, model, self.request_headers) - - def _litellm_responses_access(user_api_key_dict: UserAPIKeyAuth, request: Request) -> _LitellmResponsesAccess: - standard_jwt: Final = bool(user_api_key_dict.jwt_claims) and not user_api_key_dict.token - return _LitellmResponsesAccess( - key_admits=_key_may_call_litellm_responses(user_api_key_dict), - jwt_token=_gateway_token(request) if standard_jwt else None, - request_headers=request.headers, - ) + from litellm.proxy.auth.codex_turn_adapter import read_codex_access + + return read_codex_access(user_api_key_dict, request) def _account_slugs_cache_key(headers: Mapping[str, str]) -> str | None: @@ -324,12 +226,9 @@ async def _serves_through_litellm( headers: Mapping[str, str], cache: _SharedCache, ) -> bool: - if not isinstance(model, str) or not routes_to_litellm(model, llm_router, user_api_key_dict.team_id): - return False - if not await access.admits(model): - return False - account_slugs: Final = await _account_slugs(headers, caller, cache) - return account_slugs is not None and model not in account_slugs + from litellm.proxy.auth.codex_turn_adapter import serves_through_litellm + + return await serves_through_litellm(model, llm_router, user_api_key_dict, access, caller, headers, cache) def _gateway_token(request: Request) -> str: @@ -353,19 +252,9 @@ def _forwarded_headers(request: Request) -> Mapping[str, str]: async def native_turn_serves_through_litellm( model: str, llm_router: "Router", request: Request, user_api_key_dict: UserAPIKeyAuth ) -> bool: - chatgpt_bearer: Final = _chatgpt_bearer(request) - gateway_token: Final = _gateway_token(request) - if not chatgpt_bearer or not gateway_token or hmac.compare_digest(chatgpt_bearer.encode(), gateway_token.encode()): - return False - return await _serves_through_litellm( - model, - llm_router, - user_api_key_dict, - _litellm_responses_access(user_api_key_dict, request), - _models_etag_caller(user_api_key_dict, gateway_token), - _forwarded_headers(request), - codex_shared_cache(), - ) + from litellm.proxy.auth.codex_turn_adapter import native_turn_serves_through_litellm as native + + return await native(model, llm_router, request, user_api_key_dict) def _json_object(value: object) -> Mapping[str, object]: @@ -588,6 +477,8 @@ async def _litellm_catalog_rows( return () try: rows: Final = await codex_catalog_rows_for_key(user_api_key_dict) + except NativeAuthFailure: + raise except Exception as error: # noqa: BLE001 # a failed listing must never fail the ChatGPT answer it extends verbose_proxy_logger.warning( "Codex catalog keeps only ChatGPT's models: listing the key's LiteLLM models raised %s", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index e264363e463..8f5c3fcb809 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -70,6 +70,10 @@ def request_auth_personal_models(facts: object) -> object: ... def request_auth_owner_teams(facts: object) -> object: ... def request_auth_team_override(facts: object) -> object: ... def request_auth_model_helpers(facts: object) -> object: ... +def request_auth_codex_model(facts: object) -> object: ... +def request_auth_codex_normalization(facts: object) -> object: ... +def request_auth_codex_turn(facts: object) -> object: ... +def request_auth_route_scope(facts: object) -> object: ... def request_auth_personal_owner(facts: object) -> object: ... def request_auth_personal_model_request(facts: object) -> object: ... def request_auth_member_models(facts: object) -> object: ... diff --git a/tests/unit/_auth/support/codex_reference.py b/tests/unit/_auth/support/codex_reference.py new file mode 100644 index 00000000000..66d57b362ac --- /dev/null +++ b/tests/unit/_auth/support/codex_reference.py @@ -0,0 +1,68 @@ +"""Codex authorization frozen from main 2f430873aead937b9112e683703956eb1ab179af.""" + +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_api import RegistryReader +from litellm.types.router import Deployment +from tests.unit._auth.support.model_registry_reference import ( + _helper, # pyright: ignore[reportPrivateUsage] # independently frozen model helper, not production policy +) + +_MODEL_NAMES: Final = TypeAdapter(list[str]) + + +def _is_chatgpt(deployment: Deployment) -> bool: + return ( + deployment.litellm_params.custom_llm_provider == "chatgpt" + or deployment.litellm_params.model.partition("/")[0] == "chatgpt" + ) + + +def _serves_by_name(deployment: Deployment, names: frozenset[str]) -> bool: + return not _is_chatgpt(deployment) and ( + deployment.model_name in names or deployment.model_info.team_public_model_name in names + ) + + +def _candidates(model: str, router: RegistryReader, team_id: str | None) -> tuple[Deployment, ...]: + registry: Final = router.registry() + candidates: Final = tuple( + registry.deployment(model_id) for model_id in registry.names(model, kind="route_candidates", team_id=team_id) + ) + return tuple(deployment for deployment in candidates if deployment is not None) + + +def routes_to_litellm( + model: object, + router: RegistryReader | None, + team_id: str | None, + bundled: tuple[str, ...], + prices: Mapping[str, object], +) -> bool: + if not isinstance(model, str) or router is None or model in bundled or f"chatgpt/{model}" in prices: + return False + deployments: Final = _candidates(model, router, team_id) + names: Final = frozenset({model, router.registry().resolve(model, kind="alias") or model}) + return bool(deployments) and all(_serves_by_name(deployment, names) for deployment in deployments) + + +def _key_and_team_allow(model: str, router: RegistryReader, token: UserAPIKeyAuth) -> bool: + allowed: Final = ( + _MODEL_NAMES.validate_python(token.models), # pyright: ignore[reportUnknownMemberType] # frozen list projection from the legacy token model + _MODEL_NAMES.validate_python(token.team_models), + ) + return all(_helper(model, router, grants, token.team_id, None) for grants in allowed) + + +def native_chatgpt_model(model: str, router: RegistryReader, token: UserAPIKeyAuth) -> str | None: + if model.startswith("chatgpt/") or _key_and_team_allow(model, router, token): + return None + prefixed: Final = f"chatgpt/{model}" + deployments: Final = _candidates(prefixed, router, token.team_id) + if not deployments or not all(_is_chatgpt(deployment) for deployment in deployments): + return None + return prefixed if _key_and_team_allow(prefixed, router, token) else None diff --git a/tests/unit/_auth/support/model_reference.py b/tests/unit/_auth/support/model_reference.py index e9155f8d256..5409f2a5389 100644 --- a/tests/unit/_auth/support/model_reference.py +++ b/tests/unit/_auth/support/model_reference.py @@ -66,6 +66,8 @@ def _single(model: str, facts: ModelAccessFacts) -> ModelAccessDecision: and facts.registry.present and facts.team_id is not None and model.startswith(f"model_name_{facts.team_id}_") + and record is not None + and record.deployment_team == facts.team_id else () ) if any(helper(candidate, facts, team_aliases=False) for candidate in owning): diff --git a/tests/unit/_auth/support/model_registry_reference.py b/tests/unit/_auth/support/model_registry_reference.py index d499bab565e..996d157c179 100644 --- a/tests/unit/_auth/support/model_registry_reference.py +++ b/tests/unit/_auth/support/model_registry_reference.py @@ -1,4 +1,4 @@ -"""Model authorization and registry ordering frozen from Phase 1, 8bd7f5089a75.""" +"""Phase 1 model policy, with ownership guard updated from main 2f430873aead.""" import re from collections.abc import Mapping, Sequence @@ -64,6 +64,8 @@ def _owning( 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") + if deployment is None or deployment.model_info.team_id != team: + return () public: Final = ( deployment.model_info.team_public_model_name if deployment is not None and deployment.model_info.team_id == team diff --git a/tests/unit/_auth/test_codex_models.py b/tests/unit/_auth/test_codex_models.py new file mode 100644 index 00000000000..2bec858e576 --- /dev/null +++ b/tests/unit/_auth/test_codex_models.py @@ -0,0 +1,144 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Final, cast + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth import codex_model_adapter as native +from litellm.router_api import RouterRegistry +from litellm.rust_bridge import _native +from litellm.types.router import Deployment +from tests.unit._auth.support import codex_reference as reference + + +@dataclass(frozen=True, slots=True) +class Registry: + deployments: tuple[Deployment | None, ...] + alias: str | None = None + groups: Mapping[str, tuple[str, ...]] = field(default_factory=dict[str, tuple[str, ...]]) + calls: list[tuple[object, ...]] = field( + default_factory=list[tuple[object, ...]] + ) # mutable-ok: ordered raw registry IO trace + + def names(self, model: str | None = None, *, kind: str = "all", team_id: str | None = None) -> tuple[str, ...]: + self.calls.append(("names", model, kind, team_id)) + return ( + tuple(str(index) for index in range(len(self.deployments))) + if kind == "route_candidates" + else ("model", "chatgpt/model") + ) + + def deployment(self, model_id: str) -> Deployment | None: + self.calls.append(("deployment", model_id)) + return self.deployments[int(model_id)] + + def resolve(self, model: str, *, kind: str) -> str | None: + self.calls.append(("resolve", model, kind)) + return self.alias + + def access_groups(self, model_name: str, team_id: str | None = None) -> tuple[str, ...]: + self.calls.append(("groups", model_name, team_id)) + return self.groups.get(model_name, ()) + + +@dataclass(frozen=True, slots=True) +class Router: + reader: Registry + + def registry(self) -> RouterRegistry: + self.reader.calls.append(("registry",)) + return cast( + RouterRegistry, self.reader + ) # cast-ok: narrow raw registry double implements only the exercised IO methods + + +def deployment( + name: str, backend: str = "openai/backend", provider: str | None = None, public: str | None = None +) -> Deployment: + return Deployment.model_validate( + { + "model_name": name, + "litellm_params": {"model": backend, "custom_llm_provider": provider}, + "model_info": {"team_public_model_name": public}, + } + ) + + +_DEPLOYMENTS: Final = ( + (), + (None,), + (deployment("model"),), + (deployment("target"),), + (deployment("internal", public="model"),), + (deployment("other"),), + (deployment("model", "chatgpt/backend"),), + (deployment("model", provider="chatgpt"),), + (deployment("model"), deployment("model", "chatgpt/backend")), + (None, deployment("model")), +) + + +@pytest.mark.parametrize("model", (None, [], 5, "", "model", "bundled", "priced", "chatgpt/model")) +@pytest.mark.parametrize("alias", (None, "target", "")) +@pytest.mark.parametrize("deployments", _DEPLOYMENTS) +@pytest.mark.parametrize("router_present", (False, True)) +def test_route_candidates_match_frozen_main_and_read_order( + monkeypatch: pytest.MonkeyPatch, + model: object, + alias: str | None, + deployments: tuple[Deployment | None, ...], + router_present: bool, +) -> None: + monkeypatch.setattr(native, "bundled_codex_models", lambda: {"bundled": None}) + prices: Final[Mapping[str, object]] = {"chatgpt/priced": {}} + monkeypatch.setattr(litellm, "model_cost", prices) + old: Final = Router(Registry(deployments, alias)) + new: Final = Router(Registry(deployments, alias)) + expected: Final = reference.routes_to_litellm(model, old if router_present else None, "team", ("bundled",), prices) + assert native.routes_to_litellm(model, new if router_present else None, "team") == expected + assert new.reader.calls == old.reader.calls + + +@pytest.mark.parametrize("model", ("model", "bundled", "chatgpt/model", "")) +@pytest.mark.parametrize( + "key", ((), ("model",), ("chatgpt/model",), ("other",), ("*",), ("all-team-models",), ("group",)) +) +@pytest.mark.parametrize("team", ((), ("model",), ("chatgpt/model",), ("other",))) +@pytest.mark.parametrize("deployments", _DEPLOYMENTS) +def test_chatgpt_prefix_matches_frozen_main_and_grant_read_order( + model: str, key: tuple[str, ...], team: tuple[str, ...], deployments: tuple[Deployment | None, ...] +) -> None: + groups: Final = {"model": ("group",), "chatgpt/model": ("group",)} + old: Final = Router(Registry(deployments, groups=groups)) + new: Final = Router(Registry(deployments, groups=groups)) + token: Final = UserAPIKeyAuth(models=list(key), team_models=list(team), team_id="team") + assert native.native_chatgpt_model(model, new, token) == reference.native_chatgpt_model(model, old, token) + assert new.reader.calls == old.reader.calls + + +def test_native_model_failure_is_terminal_before_registry_reads(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("injected native failure") + + def fail(_facts: object) -> object: + raise failure + + monkeypatch.setattr(_native, "request_auth_codex_model", fail) + router: Final = Router(Registry((deployment("model"),))) + with pytest.raises(NativeAuthRuntimeError) as error: + native.routes_to_litellm("model", router, "team") + assert error.value is failure + assert router.reader.calls == [] + + +@pytest.mark.parametrize("route", ("/backend-api/codex/responses", "/v1/responses", "/messages")) +def test_native_route_guard_precedes_all_model_reads(route: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(native, "bundled_codex_models", lambda: dict[str, object]()) + monkeypatch.setattr(litellm, "model_cost", {}) + router: Final = Router(Registry((deployment("model"),))) + assert native.is_codex_native_litellm_turn(route, "model", router, "team") == ( + route == "/backend-api/codex/responses" + ) + assert bool(router.reader.calls) == (route == "/backend-api/codex/responses") diff --git a/tests/unit/_auth/test_codex_normalization.py b/tests/unit/_auth/test_codex_normalization.py new file mode 100644 index 00000000000..47ef7b1a07d --- /dev/null +++ b/tests/unit/_auth/test_codex_normalization.py @@ -0,0 +1,168 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest +from fastapi import Request + +from litellm._auth.codex_normalization import CodexNormalizationFacts, codex_normalization +from litellm._auth.codex_turn import CodexAccessFacts, CodexTurnFacts +from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.codex_normalization_adapter import resolve_codex_native_chatgpt_model +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_reference import RouteChecks as FrozenRoutes +from tests.unit._auth.test_codex_models import Registry, Router, deployment +from tests.unit._auth.test_codex_turn import ( + _KEY, # pyright: ignore[reportPrivateUsage] # immutable raw key facts shared by the composed planner + _ROUTING, # pyright: ignore[reportPrivateUsage] # immutable raw routing facts shared by the composed planner +) +from tests.unit._auth.test_codex_turn import reference as reference_turn + +_NATIVE_ROUTE: Final = "/backend-api/codex/responses" +_TURN: Final = CodexTurnFacts( + "model", access=CodexAccessFacts(_KEY, False), routing=_ROUTING, catalog_loaded=True, slugs=() +) +_ALIAS: Final = CodexNormalizationFacts( + "Alias", + "/chat/completions", + True, + True, + context=route_context("/chat/completions"), + resolved=False, + marked=True, + model_loaded=True, + model="model", + explicit=False, + provider=False, + routing=_ROUTING, + body_present=True, + settings_loaded=True, + target="target", + provider_after=False, + turn=_TURN, +) + + +def reference(facts: CodexNormalizationFacts) -> str | None: + """Completed alias reads in _resolve_router_settings_model_group_alias on main 2f430873aead.""" + if not facts.request_present or not facts.router_present or not FrozenRoutes.is_llm_api_route(facts.route): + return None + if facts.resolved: + return None + if facts.explicit or facts.provider and facts.route != _NATIVE_ROUTE: + return None + if facts.provider and facts.routing is not None and facts.model in (facts.routing.bundled or ()): + return None + if facts.model is None or not facts.body_present: + return None + if not facts.target or facts.target == facts.model: + return None + if facts.provider_after and (facts.turn is None or not reference_turn(facts.turn)): + return None + return facts.target + + +def native(facts: CodexNormalizationFacts) -> str | None: + decision: Final = codex_normalization(facts) + assert decision.kind in ("Complete", "Rewrite") + return decision.target + + +@pytest.mark.parametrize("policy", (reference, native)) +@pytest.mark.parametrize("route", ("/key/info", "/chat/completions", _NATIVE_ROUTE)) +@pytest.mark.parametrize("provider", (False, True)) +@pytest.mark.parametrize("target", (None, "", "model", "target")) +@pytest.mark.parametrize("body", (False, True)) +def test_alias_mutation_only_after_main_admission( + policy: Callable[[CodexNormalizationFacts], str | None], route: str, provider: bool, target: str | None, body: bool +) -> None: + facts: Final = replace( + _ALIAS, + route=route, + context=route_context(route), + provider=provider, + provider_after=provider, + target=target, + body_present=body, + ) + expected: Final = ( + "target" + if route != "/key/info" and (not provider or route == _NATIVE_ROUTE) and target == "target" and body + else None + ) + assert policy(facts) == expected + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (replace(_ALIAS, request_present=False, context=None), "Complete"), + (replace(_ALIAS, router_present=False, context=None), "Complete"), + (replace(_ALIAS, context=None), "ReadContext"), + (replace(_ALIAS, resolved=None), "ReadResolved"), + (replace(_ALIAS, resolved=True), "Complete"), + (replace(_ALIAS, marked=False, explicit=None), "MarkResolved"), + (replace(_ALIAS, model_loaded=False, explicit=None), "ReadModel"), + (replace(_ALIAS, explicit=True, body_present=None), "Complete"), + (replace(_ALIAS, explicit=None), "ReadExplicit"), + (replace(_ALIAS, provider=None), "ReadProvider"), + (replace(_ALIAS, body_present=None), "ReadBody"), + (replace(_ALIAS, settings_loaded=False), "ReadSettings"), + (replace(_ALIAS, provider_after=None), "ReadProviderAfter"), + ( + replace(_ALIAS, route=_NATIVE_ROUTE, context=route_context(_NATIVE_ROUTE), provider=True, routing=None), + "ReadRouting", + ), + ( + replace( + _ALIAS, + route=_NATIVE_ROUTE, + context=route_context(_NATIVE_ROUTE), + provider=True, + provider_after=True, + turn=None, + ), + "ReadTurn", + ), + ( + replace( + _ALIAS, + route=_NATIVE_ROUTE, + context=route_context(_NATIVE_ROUTE), + provider=True, + provider_after=True, + turn=replace(_TURN, slugs=("model",)), + ), + "Complete", + ), + (CodexNormalizationFacts("Prefix", _NATIVE_ROUTE, True, model="model", model_loaded=True), "ReadRouter"), + (CodexNormalizationFacts("Prefix", _NATIVE_ROUTE, True, True, model="model", model_loaded=True), "ReadPrefix"), + (CodexNormalizationFacts("Prefix", "/v1/responses", True, model="model", model_loaded=True), "Complete"), + ), +) +def test_normalization_read_and_mutation_order(facts: CodexNormalizationFacts, expected: str) -> None: + assert codex_normalization(facts).kind == expected + + +@pytest.mark.parametrize("requested_scope", (None, "client-alias")) +def test_prefix_adapter_preserves_client_model_and_rewrites_request( + monkeypatch: pytest.MonkeyPatch, requested_scope: str | None +) -> None: + router: Final = Router(Registry((deployment("chatgpt/model", "chatgpt/backend"),))) + monkeypatch.setattr(proxy_server, "llm_router", router) + token: Final = UserAPIKeyAuth(models=["chatgpt/model"], team_models=["chatgpt/model"]) + request: Final = Request({"type": "http", "path": _NATIVE_ROUTE, "headers": [], "method": "POST"}) + if requested_scope is not None: + request.scope[CLIENT_REQUESTED_MODEL_SCOPE_KEY] = requested_scope + data: Final[dict[str, object]] = { + "model": "model", + "input": "hello", + } # mutable-ok: production mutates request-local body + resolve_codex_native_chatgpt_model(data, token, request, _NATIVE_ROUTE) + assert data == {"model": "chatgpt/model", "input": "hello"} + assert request.scope[CLIENT_REQUESTED_MODEL_SCOPE_KEY] == (requested_scope or "model") + before: Final = tuple(router.reader.calls) + resolve_codex_native_chatgpt_model(data, token, request, _NATIVE_ROUTE) + assert tuple(router.reader.calls) == before diff --git a/tests/unit/_auth/test_codex_turn.py b/tests/unit/_auth/test_codex_turn.py new file mode 100644 index 00000000000..6253f2185fc --- /dev/null +++ b/tests/unit/_auth/test_codex_turn.py @@ -0,0 +1,224 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.codex_models import CodexCandidate, CodexModelFacts +from litellm._auth.codex_turn import ( + CodexAccessFacts, + CodexCredentials, + CodexKeyFacts, + CodexTurnFacts, + JwtRead, + codex_access, + codex_jwt, + codex_jwt_token, + codex_key, + codex_turn, +) +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_enable import RouteEnableFacts +from litellm._auth.route_scope import RouteScopeFacts, RouteValues +from litellm._auth.routes import VirtualRouteFacts +from litellm.proxy.auth import handle_jwt +from litellm.proxy.auth.codex_turn_adapter import read_jwt_result +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_reference import RouteChecks as FrozenRoutes + +_CONTEXT: Final = route_context("/v1/responses") +_KEY: Final = CodexKeyFacts( + RouteScopeFacts("Proxy", _CONTEXT.route, False, False), + RouteEnableFacts(_CONTEXT, failed="Other"), + VirtualRouteFacts(_CONTEXT, None), + True, + "internal_user", + RouteScopeFacts("Admin", _CONTEXT.route, False, False), +) +_ROUTING: Final = CodexModelFacts( + "Routes", + "model", + True, + bundled=(), + priced=False, + candidates=(CodexCandidate(True, "openai", "openai/backend", "model"),), + alias_loaded=True, +) +_KEYS: Final = ( + _KEY, + replace(_KEY, proxy=RouteScopeFacts("Proxy", _CONTEXT.route, True, True, values=RouteValues())), + replace(_KEY, enabled=RouteEnableFacts(_CONTEXT, failed="Http")), + replace(_KEY, virtual_route=VirtualRouteFacts(_CONTEXT, ("/other",))), + replace( + _KEY, admin=RouteScopeFacts("Admin", _CONTEXT.route, True, True, values=RouteValues(names=(_CONTEXT.route,))) + ), + replace( + _KEY, + role="proxy_admin", + admin=RouteScopeFacts("Admin", _CONTEXT.route, True, True, values=RouteValues(names=(_CONTEXT.route,))), + ), +) + + +def _scope(facts: RouteScopeFacts | None) -> bool: + assert facts is not None + if not facts.present or facts.kind == "Admin" and not facts.premium: + return True + assert facts.values is not None + values: Final = facts.values.text if facts.values.text is not None else facts.values.names + return (facts.route in values) == (facts.kind == "Proxy") + + +def reference(facts: CodexTurnFacts) -> bool: + """Completed raw facts evaluated in main 2f430873aead order, without native policy.""" + credentials: Final = facts.credentials + if facts.native and ( + credentials is None or not credentials.bearer_present or not credentials.gateway_present or credentials.equal + ): + return False + routing: Final = facts.routing + if ( + facts.model is None + or routing is None + or not routing.router_present + or routing.priced + or facts.model in (routing.bundled or ()) + ): + return False + candidates: Final = routing.candidates + names: Final = (facts.model, routing.alias or facts.model) + if not candidates or any( + candidate.provider == "chatgpt" + or (candidate.model or "").partition("/")[0] == "chatgpt" + or (candidate.name not in names and candidate.public_name not in names) + for candidate in candidates + ): + return False + access: Final = facts.access + assert access is not None + key: Final = access.key + if not _scope(key.proxy): + return False + assert key.enabled is not None and key.virtual_route is not None + if key.enabled.failed == "Http": + return False + if key.virtual_route.allowed_routes and not FrozenRoutes.check_route_access( + _CONTEXT.route, key.virtual_route.allowed_routes + ): + return False + if key.role != "proxy_admin" and not _scope(key.admin): + return False + if access.jwt_token_present and access.jwt == "Exception": + return False + return facts.slugs is not None and facts.model not in facts.slugs + + +def native(facts: CodexTurnFacts) -> bool: + decision: Final = codex_turn(facts) + assert decision.kind == "Complete" + return decision.allowed + + +@pytest.mark.parametrize("policy", (reference, native)) +@pytest.mark.parametrize("key", _KEYS) +@pytest.mark.parametrize("jwt", (None, "Returned", "Exception")) +@pytest.mark.parametrize( + "credentials", + ( + CodexCredentials(True, True, False), + CodexCredentials(False, True, False), + CodexCredentials(True, False, False), + CodexCredentials(True, True, True), + ), +) +@pytest.mark.parametrize("native_turn", (False, True)) +@pytest.mark.parametrize("routed", (False, True)) +@pytest.mark.parametrize("slugs", (None, (), ("other",), ("model",))) +def test_turn_destination_matches_main( + policy: Callable[[CodexTurnFacts], bool], + key: CodexKeyFacts, + jwt: JwtRead | None, + credentials: CodexCredentials, + native_turn: bool, + routed: bool, + slugs: tuple[str, ...] | None, +) -> None: + facts: Final = CodexTurnFacts( + "model", + native_turn, + credentials, + CodexAccessFacts(key, jwt is not None, jwt), + _ROUTING if routed else replace(_ROUTING, bundled=("model",)), + True, + slugs, + ) + expected: Final = ( + (not native_turn or credentials == CodexCredentials(True, True, False)) + and routed + and key in (_KEYS[0], _KEYS[5]) + and jwt != "Exception" + and slugs is not None + and "model" not in slugs + ) + assert policy(facts) == expected + + +@pytest.mark.parametrize( + ("facts", "stage"), + ( + (CodexTurnFacts("model", native=True), "ReadCredentials"), + (CodexTurnFacts("model", native=True, credentials=CodexCredentials(True, True, False)), "ReadAccess"), + (CodexTurnFacts("model", access=CodexAccessFacts(_KEYS[1], False)), "ReadRouting"), + (CodexTurnFacts("model", access=CodexAccessFacts(_KEY, True), routing=_ROUTING), "ReadJwt"), + (CodexTurnFacts("model", access=CodexAccessFacts(_KEY, False), routing=_ROUTING), "ReadCatalog"), + ), +) +def test_turn_retains_main_read_order(facts: CodexTurnFacts, stage: str) -> None: + assert codex_turn(facts).kind == stage + + +@pytest.mark.parametrize( + ("facts", "stage"), + ( + (CodexKeyFacts(), "ReadProxy"), + (CodexKeyFacts(proxy=_KEY.proxy), "ReadEnabled"), + (CodexKeyFacts(proxy=_KEY.proxy, enabled=_KEY.enabled), "ReadVirtual"), + (replace(_KEY, role_loaded=False), "ReadRole"), + (replace(_KEY, admin=None), "ReadAdmin"), + ), +) +def test_key_route_reads_are_selected_in_order(facts: CodexKeyFacts, stage: str) -> None: + assert codex_key(facts).kind == stage + + +@pytest.mark.parametrize("claims", (False, True)) +@pytest.mark.parametrize("token", (False, True)) +def test_standard_jwt_selection(claims: bool, token: bool) -> None: + assert codex_jwt_token(claims_present=claims, token_present=token).allowed == (claims and not token) + + +@pytest.mark.parametrize("result", ("Returned", "Exception")) +def test_jwt_validator_return_is_not_reinterpreted(result: JwtRead) -> None: + assert codex_jwt(result).allowed == (result == "Returned") + assert codex_access(CodexAccessFacts(_KEYS[1], True, result)).allowed is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", (None, ValueError("validator rejected"), NativeAuthRuntimeError("native failed"))) +async def test_jwt_wrapper_never_turns_a_native_fault_into_subscription_fallback( + monkeypatch: pytest.MonkeyPatch, failure: Exception | None +) -> None: + class Validator: + @staticmethod + async def auth_builder(**_kwargs: object) -> bool: + if failure is not None: + raise failure + return False + + monkeypatch.setattr(handle_jwt, "JWTAuthManager", Validator) + if isinstance(failure, NativeAuthRuntimeError): + with pytest.raises(NativeAuthRuntimeError) as raised: + await read_jwt_result("test-credential", "model", {}) + assert raised.value is failure + else: + assert await read_jwt_result("test-credential", "model", {}) == ("Returned" if failure is None else "Exception") diff --git a/tests/unit/_auth/test_route_scope.py b/tests/unit/_auth/test_route_scope.py new file mode 100644 index 00000000000..10fda677d03 --- /dev/null +++ b/tests/unit/_auth/test_route_scope.py @@ -0,0 +1,61 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy import proxy_server +from litellm.proxy.auth import route_scope_adapter as native + + +@dataclass(frozen=True, slots=True) +class Result: + allowed: bool + status: int | None = None + detail: object = None + error: str | None = None + + +def outcome(call: Callable[[], object]) -> Result: + try: + return Result(call() is not False) + except HTTPException as error: + return Result(False, error.status_code, error.detail) + except TypeError as error: + return Result(False, error=str(error)) + + +def reference(route: str, settings: Mapping[str, object], premium: object, admin: bool) -> object: + """RouteChecks wrapper behavior frozen from main 2f430873aead.""" + if admin: + if "admin_only_routes" not in settings or premium is not True: + return None + if route in settings["admin_only_routes"]: # pyright: ignore[reportOperatorIssue] # freeze malformed configuration TypeError as well as valid membership + raise HTTPException( + status_code=403, detail=f"user not allowed to access this route. Route={route} is an admin only route" + ) + return None + if "allowed_routes" not in settings: + return True + return route in settings["allowed_routes"] # pyright: ignore[reportOperatorIssue] # frozen Python membership preserves string, list, mapping and invalid-scalar semantics + + +@pytest.mark.parametrize("admin", (False, True)) +@pytest.mark.parametrize("premium", (False, True, 1, None)) +@pytest.mark.parametrize("route", ("/v1/responses", "/missing", "responses", "")) +@pytest.mark.parametrize("present", (False, True)) +@pytest.mark.parametrize( + "value", + (None, True, 3, (), ("/v1/responses",), "/v1/responses", {"/v1/responses": False}, (None, "/v1/responses", 3)), +) +def test_route_scope_matches_frozen_main( + monkeypatch: pytest.MonkeyPatch, admin: bool, premium: object, route: str, present: bool, value: object +) -> None: + settings: Final[Mapping[str, object]] = ( + {("admin_only_routes" if admin else "allowed_routes"): value} if present else {} + ) + monkeypatch.setattr(proxy_server, "general_settings", dict(settings)) + monkeypatch.setattr(proxy_server, "premium_user", premium) + policy: Final = native.custom_admin_only_route_check if admin else native.is_allowed_on_proxy + assert outcome(lambda: policy(route)) == outcome(lambda: reference(route, settings, premium, admin))