diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 1d5ca54308e..191bd2418e5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4567,6 +4567,7 @@ dependencies = [ name = "litellm-request-auth-core" version = "0.1.0" dependencies = [ + "fancy-regex 0.19.2", "litellm-secret-value", "rstest", "serde", diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index b7500311107..9bc7b42a37b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -63,9 +63,42 @@ mod _native { }; #[pymodule_export] use crate::request_auth::{ - NativeAuthContinuation, request_auth_budgets, request_auth_litellm_header_token, - request_auth_scheme_credential, request_auth_select_api_key, - request_auth_select_websocket_api_key, request_auth_start, + NativeAuthContinuation, request_auth_access, request_auth_admission, + request_auth_agent_access, request_auth_agent_ceiling, request_auth_agent_ceiling_source, + request_auth_agent_grants, request_auth_agent_group_load, request_auth_agent_invocation, + request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption, + request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start, + request_auth_budgets, request_auth_common_budget, request_auth_common_budget_scopes, + request_auth_common_flow, request_auth_common_read, request_auth_context, + request_auth_context_end_user, request_auth_context_member, request_auth_context_models, + request_auth_context_percentage, request_auth_context_temporary, + request_auth_context_throttle, request_auth_context_token, request_auth_custom_flow, + request_auth_end_user_apply, request_auth_end_user_assignments, + request_auth_end_user_default_id, request_auth_end_user_lookup, + request_auth_end_user_normalize_id, request_auth_end_user_registry, + request_auth_end_user_validation, request_auth_expired, request_auth_fallback_access, + request_auth_fallback_candidate, request_auth_key_model_access, + request_auth_litellm_header_token, request_auth_llm_route, request_auth_managed, + request_auth_managed_model_grants, request_auth_managed_models, request_auth_member_limit, + request_auth_member_models, request_auth_model_access, request_auth_model_attribution, + request_auth_model_budget_check, request_auth_model_budget_query, + request_auth_model_budget_resolve, request_auth_model_groups, request_auth_model_helpers, + request_auth_model_projection, request_auth_org_admin, request_auth_org_admin_raw, + request_auth_org_route, request_auth_owner_teams, request_auth_personal_budget, + request_auth_personal_model_listing, request_auth_personal_model_request, + request_auth_personal_models, request_auth_personal_owner, request_auth_reservation, + request_auth_resolved_models, request_auth_route_access, request_auth_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_select_api_key, request_auth_select_websocket_api_key, + request_auth_spend_counter_policy, request_auth_start, request_auth_team_aliases, + request_auth_team_override, request_auth_unpriced_model, + }; + #[pymodule_export] + use crate::request_auth::{ + request_auth_budget_lookup_emails, request_auth_budget_lookup_member, + request_auth_budget_lookup_organization, request_auth_budget_lookup_soft, }; #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; @@ -122,6 +155,88 @@ mod tests { Python::initialize(); Python::attach(|py| { let mut expected = vec![ + "NativeAuthContinuation", + "request_auth_access", + "request_auth_admission", + "request_auth_agent_access", + "request_auth_agent_ceiling", + "request_auth_agent_ceiling_source", + "request_auth_agent_grants", + "request_auth_agent_group_load", + "request_auth_agent_invocation", + "request_auth_budget_alert", + "request_auth_budget_alert_plan", + "request_auth_budget_exemption", + "request_auth_budget_price", + "request_auth_budget_route", + "request_auth_budget_window_start", + "request_auth_budgets", + "request_auth_common_budget", + "request_auth_common_budget_scopes", + "request_auth_common_flow", + "request_auth_common_read", + "request_auth_context", + "request_auth_context_end_user", + "request_auth_context_member", + "request_auth_context_models", + "request_auth_context_percentage", + "request_auth_context_temporary", + "request_auth_context_throttle", + "request_auth_context_token", + "request_auth_custom_flow", + "request_auth_end_user_apply", + "request_auth_end_user_assignments", + "request_auth_end_user_default_id", + "request_auth_end_user_lookup", + "request_auth_end_user_normalize_id", + "request_auth_end_user_registry", + "request_auth_end_user_validation", + "request_auth_expired", + "request_auth_fallback_access", + "request_auth_fallback_candidate", + "request_auth_key_model_access", + "request_auth_litellm_header_token", + "request_auth_llm_route", + "request_auth_managed", + "request_auth_managed_model_grants", + "request_auth_managed_models", + "request_auth_member_limit", + "request_auth_member_models", + "request_auth_model_access", + "request_auth_model_attribution", + "request_auth_model_budget_check", + "request_auth_model_budget_query", + "request_auth_model_budget_resolve", + "request_auth_model_groups", + "request_auth_model_helpers", + "request_auth_model_projection", + "request_auth_org_admin", + "request_auth_org_admin_raw", + "request_auth_org_route", + "request_auth_owner_teams", + "request_auth_personal_budget", + "request_auth_personal_model_listing", + "request_auth_personal_model_request", + "request_auth_personal_models", + "request_auth_personal_owner", + "request_auth_reservation", + "request_auth_resolved_models", + "request_auth_route_access", + "request_auth_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_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", "RustBridgeDeclined", "RustUpstreamError", "ForkedAfterNativeRuntimeStarted", diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index d96bf7b5bde..9d07c977f29 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -9,11 +9,273 @@ use pyo3::{ type ProjectedStep = (Py, Option); +#[pyfunction] +pub(crate) fn request_auth_spend_counter_policy( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::spend_counter_policy::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::spend_counter_policy::decide(facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_common_read( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::ReadFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::read(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_reservation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::reservation::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::reservation::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_common_flow( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::plan(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_common_budget_scopes( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::BudgetFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::budgets(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_custom_flow( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_flow::CustomFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_flow::custom(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_expired( + py: Python<'_>, + expires: &Bound<'_, PyAny>, + now: &Bound<'_, PyAny>, +) -> PyResult { + let expires: core::protocol::DateTime = from_py(expires)?; + let now: core::protocol::DateTime = from_py(now)?; + py.detach(move || expires.before(&now)).map_err(map_error) +} + +#[pyfunction] +pub(crate) fn request_auth_context( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::auth_context::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::auth_context::plan(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_grants( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_grants::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_grants::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_managed( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed::ManagedFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::managed::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_access( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_access::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_access::decide(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_invocation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_access::Invocation = from_py(facts)?; + to_py( + py, + &py.detach(move || core::agent_access::invocation(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_admission( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::admission::AdmissionFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::admission::plan(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_access(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::access::AccessFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::access::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_soft( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult { + let facts: core::budget_lookup::SoftFacts = from_py(facts)?; + Ok(py.detach(move || core::budget_lookup::read_soft_emails(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_organization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::OrganizationFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_lookup::organization(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_member( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::MemberFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_lookup::member(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_emails( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::EmailFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_lookup::emails(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_window_start( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_budgets::WindowFacts = from_py(facts)?; + Ok(py.detach(move || core::common_budgets::window_start(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_resolve( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_data::ResolveFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::resolve(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_query( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_data::QueryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::query(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_check( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_data::Snapshot = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::check(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_price( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::ModelPriceFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_pricing::price(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_route( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::RouteBudgetFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_pricing::route(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_common_budget( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_budgets::BudgetFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::common_budgets::check(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_alert( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult { + let facts: core::common_budgets::AlertFacts = from_py(facts)?; + Ok(py.detach(move || core::common_budgets::alert(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_member_limit( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::common_budgets::MemberLimitFacts = from_py(facts)?; + py.detach(move || core::common_budgets::member_limit(&facts)) + .map_err(map_error) +} + fn map_error(error: core::Error) -> PyErr { let message = litellm_error::Report::new(&error).to_string(); match error { core::Error::IncompatibleDateTimes => PyTypeError::new_err(message), core::Error::InvalidAnswer => PyRuntimeError::new_err(message), + core::Error::Pattern(_) => PyRuntimeError::new_err(message), } } @@ -122,3 +384,598 @@ pub(crate) fn request_auth_scheme_credential(value: &str) -> Option { pub(crate) fn request_auth_litellm_header_token(value: &str) -> String { credential::litellm_header_token(value) } + +#[pyfunction] +pub(crate) fn request_auth_route_access( + py: Python<'_>, + request: &Bound<'_, PyAny>, +) -> PyResult> { + let request: core::route_access::Request = from_py(request)?; + let decision = py + .detach(move || core::route_access::decide(request)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_org_route( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::org_routes::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::org_routes::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_org_admin( + py: Python<'_>, + memberships: &Bound<'_, PyAny>, + organizations: Vec, +) -> PyResult { + let memberships: Vec = from_py(memberships)?; + Ok(py.detach(move || core::org_routes::user_is_admin(&memberships, &organizations))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_access( + py: Python<'_>, + request: &Bound<'_, PyAny>, +) -> PyResult> { + let request: core::model_access::Request = from_py(request)?; + let decision = py + .detach(move || core::model_access::decide(request)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_model_groups( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_groups::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::model_groups::decide(facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_fallback_access( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::fallback_access::Facts = from_py(facts)?; + let decision = py + .detach(move || core::fallback_access::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_exemption( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::ExemptionFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_pricing::exemption(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_key_model_access( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::key_model_access::Facts = from_py(facts)?; + let decision = py + .detach(move || core::key_model_access::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_unpriced_model( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::unpriced_models::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::unpriced_models::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_alert_plan( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_alerts::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_alerts::plan(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_attribution( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_attribution::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_attribution::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_fallback_candidate( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::fallback_candidate::Facts = from_py(facts)?; + let decision = py + .detach(move || core::fallback_candidate::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_budget( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_budget::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_budget::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_models::Request = from_py(facts)?; + let decision = py + .detach(move || core::personal_models::decide(facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_owner_teams( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::owner_teams::Facts = from_py(facts)?; + let decision = py + .detach(move || core::owner_teams::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_context_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::ModelCheckFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::model_checks(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_end_user( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::EndUserFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::end_user_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_token( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::TokenDefaultsFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::token_assignments(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_member( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::MemberFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::member_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_temporary( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::TemporaryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::temporary(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_throttle( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::ThrottleFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::throttle(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_percentage( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::NumericFacts = from_py(facts)?; + Ok(py.detach(move || core::budget_context::percentage(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_team_override( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::team_override::Facts = from_py(facts)?; + let decision = py + .detach(move || core::team_override::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_model_helpers( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_helpers::Facts = from_py(facts)?; + let decision = py + .detach(move || core::model_helpers::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_owner( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_owner::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_owner::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_model_request( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_model_request::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_model_request::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_member_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::member_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::member_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_resolved_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::resolved_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::resolved_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_normalize_id( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::IdFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::normalize_id(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_default_id( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::DefaultFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::default_id(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_apply( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::ApplyFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::apply_default(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_registry( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::RegistryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::registry(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_assignments( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::AssignmentFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::assignments(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_lookup( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::LookupFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::lookup(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_validation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::ValidationFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::validation(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_model_listing( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_model_listing::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_model_listing::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_ceiling( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_ceilings::Facts = from_py(facts)?; + let decision = py + .detach(move || core::agent_ceilings::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_group_load( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::agent_ceilings::LoadFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_ceilings::load(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_team_aliases( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::team_aliases::Facts = from_py(facts)?; + let decision = py + .detach(move || core::team_aliases::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_managed_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed_models::Facts = from_py(facts)?; + let decision = py + .detach(move || core::managed_models::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_ceiling_source( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed_models::CeilingSource = from_py(facts)?; + to_py( + py, + &py.detach(move || core::managed_models::ceiling_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_managed_model_grants( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::managed_models::GrantFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::managed_models::grants(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::RoleFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::role(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_api( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::ApiFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::api(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_jwt( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::JwtFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::jwt(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_inside( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::InsideFacts = from_py(facts)?; + let decision = py.detach(move || core::route_roles::inside(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_expand( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::Entry = from_py(facts)?; + let decision = py.detach(move || core::route_roles::expand(&facts)); + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_role_listed( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_roles::ListedFacts = from_py(facts)?; + let decision = py + .detach(move || core::route_roles::listed(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_org_admin_raw( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::org_routes::AdminFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::org_routes::admin(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_model_projection( + py: Python<'_>, + facts: &Bound<'_, PyAny>, + reads: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_access::Facts = from_py(facts)?; + let reads: Vec = from_py(reads)?; + let decision = py + .detach(move || core::model_projection::decide(&facts, &reads)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_route_enable( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::route_enable::Facts = from_py(facts)?; + let decision = py + .detach(move || core::route_enable::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +#[pyo3(signature = (context, configured=None))] +pub(crate) fn request_auth_llm_route( + py: Python<'_>, + context: &Bound<'_, PyAny>, + configured: Option, +) -> PyResult> { + let context: core::route_access::RouteContext = from_py(context)?; + let decision = py + .detach(move || core::route_enable::llm_route(&context, configured)) + .map_err(map_error)?; + to_py(py, &decision) +} diff --git a/litellm-rust/crates/request-auth-core/Cargo.toml b/litellm-rust/crates/request-auth-core/Cargo.toml index 6b82bde399e..0b51939bb9c 100644 --- a/litellm-rust/crates/request-auth-core/Cargo.toml +++ b/litellm-rust/crates/request-auth-core/Cargo.toml @@ -5,6 +5,7 @@ edition.workspace = true license.workspace = true [dependencies] +fancy-regex.workspace = true litellm-secret-value.workspace = true serde.workspace = true sha2.workspace = true diff --git a/litellm-rust/crates/request-auth-core/src/access.rs b/litellm-rust/crates/request-auth-core/src/access.rs new file mode 100644 index 00000000000..99f2553c3a0 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/access.rs @@ -0,0 +1,480 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum AccessFacts { + ConcurrentFailures { + failures: Vec, + terminal: Vec, + }, + UserParamRequest { + enabled: bool, + }, + ClientTagRequest { + enabled: bool, + llm_route: Option, + }, + GuardrailRequest { + requested_keys: Vec, + }, + TraceRequest { + required: bool, + }, + Search { + name: String, + allowed: Vec, + strict: bool, + }, + TraceGate { + required: bool, + method: String, + llm_route: bool, + passthrough: bool, + telemetry: bool, + }, + TraceStep { + source: TraceSource, + value: Option, + present: bool, + headers_only: bool, + }, + Blocked { + blocked: bool, + }, + UserParam { + enabled: bool, + method: Option, + llm_route: bool, + mcp_route: bool, + has_user: bool, + }, + ClientTags { + enabled: bool, + llm_route: bool, + has_tags: bool, + }, + Guardrails { + requested_keys: Vec, + modify_guardrails: Option, + }, + AgentTrace { + required: bool, + header: Option, + }, + TeamTrace { + required: bool, + method: String, + llm_route: bool, + passthrough: bool, + telemetry: bool, + headers_only: bool, + header: Option, + body_trace: Option, + body_trace_present: bool, + traceparent_id: Option, + selected_trace: Option, + selected_trace_present: bool, + other_trace: Option, + }, + Tools { + requested: Vec, + key_allowed: Option>, + team_allowed: Option>, + }, + GrantLayers { + deny_by_default: bool, + identity_present: bool, + master_key: bool, + dashboard: bool, + virtual_key: bool, + session_token: bool, + team_present: bool, + }, + VectorStores { + requested: Vec, + allowed: Option>, + strict: bool, + }, + TeamLookup { + dashboard: bool, + missing: bool, + recorded_models: Vec, + allow_unavailable: Option, + }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum Denial { + Search, + Blocked, + UserParam, + ClientTags, + Guardrails, + AgentTrace, + TeamTrace, + Tools, + VectorStores, + MissingVectorGrant, + TeamLookup, +} + +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, strum::VariantArray)] +pub enum TraceSource { + Gate, + Header, + Body, + Traceparent, + Selected, + Other, +} + +#[derive(Debug, Default, PartialEq, Serialize)] +pub struct AccessDecision { + pub read: Option, + pub denial: Option, + pub rejected: Vec, + pub layers: Vec, + pub rejected_index: Option, + pub next_trace: Option, + pub failure_index: Option, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum AccessRead { + UserParam, + ClientTagRoute, + ClientTags, + GuardrailPermission, + UnavailablePolicy, +} + +fn modifies_guardrails(keys: &[String]) -> bool { + keys.iter().any(|key| { + matches!( + key.as_str(), + "guardrails" + | "disable_global_guardrails" + | "disable_global_guardrail" + | "opted_out_global_guardrails" + ) + }) +} + +fn deny_if(deny: bool, denial: Denial) -> AccessDecision { + AccessDecision { + denial: deny.then_some(denial), + ..AccessDecision::default() + } +} + +fn nonempty(value: &Option) -> bool { + value.as_ref().is_some_and(|value| !value.is_empty()) +} + +pub fn decide(facts: AccessFacts) -> AccessDecision { + match facts { + AccessFacts::ConcurrentFailures { failures, terminal } => AccessDecision { + failure_index: terminal + .iter() + .position(|failed| *failed) + .or_else(|| failures.iter().position(|failed| *failed)), + ..AccessDecision::default() + }, + AccessFacts::UserParamRequest { enabled } => AccessDecision { + read: enabled.then_some(AccessRead::UserParam), + ..AccessDecision::default() + }, + AccessFacts::ClientTagRequest { enabled, llm_route } => AccessDecision { + read: match (enabled, llm_route) { + (false, _) | (_, Some(false)) => None, + (_, None) => Some(AccessRead::ClientTagRoute), + (_, Some(true)) => Some(AccessRead::ClientTags), + }, + ..AccessDecision::default() + }, + AccessFacts::GuardrailRequest { requested_keys } => AccessDecision { + read: modifies_guardrails(&requested_keys).then_some(AccessRead::GuardrailPermission), + ..AccessDecision::default() + }, + AccessFacts::TraceRequest { required } => AccessDecision { + next_trace: required.then_some(TraceSource::Gate), + ..AccessDecision::default() + }, + AccessFacts::Search { + name, + allowed, + strict, + } => deny_if( + (strict || !allowed.is_empty()) && !allowed.contains(&name), + Denial::Search, + ), + AccessFacts::TraceGate { + required, + method, + llm_route, + passthrough, + telemetry, + } => AccessDecision { + next_trace: (required + && !matches!(method.as_str(), "GET" | "HEAD" | "OPTIONS") + && (llm_route || passthrough) + && !telemetry) + .then_some(TraceSource::Header), + ..AccessDecision::default() + }, + AccessFacts::TraceStep { + source, + value, + present, + headers_only, + } => { + if nonempty(&value) { + return AccessDecision::default(); + } + let next_trace = match source { + TraceSource::Gate => Some(TraceSource::Header), + TraceSource::Header if headers_only => None, + TraceSource::Header => Some(TraceSource::Body), + TraceSource::Body if !present => Some(TraceSource::Traceparent), + TraceSource::Body | TraceSource::Traceparent => Some(TraceSource::Selected), + TraceSource::Selected if !present => Some(TraceSource::Other), + TraceSource::Selected | TraceSource::Other => None, + }; + AccessDecision { + denial: next_trace.is_none().then_some(Denial::TeamTrace), + next_trace, + ..AccessDecision::default() + } + } + AccessFacts::Blocked { blocked } => deny_if(blocked, Denial::Blocked), + AccessFacts::UserParam { + enabled, + method, + llm_route, + mcp_route, + has_user, + } => deny_if( + enabled + && method.is_some_and(|value| value.to_uppercase() == "POST") + && llm_route + && !mcp_route + && !has_user, + Denial::UserParam, + ), + AccessFacts::ClientTags { + enabled, + llm_route, + has_tags, + } => deny_if(enabled && llm_route && has_tags, Denial::ClientTags), + AccessFacts::Guardrails { + requested_keys, + modify_guardrails, + } => deny_if( + modify_guardrails == Some(false) + && requested_keys.iter().any(|key| { + matches!( + key.as_str(), + "guardrails" + | "disable_global_guardrails" + | "disable_global_guardrail" + | "opted_out_global_guardrails" + ) + }), + Denial::Guardrails, + ), + AccessFacts::AgentTrace { required, header } => { + deny_if(required && !nonempty(&header), Denial::AgentTrace) + } + AccessFacts::TeamTrace { + required, + method, + llm_route, + passthrough, + telemetry, + headers_only, + header, + body_trace, + body_trace_present, + traceparent_id, + selected_trace, + selected_trace_present, + other_trace, + } => { + let gate = required + && !matches!(method.as_str(), "GET" | "HEAD" | "OPTIONS") + && (llm_route || passthrough) + && !telemetry; + let body = nonempty(&body_trace) + || (!body_trace_present && nonempty(&traceparent_id)) + || if selected_trace_present { + nonempty(&selected_trace) + } else { + nonempty(&other_trace) + }; + deny_if( + gate && !nonempty(&header) && (headers_only || !body), + Denial::TeamTrace, + ) + } + AccessFacts::Tools { + requested, + key_allowed, + team_allowed, + } => { + let allowed = key_allowed + .filter(|allowed| !allowed.is_empty()) + .or(team_allowed); + let rejected: Vec<_> = match allowed { + Some(allowed) if !allowed.is_empty() => requested + .into_iter() + .filter(|name| !allowed.contains(name)) + .collect(), + _ => Vec::new(), + }; + AccessDecision { + denial: (!rejected.is_empty()).then_some(Denial::Tools), + rejected, + ..AccessDecision::default() + } + } + AccessFacts::GrantLayers { + deny_by_default, + identity_present, + master_key, + dashboard, + virtual_key, + session_token, + team_present, + } => { + let layers = if !deny_by_default || !identity_present || master_key || dashboard { + Vec::new() + } else if !(virtual_key && !session_token) && !team_present { + vec!["user".into()] + } else { + [ + ("key", virtual_key && !session_token), + ("team", team_present), + ] + .into_iter() + .filter_map(|(name, required)| required.then_some(name.into())) + .collect() + }; + AccessDecision { + layers, + ..AccessDecision::default() + } + } + AccessFacts::VectorStores { + requested, + allowed, + strict, + } => match allowed { + None => deny_if(strict, Denial::MissingVectorGrant), + Some(allowed) if allowed.is_empty() => deny_if(strict, Denial::MissingVectorGrant), + Some(allowed) => { + let rejected_index = requested.iter().position(|name| !allowed.contains(name)); + AccessDecision { + denial: rejected_index.map(|_| Denial::VectorStores), + rejected_index, + ..AccessDecision::default() + } + } + }, + AccessFacts::TeamLookup { + dashboard, + missing, + recorded_models, + allow_unavailable, + } => { + if !dashboard && !missing && recorded_models.is_empty() && allow_unavailable.is_none() { + return AccessDecision { + read: Some(AccessRead::UnavailablePolicy), + ..AccessDecision::default() + }; + } + deny_if( + !dashboard + && (missing || (recorded_models.is_empty() && allow_unavailable != Some(true))), + Denial::TeamLookup, + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::disabled(false, None, None)] + #[case::read_route(true, None, Some(AccessRead::ClientTagRoute))] + #[case::non_llm(true, Some(false), None)] + #[case::read_tags(true, Some(true), Some(AccessRead::ClientTags))] + fn client_tag_reads( + #[case] enabled: bool, + #[case] llm_route: Option, + #[case] read: Option, + ) { + assert_eq!( + decide(AccessFacts::ClientTagRequest { enabled, llm_route }).read, + read + ); + } + + #[rstest] + #[case::allow(vec![false, false], vec![false, false], None)] + #[case::scope_priority(vec![true, true], vec![false, false], Some(0))] + #[case::second_scope(vec![false, true], vec![false, false], Some(1))] + #[case::terminal_fault(vec![true, true], vec![false, true], Some(1))] + fn gathered_failures_preserve_scope_order_unless_native_failed( + #[case] failures: Vec, + #[case] terminal: Vec, + #[case] expected: Option, + ) { + assert_eq!( + decide(AccessFacts::ConcurrentFailures { failures, terminal }).failure_index, + expected + ); + } + + #[rstest] + #[case::key_overrides_team(Some(vec!["key".into()]), Some(vec!["team".into()]), vec!["team".into()])] + #[case::empty_key_inherits(Some(vec![]), Some(vec!["team".into()]), vec!["key".into()])] + #[case::unrestricted(None, Some(vec![]), vec![])] + fn tools( + #[case] key: Option>, + #[case] team: Option>, + #[case] rejected: Vec, + ) { + let result = decide(AccessFacts::Tools { + requested: vec!["key".into(), "team".into()], + key_allowed: key, + team_allowed: team, + }); + assert_eq!( + result.denial, + (!rejected.is_empty()).then_some(Denial::Tools) + ); + assert_eq!(result.rejected, rejected); + } + + #[rstest] + #[case::missing_strict(None, true, Some(Denial::MissingVectorGrant))] + #[case::empty_legacy(Some(vec![]), false, None)] + #[case::empty_strict(Some(vec![]), true, Some(Denial::MissingVectorGrant))] + #[case::denied(Some(vec!["other".into()]), false, Some(Denial::VectorStores))] + #[case::granted(Some(vec!["store".into()]), true, None)] + fn vectors( + #[case] allowed: Option>, + #[case] strict: bool, + #[case] denial: Option, + ) { + let result = decide(AccessFacts::VectorStores { + requested: vec!["store".into()], + allowed, + strict, + }); + assert_eq!(result.denial, denial); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/admission.rs b/litellm-rust/crates/request-auth-core/src/admission.rs new file mode 100644 index 00000000000..8c23d47ef8e --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/admission.rs @@ -0,0 +1,183 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum AdmissionFacts { + CommonGate { + force_virtual_key_checks: bool, + public_route: bool, + additional_public_route: bool, + unauthenticated_passthrough: bool, + master_key_present: bool, + jwt_enabled: bool, + oauth_enabled: bool, + oauth_proxy_enabled: bool, + custom_auth: bool, + custom_auth_common_checks: bool, + }, + Reservation { + skip_budget_checks: bool, + disabled: bool, + apply_user_budget_to_team_keys: bool, + fail_closed: bool, + billing_agent: bool, + }, + Ownership { + virtual_key: bool, + custom_auth: bool, + jwt: bool, + oauth: bool, + }, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum CommonSkip { + PublicRoute, + UnauthenticatedPassthrough, + NoAuthDevelopment, + CustomAuth, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum Owner { + Python, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum Candidate { + VirtualKey, + CustomAuth, + Validator, + Other, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum AdmissionPlan { + Common { + skip: Option, + }, + Reserve { + reserve: bool, + apply_user_budget_to_team_keys: bool, + fail_closed: bool, + }, + Observe { + owner: Owner, + candidate: Candidate, + }, +} + +pub fn plan(facts: AdmissionFacts) -> AdmissionPlan { + match facts { + AdmissionFacts::CommonGate { + force_virtual_key_checks, + public_route, + additional_public_route, + unauthenticated_passthrough, + master_key_present, + jwt_enabled, + oauth_enabled, + oauth_proxy_enabled, + custom_auth, + custom_auth_common_checks, + } => { + let skip = if force_virtual_key_checks { + None + } else if public_route || additional_public_route { + Some(CommonSkip::PublicRoute) + } else if unauthenticated_passthrough { + Some(CommonSkip::UnauthenticatedPassthrough) + } else if !master_key_present && !jwt_enabled && !oauth_enabled && !oauth_proxy_enabled + { + Some(CommonSkip::NoAuthDevelopment) + } else if custom_auth && !custom_auth_common_checks { + Some(CommonSkip::CustomAuth) + } else { + None + }; + AdmissionPlan::Common { skip } + } + AdmissionFacts::Reservation { + skip_budget_checks, + disabled, + apply_user_budget_to_team_keys, + fail_closed, + billing_agent, + } => AdmissionPlan::Reserve { + reserve: !skip_budget_checks && !disabled, + apply_user_budget_to_team_keys, + fail_closed: fail_closed || billing_agent, + }, + AdmissionFacts::Ownership { + virtual_key, + custom_auth, + jwt, + oauth, + } => { + let candidate = if custom_auth { + Candidate::CustomAuth + } else if jwt || oauth { + Candidate::Validator + } else if virtual_key { + Candidate::VirtualKey + } else { + Candidate::Other + }; + AdmissionPlan::Observe { + owner: Owner::Python, + candidate, + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::normal(false, false, false, true, false)] + #[case::free(true, false, false, false, false)] + #[case::disabled(false, true, false, false, false)] + #[case::managed(false, false, true, true, true)] + fn reservation( + #[case] skip: bool, + #[case] disabled: bool, + #[case] billing: bool, + #[case] reserve: bool, + #[case] closed: bool, + ) { + assert_eq!( + plan(AdmissionFacts::Reservation { + skip_budget_checks: skip, + disabled, + billing_agent: billing, + fail_closed: false, + apply_user_budget_to_team_keys: true + }), + AdmissionPlan::Reserve { + reserve, + apply_user_budget_to_team_keys: true, + fail_closed: closed + } + ); + } + + #[rstest] + fn ownership_never_transfers_request() { + assert_eq!( + plan(AdmissionFacts::Ownership { + virtual_key: true, + custom_auth: false, + jwt: false, + oauth: false + }), + AdmissionPlan::Observe { + owner: Owner::Python, + candidate: Candidate::VirtualKey + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_access.rs b/litellm-rust/crates/request-auth-core/src/agent_access.rs new file mode 100644 index 00000000000..555014741fc --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_access.rs @@ -0,0 +1,269 @@ +use serde::{Deserialize, Serialize}; + +use crate::managed::Agent; + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub registry_managed: bool, + pub database_present: bool, + pub target_loaded: bool, + pub target: Option, + pub subject_present: bool, + pub key_hash_present: bool, + pub subject_managed: bool, + pub session: bool, + pub custom: bool, + pub authority_loaded: bool, + pub explicit_grants: Option>, + pub legacy_loaded: bool, + pub legacy_grants: Option>, + pub target_ids: Option>, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Decision { + ReadTarget, + ReadAuthority, + ReadExplicitGrants, + ReadLegacyGrants, + ReadTargetIds, + Allow, + Deny, +} + +pub fn decide(f: &Facts) -> Decision { + if f.registry_managed || f.database_present { + if !f.target_loaded { + return Decision::ReadTarget; + } + if f.target.is_none() && f.registry_managed { + return Decision::Deny; + } + if let Some(target) = f.target.as_ref().filter(|target| target.identity_managed) { + if !target.enabled + || !target + .identity + .as_ref() + .is_some_and(|identity| identity.active) + || !f.subject_present + { + return Decision::Deny; + } + if f.key_hash_present + && !f.subject_managed + && !f.session + && !f.custom + && !f.authority_loaded + { + return Decision::ReadAuthority; + } + return match f.explicit_grants.as_ref() { + None => Decision::ReadExplicitGrants, + Some(grants) if grants.contains(&target.agent_id) => Decision::Allow, + Some(_) => Decision::Deny, + }; + } + } + if !f.legacy_loaded { + return Decision::ReadLegacyGrants; + } + let Some(grants) = f.legacy_grants.as_ref() else { + return Decision::Allow; + }; + match f.target_ids.as_ref() { + None => Decision::ReadTargetIds, + Some(ids) if ids.iter().any(|id| grants.contains(id)) => Decision::Allow, + Some(_) => Decision::Deny, + } +} + +#[derive(Deserialize)] +pub struct Invocation { + pub registered: Option, + pub has_store: bool, + pub target_loaded: bool, + pub target: Option, + pub actor_managed: bool, + pub actor_id_present: bool, + pub permissions: Option, + pub assigned: bool, + pub fee_loaded: bool, + pub fee: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum InvocationAction { + Complete, + ReadTarget, + ReadPermissions, + Assign, + ReadFee, + Deny, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum InvocationDenial { + DatabaseRequired, + Removed, + CallerNotPermitted, + InvalidFee, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct InvocationPlan { + pub action: InvocationAction, + pub denial: Option, + pub use_target: bool, + pub assign_billing: bool, +} + +pub fn invocation(f: &Invocation) -> InvocationPlan { + let action = |action, denial| InvocationPlan { + action, + denial, + use_target: f.target.is_some(), + assign_billing: !f.actor_id_present + && f.target + .as_ref() + .or(f.registered.as_ref()) + .is_some_and(|agent| agent.identity_managed), + }; + let Some(registered) = f.registered.as_ref() else { + return action(InvocationAction::Complete, None); + }; + let registered_managed = registered.identity_managed || registered.identity.is_some(); + if !f.has_store && registered_managed { + return action( + InvocationAction::Deny, + Some(InvocationDenial::DatabaseRequired), + ); + } + if !f.target_loaded { + return action(InvocationAction::ReadTarget, None); + } + if f.target.is_none() && registered_managed { + return action(InvocationAction::Deny, Some(InvocationDenial::Removed)); + } + let effective = f.target.as_ref().unwrap_or(registered); + if !effective.identity_managed && !f.actor_managed { + return action(InvocationAction::Complete, None); + } + match f.permissions.as_ref().map(decide) { + None => return action(InvocationAction::ReadPermissions, None), + Some(Decision::Allow) => (), + Some(_) => { + return action( + InvocationAction::Deny, + Some(InvocationDenial::CallerNotPermitted), + ); + } + } + if !f.assigned { + return action(InvocationAction::Assign, None); + } + if !f.fee_loaded { + return action(InvocationAction::ReadFee, None); + } + if !f.fee.is_some_and(|fee| fee.is_finite() && fee >= 0.0) { + return action(InvocationAction::Deny, Some(InvocationDenial::InvalidFee)); + } + action(InvocationAction::Complete, None) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::free(Some(0.0), None)] + #[case::priced(Some(1.0), None)] + #[case::missing(None, Some(InvocationDenial::InvalidFee))] + #[case::negative(Some(-1.0), Some(InvocationDenial::InvalidFee))] + #[case::nan(Some(f64::NAN), Some(InvocationDenial::InvalidFee))] + #[case::infinity(Some(f64::INFINITY), Some(InvocationDenial::InvalidFee))] + fn invocation_fee_is_checked_after_assignment( + #[case] fee: Option, + #[case] denial: Option, + ) { + let facts = Invocation { + registered: Some(Agent { + agent_id: "agent".into(), + enabled: true, + identity_managed: true, + identity: None, + execution_mode: "both".into(), + }), + has_store: true, + target_loaded: true, + target: Some(Agent { + agent_id: "agent".into(), + enabled: true, + identity_managed: true, + identity: None, + execution_mode: "both".into(), + }), + actor_managed: false, + actor_id_present: false, + permissions: Some(Facts { + registry_managed: false, + database_present: false, + target_loaded: false, + target: None, + subject_present: true, + key_hash_present: false, + subject_managed: false, + session: false, + custom: false, + authority_loaded: false, + explicit_grants: None, + legacy_loaded: true, + legacy_grants: None, + target_ids: None, + }), + assigned: true, + fee_loaded: true, + fee, + }; + assert_eq!(invocation(&facts).denial, denial); + assert!(invocation(&facts).assign_billing); + assert_eq!( + invocation(&Invocation { + assigned: false, + ..facts + }) + .action, + InvocationAction::Assign + ); + } + + #[rstest] + #[case::no_database(false, false, false, Decision::ReadLegacyGrants)] + #[case::database(false, true, false, Decision::ReadTarget)] + #[case::managed_missing(true, true, true, Decision::Deny)] + #[case::ordinary_missing(false, true, true, Decision::ReadLegacyGrants)] + fn target_read_precedes_grants( + #[case] registry_managed: bool, + #[case] database_present: bool, + #[case] target_loaded: bool, + #[case] expected: Decision, + ) { + let facts = Facts { + registry_managed, + database_present, + target_loaded, + target: None, + subject_present: false, + key_hash_present: false, + subject_managed: false, + session: false, + custom: false, + authority_loaded: false, + explicit_grants: None, + legacy_loaded: false, + legacy_grants: None, + target_ids: None, + }; + assert_eq!(decide(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs b/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs new file mode 100644 index 00000000000..e270e291078 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs @@ -0,0 +1,151 @@ +use crate::Error; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Status { + Found, + Missing, + Unavailable, +} + +#[derive(Clone, Deserialize)] +pub struct Group { + pub status: Status, + pub models: Vec, + pub mcp_server_ids: Vec, + pub agent_ids: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub ids: Option>, + pub groups: Option>, + pub check_db_only: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadIds, + ReadGroups { + authoritative: bool, + }, + Complete { + present: bool, + ids: Vec, + models: BTreeSet, + mcp_server_ids: BTreeSet, + agent_ids: BTreeSet, + indices: Vec, + }, +} + +pub fn decide(facts: &Facts) -> Result { + let Some(ids) = &facts.ids else { + return Ok(Decision::ReadIds); + }; + let empty = |present| Decision::Complete { + present, + ids: vec![], + models: BTreeSet::new(), + mcp_server_ids: BTreeSet::new(), + agent_ids: BTreeSet::new(), + indices: vec![], + }; + if ids.is_empty() { + return Ok(empty(false)); + } + let Some(groups) = &facts.groups else { + return Ok(Decision::ReadGroups { + authoritative: facts.check_db_only, + }); + }; + if groups.len() != ids.len() { + return Err(Error::InvalidAnswer); + } + let indices: Vec<_> = groups + .iter() + .enumerate() + .filter_map(|(index, group)| (group.status == Status::Found).then_some(index)) + .collect(); + Ok(Decision::Complete { + present: true, + ids: ids + .iter() + .zip(groups) + .filter(|(_, group)| group.status != Status::Missing) + .map(|(id, _)| id.clone()) + .collect(), + models: indices + .iter() + .flat_map(|index| groups[*index].models.iter().cloned()) + .collect(), + mcp_server_ids: indices + .iter() + .flat_map(|index| groups[*index].mcp_server_ids.iter().cloned()) + .collect(), + agent_ids: indices + .iter() + .flat_map(|index| groups[*index].agent_ids.iter().cloned()) + .collect(), + indices, + }) +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum ReadStatus { + Success, + NotFound, + HttpError, +} + +#[derive(Clone, Deserialize)] +pub struct LoadFacts { + pub database_present: bool, + pub check_db_only: bool, + pub status: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum LoadDecision { + Read, + Return, + DatabaseUnavailable, + HttpUnavailable, + Raise, +} + +pub fn load(facts: &LoadFacts) -> LoadDecision { + if !facts.database_present { + return LoadDecision::DatabaseUnavailable; + } + match facts.status { + None => LoadDecision::Read, + Some(ReadStatus::Success) => LoadDecision::Return, + Some(ReadStatus::NotFound) => LoadDecision::Raise, + Some(ReadStatus::HttpError) if facts.check_db_only => LoadDecision::Raise, + Some(ReadStatus::HttpError) => LoadDecision::HttpUnavailable, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::soft(false, LoadDecision::HttpUnavailable)] + #[case::authoritative(true, LoadDecision::Raise)] + fn http_lookup_policy(#[case] check_db_only: bool, #[case] expected: LoadDecision) { + assert_eq!( + load(&LoadFacts { + database_present: true, + check_db_only, + status: Some(ReadStatus::HttpError) + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_grants.rs b/litellm-rust/crates/request-auth-core/src/agent_grants.rs new file mode 100644 index 00000000000..5956c7c8390 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_grants.rs @@ -0,0 +1,408 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Source { + pub team_id: Option, + pub grants: Option>, + pub unavailable: bool, +} + +#[derive(Deserialize)] +pub struct ConfiguredAgent { + pub agent_id: String, + pub groups: Vec, +} + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Resolve { + managed: bool, + strict: bool, + role: Option, + key_team_loaded: bool, + key_team: Option>, + default_deny: Option, + caller_loaded: bool, + own: Option>, + ceiling_loaded: bool, + ceiling: Option>, + }, + Failure { + strict: bool, + closed: bool, + }, + OwnerTeam { + strict: bool, + permission_id_present: bool, + permission_present: bool, + }, + OwnerUnion { + teams: Vec, + withheld: Vec, + unreadable: bool, + }, + BillingTeam { + sources: Vec, + agent_id: String, + }, + Human { + sources: Vec, + team_id: Option, + }, + Delegated { + sources: Vec, + agent_id: String, + team_id: Option, + explicit: bool, + allowed: Option>, + }, + TeamLookup { + subject: bool, + team_id: Option, + personal: Option, + database: bool, + loaded: bool, + row_present: bool, + strict: bool, + }, + GroupMatches { + agents: Vec, + groups: Vec, + }, +} + +#[derive(Default, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Action { + #[default] + Complete, + ReadManaged, + ReadKeyTeam, + ReadDefault, + ReadCaller, + ReadCeiling, + ReadPersonalEligibility, + ReadOwnerTeams, + ReadTeam, + ReadDeclared, + Unavailable, + SelectTeam, +} + +#[derive(Default, Debug, PartialEq, Eq, Serialize)] +pub struct Decision { + pub action: Action, + pub grants: Option>, + pub team_id: Option, +} + +fn action(action: Action) -> Decision { + Decision { + action, + ..Decision::default() + } +} +fn grants(grants: Option>) -> Decision { + Decision { + grants, + ..Decision::default() + } +} +fn intersection(left: Option>, right: Option>) -> Option> { + match (left, right) { + (None, right) => right, + (left, None) => left, + (Some(left), Some(right)) => { + Some(left.into_iter().filter(|id| right.contains(id)).collect()) + } + } +} + +pub fn decide(facts: Facts) -> Decision { + match facts { + Facts::Resolve { + managed, + strict, + role, + key_team_loaded, + key_team, + default_deny, + caller_loaded, + own, + ceiling_loaded, + ceiling, + } => { + if managed { + return action(Action::ReadManaged); + } + if !key_team_loaded { + return action(Action::ReadKeyTeam); + } + let admin = role.as_deref() == Some("proxy_admin"); + if !strict && !admin && default_deny.is_none() { + return action(Action::ReadDefault); + } + if (strict || (!admin && default_deny == Some(true))) && key_team.is_none() { + return grants(Some(vec![])); + } + if !caller_loaded { + return action(Action::ReadCaller); + } + if !ceiling_loaded { + return action(Action::ReadCeiling); + } + grants(intersection(own, ceiling)) + } + Facts::Failure { strict, closed } => { + if strict { + action(Action::Unavailable) + } else { + grants(closed.then(Vec::new)) + } + } + Facts::OwnerTeam { + strict, + permission_id_present, + permission_present, + } => { + if permission_id_present && !permission_present { + if strict { + action(Action::Unavailable) + } else { + grants(Some(vec![])) + } + } else { + action(Action::ReadDeclared) + } + } + Facts::OwnerUnion { + teams, + withheld, + unreadable, + } => { + if teams.iter().any(|source| source.unavailable) { + return action(Action::Unavailable); + } + let restricted = teams.is_empty() + || unreadable + || teams.iter().any(|source| source.grants.is_some()) + || withheld + .iter() + .any(|source| source.unavailable || source.grants.is_some()); + grants(restricted.then(|| { + teams + .into_iter() + .flat_map(|source| source.grants.unwrap_or_default()) + .collect::>() + .into_iter() + .collect() + })) + } + Facts::BillingTeam { sources, agent_id } => { + let selected = sources + .iter() + .filter(|source| { + source + .grants + .as_ref() + .is_some_and(|ids| ids.contains(&agent_id)) + }) + .filter_map(|source| source.team_id.clone()) + .min(); + let team_id = selected.or_else(|| { + sources + .iter() + .all(|source| source.grants.is_none() && !source.unavailable) + .then(|| { + sources + .iter() + .filter_map(|source| source.team_id.clone()) + .min() + }) + .flatten() + }); + Decision { + team_id, + ..Decision::default() + } + } + Facts::Human { sources, team_id } => grants(Some( + sources + .into_iter() + .filter(|source| source.team_id.is_none() || source.team_id == team_id) + .flat_map(|source| source.grants.unwrap_or_default()) + .collect::>() + .into_iter() + .collect(), + )), + Facts::Delegated { + sources, + agent_id, + team_id, + explicit, + allowed, + } => { + if sources.iter().any(|source| { + source.team_id.is_none() + && source + .grants + .as_ref() + .is_some_and(|ids| ids.contains(&agent_id)) + }) { + return Decision { + team_id, + ..Decision::default() + }; + } + let teams: BTreeSet<_> = sources + .iter() + .filter(|source| { + source + .grants + .as_ref() + .is_some_and(|ids| ids.contains(&agent_id)) + }) + .filter_map(|source| source.team_id.as_ref()) + .filter(|team| { + allowed + .as_ref() + .is_none_or(|allowed| allowed.contains(team)) + }) + .cloned() + .collect(); + if team_id.as_ref().is_some_and(|team| teams.contains(team)) { + return Decision { + team_id, + ..Decision::default() + }; + } + if !explicit && let Some(first) = teams.first() { + return Decision { + team_id: Some(first.clone()), + ..Decision::default() + }; + } + action(Action::SelectTeam) + } + Facts::TeamLookup { + subject, + team_id, + personal, + database, + loaded, + row_present, + strict, + } => { + if !subject { + return grants(None); + } + if team_id.is_none() { + return match personal { + None => action(Action::ReadPersonalEligibility), + Some(true) => action(Action::ReadOwnerTeams), + Some(false) => grants(None), + }; + } + if !database || (loaded && !row_present) { + return grants(strict.then(Vec::new)); + } + if !loaded { + action(Action::ReadTeam) + } else { + action(Action::ReadDeclared) + } + } + Facts::GroupMatches { agents, groups } => grants(Some( + agents + .into_iter() + .filter(|agent| agent.groups.iter().any(|id| groups.contains(id))) + .map(|agent| agent.agent_id) + .collect(), + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::legacy(false, false, Action::Complete, None)] + #[case::closed(false, true, Action::Complete, Some(vec![]))] + #[case::strict(true, false, Action::Unavailable, None)] + #[case::strict_closed(true, true, Action::Unavailable, None)] + fn lookup_failure( + #[case] strict: bool, + #[case] closed: bool, + #[case] expected: Action, + #[case] expected_grants: Option>, + ) { + assert_eq!( + decide(Facts::Failure { strict, closed }), + Decision { + action: expected, + grants: expected_grants, + team_id: None + } + ); + } + + #[rstest] + #[case::unrestricted(None, false, None)] + #[case::withheld(Some(vec!["a".into()]), false, Some(vec![]))] + #[case::unreadable(None, true, Some(vec![]))] + fn withheld_never_widens( + #[case] withheld: Option>, + #[case] unreadable: bool, + #[case] expected: Option>, + ) { + assert_eq!( + decide(Facts::OwnerUnion { + teams: vec![Source { + team_id: Some("team".into()), + grants: None, + unavailable: false + }], + withheld: vec![Source { + team_id: Some("held".into()), + grants: withheld, + unavailable: false + }], + unreadable + }) + .grants, + expected + ); + } + + #[rstest] + #[case::implicit(false, Some("a".into()), Action::Complete)] + #[case::explicit(true, None, Action::SelectTeam)] + fn delegated_requires_explicit_grant( + #[case] explicit: bool, + #[case] team_id: Option, + #[case] expected: Action, + ) { + assert_eq!( + decide(Facts::Delegated { + sources: vec![Source { + team_id: Some("a".into()), + grants: Some(vec!["agent".into()]), + unavailable: false + }], + agent_id: "agent".into(), + team_id: Some("other".into()), + explicit, + allowed: None + }), + Decision { + action: expected, + team_id, + grants: None + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/auth_context.rs b/litellm-rust/crates/request-auth-core/src/auth_context.rs new file mode 100644 index 00000000000..005a1ec57c1 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/auth_context.rs @@ -0,0 +1,169 @@ +use serde::{Deserialize, Serialize}; + +use crate::access::{self, AccessFacts, AccessRead}; + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Fetch { + team_id: Option, + ui_team_id: String, + user_id: Option, + project_id: Option, + end_user_id: Option, + }, + Team { + team_id: Option, + ui_team_id: String, + failed: bool, + missing: bool, + recorded_models: Vec, + allow_unavailable: Option, + }, + Admin { + role: Option, + user_id: Option, + admin_id: String, + user_row: bool, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum TeamSource { + None, + Token, + Loaded, + ReadUnavailablePolicy, + Failure, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Plan { + Fetch { + team: bool, + user: bool, + project: bool, + end_user: bool, + }, + Team { + source: TeamSource, + }, + Admin { + force_admin: bool, + preserve_budget: bool, + }, +} + +pub fn plan(facts: Facts) -> Plan { + match facts { + Facts::Fetch { + team_id, + ui_team_id, + user_id, + project_id, + end_user_id, + } => Plan::Fetch { + team: team_id.as_ref().is_some_and(|id| id != &ui_team_id), + user: user_id.is_some(), + project: project_id.is_some(), + end_user: end_user_id.is_some_and(|id| !id.is_empty()), + }, + Facts::Team { + team_id, + ui_team_id, + failed, + missing, + recorded_models, + allow_unavailable, + } => { + let dashboard = team_id.as_ref() == Some(&ui_team_id); + let source = if !failed { + if dashboard { + TeamSource::Token + } else { + TeamSource::Loaded + } + } else if team_id.is_none() { + TeamSource::None + } else { + let decision = access::decide(AccessFacts::TeamLookup { + dashboard, + missing, + recorded_models, + allow_unavailable, + }); + if decision.read == Some(AccessRead::UnavailablePolicy) { + TeamSource::ReadUnavailablePolicy + } else if decision.denial.is_some() { + TeamSource::Failure + } else { + TeamSource::Token + } + }; + Plan::Team { source } + } + Facts::Admin { + role, + user_id, + admin_id, + user_row, + } => Plan::Admin { + force_admin: role.as_deref() == Some("proxy_admin"), + preserve_budget: user_row && user_id.as_ref() != Some(&admin_id), + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, false)] + #[case::dashboard(Some("ui"), false)] + #[case::empty(Some(""), true)] + #[case::team(Some("team"), true)] + fn fetch_team(#[case] team: Option<&str>, #[case] expected: bool) { + assert_eq!( + plan(Facts::Fetch { + team_id: team.map(str::to_owned), + ui_team_id: "ui".into(), + user_id: Some("".into()), + project_id: Some("".into()), + end_user_id: Some("".into()) + }), + Plan::Fetch { + team: expected, + user: true, + project: true, + end_user: false + } + ); + } + + #[rstest] + #[case::master(Some("admin"), true, false)] + #[case::personal(Some("user"), true, true)] + #[case::no_row(Some("user"), false, false)] + #[case::missing_id(None, true, true)] + fn admin_budget_projection( + #[case] user: Option<&str>, + #[case] user_row: bool, + #[case] preserve_budget: bool, + ) { + assert_eq!( + plan(Facts::Admin { + role: Some("proxy_admin".into()), + user_id: user.map(str::to_owned), + admin_id: "admin".into(), + user_row + }), + Plan::Admin { + force_admin: true, + preserve_budget + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_alerts.rs b/litellm-rust/crates/request-auth-core/src/budget_alerts.rs new file mode 100644 index 00000000000..e6b6ad2a274 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_alerts.rs @@ -0,0 +1,233 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +pub struct Failure { + pub kind: String, + pub message: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ThresholdFacts { + pub key: String, + pub digit: bool, + pub digits: Option, + pub ratio: Option, + pub integer_failure: Option, + pub ratio_failure: Option, + pub text: Option, + pub items: Option>>, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ConfigFacts { + pub global_entries: Vec, + pub local_entries: Vec, + pub global_failure: Option, + pub local_failure: Option, + pub order: Option>, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub member: bool, + pub spend: Option, + pub limit: Option, + pub default_fraction: f64, + pub config: Option, +} + +#[derive(Clone, Debug, Serialize, PartialEq)] +pub struct Threshold { + pub key: String, + pub recipients: Vec, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct Decision { + pub read_config: bool, + pub read_order: bool, + pub global_keys: Vec, + pub local_keys: Vec, + pub emit: bool, + pub thresholds: Option>, + pub failure: Option, +} + +fn member_entry(entry: &ThresholdFacts) -> Result { + if !entry.digit || entry.key.chars().count() > 3 { + return Ok(false); + } + if let Some(error) = &entry.integer_failure { + return Err(error.clone()); + } + Ok(entry + .digits + .as_deref() + .and_then(|digits| digits.parse::().ok()) + .is_some_and(|n| (1..=100).contains(&n))) +} + +fn eligible(entries: &[ThresholdFacts], member: bool) -> Result, Failure> { + entries.iter().try_fold(Vec::new(), |mut selected, entry| { + if !member || member_entry(entry)? { + selected.push(entry); + } + Ok(selected) + }) +} + +fn recipients(entry: &ThresholdFacts) -> Vec { + let strings: Vec<&str> = if let Some(text) = &entry.text { + text.split(',').collect() + } else { + entry + .items + .as_deref() + .unwrap_or_default() + .iter() + .filter_map(Option::as_deref) + .collect() + }; + strings + .into_iter() + .map(|s| s.trim_matches(|c: char| c.is_whitespace() || matches!(c, '\u{1c}'..='\u{1f}'))) + .filter(|s| !s.is_empty()) + .map(str::to_owned) + .collect() +} + +fn merge(key: &str, global: &[&ThresholdFacts], local: &[&ThresholdFacts]) -> Threshold { + let recipients = global + .iter() + .chain(local) + .filter(|entry| entry.key == key) + .flat_map(|entry| recipients(entry)) + .fold(Vec::new(), |mut merged, email| { + if !merged.contains(&email) { + merged.push(email); + } + merged + }); + Threshold { + key: key.to_owned(), + recipients, + } +} + +fn evaluate(facts: &Facts) -> Result { + if !facts.member && !(facts.limit.is_some() && facts.spend.is_some_and(|spend| spend > 0.0)) { + return Ok(Decision::default()); + } + let Some(config) = &facts.config else { + return Ok(Decision { + read_config: true, + ..Decision::default() + }); + }; + if let Some(failure) = &config.global_failure { + return Err(failure.clone()); + } + if !facts.member + && let Some(failure) = &config.local_failure + { + return Err(failure.clone()); + } + let global = eligible(&config.global_entries, false)?; + let local = if facts.member && config.local_failure.is_some() { + vec![] + } else { + eligible(&config.local_entries, facts.member)? + }; + if global.is_empty() && local.is_empty() { + let emit = !facts.member + && matches!((facts.spend, facts.limit), (Some(spend), Some(limit)) if spend >= limit * facts.default_fraction && spend < limit); + return Ok(Decision { + emit, + ..Decision::default() + }); + } + let Some(order) = &config.order else { + return Ok(Decision { + read_order: true, + global_keys: global.iter().map(|entry| entry.key.clone()).collect(), + local_keys: local.iter().map(|entry| entry.key.clone()).collect(), + ..Decision::default() + }); + }; + let thresholds: Vec = order + .iter() + .map(|key| merge(key, &global, &local)) + .collect(); + if facts.member && facts.spend.is_some_and(|spend| spend <= 0.0) { + return Ok(Decision::default()); + } + let mut minimum: Option<&ThresholdFacts> = None; + for key in order { + let Some(entry) = global.iter().chain(&local).find(|entry| entry.key == *key) else { + continue; + }; + if !entry.digit { + continue; + } + if let Some(failure) = &entry.integer_failure { + return Err(failure.clone()); + } + let Some(digits) = entry.digits.as_ref() else { + continue; + }; + if minimum + .and_then(|entry| entry.digits.as_ref()) + .is_none_or(|current| (digits.len(), digits) < (current.len(), current)) + { + minimum = Some(entry); + } + } + let Some(minimum) = minimum else { + return Ok(Decision::default()); + }; + if let Some(failure) = &minimum.ratio_failure { + return Err(failure.clone()); + } + let emit = matches!((facts.spend, facts.limit, minimum.ratio), (Some(spend), Some(limit), Some(ratio)) if spend.partial_cmp(&(limit * ratio)) != Some(std::cmp::Ordering::Less)); + Ok(Decision { + emit, + thresholds: Some(thresholds), + ..Decision::default() + }) +} + +pub fn plan(facts: &Facts) -> Decision { + evaluate(facts).unwrap_or_else(|failure| Decision { + failure: Some(failure), + ..Decision::default() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::before(7.9, false)] + #[case::at(8.0, true)] + #[case::at_limit(10.0, false)] + #[case::nan(f64::NAN, false)] + #[case::zero(0.0, false)] + fn default_key_threshold(#[case] spend: f64, #[case] emit: bool) { + let facts = Facts { + member: false, + spend: Some(spend), + limit: Some(10.0), + default_fraction: 0.8, + config: Some(ConfigFacts { + global_entries: vec![], + local_entries: vec![], + global_failure: None, + local_failure: None, + order: None, + }), + }; + assert_eq!(plan(&facts).emit, emit); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_context.rs b/litellm-rust/crates/request-auth-core/src/budget_context.rs new file mode 100644 index 00000000000..5162edd0836 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_context.rs @@ -0,0 +1,484 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum ModelScope { + Team, + EndUser, + CustomEndUser, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelCheckFacts { + pub scope: ModelScope, + pub entity_id: Option, + pub budget_count: usize, + pub model_count: usize, +} + +pub fn model_checks(facts: &ModelCheckFacts) -> Vec { + if facts.budget_count == 0 + || facts.entity_id.is_none() + || (facts.scope == ModelScope::EndUser && facts.entity_id.as_deref() == Some("")) + { + return vec![]; + } + (0..facts.model_count).collect() +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct BudgetFields { + pub max_budget: Option, + pub tpm_limit: Option, + pub rpm_limit: Option, + pub tpd_limit: Option, + pub model_budget_count: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum EndUserMode { + Models, + Token, + Default, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct EndUserFacts { + pub mode: EndUserMode, + pub object_present: bool, + pub end_user_id: Option, + pub prisma_present: bool, + pub key_budget_id: Option, + pub key_read: bool, + pub key_budget: Option, + pub global_config_read: bool, + pub global_budget_id: Option, + pub global_read: bool, + pub global_budget: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum EndUserSource { + Own, + ReadKey, + Key, + ReadGlobalConfig, + ReadGlobal, + Global, + None, +} + +pub fn end_user_source(facts: &EndUserFacts) -> EndUserSource { + use EndUserSource::*; + if facts.mode != EndUserMode::Default && facts.object_present { + return Own; + } + if facts.mode == EndUserMode::Models + && (!facts.prisma_present || facts.end_user_id.as_deref().is_none_or(str::is_empty)) + { + return None; + } + if facts.key_budget_id.is_some() { + if !facts.key_read { + return ReadKey; + } + if facts.key_budget.is_some() { + return Key; + } + } + if !facts.global_config_read { + return ReadGlobalConfig; + } + if facts.global_budget_id.is_none() { + return None; + } + if !facts.global_read { + return ReadGlobal; + } + if facts.global_budget.is_some() { + Global + } else { + None + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TokenDefaultsFacts { + pub current: BudgetFields, + pub keep_token_limits: bool, + pub default_present: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum TokenField { + MaxBudget, + Tpm, + Rpm, + Tpd, + Models, +} + +pub fn token_assignments(facts: &TokenDefaultsFacts) -> Vec { + if !facts.default_present { + return vec![]; + } + [ + (TokenField::MaxBudget, facts.current.max_budget.is_none()), + (TokenField::Tpm, facts.current.tpm_limit.is_none()), + (TokenField::Rpm, facts.current.rpm_limit.is_none()), + (TokenField::Tpd, facts.current.tpd_limit.is_none()), + ( + TokenField::Models, + facts.current.model_budget_count.is_none(), + ), + ] + .into_iter() + .filter_map(|(field, missing)| (!facts.keep_token_limits || missing).then_some(field)) + .collect() +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MembershipSnapshot { + pub budget_id: Option, + pub model_budget_count: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, strum::VariantArray)] +pub enum LookupFailure { + Membership, + Default, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MemberFacts { + pub user_id: Option, + pub team_id: Option, + pub prisma_present: bool, + pub default_id: Option, + pub skip_budget_checks: bool, + pub membership_read: bool, + pub membership: Option, + pub default_read: bool, + pub default_budget: Option, + pub lookup_failure: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum MemberSource { + Done, + ReadMembership, + ReadDefault, + Member, + Default, + Clear, + RaiseLookup, + SkipLookup, +} + +pub fn member_source(facts: &MemberFacts) -> MemberSource { + use MemberSource::*; + if facts.user_id.is_none() || facts.team_id.is_none() || !facts.prisma_present { + return Done; + } + if facts.lookup_failure.is_some() { + return if facts.skip_budget_checks { + SkipLookup + } else { + RaiseLookup + }; + } + if !facts.membership_read { + return ReadMembership; + } + if facts.membership.as_ref().is_some_and(|member| { + member.budget_id.is_some() + && member.budget_id != facts.default_id + && member.model_budget_count.is_some_and(|count| count > 0) + }) { + return Member; + } + if facts.default_id.is_none() { + return Done; + } + if !facts.default_read { + return ReadDefault; + } + if facts + .default_budget + .as_ref() + .is_some_and(|budget| budget.model_budget_count.is_some()) + { + Default + } else { + Clear + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct IncreaseFacts { + pub number: Option, + pub truthy: bool, + pub type_name: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TemporaryFacts { + pub max_budget: Option, + pub metadata_read: bool, + pub has_increase: bool, + pub has_expiry: bool, + pub expires_at: Option, + pub now: Option, + pub increase: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum TemporaryRead { + Metadata, + Expiry, + Now, + Increase, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct TemporaryDecision { + pub read: Option, + pub copy: bool, + pub max_budget: Option, + pub invalid_increase_type: Option, +} + +fn temp_read(read: TemporaryRead) -> TemporaryDecision { + TemporaryDecision { + read: Some(read), + ..TemporaryDecision::default() + } +} + +pub fn temporary(facts: &TemporaryFacts) -> TemporaryDecision { + let Some(limit) = facts.max_budget else { + return TemporaryDecision::default(); + }; + if !facts.metadata_read { + return temp_read(TemporaryRead::Metadata); + } + if !facts.has_increase || !facts.has_expiry { + return TemporaryDecision::default(); + } + let Some(expiry) = facts.expires_at else { + return temp_read(TemporaryRead::Expiry); + }; + let Some(now) = facts.now else { + return temp_read(TemporaryRead::Now); + }; + if expiry <= now { + return TemporaryDecision::default(); + } + let Some(increase) = &facts.increase else { + return temp_read(TemporaryRead::Increase); + }; + if !increase.truthy { + return TemporaryDecision::default(); + } + match increase.number { + Some(amount) => TemporaryDecision { + copy: true, + max_budget: Some(limit + amount), + ..TemporaryDecision::default() + }, + None => TemporaryDecision { + invalid_increase_type: Some(increase.type_name.clone()), + ..TemporaryDecision::default() + }, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct NumericFacts { + pub number: Option, + pub boolean: bool, +} + +pub fn percentage(facts: &NumericFacts) -> Option { + if facts.boolean { + return None; + } + facts.number.filter(|pct| *pct > 0.0 && *pct <= 1.0) +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ThrottleFacts { + pub opt_in: bool, + pub tpm_limit: Option, + pub rpm_limit: Option, + pub percentage: Option, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct ThrottleDecision { + pub read_percentage: bool, + pub percentage: Option, +} + +pub fn throttle(facts: &ThrottleFacts) -> ThrottleDecision { + if !facts.opt_in || (facts.tpm_limit.is_none() && facts.rpm_limit.is_none()) { + return ThrottleDecision::default(); + } + match &facts.percentage { + None => ThrottleDecision { + read_percentage: true, + ..ThrottleDecision::default() + }, + Some(raw) => ThrottleDecision { + percentage: percentage(raw), + ..ThrottleDecision::default() + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::team_empty(ModelScope::Team, Some(""), vec![0, 1])] + #[case::team_none(ModelScope::Team, None, vec![])] + #[case::enduser_empty(ModelScope::EndUser, Some(""), vec![])] + #[case::enduser(ModelScope::EndUser, Some("user"), vec![0, 1])] + #[case::custom_empty(ModelScope::CustomEndUser, Some(""), vec![0, 1])] + #[case::custom_absent(ModelScope::CustomEndUser, None, vec![])] + fn model_identity_gates( + #[case] scope: ModelScope, + #[case] identity: Option<&str>, + #[case] expected: Vec, + ) { + assert_eq!( + model_checks(&ModelCheckFacts { + scope, + entity_id: identity.map(str::to_owned), + budget_count: 1, + model_count: 2 + }), + expected + ); + } + + #[rstest] + #[case::zero(0.0, false, None)] + #[case::negative(-0.5, false, None)] + #[case::nan(f64::NAN, false, None)] + #[case::infinite(f64::INFINITY, false, None)] + #[case::one(1.0, false, Some(1.0))] + #[case::boolean(1.0, true, None)] + fn numeric_percentages( + #[case] number: f64, + #[case] boolean: bool, + #[case] expected: Option, + ) { + assert_eq!( + percentage(&NumericFacts { + number: Some(number), + boolean + }), + expected + ); + } + + #[rstest] + #[case::expired(10, 10, 3.0, false, None)] + #[case::before(9, 10, 3.0, false, None)] + #[case::active(11, 10, 3.0, true, Some(13.0))] + #[case::negative(11, 10, -3.0, true, Some(7.0))] + #[case::zero(11, 10, 0.0, false, None)] + fn temporary_copy_boundary( + #[case] expiry: i64, + #[case] now: i64, + #[case] amount: f64, + #[case] copy: bool, + #[case] expected: Option, + ) { + let facts = TemporaryFacts { + max_budget: Some(10.0), + metadata_read: true, + has_increase: true, + has_expiry: true, + expires_at: Some(expiry), + now: Some(now), + increase: Some(IncreaseFacts { + number: Some(amount), + truthy: amount != 0.0, + type_name: "float".to_owned(), + }), + }; + let result = temporary(&facts); + assert_eq!( + (result.copy, result.max_budget, result.read), + (copy, expected, None) + ); + } + + #[rstest] + #[case::member_override(Some("member"), Some(1), Some("default"), false, MemberSource::Member)] + #[case::no_override(None, Some(1), None, false, MemberSource::Done)] + #[case::empty_override( + Some("member"), + Some(0), + Some("default"), + false, + MemberSource::ReadDefault + )] + #[case::current_default( + Some("default"), + Some(1), + Some("default"), + false, + MemberSource::ReadDefault + )] + #[case::missing_default(None, None, Some("default"), true, MemberSource::Clear)] + fn member_default_staging( + #[case] member_id: Option<&str>, + #[case] count: Option, + #[case] default_id: Option<&str>, + #[case] default_read: bool, + #[case] expected: MemberSource, + ) { + let facts = MemberFacts { + user_id: Some(String::new()), + team_id: Some(String::new()), + prisma_present: true, + default_id: default_id.map(str::to_owned), + skip_budget_checks: false, + membership_read: true, + membership: Some(MembershipSnapshot { + budget_id: member_id.map(str::to_owned), + model_budget_count: count, + }), + default_read, + default_budget: None, + lookup_failure: None, + }; + assert_eq!(member_source(&facts), expected); + } + + #[rstest] + #[case::replace(false, true, vec![TokenField::MaxBudget, TokenField::Tpm, TokenField::Rpm, TokenField::Tpd, TokenField::Models])] + #[case::keep_empty(true, true, vec![TokenField::Rpm, TokenField::Tpd])] + #[case::no_default(false, false, vec![])] + fn token_assignment_order( + #[case] keep: bool, + #[case] present: bool, + #[case] expected: Vec, + ) { + let facts = TokenDefaultsFacts { + current: BudgetFields { + max_budget: Some(0.0), + tpm_limit: Some(0), + model_budget_count: Some(0), + ..BudgetFields::default() + }, + keep_token_limits: keep, + default_present: present, + }; + assert_eq!(token_assignments(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_lookup.rs b/litellm-rust/crates/request-auth-core/src/budget_lookup.rs new file mode 100644 index 00000000000..c5c150d2bdb --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_lookup.rs @@ -0,0 +1,284 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct OrganizationFacts { + pub token_present: bool, + pub client_present: bool, + pub team_present: bool, + pub token_id: Option, + pub team_id: Option, + pub row_present: Option, + pub lookup_failed: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum OrganizationAction { + Skip, + ReadTokenId, + ReadTeamId, + ReadOrganization, + UseOrganization, +} + +pub fn organization(facts: &OrganizationFacts) -> OrganizationAction { + use OrganizationAction::*; + if !facts.token_present || !facts.client_present { + return Skip; + } + match facts.token_id { + None => return ReadTokenId, + Some(false) if !facts.team_present => return Skip, + Some(false) => match facts.team_id { + None => return ReadTeamId, + Some(false) => return Skip, + Some(true) => {} + }, + Some(true) => {} + } + if facts.lookup_failed { + return Skip; + } + match facts.row_present { + None => ReadOrganization, + Some(false) => Skip, + Some(true) => UseOrganization, + } +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct MemberFacts { + pub team_present: bool, + pub token_present: bool, + pub team_id: Option, + pub user_id: Option, + pub membership_loaded: bool, + pub limits_read: bool, + pub maximum_present: bool, + pub default_id_string: bool, + pub default_read: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum MemberAction { + Skip, + ReadTeamId, + ReadUserId, + ReadMembership, + ReadLimits, + ReadDefault, + UseLimits, +} + +pub fn member(facts: &MemberFacts) -> MemberAction { + use MemberAction::*; + if !facts.team_present { + return Skip; + } + match facts.team_id { + None => return ReadTeamId, + Some(false) => return Skip, + Some(true) => {} + } + if !facts.token_present { + return Skip; + } + match facts.user_id { + None => return ReadUserId, + Some(false) => return Skip, + Some(true) => {} + } + if !facts.membership_loaded { + return ReadMembership; + } + if !facts.limits_read { + return ReadLimits; + } + if !facts.maximum_present && facts.default_id_string && !facts.default_read { + return ReadDefault; + } + UseLimits +} + +#[derive(Clone, Debug, Deserialize)] +pub enum EmailFacts { + Text(String), + List(Vec>), + Other, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct SoftFacts { + pub limit: Option, + pub spend: Option, + pub token_present: bool, +} + +pub fn read_soft_emails(facts: &SoftFacts) -> bool { + facts.token_present + && facts + .limit + .zip(facts.spend) + .is_some_and(|(limit, spend)| spend >= limit) +} + +fn stripped(value: &str) -> &str { + value.trim_matches(|character: char| { + character.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&character) + }) +} + +pub fn emails(facts: &EmailFacts) -> Vec { + match facts { + EmailFacts::Text(value) => value + .split(',') + .map(stripped) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .collect(), + EmailFacts::List(values) => values + .iter() + .flatten() + .filter(|value| !stripped(value).is_empty()) + .cloned() + .collect(), + EmailFacts::Other => Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::below(Some(1.0), Some(0.0), true, false)] + #[case::at_limit(Some(1.0), Some(1.0), true, true)] + #[case::no_token(Some(1.0), Some(2.0), false, false)] + #[case::no_limit(None, Some(2.0), true, false)] + #[case::nan(Some(f64::NAN), Some(2.0), true, false)] + fn lazy_soft_metadata( + #[case] limit: Option, + #[case] spend: Option, + #[case] token_present: bool, + #[case] expected: bool, + ) { + assert_eq!( + read_soft_emails(&SoftFacts { + limit, + spend, + token_present + }), + expected + ); + } + + #[rstest] + fn failed_organization_lookup_is_skipped() { + assert_eq!( + organization(&OrganizationFacts { + token_present: true, + client_present: true, + token_id: Some(true), + lookup_failed: true, + ..Default::default() + }), + OrganizationAction::Skip + ); + } + + #[rstest] + #[case::missing_token(false, true, None, None, OrganizationAction::Skip)] + #[case::missing_client(true, false, None, None, OrganizationAction::Skip)] + #[case::token_first(true, true, None, None, OrganizationAction::ReadTokenId)] + #[case::token_selected(true, true, Some(true), None, OrganizationAction::ReadOrganization)] + #[case::missing_row(true, true, Some(true), Some(false), OrganizationAction::Skip)] + #[case::loaded_row( + true, + true, + Some(true), + Some(true), + OrganizationAction::UseOrganization + )] + fn organization_stages( + #[case] token_present: bool, + #[case] client_present: bool, + #[case] token_id: Option, + #[case] row_present: Option, + #[case] expected: OrganizationAction, + ) { + assert_eq!( + organization(&OrganizationFacts { + token_present, + client_present, + token_id, + row_present, + ..Default::default() + }), + expected + ); + } + + #[rstest] + #[case::no_team(false, None, OrganizationAction::Skip)] + #[case::read_team(true, None, OrganizationAction::ReadTeamId)] + #[case::no_team_id(true, Some(false), OrganizationAction::Skip)] + #[case::team_selected(true, Some(true), OrganizationAction::ReadOrganization)] + fn organization_fallback( + #[case] team_present: bool, + #[case] team_id: Option, + #[case] expected: OrganizationAction, + ) { + assert_eq!( + organization(&OrganizationFacts { + token_present: true, + client_present: true, + token_id: Some(false), + team_present, + team_id, + row_present: None, + lookup_failed: false + }), + expected + ); + } + + #[rstest] + #[case::load_membership(false, false, false, false, false, MemberAction::ReadMembership)] + #[case::read_limits(true, false, false, false, false, MemberAction::ReadLimits)] + #[case::own_maximum(true, true, true, true, false, MemberAction::UseLimits)] + #[case::default(true, true, false, true, false, MemberAction::ReadDefault)] + #[case::invalid_default(true, true, false, false, false, MemberAction::UseLimits)] + #[case::loaded_default(true, true, false, true, true, MemberAction::UseLimits)] + fn member_stages( + #[case] membership_loaded: bool, + #[case] limits_read: bool, + #[case] maximum_present: bool, + #[case] default_id_string: bool, + #[case] default_read: bool, + #[case] expected: MemberAction, + ) { + assert_eq!( + member(&MemberFacts { + team_present: true, + token_present: true, + team_id: Some(true), + user_id: Some(true), + membership_loaded, + limits_read, + maximum_present, + default_id_string, + default_read + }), + expected + ); + } + + #[rstest] + #[case::comma_list(EmailFacts::Text(" a, ,b,a ".into()), vec!["a", "b", "a"])] + #[case::list_keeps_space(EmailFacts::List(vec![Some(" a ".into()), None, Some("\u{1c}".into())]), vec![" a "])] + #[case::python_whitespace(EmailFacts::Text("\u{1c}a\u{1f},\u{85}b\u{a0}".into()), vec!["a", "b"])] + #[case::other(EmailFacts::Other, vec![])] + fn email_normalization(#[case] facts: EmailFacts, #[case] expected: Vec<&str>) { + assert_eq!(emails(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_pricing.rs b/litellm-rust/crates/request-auth-core/src/budget_pricing.rs new file mode 100644 index 00000000000..29329d5c0ea --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_pricing.rs @@ -0,0 +1,261 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize)] +pub struct DeploymentPrice { + pub has_input_price: bool, + pub has_output_price: bool, + pub input_price: Option, + pub output_price: Option, + pub query_price: Option, + pub ptu_count: Option, + pub cost_per_ptu_per_hour: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelPriceFacts { + pub cached: Option, + pub lookup_failed: bool, + pub group_found: bool, + pub input_price: Option, + pub output_price: Option, + pub deployments: Option>, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct PriceDecision { + pub zero_cost: bool, + pub cache_value: Option, + pub read_deployments: bool, +} + +pub fn price(facts: &ModelPriceFacts) -> PriceDecision { + if let Some(cached) = facts.cached { + return PriceDecision { + zero_cost: cached, + cache_value: None, + read_deployments: false, + }; + } + if facts.lookup_failed { + return PriceDecision { + zero_cost: false, + cache_value: None, + read_deployments: false, + }; + } + let group_zero = facts.group_found + && matches!((facts.input_price, facts.output_price), (Some(input), Some(output)) if !(input > 0.0 || output > 0.0)); + if group_zero && facts.deployments.is_none() { + return PriceDecision { + zero_cost: false, + cache_value: None, + read_deployments: true, + }; + } + let deployments = facts.deployments.as_deref().unwrap_or_default(); + let zero_cost = group_zero + && deployments + .iter() + .any(|d| d.has_input_price || d.has_output_price) + && !deployments.iter().any(|d| { + d.input_price.is_some_and(|p| p > 0.0) + || d.output_price.is_some_and(|p| p > 0.0) + || d.query_price.is_some_and(|p| p > 0.0) + || d.ptu_count.is_some() && d.cost_per_ptu_per_hour.is_some() + }); + PriceDecision { + zero_cost, + cache_value: Some(zero_cost), + read_deployments: false, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct RouteBudgetFacts { + pub side_effect_route: bool, + pub discovery_route: bool, + pub llm_route: bool, + pub registered_pass_through: Option, + pub route: String, + pub method: Option, + pub body_present: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum RouteBudgetDecision { + ReadPassThrough, + ReadPricing, + Skip, +} + +pub fn route(facts: &RouteBudgetFacts) -> RouteBudgetDecision { + use RouteBudgetDecision::*; + if !facts.side_effect_route { + if facts.discovery_route { + return Skip; + } + if !facts.llm_route { + match facts.registered_pass_through { + None => return ReadPassThrough, + Some(false) => return Skip, + Some(true) => (), + } + } + } + let mcp_discovery = facts.body_present + && (matches!( + facts.route.as_str(), + "/mcp-rest/tools/list" | "/v1/mcp/tools" + ) || !matches!( + facts.route.as_str(), + "/mcp/tools/call" | "/mcp-rest/tools/call" + ) && (facts.route == "/mcp" || facts.route.starts_with("/mcp/")) + && matches!( + facts.method.as_deref(), + Some("initialize" | "notifications/initialized" | "ping" | "tools/list") + )); + if mcp_discovery { Skip } else { ReadPricing } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ExemptionFacts { + pub route: RouteBudgetFacts, + pub model_input_present: bool, + pub router_present: bool, + pub model_count: usize, + #[serde(default)] + pub cache_read: bool, + pub prices: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct ExemptionDecision { + pub read_pass_through: bool, + pub read_price_cache: bool, + pub read_model: Option, + pub exempt: bool, +} + +pub fn exemption(facts: &ExemptionFacts) -> ExemptionDecision { + let done = |exempt| ExemptionDecision { + read_pass_through: false, + read_price_cache: false, + read_model: None, + exempt, + }; + match route(&facts.route) { + RouteBudgetDecision::ReadPassThrough => { + return ExemptionDecision { + read_pass_through: true, + ..done(false) + }; + } + RouteBudgetDecision::Skip => return done(true), + RouteBudgetDecision::ReadPricing => (), + } + if !facts.model_input_present || !facts.router_present { + return done(false); + } + if !facts.cache_read { + return ExemptionDecision { + read_price_cache: true, + ..done(false) + }; + } + if facts.prices.iter().any(|facts| !price(facts).zero_cost) { + return done(false); + } + if facts.prices.len() < facts.model_count { + return ExemptionDecision { + read_model: Some(facts.prices.len()), + ..done(false) + }; + } + done(true) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::zero(Some(0.0), true)] + #[case::negative(Some(-1.0), true)] + #[case::nan(Some(f64::NAN), true)] + #[case::missing(None, false)] + #[case::positive(Some(1.0), false)] + fn preserves_group_comparison(#[case] input: Option, #[case] zero_cost: bool) { + let facts = ModelPriceFacts { + cached: None, + lookup_failed: false, + group_found: true, + input_price: input, + output_price: Some(0.0), + deployments: Some(vec![DeploymentPrice { + has_input_price: true, + has_output_price: false, + input_price: None, + output_price: None, + query_price: None, + ptu_count: None, + cost_per_ptu_per_hour: None, + }]), + }; + assert_eq!( + price(&facts), + PriceDecision { + zero_cost, + cache_value: Some(zero_cost), + read_deployments: false, + } + ); + } + #[rstest] + #[case::read_first(vec![], Some(0), false)] + #[case::read_second(vec![true], Some(1), false)] + #[case::priced_first(vec![false], None, false)] + #[case::all_free(vec![true, true], None, true)] + #[case::priced_second(vec![true, false], None, false)] + fn exemption_reads_models_in_order( + #[case] cached: Vec, + #[case] read: Option, + #[case] exempt: bool, + ) { + let facts = ExemptionFacts { + route: RouteBudgetFacts { + side_effect_route: false, + discovery_route: false, + llm_route: true, + route: "/chat/completions".to_owned(), + method: None, + body_present: false, + registered_pass_through: None, + }, + model_input_present: true, + router_present: true, + model_count: 2, + cache_read: true, + prices: cached + .into_iter() + .map(|value| ModelPriceFacts { + cached: Some(value), + lookup_failed: false, + group_found: false, + input_price: None, + output_price: None, + deployments: None, + }) + .collect(), + }; + assert_eq!( + exemption(&facts), + ExemptionDecision { + read_pass_through: false, + read_price_cache: false, + read_model: read, + exempt + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/common_budgets.rs b/litellm-rust/crates/request-auth-core/src/common_budgets.rs new file mode 100644 index 00000000000..9f3463fe5c1 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/common_budgets.rs @@ -0,0 +1,230 @@ +use std::cmp::Ordering::{Equal, Greater, Less}; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Scope { + Team, + Key, + TeamWindow, + KeyWindow, + Project, + Organization, + Tag, + AccessGroup, + User, + TeamMember, + EndUser, + Proxy, + KeySoft, + TeamSoft, + ProjectSoft, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct BudgetFacts { + pub scope: Scope, + pub limit: Option, + pub spend: Option, + pub token_present: bool, + pub team_key: bool, + pub apply_user_budget_to_team_keys: bool, + pub info_route: bool, + pub llm_route: bool, + pub model_list_route: bool, + pub skip: bool, + pub alert_recipients: Vec, + pub throttle_percentage: Option, + pub throttle_opt_in: bool, + pub tpm_limit: Option, + pub rpm_limit: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize)] +pub struct BudgetDecision { + pub read_spend: bool, + pub denied: bool, + pub alert: Option<&'static str>, + pub throttle_percentage: Option, +} + +pub fn check(facts: &BudgetFacts) -> BudgetDecision { + use Scope::*; + let Some(limit) = facts.limit else { + return BudgetDecision::default(); + }; + if matches!(facts.scope, User) && facts.team_key && !facts.apply_user_budget_to_team_keys + || matches!(facts.scope, EndUser) && facts.info_route + || matches!(facts.scope, Project) && !limit.is_finite() + || matches!(facts.scope, AccessGroup) && limit <= 0.0 + || matches!(facts.scope, Proxy) + && (limit.partial_cmp(&0.0) != Some(Greater) + || facts.skip + || !facts.llm_route + || facts.model_list_route) + { + return BudgetDecision::default(); + } + let soft = matches!(facts.scope, KeySoft | TeamSoft | ProjectSoft); + let Some(spend) = facts.spend else { + return BudgetDecision { + read_spend: !soft && facts.scope != Proxy, + ..BudgetDecision::default() + }; + }; + if soft { + let alert = spend >= limit + && match facts.scope { + KeySoft => limit != 0.0, + TeamSoft => facts.token_present && !facts.alert_recipients.is_empty(), + ProjectSoft => facts.token_present, + _ => false, + }; + return BudgetDecision { + alert: alert.then_some("soft_budget"), + ..BudgetDecision::default() + }; + } + let exceeded = match facts.scope { + Team | Proxy => limit.is_finite() && spend > limit, + EndUser => spend > limit, + Tag => !matches!(spend.partial_cmp(&limit), Some(Less | Equal)), + AccessGroup => spend.partial_cmp(&limit) != Some(Less), + _ => limit.is_finite() && spend >= limit, + }; + let alert = match facts.scope { + Key => Some("token_budget"), + User => Some("user_budget"), + Team if exceeded && facts.token_present => Some("team_budget"), + Project if exceeded && facts.token_present => Some("project_budget"), + Organization if exceeded => Some("organization_budget"), + TeamMember if limit.is_finite() => Some("team_member_threshold"), + _ => None, + }; + let throttle_percentage = if exceeded + && facts.scope == Key + && facts.throttle_opt_in + && (facts.tpm_limit.is_some() || facts.rpm_limit.is_some()) + { + facts + .throttle_percentage + .filter(|pct| *pct > 0.0 && *pct <= 1.0) + } else { + None + }; + BudgetDecision { + read_spend: false, + denied: exceeded && throttle_percentage.is_none(), + alert, + throttle_percentage, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct AlertFacts { + pub spend: Option, + pub max_budget: Option, + pub configured: bool, + pub minimum_percent: Option, + pub default_ratio: f64, + pub team_member: bool, +} + +pub fn alert(facts: &AlertFacts) -> bool { + let (Some(spend), Some(limit)) = (facts.spend, facts.max_budget) else { + return false; + }; + if facts.team_member { + if !facts.configured || spend <= 0.0 { + return false; + } + } else if spend.partial_cmp(&0.0) != Some(Greater) { + return false; + } + if facts.configured { + return facts + .minimum_percent + .is_some_and(|pct| spend.partial_cmp(&(limit * (pct / 100.0))) != Some(Less)); + } + spend >= limit * facts.default_ratio && spend < limit +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MemberLimitFacts { + pub member_limit: Option, + pub default_limit: Option, + pub temporary_increase: Option, + pub expires_at: Option, + pub now: f64, + #[serde(default)] + pub now_naive: bool, +} + +pub fn member_limit(facts: &MemberLimitFacts) -> Result, crate::Error> { + let base = facts.member_limit.or_else(|| { + facts + .default_limit + .filter(|limit| !matches!(limit.partial_cmp(&0.0), Some(Less | Equal))) + }); + let Some(base) = base else { + return Ok(None); + }; + let increase = match (facts.temporary_increase, facts.expires_at) { + (Some(_), Some(_)) if facts.now_naive => return Err(crate::Error::IncompatibleDateTimes), + (Some(increase), Some(expiry)) if expiry > facts.now => increase, + _ => 0.0, + }; + Ok(Some(base + increase)) +} + +#[derive(Clone, Debug, Deserialize)] +pub struct WindowFacts { + pub duration_seconds: Option, + pub reset_at: Option, + pub now: f64, +} + +pub fn window_start(facts: &WindowFacts) -> Option { + facts + .duration_seconds + .map(|seconds| facts.reset_at.unwrap_or(facts.now) - seconds) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::team_equality(Scope::Team, 1.0, 1.0, false)] + #[case::key_equality(Scope::Key, 1.0, 1.0, true)] + #[case::tag_nan(Scope::Tag, f64::NAN, 1.0, true)] + #[case::group_nan_spend(Scope::AccessGroup, 1.0, f64::NAN, true)] + #[case::key_nan(Scope::Key, f64::NAN, 1.0, false)] + #[case::end_user_infinity(Scope::EndUser, 1.0, f64::INFINITY, true)] + fn boundary( + #[case] scope: Scope, + #[case] limit: f64, + #[case] spend: f64, + #[case] denied: bool, + ) { + let facts = BudgetFacts { + scope, + limit: Some(limit), + spend: Some(spend), + token_present: false, + team_key: false, + apply_user_budget_to_team_keys: false, + info_route: false, + llm_route: true, + model_list_route: false, + skip: false, + alert_recipients: vec![], + throttle_percentage: None, + throttle_opt_in: false, + tpm_limit: None, + rpm_limit: None, + }; + assert_eq!(check(&facts).denied, denied); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/common_flow.rs b/litellm-rust/crates/request-auth-core/src/common_flow.rs new file mode 100644 index 00000000000..2f6a4e9873d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/common_flow.rs @@ -0,0 +1,358 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, strum::VariantArray)] +pub enum ModelBudgetScope { + Key, + User, + Member, +} + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum ReadFacts { + ModelBudgets { + scope: ModelBudgetScope, + entity_present: Option, + budget_present: Option, + mapping: Option, + team_present: Option, + }, + FallbackAllowed { + router: bool, + disabled: Option, + }, + BudgetSkip { + skip: bool, + exempt: Option, + }, + Pricing { + enabled: bool, + llm_route: Option, + }, + Unpriced { + count: usize, + }, + MemberFallback { + router: bool, + multipart: Option, + resource_bound: Option, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum ReadPlan { + CheckModelBudgets, + ReadModelBudgetPresence, + ReadModelBudgetShape, + ReadModelBudgetToken, + ReadModelBudgetUser, + ReadModelBudgetTeam, + AllowFallback, + ReadFallbackDisabled, + ReadExemption, + SkipBudgets, + CheckBudgets, + ReadRoute, + ReadUnpriced, + RejectUnpriced, + ReadContentType, + ReadResourceBinding, + ReadFallbackSettings, + Continue, +} + +pub fn read(facts: ReadFacts) -> ReadPlan { + match facts { + ReadFacts::ModelBudgets { + scope, + entity_present, + budget_present, + mapping, + team_present, + } => model_budgets(scope, entity_present, budget_present, mapping, team_present), + ReadFacts::FallbackAllowed { + router: true, + disabled: Some(false), + } => ReadPlan::AllowFallback, + ReadFacts::FallbackAllowed { + router: true, + disabled: None, + } => ReadPlan::ReadFallbackDisabled, + ReadFacts::FallbackAllowed { .. } => ReadPlan::Continue, + ReadFacts::BudgetSkip { skip: true, .. } + | ReadFacts::BudgetSkip { + exempt: Some(true), .. + } => ReadPlan::SkipBudgets, + ReadFacts::BudgetSkip { + exempt: Some(false), + .. + } => ReadPlan::CheckBudgets, + ReadFacts::BudgetSkip { exempt: None, .. } => ReadPlan::ReadExemption, + ReadFacts::Pricing { enabled: false, .. } + | ReadFacts::Pricing { + llm_route: Some(false), + .. + } => ReadPlan::Continue, + ReadFacts::Pricing { + llm_route: None, .. + } => ReadPlan::ReadRoute, + ReadFacts::Pricing { + llm_route: Some(true), + .. + } => ReadPlan::ReadUnpriced, + ReadFacts::Unpriced { count: 0 } => ReadPlan::Continue, + ReadFacts::Unpriced { .. } => ReadPlan::RejectUnpriced, + ReadFacts::MemberFallback { router: false, .. } + | ReadFacts::MemberFallback { + multipart: Some(true), + .. + } + | ReadFacts::MemberFallback { + resource_bound: Some(true), + .. + } => ReadPlan::Continue, + ReadFacts::MemberFallback { + multipart: None, .. + } => ReadPlan::ReadContentType, + ReadFacts::MemberFallback { + resource_bound: None, + .. + } => ReadPlan::ReadResourceBinding, + ReadFacts::MemberFallback { .. } => ReadPlan::ReadFallbackSettings, + } +} + +fn model_budgets( + scope: ModelBudgetScope, + entity: Option, + budget: Option, + mapping: Option, + team: Option, +) -> ReadPlan { + use ReadPlan as P; + let stages = match scope { + ModelBudgetScope::Key => vec![ + (budget, P::ReadModelBudgetPresence), + (entity, P::ReadModelBudgetToken), + ], + ModelBudgetScope::User => vec![ + (entity, P::ReadModelBudgetUser), + (mapping, P::ReadModelBudgetShape), + (budget, P::ReadModelBudgetPresence), + ], + ModelBudgetScope::Member => vec![ + (entity, P::ReadModelBudgetUser), + (team, P::ReadModelBudgetTeam), + (mapping, P::ReadModelBudgetShape), + (budget, P::ReadModelBudgetPresence), + ], + }; + for (present, missing) in stages { + match present { + None => return missing, + Some(false) => return P::Continue, + Some(true) => {} + } + } + P::CheckModelBudgets +} + +#[derive(Deserialize)] +pub struct Facts { + pub model_present: bool, + pub skip_budgets: bool, + pub team_present: bool, + pub user_present: bool, + pub token_present: bool, + pub user_id: Option, + pub agent_id: Option, + pub customer_present: bool, + pub customer_models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Plan { + pub membership: bool, + pub team_models: bool, + pub member_models: bool, + pub user_models: bool, + pub customer_models: bool, + pub agent_trace: bool, +} + +pub fn plan(facts: Facts) -> Plan { + Plan { + membership: facts.token_present + && (facts.model_present || !facts.skip_budgets) + && facts.team_present + && facts.user_id.is_some(), + team_models: facts.model_present && facts.team_present, + member_models: facts.model_present + && facts.team_present + && facts.token_present + && facts.user_id.is_some_and(|id| !id.is_empty()), + user_models: facts.model_present && !facts.team_present && facts.user_present, + customer_models: facts.customer_present && !facts.customer_models.is_empty(), + agent_trace: facts.token_present && facts.agent_id.is_some_and(|id| !id.is_empty()), + } +} + +#[derive(Deserialize)] +pub struct BudgetFacts { + pub token_present: bool, + pub group_count: usize, + pub customer_present: bool, + pub customer_budget_present: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum BudgetScope { + Team, + TeamWindows, + KeyWindows, + TeamSoft, + Organization, + Tags, + Groups, + User, + Member, + Customer, +} + +pub fn budgets(facts: BudgetFacts) -> Vec { + use BudgetScope as S; + [ + Some(S::Team), + Some(S::TeamWindows), + facts.token_present.then_some(S::KeyWindows), + Some(S::TeamSoft), + Some(S::Organization), + Some(S::Tags), + (facts.group_count > 0).then_some(S::Groups), + Some(S::User), + Some(S::Member), + (facts.customer_present && facts.customer_budget_present).then_some(S::Customer), + ] + .into_iter() + .flatten() + .collect() +} + +#[derive(Deserialize)] +pub struct CustomFacts { + pub end_user_id: Option, + pub end_user_present: bool, + pub common_checks: bool, + pub model_present: bool, + pub router_present: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct CustomPlan { + pub read_end_user: bool, + pub check_end_user: bool, + pub check_model_access: bool, + pub read_pricing: bool, +} + +pub fn custom(facts: CustomFacts) -> CustomPlan { + CustomPlan { + read_end_user: facts.end_user_id.is_some(), + check_end_user: facts.end_user_present && !facts.common_checks, + check_model_access: facts.common_checks, + read_pricing: facts.model_present && facts.router_present, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::available(true, true, true, true, ReadPlan::CheckModelBudgets)] + #[case::missing_identity(false, true, true, true, ReadPlan::Continue)] + #[case::empty_budget(true, false, true, true, ReadPlan::Continue)] + #[case::malformed_member_budget(true, true, false, true, ReadPlan::Continue)] + #[case::no_team(true, true, true, false, ReadPlan::Continue)] + fn model_budget_availability( + #[case] entity_present: bool, + #[case] budget_present: bool, + #[case] mapping: bool, + #[case] team_present: bool, + #[case] expected: ReadPlan, + ) { + assert_eq!( + read(ReadFacts::ModelBudgets { + scope: ModelBudgetScope::Member, + entity_present: Some(entity_present), + budget_present: Some(budget_present), + mapping: Some(mapping), + team_present: Some(team_present) + }), + expected + ); + } + + #[rstest] + #[case::explicit_skip(ReadFacts::BudgetSkip {skip:true, exempt:None}, ReadPlan::SkipBudgets)] + #[case::exempt(ReadFacts::BudgetSkip {skip:false, exempt:Some(true)}, ReadPlan::SkipBudgets)] + #[case::check(ReadFacts::BudgetSkip {skip:false, exempt:Some(false)}, ReadPlan::CheckBudgets)] + #[case::read_exemption(ReadFacts::BudgetSkip {skip:false, exempt:None}, ReadPlan::ReadExemption)] + #[case::pricing_disabled(ReadFacts::Pricing {enabled:false, llm_route:None}, ReadPlan::Continue)] + #[case::pricing_route(ReadFacts::Pricing {enabled:true, llm_route:None}, ReadPlan::ReadRoute)] + #[case::pricing_skip(ReadFacts::Pricing {enabled:true, llm_route:Some(false)}, ReadPlan::Continue)] + #[case::pricing_check(ReadFacts::Pricing {enabled:true, llm_route:Some(true)}, ReadPlan::ReadUnpriced)] + #[case::priced(ReadFacts::Unpriced {count:0}, ReadPlan::Continue)] + #[case::unpriced(ReadFacts::Unpriced {count:1}, ReadPlan::RejectUnpriced)] + #[case::no_router(ReadFacts::MemberFallback {router:false, multipart:None, resource_bound:None}, ReadPlan::Continue)] + #[case::content_type(ReadFacts::MemberFallback {router:true, multipart:None, resource_bound:None}, ReadPlan::ReadContentType)] + #[case::multipart(ReadFacts::MemberFallback {router:true, multipart:Some(true), resource_bound:None}, ReadPlan::Continue)] + #[case::resource_read(ReadFacts::MemberFallback {router:true, multipart:Some(false), resource_bound:None}, ReadPlan::ReadResourceBinding)] + #[case::bound_resource(ReadFacts::MemberFallback {router:true, multipart:Some(false), resource_bound:Some(true)}, ReadPlan::Continue)] + #[case::settings(ReadFacts::MemberFallback {router:true, multipart:Some(false), resource_bound:Some(false)}, ReadPlan::ReadFallbackSettings)] + fn selected_reads(#[case] facts: ReadFacts, #[case] expected: ReadPlan) { + assert_eq!(read(facts), expected); + } + + #[rstest] + #[case::absent(None, false, false)] + #[case::empty(Some(""), true, false)] + #[case::user(Some("user"), true, true)] + fn member_identity(#[case] user: Option<&str>, #[case] membership: bool, #[case] models: bool) { + let actual = plan(Facts { + model_present: true, + skip_budgets: false, + team_present: true, + user_present: false, + token_present: true, + user_id: user.map(str::to_owned), + agent_id: None, + customer_present: false, + customer_models: vec![], + }); + assert_eq!(actual.membership, membership); + assert_eq!(actual.member_models, models); + } + + #[rstest] + #[case::empty(false, 0, false, false, vec![BudgetScope::Team, BudgetScope::TeamWindows, BudgetScope::TeamSoft, BudgetScope::Organization, BudgetScope::Tags, BudgetScope::User, BudgetScope::Member])] + #[case::all(true, 1, true, true, vec![BudgetScope::Team, BudgetScope::TeamWindows, BudgetScope::KeyWindows, BudgetScope::TeamSoft, BudgetScope::Organization, BudgetScope::Tags, BudgetScope::Groups, BudgetScope::User, BudgetScope::Member, BudgetScope::Customer])] + fn budget_order( + #[case] token_present: bool, + #[case] group_count: usize, + #[case] customer_present: bool, + #[case] customer_budget_present: bool, + #[case] expected: Vec, + ) { + assert_eq!( + budgets(BudgetFacts { + token_present, + group_count, + customer_present, + customer_budget_present + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/end_user_context.rs b/litellm-rust/crates/request-auth-core/src/end_user_context.rs new file mode 100644 index 00000000000..ef7a5084bc2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/end_user_context.rs @@ -0,0 +1,428 @@ +use serde::{Deserialize, Serialize}; +use strum::VariantArray; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct IdFacts { + pub value: Option, +} + +pub fn normalize_id(facts: &IdFacts) -> Option { + facts.value.clone().filter(|value| !value.is_empty()) +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DefaultFacts { + pub prisma_present: bool, + pub explicit_id: Option, + pub global_read: bool, + pub global_id: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DefaultDecision { + pub read_global: bool, + pub budget_id: Option, +} + +pub fn default_id(facts: &DefaultFacts) -> DefaultDecision { + if facts.explicit_id.is_none() && !facts.global_read { + return DefaultDecision { + read_global: true, + budget_id: None, + }; + } + if !facts.prisma_present { + return DefaultDecision { + read_global: false, + budget_id: None, + }; + } + DefaultDecision { + read_global: false, + budget_id: facts + .explicit_id + .clone() + .or_else(|| facts.global_id.clone()), + } +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ApplyFacts { + pub own_budget_id: Option, + pub own_budget_present: bool, + pub key_id: Option, + pub global_read: bool, + pub global_id: Option, + pub default_read: bool, + pub default_present: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ApplyDecision { + Keep, + ReadGlobal, + ReadDefault, + Copy, +} + +pub fn apply_default(facts: &ApplyFacts) -> ApplyDecision { + if facts.own_budget_id.is_some() && facts.own_budget_present { + return ApplyDecision::Keep; + } + if facts.key_id.is_none() { + if !facts.global_read { + return ApplyDecision::ReadGlobal; + } + if facts.global_id.is_none() { + return ApplyDecision::Keep; + } + } + if !facts.default_read { + return ApplyDecision::ReadDefault; + } + if facts.default_present { + ApplyDecision::Copy + } else { + ApplyDecision::Keep + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct RegistryFacts { + pub end_user_id: String, + pub global_id: Option, + pub key_id: Option, + pub validate_id: bool, + pub token_max_present: bool, + pub registry_read: bool, + pub registry: Option>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum RegistryDecision { + Read, + Restricted, + Unrestricted, +} + +pub fn registry(facts: &RegistryFacts) -> RegistryDecision { + if facts.global_id.is_some() + || facts.key_id.is_some() + || facts.validate_id + || facts.token_max_present + { + return RegistryDecision::Restricted; + } + if !facts.registry_read { + return RegistryDecision::Read; + } + if facts + .registry + .as_ref() + .is_some_and(|ids| !ids.contains(&facts.end_user_id)) + { + RegistryDecision::Unrestricted + } else { + RegistryDecision::Restricted + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum Field { + EndUserId, + Tpm, + Rpm, + Tpd, + MaxBudget, + Models, + Region, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct Fields { + pub tpm: bool, + pub rpm: bool, + pub tpd: bool, + pub max_budget: bool, + pub models: bool, + pub region: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum AssignmentMode { + Params, + Token, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct AssignmentFacts { + pub mode: AssignmentMode, + pub fields: Fields, +} + +pub fn assignments(facts: &AssignmentFacts) -> Vec { + let fields = &facts.fields; + match facts.mode { + AssignmentMode::Params => [ + (Field::Tpm, fields.tpm), + (Field::Rpm, fields.rpm), + (Field::Tpd, fields.tpd), + (Field::MaxBudget, fields.max_budget), + (Field::Models, fields.models), + ] + .into_iter() + .filter_map(|(field, present)| present.then_some(field)) + .collect(), + AssignmentMode::Token => [ + (Field::EndUserId, true), + (Field::Tpm, fields.tpm), + (Field::Rpm, fields.rpm), + (Field::Tpd, fields.tpd), + (Field::Region, fields.region), + (Field::Models, fields.models), + ] + .into_iter() + .filter_map(|(field, present)| present.then_some(field)) + .collect(), + } +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct LookupFacts { + pub own_budget_present: bool, + pub key_id: Option, + pub end_user_read: bool, + pub end_user_present: bool, + pub global_read: bool, + pub global_id: Option, + pub default_read: bool, + pub default_present: bool, + pub token_max_present: bool, + pub applied: bool, + pub read_failed: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum LookupAction { + ReadEndUser, + ReadGlobal, + ReadDefault, + ApplyOwn, + ApplyDefault, + Done, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct LookupDecision { + pub action: LookupAction, + pub set_max_budget: bool, + pub project_budget: bool, +} + +pub fn lookup(facts: &LookupFacts) -> LookupDecision { + let action = if facts.read_failed || facts.applied { + LookupAction::Done + } else if !facts.end_user_read { + LookupAction::ReadEndUser + } else if facts.end_user_present { + LookupAction::ApplyOwn + } else if facts.key_id.is_none() && !facts.global_read { + LookupAction::ReadGlobal + } else if facts.key_id.is_none() && facts.global_id.is_none() { + LookupAction::Done + } else if !facts.default_read { + LookupAction::ReadDefault + } else if facts.default_present { + LookupAction::ApplyDefault + } else { + LookupAction::Done + }; + LookupDecision { + project_budget: action == LookupAction::ApplyDefault + || (action == LookupAction::ApplyOwn && facts.own_budget_present), + set_max_budget: facts.applied + && !facts.read_failed + && !facts.end_user_present + && facts.default_present + && !facts.token_max_present, + action, + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ValidationMode { + #[default] + Resolve, + Exists, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ValidationFacts { + pub mode: ValidationMode, + pub end_user_id: Option, + pub enabled: bool, + pub prisma_present: bool, + pub global_id: Option, + pub key_id: Option, + pub cache_read: bool, + pub cached: Option, + pub end_user_read: bool, + pub end_user_present: bool, + pub user_read: bool, + pub user_present: bool, + pub cache_written: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ValidationRead { + Cache, + EndUser, + User, + WriteCache, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ValidationDecision { + pub read: Option, + pub end_user_id: Option, + pub valid: bool, + pub cache_value: Option, + pub cache_ttl: u64, + pub user_email: Option, +} + +fn validation_result(facts: &ValidationFacts, valid: bool) -> ValidationDecision { + let retain = valid + || facts.key_id.is_some() + || facts.global_id.as_ref().is_some_and(|id| !id.is_empty()); + ValidationDecision { + end_user_id: facts.end_user_id.clone().filter(|_| retain), + valid, + ..ValidationDecision::default() + } +} + +pub fn validation(facts: &ValidationFacts) -> ValidationDecision { + let Some(id) = &facts.end_user_id else { + return ValidationDecision::default(); + }; + if facts.mode == ValidationMode::Resolve { + if !facts.enabled || !facts.prisma_present { + return ValidationDecision { + end_user_id: Some(id.clone()), + ..ValidationDecision::default() + }; + } + if !facts.cache_read { + return ValidationDecision { + read: Some(ValidationRead::Cache), + ..ValidationDecision::default() + }; + } + match facts.cached.as_deref() { + Some("valid") => return validation_result(facts, true), + Some("invalid") => return validation_result(facts, false), + _ => {} + } + } + if !facts.end_user_read { + return ValidationDecision { + read: Some(ValidationRead::EndUser), + ..ValidationDecision::default() + }; + } + if !facts.end_user_present && !facts.user_read { + return ValidationDecision { + read: Some(ValidationRead::User), + user_email: id.contains('@').then(|| id.clone()), + ..ValidationDecision::default() + }; + } + let valid = facts.end_user_present || facts.user_present; + if facts.mode == ValidationMode::Resolve && !facts.cache_written { + return ValidationDecision { + read: Some(ValidationRead::WriteCache), + valid, + cache_value: Some(if valid { "valid" } else { "invalid" }.to_owned()), + cache_ttl: if valid { 300 } else { 60 }, + ..ValidationDecision::default() + }; + } + validation_result(facts, valid) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, None)] + #[case::empty(Some(""), None)] + #[case::space(Some(" "), Some(" "))] + #[case::id(Some("budget"), Some("budget"))] + fn normalizes_metadata_id(#[case] value: Option<&str>, #[case] expected: Option<&str>) { + assert_eq!( + normalize_id(&IdFacts { + value: value.map(str::to_owned) + }) + .as_deref(), + expected + ); + } + + #[rstest] + #[case::explicit(Some("key"), Some("global"), Some("key"))] + #[case::empty_explicit(Some(""), Some("global"), Some(""))] + #[case::global(None, Some("global"), Some("global"))] + #[case::absent(None, None, None)] + fn selects_default( + #[case] explicit: Option<&str>, + #[case] global: Option<&str>, + #[case] expected: Option<&str>, + ) { + let facts = DefaultFacts { + prisma_present: true, + explicit_id: explicit.map(str::to_owned), + global_read: true, + global_id: global.map(str::to_owned), + }; + assert_eq!(default_id(&facts).budget_id.as_deref(), expected); + } + + #[rstest] + #[case::none(None, None, None)] + #[case::empty_global(Some(""), None, None)] + #[case::global(Some("budget"), None, Some("missing"))] + #[case::key(None, Some(""), Some("missing"))] + fn retains_unknown_ids( + #[case] global: Option<&str>, + #[case] key: Option<&str>, + #[case] expected: Option<&str>, + ) { + let facts = ValidationFacts { + end_user_id: Some("missing".to_owned()), + enabled: true, + prisma_present: true, + global_id: global.map(str::to_owned), + key_id: key.map(str::to_owned), + cache_read: true, + cached: Some("invalid".to_owned()), + ..ValidationFacts::default() + }; + assert_eq!(validation(&facts).end_user_id.as_deref(), expected); + } + + #[rstest] + #[case::missing(AssignmentMode::Params, vec![])] + #[case::token_identity(AssignmentMode::Token, vec![Field::EndUserId])] + fn absent_fields(#[case] mode: AssignmentMode, #[case] expected: Vec) { + assert_eq!( + assignments(&AssignmentFacts { + mode, + fields: Fields::default() + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/error.rs b/litellm-rust/crates/request-auth-core/src/error.rs index 6656d9f409c..0f75826b4f7 100644 --- a/litellm-rust/crates/request-auth-core/src/error.rs +++ b/litellm-rust/crates/request-auth-core/src/error.rs @@ -1,5 +1,7 @@ #[derive(Debug, thiserror::Error)] pub enum Error { + #[error("invalid auth pattern")] + Pattern(#[source] fancy_regex::Error), #[error("invalid auth effect answer")] InvalidAnswer, #[error("can't compare offset-naive and offset-aware datetimes")] diff --git a/litellm-rust/crates/request-auth-core/src/fallback_access.rs b/litellm-rust/crates/request-auth-core/src/fallback_access.rs new file mode 100644 index 00000000000..6c080e12ab2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/fallback_access.rs @@ -0,0 +1,145 @@ +use serde::{Deserialize, Serialize}; + +use crate::{ + Error, + model_access::{self, Decision as ModelDecision, Facts as ModelFacts, Request, Response}, +}; + +#[derive(Clone, Deserialize)] +pub struct GroupedScope { + pub facts: ModelFacts, + pub group_ids: Vec, + pub group_models: Option>, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub refuse_agent_keys: bool, + pub agent_id: Option, + pub key: Option, + pub team: Option, + pub embedded_team: Option, + pub assigned_models: Option>, + pub assigned_scope: Option, + pub member: Option, + pub project: Option, + pub team_present: bool, + pub embedded_team_grants: Vec, + pub member_grants: Vec, + pub project_grants: Vec, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: Option }, + ReadKey, + ReadTeam, + ReadEmbeddedTeam, + ReadMember, + ReadProject, + ReadKeyGroups, + ReadTeamGroups, + ReadAssignedGroups, + ReadAssignedScope, +} + +pub fn grouped(scope: GroupedScope, merge: bool) -> Result { + let Response::Decision(decision) = model_access::decide(Request::Grouped { + facts: scope.facts, + group_ids: scope.group_ids, + group_models: scope.group_models, + merge, + })? + else { + return Err(Error::InvalidAnswer); + }; + Ok(decision) +} + +pub fn decide(facts: Facts) -> Result { + if facts.refuse_agent_keys && facts.agent_id.is_some() { + return Ok(Decision::Deny { denial: None }); + } + let Some(key) = facts.key else { + return Ok(Decision::ReadKey); + }; + let key = grouped(key, false)?; + if matches!( + key, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadKeyGroups); + } + if !key.allowed() { + return Ok(Decision::Deny { denial: Some(key) }); + } + if facts.team_present && facts.team.is_none() { + return Ok(Decision::ReadTeam); + } + if let Some(team) = facts.team { + let decision = grouped(team, true)?; + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadTeamGroups); + } + if !decision.allowed() { + if !matches!(decision, ModelDecision::Deny { .. }) { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + let Some(grants) = facts.assigned_models else { + return Ok(Decision::ReadAssignedGroups); + }; + if grants.is_empty() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + let Some(assigned) = facts.assigned_scope else { + return Ok(Decision::ReadAssignedScope); + }; + if !model_access::check(&assigned)?.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + } else if facts.embedded_team.is_none() && !facts.embedded_team_grants.is_empty() { + return Ok(Decision::ReadEmbeddedTeam); + } else if let Some(team) = facts.embedded_team { + let decision = model_access::check(&team)?; + if !decision.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + if facts.member.is_none() && !facts.member_grants.is_empty() { + return Ok(Decision::ReadMember); + } + if let Some(scope) = facts.member { + let decision = model_access::check(&scope)?; + if !decision.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + if facts.project.is_none() && !facts.project_grants.is_empty() { + return Ok(Decision::ReadProject); + } + if let Some(scope) = facts.project { + let decision = model_access::check(&scope)?; + if !decision.allowed() { + return Ok(Decision::Deny { + denial: Some(decision), + }); + } + } + Ok(Decision::Allow) +} diff --git a/litellm-rust/crates/request-auth-core/src/fallback_candidate.rs b/litellm-rust/crates/request-auth-core/src/fallback_candidate.rs new file mode 100644 index 00000000000..214c9bd5ded --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/fallback_candidate.rs @@ -0,0 +1,322 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; +use strum::VariantArray; + +use crate::{ + Error, budget_pricing, common_budgets, fallback_access, model_access, model_attribution, + model_budget_data, route_access, unpriced_models, +}; + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Entity { + Key, + User, + EndUser, + Team, + TeamMember, +} + +#[derive(Clone, Copy, Default, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Membership { + #[default] + Unread, + Loaded, + Failed, +} + +#[derive(Clone, Deserialize)] +pub struct BudgetScope { + pub entity: Entity, + pub identity_present: bool, + pub configured: bool, +} + +#[derive(Clone, Deserialize)] +pub struct BudgetRead { + pub model_index: usize, + pub entity: Entity, + pub snapshot: model_budget_data::Snapshot, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub router_present: bool, + pub user_present: bool, + pub team_present: bool, + pub membership: Membership, + pub block_unpriced: bool, + pub route: Option, + pub unpriced: Option, + pub customer_present: bool, + pub customer_grants: Option>, + pub customer: Option, + pub price: Option, + pub budget_models: Option>, + pub budget_scopes: Option>, + pub budgets: Vec, + pub authorization: Option, + pub attribution_reset: bool, + pub attribution: Option, + pub attribution_stamped: bool, + pub group_budgets_read: bool, + pub group_budgets: Vec, + pub io_failed: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { usable: bool, log_failure: bool }, + ReadMembership, + ReadRoute, + ReadUnpriced, + ReadCustomerGrants, + ReadCustomer, + ReadPrice, + ReadBudgetModels, + ReadBudgetScopes, + ReadBudget { model_index: usize, entity: Entity }, + ReadAuthorization, + ResetAttribution, + ReadAttribution, + StampAttribution { groups: Vec }, + ReadGroupBudgets { groups: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let complete = |usable, log_failure| { + Ok(Decision::Complete { + usable, + log_failure, + }) + }; + if !facts.router_present + || !facts.user_present + || !facts.team_present + || facts.membership == Membership::Failed + { + return complete(false, false); + } + if facts.membership == Membership::Unread { + return Ok(Decision::ReadMembership); + } + if facts.io_failed { + return complete(false, false); + } + if facts.block_unpriced { + let Some(route) = &facts.route else { + return Ok(Decision::ReadRoute); + }; + if route_access::decide(route_access::Request::Classify { + context: route.clone(), + kind: route_access::Classification::Llm, + })? + .allowed + { + let Some(unpriced) = &facts.unpriced else { + return Ok(Decision::ReadUnpriced); + }; + match unpriced_models::decide(unpriced) { + unpriced_models::Decision::Complete { unpriced: true } => { + return complete(false, false); + } + unpriced_models::Decision::Complete { unpriced: false } => (), + _ => return Err(Error::InvalidAnswer), + } + } + } + if facts.customer_present { + let Some(grants) = &facts.customer_grants else { + return Ok(Decision::ReadCustomerGrants); + }; + if !grants.is_empty() { + let Some(customer) = &facts.customer else { + return Ok(Decision::ReadCustomer); + }; + match model_access::decide(model_access::Request::Customer(customer.clone()))? { + model_access::Response::Decision(decision) if decision.allowed() => (), + model_access::Response::Decision(_) => return complete(false, true), + _ => return Err(Error::InvalidAnswer), + } + } + } + let Some(price) = &facts.price else { + return Ok(Decision::ReadPrice); + }; + let priced = budget_pricing::price(price); + if priced.read_deployments { + return Err(Error::InvalidAnswer); + } + if !priced.zero_cost { + let Some(models) = &facts.budget_models else { + return Ok(Decision::ReadBudgetModels); + }; + let Some(scopes) = &facts.budget_scopes else { + return Ok(Decision::ReadBudgetScopes); + }; + let mut seen = BTreeSet::new(); + for (model_index, model) in models.iter().enumerate() { + if !seen.insert(model) { + continue; + } + for entity in Entity::VARIANTS { + if !scopes.iter().any(|scope| { + scope.entity == *entity && scope.identity_present && scope.configured + }) { + continue; + } + let Some(reading) = facts + .budgets + .iter() + .find(|read| read.model_index == model_index && read.entity == *entity) + else { + return Ok(Decision::ReadBudget { + model_index, + entity: *entity, + }); + }; + let budget = model_budget_data::check(&reading.snapshot); + if budget.read.is_some() { + return Err(Error::InvalidAnswer); + } + if budget.denied { + return complete(false, false); + } + } + } + } + let Some(authorization) = &facts.authorization else { + return Ok(Decision::ReadAuthorization); + }; + match fallback_access::decide(authorization.clone())? { + fallback_access::Decision::Allow => (), + fallback_access::Decision::Deny { .. } => return complete(false, false), + _ => return Err(Error::InvalidAnswer), + } + if !facts.attribution_reset { + return Ok(Decision::ResetAttribution); + } + let Some(attribution) = &facts.attribution else { + return Ok(Decision::ReadAttribution); + }; + let model_attribution::Decision::Complete { groups, .. } = + model_attribution::decide(attribution) + else { + return Err(Error::InvalidAnswer); + }; + if !facts.attribution_stamped { + return Ok(Decision::StampAttribution { groups }); + } + if !groups.is_empty() { + if !facts.group_budgets_read { + return Ok(Decision::ReadGroupBudgets { groups }); + } + for facts in &facts.group_budgets { + let budget = common_budgets::check(facts); + if budget.read_spend { + return Err(Error::InvalidAnswer); + } + if budget.denied { + return complete(false, true); + } + } + } + complete(true, false) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn facts() -> Facts { + Facts { + router_present: true, + user_present: true, + team_present: true, + ..Facts::default() + } + } + + #[rstest] + #[case::router(false, true, true, Decision::Complete { usable: false, log_failure: false })] + #[case::user(true, false, true, Decision::Complete { usable: false, log_failure: false })] + #[case::team(true, true, false, Decision::Complete { usable: false, log_failure: false })] + #[case::ready(true, true, true, Decision::ReadMembership)] + fn prerequisites( + facts: Facts, + #[case] router: bool, + #[case] user: bool, + #[case] team: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(&Facts { + router_present: router, + user_present: user, + team_present: team, + ..facts + }) + .unwrap(), + expected + ); + } + + #[rstest] + #[case::key(0, Entity::Key)] + #[case::user(1, Entity::User)] + #[case::enduser(2, Entity::EndUser)] + #[case::team(3, Entity::Team)] + #[case::member(4, Entity::TeamMember)] + fn budgets(facts: Facts, #[case] count: usize, #[case] expected: Entity) { + let facts = Facts { + membership: Membership::Loaded, + price: Some(budget_pricing::ModelPriceFacts { + cached: Some(false), + lookup_failed: false, + group_found: false, + input_price: None, + output_price: None, + deployments: None, + }), + budget_models: Some(vec!["model".into()]), + budget_scopes: Some( + Entity::VARIANTS + .iter() + .map(|entity| BudgetScope { + entity: *entity, + identity_present: true, + configured: true, + }) + .collect(), + ), + budgets: Entity::VARIANTS + .iter() + .take(count) + .map(|entity| BudgetRead { + model_index: 0, + entity: *entity, + snapshot: model_budget_data::Snapshot { + max_budget: None, + local_spend: None, + legacy_spend: None, + shared_spend: None, + shared_started_at: None, + shared_window_seconds: None, + now: 0.0, + }, + }) + .collect(), + ..facts + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::ReadBudget { + model_index: 0, + entity: expected + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/key.rs b/litellm-rust/crates/request-auth-core/src/key.rs index b4b20621759..00eb6a82a97 100644 --- a/litellm-rust/crates/request-auth-core/src/key.rs +++ b/litellm-rust/crates/request-auth-core/src/key.rs @@ -133,9 +133,22 @@ fn check_key(check: KeyCheck) -> Result { if key.blocked { return Step::deny(Denial::KeyBlocked); } - Step::effect(Effect::EnforceModelAccess, move |(): ()| { - check_owner(check, key) - }) + Step::effect( + Effect::EnforceModelAccess, + move |facts: crate::key_model_access::Facts| match crate::key_model_access::decide( + facts, + )? { + crate::key_model_access::Decision::Allow => { + Step::effect(Effect::PrefetchAuthObjects, move |(): ()| { + check_owner(check, key) + }) + } + crate::key_model_access::Decision::Deny { denial } => { + Step::deny(Denial::ModelAccessDenied { decision: denial }) + } + _ => Err(Error::InvalidAnswer), + }, + ) }) } @@ -157,43 +170,54 @@ fn check_owner(check: KeyCheck, key: KeyFacts) -> Result { } fn check_budget_exemption(check: KeyCheck, key: KeyFacts) -> Result { - Step::effect(Effect::CheckBudgetExemption, move |exempt: bool| { - if exempt || key.team_member_spend.is_none() || !check.settings.has_db { - return check_expiry(check, key, exempt); - } - let (Some(user_id), Some(team_id)) = (key.user_id.clone(), key.team_id.clone()) else { - return check_expiry(check, key, exempt); - }; - Step::effect( - Effect::ReadTeamMemberBudget { - user_id: user_id.clone(), - team_id: team_id.clone(), - }, - move |budget: Option| { - let Some(budget) = budget else { - return check_expiry(check, key, false); - }; - if budget.spend < budget.max_budget { - return check_expiry(check, key, false); - } - Step::effect( - Effect::AlertTeamMemberBudget { - user_id: user_id.clone(), - team_id: team_id.clone(), - budget: budget.clone(), - }, - move |(): ()| { - Step::deny(Denial::TeamMemberBudgetExceeded { - user_id, - team_id, - spend: budget.spend, - max_budget: budget.max_budget, - }) - }, - ) - }, - ) - }) + Step::effect( + Effect::CheckBudgetExemption, + move |facts: crate::budget_pricing::ExemptionFacts| { + let decision = crate::budget_pricing::exemption(&facts); + if decision.read_pass_through + || decision.read_price_cache + || decision.read_model.is_some() + { + return Err(Error::InvalidAnswer); + } + let exempt = decision.exempt; + if exempt || key.team_member_spend.is_none() || !check.settings.has_db { + return check_expiry(check, key, exempt); + } + let (Some(user_id), Some(team_id)) = (key.user_id.clone(), key.team_id.clone()) else { + return check_expiry(check, key, exempt); + }; + Step::effect( + Effect::ReadTeamMemberBudget { + user_id: user_id.clone(), + team_id: team_id.clone(), + }, + move |budget: Option| { + let Some(budget) = budget else { + return check_expiry(check, key, false); + }; + if budget.spend < budget.max_budget { + return check_expiry(check, key, false); + } + Step::effect( + Effect::AlertTeamMemberBudget { + user_id: user_id.clone(), + team_id: team_id.clone(), + budget: budget.clone(), + }, + move |(): ()| { + Step::deny(Denial::TeamMemberBudgetExceeded { + user_id, + team_id, + spend: budget.spend, + max_budget: budget.max_budget, + }) + }, + ) + }, + ) + }, + ) } fn check_expiry(check: KeyCheck, key: KeyFacts, exempt: bool) -> Result { diff --git a/litellm-rust/crates/request-auth-core/src/key_model_access.rs b/litellm-rust/crates/request-auth-core/src/key_model_access.rs new file mode 100644 index 00000000000..12fffc55528 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/key_model_access.rs @@ -0,0 +1,88 @@ +use serde::{Deserialize, Serialize}; + +use crate::{ + Error, + fallback_access::{GroupedScope, grouped}, + model_access::Decision as ModelDecision, +}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub config_nonempty: bool, + pub key_models: Vec, + pub primary_loaded: bool, + pub primary: Option, + pub fallback_names: Option>, + pub fallback_scopes: Vec, + pub routing: Vec, +} + +#[derive(Clone, Deserialize)] +pub enum Routing { + Routed, + Failed, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: ModelDecision }, + ReadPrimary, + ReadPrimaryGroups, + ReadFallbackNames, + ReadFallbackScope { index: usize }, + ReadFallbackGroups { index: usize }, + RouteFallback { index: usize }, + RoutingFailure { index: usize }, +} + +pub fn decide(facts: Facts) -> Result { + if facts.config_nonempty + || facts + .key_models + .iter() + .any(|model| model == "all-team-models") + { + return Ok(Decision::Allow); + } + if !facts.primary_loaded { + return Ok(Decision::ReadPrimary); + } + if let Some(primary) = facts.primary { + let decision = grouped(primary, false)?; + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadPrimaryGroups); + } + if !decision.allowed() { + return Ok(Decision::Deny { denial: decision }); + } + } + let Some(names) = facts.fallback_names else { + return Ok(Decision::ReadFallbackNames); + }; + for index in 0..names.len() { + let Some(scope) = facts.fallback_scopes.get(index) else { + return Ok(Decision::ReadFallbackScope { index }); + }; + let decision = grouped(scope.clone(), false)?; + if matches!( + decision, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadFallbackGroups { index }); + } + if !decision.allowed() { + return Ok(Decision::Deny { denial: decision }); + } + match facts.routing.get(index) { + None => return Ok(Decision::RouteFallback { index }), + Some(Routing::Failed) => return Ok(Decision::RoutingFailure { index }), + Some(Routing::Routed) => {} + } + } + Ok(Decision::Allow) +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index a55e5ed7565..592e7268c52 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -1,10 +1,52 @@ #![forbid(unsafe_code)] +pub mod access; +pub mod admission; +pub mod agent_access; +pub mod agent_ceilings; +pub mod agent_grants; +pub mod auth_context; +pub mod budget_alerts; +pub mod budget_context; +pub mod budget_lookup; +pub mod budget_pricing; +pub mod common_budgets; +pub mod common_flow; pub mod credential; +pub mod end_user_context; mod error; +pub mod fallback_access; +pub mod fallback_candidate; mod key; +pub mod key_model_access; +pub mod managed; +pub mod managed_models; +pub mod member_models; +pub mod model_access; +pub mod model_attribution; +pub mod model_budget_data; mod model_budgets; +pub mod model_groups; +pub mod model_helpers; +pub mod model_projection; +pub mod org_routes; +pub mod owner_teams; +pub mod patterns; +pub mod personal_budget; +pub mod personal_model_listing; +pub mod personal_model_request; +pub mod personal_models; +pub mod personal_owner; pub mod protocol; +pub mod reservation; +pub mod resolved_models; +pub mod route_access; +pub mod route_enable; +pub mod route_roles; +pub mod spend_counter_policy; +pub mod team_aliases; +pub mod team_override; +pub mod unpriced_models; pub use error::Error; pub use key::{start, start_resolved_key}; diff --git a/litellm-rust/crates/request-auth-core/src/managed.rs b/litellm-rust/crates/request-auth-core/src/managed.rs new file mode 100644 index 00000000000..deb7354f601 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/managed.rs @@ -0,0 +1,348 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize)] +pub struct BoundIdentity { + pub active: bool, + pub revision: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Agent { + pub agent_id: String, + pub enabled: bool, + pub identity_managed: bool, + pub identity: Option, + pub execution_mode: String, +} + +#[derive(Debug, Deserialize)] +pub struct Context { + pub agent_id: String, + pub binding_revision: Option, + pub mode: String, + pub user_id: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum ManagedFacts { + Actor { + agent: Agent, + context: Option, + }, + ActorLookup { + agent_id: Option, + has_database: bool, + registry_loaded: bool, + registry_managed: bool, + agent_loaded: bool, + agent: Option, + retired: Option, + context: Option, + jwt_claims: bool, + delegation_verified: bool, + human_grants: Option>, + policy_assigned: bool, + }, + UnionGrants { + grants: Vec>>, + }, + IntersectGrants { + left: Option>, + right: Option>, + }, + DeclaredGrants { + direct: Vec, + declared_groups: Vec, + unified_groups: Vec, + group_agents: Option>, + unified_agents: Option>, + }, + ManagedAgentGrants { + permission_present: bool, + own: Option>, + ceilings: Option>>, + caller_loaded: bool, + caller: Option>, + mode: Option, + user_present: bool, + human: Option>, + }, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum ManagedDenial { + Disabled, + MissingContext, + ChangedIdentity, + ExecutionMode, + MissingHuman, + DatabaseRequired, + Removed, + VerifiedIdentityRequired, + HumanNotPermitted, +} + +#[derive(Debug, PartialEq, Serialize, strum::VariantArray)] +pub enum ManagedAction { + Allow, + Deny, + ReadRegistry, + ReadAgent, + ReadRetired, + AssignPolicy, + ReadHumanGrants, + Grants, + ReadDeclaredGroups, + ReadUnifiedGroups, + ReadOwn, + ReadCeilings, + ReadCaller, +} + +#[derive(Debug, PartialEq, Serialize)] +pub struct ManagedDecision { + pub action: ManagedAction, + pub denial: Option, + pub grants: Option>, +} + +fn action(action: ManagedAction) -> ManagedDecision { + ManagedDecision { + action, + denial: None, + grants: None, + } +} + +fn denied(denial: ManagedDenial) -> ManagedDecision { + ManagedDecision { + action: ManagedAction::Deny, + denial: Some(denial), + grants: None, + } +} + +fn actor(agent: &Agent, context: Option<&Context>) -> ManagedDecision { + if !agent.enabled + || !agent + .identity + .as_ref() + .is_some_and(|identity| identity.active) + { + return denied(ManagedDenial::Disabled); + } + let Some(context) = context else { + return denied(ManagedDenial::MissingContext); + }; + if context.agent_id != agent.agent_id + || agent + .identity + .as_ref() + .is_some_and(|identity| context.binding_revision.as_ref() != Some(&identity.revision)) + { + return denied(ManagedDenial::ChangedIdentity); + } + if agent.execution_mode != context.mode && agent.execution_mode != "both" { + return denied(ManagedDenial::ExecutionMode); + } + if context.mode == "delegated" + && !context + .user_id + .as_ref() + .is_some_and(|user| !user.is_empty()) + { + return denied(ManagedDenial::MissingHuman); + } + action(ManagedAction::Allow) +} + +pub fn decide(facts: ManagedFacts) -> ManagedDecision { + match facts { + ManagedFacts::Actor { agent, context } => actor(&agent, context.as_ref()), + ManagedFacts::ActorLookup { + agent_id, + has_database, + registry_loaded, + registry_managed, + agent_loaded, + agent, + retired, + context, + jwt_claims, + delegation_verified, + human_grants, + policy_assigned, + } => { + if agent_id.is_none() { + return action(ManagedAction::Allow); + } + if !has_database { + if !registry_loaded { + return action(ManagedAction::ReadRegistry); + } + return if context.is_some() || registry_managed { + denied(ManagedDenial::DatabaseRequired) + } else { + action(ManagedAction::Allow) + }; + } + if !agent_loaded { + return action(ManagedAction::ReadAgent); + } + let Some(agent) = agent else { + return match retired { + None => action(ManagedAction::ReadRetired), + Some(retired) if retired || context.is_some() => denied(ManagedDenial::Removed), + Some(_) => action(ManagedAction::Allow), + }; + }; + if !agent.identity_managed { + return action(ManagedAction::Allow); + } + if jwt_claims && context.is_none() { + return denied(ManagedDenial::VerifiedIdentityRequired); + } + let checked = actor(&agent, context.as_ref()); + if checked.denial.is_some() { + return checked; + } + if !policy_assigned { + return action(ManagedAction::AssignPolicy); + } + if context + .as_ref() + .is_some_and(|context| context.mode == "delegated") + && !delegation_verified + { + return match human_grants { + None => action(ManagedAction::ReadHumanGrants), + Some(grants) if !grants.contains(&agent.agent_id) => { + denied(ManagedDenial::HumanNotPermitted) + } + Some(_) => action(ManagedAction::Allow), + }; + } + action(ManagedAction::Allow) + } + ManagedFacts::UnionGrants { grants } => { + let grants = if !grants.is_empty() && grants.iter().all(Option::is_none) { + None + } else { + Some(grants.into_iter().flatten().flatten().collect()) + }; + ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants, + } + } + ManagedFacts::IntersectGrants { left, right } => { + let grants = match (left, right) { + (None, right) => right, + (left, None) => left, + (Some(left), Some(right)) => { + Some(left.into_iter().filter(|id| right.contains(id)).collect()) + } + }; + ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants, + } + } + ManagedFacts::DeclaredGrants { + direct, + declared_groups, + unified_groups, + group_agents, + unified_agents, + } => { + if !declared_groups.is_empty() && group_agents.is_none() { + return action(ManagedAction::ReadDeclaredGroups); + } + if !unified_groups.is_empty() && unified_agents.is_none() { + return action(ManagedAction::ReadUnifiedGroups); + } + let grants = + if direct.is_empty() && declared_groups.is_empty() && unified_groups.is_empty() { + None + } else { + Some( + direct + .into_iter() + .chain(group_agents.into_iter().flatten()) + .chain(unified_agents.into_iter().flatten()) + .collect(), + ) + }; + ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants, + } + } + ManagedFacts::ManagedAgentGrants { + permission_present, + own, + ceilings, + caller_loaded, + caller, + mode, + user_present, + human, + } => { + let grants = |ids| ManagedDecision { + action: ManagedAction::Grants, + denial: None, + grants: Some(ids), + }; + if !permission_present { + return grants(vec![]); + } + let Some(own) = own else { + return action(ManagedAction::ReadOwn); + }; + let Some(ceilings) = ceilings else { + return action(ManagedAction::ReadCeilings); + }; + if !caller_loaded { + return action(ManagedAction::ReadCaller); + } + let capped: Vec<_> = own + .into_iter() + .filter(|id| { + ceilings.iter().all(|ceiling| ceiling.contains(id)) + && caller.as_ref().is_none_or(|ids| ids.contains(id)) + }) + .collect(); + if mode.is_none() || mode.as_deref() == Some("autonomous") { + return grants(capped); + } + if !user_present { + return grants(vec![]); + } + let Some(human) = human else { + return action(ManagedAction::ReadHumanGrants); + }; + grants(capped.into_iter().filter(|id| human.contains(id)).collect()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty(vec![], Some(vec![]))] + #[case::unrestricted(vec![None], None)] + #[case::mixed(vec![None, Some(vec!["a".into()])], Some(vec!["a".into()]))] + fn union(#[case] grants: Vec>>, #[case] expected: Option>) { + assert_eq!( + decide(ManagedFacts::UnionGrants { grants }).grants, + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/managed_models.rs b/litellm-rust/crates/request-auth-core/src/managed_models.rs new file mode 100644 index 00000000000..f618b38b371 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/managed_models.rs @@ -0,0 +1,253 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + member_models, + model_access::{self, Decision as ModelDecision}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Default, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Purpose { + #[default] + All, + Ceilings, + Caller, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub purpose: Purpose, + pub model_present: bool, + pub token_present: bool, + pub managed_loaded: bool, + pub managed_present: bool, + pub managed_models_valid: bool, + pub managed_models_count: usize, + pub managed: Option, + pub ceilings: Option>>, + pub aliases_loaded: bool, + pub scopes: Vec, + pub caller_loaded: bool, + pub caller_present: bool, + pub team_loaded: bool, + pub team_present: bool, + pub team: Option, + pub member: Option, + pub user_loaded: bool, + pub user_present: bool, + pub user_grants: Vec, + pub user: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: ModelDecision, + member_models: Option>, + }, + ReadManaged, + ReadManagedScope, + ReadCeilings, + ReadAliases, + ReadCeilingScope { + index: usize, + }, + ReadCaller, + ReadTeam, + ReadTeamScope, + ReadTeamGroups, + ReadMember, + ReadUser, + ReadUserScope, +} + +pub fn decide(facts: &Facts) -> Result { + let deny = |denial| { + Ok(Decision::Deny { + denial, + member_models: None, + }) + }; + if facts.purpose == Purpose::All && !facts.managed_loaded { + return Ok(Decision::ReadManaged); + } + if !facts.model_present || !facts.token_present { + return Ok(Decision::Allow); + } + if facts.purpose == Purpose::All && facts.managed_present { + if !facts.managed_models_valid || facts.managed_models_count == 0 { + return deny(ModelDecision::NoAgentModels); + } + let Some(scope) = &facts.managed else { + return Ok(Decision::ReadManagedScope); + }; + let result = model_access::check(scope)?; + if !result.allowed() { + return deny(result); + } + } + if facts.purpose != Purpose::Caller { + let Some(ceilings) = &facts.ceilings else { + return Ok(Decision::ReadCeilings); + }; + if !facts.aliases_loaded { + return Ok(Decision::ReadAliases); + } + for (index, models) in ceilings.iter().enumerate() { + if models.is_empty() { + return deny(ModelDecision::EmptyAgentCeiling); + } + let Some(scope) = facts.scopes.get(index) else { + return Ok(Decision::ReadCeilingScope { index }); + }; + let result = model_access::check(scope)?; + if !result.allowed() { + return deny(result); + } + } + if facts.purpose == Purpose::Ceilings { + return Ok(Decision::Allow); + } + } + if !facts.caller_loaded { + return Ok(Decision::ReadCaller); + } + if !facts.caller_present { + return Ok(Decision::Allow); + } + if !facts.team_loaded { + return Ok(Decision::ReadTeam); + } + if facts.team_present { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeamScope); + }; + let result = fallback_access::grouped(team.clone(), true)?; + if matches!( + result, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadTeamGroups); + } + if !result.allowed() { + return deny(result); + } + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember); + }; + return match member_models::decide(member)? { + member_models::Decision::Allow => Ok(Decision::Allow), + member_models::Decision::Deny { + denial, + member_models, + } => Ok(Decision::Deny { + denial, + member_models, + }), + _ => Err(Error::InvalidAnswer), + }; + } + if !facts.user_loaded { + return Ok(Decision::ReadUser); + } + let precheck = model_access::decide(model_access::Request::UserPrecheck { + present: facts.user_present, + grants: facts.user_grants.clone(), + })?; + let model_access::Response::Decision(precheck) = precheck else { + return Err(Error::InvalidAnswer); + }; + if !matches!(precheck, ModelDecision::ReadScope) { + return if precheck.allowed() { + Ok(Decision::Allow) + } else { + deny(precheck) + }; + } + let Some(user) = &facts.user else { + return Ok(Decision::ReadUserScope); + }; + let result = model_access::check(user)?; + if result.allowed() { + Ok(Decision::Allow) + } else { + deny(result) + } +} + +#[derive(Clone, Default, Deserialize)] +pub struct CeilingSource { + pub agent_present: bool, + pub managed_loaded: bool, + pub managed_present: bool, + pub loaded: bool, + pub count: usize, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum SourceDecision { + ReadManaged, + ReadManagedCeilings, + ReadUnmanagedCeiling, + Complete { indices: Vec }, +} + +pub fn ceiling_source(facts: &CeilingSource) -> SourceDecision { + if !facts.agent_present { + return SourceDecision::Complete { indices: vec![] }; + } + if !facts.managed_loaded { + return SourceDecision::ReadManaged; + } + if facts.loaded { + return SourceDecision::Complete { + indices: (0..facts.count).collect(), + }; + } + if facts.managed_present { + SourceDecision::ReadManagedCeilings + } else { + SourceDecision::ReadUnmanagedCeiling + } +} + +#[derive(Clone, Deserialize)] +pub struct GrantFacts { + pub models: Option>>, + pub ceilings: Vec>, +} + +pub fn grants(facts: &GrantFacts) -> std::collections::BTreeSet { + facts + .models + .iter() + .flatten() + .flatten() + .chain(facts.ceilings.iter().flatten()) + .cloned() + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::main(Purpose::All, Decision::ReadManaged)] + #[case::groups(Purpose::Ceilings, Decision::Allow)] + #[case::caller(Purpose::Caller, Decision::Allow)] + fn empty_model_read_order(#[case] purpose: Purpose, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + purpose, + ..Facts::default() + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/member_models.rs b/litellm-rust/crates/request-auth-core/src/member_models.rs new file mode 100644 index 00000000000..a4cacc8959b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/member_models.rs @@ -0,0 +1,77 @@ +use crate::{Error, model_access}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub team_present: bool, + pub membership_loaded: bool, + pub grants: Option>, + pub scope: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: model_access::Decision, + member_models: Option>, + }, + ReadMembership, + ReadGrants, + ReadScope, +} + +pub fn decide(facts: &Facts) -> Result { + if !facts.user_present || !facts.team_present { + return Ok(Decision::Allow); + } + if !facts.membership_loaded { + return Ok(Decision::ReadMembership); + } + let Some(grants) = &facts.grants else { + return Ok(Decision::ReadGrants); + }; + if grants.is_empty() { + return Ok(Decision::Allow); + } + let Some(scope) = &facts.scope else { + return Ok(Decision::ReadScope); + }; + let denial = model_access::check(scope)?; + if denial.allowed() { + return Ok(Decision::Allow); + } + let member_models = + matches!(denial, model_access::Decision::Deny { .. }).then(|| grants.clone()); + Ok(Decision::Deny { + denial, + member_models, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(false, false, Decision::Allow)] + #[case::lookup(true, false, Decision::ReadMembership)] + #[case::preloaded(true, true, Decision::ReadGrants)] + fn membership_read( + #[case] user_present: bool, + #[case] membership_loaded: bool, + #[case] expected: Decision, + ) { + let facts = Facts { + user_present, + team_present: true, + membership_loaded, + grants: None, + scope: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_access.rs b/litellm-rust/crates/request-auth-core/src/model_access.rs new file mode 100644 index 00000000000..06211ca5cb9 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_access.rs @@ -0,0 +1,674 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive(Clone, Default, Deserialize)] +pub struct Provider { + pub stripped_model: String, + pub provider: String, +} + +#[derive(Clone, Default, Deserialize)] +pub struct ModelRecord { + pub access_groups: Vec, + pub provider: Option, + pub deployment_team: Option, + pub deployment_public_name: Option, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Registry { + pub present: bool, + pub names: Vec, + pub aliases: BTreeMap, + pub records: BTreeMap, + pub compaction_parents: BTreeMap, + pub helper_present: Option, + pub evaluations: Option>, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Aliases { + pub team: BTreeMap, + pub team_order: Option>, + pub key: BTreeMap, + #[serde(rename = "global_")] + pub global: BTreeMap, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub grants: Vec, + pub team_id: Option, + pub aliases: Aliases, + pub registry: Registry, + pub object_type: ObjectType, + pub fallback_depth: usize, + pub max_depth: usize, + pub is_list: bool, +} + +#[derive(Clone, Copy, Deserialize, Serialize, Debug, PartialEq, Eq, strum::VariantArray)] +#[serde(rename_all = "snake_case")] +pub enum ObjectType { + User, + Customer, + Team, + Key, + Org, + Project, + Agent, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + model: String, + grants: Vec, + object_type: ObjectType, + }, + Depth { + model: Option, + }, + NoDefaultModels, + NoAgentModels, + EmptyAgentCeiling, + PersonalTeam { + model: String, + }, + ReadGroups, + ReadScope, + ReadGroupScope { + grants: Vec, + }, +} + +impl Decision { + pub fn allowed(&self) -> bool { + matches!(self, Self::Allow) + } +} + +fn unique(values: impl Iterator) -> Vec { + let mut seen = BTreeSet::new(); + values.filter(|value| seen.insert(value.clone())).collect() +} + +pub fn resolve_key_grants( + grants: &[String], + team_id: Option<&str>, + team_grants: &[String], +) -> Vec { + if grants.iter().any(|grant| grant == "all-team-models") { + return if team_id.is_none() { + Vec::new() + } else { + team_grants.to_vec() + }; + } + grants.to_vec() +} + +fn expanded_grants(facts: &Facts) -> Vec { + if !facts.grants.iter().any(|grant| grant == "all-team-models") + || facts.team_id.is_none() + || !facts.registry.present + { + return facts.grants.clone(); + } + let others = facts + .grants + .iter() + .filter(|grant| grant.as_str() != "all-team-models") + .cloned(); + if facts.registry.names.is_empty() { + let values: Vec<_> = others.collect(); + return if values.is_empty() { + facts.grants.clone() + } else { + values + }; + } + unique(others.chain(facts.registry.names.iter().cloned())) +} + +fn helper(model: &str, facts: &Facts, aliases: bool) -> Result { + let grants = expanded_grants(facts); + let record = facts.registry.records.get(model); + let groups = record.map_or(&[][..], |record| record.access_groups.as_slice()); + if grants.iter().any(|grant| groups.contains(grant)) { + return Ok(true); + } + let filtered: Vec<_> = grants + .iter() + .filter(|grant| !groups.contains(grant)) + .collect(); + if aliases && facts.aliases.team.contains_key(model) { + return Ok(true); + } + for grant in &filtered { + if patterns::model_pattern(model, grant)? { + return Ok(true); + } + } + if let Some(provider) = record.and_then(|record| record.provider.as_ref()) + && (provider.stripped_model != model || !model.contains('/')) + { + for grant in &filtered { + if patterns::model_pattern( + &format!("{}/{}", provider.provider, provider.stripped_model), + grant, + )? { + return Ok(true); + } + } + } + Ok((filtered.is_empty() && grants.is_empty()) + || filtered.iter().any(|grant| { + matches!(grant.as_str(), "*" | "all-proxy-models") || grant.as_str() == model + })) +} + +pub fn candidate_names( + model: &str, + aliases: &Aliases, + registry: &Registry, + team_id: Option<&str>, +) -> (Vec, Vec) { + let after_team = aliases.team.get(model).map_or(model, String::as_str); + let after_key = aliases + .key + .get(after_team) + .map_or(after_team, String::as_str); + let after_global = aliases + .global + .get(after_key) + .map_or(after_key, String::as_str); + let dispatched = aliases + .key + .get(after_global) + .map_or(after_global, String::as_str); + if after_key != after_team || dispatched != after_global { + return (vec![dispatched.to_owned()], vec![]); + } + let candidates = std::iter::once(model.to_owned()) + .chain(registry.compaction_parents.get(model).cloned()) + .chain( + aliases + .global + .get(model) + .or_else(|| registry.aliases.get(model)) + .filter(|target| !target.is_empty()) + .cloned(), + ) + .collect(); + let Some(team) = team_id else { + return (candidates, vec![]); + }; + if !registry.present || !model.starts_with(&format!("model_name_{team}_")) { + return (candidates, vec![]); + } + let record = registry.records.get(model); + let public = record + .filter(|record| record.deployment_team.as_deref() == Some(team)) + .and_then(|record| 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(), + |names| names.iter().collect(), + ); + let legacy = ordered + .into_iter() + .filter(|alias| { + aliases + .team + .get(*alias) + .is_some_and(|target| target == model) + }) + .cloned(); + (candidates, public.into_iter().chain(legacy).collect()) +} + +pub fn live_team_target(model: &str, facts: &Facts) -> String { + let Some(target) = facts.aliases.team.get(model) else { + return model.to_owned(); + }; + let team = facts.team_id.as_deref().unwrap_or("None"); + if facts.registry.present + && target.starts_with(&format!("model_name_{team}_")) + && !facts.registry.names.contains(target) + { + return model.to_owned(); + } + target.clone() +} + +pub fn check(facts: &Facts) -> Result { + check_at(facts, 0) +} + +fn check_at(facts: &Facts, skip: usize) -> Result { + if let Some(evaluations) = &facts.registry.evaluations { + let Some(evaluation) = evaluations + .iter() + .skip(skip) + .find(|evaluation| evaluation.grants == facts.grants) + else { + return Ok(Decision::ReadGroupScope { + grants: facts.grants.clone(), + }); + }; + let scoped = Facts { + aliases: evaluation + .aliases + .clone() + .unwrap_or_else(|| facts.aliases.clone()), + ..facts.clone() + }; + return match crate::model_projection::decide(&scoped, &evaluation.reads)? { + crate::model_projection::Decision::Complete { decision } => Ok(decision), + _ => Err(Error::InvalidAnswer), + }; + } + if facts.fallback_depth >= facts.max_depth { + return Ok(Decision::Depth { model: None }); + } + for model in &facts.requested { + if facts.is_list && facts.fallback_depth + 1 >= facts.max_depth { + return Ok(Decision::Depth { + model: Some(model.clone()), + }); + } + let (candidates, owning_names) = candidate_names( + model, + &facts.aliases, + &facts.registry, + facts.team_id.as_deref(), + ); + let mut allowed = false; + for candidate in candidates { + if helper(&candidate, facts, true)? { + allowed = true; + break; + } + } + if !allowed { + for candidate in owning_names { + if helper(&candidate, facts, false)? { + allowed = true; + break; + } + } + } + if !allowed { + return Ok(Decision::Deny { + model: model.clone(), + grants: facts.grants.clone(), + object_type: facts.object_type, + }); + } + } + Ok(Decision::Allow) +} + +#[derive(Deserialize)] +pub enum Request { + Object(Facts), + Grouped { + facts: Facts, + group_ids: Vec, + group_models: Option>, + merge: bool, + }, + Helper { + model: String, + facts: Facts, + }, + User { + facts: Facts, + present: bool, + }, + UserPrecheck { + present: bool, + grants: Vec, + }, + Customer(Facts), + Agent { + facts: Facts, + managed_grants: Option>, + ceilings: Vec>, + }, + Personal { + requested: Vec, + teams: Vec, + }, + KeyGrants { + grants: Vec, + team_id: Option, + team_grants: Vec, + }, + Candidates(Facts), + PersonalApplies { + enabled: bool, + token_present: bool, + team_id: Option, + user_id: Option, + master: bool, + role: Option, + api_key_present: bool, + inheritance: bool, + }, + AssignedGroups { + groups: Vec, + team_id: Option, + key_id: Option, + }, +} + +#[derive(Deserialize)] +pub struct AssignedGroup { + pub assigned_team_ids: Vec, + pub assigned_key_ids: Vec, + pub models: Vec, +} + +#[derive(Deserialize)] +pub struct PersonalTeam { + pub blocked: bool, + pub available: bool, + pub team: Facts, + pub member: Option, +} + +#[derive(Serialize)] +#[serde(untagged)] +pub enum Response { + Decision(Decision), + Models(Vec), + Predicate(bool), +} + +pub fn decide(request: Request) -> Result { + let decision = match request { + Request::Object(facts) => check(&facts)?, + Request::Grouped { + facts, + group_ids, + group_models, + merge, + } => { + let primary = check(&facts)?; + if primary.allowed() + || !matches!(primary, Decision::Deny { .. }) + || group_ids.is_empty() + { + primary + } else if let Some(models) = group_models { + if models.is_empty() { + primary + } else { + check_at( + &Facts { + grants: if merge { + unique(facts.grants.iter().cloned().chain(models)) + } else { + models + }, + ..facts + }, + 1, + )? + } + } else { + Decision::ReadGroups + } + } + Request::Helper { model, facts } => { + return Ok(Response::Predicate(helper(&model, &facts, true)?)); + } + Request::User { facts, present } => { + if !present { + Decision::Allow + } else if facts + .grants + .iter() + .any(|grant| grant == "no-default-models") + { + Decision::NoDefaultModels + } else { + check(&facts)? + } + } + Request::UserPrecheck { present, grants } => { + if !present { + Decision::Allow + } else if grants.iter().any(|grant| grant == "no-default-models") { + Decision::NoDefaultModels + } else { + Decision::ReadScope + } + } + Request::Customer(facts) => { + for name in &facts.requested { + let target = live_team_target(name, &facts); + if target != *name && facts.grants.contains(name) { + continue; + } + let customer = Facts { + requested: vec![target], + is_list: false, + team_id: None, + aliases: Aliases { + key: facts.aliases.key.clone(), + global: facts.aliases.global.clone(), + ..Default::default() + }, + ..facts.clone() + }; + let result = check(&customer)?; + if !result.allowed() { + return Ok(Response::Decision(result)); + } + } + Decision::Allow + } + Request::Agent { + facts, + managed_grants, + ceilings, + } => { + if facts.requested.is_empty() + || (facts.requested.len() == 1 && !facts.is_list && facts.requested[0].is_empty()) + { + return Ok(Response::Decision(Decision::Allow)); + } + let dispatched = facts + .requested + .iter() + .map(|name| live_team_target(name, &facts)) + .collect(); + let scope = Facts { + requested: dispatched, + aliases: Aliases { + key: facts.aliases.key.clone(), + global: facts.aliases.global.clone(), + ..Default::default() + }, + object_type: ObjectType::Agent, + ..facts + }; + if let Some(grants) = managed_grants { + if grants.is_empty() { + return Ok(Response::Decision(Decision::NoAgentModels)); + } + let result = check(&Facts { + grants, + ..scope.clone() + })?; + if !result.allowed() { + return Ok(Response::Decision(result)); + } + } + for grants in ceilings { + if grants.is_empty() { + return Ok(Response::Decision(Decision::EmptyAgentCeiling)); + } + let result = check(&Facts { + grants, + ..scope.clone() + })?; + if !result.allowed() { + return Ok(Response::Decision(result)); + } + } + Decision::Allow + } + Request::Personal { requested, teams } => { + for name in requested { + let mut allowed = false; + for team in &teams { + if team.blocked || !team.available { + continue; + } + if !check(&Facts { + requested: vec![name.clone()], + is_list: false, + ..team.team.clone() + })? + .allowed() + { + continue; + } + if let Some(member) = &team.member + && !check(&Facts { + requested: vec![name.clone()], + is_list: false, + ..member.clone() + })? + .allowed() + { + continue; + } + allowed = true; + break; + } + if !allowed { + return Ok(Response::Decision(Decision::PersonalTeam { model: name })); + } + } + Decision::Allow + } + Request::KeyGrants { + grants, + team_id, + team_grants, + } => { + return Ok(Response::Models(resolve_key_grants( + &grants, + team_id.as_deref(), + &team_grants, + ))); + } + Request::Candidates(facts) => { + return Ok(Response::Models(unique(facts.requested.iter().flat_map( + |model| { + let (candidates, owned) = candidate_names( + model, + &facts.aliases, + &facts.registry, + facts.team_id.as_deref(), + ); + candidates.into_iter().chain(owned) + }, + )))); + } + Request::PersonalApplies { + enabled, + token_present, + team_id, + user_id, + master, + role, + api_key_present, + inheritance, + } => { + return Ok(Response::Predicate( + enabled + && token_present + && team_id.is_none() + && user_id.is_some() + && !master + && role.as_deref() != Some("proxy_admin") + && (!inheritance || api_key_present), + )); + } + Request::AssignedGroups { + groups, + team_id, + key_id, + } => { + return Ok(Response::Models(unique( + groups + .into_iter() + .filter(|group| { + team_id.as_ref().is_some_and(|id| { + !id.is_empty() && group.assigned_team_ids.contains(id) + }) || key_id + .as_ref() + .is_some_and(|id| !id.is_empty() && group.assigned_key_ids.contains(id)) + }) + .flat_map(|group| group.models), + ))); + } + }; + Ok(Response::Decision(decision)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(false, vec!["no-default-models"], Decision::Allow)] + #[case::restricted(true, vec!["no-default-models"], Decision::NoDefaultModels)] + #[case::empty(true, vec![], Decision::ReadScope)] + #[case::grants(true, vec!["model"], Decision::ReadScope)] + fn user_precheck(#[case] present: bool, #[case] grants: Vec<&str>, #[case] expected: Decision) { + let response = decide(Request::UserPrecheck { + present, + grants: grants.into_iter().map(str::to_owned).collect(), + }) + .unwrap(); + let Response::Decision(decision) = response else { + panic!("expected decision") + }; + assert_eq!(decision, expected); + } + + #[rstest] + #[case::empty(vec![], true)] + #[case::wildcard(vec!["*"], true)] + #[case::proxy(vec!["all-proxy-models"], true)] + #[case::other(vec!["other"], false)] + #[case::exact(vec!["model"], true)] + fn model_grants(#[case] grants: Vec<&str>, #[case] expected: bool) { + let facts = Facts { + requested: vec!["model".into()], + grants: grants.into_iter().map(str::to_owned).collect(), + team_id: None, + aliases: Aliases::default(), + registry: Registry::default(), + object_type: ObjectType::Key, + fallback_depth: 0, + max_depth: 10, + is_list: false, + }; + assert_eq!(check(&facts).unwrap().allowed(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_attribution.rs b/litellm-rust/crates/request-auth-core/src/model_attribution.rs new file mode 100644 index 00000000000..7cea0a344b0 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_attribution.rs @@ -0,0 +1,147 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub model_present: bool, + pub token_present: bool, + pub router_present: bool, + pub database_present: bool, + pub team_present: bool, + pub user_present: bool, + pub membership_loaded: bool, + pub token_org: Option, + pub team_org: Option, + pub registry_loaded: bool, + pub registry: Option>, + pub covering: Option>, + pub grants: Vec>, + pub org_failed: bool, + pub io_failed: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { groups: Vec, stamp: bool }, + ReadRegistry, + ReadCovering, + ReadKey, + ReadTeam, + ReadMember { lookup: bool, include: bool }, + ReadProject, + ReadOrg { id: Option }, +} + +pub fn decide(facts: &Facts) -> Decision { + let empty = || Decision::Complete { + groups: vec![], + stamp: false, + }; + if facts.io_failed + || !facts.model_present + || !facts.token_present + || !facts.router_present + || !facts.database_present + { + return empty(); + } + if !facts.registry_loaded { + return Decision::ReadRegistry; + } + if facts.registry.as_ref().is_some_and(Vec::is_empty) { + return empty(); + } + let Some(covering) = &facts.covering else { + return Decision::ReadCovering; + }; + let budgeted: BTreeSet<_> = covering + .iter() + .filter(|group| { + facts + .registry + .as_ref() + .is_none_or(|registry| registry.contains(group)) + }) + .collect(); + if budgeted.is_empty() { + return empty(); + } + match facts.grants.len() { + 0 => return Decision::ReadKey, + 1 => return Decision::ReadTeam, + 2 => { + return Decision::ReadMember { + lookup: facts.team_present && facts.user_present && !facts.membership_loaded, + include: facts.team_present && facts.user_present, + }; + } + 3 => return Decision::ReadProject, + 4 => { + return Decision::ReadOrg { + id: facts + .token_org + .as_ref() + .filter(|id| !id.is_empty()) + .or(facts.team_org.as_ref()) + .cloned(), + }; + } + _ => (), + } + let granted: BTreeSet<_> = facts + .grants + .iter() + .take(if facts.org_failed { 4 } else { 5 }) + .flatten() + .collect(); + let groups: Vec<_> = budgeted + .intersection(&granted) + .map(|group| (*group).clone()) + .collect(); + Decision::Complete { + stamp: !groups.is_empty(), + groups, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::unknown(None, vec!["a", "b"])] + #[case::known(Some(vec!["a".into()]), vec!["a"])] + #[case::empty(Some(vec![]), vec![])] + fn intersects_budgeted_covering_and_granted( + #[case] registry: Option>, + #[case] expected: Vec<&str>, + ) { + let facts = Facts { + model_present: true, + token_present: true, + router_present: true, + database_present: true, + registry_loaded: true, + registry, + covering: Some(vec!["b".into(), "a".into(), "b".into()]), + grants: vec![ + vec!["a".into()], + vec![], + vec!["b".into()], + vec![], + vec!["c".into()], + ], + ..Default::default() + }; + assert_eq!( + decide(&facts), + Decision::Complete { + stamp: !expected.is_empty(), + groups: expected.into_iter().map(str::to_owned).collect() + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_budget_data.rs b/litellm-rust/crates/request-auth-core/src/model_budget_data.rs new file mode 100644 index 00000000000..51ebd2fc53c --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_budget_data.rs @@ -0,0 +1,263 @@ +use std::cmp::Ordering::{Equal, Greater, Less}; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Entry { + pub name: String, + pub max_budget: Option, + pub duration: Option, + pub duration_seconds: Option, +} + +impl Entry { + fn usable(&self) -> bool { + self.duration.is_some() && self.duration_seconds.is_some() + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ResolveFacts { + pub model: String, + pub bedrock_base: String, + pub bedrock_provider: String, + pub entries: Vec, + pub key_entries: Vec, + pub team: bool, + pub access_groups: Option>, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Resolution { + pub read_groups: bool, + pub entry: Option, + pub named_model: bool, +} + +fn candidates(facts: &ResolveFacts) -> Vec<&str> { + let base_names = if facts.bedrock_provider.starts_with("bedrock") { + vec![ + facts.bedrock_base.as_str(), + facts.bedrock_base.split_once('.').map_or("", |(_, s)| s), + ] + } else { + vec![] + }; + std::iter::once(facts.model.as_str()) + .chain(std::iter::once( + facts.model.rsplit('/').next().unwrap_or(&facts.model), + )) + .chain(base_names.into_iter().filter(|s| !s.is_empty())) + .collect() +} + +fn exact<'a>(entries: &'a [Entry], names: &[&str]) -> Option<&'a Entry> { + names.iter().find_map(|name| { + entries + .iter() + .find(|entry| entry.name == *name && entry.usable()) + }) +} + +pub fn resolve(facts: &ResolveFacts) -> Resolution { + let names = candidates(facts); + let key_replaces = facts.team + && exact(&facts.key_entries, &names) + .is_some_and(|entry| entry.max_budget.is_some_and(|limit| limit >= 0.0)); + if !key_replaces && let Some(entry) = exact(&facts.entries, &names) { + return Resolution { + read_groups: false, + entry: Some(entry.clone()), + named_model: true, + }; + } + let group_entries: Vec<&Entry> = facts + .entries + .iter() + .filter(|entry| !entry.name.contains('*') && !names.contains(&entry.name.as_str())) + .collect(); + if !group_entries.is_empty() && facts.access_groups.is_none() { + return Resolution { + read_groups: true, + entry: None, + named_model: false, + }; + } + let group = group_entries + .into_iter() + .filter(|entry| { + entry.usable() + && facts + .access_groups + .as_ref() + .is_some_and(|groups| groups.contains(&entry.name)) + }) + .min_by(|a, b| a.name.cmp(&b.name)); + if let Some(entry) = group { + return Resolution { + read_groups: false, + entry: Some(entry.clone()), + named_model: false, + }; + } + let entry = facts + .entries + .iter() + .filter(|entry| { + entry.usable() + && !(key_replaces && names.contains(&entry.name.as_str())) + && entry.name.contains('*') + && names.iter().any(|name| glob(&entry.name, name)) + }) + .max_by_key(|entry| { + let wildcards = entry.name.chars().filter(|c| *c == '*').count(); + ( + entry.name.chars().count() - wildcards, + std::cmp::Reverse(wildcards), + &entry.name, + ) + }) + .cloned(); + Resolution { + read_groups: false, + entry, + named_model: false, + } +} + +fn glob(pattern: &str, name: &str) -> bool { + let parts: Vec<&str> = pattern.split('*').collect(); + if parts.len() == 1 { + return pattern == name; + } + let head = parts[0]; + let tail = parts[parts.len() - 1]; + if !name.starts_with(head) || !name.ends_with(tail) || head.len() + tail.len() > name.len() { + return false; + } + parts[1..parts.len() - 1] + .iter() + .try_fold(head.len(), |position, piece| { + name[position..name.len() - tail.len()] + .find(piece) + .map(|index| position + index + piece.len()) + }) + .is_some() +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Entity { + Key, + User, + Team, + TeamMember, + EndUser, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Snapshot { + pub max_budget: Option, + pub local_spend: Option, + pub legacy_spend: Option, + pub shared_spend: Option, + pub shared_started_at: Option, + pub shared_window_seconds: Option, + pub now: f64, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct QueryFacts { + pub entity: Entity, + pub snapshot: Snapshot, + pub duration_seconds: Option, + pub process_uptime: f64, + pub same_counter: bool, + pub redis_present: bool, + pub shared_read: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Read { + Local, + Legacy, + Shared, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct Decision { + pub read: Option, + pub denied: bool, + pub current_spend: f64, +} + +pub fn check(snapshot: &Snapshot) -> Decision { + let local = snapshot.local_spend.unwrap_or(0.0) + snapshot.legacy_spend.unwrap_or(0.0); + let current_spend = match ( + snapshot.shared_spend, + snapshot.shared_started_at, + snapshot.shared_window_seconds, + ) { + (Some(spend), Some(start), Some(seconds)) + if snapshot.now < start + seconds && spend > local => + { + spend + } + _ => local, + }; + Decision { + read: None, + denied: snapshot + .max_budget + .is_some_and(|limit| limit.partial_cmp(&0.0) != Some(Less) && current_spend >= limit), + current_spend, + } +} + +pub fn query(facts: &QueryFacts) -> Decision { + let done = check(&facts.snapshot); + let Some(limit) = facts.snapshot.max_budget else { + return done; + }; + if limit < 0.0 { + return done; + } + let read = if facts.snapshot.local_spend.is_none() { + Some(Read::Local) + } else if matches!(facts.entity, Entity::Key | Entity::EndUser) + && facts + .duration_seconds + .is_some_and(|seconds| facts.process_uptime < seconds) + && !facts.same_counter + && facts.snapshot.legacy_spend.is_none() + { + Some(Read::Legacy) + } else if !matches!( + (facts.snapshot.local_spend.unwrap_or(0.0) + facts.snapshot.legacy_spend.unwrap_or(0.0)) + .partial_cmp(&limit), + Some(Greater | Equal) + ) && !facts.redis_present + && facts.entity != Entity::EndUser + && !facts.shared_read + { + Some(Read::Shared) + } else { + None + }; + Decision { read, ..done } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty("*", "", true)] + #[case::unicode("é*界", "é世世界", true)] + #[case::overlap("ab*bc", "abc", false)] + #[case::middle("a*b*c", "axybzc", true)] + #[case::not_regex("a?*", "ab", false)] + fn wildcard(#[case] pattern: &str, #[case] name: &str, #[case] matches: bool) { + assert_eq!(glob(pattern, name), matches); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_budgets.rs b/litellm-rust/crates/request-auth-core/src/model_budgets.rs index b02ed6d62c4..709a13c8a43 100644 --- a/litellm-rust/crates/request-auth-core/src/model_budgets.rs +++ b/litellm-rust/crates/request-auth-core/src/model_budgets.rs @@ -53,7 +53,8 @@ fn key_models( Effect::CheckKeyModelBudget { model: model.clone(), }, - move |exceeded: Option| { + move |reading: ModelBudgetRead| { + let exceeded = reading.exceeded(); let Some(exceeded) = exceeded else { return key_models(fallbacks, rest, rewritable, then); }; @@ -86,7 +87,8 @@ fn first_affordable( Effect::CheckKeyModelBudget { model: candidate.clone(), }, - move |over: Option| { + move |reading: ModelBudgetRead| { + let over = reading.exceeded(); if over.is_some() { return first_affordable(rest, exceeded, found); } @@ -94,9 +96,16 @@ fn first_affordable( Effect::AuthorizeBudgetFallback { model: candidate.clone(), }, - move |authorized: bool| { - if !authorized { - return Step::deny(exceeded); + move |facts: crate::fallback_access::Facts| { + match crate::fallback_access::decide(facts)? { + crate::fallback_access::Decision::Allow => (), + crate::fallback_access::Decision::Deny { + denial: Some(denial @ crate::model_access::Decision::Depth { .. }), + } => return Step::deny(Denial::ModelAccessDenied { decision: denial }), + crate::fallback_access::Decision::Deny { .. } => { + return Step::deny(exceeded); + } + _ => return Err(Error::InvalidAnswer), } Step::effect( Effect::RewriteRequestModel { @@ -122,7 +131,7 @@ fn user_models(enabled: bool, models: Vec, then: Then) -> Result| match exceeded { + move |reading: ModelBudgetRead| match reading.exceeded() { Some(exceeded) => Step::deny(exceeded), None => user_models(enabled, rest, then), }, @@ -155,7 +164,8 @@ fn member_models( Effect::CheckTeamMemberModelBudget { model: model.clone(), }, - move |exceeded: Option| { + move |reading: ModelBudgetRead| { + let exceeded = reading.exceeded(); let Some(exceeded) = exceeded else { return member_models(rest, rewritable, can_fall_back); }; @@ -197,7 +207,12 @@ fn first_usable( Effect::VetTeamMemberFallback { target: target.clone(), }, - move |usable: bool| { + move |facts: crate::fallback_candidate::Facts| { + let crate::fallback_candidate::Decision::Complete { usable, .. } = + crate::fallback_candidate::decide(&facts)? + else { + return Err(Error::InvalidAnswer); + }; if !usable { return first_usable(rest, exceeded, found); } @@ -234,11 +249,42 @@ mod tests { assert!( matches!(read.node, Node::Effect(Effect::CheckKeyModelBudget { ref model }) if model == "first") ); - let authorize = (read.resume.unwrap())(Answer::Exceeded(None)).unwrap(); + let authorize = (read.resume.unwrap())(Answer::ModelBudgetData(ModelBudgetRead { + snapshot: crate::model_budget_data::Snapshot { + max_budget: None, + local_spend: None, + legacy_spend: None, + shared_spend: None, + shared_started_at: None, + shared_window_seconds: None, + now: 0.0, + }, + message: String::new(), + entity_type: None, + entity_id: None, + })) + .unwrap(); assert!( matches!(authorize.node, Node::Effect(Effect::AuthorizeBudgetFallback { ref model }) if model == "first") ); - let result = (authorize.resume.unwrap())(Answer::Bool(false)).unwrap(); + let result = (authorize.resume.unwrap())(Answer::FallbackAccess(Box::new( + crate::fallback_access::Facts { + refuse_agent_keys: true, + agent_id: Some("agent".into()), + key: None, + team: None, + embedded_team: None, + assigned_models: None, + assigned_scope: None, + member: None, + project: None, + team_present: false, + embedded_team_grants: vec![], + member_grants: vec![], + project_grants: vec![], + }, + ))) + .unwrap(); assert!(matches!( result.node, Node::Outcome(Outcome::Deny { diff --git a/litellm-rust/crates/request-auth-core/src/model_groups.rs b/litellm-rust/crates/request-auth-core/src/model_groups.rs new file mode 100644 index 00000000000..7853d4b2c5f --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_groups.rs @@ -0,0 +1,73 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +pub struct Group { + pub assigned_team_ids: Vec, + pub assigned_key_ids: Vec, + pub resources: Vec, +} + +#[derive(Deserialize)] +#[serde(tag = "kind", content = "group")] +pub enum Lookup { + Found(Group), + Missing, + Failed, +} + +#[derive(Deserialize)] +pub struct Facts { + pub subject_present: bool, + pub database_present: bool, + pub cache_present: bool, + pub group_ids: Vec, + pub team_id: Option, + pub key_id: Option, + pub check_db_only: bool, + pub strict_errors: bool, + pub loaded: Vec, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Read { index: usize }, + Raise { index: usize }, + Complete { resources: Vec }, +} + +pub fn decide(facts: Facts) -> Decision { + if !facts.subject_present || !facts.database_present || !facts.cache_present { + return Decision::Complete { resources: vec![] }; + } + let mut resources = BTreeSet::new(); + for index in 0..facts.group_ids.len() { + let Some(loaded) = facts.loaded.get(index) else { + return Decision::Read { index }; + }; + match loaded { + Lookup::Missing | Lookup::Failed if facts.check_db_only || facts.strict_errors => { + return Decision::Raise { index }; + } + Lookup::Missing | Lookup::Failed => {} + Lookup::Found(group) => { + if facts + .team_id + .as_ref() + .is_some_and(|id| !id.is_empty() && group.assigned_team_ids.contains(id)) + || facts + .key_id + .as_ref() + .is_some_and(|id| !id.is_empty() && group.assigned_key_ids.contains(id)) + { + resources.extend(group.resources.iter().cloned()); + } + } + } + } + Decision::Complete { + resources: resources.into_iter().collect(), + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_helpers.rs b/litellm-rust/crates/request-auth-core/src/model_helpers.rs new file mode 100644 index 00000000000..9e4ac468a87 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_helpers.rs @@ -0,0 +1,148 @@ +use crate::{Error, model_access::Provider, patterns}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Group, + Helper, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub model: String, + pub grants: Option>, + pub team_id: Option, + pub team_aliases: Vec, + pub router_present: bool, + pub groups: Option>, + pub names: Option>, + pub provider_loaded: bool, + pub provider: Option, + pub provider_reads: Option>>, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + ReadGroups, + ReadNames, + ReadProvider, +} + +pub fn decide(facts: &Facts) -> Result { + let done = |allowed| Ok(Decision::Complete { allowed }); + if facts.kind == Kind::Group + && facts + .grants + .as_ref() + .is_none_or(|grants| grants.contains(&facts.model)) + { + return done(true); + } + let grants = facts.grants.as_ref().ok_or(Error::InvalidAnswer)?; + if facts.router_present && facts.groups.is_none() { + return Ok(Decision::ReadGroups); + } + let groups = facts.groups.as_deref().unwrap_or_default(); + if facts.kind == Kind::Group { + return done(grants.iter().any(|grant| groups.contains(grant))); + } + let expanded = if facts.router_present + && facts.team_id.is_some() + && grants.iter().any(|grant| grant == "all-team-models") + { + let Some(names) = &facts.names else { + return Ok(Decision::ReadNames); + }; + let mut seen = BTreeSet::new(); + let values: Vec<_> = grants + .iter() + .filter(|grant| *grant != "all-team-models") + .chain(names) + .filter(|grant| seen.insert(*grant)) + .cloned() + .collect(); + if values.is_empty() { + grants.clone() + } else { + values + } + } else { + grants.clone() + }; + if expanded.iter().any(|grant| groups.contains(grant)) + || facts.team_aliases.contains(&facts.model) + { + return done(true); + } + let filtered: Vec<_> = expanded + .iter() + .filter(|grant| !groups.contains(grant)) + .collect(); + let wildcards: Vec<_> = filtered + .iter() + .filter(|grant| grant.contains('*')) + .collect(); + for pattern in &wildcards { + if patterns::model_pattern(&facts.model, pattern)? { + return done(true); + } + } + for (index, pattern) in wildcards.iter().enumerate() { + let provider = if let Some(providers) = &facts.provider_reads { + let Some(provider) = providers.get(index) else { + return Ok(Decision::ReadProvider); + }; + provider + } else { + if !facts.provider_loaded { + return Ok(Decision::ReadProvider); + } + &facts.provider + }; + if let Some(provider) = provider + && (provider.stripped_model != facts.model || !facts.model.contains('/')) + && patterns::model_pattern( + &format!("{}/{}", provider.provider, provider.stripped_model), + pattern, + )? + { + return done(true); + } + } + done( + (filtered.is_empty() && expanded.is_empty()) + || filtered.iter().any(|grant| { + matches!(grant.as_str(), "*" | "all-proxy-models") || **grant == facts.model + }), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::group(Kind::Group, Decision::Complete { allowed: true })] + #[case::helper(Kind::Helper, Decision::ReadGroups)] + fn direct_match_read_order(#[case] kind: Kind, #[case] expected: Decision) { + let facts = Facts { + kind, + model: "model".into(), + grants: Some(vec!["model".into()]), + team_id: None, + team_aliases: vec![], + router_present: true, + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_projection.rs b/litellm-rust/crates/request-auth-core/src/model_projection.rs new file mode 100644 index 00000000000..829b72c703b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_projection.rs @@ -0,0 +1,237 @@ +use crate::{ + Error, + model_access::{self, Decision as ModelDecision, Facts, ModelRecord, Provider, Registry}, + model_helpers, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum ReadKind { + Compaction, + Aliases, + AliasTarget, + Deployment, + Groups, + Names, + Provider, +} + +#[derive(Clone, Deserialize)] +pub struct Reading { + pub kind: ReadKind, + pub model: String, + pub values: Vec, + pub value: Option, + pub team: Option, + pub provider: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Evaluation { + pub grants: Vec, + pub reads: Vec, + pub aliases: Option, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { decision: ModelDecision }, + Read { read: ReadKind, model: String }, +} + +enum Pending { + Read(ReadKind, String), + Invalid(Error), +} +impl From for Pending { + fn from(error: Error) -> Self { + Self::Invalid(error) + } +} + +struct Reader<'a> { + reads: &'a [Reading], + index: usize, +} +impl<'a> Reader<'a> { + fn read(&mut self, kind: ReadKind, model: &str) -> Result<&'a Reading, Pending> { + let Some(value) = self.reads.get(self.index) else { + return Err(Pending::Read(kind, model.to_owned())); + }; + if value.kind != kind || value.model != model { + return Err(Pending::Invalid(Error::InvalidAnswer)); + } + self.index += 1; + Ok(value) + } +} + +fn helper( + model: &str, + facts: &Facts, + aliases: bool, + reader: &mut Reader<'_>, +) -> Result { + let mut helper = model_helpers::Facts { + kind: model_helpers::Kind::Helper, + model: model.to_owned(), + grants: Some(facts.grants.clone()), + team_id: facts.team_id.clone(), + team_aliases: if aliases { + facts.aliases.team.keys().cloned().collect() + } else { + vec![] + }, + router_present: facts + .registry + .helper_present + .unwrap_or(facts.registry.present), + groups: None, + names: None, + provider_loaded: false, + provider: None, + provider_reads: Some(vec![]), + }; + loop { + match model_helpers::decide(&helper)? { + model_helpers::Decision::Complete { allowed } => return Ok(allowed), + model_helpers::Decision::ReadGroups => { + helper.groups = Some(reader.read(ReadKind::Groups, model)?.values.clone()) + } + model_helpers::Decision::ReadNames => { + helper.names = Some(reader.read(ReadKind::Names, model)?.values.clone()) + } + model_helpers::Decision::ReadProvider => { + helper + .provider_reads + .as_mut() + .ok_or(Error::InvalidAnswer)? + .push(reader.read(ReadKind::Provider, model)?.provider.clone()); + } + } + } +} + +fn check_one(model: &str, facts: &Facts, reader: &mut Reader<'_>) -> Result { + let compaction = reader.read(ReadKind::Compaction, model)?.value.clone(); + let mut registry = Registry { + present: facts.registry.present, + ..Registry::default() + }; + if let Some(parent) = compaction { + registry.compaction_parents.insert(model.to_owned(), parent); + } + if !facts.aliases.global.contains_key(model) && registry.present { + let aliases = reader.read(ReadKind::Aliases, model)?; + if aliases.values.iter().any(|alias| alias == model) + && let Some(target) = &reader.read(ReadKind::AliasTarget, model)?.value + { + registry.aliases.insert(model.to_owned(), target.clone()); + } + } + let after_team = facts.aliases.team.get(model).map_or(model, String::as_str); + let after_key = facts + .aliases + .key + .get(after_team) + .map_or(after_team, String::as_str); + let after_global = facts + .aliases + .global + .get(after_key) + .map_or(after_key, String::as_str); + let dispatched = facts + .aliases + .key + .get(after_global) + .map_or(after_global, String::as_str); + let key_alias = after_key != after_team || dispatched != after_global; + if !key_alias + && registry.present + && facts + .team_id + .as_ref() + .is_some_and(|team| model.starts_with(&format!("model_name_{team}_"))) + { + let deployment = reader.read(ReadKind::Deployment, model)?; + registry.records.insert( + model.to_owned(), + ModelRecord { + deployment_team: deployment.team.clone(), + deployment_public_name: deployment.value.clone(), + ..ModelRecord::default() + }, + ); + } + let (candidates, owning) = + model_access::candidate_names(model, &facts.aliases, ®istry, facts.team_id.as_deref()); + for candidate in candidates { + if helper(&candidate, facts, true, reader)? { + return Ok(true); + } + } + for candidate in owning { + if helper(&candidate, facts, false, reader)? { + return Ok(true); + } + } + Ok(false) +} + +fn check(facts: &Facts, reader: &mut Reader<'_>) -> Result { + if facts.fallback_depth >= facts.max_depth { + return Ok(ModelDecision::Depth { model: None }); + } + for model in &facts.requested { + if facts.is_list && facts.fallback_depth + 1 >= facts.max_depth { + return Ok(ModelDecision::Depth { + model: Some(model.clone()), + }); + } + if !check_one(model, facts, reader)? { + return Ok(ModelDecision::Deny { + model: model.clone(), + grants: facts.grants.clone(), + object_type: facts.object_type, + }); + } + } + Ok(ModelDecision::Allow) +} + +pub fn decide(facts: &Facts, reads: &[Reading]) -> Result { + let mut reader = Reader { reads, index: 0 }; + match check(facts, &mut reader) { + Ok(decision) if reader.index == reads.len() => Ok(Decision::Complete { decision }), + Ok(_) => Err(Error::InvalidAnswer), + Err(Pending::Read(read, model)) => Ok(Decision::Read { read, model }), + Err(Pending::Invalid(error)) => Err(error), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + #[rstest] + #[case::empty(vec![], true)] + #[case::model(vec!["model".to_owned()], false)] + fn empty_list_never_reads_registry(#[case] requested: Vec, #[case] complete: bool) { + let facts = Facts { + requested, + grants: vec![], + team_id: None, + aliases: Default::default(), + registry: Default::default(), + object_type: model_access::ObjectType::Key, + fallback_depth: 0, + max_depth: 5, + is_list: true, + }; + assert_eq!( + matches!(decide(&facts, &[]).unwrap(), Decision::Complete { .. }), + complete + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/org_routes.rs b/litellm-rust/crates/request-auth-core/src/org_routes.rs new file mode 100644 index 00000000000..d361e46a25e --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/org_routes.rs @@ -0,0 +1,304 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Membership { + pub organization_id: Option, + pub role: Option, +} + +#[derive(Deserialize)] +pub struct Facts { + pub route: String, + pub user_present: bool, + pub user_role: Option, + pub memberships: Option>, + pub organization_id: Option, + pub org_admin_routes: Vec, + #[serde(default)] + pub organization_id_present: bool, + #[serde(default)] + pub identifier_pending: bool, + #[serde(default)] + pub hash_failure: Option, + #[serde(default)] + pub request_pending: bool, + #[serde(default)] + pub memberships_pending: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Denial { + CreateOrganization, + NoMemberships, + MissingOrganization, + MissingRole, + WrongRole, + TeamMissingOrganization, + TeamWrongRole, +} + +#[derive(Debug, Default, PartialEq, Eq, Serialize)] +pub struct Decision { + pub denial: Option, + pub organization_role: Option, + pub organization_ids: Vec, + pub read_identifier: bool, + pub hash_failure: Option, + pub read_request: bool, + pub read_memberships: bool, +} + +pub fn decide(facts: Facts) -> Decision { + if facts.user_present && facts.request_pending { + return Decision { + read_request: true, + ..Decision::default() + }; + } + if facts.user_present + && facts.user_role.as_deref() != Some("proxy_admin") + && facts.route != "/organization/new" + && (facts.org_admin_routes.contains(&facts.route) || facts.route == "/team/new") + && facts.memberships_pending + { + return Decision { + read_memberships: true, + ..Decision::default() + }; + } + let memberships = facts.memberships.as_deref().unwrap_or(&[]); + let ids = memberships + .iter() + .filter_map(|item| item.organization_id.clone()) + .collect(); + let role = memberships + .iter() + .rev() + .find(|item| { + item.organization_id.is_some() && item.organization_id == facts.organization_id + }) + .and_then(|item| item.role.clone()); + let present = facts.organization_id_present || facts.organization_id.is_some(); + let needs_lookup = facts.user_present + && facts.user_role.as_deref() != Some("proxy_admin") + && facts.route != "/organization/new" + && present + && ((facts.org_admin_routes.contains(&facts.route) && facts.memberships.is_some()) + || facts.route == "/team/new" && !memberships.is_empty()); + if needs_lookup && (facts.identifier_pending || facts.hash_failure.is_some()) { + return Decision { + denial: None, + organization_role: role, + organization_ids: ids, + read_identifier: facts.identifier_pending, + hash_failure: facts.hash_failure, + ..Decision::default() + }; + } + let denial = if !facts.user_present || facts.user_role.as_deref() == Some("proxy_admin") { + None + } else if facts.route == "/organization/new" { + Some(Denial::CreateOrganization) + } else if facts.org_admin_routes.contains(&facts.route) { + if facts.memberships.is_none() { + Some(Denial::NoMemberships) + } else if !present { + Some(Denial::MissingOrganization) + } else if role.is_none() { + Some(Denial::MissingRole) + } else if role.as_deref() != Some("org_admin") { + Some(Denial::WrongRole) + } else { + None + } + } else if facts.route == "/team/new" && !memberships.is_empty() { + if !present { + Some(Denial::TeamMissingOrganization) + } else if role.as_deref() != Some("org_admin") { + Some(Denial::TeamWrongRole) + } else { + None + } + } else { + None + }; + Decision { + denial, + organization_role: role, + organization_ids: ids, + read_identifier: false, + hash_failure: None, + ..Decision::default() + } +} + +pub fn user_is_admin(memberships: &[Membership], organizations: &[String]) -> bool { + !organizations.is_empty() + && organizations.iter().all(|id| { + memberships.iter().any(|item| { + item.organization_id.as_ref() == Some(id) + && item.role.as_deref() == Some("org_admin") + }) + }) +} + +#[derive(Deserialize)] +pub struct Identifier { + pub text: Option, + pub hash_failure: Option, +} + +#[derive(Deserialize)] +pub struct AdminFacts { + pub memberships: Option>, + pub request_read: bool, + pub singular_present: bool, + pub plural_count: usize, + pub candidate_index: usize, + pub candidate: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum AdminAction { + ReadRequest, + ReadSingular, + ReadPlural, + Allow, + Deny, + HashFailure, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct AdminDecision { + pub action: AdminAction, + pub candidate_index: usize, + pub plural_index: usize, + pub hash_failure: Option, +} + +pub fn admin(facts: &AdminFacts) -> AdminDecision { + let result = |action| AdminDecision { + action, + candidate_index: 0, + plural_index: 0, + hash_failure: None, + }; + let Some(memberships) = &facts.memberships else { + return result(AdminAction::Deny); + }; + if !facts.request_read { + return result(AdminAction::ReadRequest); + } + let count = usize::from(facts.singular_present) + facts.plural_count; + if count == 0 { + return result(AdminAction::Deny); + } + let next = if let Some(candidate) = &facts.candidate { + if let Some(message) = &candidate.hash_failure { + return AdminDecision { + hash_failure: Some(message.clone()), + ..result(AdminAction::HashFailure) + }; + } + if !candidate.text.as_ref().is_some_and(|id| { + memberships.iter().any(|item| { + item.organization_id.as_ref() == Some(id) + && item.role.as_deref() == Some("org_admin") + }) + }) { + return result(AdminAction::Deny); + } + facts.candidate_index + 1 + } else { + 0 + }; + if next == count { + return result(AdminAction::Allow); + } + AdminDecision { + action: if facts.singular_present && next == 0 { + AdminAction::ReadSingular + } else { + AdminAction::ReadPlural + }, + candidate_index: next, + plural_index: next.saturating_sub(usize::from(facts.singular_present)), + hash_failure: None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::no_memberships(false, false, None, None, AdminAction::Deny)] + #[case::read_request(true, false, None, None, AdminAction::ReadRequest)] + #[case::read_singular(true, true, None, None, AdminAction::ReadSingular)] + #[case::allowed(true, true, Some("org"), None, AdminAction::Allow)] + #[case::wrong_org(true, true, Some("other"), None, AdminAction::Deny)] + #[case::unhashable( + true, + true, + None, + Some("unhashable type: 'list'"), + AdminAction::HashFailure + )] + fn raw_admin_stages( + #[case] present: bool, + #[case] request_read: bool, + #[case] id: Option<&str>, + #[case] failure: Option<&str>, + #[case] expected: AdminAction, + ) { + let memberships = present.then(|| { + vec![Membership { + organization_id: Some("org".into()), + role: Some("org_admin".into()), + }] + }); + let candidate = (id.is_some() || failure.is_some()).then(|| Identifier { + text: id.map(str::to_owned), + hash_failure: failure.map(str::to_owned), + }); + assert_eq!( + admin(&AdminFacts { + memberships, + request_read, + singular_present: true, + plural_count: 0, + candidate_index: 0, + candidate + }) + .action, + expected + ); + } + + #[rstest] + #[case::none(None, Some(Denial::NoMemberships))] + #[case::empty(Some(vec![]), Some(Denial::MissingOrganization))] + fn membership_presence( + #[case] memberships: Option>, + #[case] expected: Option, + ) { + assert_eq!( + decide(Facts { + route: "/organization/update".into(), + user_present: true, + user_role: None, + memberships, + organization_id: None, + org_admin_routes: vec!["/organization/update".into()], + organization_id_present: false, + identifier_pending: false, + hash_failure: None, + request_pending: false, + memberships_pending: false, + }) + .denial, + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/owner_teams.rs b/litellm-rust/crates/request-auth-core/src/owner_teams.rs new file mode 100644 index 00000000000..d8816cf5407 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/owner_teams.rs @@ -0,0 +1,173 @@ +use crate::{Error, personal_budget}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Purpose { + Models, + Grants, +} + +#[derive(Clone, Deserialize)] +pub struct Load { + pub present: bool, + pub failed: bool, + pub blocked: bool, + pub org: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Budget { + pub index: usize, + pub facts: personal_budget::Facts, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub purpose: Purpose, + pub user_present: bool, + pub database_present: bool, + pub ids: Vec, + pub ui_team_id: String, + pub org: Option, + pub loaded: bool, + pub loads: Vec, + pub budgets: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadTeams { + ids: Vec, + }, + ReadBudgets { + indices: Vec, + }, + Complete { + selected: Vec, + withheld: Vec, + unreadable: bool, + }, +} + +pub fn decide(facts: &Facts) -> Result { + if facts.purpose == Purpose::Grants && (!facts.user_present || !facts.database_present) { + return Ok(Decision::Complete { + selected: vec![], + withheld: vec![], + unreadable: false, + }); + } + if !facts.loaded { + let mut seen = BTreeSet::new(); + let ids = facts + .ids + .iter() + .filter(|id| { + facts.purpose == Purpose::Models || (!id.is_empty() && **id != facts.ui_team_id) + }) + .filter(|id| seen.insert(*id)) + .cloned() + .collect(); + return Ok(Decision::ReadTeams { ids }); + } + let eligible: Vec<_> = facts + .loads + .iter() + .enumerate() + .filter(|(_, load)| load.present && !load.blocked) + .filter(|(_, load)| { + facts.purpose == Purpose::Models + || facts.org.is_none() + || load.org.is_none() + || load.org == facts.org + }) + .map(|(index, _)| index) + .collect(); + let unreadable = facts.loads.iter().any(|load| load.failed); + if facts.purpose == Purpose::Models { + return Ok(Decision::Complete { + selected: eligible, + withheld: vec![], + unreadable, + }); + } + let missing: Vec<_> = eligible + .iter() + .filter(|index| !facts.budgets.iter().any(|budget| budget.index == **index)) + .copied() + .collect(); + if !missing.is_empty() { + return Ok(Decision::ReadBudgets { indices: missing }); + } + let results: Result, _> = eligible + .iter() + .map(|index| { + let budget = facts + .budgets + .iter() + .find(|budget| budget.index == *index) + .ok_or(Error::InvalidAnswer)?; + match personal_budget::decide(&budget.facts)? { + personal_budget::Decision::Complete { allowed, .. } => Ok((*index, allowed)), + _ => Err(Error::InvalidAnswer), + } + }) + .collect(); + let results = results?; + let blocked = facts + .loads + .iter() + .enumerate() + .filter(|(_, load)| { + load.present + && load.blocked + && (facts.org.is_none() || load.org.is_none() || load.org == facts.org) + }) + .map(|(index, _)| index); + let withheld = blocked + .chain( + results + .iter() + .filter_map(|(index, allowed)| (!allowed).then_some(*index)), + ) + .collect(); + Ok(Decision::Complete { + selected: results + .into_iter() + .filter_map(|(index, allowed)| allowed.then_some(index)) + .collect(), + withheld, + unreadable, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::models(Purpose::Models, vec!["", "ui", "a"])] + #[case::grants(Purpose::Grants, vec!["a"])] + fn id_filter(#[case] purpose: Purpose, #[case] expected: Vec<&str>) { + let facts = Facts { + purpose, + user_present: true, + database_present: true, + ids: vec!["".into(), "ui".into(), "a".into(), "a".into()], + ui_team_id: "ui".into(), + org: None, + loaded: false, + loads: vec![], + budgets: vec![], + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::ReadTeams { + ids: expected.into_iter().map(str::to_owned).collect() + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/patterns.rs b/litellm-rust/crates/request-auth-core/src/patterns.rs new file mode 100644 index 00000000000..dc378e0f01b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/patterns.rs @@ -0,0 +1,134 @@ +use fancy_regex::Regex; +use std::{ + collections::BTreeMap, + sync::{Arc, LazyLock, Mutex}, +}; + +use crate::Error; + +static EXPRESSIONS: LazyLock>>> = + LazyLock::new(|| Mutex::new(BTreeMap::new())); + +fn expression(pattern: &str) -> Result, Error> { + if let Some(compiled) = EXPRESSIONS + .lock() + .unwrap_or_else(|error| error.into_inner()) + .get(pattern) + { + return Ok(Arc::clone(compiled)); + } + let compiled = Arc::new(Regex::new(pattern).map_err(Error::Pattern)?); + let mut expressions = EXPRESSIONS + .lock() + .unwrap_or_else(|error| error.into_inner()); + if expressions.len() >= 1024 { + expressions.clear(); + } + expressions.insert(pattern.to_owned(), Arc::clone(&compiled)); + Ok(compiled) +} + +pub fn regex_match(pattern: &str, value: &str) -> Result { + let expression = expression(pattern)?; + Ok(expression + .find(value) + .map_err(Error::Pattern)? + .is_some_and(|found| found.start() == 0)) +} + +pub fn model_pattern(model: &str, pattern: &str) -> Result { + if !pattern.contains('*') { + return Ok(false); + } + regex_match( + &format!("^{}(?=\\n?\\z)", pattern.replace('*', ".*")), + model, + ) +} + +pub fn route_pattern(route: &str, pattern: &str) -> Result { + let mut expression = String::from("^"); + let mut rest = pattern; + while let Some(start) = rest.find('{') { + expression.push_str(&rest[..start]); + let Some(end) = rest[start + 1..].find('}').map(|index| index + start + 1) else { + expression.push_str(&rest[start..]); + rest = ""; + break; + }; + let name = &rest[start + 1..end]; + if name.is_empty() { + expression.push_str("{}"); + } else if name.ends_with(":path") { + expression.push_str(if rest[end + 1..].contains(':') { + "[^:]+" + } else { + "[\\s\\S]+" + }); + } else { + expression.push_str("[^/]+"); + } + rest = &rest[end + 1..]; + } + expression.push_str(rest); + expression.push_str("(?=\\n?\\z)"); + regex_match(&expression, route) +} + +pub fn wildcard(route: &str, pattern: &str) -> bool { + pattern + .strip_suffix('*') + .map_or(route == pattern, |prefix| route.starts_with(prefix)) +} + +pub fn prefix(route: &str, pattern: &str) -> bool { + route == pattern + || route + .strip_prefix(pattern) + .is_some_and(|suffix| suffix.starts_with('/')) +} + +pub fn route_access(route: &str, patterns: &[String]) -> Result { + if patterns + .iter() + .any(|pattern| pattern == route || (pattern.ends_with('*') && wildcard(route, pattern))) + { + return Ok(true); + } + for pattern in patterns { + if route_pattern(route, pattern)? { + return Ok(true); + } + } + Ok(false) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + #[rstest] + #[case::segment("/key/k/reset", "/key/{id}/reset", true)] + #[case::segment_slash("/key/a/b/reset", "/key/{id}/reset", false)] + #[case::path_slash("/credentials/a/b/start", "/credentials/{id:path}/start", true)] + #[case::path_newline("/credentials/a\nb/start", "/credentials/{id:path}/start", true)] + #[case::google_suffix( + "/models/a:b:generateContent", + "/models/{id:path}:generateContent", + false + )] + #[case::final_newline("/models/a\n", "/models/{id}", true)] + fn route_templates(#[case] route: &str, #[case] pattern: &str, #[case] expected: bool) { + assert_eq!(super::route_pattern(route, pattern).unwrap(), expected); + } + + #[rstest] + #[case::provider("openai/a", "openai/*", true)] + #[case::exact_not_pattern("a", "a", false)] + #[case::regex_dot("aa-1", "a.-*", true)] + #[case::newline("a\nb", "*", false)] + #[case::final_newline("a\n", "a*", true)] + fn model_patterns(#[case] model: &str, #[case] pattern: &str, #[case] expected: bool) { + assert_eq!(super::model_pattern(model, pattern).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_budget.rs b/litellm-rust/crates/request-auth-core/src/personal_budget.rs new file mode 100644 index 00000000000..7f2656f401d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_budget.rs @@ -0,0 +1,186 @@ +use serde::{Deserialize, Serialize}; + +use crate::{ + Error, + common_budgets::{self, BudgetFacts, MemberLimitFacts}, +}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub team_id_present: bool, + pub check_db_only: bool, + pub token_org: Option, + pub team_org: Option, + pub source_ready: bool, + pub team: Option, + pub team_alerted: bool, + pub window_count: Option, + pub windows: Vec, + pub membership_loaded: bool, + pub member_limit: Option, + pub default_id: Option, + pub default_loaded: bool, + pub member: Option, + pub member_alerted: bool, + pub org_loaded: bool, + pub org_failed: bool, + pub org: Option, + pub org_alerted: bool, + pub io_failed: bool, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool, log_failure: bool }, + Source { org_id: Option }, + ReadTeam, + AlertTeam, + ReadWindows, + ReadWindow { index: usize }, + ReadMembership, + ReadMemberLimit, + ReadDefault { id: String }, + ReadMember { limit: Option }, + AlertMember, + ReadOrg { id: String }, + AlertOrg { id: String }, +} + +pub fn decide(facts: &Facts) -> Result { + let complete = |allowed, log_failure| { + Ok(Decision::Complete { + allowed, + log_failure, + }) + }; + let org_id = facts + .token_org + .as_ref() + .filter(|id| !id.is_empty()) + .or(facts.team_org.as_ref()); + if !facts.source_ready { + return Ok(Decision::Source { + org_id: org_id.cloned(), + }); + } + if facts.io_failed { + return complete(false, false); + } + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeam); + }; + let team = common_budgets::check(team); + if team.read_spend { + return Err(Error::InvalidAnswer); + } + if team.alert.is_some() && !facts.team_alerted { + return Ok(Decision::AlertTeam); + } + if team.denied { + return complete(false, false); + } + let Some(count) = facts.window_count else { + return Ok(Decision::ReadWindows); + }; + for index in 0..count { + let Some(window) = facts.windows.get(index) else { + return Ok(Decision::ReadWindow { index }); + }; + let budget = common_budgets::check(window); + if budget.read_spend { + return Err(Error::InvalidAnswer); + } + if budget.denied { + return complete(false, false); + } + } + if facts.check_db_only && !facts.user_present { + return complete(false, false); + } + if facts.user_present + && (facts.team_id_present || facts.check_db_only) + && !facts.membership_loaded + { + return Ok(Decision::ReadMembership); + } + if facts.user_present && facts.team_id_present { + let Some(limit) = &facts.member_limit else { + return Ok(Decision::ReadMemberLimit); + }; + if limit.member_limit.is_none() + && !facts.default_loaded + && let Some(id) = &facts.default_id + { + return Ok(Decision::ReadDefault { id: id.clone() }); + } + let limit = match common_budgets::member_limit(limit) { + Ok(limit) => limit, + Err(Error::IncompatibleDateTimes) => return complete(false, true), + Err(error) => return Err(error), + }; + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember { limit }); + }; + let member = common_budgets::check(member); + if member.read_spend { + return Err(Error::InvalidAnswer); + } + if member.alert.is_some() && !facts.member_alerted { + return Ok(Decision::AlertMember); + } + if member.denied { + return complete(false, false); + } + } + if let Some(id) = org_id { + if !facts.org_loaded { + return Ok(Decision::ReadOrg { id: id.clone() }); + } + if !facts.org_failed + && let Some(org) = &facts.org + { + let org = common_budgets::check(org); + if org.read_spend { + return Err(Error::InvalidAnswer); + } + if org.alert.is_some() && !facts.org_alerted { + return Ok(Decision::AlertOrg { id: id.clone() }); + } + if org.denied { + return complete(false, false); + } + } + } + complete(true, false) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::token(Some("token"), Some("team"), Some("token"))] + #[case::empty_token(Some(""), Some("team"), Some("team"))] + #[case::empty_team(None, Some(""), Some(""))] + #[case::absent(None, None, None)] + fn source_org( + #[case] token: Option<&str>, + #[case] team: Option<&str>, + #[case] selected: Option<&str>, + ) { + let facts = Facts { + token_org: token.map(str::to_owned), + team_org: team.map(str::to_owned), + ..Facts::default() + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Source { + org_id: selected.map(str::to_owned) + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs b/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs new file mode 100644 index 00000000000..98c66742bbe --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs @@ -0,0 +1,81 @@ +use crate::{Error, personal_model_request::Loaded, personal_models, personal_owner}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub owner: personal_owner::Facts, + pub names: Option>, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadOwner, + ReadNames, + ReadModel { index: usize, teams: Vec }, + Complete { selected: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let teams = match personal_owner::decide(&facts.owner)? { + personal_owner::Decision::Complete { selected, .. } => selected, + _ => return Ok(Decision::ReadOwner), + }; + if facts + .owner + .teams + .as_ref() + .is_none_or(|teams| teams.loads.is_empty()) + { + return Ok(Decision::Complete { selected: vec![] }); + } + let Some(names) = &facts.names else { + return Ok(Decision::ReadNames); + }; + let mut selected = vec![]; + for index in 0..names.len() { + let Some(loaded) = facts.models.get(index) else { + return Ok(Decision::ReadModel { index, teams }); + }; + match personal_models::decide(personal_models::Request::Loaded { + present: loaded.present.clone(), + scopes: loaded.scopes.clone(), + })? { + personal_models::Decision::Complete { allowed } => { + if allowed { + selected.push(index); + } + } + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Complete { selected }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing_database(false, Decision::Complete { selected: vec![] })] + #[case::load(true, Decision::ReadOwner)] + fn owner_before_names(#[case] database_present: bool, #[case] expected: Decision) { + let facts = Facts { + owner: personal_owner::Facts { + user_present: true, + database_present, + cache_present: true, + logging_present: true, + owner_loaded: false, + owner_present: false, + owner_failed: false, + teams: None, + }, + names: None, + models: vec![], + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_model_request.rs b/litellm-rust/crates/request-auth-core/src/personal_model_request.rs new file mode 100644 index 00000000000..9db1e2a03cb --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_model_request.rs @@ -0,0 +1,86 @@ +use crate::{Error, owner_teams, personal_models}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Loaded { + pub present: Vec, + pub scopes: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub owner_present: bool, + pub teams: Option, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { model: String }, + ReadTeams, + ReadModel { index: usize, teams: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let teams = if facts.owner_present { + let Some(teams) = &facts.teams else { + return Ok(Decision::ReadTeams); + }; + match owner_teams::decide(teams)? { + owner_teams::Decision::Complete { selected, .. } => selected, + _ => return Err(Error::InvalidAnswer), + } + } else { + vec![] + }; + for (index, model) in facts.requested.iter().enumerate() { + if !facts.owner_present { + return Ok(Decision::Deny { + model: model.clone(), + }); + } + let Some(loaded) = facts.models.get(index) else { + return Ok(Decision::ReadModel { index, teams }); + }; + match personal_models::decide(personal_models::Request::Loaded { + present: loaded.present.clone(), + scopes: loaded.scopes.clone(), + })? { + personal_models::Decision::Complete { allowed: true } => (), + personal_models::Decision::Complete { allowed: false } => { + return Ok(Decision::Deny { + model: model.clone(), + }); + } + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Allow) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty(false, vec![], Decision::Allow)] + #[case::owner_lookup(true, vec![], Decision::ReadTeams)] + #[case::absent(false, vec!["model".into()], Decision::Deny { model: "model".into() })] + fn owner_and_empty_requests( + #[case] owner_present: bool, + #[case] requested: Vec, + #[case] expected: Decision, + ) { + let facts = Facts { + requested, + owner_present, + teams: None, + models: vec![], + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_models.rs b/litellm-rust/crates/request-auth-core/src/personal_models.rs new file mode 100644 index 00000000000..1fce2d6bd88 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_models.rs @@ -0,0 +1,120 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + model_access, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Scope { + pub user_present: bool, + pub team_id_present: bool, + pub team: Option, + pub member_grants: Option>, + pub member: Option, + pub io_failed: bool, +} + +#[derive(Deserialize)] +pub enum Request { + Scope(Box), + Loaded { + present: Vec, + scopes: Vec, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + ReadTeam, + ReadGroups, + ReadMembership, + ReadMember, + ReadScopes { indices: Vec }, +} + +fn scope(facts: Scope) -> Result { + if facts.io_failed { + return Ok(Decision::Complete { allowed: false }); + } + let Some(team) = facts.team else { + return Ok(Decision::ReadTeam); + }; + let team = fallback_access::grouped(team, true)?; + if matches!( + team, + model_access::Decision::ReadGroups | model_access::Decision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadGroups); + } + if !team.allowed() { + return Ok(Decision::Complete { allowed: false }); + } + if !facts.user_present || !facts.team_id_present { + return Ok(Decision::Complete { allowed: true }); + } + let Some(grants) = facts.member_grants else { + return Ok(Decision::ReadMembership); + }; + if grants.is_empty() { + return Ok(Decision::Complete { allowed: true }); + } + let Some(member) = facts.member else { + return Ok(Decision::ReadMember); + }; + Ok(Decision::Complete { + allowed: model_access::check(&member)?.allowed(), + }) +} + +pub fn decide(request: Request) -> Result { + match request { + Request::Scope(facts) => scope(*facts), + Request::Loaded { present, scopes } => { + let indices: Vec<_> = present + .iter() + .enumerate() + .filter_map(|(index, present)| present.then_some(index)) + .collect(); + if scopes.is_empty() && !indices.is_empty() { + return Ok(Decision::ReadScopes { indices }); + } + if scopes.len() != indices.len() { + return Err(Error::InvalidAnswer); + } + let results: Result, _> = scopes + .into_iter() + .map(|facts| match scope(facts)? { + Decision::Complete { allowed } => Ok(allowed), + _ => Err(Error::InvalidAnswer), + }) + .collect(); + Ok(Decision::Complete { + allowed: results?.into_iter().any(|allowed| allowed), + }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::none(vec![], Decision::Complete { allowed: false })] + #[case::missing(vec![false], Decision::Complete { allowed: false })] + #[case::read_all(vec![true, false, true], Decision::ReadScopes { indices: vec![0, 2] })] + fn all_team_reads(#[case] present: Vec, #[case] expected: Decision) { + assert_eq!( + decide(Request::Loaded { + present, + scopes: vec![] + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_owner.rs b/litellm-rust/crates/request-auth-core/src/personal_owner.rs new file mode 100644 index 00000000000..bc821d2d509 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_owner.rs @@ -0,0 +1,87 @@ +use crate::{Error, owner_teams}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub database_present: bool, + pub cache_present: bool, + pub logging_present: bool, + pub owner_loaded: bool, + pub owner_present: bool, + pub owner_failed: bool, + pub teams: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadOwner, + ReadTeams, + Complete { + selected: Vec, + withheld: Vec, + unreadable: bool, + }, +} + +pub fn decide(facts: &Facts) -> Result { + let empty = || { + Ok(Decision::Complete { + selected: vec![], + withheld: vec![], + unreadable: false, + }) + }; + if !facts.user_present + || !facts.database_present + || !facts.cache_present + || !facts.logging_present + { + return empty(); + } + if !facts.owner_loaded { + return Ok(Decision::ReadOwner); + } + if facts.owner_failed || !facts.owner_present { + return empty(); + } + let Some(teams) = &facts.teams else { + return Ok(Decision::ReadTeams); + }; + match owner_teams::decide(teams)? { + owner_teams::Decision::Complete { + selected, + withheld, + unreadable, + } => Ok(Decision::Complete { + selected, + withheld, + unreadable, + }), + _ => Err(Error::InvalidAnswer), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::no_database(false, Decision::Complete { selected: vec![], withheld: vec![], unreadable: false })] + #[case::lookup(true, Decision::ReadOwner)] + fn prerequisites(#[case] database_present: bool, #[case] expected: Decision) { + let facts = Facts { + user_present: true, + database_present, + cache_present: true, + logging_present: true, + owner_loaded: false, + owner_present: false, + owner_failed: false, + teams: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/protocol.rs b/litellm-rust/crates/request-auth-core/src/protocol.rs index c5a139f4402..96169a7be06 100644 --- a/litellm-rust/crates/request-auth-core/src/protocol.rs +++ b/litellm-rust/crates/request-auth-core/src/protocol.rs @@ -63,6 +63,9 @@ pub enum Grant { #[derive(Clone, veil::Redact, Serialize)] #[serde(tag = "kind")] pub enum Denial { + ModelAccessDenied { + decision: crate::model_access::Decision, + }, MissingKey, MalformedKey, UserAuthDisabled, @@ -108,6 +111,30 @@ pub struct ModelBudgetExceeded { pub entity_id: Option, } +#[derive(Clone, Debug, Deserialize)] +pub struct ModelBudgetRead { + pub snapshot: crate::model_budget_data::Snapshot, + pub message: String, + pub entity_type: Option, + pub entity_id: Option, +} + +impl ModelBudgetRead { + pub fn exceeded(self) -> Option { + let result = crate::model_budget_data::check(&self.snapshot); + if !result.denied { + return None; + } + Some(ModelBudgetExceeded { + message: self.message, + current_cost: result.current_spend, + max_budget: self.snapshot.max_budget?, + entity_type: self.entity_type, + entity_id: self.entity_id, + }) + } +} + impl From for Denial { fn from(value: ModelBudgetExceeded) -> Self { Self::ModelBudgetExceeded { @@ -196,6 +223,7 @@ pub enum Effect { }, InspectKey, EnforceModelAccess, + PrefetchAuthObjects, LoadKeyOwner { user_id: String, }, @@ -267,8 +295,12 @@ pub enum Node { #[derive(Deserialize)] #[serde(tag = "kind", content = "value")] pub enum Answer { + ModelAccess(Box), + FallbackAccess(Box), + Candidate(Box), Unit, Bool(bool), + Exemption(crate::budget_pricing::ExemptionFacts), Presented(Option), Key(KeyFacts), Owner(Option), @@ -277,6 +309,7 @@ pub enum Answer { ProxySpend(Option), ModelScope(ModelBudgetScope), Exceeded(Option), + ModelBudgetData(ModelBudgetRead), TeamScope(TeamMemberModelBudgetScope), Targets(Vec), } @@ -298,6 +331,7 @@ macro_rules! answer { }; } answer!(bool, Bool); +answer!(crate::budget_pricing::ExemptionFacts, Exemption); answer!(Option, Presented); answer!(KeyFacts, Key); answer!(Option, Owner); @@ -306,9 +340,37 @@ answer!(f64, Spend); answer!(Option, ProxySpend); answer!(ModelBudgetScope, ModelScope); answer!(Option, Exceeded); +answer!(ModelBudgetRead, ModelBudgetData); answer!(TeamMemberModelBudgetScope, TeamScope); answer!(Vec, Targets); +impl FromAnswer for crate::fallback_candidate::Facts { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::Candidate(value) => Ok(*value), + _ => Err(Error::InvalidAnswer), + } + } +} + +impl FromAnswer for crate::key_model_access::Facts { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::ModelAccess(value) => Ok(*value), + _ => Err(Error::InvalidAnswer), + } + } +} + +impl FromAnswer for crate::fallback_access::Facts { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::FallbackAccess(value) => Ok(*value), + _ => Err(Error::InvalidAnswer), + } + } +} + impl FromAnswer for () { fn from_answer(answer: Answer) -> Result { match answer { diff --git a/litellm-rust/crates/request-auth-core/src/reservation.rs b/litellm-rust/crates/request-auth-core/src/reservation.rs new file mode 100644 index 00000000000..185cfd19bfc --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/reservation.rs @@ -0,0 +1,312 @@ +use serde::{Deserialize, Serialize}; +use std::cmp::Ordering::Greater; + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Gate { + token_present: bool, + llm_route: Option, + unbilled: Option, + model_present: Option, + counter_count: Option, + cost_read: bool, + cost: Option, + }, + Sources { + key_present: bool, + team_present: bool, + team_id: Option, + user_present: bool, + user_id: Option, + apply_user_budget_to_team_keys: bool, + }, + Counter { + scope: Scope, + limit: Option, + }, + Window { + duration_present: bool, + limit: Option, + start_present: Option, + }, + Identity { + scope: IdentityScope, + primary: Option, + fallback: Option, + }, + MemberGate { + team_present: bool, + team_id: Option, + user_present: bool, + user_id: Option, + }, + MemberCap { + member_limit: Option, + default_read: bool, + }, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum IdentityScope { + EndUser, + Organization, + Project, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum Scope { + Key, + Team, + User, + EndUser, + Tag, + AccessGroup, + Member, + Organization, + Project, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadRoute, + ReadUnbilled, + ReadModel, + ReadCounters, + ReadEstimate, + Reserve, + Skip, + ReadStart, + WarnWindow, + AppendWindow, + ReadMember, + ReadDefault, + ReadCap, + Identity { id: Option }, + Sources { key: bool, team: bool, user: bool }, + Counter { include: bool }, +} + +pub fn decide(facts: Facts) -> Decision { + match facts { + Facts::Identity { + scope, + primary, + fallback, + } => Decision::Identity { + id: match scope { + IdentityScope::EndUser => primary.filter(|id| !id.is_empty()).or(fallback), + IdentityScope::Organization => primary.or(fallback), + IdentityScope::Project => primary, + }, + }, + Facts::MemberGate { + team_present, + team_id, + user_present, + user_id, + } => { + if team_present && team_id.is_some() && user_present && user_id.is_some() { + Decision::ReadMember + } else { + Decision::Skip + } + } + Facts::MemberCap { + member_limit, + default_read, + } => { + if member_limit.is_none() && !default_read { + Decision::ReadDefault + } else { + Decision::ReadCap + } + } + Facts::Gate { + token_present, + llm_route, + unbilled, + model_present, + counter_count, + cost_read, + cost, + } => { + if !token_present { + return Decision::Skip; + } + match llm_route { + None => return Decision::ReadRoute, + Some(false) => return Decision::Skip, + Some(true) => {} + } + match unbilled { + None => return Decision::ReadUnbilled, + Some(true) => return Decision::Skip, + Some(false) => {} + } + match model_present { + None => return Decision::ReadModel, + Some(false) => return Decision::Skip, + Some(true) => {} + } + match counter_count { + None => return Decision::ReadCounters, + Some(0) => return Decision::Skip, + Some(_) => {} + } + if !cost_read { + return Decision::ReadEstimate; + } + match cost { + None => Decision::Skip, + Some(value) if value <= 0.0 => Decision::Skip, + Some(_) => Decision::Reserve, + } + } + Facts::Sources { + key_present, + team_present, + team_id, + user_present, + user_id, + apply_user_budget_to_team_keys, + } => { + let team = team_present && team_id.is_some(); + Decision::Sources { + key: key_present, + team, + user: (!team || apply_user_budget_to_team_keys) + && user_present + && user_id.is_some(), + } + } + Facts::Counter { scope, limit } => Decision::Counter { + include: limit.is_some_and(|limit| match scope { + Scope::Key | Scope::Team | Scope::User => limit.partial_cmp(&0.0) == Some(Greater), + Scope::Project => limit.is_finite() && limit > 0.0, + _ => matches!(limit.partial_cmp(&0.0), None | Some(Greater)), + }), + }, + Facts::Window { + duration_present, + limit, + start_present, + } => { + if !duration_present || limit.is_none_or(|limit| limit <= 0.0) { + return Decision::Skip; + } + match start_present { + None => Decision::ReadStart, + Some(false) => Decision::WarnWindow, + Some(true) => Decision::AppendWindow, + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::customer_empty(IdentityScope::EndUser, Some(""), Some("fallback"), Some("fallback"))] + #[case::org_empty(IdentityScope::Organization, Some(""), Some("fallback"), Some(""))] + #[case::org_fallback(IdentityScope::Organization, None, Some("fallback"), Some("fallback"))] + #[case::project_no_fallback(IdentityScope::Project, None, Some("fallback"), None)] + fn identities( + #[case] scope: IdentityScope, + #[case] primary: Option<&str>, + #[case] fallback: Option<&str>, + #[case] expected: Option<&str>, + ) { + assert_eq!( + decide(Facts::Identity { + scope, + primary: primary.map(str::to_owned), + fallback: fallback.map(str::to_owned) + }), + Decision::Identity { + id: expected.map(str::to_owned) + } + ); + } + + #[rstest] + #[case::member_cap(Some(0.0), false, Decision::ReadCap)] + #[case::default_needed(None, false, Decision::ReadDefault)] + #[case::default_read(None, true, Decision::ReadCap)] + fn member_cap( + #[case] member_limit: Option, + #[case] default_read: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::MemberCap { + member_limit, + default_read + }), + expected + ); + } + + #[rstest] + #[case::all_present(true, Some(""), true, Some(""), Decision::ReadMember)] + #[case::missing_team(true, None, true, Some("user"), Decision::Skip)] + #[case::missing_user(true, Some("team"), false, Some("user"), Decision::Skip)] + fn member_gate( + #[case] team_present: bool, + #[case] team_id: Option<&str>, + #[case] user_present: bool, + #[case] user_id: Option<&str>, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::MemberGate { + team_present, + team_id: team_id.map(str::to_owned), + user_present, + user_id: user_id.map(str::to_owned) + }), + expected + ); + } + + #[rstest] + #[case::key_nan(Scope::Key, f64::NAN, false)] + #[case::customer_nan(Scope::EndUser, f64::NAN, true)] + #[case::project_infinite(Scope::Project, f64::INFINITY, false)] + #[case::team_infinite(Scope::Team, f64::INFINITY, true)] + #[case::zero(Scope::Organization, 0.0, false)] + #[case::positive(Scope::User, 10.0, true)] + fn limits(#[case] scope: Scope, #[case] limit: f64, #[case] include: bool) { + assert_eq!( + decide(Facts::Counter { + scope, + limit: Some(limit) + }), + Decision::Counter { include } + ); + } + + #[rstest] + #[case::no_duration(false, None, Decision::Skip)] + #[case::unread(true, None, Decision::ReadStart)] + #[case::invalid(true, Some(false), Decision::WarnWindow)] + #[case::valid(true, Some(true), Decision::AppendWindow)] + fn windows( + #[case] duration_present: bool, + #[case] start_present: Option, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::Window { + duration_present, + limit: Some(1.0), + start_present + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/resolved_models.rs b/litellm-rust/crates/request-auth-core/src/resolved_models.rs new file mode 100644 index 00000000000..d048a49745d --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/resolved_models.rs @@ -0,0 +1,219 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + member_models, + model_access::{self, Decision as ModelDecision}, + personal_model_request, team_override, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct TeamRead { + pub found: bool, + pub failed: bool, +} + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub model: String, + pub config_nonempty: bool, + pub key_models: Vec, + pub team_id: Option, + pub user_id: Option, + pub project_present: bool, + pub customer_present: bool, + pub database_present: bool, + pub personal_enabled: bool, + pub master: bool, + pub role: Option, + pub key: Option, + pub team_read: Option, + pub embedded_ready: bool, + pub team: Option, + pub member: Option, + pub owner_loaded: bool, + pub owner_present: bool, + pub owner_failed: bool, + pub personal: Option, + pub project_loaded: bool, + pub project_grants: Vec, + pub project: Option, + pub customer_loaded: bool, + pub customer_grants: Vec, + pub customer: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { + denial: ModelDecision, + member_models: Option>, + }, + ReadKey, + ReadKeyGroups, + ReadTeamObject, + ReadEmbeddedTeam, + ReadTeamScope, + ReadMember, + ReadOwner, + ReadPersonal, + ReadProjectObject, + ReadProjectScope, + ReadCustomerObject, + ReadCustomerScope, +} + +pub fn decide(facts: &Facts) -> Result { + let deny = |denial| { + Ok(Decision::Deny { + denial, + member_models: None, + }) + }; + if !facts.config_nonempty + && !facts + .key_models + .iter() + .any(|model| model == "all-team-models") + { + let Some(key) = &facts.key else { + return Ok(Decision::ReadKey); + }; + let key = fallback_access::grouped(key.clone(), false)?; + if matches!( + key, + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } + ) { + return Ok(Decision::ReadKeyGroups); + } + if !key.allowed() { + return deny(key); + } + } + if facts.team_id.is_some() { + let Some(read) = &facts.team_read else { + return Ok(Decision::ReadTeamObject); + }; + if read.failed && !facts.embedded_ready { + return Ok(Decision::ReadEmbeddedTeam); + } + if read.found || read.failed { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeamScope); + }; + match team_override::decide(team)? { + team_override::Decision::Allow => (), + team_override::Decision::Deny { denial } => return deny(denial), + _ => return Err(Error::InvalidAnswer), + } + if read.found && facts.user_id.is_some() { + let Some(member) = &facts.member else { + return Ok(Decision::ReadMember); + }; + match member_models::decide(member)? { + member_models::Decision::Allow => (), + member_models::Decision::Deny { + denial, + member_models, + } => { + return Ok(Decision::Deny { + denial, + member_models, + }); + } + _ => return Err(Error::InvalidAnswer), + } + } + } + } + let model_access::Response::Predicate(personal) = + model_access::decide(model_access::Request::PersonalApplies { + enabled: facts.personal_enabled, + token_present: true, + team_id: facts.team_id.clone(), + user_id: facts.user_id.clone(), + master: facts.master, + role: facts.role.clone(), + api_key_present: true, + inheritance: false, + })? + else { + return Err(Error::InvalidAnswer); + }; + if personal { + if !facts.owner_loaded { + return Ok(Decision::ReadOwner); + } + if facts.owner_failed || !facts.owner_present { + return deny(ModelDecision::PersonalTeam { + model: facts.model.clone(), + }); + } + let Some(personal) = &facts.personal else { + return Ok(Decision::ReadPersonal); + }; + match personal_model_request::decide(personal)? { + personal_model_request::Decision::Allow => (), + personal_model_request::Decision::Deny { model } => { + return deny(ModelDecision::PersonalTeam { model }); + } + _ => return Err(Error::InvalidAnswer), + } + } + if facts.project_present { + if !facts.project_loaded { + return Ok(Decision::ReadProjectObject); + } + if !facts.project_grants.is_empty() { + let Some(project) = &facts.project else { + return Ok(Decision::ReadProjectScope); + }; + let decision = model_access::check(project)?; + if !decision.allowed() { + return deny(decision); + } + } + } + if facts.customer_present && facts.database_present { + if !facts.customer_loaded { + return Ok(Decision::ReadCustomerObject); + } + if !facts.customer_grants.is_empty() { + let Some(customer) = &facts.customer else { + return Ok(Decision::ReadCustomerScope); + }; + let model_access::Response::Decision(decision) = + model_access::decide(model_access::Request::Customer(customer.clone()))? + else { + return Err(Error::InvalidAnswer); + }; + if !decision.allowed() { + return deny(decision); + } + } + } + Ok(Decision::Allow) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::key(false, Decision::ReadKey)] + #[case::team(true, Decision::ReadTeamObject)] + fn key_before_team(#[case] config_nonempty: bool, #[case] expected: Decision) { + assert_eq!( + decide(&Facts { + config_nonempty, + team_id: Some("team".into()), + ..Facts::default() + }) + .unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_access.rs b/litellm-rust/crates/request-auth-core/src/route_access.rs new file mode 100644 index 00000000000..b1eb51177bf --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_access.rs @@ -0,0 +1,567 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive(Clone, Default, Deserialize)] +pub struct RouteCatalog { + pub groups: BTreeMap>, + pub viewer_blocked: Vec, + pub viewer_write_blocked: Vec, + pub self_credentials: Vec<(String, String)>, + pub managed_models: Vec, + pub managed_paths: Vec, + pub managed_realtime: Vec, + pub managed_live: Vec, + pub managed_mcp: Vec, + pub unbilled: Vec, +} + +impl RouteCatalog { + fn group(&self, name: &str) -> &[String] { + self.groups.get(name).map_or(&[], Vec::as_slice) + } + + fn access(&self, route: &str, name: &str) -> Result { + patterns::route_access(route, self.group(name)) + } + + fn contains(&self, route: &str, name: &str) -> bool { + self.group(name).iter().any(|entry| entry == route) + } +} + +#[derive(Clone, Default, Deserialize)] +pub struct RouteContext { + pub route: String, + pub method: Option, + pub catalog: RouteCatalog, + pub registered: bool, + pub registered_auth: Option, + pub configured_eu_jev: bool, +} + +#[derive(Clone, Default, Deserialize)] +pub struct PassThroughFacts { + pub key_allowed: Vec, + pub team_allowed: Vec, + pub denied: Vec, + pub forwarded: Vec, + pub jwt_team_routes: Vec, +} + +#[derive(Deserialize)] +pub struct VirtualRouteFacts { + pub context: RouteContext, + pub allowed_routes: Option>, + pub passthrough: PassThroughFacts, + pub password_reset_required: bool, +} + +#[derive(Deserialize)] +pub struct ViewerFacts { + pub context: RouteContext, + pub request_keys: Vec, +} + +#[derive(Deserialize)] +pub struct NonAdminFacts { + pub viewer: ViewerFacts, + pub role: Option, + pub custom_admin_routes: Vec, + pub premium: bool, + pub query_user_id: Option, + pub token_user_id: Option, + pub permissions: Vec, + pub organization_ids: Vec, + pub admin_organization_ids: Vec, + pub team_service_account: bool, + pub allowed_routes: Option>, + pub passthrough: PassThroughFacts, +} + +#[derive(Deserialize)] +pub struct MatchFacts { + pub route: String, + pub patterns: Vec, + pub mode: MatchMode, +} + +#[derive(Deserialize, strum::VariantArray)] +pub enum MatchMode { + Access, + Template, + Prefix, + Wildcard, + Denied, + Jwt, +} + +#[derive(Deserialize, strum::VariantArray)] +pub enum Classification { + Llm, + Unbilled, + Discovery, + SelfCredential, + Azure, + Managed, +} + +#[derive(Deserialize)] +pub enum Request { + Match(MatchFacts), + Classify { + context: RouteContext, + kind: Classification, + }, + Virtual(VirtualRouteFacts), + Viewer(ViewerFacts), + NonAdmin(NonAdminFacts), + PassThrough { + context: RouteContext, + facts: PassThroughFacts, + required: bool, + }, + DeniedPassThrough { + route: String, + facts: PassThroughFacts, + }, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Decision { + pub allowed: bool, + pub denial: Option, + pub detail: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Denial { + VirtualRoute, + PasswordReset, + PassThrough, + DeniedPassThrough, + CustomAdmin, + UserInfo, + AdminOnly, + ViewerInference, + ViewerWrite, + ViewerParameter, +} + +impl Decision { + fn allow() -> Self { + Self { + allowed: true, + denial: None, + detail: None, + } + } + fn predicate(value: bool) -> Self { + Self { + allowed: value, + denial: None, + detail: None, + } + } + fn deny(denial: Denial) -> Self { + Self { + allowed: false, + denial: Some(denial), + detail: None, + } + } + fn detail(denial: Denial, detail: String) -> Self { + Self { + allowed: false, + denial: Some(denial), + detail: Some(detail), + } + } +} + +fn denied_pattern(route: &str, pattern: &str) -> bool { + let normalized = pattern.trim_end_matches('/'); + normalized.is_empty() + || patterns::prefix(route, normalized) + || patterns::wildcard(route, pattern) +} + +fn jwt_grant(route: &str, entries: &[String]) -> bool { + entries.iter().any(|entry| { + !entry.trim_end_matches('*').trim_matches('/').is_empty() + && patterns::wildcard(route, entry) + }) +} + +fn passthrough_grant(route: &str, facts: &PassThroughFacts) -> bool { + let entries = if facts.key_allowed.is_empty() { + &facts.team_allowed + } else { + &facts.key_allowed + }; + entries.iter().any(|entry| patterns::prefix(route, entry)) +} + +fn denied_passthrough(facts: &PassThroughFacts) -> Option<&str> { + facts + .denied + .iter() + .find(|denied| { + facts + .forwarded + .iter() + .any(|route| denied_pattern(route, denied)) + }) + .map(String::as_str) +} + +fn require_passthrough(route: &str, facts: &PassThroughFacts) -> Decision { + if let Some(denied) = denied_passthrough(facts) { + return Decision::detail(Denial::DeniedPassThrough, denied.to_owned()); + } + if passthrough_grant(route, facts) || jwt_grant(route, &facts.jwt_team_routes) { + return Decision::allow(); + } + Decision::deny(Denial::PassThrough) +} + +fn azure(route: &str) -> Result { + Ok(patterns::regex_match( + "^/openai/deployments/[^/]+/[^/]+/chat/completions(?=\\n?\\z)", + route, + )? || patterns::regex_match("^/engines/[^/]+/chat/completions(?=\\n?\\z)", route)?) +} + +fn llm(context: &RouteContext) -> Result { + let route = &context.route; + let catalog = &context.catalog; + for group in ["openai_routes", "anthropic_routes", "google_routes"] { + if catalog.contains(route, group) { + return Ok(true); + } + } + for group in [ + "mcp_inference_routes", + "agent_inference_routes", + "litellm_native_routes", + ] { + if catalog.access(route, group)? { + return Ok(true); + } + } + for group in ["openai_routes", "google_routes", "anthropic_routes"] { + for pattern in catalog.group(group) { + if pattern.contains('{') && patterns::route_pattern(route, pattern)? { + return Ok(true); + } + if group == "openai_routes" + && pattern.ends_with('*') + && patterns::wildcard(route, pattern) + { + return Ok(true); + } + } + } + Ok(azure(route)? + || catalog + .group("mapped_pass_through_routes") + .iter() + .any(|pattern| patterns::prefix(route, pattern)) + || (route.trim_end_matches('/') == "/eu_jev/v1/systemone" && !context.configured_eu_jev)) +} + +fn discovery(context: &RouteContext) -> bool { + context.method.as_deref() == Some("GET") + && (context.route == "/v1/mcp/server" + || context + .route + .strip_prefix("/v1/mcp/server/") + .is_some_and(|tail| !tail.is_empty() && !tail.contains('/'))) +} + +fn self_credential(context: &RouteContext) -> Result { + for (method, pattern) in &context.catalog.self_credentials { + if context.method.as_ref() == Some(method) + && patterns::route_access(&context.route, std::slice::from_ref(pattern))? + { + return Ok(true); + } + } + Ok(false) +} + +fn managed(context: &RouteContext) -> Result { + let route = &context.route; + let method = context.method.as_deref(); + let catalog = &context.catalog; + if route == "/agents" || route == "/v1/agents" { + return Ok(matches!(method, None | Some("GET" | "HEAD"))); + } + if catalog.managed_realtime.contains(route) || catalog.managed_live.contains(route) { + return Ok(matches!(method, None | Some("GET"))); + } + if catalog.managed_models.contains(route) + || patterns::route_access(route, &catalog.managed_paths)? + { + return Ok(matches!(method, None | Some("POST"))); + } + Ok(patterns::route_access(route, &catalog.managed_mcp)? + || catalog.access(route, "agent_inference_routes")?) +} + +fn virtual_route(facts: &VirtualRouteFacts) -> Result { + let Some(allowed) = &facts.allowed_routes else { + return Ok(Decision::allow()); + }; + if allowed.is_empty() { + return Ok(Decision::allow()); + } + let context = &facts.context; + let route = &context.route; + if allowed + .iter() + .any(|pattern| patterns::prefix(route, pattern)) + { + return Ok(Decision::allow()); + } + let mut denied_passthrough = false; + for group in allowed { + if !context.catalog.groups.contains_key(group) { + continue; + } + let is_llm = matches!( + group.as_str(), + "llm_api_routes" | "llm_api_routes_without_mcp" + ); + if context.catalog.access(route, group)? { + if (is_llm || group == "openai_routes") && context.registered_auth == Some(true) { + if passthrough_grant(route, &facts.passthrough) { + return Ok(Decision::allow()); + } + denied_passthrough = true; + } else { + return Ok(Decision::allow()); + } + } + if !is_llm { + continue; + } + if route == "/auto_router/session" && context.method.as_deref() == Some("GET") { + return Ok(Decision::allow()); + } + if context.registered { + if context.registered_auth == Some(true) { + if passthrough_grant(route, &facts.passthrough) { + return Ok(Decision::allow()); + } + denied_passthrough = true; + } else { + return Ok(Decision::allow()); + } + } + if group == "llm_api_routes" + && (discovery(context) || context.catalog.access(route, "agent_management_routes")?) + { + return Ok(Decision::allow()); + } + } + if allowed + .iter() + .any(|pattern| patterns::wildcard(route, pattern)) + { + return Ok(Decision::allow()); + } + Ok(Decision::deny(if denied_passthrough { + Denial::PassThrough + } else if facts.password_reset_required { + Denial::PasswordReset + } else { + Denial::VirtualRoute + })) +} + +fn viewer(facts: &ViewerFacts) -> Result { + let context = &facts.context; + let route = &context.route; + let catalog = &context.catalog; + if llm(context)? { + return Ok(Decision::deny(Denial::ViewerInference)); + } + let invalid = facts + .request_keys + .iter() + .find(|key| key.as_str() != "user_email"); + if catalog.access(route, "management_routes")? { + if route == "/user/update" { + if let Some(key) = invalid { + return Ok(Decision::detail(Denial::ViewerParameter, key.clone())); + } + } else if patterns::route_access(route, &catalog.viewer_blocked)? + || (route.starts_with("/key/") + && (route.ends_with("/regenerate") || route.ends_with("/reset_spend"))) + { + return Ok(Decision::deny(Denial::ViewerWrite)); + } + } + if matches!( + context.method.as_deref().unwrap_or("GET"), + "GET" | "HEAD" | "OPTIONS" + ) { + return Ok(Decision::allow()); + } + if route == "/user/update" { + return Ok(invalid.map_or_else(Decision::allow, |key| { + Decision::detail(Denial::ViewerParameter, key.clone()) + })); + } + if route == "/user/password/change" || route == "/session/logout" || self_credential(context)? { + return Ok(Decision::allow()); + } + if patterns::route_access(route, &catalog.viewer_write_blocked)? + || (route.starts_with("/key/") && route.ends_with("/regenerate")) + { + return Ok(Decision::deny(Denial::ViewerWrite)); + } + if catalog.access(route, "admin_viewer_routes")? + || catalog.access(route, "global_spend_tracking_routes")? + { + return Ok(Decision::allow()); + } + Ok(Decision::deny(Denial::ViewerWrite)) +} + +fn non_admin(facts: &NonAdminFacts) -> Result { + let context = &facts.viewer.context; + let route = &context.route; + let catalog = &context.catalog; + let role = facts.role.as_deref(); + if facts.premium && facts.custom_admin_routes.contains(route) { + return Ok(Decision::deny(Denial::CustomAdmin)); + } + if context.registered_auth == Some(true) { + return Ok(require_passthrough(route, &facts.passthrough)); + } + if llm(context)? { + return Ok(Decision::allow()); + } + if catalog.contains(route, "info_routes") { + if route == "/user/info" + && facts + .query_user_id + .as_ref() + .is_some_and(|id| !id.is_empty()) + && facts.query_user_id != facts.token_user_id + && role != Some("proxy_admin_viewer") + { + return Ok(Decision::deny(Denial::UserInfo)); + } + return Ok(Decision::allow()); + } + if catalog.contains(route, "global_spend_tracking_routes") + && facts + .permissions + .iter() + .any(|entry| entry == "get_spend_routes") + { + return Ok(Decision::allow()); + } + if patterns::prefix(route, "/lens") { + return Ok(Decision::allow()); + } + if role == Some("proxy_admin_viewer") { + return viewer(&facts.viewer); + } + let org_admin = !facts.organization_ids.is_empty() + && facts + .organization_ids + .iter() + .all(|id| facts.admin_organization_ids.contains(id)); + if (role == Some("internal_user") + && (catalog.access(route, "internal_user_routes")? || self_credential(context)?)) + || (org_admin && catalog.access(route, "org_admin_allowed_routes")?) + || (role == Some("internal_user_viewer") + && catalog.access(route, "internal_user_view_only_routes")?) + || catalog.access(route, "self_managed_routes")? + || route.starts_with("/v1/mcp/") + || route.starts_with("/mcp-rest/") + || passthrough_grant(route, &facts.passthrough) + || (facts.team_service_account + && catalog.access(route, "team_service_account_key_routes")?) + { + return Ok(Decision::allow()); + } + if facts.allowed_routes.as_ref().is_some_and(|entries| { + entries + .iter() + .any(|entry| patterns::prefix(route, entry) || patterns::wildcard(route, entry)) + }) { + return Ok(Decision::allow()); + } + Ok(Decision::deny(Denial::AdminOnly)) +} + +pub fn decide(request: Request) -> Result { + match request { + Request::Virtual(facts) => virtual_route(&facts), + Request::Viewer(facts) => viewer(&facts), + Request::NonAdmin(facts) => non_admin(&facts), + Request::PassThrough { + context, + facts, + required, + } => Ok(if required { + require_passthrough(&context.route, &facts) + } else { + Decision::predicate(passthrough_grant(&context.route, &facts)) + }), + Request::DeniedPassThrough { route: _, facts } => Ok(denied_passthrough(&facts) + .map_or_else(Decision::allow, |entry| { + Decision::detail(Denial::DeniedPassThrough, entry.to_owned()) + })), + Request::Match(facts) => { + let matched = match facts.mode { + MatchMode::Access => patterns::route_access(&facts.route, &facts.patterns)?, + MatchMode::Jwt => jwt_grant(&facts.route, &facts.patterns), + MatchMode::Template => { + let mut found = false; + for pattern in &facts.patterns { + if patterns::route_pattern(&facts.route, pattern)? { + found = true; + break; + } + } + found + } + MatchMode::Prefix => facts + .patterns + .iter() + .any(|pattern| patterns::prefix(&facts.route, pattern)), + MatchMode::Wildcard => facts + .patterns + .iter() + .any(|pattern| patterns::wildcard(&facts.route, pattern)), + MatchMode::Denied => facts + .patterns + .iter() + .any(|pattern| denied_pattern(&facts.route, pattern)), + }; + Ok(Decision::predicate(matched)) + } + Request::Classify { context, kind } => Ok(Decision::predicate(match kind { + Classification::Llm => llm(&context)?, + Classification::Azure => azure(&context.route)?, + Classification::Discovery => discovery(&context), + Classification::SelfCredential => self_credential(&context)?, + Classification::Managed => managed(&context)?, + Classification::Unbilled => { + let tail = context.route.rsplit('/').next().unwrap_or(""); + let (resource, action) = tail.split_once(':').unwrap_or((tail, "")); + context.catalog.unbilled.contains(&context.route) + || matches!(resource, "count_tokens" | "count-tokens") + || action == "countTokens" + } + })), + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_enable.rs b/litellm-rust/crates/request-auth-core/src/route_enable.rs new file mode 100644 index 00000000000..4f4a42294d2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_enable.rs @@ -0,0 +1,180 @@ +use crate::{ + Error, patterns, + route_access::{self, Classification, RouteContext}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum Feature { + Admin, + Llm, +} + +#[derive(Clone, Deserialize)] +pub struct Setting { + pub environment_present: bool, + pub premium: bool, + pub secret: Option, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub context: RouteContext, + pub imported: bool, + pub failed: Option, + pub admin: Option, + pub llm: Option, + pub exemptions: Option>, + pub eu_jev_configured: Option, +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum Failure { + Http, + Other, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Reason { + License, + Disabled, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + PropagateHttp, + ReadImport, + ReadExemptions, + ReadEuJev, + ReadSettings { feature: Feature }, + ReadSecret { feature: Feature }, + Deny { feature: Feature, reason: Reason }, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum LlmDecision { + Allow, + Deny, + ReadEuJev, +} + +pub fn llm_route(context: &RouteContext, configured: Option) -> Result { + let static_context = RouteContext { + configured_eu_jev: true, + ..context.clone() + }; + if route_access::decide(route_access::Request::Classify { + context: static_context, + kind: Classification::Llm, + })? + .allowed + { + return Ok(LlmDecision::Allow); + } + if context.route.trim_end_matches('/') != "/eu_jev/v1/systemone" { + return Ok(LlmDecision::Deny); + } + Ok(match configured { + None => LlmDecision::ReadEuJev, + Some(true) => LlmDecision::Deny, + Some(false) => LlmDecision::Allow, + }) +} + +fn setting(setting: Option<&Setting>, feature: Feature) -> Option { + let Some(setting) = setting else { + return Some(Decision::ReadSettings { feature }); + }; + if setting.environment_present && !setting.premium { + return Some(Decision::Deny { + feature, + reason: Reason::License, + }); + } + match setting.secret { + None => Some(Decision::ReadSecret { feature }), + Some(true) => Some(Decision::Deny { + feature, + reason: Reason::Disabled, + }), + Some(false) => None, + } +} + +pub fn decide(facts: &Facts) -> Result { + if let Some(failed) = facts.failed { + return Ok(match failed { + Failure::Http => Decision::PropagateHttp, + Failure::Other => Decision::Allow, + }); + } + if !facts.imported { + return Ok(Decision::ReadImport); + } + let management = facts + .context + .catalog + .groups + .get("management_routes") + .map(Vec::as_slice) + .unwrap_or_default(); + if patterns::route_access(&facts.context.route, management)? + && let Some(decision) = setting(facts.admin.as_ref(), Feature::Admin) + { + return Ok(decision); + } + match llm_route(&facts.context, facts.eu_jev_configured)? { + LlmDecision::Deny => return Ok(Decision::Allow), + LlmDecision::ReadEuJev => return Ok(Decision::ReadEuJev), + LlmDecision::Allow => (), + } + let Some(exemptions) = &facts.exemptions else { + return Ok(Decision::ReadExemptions); + }; + if exemptions.contains(&facts.context.route) { + return Ok(Decision::Allow); + } + Ok(setting(facts.llm.as_ref(), Feature::Llm).unwrap_or(Decision::Allow)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(None, LlmDecision::ReadEuJev)] + #[case::builtin(Some(false), LlmDecision::Allow)] + #[case::configured(Some(true), LlmDecision::Deny)] + fn builtin_configuration(#[case] configured: Option, #[case] expected: LlmDecision) { + let context = RouteContext { + route: "/eu_jev/v1/systemone/".into(), + ..RouteContext::default() + }; + assert_eq!(llm_route(&context, configured).unwrap(), expected); + } + + #[rstest] + #[case::license(true, false, Decision::Deny { feature: Feature::Admin, reason: Reason::License })] + #[case::absent(false, false, Decision::ReadSecret { feature: Feature::Admin })] + #[case::premium(true, true, Decision::ReadSecret { feature: Feature::Admin })] + fn license_precedes_secret( + #[case] environment_present: bool, + #[case] premium: bool, + #[case] expected: Decision, + ) { + assert_eq!( + setting( + Some(&Setting { + environment_present, + premium, + secret: None + }), + Feature::Admin + ), + Some(expected) + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/route_roles.rs b/litellm-rust/crates/request-auth-core/src/route_roles.rs new file mode 100644 index 00000000000..059fa47075b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/route_roles.rs @@ -0,0 +1,325 @@ +use serde::{Deserialize, Serialize}; + +use crate::{Error, patterns}; + +#[derive( + Clone, Copy, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr, strum::VariantArray, +)] +enum Role { + #[strum(serialize = "proxy_admin")] + Admin, + #[strum(serialize = "proxy_admin_viewer")] + Viewer, + #[strum(serialize = "org_admin")] + Organization, + #[strum(serialize = "internal_user")] + Internal, + #[strum(serialize = "internal_user_viewer")] + InternalViewer, + #[strum(serialize = "team")] + Team, + #[strum(serialize = "customer")] + Customer, +} + +#[derive(Deserialize)] +pub struct RoleFacts { + pub user_present: bool, + pub role: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct RoleDecision { + pub effective: &'static str, + pub user_role: Option<&'static str>, + pub admin: bool, +} + +pub fn role(facts: &RoleFacts) -> RoleDecision { + let parsed = facts + .role + .as_deref() + .and_then(|value| value.parse::().ok()); + let effective = parsed.unwrap_or(Role::Internal).into(); + RoleDecision { + effective, + user_role: facts.user_present.then_some(effective), + admin: facts.user_present && parsed == Some(Role::Admin), + } +} + +#[derive(Deserialize)] +pub struct ApiFacts { + pub token_present: bool, + pub admin_read: bool, + pub user: RoleFacts, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum ApiDecision { + InvalidToken, + ReadAdmin, + NonAdmin, + Allow, +} + +pub fn api(facts: &ApiFacts) -> ApiDecision { + if !facts.token_present { + ApiDecision::InvalidToken + } else if !facts.admin_read { + ApiDecision::ReadAdmin + } else if role(&facts.user).admin { + ApiDecision::Allow + } else { + ApiDecision::NonAdmin + } +} + +#[derive(Deserialize)] +pub struct JwtFacts { + pub role: Option, + pub config_read: bool, + pub config_missing: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum JwtDecision { + ReadAdmin, + ReadTeam, + DefaultTeam, + Check, + Deny, +} + +pub fn jwt(facts: &JwtFacts) -> JwtDecision { + match facts.role.as_deref() { + Some("proxy_admin") if !facts.config_read => JwtDecision::ReadAdmin, + Some("team") if !facts.config_read => JwtDecision::ReadTeam, + Some("team") if facts.config_missing => JwtDecision::DefaultTeam, + Some("proxy_admin" | "team") => JwtDecision::Check, + _ => JwtDecision::Deny, + } +} + +#[derive(Deserialize)] +pub struct InsideFacts { + pub role: Option, + pub user_id: Option, + pub requested_user_id: Option, +} + +pub fn inside(facts: &InsideFacts) -> bool { + matches!( + facts.role.as_deref(), + Some("proxy_admin" | "proxy_admin_viewer") + ) || facts.user_id.is_some() && facts.user_id == facts.requested_user_id +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Entry { + pub text: Option, + pub python_type: String, + pub hash_failure: Option, + pub templates: Option>, +} + +#[derive(Deserialize)] +pub struct ListedFacts { + pub route: String, + pub entry: Option, + pub exhausted: bool, + pub template_index: usize, + pub compiled: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Action { + ReadEntry, + ReadTemplate, + Allow, + Deny, + Keep, + Expand, + Failure, +} + +#[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum FailureKind { + TypeError, + AttributeError, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct Failure { + pub kind: FailureKind, + pub message: String, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +pub struct EntryDecision { + pub action: Action, + pub template_index: usize, + pub failure: Option, +} + +impl EntryDecision { + fn action(action: Action) -> Self { + Self { + action, + template_index: 0, + failure: None, + } + } + + fn failure(kind: FailureKind, message: String) -> Self { + Self { + action: Action::Failure, + template_index: 0, + failure: Some(Failure { kind, message }), + } + } +} + +pub fn expand(entry: &Entry) -> EntryDecision { + if let Some(message) = &entry.hash_failure { + EntryDecision::failure(FailureKind::TypeError, message.clone()) + } else if entry.templates.is_some() { + EntryDecision::action(Action::Expand) + } else { + EntryDecision::action(Action::Keep) + } +} + +pub fn listed(facts: &ListedFacts) -> Result { + if facts.exhausted { + return Ok(EntryDecision::action(Action::Deny)); + } + let Some(entry) = &facts.entry else { + return Ok(EntryDecision::action(Action::ReadEntry)); + }; + if let Some(message) = &entry.hash_failure { + return Ok(EntryDecision::failure( + FailureKind::TypeError, + message.clone(), + )); + } + let Some(text) = &entry.text else { + return Ok(EntryDecision::failure( + FailureKind::AttributeError, + format!("'{}' object has no attribute 'endswith'", entry.python_type), + )); + }; + let Some(templates) = &entry.templates else { + return Ok(EntryDecision::action( + if patterns::wildcard(&facts.route, text) { + Action::Allow + } else { + Action::ReadEntry + }, + )); + }; + let next = if let Some(compiled) = &facts.compiled { + let expression = compiled + .strip_suffix('$') + .map_or_else(|| compiled.clone(), |prefix| format!("{prefix}(?=\\n?\\z)")); + if patterns::regex_match(&expression, &facts.route)? { + return Ok(EntryDecision::action(Action::Allow)); + } + facts.template_index + 1 + } else { + facts.template_index + }; + Ok(if next < templates.len() { + EntryDecision { + action: Action::ReadTemplate, + template_index: next, + failure: None, + } + } else { + EntryDecision::action(Action::ReadEntry) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(false, None, None, false)] + #[case::default(true, None, Some("internal_user"), false)] + #[case::unknown(true, Some("ADMIN"), Some("internal_user"), false)] + #[case::admin(true, Some("proxy_admin"), Some("proxy_admin"), true)] + #[case::viewer(true, Some("proxy_admin_viewer"), Some("proxy_admin_viewer"), false)] + fn roles( + #[case] present: bool, + #[case] raw: Option<&str>, + #[case] expected: Option<&str>, + #[case] admin: bool, + ) { + let result = role(&RoleFacts { + user_present: present, + role: raw.map(str::to_owned), + }); + assert_eq!(result.user_role, expected); + assert_eq!(result.admin, admin); + } + + #[rstest] + #[case::invalid(false, false, ApiDecision::InvalidToken)] + #[case::read(true, false, ApiDecision::ReadAdmin)] + #[case::nonadmin(true, true, ApiDecision::NonAdmin)] + fn api_order(#[case] token: bool, #[case] read: bool, #[case] expected: ApiDecision) { + assert_eq!( + api(&ApiFacts { + token_present: token, + admin_read: read, + user: RoleFacts { + user_present: false, + role: None + } + }), + expected + ); + } + + #[rstest] + #[case::read("team", false, false, JwtDecision::ReadTeam)] + #[case::default("team", true, true, JwtDecision::DefaultTeam)] + #[case::empty("team", true, false, JwtDecision::Check)] + #[case::other("internal_user", false, false, JwtDecision::Deny)] + fn jwt_selection( + #[case] role: &str, + #[case] read: bool, + #[case] missing: bool, + #[case] expected: JwtDecision, + ) { + assert_eq!( + jwt(&JwtFacts { + role: Some(role.into()), + config_read: read, + config_missing: missing + }), + expected + ); + } + + #[rstest] + #[case::self_empty(Some(""), Some(""), true)] + #[case::absent(None, None, false)] + #[case::different(Some("a"), Some("b"), false)] + fn inside_identity( + #[case] user: Option<&str>, + #[case] requested: Option<&str>, + #[case] expected: bool, + ) { + assert_eq!( + inside(&InsideFacts { + role: None, + user_id: user.map(str::to_owned), + requested_user_id: requested.map(str::to_owned) + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs b/litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs new file mode 100644 index 00000000000..12213f6b87b --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/spend_counter_policy.rs @@ -0,0 +1,299 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +#[serde(tag = "kind")] +pub enum Facts { + Spend { + current: f64, + verified: bool, + fallback: f64, + limit: Option, + window: bool, + fallback_authoritative: bool, + strict_before: Option, + authoritative_read: bool, + authoritative: Option, + strict_after: Option, + }, + ThrottleKey { + counter_key: String, + token_present: bool, + token: Option, + }, + Overflow { + current: f64, + limit: f64, + }, + OverBudget { + throttled: bool, + current: f64, + limit: f64, + reservation: f64, + strict: bool, + }, + ReservationFailure { + strict: bool, + }, + ReservationGroups { + currents: Vec, + limits: Vec, + reservation: f64, + }, + Adjustment { + current: Option, + adjustment: f64, + }, + ResizeConsistency { + inconsistent: bool, + reseed: bool, + }, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Value { spend: f64 }, + Repair { spend: f64 }, + ReadStrictBefore, + ReadStrictAfter, + ReadFloor, + Unverifiable, + ReadThrottle, + Skip, + Release, + Reject { current: f64, estimate: Option }, + Resize { cost: f64 }, + CheckOverflow, + Unavailable, + Compatible, + Inconsistent, + RejectResize, + ReservationGroups { groups: Vec> }, +} + +pub fn decide(facts: Facts) -> Decision { + match facts { + Facts::Adjustment { + current, + adjustment, + } => { + if current.is_some_and(|current| !(adjustment < 0.0 && current + adjustment < -1e-12)) { + Decision::Compatible + } else { + Decision::Inconsistent + } + } + Facts::ResizeConsistency { + inconsistent, + reseed, + } => { + if inconsistent && !reseed { + Decision::RejectResize + } else { + Decision::Compatible + } + } + Facts::Spend { + current, + verified, + fallback, + limit, + window, + fallback_authoritative, + strict_before, + authoritative_read, + authoritative, + strict_after, + } => { + if limit.is_none_or(|limit| current >= limit) { + return Decision::Value { spend: current }; + } + let obvious_floor = fallback > current || window; + if !obvious_floor && strict_before.is_none() { + return Decision::ReadStrictBefore; + } + if obvious_floor || strict_before == Some(true) { + if !authoritative_read { + return Decision::ReadFloor; + } + if let Some(spend) = authoritative { + if spend > current { + return Decision::Repair { spend }; + } + } else if fallback > current { + return Decision::Value { spend: fallback }; + } + } + if !(verified || fallback_authoritative || authoritative.is_some()) { + return match strict_after { + None => Decision::ReadStrictAfter, + Some(true) => Decision::Unverifiable, + Some(false) => Decision::Value { spend: current }, + }; + } + Decision::Value { spend: current } + } + Facts::ThrottleKey { + counter_key, + token_present, + token, + } => { + if token_present + && counter_key == format!("spend:key:{}", token.as_deref().unwrap_or("None")) + { + Decision::ReadThrottle + } else { + Decision::Skip + } + } + Facts::Overflow { current, limit } => { + if current > limit { + Decision::CheckOverflow + } else { + Decision::Skip + } + } + Facts::OverBudget { + throttled, + current, + limit, + reservation, + strict, + } => { + if throttled { + return Decision::Release; + } + let remaining = limit - (current - reservation); + if remaining <= 1e-12 { + return Decision::Reject { + current, + estimate: None, + }; + } + if strict && current - limit > 1e-12 { + return Decision::Reject { + current: current - reservation, + estimate: Some(reservation), + }; + } + Decision::Resize { cost: remaining } + } + Facts::ReservationFailure { strict } => { + if strict { + Decision::Unavailable + } else { + Decision::Skip + } + } + Facts::ReservationGroups { + currents, + limits, + reservation, + } => { + let indices: Vec = (0..currents.len()).collect(); + let groups = if currents.is_empty() { + vec![] + } else if currents + .iter() + .zip(&limits) + .all(|(current, limit)| current + reservation <= *limit) + { + vec![indices] + } else { + indices.into_iter().map(|index| vec![index]).collect() + }; + Decision::ReservationGroups { groups } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(None, 1.0, Decision::Inconsistent)] + #[case::positive(Some(1.0), 1.0, Decision::Compatible)] + #[case::negative(Some(1.0), -2.0, Decision::Inconsistent)] + #[case::zero(Some(1.0), -1.0, Decision::Compatible)] + fn counter_adjustment( + #[case] current: Option, + #[case] adjustment: f64, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::Adjustment { + current, + adjustment + }), + expected + ); + } + + #[rstest] + #[case::resize(true, false, Decision::RejectResize)] + #[case::settlement(true, true, Decision::Compatible)] + #[case::consistent(false, false, Decision::Compatible)] + fn consistency(#[case] inconsistent: bool, #[case] reseed: bool, #[case] expected: Decision) { + assert_eq!( + decide(Facts::ResizeConsistency { + inconsistent, + reseed + }), + expected + ); + } + + #[rstest] + #[case::throttled(true, 12.0, true, Decision::Release)] + #[case::empty(false, 12.0, false, Decision::Reject { current:12.0, estimate:None })] + #[case::partial(false, 11.0, false, Decision::Resize { cost:1.0 })] + #[case::strict(false, 11.0, true, Decision::Reject { current:9.0, estimate:Some(2.0) })] + fn reservation_outcome( + #[case] throttled: bool, + #[case] current: f64, + #[case] strict: bool, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::OverBudget { + throttled, + current, + limit: 10.0, + reservation: 2.0, + strict + }), + expected + ); + } + + #[rstest] + #[case::unchecked(None, false, None, Decision::Value { spend:2.0 })] + #[case::strict_read(Some(3.0), false, None, Decision::ReadStrictBefore)] + #[case::window(Some(3.0), true, None, Decision::ReadFloor)] + #[case::strict_floor(Some(3.0), false, Some(true), Decision::ReadFloor)] + #[case::second_read(Some(3.0), false, Some(false), Decision::ReadStrictAfter)] + fn spend_read_order( + #[case] limit: Option, + #[case] window: bool, + #[case] strict_before: Option, + #[case] expected: Decision, + ) { + assert_eq!( + decide(Facts::Spend { + current: 2.0, + verified: false, + fallback: 0.0, + limit, + window, + fallback_authoritative: false, + strict_before, + authoritative_read: false, + authoritative: None, + strict_after: None + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_aliases.rs b/litellm-rust/crates/request-auth-core/src/team_aliases.rs new file mode 100644 index 00000000000..6abd1e477b8 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/team_aliases.rs @@ -0,0 +1,77 @@ +use crate::Error; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub aliases: BTreeMap, + pub team_id: Option, + pub router_present: bool, + pub names: Option>>, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadNames { count: usize }, + Complete { models: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let prefix = format!("model_name_{}_", facts.team_id.as_deref().unwrap_or("None")); + let needs_names = |target: &str| facts.router_present && target.starts_with(&prefix); + let count = facts + .requested + .iter() + .filter_map(|model| facts.aliases.get(model)) + .filter(|target| needs_names(target)) + .count(); + let names = match &facts.names { + None if count != 0 => return Ok(Decision::ReadNames { count }), + Some(names) if names.len() != count => return Err(Error::InvalidAnswer), + names => names.as_deref().unwrap_or_default(), + }; + let mut reads = names.iter(); + let models = facts + .requested + .iter() + .map(|model| { + let Some(target) = facts.aliases.get(model) else { + return model.clone(); + }; + if needs_names(target) && !reads.next().is_some_and(|names| names.contains(target)) { + model.clone() + } else { + target.clone() + } + }) + .collect(); + Ok(Decision::Complete { models }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing(false, "alias")] + #[case::live(true, "model_name_team_target")] + fn deleted_target(#[case] live: bool, #[case] expected: &str) { + let target = "model_name_team_target".to_owned(); + let facts = Facts { + requested: vec!["alias".into()], + aliases: BTreeMap::from([("alias".into(), target.clone())]), + team_id: Some("team".into()), + router_present: true, + names: Some(vec![if live { vec![target] } else { vec![] }]), + }; + assert_eq!( + decide(&facts).unwrap(), + Decision::Complete { + models: vec![expected.into()] + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_override.rs b/litellm-rust/crates/request-auth-core/src/team_override.rs new file mode 100644 index 00000000000..08e0cc20481 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/team_override.rs @@ -0,0 +1,70 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + model_access::{self, Decision as ModelDecision, Facts as ModelFacts, ObjectType}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub team: Option, + pub assigned_models: Option>, + pub assigned_scope: Option, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: ModelDecision }, + ReadTeam, + ReadTeamGroups, + ReadAssignedGroups, + ReadAssignedScope, +} + +pub fn decide(facts: &Facts) -> Result { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeam); + }; + let denial = fallback_access::grouped(team.clone(), true)?; + match denial { + ModelDecision::Allow => return Ok(Decision::Allow), + ModelDecision::ReadGroups | ModelDecision::ReadGroupScope { .. } => { + return Ok(Decision::ReadTeamGroups); + } + ModelDecision::Deny { + object_type: ObjectType::Team, + .. + } => (), + _ => return Ok(Decision::Deny { denial }), + } + let Some(assigned) = &facts.assigned_models else { + return Ok(Decision::ReadAssignedGroups); + }; + if assigned.is_empty() { + return Ok(Decision::Deny { denial }); + } + let Some(scope) = &facts.assigned_scope else { + return Ok(Decision::ReadAssignedScope); + }; + match model_access::check(scope)? { + ModelDecision::Allow => Ok(Decision::Allow), + ModelDecision::Deny { .. } => Ok(Decision::Deny { denial }), + error => Ok(Decision::Deny { denial: error }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn team_is_read_before_override_groups() { + assert!(matches!( + decide(&Facts::default()).unwrap(), + Decision::ReadTeam + )); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/unpriced_models.rs b/litellm-rust/crates/request-auth-core/src/unpriced_models.rs new file mode 100644 index 00000000000..bc6073555c3 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/unpriced_models.rs @@ -0,0 +1,144 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Metric { + pub name: String, + pub scalar: Option, + pub nested: Vec>, +} + +#[derive(Clone, Deserialize)] +pub struct Entry { + pub tiered_present: bool, + pub metrics: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Deployment { + pub parameter_keys: Vec, + pub has_id: bool, + pub info_loaded: bool, + pub info: Option, + pub cost_keys: Option>, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub model: Option, + pub router_present: bool, + pub group_found: Option, + pub deployments: Option>, + pub cost_deployments: Option>, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { unpriced: bool }, + ReadGroup, + ReadDeployments, + ReadCostDeployments, + ReadInfo { index: usize }, + ReadCost { index: usize }, +} + +pub fn entry_has_priced_metric(entry: &Entry) -> bool { + entry.tiered_present + || entry.metrics.iter().any(|metric| { + metric.name.contains("cost_per") + && (metric.scalar.is_some_and(|price| price > 0.0) + || metric + .nested + .iter() + .any(|price| price.is_some_and(|value| value > 0.0))) + }) +} + +pub fn decide(facts: &Facts) -> Decision { + if !facts.router_present || facts.model.as_deref().is_none_or(str::is_empty) { + return Decision::Complete { unpriced: false }; + } + match facts.group_found { + None => return Decision::ReadGroup, + Some(false) => return Decision::Complete { unpriced: false }, + Some(true) => {} + } + let Some(deployments) = &facts.deployments else { + return Decision::ReadDeployments; + }; + for (index, deployment) in deployments.iter().enumerate() { + if deployment + .parameter_keys + .iter() + .any(|key| key.contains("cost_per") || key == "tiered_pricing") + { + return Decision::Complete { unpriced: false }; + } + if !deployment.has_id { + continue; + } + if !deployment.info_loaded { + return Decision::ReadInfo { index }; + } + if deployment + .info + .as_ref() + .is_some_and(entry_has_priced_metric) + { + return Decision::Complete { unpriced: false }; + } + } + let Some(cost_deployments) = &facts.cost_deployments else { + return Decision::ReadCostDeployments; + }; + for (index, deployment) in cost_deployments.iter().enumerate() { + if !deployment.has_id { + continue; + } + let Some(keys) = &deployment.cost_keys else { + return Decision::ReadCost { index }; + }; + if keys.iter().any(|key| { + matches!( + key.as_str(), + "input_cost_per_token" | "output_cost_per_token" + ) + }) { + return Decision::Complete { unpriced: false }; + } + } + Decision::Complete { unpriced: true } +} + +#[cfg(test)] +mod tests { + use super::{Entry, Metric, entry_has_priced_metric}; + use rstest::rstest; + + #[rstest] + #[case::positive("input_cost_per_character", Some(0.1), vec![], false, true)] + #[case::zero("input_cost_per_token", Some(0.0), vec![], false, false)] + #[case::negative("cost_per_query", Some(-1.0), vec![], false, false)] + #[case::nested("cost_per_image", None, vec![None, Some(0.2)], false, true)] + #[case::non_price("context_window", Some(1000.0), vec![], false, false)] + #[case::tiered("context_window", None, vec![], true, true)] + fn metrics( + #[case] name: &str, + #[case] scalar: Option, + #[case] nested: Vec>, + #[case] tiered_present: bool, + #[case] expected: bool, + ) { + assert_eq!( + entry_has_priced_metric(&Entry { + tiered_present, + metrics: vec![Metric { + name: name.to_owned(), + scalar, + nested + }] + }), + expected + ); + } +} diff --git a/litellm/_auth/access.py b/litellm/_auth/access.py new file mode 100644 index 00000000000..2e3df005b63 --- /dev/null +++ b/litellm/_auth/access.py @@ -0,0 +1,200 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Blocked: + blocked: bool + + +@dataclass(frozen=True, slots=True) +class UserParamRequest: + enabled: bool + + +@dataclass(frozen=True, slots=True) +class ClientTagRequest: + enabled: bool + llm_route: bool | None = None + + +@dataclass(frozen=True, slots=True) +class GuardrailRequest: + requested_keys: tuple[str, ...] + + +TraceSource: TypeAlias = Literal["Gate", "Header", "Body", "Traceparent", "Selected", "Other"] + + +@dataclass(frozen=True, slots=True) +class TraceRequest: + required: bool + + +@dataclass(frozen=True, slots=True) +class Search: + name: str + allowed: tuple[str, ...] + strict: bool + + +@dataclass(frozen=True, slots=True) +class TraceGate: + required: bool + method: str + llm_route: bool + passthrough: bool + telemetry: bool + + +@dataclass(frozen=True, slots=True) +class TraceStep: + source: TraceSource + value: str | None + present: bool + headers_only: bool + + +@dataclass(frozen=True, slots=True) +class UserParam: + enabled: bool + method: str | None + llm_route: bool + mcp_route: bool + has_user: bool + + +@dataclass(frozen=True, slots=True) +class ClientTags: + enabled: bool + llm_route: bool + has_tags: bool + + +@dataclass(frozen=True, slots=True) +class Guardrails: + requested_keys: tuple[str, ...] + modify_guardrails: bool | None + + +@dataclass(frozen=True, slots=True) +class AgentTrace: + required: bool + header: str | None + + +@dataclass(frozen=True, slots=True) +class ConcurrentFailures: + failures: tuple[bool, ...] + terminal: tuple[bool, ...] + + +@dataclass(frozen=True, slots=True) +class TeamTrace: + required: bool + method: str + llm_route: bool + passthrough: bool + telemetry: bool + headers_only: bool + header: str | None + body_trace: str | None + body_trace_present: bool + traceparent_id: str | None + selected_trace: str | None + selected_trace_present: bool + other_trace: str | None + + +@dataclass(frozen=True, slots=True) +class Tools: + requested: tuple[str, ...] + key_allowed: tuple[str, ...] | None + team_allowed: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class GrantLayers: + deny_by_default: bool + identity_present: bool + master_key: bool + dashboard: bool + virtual_key: bool + session_token: bool + team_present: bool + + +@dataclass(frozen=True, slots=True) +class VectorStores: + requested: tuple[str, ...] + allowed: tuple[str, ...] | None + strict: bool + + +@dataclass(frozen=True, slots=True) +class TeamLookup: + dashboard: bool + missing: bool + recorded_models: tuple[str, ...] + allow_unavailable: bool | None + + +AccessFacts: TypeAlias = ( + ConcurrentFailures + | UserParamRequest + | ClientTagRequest + | GuardrailRequest + | TraceRequest + | Search + | Blocked + | TraceGate + | TraceStep + | UserParam + | ClientTags + | Guardrails + | AgentTrace + | TeamTrace + | Tools + | GrantLayers + | VectorStores + | TeamLookup +) +AccessDenial: TypeAlias = Literal[ + "Search", + "Blocked", + "UserParam", + "ClientTags", + "Guardrails", + "AgentTrace", + "TeamTrace", + "Tools", + "VectorStores", + "MissingVectorGrant", + "TeamLookup", +] +GrantLayer: TypeAlias = Literal["key", "team", "user"] + + +@dataclass(frozen=True, slots=True) +class AccessDecision: + read: Literal["UserParam", "ClientTagRoute", "ClientTags", "GuardrailPermission", "UnavailablePolicy"] | None = None + denial: AccessDenial | None = None + rejected: tuple[str, ...] = () + layers: tuple[GrantLayer, ...] = () + rejected_index: int | None = None + next_trace: TraceSource | None = None + failure_index: int | None = None + + +_DECISION: Final = TypeAdapter(AccessDecision) + + +@native_boundary +def decide(facts: AccessFacts) -> AccessDecision: + return _DECISION.validate_python(_native.request_auth_access(encode_facts(facts))) diff --git a/litellm/_auth/admission.py b/litellm/_auth/admission.py new file mode 100644 index 00000000000..691cb6a37ff --- /dev/null +++ b/litellm/_auth/admission.py @@ -0,0 +1,77 @@ +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class CommonGate: + force_virtual_key_checks: bool + public_route: bool + additional_public_route: bool + unauthenticated_passthrough: bool + master_key_present: bool + jwt_enabled: bool + oauth_enabled: bool + oauth_proxy_enabled: bool + custom_auth: bool + custom_auth_common_checks: bool + + +@dataclass(frozen=True, slots=True) +class Reservation: + skip_budget_checks: bool + disabled: bool + apply_user_budget_to_team_keys: bool + fail_closed: bool + billing_agent: bool + + +@dataclass(frozen=True, slots=True) +class Ownership: + virtual_key: bool + custom_auth: bool + jwt: bool + oauth: bool + + +@dataclass(frozen=True, slots=True) +class CommonPlan: + skip: Literal["PublicRoute", "UnauthenticatedPassthrough", "NoAuthDevelopment", "CustomAuth"] | None + + +@dataclass(frozen=True, slots=True) +class ReservationPlan: + reserve: bool + apply_user_budget_to_team_keys: bool + fail_closed: bool + + +@dataclass(frozen=True, slots=True) +class OwnershipPlan: + owner: Literal["Python"] + candidate: Literal["VirtualKey", "CustomAuth", "Validator", "Other"] + + +_COMMON: Final = TypeAdapter(CommonPlan) +_RESERVATION: Final = TypeAdapter(ReservationPlan) +_OWNERSHIP: Final = TypeAdapter(OwnershipPlan) + + +@native_boundary +def common_plan(facts: CommonGate) -> CommonPlan: + return _COMMON.validate_python(_native.request_auth_admission(encode_facts(facts))) + + +@native_boundary +def reservation_plan(facts: Reservation) -> ReservationPlan: + return _RESERVATION.validate_python(_native.request_auth_admission(encode_facts(facts))) + + +@native_boundary +def ownership_plan(facts: Ownership) -> OwnershipPlan: + return _OWNERSHIP.validate_python(_native.request_auth_admission(encode_facts(facts))) diff --git a/litellm/_auth/agent_access.py b/litellm/_auth/agent_access.py new file mode 100644 index 00000000000..1655dc3445a --- /dev/null +++ b/litellm/_auth/agent_access.py @@ -0,0 +1,68 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.managed import Agent +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class AgentAccessFacts: + registry_managed: bool + database_present: bool + target_loaded: bool = False + target: Agent | None = None + subject_present: bool = False + key_hash_present: bool = False + subject_managed: bool = False + session: bool = False + custom: bool = False + authority_loaded: bool = False + explicit_grants: tuple[str, ...] | None = None + legacy_loaded: bool = False + legacy_grants: tuple[str, ...] | None = None + target_ids: tuple[str, ...] | None = None + + +AgentAccessDecision: TypeAlias = Literal[ + "ReadTarget", "ReadAuthority", "ReadExplicitGrants", "ReadLegacyGrants", "ReadTargetIds", "Allow", "Deny" +] +_DECISION: Final[TypeAdapter[AgentAccessDecision]] = TypeAdapter(AgentAccessDecision) + + +@native_boundary +def decide(facts: AgentAccessFacts) -> AgentAccessDecision: + return _DECISION.validate_python(_native.request_auth_agent_access(encode_facts(facts))) + + +@dataclass(frozen=True, slots=True) +class InvocationFacts: + registered: Agent | None + has_store: bool + target_loaded: bool = False + target: Agent | None = None + actor_managed: bool = False + actor_id_present: bool = False + permissions: AgentAccessFacts | None = None + assigned: bool = False + fee_loaded: bool = False + fee: float | None = None + + +@dataclass(frozen=True, slots=True) +class InvocationPlan: + action: Literal["Complete", "ReadTarget", "ReadPermissions", "Assign", "ReadFee", "Deny"] + denial: Literal["DatabaseRequired", "Removed", "CallerNotPermitted", "InvalidFee"] | None + use_target: bool + assign_billing: bool + + +_INVOCATION: Final = TypeAdapter(InvocationPlan) + + +@native_boundary +def invocation(facts: InvocationFacts) -> InvocationPlan: + return _INVOCATION.validate_python(_native.request_auth_agent_invocation(encode_facts(facts))) diff --git a/litellm/_auth/agent_ceilings.py b/litellm/_auth/agent_ceilings.py new file mode 100644 index 00000000000..544bf05d556 --- /dev/null +++ b/litellm/_auth/agent_ceilings.py @@ -0,0 +1,60 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class AgentCeilingGroup: + status: Literal["Found", "Missing", "Unavailable"] + models: tuple[str, ...] = () + mcp_server_ids: tuple[str, ...] = () + agent_ids: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class AgentCeilingFacts: + ids: tuple[str, ...] | None = None + groups: tuple[AgentCeilingGroup, ...] | None = None + check_db_only: bool = False + + +@dataclass(frozen=True, slots=True) +class AgentCeilingDecision: + kind: Literal["ReadIds", "ReadGroups", "Complete"] + authoritative: bool = False + present: bool = False + ids: tuple[str, ...] = () + models: frozenset[str] = frozenset() + mcp_server_ids: frozenset[str] = frozenset() + agent_ids: frozenset[str] = frozenset() + indices: tuple[int, ...] = () + + +@dataclass(frozen=True, slots=True) +class AgentGroupLoadFacts: + database_present: bool + check_db_only: bool = False + status: Literal["Success", "NotFound", "HttpError"] | None = None + + +@dataclass(frozen=True, slots=True) +class AgentGroupLoadDecision: + kind: Literal["Read", "Return", "DatabaseUnavailable", "HttpUnavailable", "Raise"] + + +_DECISION: Final = TypeAdapter(AgentCeilingDecision) +_LOAD: Final = TypeAdapter(AgentGroupLoadDecision) + + +@native_boundary +def agent_ceiling(facts: AgentCeilingFacts) -> AgentCeilingDecision: + return _DECISION.validate_python(_native.request_auth_agent_ceiling(asdict(facts))) + + +@native_boundary +def agent_group_load(facts: AgentGroupLoadFacts) -> AgentGroupLoadDecision: + return _LOAD.validate_python(_native.request_auth_agent_group_load(asdict(facts))) diff --git a/litellm/_auth/agent_grants.py b/litellm/_auth/agent_grants.py new file mode 100644 index 00000000000..a0af67f5273 --- /dev/null +++ b/litellm/_auth/agent_grants.py @@ -0,0 +1,128 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Source: + team_id: str | None + grants: tuple[str, ...] | None + unavailable: bool = False + + +@dataclass(frozen=True, slots=True) +class Resolve: + managed: bool + strict: bool + role: str | None + key_team_loaded: bool = False + key_team: tuple[str, ...] | None = None + default_deny: bool | None = None + caller_loaded: bool = False + own: tuple[str, ...] | None = None + ceiling_loaded: bool = False + ceiling: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class Failure: + strict: bool + closed: bool + + +@dataclass(frozen=True, slots=True) +class OwnerTeam: + strict: bool + permission_id_present: bool + permission_present: bool + + +@dataclass(frozen=True, slots=True) +class OwnerUnion: + teams: tuple[Source, ...] + withheld: tuple[Source, ...] + unreadable: bool + + +@dataclass(frozen=True, slots=True) +class BillingTeam: + sources: tuple[Source, ...] + agent_id: str + + +@dataclass(frozen=True, slots=True) +class Human: + sources: tuple[Source, ...] + team_id: str | None + + +@dataclass(frozen=True, slots=True) +class Delegated: + sources: tuple[Source, ...] + agent_id: str + team_id: str | None + explicit: bool + allowed: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class TeamLookup: + subject: bool + team_id: str | None + database: bool + strict: bool + personal: bool | None = None + loaded: bool = False + row_present: bool = False + + +@dataclass(frozen=True, slots=True) +class ConfiguredAgent: + agent_id: str + groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class GroupMatches: + agents: tuple[ConfiguredAgent, ...] + groups: tuple[str, ...] + + +Action: TypeAlias = Literal[ + "Complete", + "ReadManaged", + "ReadKeyTeam", + "ReadDefault", + "ReadCaller", + "ReadCeiling", + "ReadPersonalEligibility", + "ReadOwnerTeams", + "ReadTeam", + "ReadDeclared", + "Unavailable", + "SelectTeam", +] + + +@dataclass(frozen=True, slots=True) +class Decision: + action: Action + grants: tuple[str, ...] | None + team_id: str | None + + +_DECISION: Final = TypeAdapter(Decision) + +Facts: TypeAlias = ( + Resolve | Failure | OwnerTeam | OwnerUnion | BillingTeam | Human | Delegated | TeamLookup | GroupMatches +) + + +@native_boundary +def decide(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_agent_grants(encode_facts(facts))) diff --git a/litellm/_auth/auth_context.py b/litellm/_auth/auth_context.py new file mode 100644 index 00000000000..a0754b1ed11 --- /dev/null +++ b/litellm/_auth/auth_context.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Fetch: + team_id: str | None + ui_team_id: str + user_id: str | None + project_id: str | None + end_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Team: + team_id: str | None + ui_team_id: str + failed: bool + missing: bool + recorded_models: tuple[str, ...] + allow_unavailable: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Admin: + role: str | None + user_id: str | None + admin_id: str + user_row: bool + + +@dataclass(frozen=True, slots=True) +class FetchPlan: + team: bool + user: bool + project: bool + end_user: bool + + +@dataclass(frozen=True, slots=True) +class TeamPlan: + source: Literal["None", "Token", "Loaded", "ReadUnavailablePolicy", "Failure"] + + +@dataclass(frozen=True, slots=True) +class AdminPlan: + force_admin: bool + preserve_budget: bool + + +_FETCH: Final = TypeAdapter(FetchPlan) +_TEAM: Final = TypeAdapter(TeamPlan) +_ADMIN: Final = TypeAdapter(AdminPlan) + + +@native_boundary +def fetch(facts: Fetch) -> FetchPlan: + return _FETCH.validate_python(_native.request_auth_context(encode_facts(facts))) + + +@native_boundary +def team(facts: Team) -> TeamPlan: + return _TEAM.validate_python(_native.request_auth_context(encode_facts(facts))) + + +@native_boundary +def admin(facts: Admin) -> AdminPlan: + return _ADMIN.validate_python(_native.request_auth_context(encode_facts(facts))) diff --git a/litellm/_auth/budget_alerts.py b/litellm/_auth/budget_alerts.py new file mode 100644 index 00000000000..d5546243790 --- /dev/null +++ b/litellm/_auth/budget_alerts.py @@ -0,0 +1,69 @@ +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: str + message: str + + +@dataclass(frozen=True, slots=True) +class ThresholdFacts: + key: str + digit: bool + digits: str | None = None + ratio: float | None = None + integer_failure: Failure | None = None + ratio_failure: Failure | None = None + text: str | None = None + items: tuple[str | None, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ConfigFacts: + global_entries: tuple[ThresholdFacts, ...] = () + local_entries: tuple[ThresholdFacts, ...] = () + global_failure: Failure | None = None + local_failure: Failure | None = None + order: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class Facts: + member: bool + spend: float | None + limit: float | None + default_fraction: float + config: ConfigFacts | None = None + + +@dataclass(frozen=True, slots=True) +class Threshold: + key: str + recipients: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Decision: + read_config: bool = False + read_order: bool = False + global_keys: tuple[str, ...] = () + local_keys: tuple[str, ...] = () + emit: bool = False + thresholds: tuple[Threshold, ...] | None = None + failure: Failure | None = None + + +_FACTS: Final = TypeAdapter(Facts) +_DECISION: Final = TypeAdapter(Decision) + + +@native_boundary +def plan(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_budget_alert_plan(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/budget_context.py b/litellm/_auth/budget_context.py new file mode 100644 index 00000000000..ae052b3605f --- /dev/null +++ b/litellm/_auth/budget_context.py @@ -0,0 +1,174 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelCheckFacts: + scope: Literal["Team", "EndUser", "CustomEndUser"] + entity_id: str | None + budget_count: int + model_count: int + + +@dataclass(frozen=True, slots=True) +class BudgetFields: + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + tpd_limit: int | None = None + model_budget_count: int | None = None + + +@dataclass(frozen=True, slots=True) +class EndUserFacts: + mode: Literal["Models", "Token", "Default"] + object_present: bool + end_user_id: str | None + prisma_present: bool + key_budget_id: str | None + key_read: bool = False + key_budget: BudgetFields | None = None + global_config_read: bool = False + global_budget_id: str | None = None + global_read: bool = False + global_budget: BudgetFields | None = None + + +EndUserSource: TypeAlias = Literal["Own", "ReadKey", "Key", "ReadGlobalConfig", "ReadGlobal", "Global", "None"] + + +@dataclass(frozen=True, slots=True) +class TokenDefaultsFacts: + current: BudgetFields + keep_token_limits: bool + default_present: bool + + +TokenField: TypeAlias = Literal["MaxBudget", "Tpm", "Rpm", "Tpd", "Models"] + + +@dataclass(frozen=True, slots=True) +class MembershipSnapshot: + budget_id: str | None + model_budget_count: int | None + + +@dataclass(frozen=True, slots=True) +class MemberFacts: + user_id: str | None + team_id: str | None + prisma_present: bool + default_id: str | None + skip_budget_checks: bool + membership_read: bool = False + membership: MembershipSnapshot | None = None + default_read: bool = False + default_budget: BudgetFields | None = None + lookup_failure: Literal["Membership", "Default"] | None = None + + +MemberSource: TypeAlias = Literal[ + "Done", "ReadMembership", "ReadDefault", "Member", "Default", "Clear", "RaiseLookup", "SkipLookup" +] + + +@dataclass(frozen=True, slots=True) +class IncreaseFacts: + number: float | None + truthy: bool + type_name: str + + +@dataclass(frozen=True, slots=True) +class TemporaryFacts: + max_budget: float | None + metadata_read: bool = False + has_increase: bool = False + has_expiry: bool = False + expires_at: int | None = None + now: int | None = None + increase: IncreaseFacts | None = None + + +@dataclass(frozen=True, slots=True) +class TemporaryDecision: + read: Literal["Metadata", "Expiry", "Now", "Increase"] | None = None + copy: bool = False + max_budget: float | None = None + invalid_increase_type: str | None = None + + +@dataclass(frozen=True, slots=True) +class NumericFacts: + number: float | None + boolean: bool + + +@dataclass(frozen=True, slots=True) +class ThrottleFacts: + opt_in: bool + tpm_limit: int | None + rpm_limit: int | None + percentage: NumericFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ThrottleDecision: + read_percentage: bool = False + percentage: float | None = None + + +_MODELS: Final = TypeAdapter(ModelCheckFacts) +_INDICES: Final = TypeAdapter(tuple[int, ...]) +_END: Final = TypeAdapter(EndUserFacts) +_END_SOURCE: Final = TypeAdapter[EndUserSource](EndUserSource) +_TOKEN: Final = TypeAdapter(TokenDefaultsFacts) +_FIELDS: Final = TypeAdapter(tuple[TokenField, ...]) +_MEMBER: Final = TypeAdapter(MemberFacts) +_MEMBER_SOURCE: Final = TypeAdapter[MemberSource](MemberSource) +_TEMP: Final = TypeAdapter(TemporaryFacts) +_TEMP_DECISION: Final = TypeAdapter(TemporaryDecision) +_THROTTLE: Final = TypeAdapter(ThrottleFacts) +_THROTTLE_DECISION: Final = TypeAdapter(ThrottleDecision) +_NUMBER: Final = TypeAdapter(NumericFacts) +_PERCENTAGE: Final = TypeAdapter[float | None](float | None) + + +@native_boundary +def model_checks(facts: ModelCheckFacts) -> tuple[int, ...]: + return _INDICES.validate_python(_native.request_auth_context_models(_MODELS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def end_user_source(facts: EndUserFacts) -> EndUserSource: + return _END_SOURCE.validate_python(_native.request_auth_context_end_user(_END.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def token_assignments(facts: TokenDefaultsFacts) -> tuple[TokenField, ...]: + return _FIELDS.validate_python(_native.request_auth_context_token(_TOKEN.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def member_source(facts: MemberFacts) -> MemberSource: + return _MEMBER_SOURCE.validate_python(_native.request_auth_context_member(_MEMBER.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def temporary(facts: TemporaryFacts) -> TemporaryDecision: + return _TEMP_DECISION.validate_python(_native.request_auth_context_temporary(_TEMP.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def percentage(facts: NumericFacts) -> float | None: + return _PERCENTAGE.validate_python(_native.request_auth_context_percentage(_NUMBER.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def throttle(facts: ThrottleFacts) -> ThrottleDecision: + return _THROTTLE_DECISION.validate_python(_native.request_auth_context_throttle(_THROTTLE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/budget_lookup.py b/litellm/_auth/budget_lookup.py new file mode 100644 index 00000000000..a980a1dee26 --- /dev/null +++ b/litellm/_auth/budget_lookup.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class OrganizationFacts: + token_present: bool + client_present: bool + team_present: bool + token_id: bool | None = None + team_id: bool | None = None + row_present: bool | None = None + lookup_failed: bool = False + + +OrganizationAction: TypeAlias = Literal["Skip", "ReadTokenId", "ReadTeamId", "ReadOrganization", "UseOrganization"] + + +@dataclass(frozen=True, slots=True) +class MemberFacts: + team_present: bool + token_present: bool + membership_loaded: bool + team_id: bool | None = None + user_id: bool | None = None + limits_read: bool = False + maximum_present: bool = False + default_id_string: bool = False + default_read: bool = False + + +MemberAction: TypeAlias = Literal[ + "Skip", "ReadTeamId", "ReadUserId", "ReadMembership", "ReadLimits", "ReadDefault", "UseLimits" +] +_ORGANIZATION: Final = TypeAdapter[OrganizationAction](OrganizationAction) +_MEMBER: Final = TypeAdapter[MemberAction](MemberAction) +_ORGANIZATION_FACTS: Final = TypeAdapter(OrganizationFacts) +_MEMBER_FACTS: Final = TypeAdapter(MemberFacts) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_EMAILS: Final = TypeAdapter(tuple[str, ...]) +_METADATA: Final = TypeAdapter(Mapping[object, object]) + + +def soft_email_value(raw: object) -> object: + return _METADATA.validate_python(raw).get("soft_budget_alerting_emails") if isinstance(raw, dict) else None + + +@dataclass(frozen=True, slots=True) +class SoftFacts: + limit: float | None + spend: float | None + token_present: bool + + +_SOFT: Final = TypeAdapter(SoftFacts) + + +@native_boundary +def read_soft_emails(facts: SoftFacts) -> bool: + return _native.request_auth_budget_lookup_soft(_SOFT.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def organization(facts: OrganizationFacts) -> OrganizationAction: + return _ORGANIZATION.validate_python( + _native.request_auth_budget_lookup_organization(_ORGANIZATION_FACTS.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + ) + + +@native_boundary +def member(facts: MemberFacts) -> MemberAction: + return _MEMBER.validate_python(_native.request_auth_budget_lookup_member(_MEMBER_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def emails(raw: object) -> tuple[str, ...]: + facts: Final = ( + {"Text": raw} + if isinstance(raw, str) + else {"List": tuple(value if isinstance(value, str) else None for value in _ITEMS.validate_python(raw))} + if isinstance(raw, list) + else "Other" + ) + return _EMAILS.validate_python(_native.request_auth_budget_lookup_emails(facts)) diff --git a/litellm/_auth/budget_pricing.py b/litellm/_auth/budget_pricing.py new file mode 100644 index 00000000000..fda075c1730 --- /dev/null +++ b/litellm/_auth/budget_pricing.py @@ -0,0 +1,90 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class DeploymentPrice: + has_input_price: bool = False + has_output_price: bool = False + input_price: float | None = None + output_price: float | None = None + query_price: float | None = None + ptu_count: float | None = None + cost_per_ptu_per_hour: float | None = None + + +@dataclass(frozen=True, slots=True) +class ModelPriceFacts: + cached: bool | None = None + lookup_failed: bool = False + group_found: bool = False + input_price: float | None = None + output_price: float | None = None + deployments: tuple[DeploymentPrice, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class PriceDecision: + zero_cost: bool + cache_value: bool | None + read_deployments: bool = False + + +@dataclass(frozen=True, slots=True) +class RouteBudgetFacts: + side_effect_route: bool = False + discovery_route: bool = False + llm_route: bool = False + registered_pass_through: bool | None = None + route: str = "" + method: str | None = None + body_present: bool = False + + +RouteBudgetDecision: TypeAlias = Literal["ReadPassThrough", "ReadPricing", "Skip"] +_PRICE: Final = TypeAdapter(ModelPriceFacts) +_PRICE_DECISION: Final = TypeAdapter(PriceDecision) +_ROUTE: Final = TypeAdapter(RouteBudgetFacts) +_ROUTE_DECISION: Final = TypeAdapter[RouteBudgetDecision](RouteBudgetDecision) + + +@native_boundary +def price(facts: ModelPriceFacts) -> PriceDecision: + return _PRICE_DECISION.validate_python(_native.request_auth_budget_price(_PRICE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def route(facts: RouteBudgetFacts) -> RouteBudgetDecision: + return _ROUTE_DECISION.validate_python(_native.request_auth_budget_route(_ROUTE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@dataclass(frozen=True, slots=True) +class ExemptionFacts: + route: RouteBudgetFacts + model_input_present: bool + router_present: bool + model_count: int + prices: tuple[ModelPriceFacts, ...] = () + cache_read: bool = False + + +@dataclass(frozen=True, slots=True) +class ExemptionDecision: + read_pass_through: bool + read_price_cache: bool + read_model: int | None + exempt: bool + + +_EXEMPTION: Final = TypeAdapter(ExemptionFacts) +_EXEMPTION_DECISION: Final = TypeAdapter(ExemptionDecision) + + +@native_boundary +def exemption(facts: ExemptionFacts) -> ExemptionDecision: + return _EXEMPTION_DECISION.validate_python(_native.request_auth_budget_exemption(_EXEMPTION.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/common_budgets.py b/litellm/_auth/common_budgets.py new file mode 100644 index 00000000000..54e6a020a20 --- /dev/null +++ b/litellm/_auth/common_budgets.py @@ -0,0 +1,115 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +BudgetScope: TypeAlias = Literal[ + "Team", + "Key", + "TeamWindow", + "KeyWindow", + "Project", + "Organization", + "Tag", + "AccessGroup", + "User", + "TeamMember", + "EndUser", + "Proxy", + "KeySoft", + "TeamSoft", + "ProjectSoft", +] +BudgetAlert: TypeAlias = Literal[ + "token_budget", + "user_budget", + "team_budget", + "project_budget", + "organization_budget", + "team_member_threshold", + "soft_budget", +] + + +@dataclass(frozen=True, slots=True) +class BudgetFacts: + scope: BudgetScope + limit: float | None + spend: float | None = None + token_present: bool = False + team_key: bool = False + apply_user_budget_to_team_keys: bool = False + info_route: bool = False + llm_route: bool = True + model_list_route: bool = False + skip: bool = False + alert_recipients: tuple[str, ...] = () + throttle_percentage: float | None = None + throttle_opt_in: bool = False + tpm_limit: int | None = None + rpm_limit: int | None = None + + +@dataclass(frozen=True, slots=True) +class BudgetDecision: + read_spend: bool = False + denied: bool = False + alert: BudgetAlert | None = None + throttle_percentage: float | None = None + + +@dataclass(frozen=True, slots=True) +class AlertFacts: + spend: float | None + max_budget: float | None + configured: bool = False + minimum_percent: float | None = None + default_ratio: float = 0.8 + team_member: bool = False + + +@dataclass(frozen=True, slots=True) +class MemberLimitFacts: + member_limit: float | None + default_limit: float | None + temporary_increase: float | None + expires_at: float | None + now: float + now_naive: bool = False + + +@dataclass(frozen=True, slots=True) +class WindowFacts: + duration_seconds: float | None + reset_at: float | None + now: float + + +_FACTS: Final = TypeAdapter(BudgetFacts) +_DECISION: Final = TypeAdapter(BudgetDecision) +_ALERT: Final = TypeAdapter(AlertFacts) +_MEMBER: Final = TypeAdapter(MemberLimitFacts) +_WINDOW: Final = TypeAdapter(WindowFacts) + + +@native_boundary +def check(facts: BudgetFacts) -> BudgetDecision: + return _DECISION.validate_python(_native.request_auth_common_budget(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def alert(facts: AlertFacts) -> bool: + return _native.request_auth_budget_alert(_ALERT.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def member_limit(facts: MemberLimitFacts) -> float | None: + return _native.request_auth_member_limit(_MEMBER.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def window_start(facts: WindowFacts) -> float | None: + return _native.request_auth_budget_window_start(_WINDOW.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary diff --git a/litellm/_auth/common_flow.py b/litellm/_auth/common_flow.py new file mode 100644 index 00000000000..208f68cd253 --- /dev/null +++ b/litellm/_auth/common_flow.py @@ -0,0 +1,148 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class BudgetSkip: + skip: bool + exempt: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Pricing: + enabled: bool + llm_route: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Unpriced: + count: int + + +@dataclass(frozen=True, slots=True) +class MemberFallback: + router: bool + multipart: bool | None = None + resource_bound: bool | None = None + + +@dataclass(frozen=True, slots=True) +class ModelBudgets: + scope: Literal["Key", "User", "Member"] + entity_present: bool | None = None + budget_present: bool | None = None + mapping: bool | None = None + team_present: bool | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAllowed: + router: bool + disabled: bool | None = None + + +ReadPlan: TypeAlias = Literal[ + "ReadExemption", + "SkipBudgets", + "CheckBudgets", + "ReadRoute", + "ReadUnpriced", + "RejectUnpriced", + "ReadContentType", + "ReadResourceBinding", + "ReadFallbackSettings", + "CheckModelBudgets", + "ReadModelBudgetPresence", + "ReadModelBudgetShape", + "ReadModelBudgetToken", + "ReadModelBudgetUser", + "ReadModelBudgetTeam", + "AllowFallback", + "ReadFallbackDisabled", + "Continue", +] +_READ: Final = TypeAdapter(ReadPlan) + + +@native_boundary +def read(facts: BudgetSkip | Pricing | Unpriced | MemberFallback | ModelBudgets | FallbackAllowed) -> ReadPlan: + return _READ.validate_python(_native.request_auth_common_read(encode_facts(facts))) + + +@dataclass(frozen=True, slots=True) +class Facts: + model_present: bool + skip_budgets: bool + team_present: bool + user_present: bool + token_present: bool + user_id: str | None + agent_id: str | None + customer_present: bool + customer_models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Plan: + membership: bool + team_models: bool + member_models: bool + user_models: bool + customer_models: bool + agent_trace: bool + + +@dataclass(frozen=True, slots=True) +class BudgetFacts: + token_present: bool + group_count: int + customer_present: bool + customer_budget_present: bool + + +BudgetScope: TypeAlias = Literal[ + "Team", "TeamWindows", "KeyWindows", "TeamSoft", "Organization", "Tags", "Groups", "User", "Member", "Customer" +] + + +@dataclass(frozen=True, slots=True) +class CustomFacts: + end_user_id: str | None + end_user_present: bool + common_checks: bool + model_present: bool = False + router_present: bool = False + + +@dataclass(frozen=True, slots=True) +class CustomPlan: + read_end_user: bool + check_end_user: bool + check_model_access: bool + read_pricing: bool + + +_PLAN: Final = TypeAdapter(Plan) +_BUDGETS: Final = TypeAdapter(tuple[BudgetScope, ...]) +_CUSTOM: Final = TypeAdapter(CustomPlan) + + +@native_boundary +def plan(facts: Facts) -> Plan: + return _PLAN.validate_python(_native.request_auth_common_flow(asdict(facts))) + + +@native_boundary +def budgets(facts: BudgetFacts) -> tuple[BudgetScope, ...]: + return _BUDGETS.validate_python(_native.request_auth_common_budget_scopes(asdict(facts))) + + +@native_boundary +def custom(facts: CustomFacts) -> CustomPlan: + return _CUSTOM.validate_python(_native.request_auth_custom_flow(asdict(facts))) diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py index b002d28b701..8703d126772 100644 --- a/litellm/_auth/decision.py +++ b/litellm/_auth/decision.py @@ -3,6 +3,13 @@ from dataclasses import dataclass from datetime import datetime from typing import Protocol, TypeAlias +from litellm._auth.budget_pricing import ExemptionFacts +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from litellm._auth.model_access import ModelAccessDecision +from litellm._auth.model_budget_data import ModelBudgetSnapshot + @dataclass(frozen=True, slots=True) class RequestFacts: @@ -128,8 +135,22 @@ class ModelBudgetExceeded: entity_id: str | None +@dataclass(frozen=True, slots=True) +class ModelBudgetRead: + snapshot: ModelBudgetSnapshot + message: str = "" + entity_type: str | None = None + entity_id: str | None = None + + +@dataclass(frozen=True, slots=True) +class ModelAccessDenied: + decision: ModelAccessDecision + + Denial: TypeAlias = ( - MissingKey + ModelAccessDenied + | MissingKey | MalformedKey | UserAuthDisabled | ExpiredKey @@ -257,6 +278,11 @@ class InspectKey: @dataclass(frozen=True, slots=True) class EnforceModelAccess: + resume: "Callable[[KeyAndFallbackModelAccessFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class PrefetchAuthObjects: resume: "Callable[[], Step]" @@ -268,7 +294,7 @@ class LoadKeyOwner: @dataclass(frozen=True, slots=True) class CheckBudgetExemption: - resume: "Callable[[bool], Step]" + resume: "Callable[[ExemptionFacts], Step]" @dataclass(frozen=True, slots=True) @@ -332,13 +358,13 @@ class ReadModelBudgetScope: @dataclass(frozen=True, slots=True) class CheckKeyModelBudget: model: str - resume: "Callable[[ModelBudgetExceeded | None], Step]" + resume: "Callable[[ModelBudgetRead], Step]" @dataclass(frozen=True, slots=True) class AuthorizeBudgetFallback: model: str - resume: "Callable[[bool], Step]" + resume: "Callable[[FallbackAuthorizationFacts], Step]" @dataclass(frozen=True, slots=True) @@ -350,7 +376,7 @@ class RewriteRequestModel: @dataclass(frozen=True, slots=True) class CheckUserModelBudget: model: str - resume: "Callable[[ModelBudgetExceeded | None], Step]" + resume: "Callable[[ModelBudgetRead], Step]" IdentityEffect: TypeAlias = ( @@ -361,6 +387,7 @@ IdentityEffect: TypeAlias = ( | PrepareKey | InspectKey | EnforceModelAccess + | PrefetchAuthObjects | LoadKeyOwner | CheckBudgetExemption | LoadKeyContext @@ -390,7 +417,7 @@ class ReadTeamMemberModelBudgetScope: @dataclass(frozen=True, slots=True) class CheckTeamMemberModelBudget: model: str - resume: "Callable[[ModelBudgetExceeded | None], TeamMemberStep]" + resume: "Callable[[ModelBudgetRead], TeamMemberStep]" @dataclass(frozen=True, slots=True) @@ -402,7 +429,7 @@ class ReadTeamMemberFallbacks: @dataclass(frozen=True, slots=True) class VetTeamMemberFallback: target: FallbackTarget - resume: "Callable[[bool], TeamMemberStep]" + resume: "Callable[[FallbackCandidateFacts], TeamMemberStep]" @dataclass(frozen=True, slots=True) diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py index e447d7c0d04..00cfc9e2333 100644 --- a/litellm/_auth/driver.py +++ b/litellm/_auth/driver.py @@ -1,5 +1,6 @@ from typing import Protocol, assert_never +from litellm._auth.budget_pricing import ExemptionFacts from litellm._auth.decision import ( AlertKeyBudget, AlertKeySoftBudget, @@ -24,10 +25,11 @@ from litellm._auth.decision import ( LoadKeyContext, LoadKeyOwner, ModelBudgetEffect, - ModelBudgetExceeded, + ModelBudgetRead, ModelBudgetScope, Outcome, OwnerFacts, + PrefetchAuthObjects, PrepareKey, PresentedKey, ReadKeySpend, @@ -48,28 +50,31 @@ from litellm._auth.decision import ( TeamMemberStep, VetTeamMemberFallback, ) +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts class ModelBudgetHandler(Protocol): async def read_model_budget_scope(self) -> ModelBudgetScope: ... - async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: ... - async def authorize_budget_fallback(self, model: str) -> bool: ... + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: ... async def rewrite_request_model(self, model: str) -> None: ... - async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: ... class TeamMemberModelBudgetHandler(Protocol): async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: ... - async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: ... async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: ... - async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: ... + async def vet_team_member_fallback(self, target: FallbackTarget) -> FallbackCandidateFacts: ... async def reroute_team_member_request(self, model: str) -> None: ... @@ -87,11 +92,13 @@ class EffectHandler(ModelBudgetHandler, Protocol): async def inspect_key(self) -> KeyFacts: ... - async def enforce_model_access(self) -> None: ... + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: ... + + async def prefetch_auth_objects(self) -> None: ... async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ... - async def check_budget_exemption(self) -> bool: ... + async def check_budget_exemption(self) -> ExemptionFacts: ... async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: ... @@ -122,6 +129,7 @@ async def _answer(step: Effect, handler: EffectHandler) -> Step: | PrepareKey() | InspectKey() | EnforceModelAccess() + | PrefetchAuthObjects() | LoadKeyOwner() | CheckBudgetExemption() | LoadKeyContext() @@ -158,7 +166,9 @@ async def _answer_identity(step: IdentityEffect, handler: EffectHandler) -> Step case InspectKey(resume=resume): return resume(await handler.inspect_key()) case EnforceModelAccess(resume=resume): - await handler.enforce_model_access() + return resume(await handler.enforce_model_access()) + case PrefetchAuthObjects(resume=resume): + await handler.prefetch_auth_objects() return resume() case LoadKeyOwner(user_id=user_id, resume=resume): return resume(await handler.load_key_owner(user_id)) diff --git a/litellm/_auth/end_user_context.py b/litellm/_auth/end_user_context.py new file mode 100644 index 00000000000..457cf5e6af6 --- /dev/null +++ b/litellm/_auth/end_user_context.py @@ -0,0 +1,171 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class IdFacts: + value: str | None + + +@dataclass(frozen=True, slots=True) +class DefaultFacts: + prisma_present: bool + explicit_id: str | None + global_id: str | None + global_read: bool = False + + +@dataclass(frozen=True, slots=True) +class DefaultDecision: + read_global: bool + budget_id: str | None + + +@dataclass(frozen=True, slots=True) +class ApplyFacts: + own_budget_id: str | None + own_budget_present: bool + key_id: str | None + global_read: bool = False + global_id: str | None = None + default_read: bool = False + default_present: bool = False + + +ApplyDecision: TypeAlias = Literal["Keep", "ReadGlobal", "ReadDefault", "Copy"] + + +@dataclass(frozen=True, slots=True) +class RegistryFacts: + end_user_id: str + global_id: str | None + key_id: str | None + validate_id: bool + token_max_present: bool + registry_read: bool = False + registry: tuple[str, ...] | None = None + + +RegistryDecision: TypeAlias = Literal["Read", "Restricted", "Unrestricted"] +Field: TypeAlias = Literal["EndUserId", "Tpm", "Rpm", "Tpd", "MaxBudget", "Models", "Region"] + + +@dataclass(frozen=True, slots=True) +class Fields: + tpm: bool = False + rpm: bool = False + tpd: bool = False + max_budget: bool = False + models: bool = False + region: bool = False + + +@dataclass(frozen=True, slots=True) +class AssignmentFacts: + mode: Literal["Params", "Token"] + fields: Fields + + +@dataclass(frozen=True, slots=True) +class LookupFacts: + key_id: str | None + own_budget_present: bool = False + end_user_read: bool = False + end_user_present: bool = False + global_read: bool = False + global_id: str | None = None + default_read: bool = False + default_present: bool = False + token_max_present: bool = False + applied: bool = False + read_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class LookupDecision: + action: Literal["ReadEndUser", "ReadGlobal", "ReadDefault", "ApplyOwn", "ApplyDefault", "Done"] + set_max_budget: bool + project_budget: bool + + +@dataclass(frozen=True, slots=True) +class ValidationFacts: + mode: Literal["Resolve", "Exists"] + end_user_id: str | None + enabled: bool = False + prisma_present: bool = False + global_id: str | None = None + key_id: str | None = None + cache_read: bool = False + cached: str | None = None + end_user_read: bool = False + end_user_present: bool = False + user_read: bool = False + user_present: bool = False + cache_written: bool = False + + +@dataclass(frozen=True, slots=True) +class ValidationDecision: + read: Literal["Cache", "EndUser", "User", "WriteCache"] | None + end_user_id: str | None + valid: bool + cache_value: str | None + cache_ttl: int + user_email: str | None + + +_ID: Final = TypeAdapter(IdFacts) +_STRING: Final = TypeAdapter[str | None](str | None) +_DEFAULT: Final = TypeAdapter(DefaultFacts) +_DEFAULT_RESULT: Final = TypeAdapter(DefaultDecision) +_APPLY: Final = TypeAdapter(ApplyFacts) +_APPLY_RESULT: Final = TypeAdapter[ApplyDecision](ApplyDecision) +_REGISTRY: Final = TypeAdapter(RegistryFacts) +_REGISTRY_RESULT: Final = TypeAdapter[RegistryDecision](RegistryDecision) +_ASSIGNMENTS: Final = TypeAdapter(AssignmentFacts) +_FIELDS: Final = TypeAdapter(tuple[Field, ...]) +_LOOKUP: Final = TypeAdapter(LookupFacts) +_LOOKUP_RESULT: Final = TypeAdapter(LookupDecision) +_VALIDATION: Final = TypeAdapter(ValidationFacts) +_VALIDATION_RESULT: Final = TypeAdapter(ValidationDecision) + + +@native_boundary +def normalize_id(facts: IdFacts) -> str | None: + return _STRING.validate_python(_native.request_auth_end_user_normalize_id(_ID.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def default_id(facts: DefaultFacts) -> DefaultDecision: + return _DEFAULT_RESULT.validate_python(_native.request_auth_end_user_default_id(_DEFAULT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def apply_default(facts: ApplyFacts) -> ApplyDecision: + return _APPLY_RESULT.validate_python(_native.request_auth_end_user_apply(_APPLY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def registry(facts: RegistryFacts) -> RegistryDecision: + return _REGISTRY_RESULT.validate_python(_native.request_auth_end_user_registry(_REGISTRY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def assignments(facts: AssignmentFacts) -> tuple[Field, ...]: + return _FIELDS.validate_python(_native.request_auth_end_user_assignments(_ASSIGNMENTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def lookup(facts: LookupFacts) -> LookupDecision: + return _LOOKUP_RESULT.validate_python(_native.request_auth_end_user_lookup(_LOOKUP.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def validation(facts: ValidationFacts) -> ValidationDecision: + return _VALIDATION_RESULT.validate_python(_native.request_auth_end_user_validation(_VALIDATION.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts diff --git a/litellm/_auth/failure.py b/litellm/_auth/failure.py new file mode 100644 index 00000000000..66b13582096 --- /dev/null +++ b/litellm/_auth/failure.py @@ -0,0 +1,39 @@ +from collections.abc import Callable +from functools import wraps +from typing import ParamSpec, TypeVar + +_P = ParamSpec("_P") # rebind-ok: ParamSpec declarations cannot be annotated Final +_R = TypeVar("_R") # rebind-ok: TypeVar declarations cannot be annotated Final + + +class NativeAuthFailure(Exception): + pass + + +class NativeAuthTypeError(NativeAuthFailure, TypeError): + pass + + +class NativeAuthValueError(NativeAuthFailure, ValueError): + pass + + +class NativeAuthRuntimeError(NativeAuthFailure, RuntimeError): + pass + + +def native_boundary(function: Callable[_P, _R]) -> Callable[_P, _R]: + @wraps(function) + def call(*args: _P.args, **kwargs: _P.kwargs) -> _R: # kwargs-ok: ParamSpec preserves the wrapped signature + try: + return function(*args, **kwargs) + except NativeAuthFailure: + raise + except TypeError as error: + raise NativeAuthTypeError(str(error)) from error + except ValueError as error: + raise NativeAuthValueError(str(error)) from error + except Exception as error: # noqa: BLE001 # native faults must bypass legacy IO fail-open catches + raise NativeAuthRuntimeError(str(error)) from error + + return call diff --git a/litellm/_auth/fallback_access.py b/litellm/_auth/fallback_access.py new file mode 100644 index 00000000000..057f3a47745 --- /dev/null +++ b/litellm/_auth/fallback_access.py @@ -0,0 +1,58 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class GroupedModelScope: + facts: ModelAccessFacts + group_ids: tuple[str, ...] = () + group_models: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAuthorizationFacts: + key: GroupedModelScope | None = None + refuse_agent_keys: bool = False + agent_id: str | None = None + team: GroupedModelScope | None = None + embedded_team: ModelAccessFacts | None = None + assigned_models: tuple[str, ...] | None = None + member: ModelAccessFacts | None = None + project: ModelAccessFacts | None = None + team_present: bool = False + embedded_team_grants: tuple[str, ...] = () + member_grants: tuple[str, ...] = () + project_grants: tuple[str, ...] = () + assigned_scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAuthorizationDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadKey", + "ReadTeam", + "ReadEmbeddedTeam", + "ReadMember", + "ReadProject", + "ReadKeyGroups", + "ReadTeamGroups", + "ReadAssignedGroups", + "ReadAssignedScope", + ] + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(FallbackAuthorizationDecision) + + +@native_boundary +def fallback_authorization(facts: FallbackAuthorizationFacts) -> FallbackAuthorizationDecision: + return _DECISION.validate_python(_native.request_auth_fallback_access(asdict(facts))) diff --git a/litellm/_auth/fallback_candidate.py b/litellm/_auth/fallback_candidate.py new file mode 100644 index 00000000000..39c50cdfd63 --- /dev/null +++ b/litellm/_auth/fallback_candidate.py @@ -0,0 +1,88 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.common_budgets import BudgetFacts +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.model_budget_data import ModelBudgetEntity, ModelBudgetSnapshot +from litellm._auth.routes import RouteContext +from litellm._auth.unpriced_models import UnpricedModelFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class CandidateBudgetScope: + entity: ModelBudgetEntity + identity_present: bool + configured: bool + + +@dataclass(frozen=True, slots=True) +class CandidateBudgetRead: + model_index: int + entity: ModelBudgetEntity + snapshot: ModelBudgetSnapshot + + +@dataclass(frozen=True, slots=True) +class FallbackCandidateFacts: + router_present: bool + user_present: bool + team_present: bool + membership: Literal["Unread", "Loaded", "Failed"] = "Unread" + block_unpriced: bool = False + route: RouteContext | None = None + unpriced: UnpricedModelFacts | None = None + customer_present: bool = False + customer_grants: tuple[str, ...] | None = None + customer: ModelAccessFacts | None = None + price: ModelPriceFacts | None = None + budget_models: tuple[str, ...] | None = None + budget_scopes: tuple[CandidateBudgetScope, ...] | None = None + budgets: tuple[CandidateBudgetRead, ...] = () + authorization: FallbackAuthorizationFacts | None = None + attribution_reset: bool = False + attribution: ModelGroupAttributionFacts | None = None + attribution_stamped: bool = False + group_budgets_read: bool = False + group_budgets: tuple[BudgetFacts, ...] = () + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class FallbackCandidateDecision: + kind: Literal[ + "Complete", + "ReadMembership", + "ReadRoute", + "ReadUnpriced", + "ReadCustomerGrants", + "ReadCustomer", + "ReadPrice", + "ReadBudgetModels", + "ReadBudgetScopes", + "ReadBudget", + "ReadAuthorization", + "ResetAttribution", + "ReadAttribution", + "StampAttribution", + "ReadGroupBudgets", + ] + usable: bool = False + log_failure: bool = False + model_index: int = 0 + entity: ModelBudgetEntity = "Key" + groups: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(FallbackCandidateDecision) + + +@native_boundary +def fallback_candidate(facts: FallbackCandidateFacts) -> FallbackCandidateDecision: + return _DECISION.validate_python(_native.request_auth_fallback_candidate(asdict(facts))) diff --git a/litellm/_auth/key_model_access.py b/litellm/_auth/key_model_access.py new file mode 100644 index 00000000000..00666dd2a57 --- /dev/null +++ b/litellm/_auth/key_model_access.py @@ -0,0 +1,45 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class KeyAndFallbackModelAccessFacts: + config_nonempty: bool + key_models: tuple[str, ...] + primary_loaded: bool = False + primary: GroupedModelScope | None = None + fallback_names: tuple[str, ...] | None = None + fallback_scopes: tuple[GroupedModelScope, ...] = () + routing: tuple[Literal["Routed", "Failed"], ...] = () + + +@dataclass(frozen=True, slots=True) +class KeyAndFallbackModelAccessDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadPrimary", + "ReadPrimaryGroups", + "ReadFallbackNames", + "ReadFallbackScope", + "ReadFallbackGroups", + "RouteFallback", + "RoutingFailure", + ] + index: int = 0 + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(KeyAndFallbackModelAccessDecision) + + +@native_boundary +def key_and_fallback_model_access(facts: KeyAndFallbackModelAccessFacts) -> KeyAndFallbackModelAccessDecision: + return _DECISION.validate_python(_native.request_auth_key_model_access(asdict(facts))) diff --git a/litellm/_auth/managed.py b/litellm/_auth/managed.py new file mode 100644 index 00000000000..19c9c64c960 --- /dev/null +++ b/litellm/_auth/managed.py @@ -0,0 +1,129 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class BoundIdentity: + active: bool + revision: str + + +@dataclass(frozen=True, slots=True) +class Agent: + agent_id: str + enabled: bool + identity_managed: bool + identity: BoundIdentity | None + execution_mode: str + + +@dataclass(frozen=True, slots=True) +class Context: + agent_id: str + binding_revision: str | None + mode: str + user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Actor: + agent: Agent + context: Context | None + + +@dataclass(frozen=True, slots=True) +class ActorLookup: + agent_id: str | None + has_database: bool + registry_loaded: bool = False + registry_managed: bool = False + agent_loaded: bool = False + agent: Agent | None = None + retired: bool | None = None + context: Context | None = None + jwt_claims: bool = False + delegation_verified: bool = False + human_grants: tuple[str, ...] | None = None + policy_assigned: bool = False + + +@dataclass(frozen=True, slots=True) +class UnionGrants: + grants: tuple[tuple[str, ...] | None, ...] + + +@dataclass(frozen=True, slots=True) +class IntersectGrants: + left: tuple[str, ...] | None + right: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class DeclaredGrants: + direct: tuple[str, ...] + declared_groups: tuple[str, ...] + unified_groups: tuple[str, ...] + group_agents: tuple[str, ...] | None = None + unified_agents: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ManagedAgentGrants: + permission_present: bool + mode: str | None + user_present: bool + own: tuple[str, ...] | None = None + ceilings: tuple[tuple[str, ...], ...] | None = None + caller_loaded: bool = False + caller: tuple[str, ...] | None = None + human: tuple[str, ...] | None = None + + +ManagedDenial: TypeAlias = Literal[ + "Disabled", + "MissingContext", + "ChangedIdentity", + "ExecutionMode", + "MissingHuman", + "DatabaseRequired", + "Removed", + "VerifiedIdentityRequired", + "HumanNotPermitted", +] + + +@dataclass(frozen=True, slots=True) +class ManagedDecision: + action: Literal[ + "Allow", + "Deny", + "ReadRegistry", + "ReadAgent", + "ReadRetired", + "AssignPolicy", + "ReadHumanGrants", + "Grants", + "ReadDeclaredGroups", + "ReadUnifiedGroups", + "ReadOwn", + "ReadCeilings", + "ReadCaller", + ] + denial: ManagedDenial | None + grants: tuple[str, ...] | None + + +_DECISION: Final = TypeAdapter(ManagedDecision) + + +@native_boundary +def decide( + facts: Actor | ActorLookup | UnionGrants | IntersectGrants | DeclaredGrants | ManagedAgentGrants, +) -> ManagedDecision: + return _DECISION.validate_python(_native.request_auth_managed(encode_facts(facts))) diff --git a/litellm/_auth/managed_models.py b/litellm/_auth/managed_models.py new file mode 100644 index 00000000000..2a41347d329 --- /dev/null +++ b/litellm/_auth/managed_models.py @@ -0,0 +1,95 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ManagedModelFacts: + purpose: Literal["All", "Ceilings", "Caller"] + model_present: bool + token_present: bool + managed_loaded: bool = False + managed_present: bool = False + managed_models_valid: bool = False + managed_models_count: int = 0 + managed: ModelAccessFacts | None = None + ceilings: tuple[tuple[str, ...], ...] | None = None + aliases_loaded: bool = False + scopes: tuple[ModelAccessFacts, ...] = () + caller_loaded: bool = False + caller_present: bool = False + team_loaded: bool = False + team_present: bool = False + team: GroupedModelScope | None = None + member: MemberModelFacts | None = None + user_loaded: bool = False + user_present: bool = False + user_grants: tuple[str, ...] = () + user: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ManagedModelDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadManaged", + "ReadManagedScope", + "ReadCeilings", + "ReadAliases", + "ReadCeilingScope", + "ReadCaller", + "ReadTeam", + "ReadTeamScope", + "ReadTeamGroups", + "ReadMember", + "ReadUser", + "ReadUserScope", + ] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + index: int = 0 + + +@dataclass(frozen=True, slots=True) +class AgentCeilingSourceFacts: + agent_present: bool + managed_loaded: bool = False + managed_present: bool = False + loaded: bool = False + count: int = 0 + + +@dataclass(frozen=True, slots=True) +class AgentCeilingSourceDecision: + kind: Literal["ReadManaged", "ReadManagedCeilings", "ReadUnmanagedCeiling", "Complete"] + indices: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(ManagedModelDecision) +_CEILINGS: Final = TypeAdapter(AgentCeilingSourceDecision) +_GRANTS: Final = TypeAdapter(frozenset[str]) + + +@native_boundary +def managed_models(facts: ManagedModelFacts) -> ManagedModelDecision: + return _DECISION.validate_python(_native.request_auth_managed_models(asdict(facts))) + + +@native_boundary +def agent_ceiling_source(facts: AgentCeilingSourceFacts) -> AgentCeilingSourceDecision: + return _CEILINGS.validate_python(_native.request_auth_agent_ceiling_source(asdict(facts))) + + +@native_boundary +def managed_model_grants( + models: tuple[str | None, ...] | None, ceilings: tuple[tuple[str, ...], ...] +) -> frozenset[str]: + return _GRANTS.validate_python(_native.request_auth_managed_model_grants({"models": models, "ceilings": ceilings})) diff --git a/litellm/_auth/member_models.py b/litellm/_auth/member_models.py new file mode 100644 index 00000000000..a2e845b69f4 --- /dev/null +++ b/litellm/_auth/member_models.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class MemberModelFacts: + user_present: bool + team_present: bool + membership_loaded: bool = False + grants: tuple[str, ...] | None = None + scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class MemberModelDecision: + kind: Literal["Allow", "Deny", "ReadMembership", "ReadGrants", "ReadScope"] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + + +_DECISION: Final = TypeAdapter(MemberModelDecision) + + +@native_boundary +def member_models(facts: MemberModelFacts) -> MemberModelDecision: + return _DECISION.validate_python(_native.request_auth_member_models(asdict(facts))) diff --git a/litellm/_auth/model_access.py b/litellm/_auth/model_access.py new file mode 100644 index 00000000000..e3c29ad4c64 --- /dev/null +++ b/litellm/_auth/model_access.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import asdict, dataclass, field +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +ModelObject: TypeAlias = Literal["user", "customer", "team", "key", "org", "project", "agent"] +ModelReadKind: TypeAlias = Literal["Compaction", "Aliases", "AliasTarget", "Deployment", "Groups", "Names", "Provider"] + + +@dataclass(frozen=True, slots=True) +class ModelProvider: + stripped_model: str + provider: str + + +@dataclass(frozen=True, slots=True) +class ModelRecord: + access_groups: tuple[str, ...] = () + provider: ModelProvider | None = None + deployment_team: str | None = None + deployment_public_name: str | None = None + + +@dataclass(frozen=True, slots=True) +class ModelReading: + kind: ModelReadKind + model: str + values: tuple[str, ...] = () + value: str | None = None + team: str | None = None + provider: ModelProvider | None = None + + +@dataclass(frozen=True, slots=True) +class ModelEvaluation: + grants: tuple[str, ...] + reads: tuple[ModelReading, ...] + aliases: ModelAliases | None = None + + +@dataclass(frozen=True, slots=True) +class ModelRegistry: + present: bool = False + names: tuple[str, ...] = () + aliases: Mapping[str, str] = field(default_factory=dict[str, str]) + records: Mapping[str, ModelRecord] = field(default_factory=dict[str, ModelRecord]) + compaction_parents: Mapping[str, str] = field(default_factory=dict[str, str]) + helper_present: bool | None = None + evaluations: tuple[ModelEvaluation, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelAliases: + team: Mapping[str, str] = field(default_factory=dict[str, str]) + key: Mapping[str, str] = field(default_factory=dict[str, str]) + global_: Mapping[str, str] = field(default_factory=dict[str, str]) + team_order: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelAccessFacts: + requested: tuple[str, ...] + grants: tuple[str, ...] = () + team_id: str | None = None + aliases: ModelAliases = field(default_factory=ModelAliases) + registry: ModelRegistry = field(default_factory=ModelRegistry) + object_type: ModelObject = "user" + fallback_depth: int = 0 + max_depth: int = 5 + is_list: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelAccessDecision: + kind: Literal[ + "Allow", + "Deny", + "Depth", + "NoDefaultModels", + "NoAgentModels", + "EmptyAgentCeiling", + "PersonalTeam", + "ReadGroups", + "ReadScope", + "ReadGroupScope", + ] + model: str | None = None + grants: tuple[str, ...] = () + object_type: ModelObject = "user" + + +@dataclass(frozen=True, slots=True) +class AssignedModelGroup: + assigned_team_ids: tuple[str, ...] + assigned_key_ids: tuple[str, ...] + models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class PersonalModelTeam: + blocked: bool + available: bool + team: ModelAccessFacts + member: ModelAccessFacts | None = None + + +_DECISION: Final = TypeAdapter(ModelAccessDecision) +_MODELS: Final = TypeAdapter(tuple[str, ...]) +_BOOL: Final = TypeAdapter(bool) + + +@native_boundary +def model_access(facts: ModelAccessFacts, *, kind: Literal["Object", "Customer"] = "Object") -> ModelAccessDecision: + return _DECISION.validate_python(_native.request_auth_model_access({kind: asdict(facts)})) + + +@native_boundary +def grouped_model_access( + facts: ModelAccessFacts, + group_ids: tuple[str, ...], + group_models: tuple[str, ...] | None = None, + *, + merge: bool = False, +) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access( + { + "Grouped": { + "facts": asdict(facts), + "group_ids": group_ids, + "group_models": group_models, + "merge": merge, + } + } + ) + ) + + +@native_boundary +def user_model_access(facts: ModelAccessFacts, *, present: bool) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access({"User": {"facts": asdict(facts), "present": present}}) + ) + + +@native_boundary +def user_model_precheck(*, present: bool, grants: tuple[str, ...]) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access({"UserPrecheck": {"present": present, "grants": grants}}) + ) + + +@native_boundary +def agent_model_access( + facts: ModelAccessFacts, managed_grants: tuple[str, ...] | None, ceilings: tuple[tuple[str, ...], ...] +) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access( + { + "Agent": { + "facts": asdict(facts), + "managed_grants": managed_grants, + "ceilings": ceilings, + } + } + ) + ) + + +@native_boundary +def personal_model_access(requested: tuple[str, ...], teams: tuple[PersonalModelTeam, ...]) -> ModelAccessDecision: + return _DECISION.validate_python( + _native.request_auth_model_access( + { + "Personal": { + "requested": requested, + "teams": tuple(asdict(team) for team in teams), + } + } + ) + ) + + +@native_boundary +def model_candidates(facts: ModelAccessFacts) -> tuple[str, ...]: + return _MODELS.validate_python(_native.request_auth_model_access({"Candidates": asdict(facts)})) + + +@native_boundary +def model_access_helper(model: str, facts: ModelAccessFacts) -> bool: + return _BOOL.validate_python( + _native.request_auth_model_access({"Helper": {"model": model, "facts": asdict(facts)}}) + ) + + +@native_boundary +def key_model_grants(grants: tuple[str, ...], team_id: str | None, team_grants: tuple[str, ...]) -> tuple[str, ...]: + return _MODELS.validate_python( + _native.request_auth_model_access( + { + "KeyGrants": { + "grants": grants, + "team_id": team_id, + "team_grants": team_grants, + } + } + ) + ) + + +@native_boundary +def assigned_model_grants( + groups: tuple[AssignedModelGroup, ...], team_id: str | None, key_id: str | None +) -> tuple[str, ...]: + return _MODELS.validate_python( + _native.request_auth_model_access( + { + "AssignedGroups": { + "groups": tuple(asdict(group) for group in groups), + "team_id": team_id, + "key_id": key_id, + } + } + ) + ) + + +@native_boundary +def personal_model_access_applies( + *, + enabled: bool, + token_present: bool, + team_id: str | None, + user_id: str | None, + master: bool, + role: str | None, + api_key_present: bool, + inheritance: bool = False, +) -> bool: + return _BOOL.validate_python( + _native.request_auth_model_access( + { + "PersonalApplies": { + "enabled": enabled, + "token_present": token_present, + "team_id": team_id, + "user_id": user_id, + "master": master, + "role": role, + "api_key_present": api_key_present, + "inheritance": inheritance, + } + } + ) + ) diff --git a/litellm/_auth/model_attribution.py b/litellm/_auth/model_attribution.py new file mode 100644 index 00000000000..aa2dd98fdd4 --- /dev/null +++ b/litellm/_auth/model_attribution.py @@ -0,0 +1,46 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelGroupAttributionFacts: + model_present: bool = False + token_present: bool = False + router_present: bool = False + database_present: bool = False + team_present: bool = False + user_present: bool = False + membership_loaded: bool = False + token_org: str | None = None + team_org: str | None = None + registry_loaded: bool = False + registry: tuple[str, ...] | None = None + covering: tuple[str, ...] | None = None + grants: tuple[tuple[str, ...], ...] = () + org_failed: bool = False + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelGroupAttributionDecision: + kind: Literal[ + "Complete", "ReadRegistry", "ReadCovering", "ReadKey", "ReadTeam", "ReadMember", "ReadProject", "ReadOrg" + ] + groups: tuple[str, ...] = () + stamp: bool = False + lookup: bool = False + include: bool = False + id: str | None = None + + +_DECISION: Final = TypeAdapter(ModelGroupAttributionDecision) + + +@native_boundary +def model_group_attribution(facts: ModelGroupAttributionFacts) -> ModelGroupAttributionDecision: + return _DECISION.validate_python(_native.request_auth_model_attribution(asdict(facts))) diff --git a/litellm/_auth/model_budget_data.py b/litellm/_auth/model_budget_data.py new file mode 100644 index 00000000000..b0b6f95c770 --- /dev/null +++ b/litellm/_auth/model_budget_data.py @@ -0,0 +1,86 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + +ModelBudgetEntity: TypeAlias = Literal["Key", "User", "Team", "TeamMember", "EndUser"] + + +@dataclass(frozen=True, slots=True) +class ModelBudgetEntry: + name: str + max_budget: float | None + duration: str | None + duration_seconds: float | None + + +@dataclass(frozen=True, slots=True) +class ModelBudgetResolveFacts: + model: str + bedrock_base: str + bedrock_provider: str + entries: tuple[ModelBudgetEntry, ...] + key_entries: tuple[ModelBudgetEntry, ...] = () + team: bool = False + access_groups: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelBudgetResolution: + read_groups: bool + entry: ModelBudgetEntry | None + named_model: bool + + +@dataclass(frozen=True, slots=True) +class ModelBudgetSnapshot: + max_budget: float | None = None + local_spend: float | None = None + legacy_spend: float | None = None + shared_spend: float | None = None + shared_started_at: float | None = None + shared_window_seconds: float | None = None + now: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class ModelBudgetQueryFacts: + entity: ModelBudgetEntity + snapshot: ModelBudgetSnapshot + duration_seconds: float | None = None + process_uptime: float = 0.0 + same_counter: bool = False + redis_present: bool = False + shared_read: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelBudgetDecision: + read: Literal["Local", "Legacy", "Shared"] | None + denied: bool + current_spend: float + + +_RESOLVE: Final = TypeAdapter(ModelBudgetResolveFacts) +_RESOLUTION: Final = TypeAdapter(ModelBudgetResolution) +_QUERY: Final = TypeAdapter(ModelBudgetQueryFacts) +_SNAPSHOT: Final = TypeAdapter(ModelBudgetSnapshot) +_DECISION: Final = TypeAdapter(ModelBudgetDecision) + + +@native_boundary +def resolve(facts: ModelBudgetResolveFacts) -> ModelBudgetResolution: + return _RESOLUTION.validate_python(_native.request_auth_model_budget_resolve(_RESOLVE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def query(facts: ModelBudgetQueryFacts) -> ModelBudgetDecision: + return _DECISION.validate_python(_native.request_auth_model_budget_query(_QUERY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@native_boundary +def check(facts: ModelBudgetSnapshot) -> ModelBudgetDecision: + return _DECISION.validate_python(_native.request_auth_model_budget_check(_SNAPSHOT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary diff --git a/litellm/_auth/model_groups.py b/litellm/_auth/model_groups.py new file mode 100644 index 00000000000..43a92d972a2 --- /dev/null +++ b/litellm/_auth/model_groups.py @@ -0,0 +1,48 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class GroupResources: + assigned_team_ids: tuple[str, ...] + assigned_key_ids: tuple[str, ...] + resources: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class GroupLookup: + kind: Literal["Found", "Missing", "Failed"] + group: GroupResources | None = None + + +@dataclass(frozen=True, slots=True) +class ResourceGroupFacts: + subject_present: bool + database_present: bool + cache_present: bool + group_ids: tuple[str, ...] + team_id: str | None + key_id: str | None + check_db_only: bool = False + strict_errors: bool = False + loaded: tuple[GroupLookup, ...] = () + + +@dataclass(frozen=True, slots=True) +class ResourceGroupDecision: + kind: Literal["Read", "Raise", "Complete"] + index: int = 0 + resources: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(ResourceGroupDecision) + + +@native_boundary +def resource_group_decision(facts: ResourceGroupFacts) -> ResourceGroupDecision: + return _DECISION.validate_python(_native.request_auth_model_groups(asdict(facts))) diff --git a/litellm/_auth/model_helpers.py b/litellm/_auth/model_helpers.py new file mode 100644 index 00000000000..8de363b28ac --- /dev/null +++ b/litellm/_auth/model_helpers.py @@ -0,0 +1,37 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelProvider +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelHelperFacts: + kind: Literal["Group", "Helper"] + model: str + grants: tuple[str, ...] | None + team_id: str | None = None + team_aliases: tuple[str, ...] = () + router_present: bool = False + groups: tuple[str, ...] | None = None + names: tuple[str, ...] | None = None + provider_loaded: bool = False + provider: ModelProvider | None = None + provider_reads: tuple[ModelProvider | None, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ModelHelperDecision: + kind: Literal["Complete", "ReadGroups", "ReadNames", "ReadProvider"] + allowed: bool = False + + +_DECISION: Final = TypeAdapter(ModelHelperDecision) + + +@native_boundary +def model_helper(facts: ModelHelperFacts) -> ModelHelperDecision: + return _DECISION.validate_python(_native.request_auth_model_helpers(asdict(facts))) diff --git a/litellm/_auth/model_projection.py b/litellm/_auth/model_projection.py new file mode 100644 index 00000000000..63022394416 --- /dev/null +++ b/litellm/_auth/model_projection.py @@ -0,0 +1,26 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts, ModelReading, ModelReadKind +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelProjectionDecision: + kind: Literal["Complete", "Read"] + read: ModelReadKind | None = None + model: str | None = None + decision: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(ModelProjectionDecision) + + +@native_boundary +def model_projection(facts: ModelAccessFacts, reads: tuple[ModelReading, ...]) -> ModelProjectionDecision: + return _DECISION.validate_python( + _native.request_auth_model_projection(asdict(facts), tuple(asdict(read) for read in reads)) + ) diff --git a/litellm/_auth/org_routes.py b/litellm/_auth/org_routes.py new file mode 100644 index 00000000000..66563e6a5c1 --- /dev/null +++ b/litellm/_auth/org_routes.py @@ -0,0 +1,95 @@ +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 + + +@dataclass(frozen=True, slots=True) +class OrganizationMembership: + organization_id: str | None + role: str | None + + +@dataclass(frozen=True, slots=True) +class OrganizationRouteFacts: + route: str + user_present: bool + user_role: str | None + memberships: tuple[OrganizationMembership, ...] | None + organization_id: str | None + org_admin_routes: tuple[str, ...] + organization_id_present: bool = False + identifier_pending: bool = False + hash_failure: str | None = None + request_pending: bool = False + memberships_pending: bool = False + + +OrgRouteDenial: TypeAlias = Literal[ + "CreateOrganization", + "NoMemberships", + "MissingOrganization", + "MissingRole", + "WrongRole", + "TeamMissingOrganization", + "TeamWrongRole", +] + + +@dataclass(frozen=True, slots=True) +class OrganizationRouteDecision: + denial: OrgRouteDenial | None + organization_role: str | None + organization_ids: tuple[str, ...] + read_identifier: bool = False + hash_failure: str | None = None + read_request: bool = False + read_memberships: bool = False + + +_DECISION: Final = TypeAdapter(OrganizationRouteDecision) + + +@native_boundary +def organization_route(facts: OrganizationRouteFacts) -> OrganizationRouteDecision: + return _DECISION.validate_python(_native.request_auth_org_route(asdict(facts))) + + +@native_boundary +def user_is_org_admin(memberships: tuple[OrganizationMembership, ...], organizations: tuple[str, ...]) -> bool: + return _native.request_auth_org_admin(tuple(asdict(item) for item in memberships), organizations) + + +@dataclass(frozen=True, slots=True) +class Identifier: + text: str | None + hash_failure: str | None = None + + +@dataclass(frozen=True, slots=True) +class AdminFacts: + memberships: tuple[OrganizationMembership, ...] | None + request_read: bool = False + singular_present: bool = False + plural_count: int = 0 + candidate_index: int = 0 + candidate: Identifier | None = None + + +@dataclass(frozen=True, slots=True) +class AdminDecision: + action: Literal["ReadRequest", "ReadSingular", "ReadPlural", "Allow", "Deny", "HashFailure"] + candidate_index: int + plural_index: int + hash_failure: str | None + + +_ADMIN: Final = TypeAdapter(AdminDecision) + + +@native_boundary +def admin(facts: AdminFacts) -> AdminDecision: + return _ADMIN.validate_python(_native.request_auth_org_admin_raw(asdict(facts))) diff --git a/litellm/_auth/owner_teams.py b/litellm/_auth/owner_teams.py new file mode 100644 index 00000000000..be1919eb626 --- /dev/null +++ b/litellm/_auth/owner_teams.py @@ -0,0 +1,53 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.personal_budget import PersonalTeamBudgetFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class OwnerTeamLoad: + present: bool + failed: bool = False + blocked: bool = False + org: str | None = None + + +@dataclass(frozen=True, slots=True) +class OwnerTeamBudget: + index: int + facts: PersonalTeamBudgetFacts + + +@dataclass(frozen=True, slots=True) +class OwnerTeamFacts: + purpose: Literal["Models", "Grants"] + ids: tuple[str, ...] + ui_team_id: str + user_present: bool = True + database_present: bool = True + org: str | None = None + loaded: bool = False + loads: tuple[OwnerTeamLoad, ...] = () + budgets: tuple[OwnerTeamBudget, ...] = () + + +@dataclass(frozen=True, slots=True) +class OwnerTeamDecision: + kind: Literal["ReadTeams", "ReadBudgets", "Complete"] + ids: tuple[str, ...] = () + indices: tuple[int, ...] = () + selected: tuple[int, ...] = () + withheld: tuple[int, ...] = () + unreadable: bool = False + + +_DECISION: Final = TypeAdapter(OwnerTeamDecision) + + +@native_boundary +def owner_teams(facts: OwnerTeamFacts) -> OwnerTeamDecision: + return _DECISION.validate_python(_native.request_auth_owner_teams(asdict(facts))) diff --git a/litellm/_auth/personal_budget.py b/litellm/_auth/personal_budget.py new file mode 100644 index 00000000000..b82fe1d1937 --- /dev/null +++ b/litellm/_auth/personal_budget.py @@ -0,0 +1,66 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.common_budgets import BudgetFacts, MemberLimitFacts +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalTeamBudgetFacts: + user_present: bool = False + team_id_present: bool = True + check_db_only: bool = False + token_org: str | None = None + team_org: str | None = None + source_ready: bool = False + team: BudgetFacts | None = None + team_alerted: bool = False + window_count: int | None = None + windows: tuple[BudgetFacts, ...] = () + membership_loaded: bool = False + member_limit: MemberLimitFacts | None = None + default_id: str | None = None + default_loaded: bool = False + member: BudgetFacts | None = None + member_alerted: bool = False + org_loaded: bool = False + org_failed: bool = False + org: BudgetFacts | None = None + org_alerted: bool = False + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class PersonalTeamBudgetDecision: + kind: Literal[ + "Complete", + "Source", + "ReadTeam", + "AlertTeam", + "ReadWindows", + "ReadWindow", + "ReadMembership", + "ReadMemberLimit", + "ReadDefault", + "ReadMember", + "AlertMember", + "ReadOrg", + "AlertOrg", + ] + allowed: bool = False + log_failure: bool = False + org_id: str | None = None + id: str = "" + index: int = 0 + limit: float | None = None + + +_DECISION: Final = TypeAdapter(PersonalTeamBudgetDecision) + + +@native_boundary +def personal_team_budget(facts: PersonalTeamBudgetFacts) -> PersonalTeamBudgetDecision: + return _DECISION.validate_python(_native.request_auth_personal_budget(asdict(facts))) diff --git a/litellm/_auth/personal_model_listing.py b/litellm/_auth/personal_model_listing.py new file mode 100644 index 00000000000..e90cc43c388 --- /dev/null +++ b/litellm/_auth/personal_model_listing.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm._auth.personal_owner import PersonalOwnerFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelListingFacts: + owner: PersonalOwnerFacts + names: tuple[str, ...] | None = None + models: tuple[LoadedTeamModelFacts, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelListingDecision: + kind: Literal["ReadOwner", "ReadNames", "ReadModel", "Complete"] + index: int = 0 + teams: tuple[int, ...] = () + selected: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelListingDecision) + + +@native_boundary +def personal_model_listing(facts: PersonalModelListingFacts) -> PersonalModelListingDecision: + return _DECISION.validate_python(_native.request_auth_personal_model_listing(asdict(facts))) diff --git a/litellm/_auth/personal_model_request.py b/litellm/_auth/personal_model_request.py new file mode 100644 index 00000000000..aa885980426 --- /dev/null +++ b/litellm/_auth/personal_model_request.py @@ -0,0 +1,33 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.owner_teams import OwnerTeamFacts +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelRequestFacts: + requested: tuple[str, ...] + owner_present: bool + teams: OwnerTeamFacts | None = None + models: tuple[LoadedTeamModelFacts, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelRequestDecision: + kind: Literal["Allow", "Deny", "ReadTeams", "ReadModel"] + model: str = "" + index: int = 0 + teams: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelRequestDecision) + + +@native_boundary +def personal_model_request(facts: PersonalModelRequestFacts) -> PersonalModelRequestDecision: + return _DECISION.validate_python(_native.request_auth_personal_model_request(asdict(facts))) diff --git a/litellm/_auth/personal_models.py b/litellm/_auth/personal_models.py new file mode 100644 index 00000000000..88a13391eec --- /dev/null +++ b/litellm/_auth/personal_models.py @@ -0,0 +1,45 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelScope: + user_present: bool + team_id_present: bool + team: GroupedModelScope | None = None + member_grants: tuple[str, ...] | None = None + member: ModelAccessFacts | None = None + io_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class LoadedTeamModelFacts: + present: tuple[bool, ...] + scopes: tuple[PersonalModelScope, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelDecision: + kind: Literal["Complete", "ReadTeam", "ReadGroups", "ReadMembership", "ReadMember", "ReadScopes"] + allowed: bool = False + indices: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelDecision) + + +@native_boundary +def personal_model_scope(facts: PersonalModelScope) -> PersonalModelDecision: + return _DECISION.validate_python(_native.request_auth_personal_models({"Scope": asdict(facts)})) + + +@native_boundary +def loaded_team_models(facts: LoadedTeamModelFacts) -> PersonalModelDecision: + return _DECISION.validate_python(_native.request_auth_personal_models({"Loaded": asdict(facts)})) diff --git a/litellm/_auth/personal_owner.py b/litellm/_auth/personal_owner.py new file mode 100644 index 00000000000..7ec4c9a16c8 --- /dev/null +++ b/litellm/_auth/personal_owner.py @@ -0,0 +1,36 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.owner_teams import OwnerTeamFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerFacts: + user_present: bool + database_present: bool + cache_present: bool + logging_present: bool + owner_loaded: bool = False + owner_present: bool = False + owner_failed: bool = False + teams: OwnerTeamFacts | None = None + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerDecision: + kind: Literal["ReadOwner", "ReadTeams", "Complete"] + selected: tuple[int, ...] = () + withheld: tuple[int, ...] = () + unreadable: bool = False + + +_DECISION: Final = TypeAdapter(PersonalOwnerDecision) + + +@native_boundary +def personal_owner(facts: PersonalOwnerFacts) -> PersonalOwnerDecision: + return _DECISION.validate_python(_native.request_auth_personal_owner(asdict(facts))) diff --git a/litellm/_auth/reservation.py b/litellm/_auth/reservation.py new file mode 100644 index 00000000000..3fe0ac82678 --- /dev/null +++ b/litellm/_auth/reservation.py @@ -0,0 +1,131 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Gate: + token_present: bool + llm_route: bool | None = None + unbilled: bool | None = None + model_present: bool | None = None + counter_count: int | None = None + cost_read: bool = False + cost: float | None = None + + +@dataclass(frozen=True, slots=True) +class Sources: + key_present: bool + team_present: bool + team_id: str | None + user_present: bool + user_id: str | None + apply_user_budget_to_team_keys: bool + + +Scope: TypeAlias = Literal["Key", "Team", "User", "EndUser", "Tag", "AccessGroup", "Member", "Organization", "Project"] + + +@dataclass(frozen=True, slots=True) +class Counter: + scope: Scope + limit: float | None + + +@dataclass(frozen=True, slots=True) +class Window: + duration_present: bool + limit: float | None + start_present: bool | None = None + + +@dataclass(frozen=True, slots=True) +class Identity: + scope: Literal["EndUser", "Organization", "Project"] + primary: str | None + fallback: str | None = None + + +@dataclass(frozen=True, slots=True) +class MemberGate: + team_present: bool + team_id: str | None + user_present: bool + user_id: str | None + + +@dataclass(frozen=True, slots=True) +class MemberCap: + member_limit: float | None + default_read: bool = False + + +@dataclass(frozen=True, slots=True) +class IdentityPlan: + id: str | None + kind: Literal["Identity"] = "Identity" + + +@dataclass(frozen=True, slots=True) +class Action: + kind: Literal[ + "ReadRoute", + "ReadUnbilled", + "ReadModel", + "ReadCounters", + "ReadEstimate", + "Reserve", + "Skip", + "ReadStart", + "WarnWindow", + "AppendWindow", + "ReadMember", + "ReadDefault", + "ReadCap", + ] + + +@dataclass(frozen=True, slots=True) +class SourcePlan: + key: bool + team: bool + user: bool + kind: Literal["Sources"] = "Sources" + + +@dataclass(frozen=True, slots=True) +class CounterPlan: + include: bool + kind: Literal["Counter"] = "Counter" + + +_ACTION: Final = TypeAdapter(Action) +_SOURCES: Final = TypeAdapter(SourcePlan) +_COUNTER: Final = TypeAdapter(CounterPlan) +_IDENTITY: Final = TypeAdapter(IdentityPlan) + + +@native_boundary +def action(facts: Gate | Window | MemberGate | MemberCap) -> Action: + return _ACTION.validate_python(_native.request_auth_reservation(encode_facts(facts))) + + +@native_boundary +def sources(facts: Sources) -> SourcePlan: + return _SOURCES.validate_python(_native.request_auth_reservation(encode_facts(facts))) + + +@native_boundary +def counter(scope: Scope, limit: float | None) -> bool: + return _COUNTER.validate_python(_native.request_auth_reservation(encode_facts(Counter(scope, limit)))).include + + +@native_boundary +def identity(facts: Identity) -> str | None: + return _IDENTITY.validate_python(_native.request_auth_reservation(encode_facts(facts))).id diff --git a/litellm/_auth/resolved_models.py b/litellm/_auth/resolved_models.py new file mode 100644 index 00000000000..5e784b035d8 --- /dev/null +++ b/litellm/_auth/resolved_models.py @@ -0,0 +1,78 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm._auth.personal_model_request import PersonalModelRequestFacts +from litellm._auth.team_override import TeamOverrideFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ResolvedTeamRead: + found: bool = False + failed: bool = False + + +@dataclass(frozen=True, slots=True) +class ResolvedModelFacts: + model: str = "" + config_nonempty: bool = False + key_models: tuple[str, ...] = () + team_id: str | None = None + user_id: str | None = None + project_present: bool = False + customer_present: bool = False + database_present: bool = False + personal_enabled: bool = False + master: bool = False + role: str | None = None + key: GroupedModelScope | None = None + team_read: ResolvedTeamRead | None = None + embedded_ready: bool = False + team: TeamOverrideFacts | None = None + member: MemberModelFacts | None = None + owner_loaded: bool = False + owner_present: bool = False + owner_failed: bool = False + personal: PersonalModelRequestFacts | None = None + project_loaded: bool = False + project_grants: tuple[str, ...] = () + project: ModelAccessFacts | None = None + customer_loaded: bool = False + customer_grants: tuple[str, ...] = () + customer: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ResolvedModelDecision: + kind: Literal[ + "Allow", + "Deny", + "ReadKey", + "ReadKeyGroups", + "ReadTeamObject", + "ReadEmbeddedTeam", + "ReadTeamScope", + "ReadMember", + "ReadOwner", + "ReadPersonal", + "ReadProjectObject", + "ReadProjectScope", + "ReadCustomerObject", + "ReadCustomerScope", + ] + denial: ModelAccessDecision | None = None + member_models: tuple[str, ...] | None = None + + +_DECISION: Final = TypeAdapter(ResolvedModelDecision) + + +@native_boundary +def resolved_models(facts: ResolvedModelFacts) -> ResolvedModelDecision: + return _DECISION.validate_python(_native.request_auth_resolved_models(asdict(facts))) diff --git a/litellm/_auth/route_enable.py b/litellm/_auth/route_enable.py new file mode 100644 index 00000000000..828bc618cfc --- /dev/null +++ b/litellm/_auth/route_enable.py @@ -0,0 +1,52 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.routes import RouteContext +from litellm.rust_bridge import _native + +RouteFeature: TypeAlias = Literal["Admin", "Llm"] +LlmRouteDecision: TypeAlias = Literal["Allow", "Deny", "ReadEuJev"] + + +@dataclass(frozen=True, slots=True) +class RouteSetting: + environment_present: bool + premium: bool + secret: bool | None = None + + +@dataclass(frozen=True, slots=True) +class RouteEnableFacts: + context: RouteContext + imported: bool = False + failed: Literal["Http", "Other"] | None = None + admin: RouteSetting | None = None + llm: RouteSetting | None = None + exemptions: tuple[str, ...] | None = None + eu_jev_configured: bool | None = None + + +@dataclass(frozen=True, slots=True) +class RouteEnableDecision: + kind: Literal[ + "Allow", "PropagateHttp", "ReadImport", "ReadExemptions", "ReadEuJev", "ReadSettings", "ReadSecret", "Deny" + ] + feature: RouteFeature | None = None + reason: Literal["License", "Disabled"] | None = None + + +_DECISION: Final = TypeAdapter(RouteEnableDecision) +_LLM: Final = TypeAdapter[LlmRouteDecision](LlmRouteDecision) + + +@native_boundary +def route_enable(facts: RouteEnableFacts) -> RouteEnableDecision: + return _DECISION.validate_python(_native.request_auth_route_enable(asdict(facts))) + + +@native_boundary +def llm_route(context: RouteContext, configured: bool | None = None) -> LlmRouteDecision: + return _LLM.validate_python(_native.request_auth_llm_route(asdict(context), configured)) diff --git a/litellm/_auth/route_roles.py b/litellm/_auth/route_roles.py new file mode 100644 index 00000000000..4ad919ca8e0 --- /dev/null +++ b/litellm/_auth/route_roles.py @@ -0,0 +1,120 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class RoleFacts: + user_present: bool + role: str | None + + +@dataclass(frozen=True, slots=True) +class RoleDecision: + effective: str + user_role: str | None + admin: bool + + +@dataclass(frozen=True, slots=True) +class ApiFacts: + token_present: bool + admin_read: bool + user: RoleFacts + + +ApiDecision: TypeAlias = Literal["InvalidToken", "ReadAdmin", "NonAdmin", "Allow"] + + +@dataclass(frozen=True, slots=True) +class JwtFacts: + role: str | None + config_read: bool = False + config_missing: bool = False + + +JwtDecision: TypeAlias = Literal["ReadAdmin", "ReadTeam", "DefaultTeam", "Check", "Deny"] + + +@dataclass(frozen=True, slots=True) +class InsideFacts: + role: str | None + user_id: str | None + requested_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class Entry: + text: str | None + python_type: str + hash_failure: str | None = None + templates: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ListedFacts: + route: str + entry: Entry | None = None + exhausted: bool = False + template_index: int = 0 + compiled: str | None = None + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: Literal["TypeError", "AttributeError"] + message: str + + +@dataclass(frozen=True, slots=True) +class EntryDecision: + action: Literal["ReadEntry", "ReadTemplate", "Allow", "Deny", "Keep", "Expand", "Failure"] + template_index: int + failure: Failure | None + + +_ROLE: Final = TypeAdapter(RoleFacts) +_ROLE_RESULT: Final = TypeAdapter(RoleDecision) +_API: Final = TypeAdapter(ApiFacts) +_API_RESULT: Final = TypeAdapter[ApiDecision](ApiDecision) +_JWT: Final = TypeAdapter(JwtFacts) +_JWT_RESULT: Final = TypeAdapter[JwtDecision](JwtDecision) +_INSIDE: Final = TypeAdapter(InsideFacts) +_BOOL: Final = TypeAdapter(bool) +_ENTRY: Final = TypeAdapter(Entry) +_LISTED: Final = TypeAdapter(ListedFacts) +_ENTRY_RESULT: Final = TypeAdapter(EntryDecision) + + +@native_boundary +def role(facts: RoleFacts) -> RoleDecision: + return _ROLE_RESULT.validate_python(_native.request_auth_route_role(_ROLE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def api(facts: ApiFacts) -> ApiDecision: + return _API_RESULT.validate_python(_native.request_auth_route_role_api(_API.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def jwt(facts: JwtFacts) -> JwtDecision: + return _JWT_RESULT.validate_python(_native.request_auth_route_role_jwt(_JWT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def inside(facts: InsideFacts) -> bool: + return _BOOL.validate_python(_native.request_auth_route_role_inside(_INSIDE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def listed(facts: ListedFacts) -> EntryDecision: + return _ENTRY_RESULT.validate_python(_native.request_auth_route_role_listed(_LISTED.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def expand(facts: Entry) -> EntryDecision: + return _ENTRY_RESULT.validate_python(_native.request_auth_route_role_expand(_ENTRY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/_auth/routes.py b/litellm/_auth/routes.py new file mode 100644 index 00000000000..70e28eb404d --- /dev/null +++ b/litellm/_auth/routes.py @@ -0,0 +1,146 @@ +from collections.abc import Mapping +from dataclasses import asdict, dataclass, field +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class RouteCatalog: + groups: Mapping[str, tuple[str, ...]] = field(default_factory=dict[str, tuple[str, ...]]) + viewer_blocked: tuple[str, ...] = () + viewer_write_blocked: tuple[str, ...] = () + self_credentials: tuple[tuple[str, str], ...] = () + managed_models: tuple[str, ...] = () + managed_paths: tuple[str, ...] = () + managed_realtime: tuple[str, ...] = () + managed_live: tuple[str, ...] = () + managed_mcp: tuple[str, ...] = () + unbilled: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class RouteContext: + route: str + method: str | None = None + catalog: RouteCatalog = field(default_factory=RouteCatalog) + registered: bool = False + registered_auth: bool | None = None + configured_eu_jev: bool = False + + +@dataclass(frozen=True, slots=True) +class PassThroughFacts: + key_allowed: tuple[str, ...] = () + team_allowed: tuple[str, ...] = () + denied: tuple[str, ...] = () + forwarded: tuple[str, ...] = () + jwt_team_routes: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class VirtualRouteFacts: + context: RouteContext + allowed_routes: tuple[str, ...] | None + passthrough: PassThroughFacts = field(default_factory=PassThroughFacts) + password_reset_required: bool = False + + +@dataclass(frozen=True, slots=True) +class ViewerFacts: + context: RouteContext + request_keys: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class NonAdminFacts: + viewer: ViewerFacts + role: str | None = None + custom_admin_routes: tuple[str, ...] = () + premium: bool = False + query_user_id: str | None = None + token_user_id: str | None = None + permissions: tuple[str, ...] = () + organization_ids: tuple[str, ...] = () + admin_organization_ids: tuple[str, ...] = () + team_service_account: bool = False + allowed_routes: tuple[str, ...] | None = None + passthrough: PassThroughFacts = field(default_factory=PassThroughFacts) + + +RouteDenial: TypeAlias = Literal[ + "VirtualRoute", + "PasswordReset", + "PassThrough", + "DeniedPassThrough", + "CustomAdmin", + "UserInfo", + "AdminOnly", + "ViewerInference", + "ViewerWrite", + "ViewerParameter", +] + + +@dataclass(frozen=True, slots=True) +class RouteDecision: + allowed: bool + denial: RouteDenial | None = None + detail: str | None = None + + +_DECISION: Final = TypeAdapter(RouteDecision) + + +@native_boundary +def match_route( + route: str, + patterns: tuple[str, ...], + mode: Literal["Access", "Template", "Prefix", "Wildcard", "Denied", "Jwt"] = "Access", +) -> bool: + return _DECISION.validate_python( + _native.request_auth_route_access({"Match": {"route": route, "patterns": patterns, "mode": mode}}) + ).allowed + + +@native_boundary +def classify_route( + context: RouteContext, kind: Literal["Llm", "Unbilled", "Discovery", "SelfCredential", "Azure", "Managed"] +) -> bool: + return _DECISION.validate_python( + _native.request_auth_route_access({"Classify": {"context": asdict(context), "kind": kind}}) + ).allowed + + +@native_boundary +def virtual_route(facts: VirtualRouteFacts) -> RouteDecision: + return _DECISION.validate_python(_native.request_auth_route_access({"Virtual": asdict(facts)})) + + +@native_boundary +def non_admin_route(facts: NonAdminFacts) -> RouteDecision: + return _DECISION.validate_python(_native.request_auth_route_access({"NonAdmin": asdict(facts)})) + + +@native_boundary +def viewer_route(facts: ViewerFacts) -> RouteDecision: + return _DECISION.validate_python(_native.request_auth_route_access({"Viewer": asdict(facts)})) + + +@native_boundary +def passthrough_route(context: RouteContext, facts: PassThroughFacts, *, required: bool = False) -> RouteDecision: + return _DECISION.validate_python( + _native.request_auth_route_access( + {"PassThrough": {"context": asdict(context), "facts": asdict(facts), "required": required}} + ) + ) + + +@native_boundary +def denied_passthrough_route(route: str, facts: PassThroughFacts) -> RouteDecision: + return _DECISION.validate_python( + _native.request_auth_route_access({"DeniedPassThrough": {"route": route, "facts": asdict(facts)}}) + ) diff --git a/litellm/_auth/runtime.py b/litellm/_auth/runtime.py index 63544108511..58b83bc7296 100644 --- a/litellm/_auth/runtime.py +++ b/litellm/_auth/runtime.py @@ -1,88 +1,23 @@ -import os -from typing import Final, Literal, TypeAlias - -from litellm._auth import credential +from litellm._auth import rust_core from litellm._auth.credential import PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys from litellm._auth.decision import AuthCore -from litellm._auth.python_core import PythonCore - -Backend: TypeAlias = Literal["python", "rust", "shadow"] -def backend() -> Backend: - selected: Final = os.environ.get("LITELLM_AUTH_CORE", "rust") - match selected: - case "python" | "rust" | "shadow": - return selected - case _: - raise ValueError("LITELLM_AUTH_CORE must be python, rust or shadow") - - -def get_auth_core(selected: Backend | None = None) -> AuthCore: - match selected if selected is not None else backend(): - case "python": - return PythonCore() - case "rust": - from litellm._auth.rust_core import RustCore - - return RustCore() - case "shadow": - from litellm._auth.shadow import ShadowCore - - return ShadowCore() +def get_auth_core() -> AuthCore: + return rust_core.RustCore() def scheme_credential(value: str) -> str | None: - selected: Final = backend() - if selected == "python": - return credential.scheme_credential(value) - from litellm._auth import rust_core - - result: Final = rust_core.scheme_credential(value) - if selected == "shadow": - from litellm._auth.shadow import compare - - compare(result, credential.scheme_credential(value)) - return result + return rust_core.scheme_credential(value) def litellm_header_token(value: str) -> str: - selected: Final = backend() - if selected == "python": - return credential.litellm_header_token(value) - from litellm._auth import rust_core - - result: Final = rust_core.litellm_header_token(value) - if selected == "shadow": - from litellm._auth.shadow import compare - - compare(result, credential.litellm_header_token(value)) - return result + return rust_core.litellm_header_token(value) def select_api_key(keys: PresentedKeys) -> SelectedKey | None: - selected: Final = backend() - if selected == "python": - return credential.select_api_key(keys) - from litellm._auth import rust_core - - result: Final = rust_core.select_api_key(keys) - if selected == "shadow": - from litellm._auth.shadow import compare - - compare(result, credential.select_api_key(keys)) - return result + return rust_core.select_api_key(keys) def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: - selected: Final = backend() - if selected == "python": - return credential.select_websocket_api_key(keys) - from litellm._auth import rust_core - - result: Final = rust_core.select_websocket_api_key(keys) - if selected == "shadow": - from litellm._auth.shadow import compare - - compare(result, credential.select_websocket_api_key(keys)) - return result + return rust_core.select_websocket_api_key(keys) diff --git a/litellm/_auth/rust_core.py b/litellm/_auth/rust_core.py index 1f2a3cf5bd8..87b2ea3ad66 100644 --- a/litellm/_auth/rust_core.py +++ b/litellm/_auth/rust_core.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter from litellm._auth import decision from litellm._auth.credential import KeySource, PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys +from litellm._auth.failure import native_boundary from litellm.rust_bridge import _native Wire: TypeAlias = None | bool | int | float | str | list["Wire"] | dict[str, "Wire"] @@ -26,6 +27,7 @@ TEAM_BUDGET: Final = TypeAdapter(decision.TeamMemberBudget) FALLBACK_TARGET: Final = TypeAdapter(decision.FallbackTarget) TEAM_ALERT: Final = TypeAdapter(decision.AlertTeamMemberBudget) MEMBER_FALLBACK: Final = TypeAdapter(decision.VetTeamMemberFallback) +MODEL_DENIAL: Final = TypeAdapter(decision.ModelAccessDenied) _PUBLIC_TYPES: Final[Mapping[str, Callable[..., object]]] = { cls.__name__: cls @@ -49,6 +51,7 @@ _PUBLIC_TYPES: Final[Mapping[str, Callable[..., object]]] = { decision.KeyBudgetExceeded, decision.KeyWithoutToken, decision.ModelBudgetExceeded, + decision.ModelAccessDenied, decision.Allow, decision.Deny, decision.ResolveEndUser, @@ -58,6 +61,7 @@ _PUBLIC_TYPES: Final[Mapping[str, Callable[..., object]]] = { decision.PrepareKey, decision.InspectKey, decision.EnforceModelAccess, + decision.PrefetchAuthObjects, decision.LoadKeyOwner, decision.CheckBudgetExemption, decision.ReadTeamMemberBudget, @@ -82,21 +86,22 @@ _PUBLIC_TYPES: Final[Mapping[str, Callable[..., object]]] = { ) } _ANSWER_TYPES: Final = { + "EnforceModelAccess": "ModelAccess", "ResolvePresentedKey": "Presented", "InspectKey": "Key", "LoadKeyOwner": "Owner", - "CheckBudgetExemption": "Bool", + "CheckBudgetExemption": "Exemption", "ReadTeamMemberBudget": "TeamBudget", "ReadKeySpend": "Spend", "ReadProxySpend": "ProxySpend", "ReadModelBudgetScope": "ModelScope", - "CheckKeyModelBudget": "Exceeded", - "AuthorizeBudgetFallback": "Bool", - "CheckUserModelBudget": "Exceeded", + "CheckKeyModelBudget": "ModelBudgetData", + "AuthorizeBudgetFallback": "FallbackAccess", + "CheckUserModelBudget": "ModelBudgetData", "ReadTeamMemberModelBudgetScope": "TeamScope", - "CheckTeamMemberModelBudget": "Exceeded", + "CheckTeamMemberModelBudget": "ModelBudgetData", "ReadTeamMemberFallbacks": "Targets", - "VetTeamMemberFallback": "Bool", + "VetTeamMemberFallback": "Candidate", } @@ -104,6 +109,7 @@ class _Projection: def __init__(self) -> None: self.datetimes: tuple[datetime, ...] = () + @native_boundary def encode(self, value: object) -> Wire: if isinstance(value, datetime): offset: Final = value.utcoffset() @@ -145,6 +151,8 @@ class _Projection: raise TypeError("invalid auth datetime reference") return self.datetimes[reference] kind: Final = data.get("kind") + if kind == "ModelAccessDenied": + return MODEL_DENIAL.validate_python({"decision": data["decision"]}) if kind is None: return {key: self.decode(item) for key, item in data.items()} if not isinstance(kind, str): @@ -152,6 +160,7 @@ class _Projection: constructor: Final = public_type(kind) return constructor(**{key: self.decode(item) for key, item in data.items() if key != "kind"}) + @native_boundary def step(self, result: tuple[Mapping[str, object], _native.NativeAuthContinuation | None]) -> CoreStep: data, continuation = result kind: Final = data["kind"] @@ -160,6 +169,7 @@ class _Projection: if continuation is None: return CORE_STEP.validate_python(self.decode(data)) + @native_boundary def resume(*answers: object) -> CoreStep: answer: Final = { "kind": _ANSWER_TYPES.get(kind, "Unit"), @@ -184,11 +194,22 @@ def field_value(value: object, name: str) -> object: return cast(object, getattr(value, name)) # cast-ok: dataclass fields are discovered dynamically +@native_boundary +def encode_facts(value: object) -> Wire: + return _Projection().encode(value) + + +@native_boundary +def expired(expires: datetime, now: datetime) -> bool: + return _native.request_auth_expired(encode_facts(expires), encode_facts(now)) + + def public_type(kind: str) -> Callable[..., object]: return _PUBLIC_TYPES[kind] class RustCore: + @native_boundary def start(self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime) -> decision.Step: projection: Final = _Projection() return AUTH_STEP.validate_python( @@ -199,6 +220,7 @@ class RustCore: ), ) + @native_boundary def start_resolved_key( self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime ) -> decision.Step: @@ -211,16 +233,20 @@ class RustCore: ), ) + @native_boundary def check_model_budgets(self) -> decision.Step: return AUTH_STEP.validate_python(_Projection().step(_native.request_auth_budgets("Key"))) + @native_boundary def check_user_model_budgets(self) -> decision.Step: return AUTH_STEP.validate_python(_Projection().step(_native.request_auth_budgets("User"))) + @native_boundary def check_team_member_model_budgets(self) -> decision.TeamMemberStep: return MEMBER_STEP.validate_python(_Projection().step(_native.request_auth_budgets("TeamMember"))) +@native_boundary def select_api_key(keys: PresentedKeys) -> SelectedKey | None: selected: Final = _native.request_auth_select_api_key(_Projection().encode(keys)) if selected is None: @@ -228,14 +254,17 @@ def select_api_key(keys: PresentedKeys) -> SelectedKey | None: return SelectedKey(source=KeySource(selected["source"]), presented=selected["presented"], token=selected["token"]) +@native_boundary def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: selected: Final = _native.request_auth_select_websocket_api_key(_Projection().encode(keys)) return WebSocketDenial(selected["value"]) if selected["kind"] == "Denial" else selected["value"] +@native_boundary def scheme_credential(value: str) -> str | None: return _native.request_auth_scheme_credential(value) +@native_boundary def litellm_header_token(value: str) -> str: return _native.request_auth_litellm_header_token(value) diff --git a/litellm/_auth/spend_counter_policy.py b/litellm/_auth/spend_counter_policy.py new file mode 100644 index 00000000000..924cc58ca21 --- /dev/null +++ b/litellm/_auth/spend_counter_policy.py @@ -0,0 +1,129 @@ +from dataclasses import dataclass +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import Field, TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.rust_core import encode_facts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Spend: + current: float + verified: bool + fallback: float + limit: float | None + window: bool + fallback_authoritative: bool + strict_before: bool | None = None + authoritative_read: bool = False + authoritative: float | None = None + strict_after: bool | None = None + + +@dataclass(frozen=True, slots=True) +class ThrottleKey: + counter_key: str + token_present: bool + token: str | None + + +@dataclass(frozen=True, slots=True) +class Overflow: + current: float + limit: float + + +@dataclass(frozen=True, slots=True) +class OverBudget: + throttled: bool + current: float + limit: float + reservation: float + strict: bool + + +@dataclass(frozen=True, slots=True) +class ReservationFailure: + strict: bool + + +@dataclass(frozen=True, slots=True) +class ReservationGroups: + currents: tuple[float, ...] + limits: tuple[float, ...] + reservation: float + + +@dataclass(frozen=True, slots=True) +class SpendValue: + kind: Literal["Value", "Repair"] + spend: float + + +@dataclass(frozen=True, slots=True) +class Action: + kind: Literal[ + "ReadStrictBefore", + "ReadStrictAfter", + "ReadFloor", + "Unverifiable", + "ReadThrottle", + "Skip", + "Release", + "CheckOverflow", + "Unavailable", + "Compatible", + "Inconsistent", + "RejectResize", + ] + + +@dataclass(frozen=True, slots=True) +class Rejected: + kind: Literal["Reject"] + current: float + estimate: float | None + + +@dataclass(frozen=True, slots=True) +class Resize: + kind: Literal["Resize"] + cost: float + + +@dataclass(frozen=True, slots=True) +class Groups: + kind: Literal["ReservationGroups"] + groups: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True, slots=True) +class Adjustment: + current: float | None + adjustment: float + + +@dataclass(frozen=True, slots=True) +class ResizeConsistency: + inconsistent: bool + reseed: bool + + +Decision: TypeAlias = Annotated[SpendValue | Action | Rejected | Resize | Groups, Field(discriminator="kind")] +_DECISION: Final = TypeAdapter[Decision](Decision) + + +@native_boundary +def decide( + facts: Spend + | ThrottleKey + | Overflow + | OverBudget + | ReservationFailure + | ReservationGroups + | Adjustment + | ResizeConsistency, +) -> Decision: + return _DECISION.validate_python(_native.request_auth_spend_counter_policy(encode_facts(facts))) diff --git a/litellm/_auth/team_aliases.py b/litellm/_auth/team_aliases.py new file mode 100644 index 00000000000..c72ff652111 --- /dev/null +++ b/litellm/_auth/team_aliases.py @@ -0,0 +1,34 @@ +from collections.abc import Mapping +from dataclasses import asdict, dataclass, replace +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class TeamAliasFacts: + requested: tuple[str, ...] + aliases: Mapping[str, str] + team_id: str | None + router_present: bool + names: tuple[tuple[str, ...], ...] | None = None + + +@dataclass(frozen=True, slots=True) +class TeamAliasDecision: + kind: Literal["ReadNames", "Complete"] + count: int = 0 + models: tuple[str, ...] = () + + +_DECISION: Final = TypeAdapter(TeamAliasDecision) + + +@native_boundary +def team_aliases(facts: TeamAliasFacts) -> TeamAliasDecision: + return _DECISION.validate_python( + _native.request_auth_team_aliases(asdict(replace(facts, aliases=dict(facts.aliases)))) + ) diff --git a/litellm/_auth/team_override.py b/litellm/_auth/team_override.py new file mode 100644 index 00000000000..504f5cc83b7 --- /dev/null +++ b/litellm/_auth/team_override.py @@ -0,0 +1,30 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class TeamOverrideFacts: + team: GroupedModelScope | None = None + assigned_models: tuple[str, ...] | None = None + assigned_scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class TeamOverrideDecision: + kind: Literal["Allow", "Deny", "ReadTeam", "ReadTeamGroups", "ReadAssignedGroups", "ReadAssignedScope"] + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(TeamOverrideDecision) + + +@native_boundary +def team_override(facts: TeamOverrideFacts) -> TeamOverrideDecision: + return _DECISION.validate_python(_native.request_auth_team_override(asdict(facts))) diff --git a/litellm/_auth/unpriced_models.py b/litellm/_auth/unpriced_models.py new file mode 100644 index 00000000000..2944bc8294b --- /dev/null +++ b/litellm/_auth/unpriced_models.py @@ -0,0 +1,53 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PriceMetric: + name: str + scalar: float | None = None + nested: tuple[float | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class PriceEntry: + tiered_present: bool + metrics: tuple[PriceMetric, ...] + + +@dataclass(frozen=True, slots=True) +class DeploymentPricing: + parameter_keys: tuple[str, ...] + has_id: bool + info_loaded: bool = False + info: PriceEntry | None = None + cost_keys: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class UnpricedModelFacts: + model: str | None + router_present: bool + group_found: bool | None = None + deployments: tuple[DeploymentPricing, ...] | None = None + cost_deployments: tuple[DeploymentPricing, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class UnpricedModelDecision: + kind: Literal["Complete", "ReadGroup", "ReadDeployments", "ReadInfo", "ReadCostDeployments", "ReadCost"] + unpriced: bool = False + index: int = 0 + + +_DECISION: Final = TypeAdapter(UnpricedModelDecision) + + +@native_boundary +def unpriced_model(facts: UnpricedModelFacts) -> UnpricedModelDecision: + return _DECISION.validate_python(_native.request_auth_unpriced_model(asdict(facts))) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 557d1c7bc0a..34bb6781df4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,11 +1,7 @@ -import asyncio from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Final, TypeAlias -from fastapi import HTTPException - -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_AccessGroupTable if TYPE_CHECKING: @@ -31,13 +27,6 @@ class AgentAccessGroupCeiling: CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] -class _MissingGroup: - __slots__ = () - - -_MISSING: Final = _MissingGroup() - - async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through @@ -46,29 +35,9 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup: - from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + from litellm.proxy.auth.agent_ceiling_adapter import load_agent_access_group - if prisma_client is None: - verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) - return None - try: - return await get_access_object( - access_group_id=access_group_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except AccessGroupNotFoundError: - raise - except HTTPException as e: - if check_db_only: - raise - verbose_proxy_logger.warning( - "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail - ) - return None + return await load_agent_access_group(access_group_id, check_db_only=check_db_only) async def resolve_agent_access_group_ceiling( @@ -82,31 +51,10 @@ async def resolve_agent_access_group_ceiling( provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once any group is named, so an agent whose every group is missing is still capped to nothing.""" - named_ids: Final = await load_access_group_ids(agent_id) - if not named_ids: - return None + from litellm.proxy.auth.agent_ceiling_adapter import project_agent_ceiling, read_agent_ceiling_facts - from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError - - async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup: - try: - return await ( - _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id) - ) - except AccessGroupNotFoundError: - return _MISSING - - results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids)) - access_group_ids: Final = tuple( - group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup) - ) - groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable)) - return AgentAccessGroupCeiling( - access_group_ids=access_group_ids, - models=frozenset(model for group in groups for model in group.access_model_names), - mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), - agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), - access_groups=groups, + return project_agent_ceiling( + await read_agent_ceiling_facts(agent_id, load_access_group_ids, load_access_group, check_db_only=check_db_only) ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 9be9d4c7121..d7e036b2915 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -7,14 +7,15 @@ Follows the same pattern as MCP permission handling. import asyncio from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass -from itertools import chain +from dataclasses import dataclass, replace from types import MappingProxyType -from typing import Final, Protocol, TypeAlias +from typing import Final, Protocol, TypeAlias, assert_never from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError +from litellm._auth import agent_access, agent_grants, managed +from litellm._auth.failure import NativeAuthFailure from litellm._logging import verbose_logger from litellm.proxy._types import ( UI_TEAM_ID, @@ -30,6 +31,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.mcp_server.ui_session_utils import build_effective_auth_contexts +from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -58,6 +60,8 @@ class AgentAccessLookup(Protocol): def __call__(self, user_api_key_auth: UserAPIKeyAuth | None, /, *, strict: bool) -> Awaitable[AgentAccess]: ... +_AGENT_GROUPS: Final = TypeAdapter(tuple[str, ...]) + _DEFAULT_DENY_FLAG: Final = TypeAdapter(bool) @@ -70,13 +74,6 @@ def agent_access_default_deny() -> bool: return True -def _is_proxy_admin(user_api_key_auth: UserAPIKeyAuth | None) -> bool: - return user_api_key_auth is not None and user_api_key_auth.user_role in ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN.value, - ) - - def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry @@ -89,24 +86,20 @@ def _restricted_ids(access: AgentAccess) -> frozenset[str] | None: return _to_stable_ids(access.agent_ids) -def _union_agent_access(accesses: Sequence[AgentAccess]) -> AgentAccess: - if accesses and all(isinstance(access, UnrestrictedAgentAccess) for access in accesses): - return UnrestrictedAgentAccess() - return RestrictedAgentAccess( - frozenset( - chain.from_iterable(access.agent_ids for access in accesses if isinstance(access, RestrictedAgentAccess)) - ) - ) +def _project_agent_grants(decision: managed.ManagedDecision | agent_grants.Decision) -> AgentAccess: + return UnrestrictedAgentAccess() if decision.grants is None else RestrictedAgentAccess(frozenset(decision.grants)) def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess: key_ids: Final = _restricted_ids(key_access) team_ids: Final = _restricted_ids(team_access) - if key_ids is None: - return UnrestrictedAgentAccess() if team_ids is None else RestrictedAgentAccess(team_ids) - if team_ids is None: - return RestrictedAgentAccess(key_ids) - return RestrictedAgentAccess(key_ids & team_ids) + return _project_agent_grants( + managed.decide( + managed.IntersectGrants( + tuple(key_ids) if key_ids is not None else None, tuple(team_ids) if team_ids is not None else None + ) + ) + ) class AgentRequestHandler: @@ -135,27 +128,18 @@ class AgentRequestHandler: ) -> AgentAccess: """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" - if managed_agent_policy(user_api_key_auth) is not None: - return await _managed_actor_agent_access(user_api_key_auth) - key_team_access: Final = await AgentRequestHandler.resolve_key_team_agent_access( + return await _read_resolved_agent_grants( user_api_key_auth, - strict=strict, - lookup_key_access=lookup_key_access, - lookup_team_access=lookup_team_access, + resolve_ceiling, + default_deny, + lookup_key_access, + lookup_team_access, + agent_grants.Resolve( + managed=managed_agent_policy(user_api_key_auth) is not None, + strict=strict, + role=user_api_key_auth.user_role if user_api_key_auth is not None else None, + ), ) - deny_undefined: Final = strict or (not _is_proxy_admin(user_api_key_auth) and default_deny()) - if deny_undefined and isinstance(key_team_access, UnrestrictedAgentAccess): - return RestrictedAgentAccess(frozenset()) - caller_access: Final = await AgentRequestHandler.agent_caller_access(user_api_key_auth, strict=strict) - own_access: Final = _intersect_agent_access(key_team_access, caller_access) - agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling( - user_api_key_auth, resolve_ceiling, strict=strict - ) - if agent_ceiling is None: - return own_access - if isinstance(own_access, UnrestrictedAgentAccess): - return RestrictedAgentAccess(agent_ceiling) - return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) @staticmethod async def agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None, *, strict: bool = False) -> AgentAccess: @@ -181,11 +165,14 @@ class AgentRequestHandler: ) key_access: Final = await key_lookup(user_api_key_auth, strict=strict) team_access: Final = await team_lookup(user_api_key_auth, strict=strict) + except NativeAuthFailure: + raise except Exception as e: - if strict: + failure: Final = agent_grants.decide(agent_grants.Failure(strict, closed=False)) + if failure.action == "Unavailable": raise HTTPException(503, "Agent invocation policy is unavailable") from e verbose_logger.warning("Failed to get allowed agents: %s", e) - return UnrestrictedAgentAccess() + return _project_agent_grants(failure) return _intersect_agent_access(key_access, team_access) @staticmethod @@ -212,71 +199,8 @@ class AgentRequestHandler: user_api_key_auth: UserAPIKeyAuth | None = None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> bool: - """ - Check if a specific agent is allowed for the given user/key. - - Args: - agent_id: The agent ID to check - user_api_key_auth: User authentication info - - Returns: - bool: True if agent is allowed, False otherwise - """ - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore - from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure - from litellm.proxy.proxy_server import prisma_client - from litellm.types.proxy.agent_identity import AgentIdentityFailure - - registered: Final = global_agent_registry.get_agent_by_id(agent_id) - registry_managed: Final = isinstance(registered, AgentResponse) and registered.identity_managed - if registry_managed or prisma_client is not None: - target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id) - if isinstance(target, AgentIdentityFailure): - raise_identity_failure(target) - elif target is None and registry_managed: - return False - elif isinstance(target, AgentResponse) and target.identity_managed: - if ( - not target.enabled - or target.identity is None - or not target.identity.active - or user_api_key_auth is None - ): - return False - from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler - - key_hash: Final = user_api_key_auth.api_key or user_api_key_auth.token - authority: Final = ( - await MCPRequestHandler._reload_admitted_key(key_hash, check_db_only=True) # pyright: ignore[reportPrivateUsage] # the authoritative key reload has no public seam - if key_hash - and managed_agent_policy(user_api_key_auth) is None - and not user_api_key_auth.is_session_token - and not user_api_key_auth.authenticated_by_custom_auth - else user_api_key_auth - ) - fresh_auth: Final = authority.model_copy( - update=MappingProxyType( - {"requires_fresh_policy": True, "agent_caller": user_api_key_auth.agent_caller} - ) - ) - explicit: Final = await _granted_agent_ids( - fresh_auth, - _strict_agent_access, - build_effective_auth_contexts, - ) - return target.agent_id in explicit - - match await AgentRequestHandler.resolve_agent_access( - user_api_key_auth, - resolve_ceiling, - strict=bool(user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy), - ): - case UnrestrictedAgentAccess(): - return True - case RestrictedAgentAccess(allowed_agent_ids): - stable_id: Final = global_agent_registry.stable_agent_id(agent_id) - return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids) + facts: Final = await read_agent_access_facts(agent_id, user_api_key_auth, resolve_ceiling) + return agent_access.decide(facts) == "Allow" @staticmethod def _get_key_object_permission( @@ -361,29 +285,11 @@ class AgentRequestHandler: # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ()) - if not direct_agents and not declared_access_groups and not key_access_group_ids: - return UnrestrictedAgentAccess() - - access_group_agents: Final = ( - tuple( - await AgentRequestHandler._get_agents_from_access_groups( - declared_access_groups, check_db_only=strict - ) - ) - if declared_access_groups - else () + return await _read_declared_grants( + managed.DeclaredGrants(direct_agents, declared_access_groups, key_access_group_ids), strict ) - unified_agents: Final = ( - tuple( - await AgentRequestHandler._get_unified_access_group_agents( - key_access_group_ids, check_db_only=strict - ) - ) - if key_access_group_ids - else () - ) - - return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) + except NativeAuthFailure: + raise except Exception as e: if strict: raise HTTPException(503, "Agent invocation policy is unavailable") from e @@ -408,6 +314,8 @@ class AgentRequestHandler: proxy_logging_obj=proxy_logging_obj, check_db_only=strict, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # fail closed: an owner whose teams cannot be read grants nothing if strict: raise HTTPException(503, "Agent invocation policy is unavailable") from e @@ -419,22 +327,16 @@ class AgentRequestHandler: withheld_results: Final = await asyncio.gather( *(AgentRequestHandler._owner_team_agent_grants(team, strict=strict) for team in owner_teams.withheld) ) - if any(isinstance(access, AgentPolicyUnavailable) for access in team_results): + decision: Final = agent_grants.decide( + agent_grants.OwnerUnion( + tuple(_agent_grant_source(None, access) for access in team_results), + tuple(_agent_grant_source(None, access) for access in withheld_results), + owner_teams.unreadable, + ) + ) + if decision.action == "Unavailable": raise HTTPException(503, "Agent invocation policy is unavailable") - team_accesses: Final = tuple( - access for access in team_results if not isinstance(access, AgentPolicyUnavailable) - ) - withheld_restrictions: Final = tuple( - RestrictedAgentAccess(frozenset()) - for access in withheld_results - if isinstance(access, (AgentPolicyUnavailable, RestrictedAgentAccess)) - ) - accesses: Final = ( - (*team_accesses, *withheld_restrictions, RestrictedAgentAccess(frozenset())) - if owner_teams.unreadable - else (*team_accesses, *withheld_restrictions) - ) - return _union_agent_access(accesses) + return _project_agent_grants(decision) @staticmethod async def personal_key_agent_billing_auth( @@ -469,23 +371,27 @@ class AgentRequestHandler: ) stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id) - restricted_sources: Final = tuple( - (team, access) - for team, access in zip(teams, accesses, strict=True) - if isinstance(access, RestrictedAgentAccess) and stable_agent_id in _to_stable_ids(access.agent_ids) + selection: Final = agent_grants.decide( + agent_grants.BillingTeam( + tuple( + agent_grants.Source( + team.team_id, + tuple(_to_stable_ids(access.agent_ids)) + if isinstance(access, RestrictedAgentAccess) + else None, + isinstance(access, AgentPolicyUnavailable), + ) + for team, access in zip(teams, accesses, strict=True) + ), + stable_agent_id, + ) ) - unrestricted_sources: Final = ( - tuple(zip(teams, accesses, strict=True)) - if not restricted_sources - and teams - and all(isinstance(access, UnrestrictedAgentAccess) for access in accesses) - else () - ) - sources: Final = restricted_sources or unrestricted_sources - if not sources: + if selection.team_id is None: verbose_logger.warning("No owner team grants agent %s; billing against the personal key", agent_id) return auth - team: Final = min((source[0] for source in sources), key=lambda source: source.team_id) + team: Final = next(team for team in teams if team.team_id == selection.team_id) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # attribution must never fail an authorized call verbose_logger.warning("Failed to attribute agent %s spend to an owner team: %s", agent_id, e) return auth @@ -522,6 +428,8 @@ class AgentRequestHandler: else None ) team_auth_update: Final = await team_grant_billing_auth_update(auth, team, membership) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # a chosen team's controls must survive billing lookup failures verbose_logger.warning( "Failed to load billing controls for agent %s from owner team %s: %s", @@ -538,18 +446,27 @@ class AgentRequestHandler: *, strict: bool, ) -> AgentAccess | AgentPolicyUnavailable: - if team.object_permission_id is not None and team.object_permission is None: - if strict: - return AgentPolicyUnavailable(team_id=team.team_id) - verbose_logger.warning("Owner team %s has an unreadable object_permission, granting nothing", team.team_id) - return RestrictedAgentAccess(frozenset()) + decision: Final = agent_grants.decide( + agent_grants.OwnerTeam( + strict, + team.object_permission_id is not None, + team.object_permission is not None, + ) + ) + if decision.action == "Unavailable": + return AgentPolicyUnavailable(team_id=team.team_id) + if decision.action == "Complete": + return _project_agent_grants(decision) try: return await AgentRequestHandler._team_agent_grants(team, strict=strict) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # per-team isolation: a team whose grants cannot be read grants nothing - if strict: + failure: Final = agent_grants.decide(agent_grants.Failure(strict, closed=True)) + if failure.action == "Unavailable": return AgentPolicyUnavailable(team_id=team.team_id) verbose_logger.warning("Failed to get allowed agents for owner team %s: %s", team.team_id, e) - return RestrictedAgentAccess(frozenset()) + return _project_agent_grants(failure) @staticmethod async def _team_agent_grants(team: LiteLLM_TeamTable, *, strict: bool) -> AgentAccess: @@ -564,25 +481,9 @@ class AgentRequestHandler: # 2. Also include agents from team's access_group_ids (unified access groups) team_access_group_ids: Final = tuple(team.access_group_ids or ()) - if not direct_agents and not declared_access_groups and not team_access_group_ids: - return UnrestrictedAgentAccess() - - access_group_agents: Final = ( - tuple( - await AgentRequestHandler._get_agents_from_access_groups(declared_access_groups, check_db_only=strict) - ) - if declared_access_groups - else () + return await _read_declared_grants( + managed.DeclaredGrants(direct_agents, declared_access_groups, team_access_group_ids), strict ) - unified_agents: Final = ( - tuple( - await AgentRequestHandler._get_unified_access_group_agents(team_access_group_ids, check_db_only=strict) - ) - if team_access_group_ids - else () - ) - - return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) @staticmethod async def _get_allowed_agents_for_team( @@ -599,14 +500,27 @@ class AgentRequestHandler: Fetches the team object once and reuses it for both permission sources. Declared-but-empty grants stay restricted; see `get_allowed_agents_for_key`. """ - if user_api_key_auth is None: - return UnrestrictedAgentAccess() - - if user_api_key_auth.team_id is None: + facts: Final = agent_grants.TeamLookup( + subject=user_api_key_auth is not None, + team_id=user_api_key_auth.team_id if user_api_key_auth is not None else None, + database=True, + strict=strict, + ) + initial: Final = agent_grants.decide(facts) + if initial.action == "Complete": + return _project_agent_grants(initial) + assert user_api_key_auth is not None + if initial.action == "ReadPersonalEligibility": from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants - if not personal_key_inherits_team_grants(user_api_key_auth): - return UnrestrictedAgentAccess() + personal: Final = agent_grants.decide( + replace( + facts, + personal=personal_key_inherits_team_grants(user_api_key_auth), + ) + ) + if personal.action == "Complete": + return _project_agent_grants(personal) return await AgentRequestHandler._personal_key_owner_team_agent_access(user_api_key_auth, strict=strict) try: @@ -617,11 +531,13 @@ class AgentRequestHandler: user_api_key_cache, ) - if not prisma_client: - return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() + database: Final = agent_grants.decide(replace(facts, database=bool(prisma_client))) + if database.action == "Complete": + return _project_agent_grants(database) # Fetch the team object once for both permission sources - team_obj: Final = await get_team_object( + assert user_api_key_auth.team_id is not None + team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object( team_id=user_api_key_auth.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -630,18 +546,23 @@ class AgentRequestHandler: check_db_only=strict, ) - if team_obj is None: - return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() + loaded: Final = agent_grants.decide(replace(facts, loaded=True, row_present=team_obj is not None)) # pyright: ignore[reportUnnecessaryComparison] # legacy adapters may return no row + if loaded.action == "Complete": + return _project_agent_grants(loaded) + assert team_obj is not None return await AgentRequestHandler._team_agent_grants(team_obj, strict=strict) + except NativeAuthFailure: + raise except Exception as e: - if strict: + failure: Final = agent_grants.decide(agent_grants.Failure(strict, closed=False)) + if failure.action == "Unavailable": raise HTTPException(503, "Agent invocation policy is unavailable") from e # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: verbose_logger.warning("Failed to get allowed agents for team: %s", e) - return UnrestrictedAgentAccess() + return _project_agent_grants(failure) @staticmethod def _get_config_agent_ids_for_access_groups( @@ -650,17 +571,22 @@ class AgentRequestHandler: """ Helper to get agent_ids from config-loaded agents that match any of the given access groups. """ - server_ids: Final[set[str]] = set() - for agent in config_agents: - agent_access_groups = getattr(agent, "agent_access_groups", None) - if agent_access_groups: - if any(group in agent_access_groups for group in access_groups): - server_ids.add(agent.agent_id) - return server_ids + decision: Final = agent_grants.decide( + agent_grants.GroupMatches( + tuple( + agent_grants.ConfiguredAgent( + agent.agent_id, _AGENT_GROUPS.validate_python(getattr(agent, "agent_access_groups", None) or ()) + ) + for agent in config_agents + ), + tuple(access_groups), + ) + ) + return set(decision.grants or ()) @staticmethod async def _get_db_agent_ids_for_access_groups( - prisma_client, access_groups: Sequence[str], *, check_db_only: bool = False + prisma_client: PrismaClient | None, access_groups: Sequence[str], *, check_db_only: bool = False ) -> set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. @@ -685,7 +611,9 @@ class AgentRequestHandler: """ from litellm.proxy.auth.auth_checks import get_agent_ids_from_access_groups - return await get_agent_ids_from_access_groups(access_group_ids=access_group_ids, check_db_only=check_db_only) + return await get_agent_ids_from_access_groups( + access_group_ids=list(access_group_ids), check_db_only=check_db_only + ) @staticmethod async def _get_agents_from_access_groups( @@ -765,6 +693,8 @@ class AgentRequestHandler: return [] return key_object_permission.agent_access_groups or [] + except NativeAuthFailure: + raise except Exception as e: verbose_logger.warning("Failed to get agent access groups for key: %s", e) return [] @@ -799,7 +729,7 @@ class AgentRequestHandler: parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - if team_obj is None: + if team_obj is None: # pyright: ignore[reportUnnecessaryComparison] # legacy adapters may return no row verbose_logger.debug("team_obj is None") return [] @@ -808,6 +738,8 @@ class AgentRequestHandler: return [] return object_permissions.agent_access_groups or [] + except NativeAuthFailure: + raise except Exception as e: verbose_logger.warning("Failed to get agent access groups for team: %s", e) return [] @@ -836,7 +768,7 @@ async def _granted_agent_ids( accesses: Final = await asyncio.gather( *(resolve_access(auth_context) for auth_context in await effective_contexts(user_api_key_auth)) ) - return frozenset().union(*(_granted_ids(access) for access in accesses)) + return frozenset[str]().union(*(_granted_ids(access) for access in accesses)) async def accessible_agents( @@ -881,25 +813,76 @@ async def _strict_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: agent: Final = managed_agent_policy(auth) - if agent is None or not agent.object_permission: - return RestrictedAgentAccess(frozenset()) - permission: Final = LiteLLM_ObjectPermissionTable.model_validate(agent.object_permission or MappingProxyType({})) - own_auth: Final = UserAPIKeyAuth(object_permission=permission) - own: Final = _granted_ids(await AgentRequestHandler.get_allowed_agents_for_key(own_auth, strict=True)) - - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - ceilings: Final = await resolve_managed_agent_ceilings(agent) - grouped: Final = frozenset(target for target in own if all(target in ceiling.agent_ids for ceiling in ceilings)) - caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=True) - capped: Final = grouped if isinstance(caller, UnrestrictedAgentAccess) else grouped & caller.agent_ids context: Final = auth.managed_agent_context - if context is None or context.mode == "autonomous": - return RestrictedAgentAccess(capped) - if context.user_id is None: - return RestrictedAgentAccess(frozenset()) - human_ids: Final = await verified_human_agent_grants(context.user_id, auth.team_id) - return RestrictedAgentAccess(capped.intersection(human_ids)) + return await _read_managed_agent_grants( + auth, + managed.ManagedAgentGrants( + permission_present=bool(agent is not None and agent.object_permission), + mode=context.mode if context is not None else None, + user_present=context is not None and context.user_id is not None, + ), + ) + + +async def _read_managed_agent_grants(auth: UserAPIKeyAuth, facts: managed.ManagedAgentGrants) -> AgentAccess: + decision: Final = managed.decide(facts) + agent: Final = managed_agent_policy(auth) + match decision.action: + case "Grants": + return _project_agent_grants(decision) + case "ReadOwn": + permission: Final = LiteLLM_ObjectPermissionTable.model_validate( + agent.object_permission if agent is not None else MappingProxyType({}) + ) + own_auth: Final = UserAPIKeyAuth(object_permission=permission) + own: Final = await AgentRequestHandler.get_allowed_agents_for_key(own_auth, strict=True) + return await _read_managed_agent_grants(auth, replace(facts, own=tuple(_granted_ids(own)))) + case "ReadCeilings": + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + if agent is None: + raise RuntimeError("native requested absent managed agent") + ceilings: Final = await resolve_managed_agent_ceilings(agent) + return await _read_managed_agent_grants( + auth, replace(facts, ceilings=tuple(tuple(ceiling.agent_ids) for ceiling in ceilings)) + ) + case "ReadCaller": + caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=True) + return await _read_managed_agent_grants( + auth, + replace( + facts, + caller_loaded=True, + caller=tuple(caller.agent_ids) if isinstance(caller, RestrictedAgentAccess) else None, + ), + ) + case "ReadHumanGrants": + context: Final = auth.managed_agent_context + human: Final = await verified_human_agent_grants( + context.user_id if context is not None else None, auth.team_id + ) + return await _read_managed_agent_grants(auth, replace(facts, human=tuple(human))) + case _: + raise RuntimeError("invalid native managed grants transition") + + +async def _read_declared_grants(facts: managed.DeclaredGrants, strict: bool) -> AgentAccess: + decision: Final = managed.decide(facts) + match decision.action: + case "Grants": + return _project_agent_grants(decision) + case "ReadDeclaredGroups": + groups: Final = await AgentRequestHandler._get_agents_from_access_groups( # pyright: ignore[reportPrivateUsage] # module-owned staged reader + facts.declared_groups, check_db_only=strict + ) + return await _read_declared_grants(replace(facts, group_agents=tuple(groups)), strict) + case "ReadUnifiedGroups": + unified: Final = await AgentRequestHandler._get_unified_access_group_agents( # pyright: ignore[reportPrivateUsage] # module-owned staged reader + facts.unified_groups, check_db_only=strict + ) + return await _read_declared_grants(replace(facts, unified_agents=tuple(unified)), strict) + case _: + raise RuntimeError("invalid native declared grants transition") async def _verified_human_agent_sources( @@ -919,7 +902,13 @@ async def verified_human_agent_grants(user_id: str | None, team_id: str | None = sources: Final = await _verified_human_agent_sources( user_id, allowed_team_ids=frozenset((team_id,)) if team_id else frozenset() ) - return frozenset().union(*(grants for source, grants in sources if source is None or source == team_id)) + decision: Final = agent_grants.decide( + agent_grants.Human( + tuple(agent_grants.Source(source, tuple(grants)) for source, grants in sources), + team_id, + ) + ) + return frozenset(decision.grants or ()) async def resolve_delegated_agent_team( @@ -931,15 +920,202 @@ async def resolve_delegated_agent_team( allowed_team_ids: frozenset[str] | None = None, ) -> str | None: sources: Final = await _verified_human_agent_sources(user_id) - if any(source is None and agent_id in grants for source, grants in sources): - return team_id - granting_teams: Final = frozenset( - source - for source, grants in sources - if source is not None and agent_id in grants and (allowed_team_ids is None or source in allowed_team_ids) + decision: Final = agent_grants.decide( + agent_grants.Delegated( + tuple(agent_grants.Source(source, tuple(grants)) for source, grants in sources), + agent_id, + team_id, + explicit_team, + tuple(allowed_team_ids) if allowed_team_ids is not None else None, + ) ) - if team_id in granting_teams: - return team_id - if not explicit_team and granting_teams: - return min(granting_teams) - raise HTTPException(403, "Select a team that grants access to this agent using x-litellm-team-id") + if decision.action == "SelectTeam": + raise HTTPException(403, "Select a team that grants access to this agent using x-litellm-team-id") + return decision.team_id + + +def _agent_grant_source(team_id: str | None, access: AgentAccess | AgentPolicyUnavailable) -> agent_grants.Source: + return agent_grants.Source( + team_id, + tuple(access.agent_ids) if isinstance(access, RestrictedAgentAccess) else None, + isinstance(access, AgentPolicyUnavailable), + ) + + +async def _read_resolved_agent_grants( + auth: UserAPIKeyAuth | None, + ceiling: CeilingResolver, + default_deny: Callable[[], bool], + key_lookup: AgentAccessLookup | None, + team_lookup: AgentAccessLookup | None, + facts: agent_grants.Resolve, +) -> AgentAccess: + decision: Final = agent_grants.decide(facts) + match decision.action: + case "Complete": + return _project_agent_grants(decision) + case "ReadManaged": + assert auth is not None + return await _managed_actor_agent_access(auth) + case "ReadKeyTeam": + access: Final = await AgentRequestHandler.resolve_key_team_agent_access( + auth, + strict=facts.strict, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + key_team_loaded=True, + key_team=_agent_grant_source(None, access).grants, + ), + ) + case "ReadDefault": + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + default_deny=default_deny(), + ), + ) + case "ReadCaller": + caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=facts.strict) + own: Final = managed.decide( + managed.IntersectGrants( + facts.key_team, + _agent_grant_source(None, caller).grants, + ) + ) + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + caller_loaded=True, + own=own.grants, + ), + ) + case "ReadCeiling": + ids: Final = await AgentRequestHandler._agent_access_group_ceiling(auth, ceiling, strict=facts.strict) # pyright: ignore[reportPrivateUsage] # module-owned staged reader + return await _read_resolved_agent_grants( + auth, + ceiling, + default_deny, + key_lookup, + team_lookup, + replace( + facts, + ceiling_loaded=True, + ceiling=tuple(ids) if ids is not None else None, + ), + ) + case _: + raise NativeAuthFailure(f"unexpected agent grant directive {decision.action}") + + +async def read_agent_access_facts( + agent_id: str, + auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> agent_access.AgentAccessFacts: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.proxy_server import prisma_client + + registered: Final = global_agent_registry.get_agent_by_id(agent_id) + facts: Final = agent_access.AgentAccessFacts( + registry_managed=isinstance(registered, AgentResponse) and registered.identity_managed, + database_present=prisma_client is not None, + subject_present=auth is not None, + key_hash_present=bool(auth.api_key or auth.token) if auth is not None else False, + subject_managed=managed_agent_policy(auth) is not None, + session=auth.is_session_token if auth is not None else False, + custom=auth.authenticated_by_custom_auth if auth is not None else False, + ) + return await _read_agent_access(agent_id, auth, resolve_ceiling, facts) + + +async def _read_agent_access( + agent_id: str, + auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver, + facts: agent_access.AgentAccessFacts, + authority: UserAPIKeyAuth | None = None, +) -> agent_access.AgentAccessFacts: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.auth.managed_authorization import project_agent + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure + from litellm.proxy.proxy_server import prisma_client + from litellm.types.proxy.agent_identity import AgentIdentityFailure + + action: Final = agent_access.decide(facts) + match action: + case "Allow" | "Deny": + return facts + case "ReadTarget": + target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id) + if isinstance(target, AgentIdentityFailure): + raise_identity_failure(target) + return await _read_agent_access( + agent_id, + auth, + resolve_ceiling, + replace(facts, target_loaded=True, target=project_agent(target) if target is not None else None), + authority, + ) + case "ReadAuthority": + from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + assert auth is not None + key_hash: Final = auth.api_key or auth.token + assert key_hash is not None + reloaded: Final = await MCPRequestHandler._reload_admitted_key(key_hash, check_db_only=True) # pyright: ignore[reportPrivateUsage] # authoritative key reload has no public seam + return await _read_agent_access( + agent_id, auth, resolve_ceiling, replace(facts, authority_loaded=True), reloaded + ) + case "ReadExplicitGrants": + assert auth is not None + source: Final = authority if authority is not None else auth + fresh: Final = source.model_copy( + update=MappingProxyType({"requires_fresh_policy": True, "agent_caller": auth.agent_caller}) + ) + explicit: Final = await _granted_agent_ids(fresh, _strict_agent_access, build_effective_auth_contexts) + return await _read_agent_access( + agent_id, auth, resolve_ceiling, replace(facts, explicit_grants=tuple(explicit)), authority + ) + case "ReadLegacyGrants": + grants: Final = await AgentRequestHandler.resolve_agent_access( + auth, resolve_ceiling, strict=bool(auth is not None and auth.requires_fresh_policy) + ) + return await _read_agent_access( + agent_id, + auth, + resolve_ceiling, + replace( + facts, + legacy_loaded=True, + legacy_grants=tuple(grants.agent_ids) if isinstance(grants, RestrictedAgentAccess) else None, + ), + authority, + ) + case "ReadTargetIds": + stable: Final = global_agent_registry.stable_agent_id(agent_id) + ids: Final = global_agent_registry.ids_for_agent(stable) + return await _read_agent_access( + agent_id, auth, resolve_ceiling, replace(facts, target_ids=tuple(ids)), authority + ) + case _: + assert_never(action) diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index 57e99edecba..b0aeeeb0164 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -1,10 +1,12 @@ from collections.abc import Mapping +from dataclasses import replace from itertools import product from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import Final, Literal, assert_never -from pydantic import Field, TypeAdapter, ValidationError +from pydantic import TypeAdapter, ValidationError +from litellm._auth import agent_access, managed from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure @@ -94,17 +96,18 @@ _MODEL_ROUTE_KINDS: Final[ def managed_agent_route_allowed(route: str, method: str | None) -> bool: - from litellm.proxy.auth.route_checks import RouteChecks + from litellm.proxy.auth import route_policy_adapter - if route in ("/agents", "/v1/agents"): - return method in (None, "GET", "HEAD") - if route in _MANAGED_REALTIME_ROUTES or route in _MANAGED_LIVE_ROUTES: - return method in (None, "GET") - if route in _MANAGED_MODEL_ROUTES or RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS): - return method in (None, "POST") - return RouteChecks.check_route_access(route, _MANAGED_MCP_ROUTES) or RouteChecks.check_route_access( - route, LiteLLMRoutes.agent_inference_routes.value - ) + return route_policy_adapter.classify(route, "Managed", method) + + +def _model_route_kind( + route: str, +) -> Literal["completion", "image_generation", "image_edit", "moderation", "speech", "body", "path"]: + for suffix, kind in _MODEL_ROUTE_KINDS.items(): + if route.endswith(suffix): + return kind + return "completion" def managed_inference_request( @@ -134,7 +137,7 @@ def managed_inference_request( return dict(body) from litellm.proxy.common_utils.http_parsing_utils import resolve_inference_model - kind: Final = next((kind for suffix, kind in _MODEL_ROUTE_KINDS.items() if route.endswith(suffix)), "completion") + kind: Final = _model_route_kind(route) endpoint_model: Final = path_model or ( query_model if route.endswith(("/completions", "/embeddings", "/images/generations", "/images/edits")) else None ) @@ -159,71 +162,131 @@ def managed_agent_policy(auth: "UserAPIKeyAuth | None") -> AgentResponse | None: async def admit_managed_actor(auth: UserAPIKeyAuth, store: AgentIdentityStore | None) -> None: delegation_verified: Final = auth._managed_delegation_verified # pyright: ignore[reportPrivateUsage] # the one-shot delegation marker is a PrivateAttr by design auth._managed_delegation_verified = False # pyright: ignore[reportPrivateUsage] # consumed here so a replayed token cannot reuse it - if auth.agent_id is None: + facts: Final = managed.ActorLookup( + agent_id=auth.agent_id, + has_database=store is not None, + context=_managed_context(auth.managed_agent_context), + jwt_claims=bool(auth.jwt_claims), + delegation_verified=delegation_verified, + ) + await _drive_actor(auth, store, facts) + + +async def _drive_actor( + auth: UserAPIKeyAuth, + store: AgentIdentityStore | None, + facts: managed.ActorLookup, + agent: AgentResponse | None = None, +) -> None: + decision: Final = managed.decide(facts) + if decision.denial is not None: + raise_identity_failure(_managed_failure(decision.denial)) + if decision.action == "Allow": return - if store is None: + if decision.action == "ReadRegistry": from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + assert auth.agent_id is not None registered: Final = global_agent_registry.get_agent_by_id(auth.agent_id) - if auth.managed_agent_context is not None or ( - registered is not None and (registered.identity_managed or registered.identity is not None) - ): - raise_identity_failure( - AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") - ) + await _drive_actor( + auth, + store, + replace( + facts, + registry_loaded=True, + registry_managed=registered is not None + and (registered.identity_managed or registered.identity is not None), + ), + ) return - agent: Final = await store.agent(auth.agent_id) - if isinstance(agent, AgentIdentityFailure): - raise_identity_failure(agent) - if agent is None: + if decision.action == "ReadAgent": + assert store is not None and auth.agent_id is not None + loaded: Final = await store.agent(auth.agent_id) + if isinstance(loaded, AgentIdentityFailure): + raise_identity_failure(loaded) + await _drive_actor( + auth, + store, + replace(facts, agent_loaded=True, agent=project_agent(loaded) if loaded is not None else None), + loaded, + ) + return + if decision.action == "ReadRetired": + assert store is not None and auth.agent_id is not None retired: Final = await store.retired_agent(auth.agent_id) if isinstance(retired, AgentIdentityFailure): raise_identity_failure(retired) - if auth.managed_agent_context is not None or retired: - raise_identity_failure(AgentIdentityFailure(message="Agent no longer exists")) + await _drive_actor(auth, store, replace(facts, retired=retired), agent) return - if not agent.identity_managed: + if decision.action == "AssignPolicy": + assert agent is not None + auth.managed_agent_policy = agent + auth.billing_agent_policy = agent + auth.requires_fresh_policy = True + await _drive_actor(auth, store, replace(facts, policy_assigned=True), agent) return - if auth.jwt_claims and auth.managed_agent_context is None: - raise_identity_failure(AgentIdentityFailure(message="A managed agent requires a matching verified identity")) - failure: Final = actor_admission_failure(agent, auth.managed_agent_context) - if failure is not None: - raise_identity_failure(failure) - auth.managed_agent_policy = agent - auth.billing_agent_policy = agent - auth.requires_fresh_policy = True - if ( - auth.managed_agent_context is not None - and auth.managed_agent_context.mode == "delegated" - and not delegation_verified - ): + if decision.action == "ReadHumanGrants": from litellm.proxy.agent_endpoints.auth.agent_permission_handler import verified_human_agent_grants + assert auth.managed_agent_context is not None grants: Final = await verified_human_agent_grants(auth.managed_agent_context.user_id, auth.team_id) - if agent.agent_id not in grants: - raise_identity_failure( - AgentIdentityFailure(message="The delegated user is not permitted to invoke this agent") - ) + await _drive_actor(auth, store, replace(facts, human_grants=tuple(grants)), agent) + return + raise RuntimeError("Invalid managed actor decision") + + +def project_agent(agent: AgentResponse) -> managed.Agent: + return managed.Agent( + agent_id=agent.agent_id, + enabled=agent.enabled, + identity_managed=agent.identity_managed, + identity=managed.BoundIdentity(active=agent.identity.active, revision=agent.identity.revision) + if agent.identity is not None + else None, + execution_mode=agent.execution_mode, + ) + + +def _managed_context(context: ManagedAgentContext | None) -> managed.Context | None: + return ( + managed.Context( + agent_id=context.agent_id, + binding_revision=context.binding_revision, + mode=context.mode, + user_id=context.user_id, + ) + if context is not None + else None + ) + + +def _managed_failure(denial: managed.ManagedDenial) -> AgentIdentityFailure: + if denial == "DatabaseRequired": + return AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") + messages: Final[Mapping[managed.ManagedDenial, str]] = MappingProxyType( + { + "Disabled": "Agent execution is disabled", + "MissingContext": "This agent requires its bound identity provider token", + "ChangedIdentity": "Agent identity changed during authentication; retry", + "ExecutionMode": "Agent is not enabled for this execution mode", + "MissingHuman": "A verified human subject is required", + "Removed": "Agent no longer exists", + "VerifiedIdentityRequired": "A managed agent requires a matching verified identity", + "HumanNotPermitted": "The delegated user is not permitted to invoke this agent", + } + ) + return AgentIdentityFailure(message=messages[denial]) def actor_admission_failure( agent: AgentResponse, context: ManagedAgentContext | None, ) -> AgentIdentityFailure | None: - if not agent.enabled or agent.identity is None or not agent.identity.active: - return AgentIdentityFailure(message="Agent execution is disabled") - if context is None: - return AgentIdentityFailure(message="This agent requires its bound identity provider token") - if context.agent_id != agent.agent_id or context.binding_revision != agent.identity.revision: - return AgentIdentityFailure(message="Agent identity changed during authentication; retry") - if agent.execution_mode not in (context.mode, "both"): - return AgentIdentityFailure(message="Agent is not enabled for this execution mode") - if context.mode == "delegated" and not context.user_id: - return AgentIdentityFailure(message="A verified human subject is required") - return None + decision: Final = managed.decide(managed.Actor(project_agent(agent), _managed_context(context))) + return _managed_failure(decision.denial) if decision.denial is not None else None -_INVOCATION_COST: Final = TypeAdapter(Annotated[float, Field(ge=0, allow_inf_nan=False)]) +_INVOCATION_COST: Final = TypeAdapter(float) def invocation_target(route: str, body: Mapping[str, object]) -> str | None: @@ -238,36 +301,93 @@ def invocation_target(route: str, body: Mapping[str, object]) -> str | None: async def prepare_agent_invocation( auth: UserAPIKeyAuth, target_name: str, store: AgentIdentityStore | None, *, billable: bool = True ) -> None: - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentRequestHandler from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through registered: Final = await get_agent_with_read_through(target_name) - if registered is None: - return - registered_managed: Final = registered.identity_managed or registered.identity is not None - if store is None and registered_managed: - raise_identity_failure( - AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") - ) - target: Final = await store.agent(registered.agent_id) if store is not None else None - if isinstance(target, AgentIdentityFailure): - raise_identity_failure(target) - if target is None and registered_managed: - raise_identity_failure(AgentIdentityFailure(message="Invoked agent no longer exists")) - effective: Final = target if target is not None else registered - if not effective.identity_managed and auth.managed_agent_policy is None: - return - if not await AgentRequestHandler.is_agent_allowed(effective.agent_id, auth): - raise_identity_failure(AgentIdentityFailure(message="The caller is not permitted to invoke this agent")) - auth.invoked_agent_id = effective.agent_id - auth.invoked_agent_policy = effective - if auth.agent_id is None and effective.identity_managed: - auth.billing_agent_policy = effective - raw_fee: Final = (effective.litellm_params or MappingProxyType({})).get("cost_per_query", 0.0) if billable else 0.0 + facts: Final = agent_access.InvocationFacts( + registered=project_agent(registered) if registered is not None else None, + has_store=store is not None, + actor_managed=auth.managed_agent_policy is not None, + actor_id_present=auth.agent_id is not None, + ) + await _drive_invocation(auth, store, registered, None, billable, facts) + + +async def _drive_invocation( + auth: UserAPIKeyAuth, + store: AgentIdentityStore | None, + registered: AgentResponse | None, + target: AgentResponse | None, + billable: bool, + facts: agent_access.InvocationFacts, +) -> None: + plan: Final = agent_access.invocation(facts) + effective: Final = target if plan.use_target else registered + match plan.action: + case "Complete": + if facts.fee_loaded: + assert facts.fee is not None + auth.agent_invocation_cost = facts.fee + return + case "Deny": + failures: Final = MappingProxyType( + { + "DatabaseRequired": AgentIdentityFailure( + code="policy_unavailable", message="Managed agent policy requires a database" + ), + "Removed": AgentIdentityFailure(message="Invoked agent no longer exists"), + "CallerNotPermitted": AgentIdentityFailure( + message="The caller is not permitted to invoke this agent" + ), + "InvalidFee": AgentIdentityFailure( + code="policy_unavailable", message="Agent invocation price is invalid" + ), + } + ) + assert plan.denial is not None + raise_identity_failure(failures[plan.denial]) + case "ReadTarget": + assert registered is not None + loaded: Final = await store.agent(registered.agent_id) if store is not None else None + if isinstance(loaded, AgentIdentityFailure): + raise_identity_failure(loaded) + await _drive_invocation( + auth, + store, + registered, + loaded, + billable, + replace(facts, target_loaded=True, target=project_agent(loaded) if loaded is not None else None), + ) + case "ReadPermissions": + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import read_agent_access_facts + + assert effective is not None + permissions: Final = await read_agent_access_facts(effective.agent_id, auth) + await _drive_invocation(auth, store, registered, target, billable, replace(facts, permissions=permissions)) + case "Assign": + assert effective is not None + auth.invoked_agent_id = effective.agent_id + auth.invoked_agent_policy = effective + if plan.assign_billing: + auth.billing_agent_policy = effective + await _drive_invocation(auth, store, registered, target, billable, replace(facts, assigned=True)) + case "ReadFee": + assert effective is not None + raw: Final = ( + (effective.litellm_params or MappingProxyType[str, object]({})).get("cost_per_query", 0.0) + if billable + else 0.0 + ) + await _drive_invocation( + auth, store, registered, target, billable, replace(facts, fee_loaded=True, fee=_invocation_fee(raw)) + ) + case _: + assert_never(plan.action) + + +def _invocation_fee(raw: object) -> float | None: try: - fee: Final = _INVOCATION_COST.validate_python(raw_fee) + return _INVOCATION_COST.validate_python(raw) except ValidationError: - raise_identity_failure( - AgentIdentityFailure(code="policy_unavailable", message="Agent invocation price is invalid") - ) - auth.agent_invocation_cost = fee + return None diff --git a/litellm/proxy/auth/agent_ceiling_adapter.py b/litellm/proxy/auth/agent_ceiling_adapter.py new file mode 100644 index 00000000000..aee14fc9c6b --- /dev/null +++ b/litellm/proxy/auth/agent_ceiling_adapter.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Final + +from fastapi import HTTPException + +from litellm._auth.agent_ceilings import ( + AgentCeilingFacts, + AgentCeilingGroup, + AgentGroupLoadFacts, + agent_ceiling, + agent_group_load, +) +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable + +if TYPE_CHECKING: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AccessGroupIdsLoader, + AccessGroupLoader, + AgentAccessGroupCeiling, + ) + + +@dataclass(frozen=True, slots=True) +class AgentCeilingReading: + facts: AgentCeilingFacts + rows: tuple[LiteLLM_AccessGroupTable | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + agent_id: str + ids: AccessGroupIdsLoader + group: AccessGroupLoader + + +async def read_agent_ceiling_facts( + agent_id: str, + load_access_group_ids: AccessGroupIdsLoader, + load_access_group: AccessGroupLoader, + *, + check_db_only: bool = False, +) -> AgentCeilingReading: + return await _read( + AgentCeilingReading(AgentCeilingFacts(check_db_only=check_db_only)), + _Source(agent_id, load_access_group_ids, load_access_group), + ) + + +def project_agent_ceiling(reading: AgentCeilingReading) -> AgentAccessGroupCeiling | None: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + decision: Final = agent_ceiling(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete agent ceiling facts") + if not decision.present: + return None + return AgentAccessGroupCeiling( + decision.ids, + decision.models, + decision.mcp_server_ids, + decision.agent_ids, + tuple(_row(reading, index) for index in decision.indices), + ) + + +def _row(reading: AgentCeilingReading, index: int) -> LiteLLM_AccessGroupTable: + row: Final = reading.rows[index] + if row is None: + raise NativeAuthRuntimeError("agent ceiling selected a missing row") + return row + + +async def _group(group_id: str, loader: AccessGroupLoader) -> tuple[AgentCeilingGroup, LiteLLM_AccessGroupTable | None]: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + try: + row: Final = await loader(group_id) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError: + return AgentCeilingGroup("Missing"), None + if not isinstance(row, LiteLLM_AccessGroupTable): + return AgentCeilingGroup("Unavailable"), None + return AgentCeilingGroup( + "Found", tuple(row.access_model_names), tuple(row.access_mcp_server_ids), tuple(row.access_agent_ids) + ), row + + +async def _authoritative(group_id: str) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + _load_access_group, # pyright: ignore[reportPrivateUsage] # authoritative raw reader preserves the public injected-loader contract + ) + + return await _load_access_group(group_id, check_db_only=True) + + +async def _read(reading: AgentCeilingReading, source: _Source) -> AgentCeilingReading: + decision: Final = agent_ceiling(reading.facts) + match decision.kind: + case "Complete": + return reading + case "ReadIds": + ids: Final = await source.ids(source.agent_id) + return await _read(replace(reading, facts=replace(reading.facts, ids=ids)), source) + case "ReadGroups": + if reading.facts.ids is None: + raise NativeAuthRuntimeError("agent group reads without named IDs") + loader: Final = _authoritative if decision.authoritative else source.group + rows: Final = tuple(await asyncio.gather(*(_group(group_id, loader) for group_id in reading.facts.ids))) + return AgentCeilingReading( + replace(reading.facts, groups=tuple(row[0] for row in rows)), tuple(row[1] for row in rows) + ) + + +async def load_agent_access_group( + access_group_id: str, *, check_db_only: bool = False +) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + initial: Final = AgentGroupLoadFacts(prisma_client is not None, check_db_only) + initial_decision: Final = agent_group_load(initial) + if initial_decision.kind == "DatabaseUnavailable": + verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) + return None + if initial_decision.kind != "Read": + raise NativeAuthRuntimeError("invalid initial agent group read decision") + try: + return await get_access_object( + access_group_id=access_group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError: + if agent_group_load(replace(initial, status="NotFound")).kind != "Raise": + raise NativeAuthRuntimeError("invalid missing agent group decision") from None + raise + except HTTPException as error: + decision: Final = agent_group_load(replace(initial, status="HttpError")) + if decision.kind == "Raise": + raise + if decision.kind != "HttpUnavailable": + raise NativeAuthRuntimeError("invalid failed agent group decision") from error + verbose_proxy_logger.warning( + "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, error.detail + ) + return None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 71a47df48ca..4ba261368a3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -10,22 +10,24 @@ Run checks for: """ import asyncio -import math import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, Protocol, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, Protocol, TypeAlias, assert_never from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack import litellm +from litellm._auth import access as access_policy +from litellm._auth import common_flow +from litellm._auth.failure import NativeAuthFailure from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict @@ -34,8 +36,6 @@ from litellm.constants import ( CLI_SESSION_KEY_PREFIX, DEFAULT_ACCESS_GROUP_CACHE_TTL, DEFAULT_IN_MEMORY_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, INVALID_API_KEY_ERROR_MESSAGE, LITELLM_PROXY_MASTER_KEY_ALIAS, @@ -43,9 +43,7 @@ from litellm.constants import ( REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.dd_tracing import tracer -from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.tool_fields import DropPolicy @@ -77,7 +75,6 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, - ModelAccessDeniedProxyException, NewTeamRequest, ProxyErrorTypes, ProxyException, @@ -91,19 +88,9 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( resolve_agent_access_group_ceiling, ) from litellm.proxy.agent_endpoints.auth.agent_caller import ( - agent_caller_auth, load_agent_caller_team, load_agent_caller_user, ) -from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy -from litellm.proxy.auth.budget_throttle import ( - budget_throttle_percentage, - should_throttle_budget_exceeded, -) -from litellm.proxy.auth.model_access_denied import ( - customer_model_access_denied_client_message, - model_access_denied_client_message, -) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( broadcast_evictions, @@ -147,22 +134,10 @@ from litellm.proxy.guardrails.tool_name_extraction import ( extract_request_tool_names, ) from litellm.proxy.route_llm_request import route_request -from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start -from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, - end_user_counter, - key_counter, - key_window_counter, - model_access_group_counter, - org_counter, - project_counter, - tag_counter, - team_counter, team_member_counter, - team_window_counter, - user_counter, ) from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository @@ -401,22 +376,6 @@ def _model_access_group_budget_table( return repo.table -class _MemberModelScope(Protocol): - @property - def allowed_models(self) -> Sequence[str] | None: ... - - -class _TeamMembershipModelScope(Protocol): - @property - def litellm_budget_table(self) -> _MemberModelScope | None: ... - - -def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[str]: - """The member's own model scope, read through a narrowed view of the membership row.""" - budget_table: Final = membership.litellm_budget_table - return () if budget_table is None else (budget_table.allowed_models or ()) - - class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... @@ -435,18 +394,6 @@ typed_general_settings: Final = _typed_request_body _TRACE_ID_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -def _metadata_mapping_has_trace_id(metadata: Mapping[str, object]) -> bool: - trace_id: Final[object] = metadata.get("trace_id") - return isinstance(trace_id, str) and bool(trace_id) - - -def _metadata_has_trace_id(metadata: object) -> bool: - if not isinstance(metadata, dict): - return False - metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(metadata) - return _metadata_mapping_has_trace_id(metadata_dict) - - def _team_metadata_requires_trace_id(team_object: LiteLLM_TeamTable) -> bool: if not isinstance( team_object.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # narrow the untyped field @@ -460,31 +407,45 @@ def _team_metadata_requires_trace_id(team_object: LiteLLM_TeamTable) -> bool: def _request_has_trace_id(request_body: Mapping[str, object], request: Request, *, headers_only: bool = False) -> bool: + return _read_trace_id(request_body, request, "Header", headers_only=headers_only) + + +def _read_trace_id( + request_body: Mapping[str, object], + request: Request, + source: access_policy.TraceSource, + *, + headers_only: bool, +) -> bool: from litellm.proxy.litellm_pre_call_utils import ( _trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser get_chain_id_from_headers, metadata_variable_name_for_request, ) - if get_chain_id_from_headers(dict(request.headers)): - return True - if headers_only: - return False - litellm_trace_id: Final[object] = request_body.get("litellm_trace_id") - if isinstance(litellm_trace_id, str) and litellm_trace_id: - return True - if "litellm_trace_id" not in request_body: - traceparent: Final[str | None] = request.headers.get("traceparent") - if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent): - return True - metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body) - selected_metadata: Final[object] = request_body.get(metadata_variable_name) - if isinstance(selected_metadata, dict): - selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata) - if "trace_id" in selected_metadata_dict: - return _metadata_mapping_has_trace_id(selected_metadata_dict) - other_metadata_variable_name: Final = "litellm_metadata" if metadata_variable_name == "metadata" else "metadata" - return _metadata_has_trace_id(request_body.get(other_metadata_variable_name)) + if source == "Header": + value, present = get_chain_id_from_headers(dict(request.headers)), True + elif source == "Body": + value, present = request_body.get("litellm_trace_id"), "litellm_trace_id" in request_body + elif source == "Traceparent": + traceparent: Final = request.headers.get("traceparent") + value, present = _trace_id_from_traceparent(traceparent) if isinstance(traceparent, str) else None, True + else: + selected_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body) + name: Final = ( + selected_name + if source == "Selected" + else ("litellm_metadata" if selected_name == "metadata" else "metadata") + ) + metadata: Final = request_body.get(name) + mapping: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(metadata) if isinstance(metadata, dict) else {} + value, present = mapping.get("trace_id"), "trace_id" in mapping + result: Final = access_policy.decide( + access_policy.TraceStep(source, value if isinstance(value, str) else None, present, headers_only) + ) + if result.next_trace is not None: + return _read_trace_id(request_body, request, result.next_trace, headers_only=headers_only) + return result.denial is None class _JsonLoadsObj(Protocol): @@ -541,138 +502,10 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None: ) -def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None: - """ - Return the router's per-instance zero-cost cache, or ``None`` for objects - that don't expose one (e.g. ``MagicMock`` stand-ins in unit tests). - - The cache lives on the ``Router`` instance so it: - * is invalidated by ``Router._invalidate_model_group_info_cache`` on - any model add/remove/upsert (including in-place pricing changes via - ``/model/update``, which go through ``upsert_deployment``); - * dies with the router itself — no risk of CPython reusing the - previous router's ``id()`` and serving its cached entries. - """ - cache: Final = llm_router.registry().zero_cost() - return cache if isinstance(cache, dict) else None - - def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool: - """ - Check if a model has zero cost (no configured pricing). + from litellm.proxy.auth import budget_pricing_adapter - Uses the router registry's group info to get pricing information. - - Args: - model: The model name or list of model names - llm_router: The LiteLLM router instance - - Returns: - bool: True if all costs for the model are zero, False otherwise - """ - if model is None or llm_router is None: - return False - - # Handle list of models - model_list: Final = [model] if isinstance(model, str) else model - - zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) - - for model_name in model_list: - if zero_cost_cache is not None: - cached = zero_cost_cache.get(model_name) - if cached is not None: - if cached is False: - return False - continue - try: - model_group_info = llm_router.registry().group(model_name, include_hidden=True) - - if model_group_info is None: - # Model not found or no pricing info available - # Conservative approach: assume it has cost - verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - # Check costs for this model - # Only allow bypass if BOTH costs are explicitly set to 0 (not None) - input_cost = model_group_info.input_cost_per_token - output_cost = model_group_info.output_cost_per_token - - # If costs are not explicitly configured (None), assume it has cost - if input_cost is None or output_cost is None: - verbose_proxy_logger.debug( - "Model %s has undefined cost (input: %s, output: %s), assuming it has cost", - model_name, - input_cost, - output_cost, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - # If either cost is non-zero, return False - if input_cost > 0 or output_cost > 0: - verbose_proxy_logger.debug( - "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - # Costs are 0 — verify this is from explicit configuration, - # not from defaulted sparse auto-registration entries. - # See: https://github.com/BerriAI/litellm/issues/24770 - safe_name = str(model_name).replace("\n", "").replace("\r", "") - served_deployments = _deployments_served_for(model_name, llm_router) - if not _is_cost_explicitly_configured(served_deployments): - verbose_proxy_logger.debug( - "Model %s has zero cost but no explicit cost " - "configuration in model_cost entry — treating as unknown " - "cost (enforce budget)", - safe_name, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - if _has_explicit_positive_cost(served_deployments): - verbose_proxy_logger.debug( - "Model %s routes to a deployment with an explicit positive per-token price (enforce budget)", - safe_name, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - if _has_ptu_flat_cost(served_deployments): - verbose_proxy_logger.debug( - "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " - "rate is not a free model (enforce budget)", - safe_name, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = False - return False - - verbose_proxy_logger.debug( - "Model %s has zero cost explicitly configured (input: %s, output: %s)", - safe_name, - input_cost, - output_cost, - ) - if zero_cost_cache is not None: - zero_cost_cache[model_name] = True - - except Exception as e: - # If we can't determine the cost, assume it has cost (conservative approach) - verbose_proxy_logger.debug("Error checking cost for model %s: %s, assuming it has cost", model_name, e) - return False - - # All models checked have zero cost - return True + return budget_pricing_adapter.is_model_cost_zero(model, llm_router) _is_model_cost_zero: Final = is_model_cost_zero @@ -682,7 +515,7 @@ _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) -def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: +def deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: """ The deployments a request to ``model`` routes to, for the gates that waive budget checks. @@ -706,128 +539,16 @@ def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[ ) -def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: - """Whether any of the deployments bills reserved PTU capacity as a flat cost. - - Such a deployment carries an explicit zero per-token price so the flat cost is not charged - twice, which otherwise reads here as a free model and waives every budget check for it. - """ - for deployment in deployments: - model_info = deployment.get("model_info") or _NO_MODEL_INFO - if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: - return True - return False - - -def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool: - """ - Check if any of the deployments has cost fields explicitly set in its - litellm.model_cost entry. - - When Router._create_deployment() registers a model not in the global - cost map, it creates a sparse entry like {"id": ""} with no cost - fields. _get_model_info_helper() then defaults missing costs to 0. - This function detects that scenario by checking the raw model_cost entry. - - It also reaches a deployment that prices itself through its ``model_info`` block, whose entry - lands in the cost map under the deployment id. - """ - for deployment in deployments: - model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") - if model_id is None: - continue - raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) - if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: - return True - return False - - -def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: - """Whether any of the deployments carries an explicit positive per-token or per-query price. - - The group's price is read from the deployments named after the routed group and from - that group's own alias target, while an alias chain is served from the deployments named - after its routed group or from a wildcard route; a priced route among those serves the - request at its own rate whatever the group's price reads. - """ - for deployment in deployments: - model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") - if model_id is None: - continue - raw_entry = as_str_mapping(litellm.model_cost.get(model_id)) or _EMPTY_COST_ENTRY - if ( - _is_positive_cost(raw_entry.get("input_cost_per_token")) - or _is_positive_cost(raw_entry.get("output_cost_per_token")) - or _is_positive_cost(_deployment_query_price(deployment["litellm_params"], raw_entry)) - ): - return True - return False - - -def _deployment_query_price(litellm_params: Mapping[str, object], cost_entry: Mapping[str, object]) -> float | None: - provider: Final = litellm_params.get("custom_llm_provider") or str(litellm_params.get("model")).partition("/")[0] - return decisions_request_price({**cost_entry, "litellm_provider": provider}, decisions_pricing_fields(cost_entry)) +_deployments_served_for: Final = deployments_served_for _EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) -def _is_positive_cost(value: object) -> bool: - return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 - - -def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: - if entry.get("tiered_pricing") is not None: - return True - for key, value in entry.items(): - if "cost_per" not in key: - continue - if _is_positive_cost(value): - return True - if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()): - return True - return False - - -def _entry_declares_price(entry: Mapping[str, object]) -> bool: - return any("cost_per" in key or key == "tiered_pricing" for key in entry) - - -def _model_group_has_pricing(model: str, llm_router: "Router") -> bool: - """ - A model group counts as priced when a deployment overrides any *cost_per* field or - tiered_pricing in its litellm_params, even at zero, or when its resolved model info carries - tiered_pricing or a positive price on any billed metric (tokens, characters, seconds, pages, - images, queries, ...), so models billed by a non-token metric are not treated as unpriced. - """ - registry: Final = llm_router.registry() - for deployment in registry.deployments(model) or (): - litellm_params = deployment.get("litellm_params") or _EMPTY_COST_ENTRY - if _entry_declares_price(litellm_params): - return True - - model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") - if model_id is None: - continue - - model_info = registry.deployment_model_info(model_id, litellm_params.get("model") or "") - if model_info is not None and _entry_has_priced_metric(model_info): - return True - - return False - - def model_has_no_cost_mapping(model: str | None, llm_router: Router | None) -> bool: - if not model or llm_router is None: - return False + from litellm.proxy.auth.unpriced_model_adapter import model_has_no_cost_mapping as native_model_has_no_cost_mapping - if llm_router.registry().group(model) is None: - return False - - if _model_group_has_pricing(model=model, llm_router=llm_router): - return False - - return not _is_cost_explicitly_configured(llm_router.registry().deployments(model) or ()) + return native_model_has_no_cost_mapping(model, llm_router) def _unpriced_models_in_request(model: str | list[str] | None, llm_router: Router | None) -> tuple[str, ...]: @@ -863,7 +584,7 @@ async def _run_project_checks( return # 1.1. If project is blocked - if project_object.blocked is True: + if access_policy.decide(access_policy.Blocked(project_object.blocked is True)).denial is not None: raise Exception( f"Project={project_object.project_id} is blocked. Update via `/project/update` if you're an admin." ) @@ -894,26 +615,43 @@ async def _run_project_checks( def _enforce_user_param_check(general_settings: dict, request: Request, request_body: dict, route: str) -> None: - if not general_settings.get("enforce_user_param", False): + if ( + access_policy.decide( + access_policy.UserParamRequest(bool(general_settings.get("enforce_user_param", False))) + ).read + is None + ): return - http_method: Final = request.method if hasattr(request, "method") else None - is_post_method: Final = http_method and http_method.upper() == "POST" is_openai_route: Final = RouteChecks.is_llm_api_route(route=route) is_mcp_route: Final = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value ) - if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body: + decision: Final = access_policy.decide( + access_policy.UserParam( + enabled=bool(general_settings.get("enforce_user_param", False)), + method=http_method, + llm_route=is_openai_route, + mcp_route=is_mcp_route, + has_user="user" in request_body, + ) + ) + if decision.denial is not None: raise Exception(f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}") def _reject_clientside_metadata_tags_check( general_settings: dict, request_body: dict, route: str, request: Request | None = None ) -> None: - if not general_settings.get("reject_clientside_metadata_tags", False): + initial: Final = access_policy.ClientTagRequest( + bool(general_settings.get("reject_clientside_metadata_tags", False)) + ) + if access_policy.decide(initial).read is None: + return + routed: Final = replace(initial, llm_route=RouteChecks.is_llm_api_route(route=route)) + if access_policy.decide(routed).read is None: return - original_client_tags: Final = ( getattr(request.state, "litellm_baggage_client_metadata_tags", None) if request is not None else None ) @@ -926,7 +664,14 @@ def _reject_clientside_metadata_tags_check( and "tags" in request_body["metadata"] ) ) - if RouteChecks.is_llm_api_route(route=route) and has_client_tags: + decision: Final = access_policy.decide( + access_policy.ClientTags( + enabled=initial.enabled, + llm_route=bool(routed.llm_route), + has_tags=has_client_tags, + ) + ) + if decision.denial is not None: raise ProxyException( message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via key or team metadata, or by configured request_tag_mappings.", type=ProxyErrorTypes.bad_request_error, @@ -936,19 +681,9 @@ def _reject_clientside_metadata_tags_check( def _global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_checks: bool, route: str) -> None: - if ( - litellm.max_budget > 0 - and not skip_budget_checks - and global_proxy_spend is not None - and RouteChecks.is_llm_api_route(route=route) - and route not in MODEL_LIST_ROUTES - ): - if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget: - raise litellm.BudgetExceededError( - current_cost=global_proxy_spend, - max_budget=litellm.max_budget, - entity_type=Litellm_EntityType.PROXY.value, - ) + from litellm.proxy.auth import common_budget_adapter + + return common_budget_adapter.global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route) _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( @@ -964,17 +699,22 @@ def requests_guardrail_modification(request_body: Mapping[str, object]) -> bool: default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``), at the top level or in either metadata key""" return ( - _requests_guardrail_modification_in(request_body.get("metadata")) - or _requests_guardrail_modification_in(request_body.get("litellm_metadata")) - or _requests_guardrail_modification_in(request_body) + access_policy.decide(access_policy.GuardrailRequest(_requested_guardrail_keys(request_body))).read is not None ) -def _requests_guardrail_modification_in(container: object) -> bool: - coerced: Final = _guardrail_container_as_dict(container) - if coerced is None: - return False - return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) +def _requested_guardrail_keys(request_body: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + chain.from_iterable( + tuple(container) + for container in ( + _guardrail_container_as_dict(request_body.get("metadata")), + _guardrail_container_as_dict(request_body.get("litellm_metadata")), + request_body, + ) + if container is not None + ) + ) def _guardrail_container_as_dict(container: object) -> Mapping[str, object] | None: @@ -1008,10 +748,19 @@ def guardrail_modification_check(request_body: Mapping[str, object], team_object require_guardrail_modification_permission(team_object) -def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails +def _team_modify_guardrails(team_object: LiteLLM_TeamTable | None) -> bool | None: + metadata: Final = _typed_request_body(team_object.metadata or {}) if team_object is not None else {} + guardrails: Final = metadata.get("guardrails") + values: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(guardrails) if isinstance(guardrails, dict) else {} + value: Final = values.get("modify_guardrails") + return value if isinstance(value, bool) else None - if not can_modify_guardrails(team_object): + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: + if ( + access_policy.decide(access_policy.Guardrails(("guardrails",), _team_modify_guardrails(team_object))).denial + is not None + ): raise HTTPException( status_code=403, detail={"error": "Your team does not have permission to modify guardrails."}, @@ -1048,11 +797,16 @@ async def check_tools_allowlist( tool_names: Final = extract_request_tool_names(route, request_body) if not tool_names: return - allowed_set: Final = {str(t) for t in effective} - disallowed: Final = [n for n in tool_names if n not in allowed_set] - if disallowed: + decision: Final = access_policy.decide( + access_policy.Tools( + requested=tuple(tool_names), + key_allowed=tuple(str(t) for t in effective), + team_allowed=None, + ) + ) + if decision.denial is not None: raise ProxyException( - message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.", + message=f"Tool(s) {list(decision.rejected)} are not in the allowed tools list for this key/team.", type=ProxyErrorTypes.tool_access_denied, param="tools", code=status.HTTP_403_FORBIDDEN, @@ -1107,17 +861,9 @@ def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | def route_skips_budget_checks(route: str) -> bool: - if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: - return False - if route in MODEL_DISCOVERY_ROUTES: - return True - if RouteChecks.is_llm_api_route(route=route): - return False - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack - InitPassThroughEndpointHelpers, - ) + from litellm.proxy.auth import budget_pricing_adapter - return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + return budget_pricing_adapter.route_skips_budget_checks(route) def request_skips_budget_checks( @@ -1127,11 +873,9 @@ def request_skips_budget_checks( *, request_body: Mapping[str, object] | None = None, ) -> bool: - if route_skips_budget_checks(route=route): - return True - if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body): - return True - return is_model_cost_zero(model=model, llm_router=llm_router) + from litellm.proxy.auth import budget_pricing_adapter + + return budget_pricing_adapter.request_skips_budget_checks(route, model, llm_router, request_body=request_body) async def common_checks( @@ -1181,35 +925,60 @@ async def common_checks( team_id=valid_token.team_id if valid_token is not None else None, ) - skip_all_budget_checks: Final = skip_budget_checks or request_skips_budget_checks( - route=route, - model=_model, - llm_router=llm_router, - request_body=_typed_request_body(request_body), + budget_skip: Final = common_flow.BudgetSkip(skip_budget_checks) + budget_gate: Final = ( + common_flow.read( + replace( + budget_skip, + exempt=request_skips_budget_checks( + route=route, model=_model, llm_router=llm_router, request_body=_typed_request_body(request_body) + ), + ) + ) + if common_flow.read(budget_skip) == "ReadExemption" + else common_flow.read(budget_skip) ) + skip_all_budget_checks: Final = budget_gate == "SkipBudgets" - membership_user_id: Final = ( - valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None + flow_facts: Final = common_flow.Facts( + model_present=bool(_model), + skip_budgets=skip_all_budget_checks, + team_present=team_object is not None, + user_present=user_object is not None, + token_present=valid_token is not None, + user_id=valid_token.user_id if valid_token is not None else None, + agent_id=None, + customer_present=False, + customer_models=(), ) - team_membership_loaded: Final = team_object is not None and membership_user_id is not None - loaded_team_membership: Final = ( - await get_team_membership( + flow: Final = common_flow.plan(flow_facts) + membership_user_id: Final = valid_token.user_id if valid_token is not None else None + team_membership_loaded: Final = flow.membership + + async def read_membership() -> LiteLLM_TeamMembership | None: + if not flow.membership: + return None + assert membership_user_id is not None and team_object is not None + return await get_team_membership( user_id=membership_user_id, team_id=team_object.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - if team_object is not None and membership_user_id is not None - else None - ) - unpriced_models: Final = ( - _unpriced_models_in_request(model=_model, llm_router=llm_router) - if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route) - else () + loaded_team_membership: Final = await read_membership() + + pricing: Final = common_flow.Pricing(litellm.block_requests_for_models_without_pricing) + pricing_read: Final = ( + common_flow.read(replace(pricing, llm_route=RouteChecks.is_llm_api_route(route=route))) + if common_flow.read(pricing) == "ReadRoute" + else common_flow.read(pricing) ) - if unpriced_models: + unpriced_models: Final = ( + _unpriced_models_in_request(model=_model, llm_router=llm_router) if pricing_read == "ReadUnpriced" else () + ) + if common_flow.read(common_flow.Unpriced(len(unpriced_models))) == "RejectUnpriced": raise ProxyException( message=_unpriced_models_block_message(unpriced_models), type=ProxyErrorTypes.model_cost_map_missing, @@ -1218,34 +987,30 @@ async def common_checks( ) # 1. If team is blocked - if team_object is not None and team_object.blocked is True: + if ( + team_object is not None + and access_policy.decide(access_policy.Blocked(team_object.blocked is True)).denial is not None + ): raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.") # 2. If team can call model (or key's access_group_ids grant it) - if _model and team_object: + if flow.team_models: + assert _model is not None and team_object is not None with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"): - try: - await can_team_access_model( - model=_model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=(valid_token.team_model_aliases if valid_token else None), - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - except ProxyException as team_denial: - if team_denial.type != ProxyErrorTypes.team_model_access_denied: - raise - if not await key_access_group_grants_model( - model=_model, - valid_token=valid_token, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ): - raise + from litellm.proxy.auth.team_override_adapter import can_team_call_model_with_override + + await can_team_call_model_with_override( + model=_model, + valid_token=valid_token, + team_object=team_object, + llm_router=llm_router, + team_model_aliases=(valid_token.team_model_aliases if valid_token else None), + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) # 2.2. If team member has per-member model scope, enforce it - if _model and team_object and valid_token and valid_token.user_id: + if flow.member_models: + assert _model is not None and team_object is not None and valid_token is not None with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_model_access"): await check_team_member_model_access( model=_model, @@ -1261,17 +1026,20 @@ async def common_checks( ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent - if valid_token is not None and valid_token.agent_id: + if common_flow.plan( + replace(flow_facts, agent_id=valid_token.agent_id if valid_token is not None else None) + ).agent_trace: + assert valid_token is not None and valid_token.agent_id is not None from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers agent: Final = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) if agent is not None: require_trace_id: Final = (agent.litellm_params or {}).get("require_trace_id_on_calls_by_agent") - if require_trace_id: + if access_policy.decide(access_policy.TraceRequest(bool(require_trace_id))).next_trace is not None: headers_dict: Final = dict(request.headers) trace_id: Final = get_chain_id_from_headers(headers_dict) - if not trace_id: + if access_policy.decide(access_policy.AgentTrace(bool(require_trace_id), trace_id)).denial is not None: raise ProxyException( message="Requests made with this agent's key must include the x-litellm-trace-id header.", type=ProxyErrorTypes.bad_request_error, @@ -1294,32 +1062,28 @@ async def common_checks( ) or pass_through_route ) - team_trace_id_gate_applies: Final = ( - team_object is not None - and _team_metadata_requires_trace_id(team_object) - and request.method not in ("GET", "HEAD", "OPTIONS") - and (RouteChecks.is_llm_api_route(route=route) or pass_through_route) - and not RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.trace_telemetry_routes.value, - ) + trace_request: Final = access_policy.decide( + access_policy.TraceRequest(team_object is not None and _team_metadata_requires_trace_id(team_object)) ) - team_metadata: Final[Mapping[str, object] | None] = ( - _TRACE_ID_METADATA_ADAPTER.validate_python( - team_object.metadata # pyright: ignore[reportUnknownMemberType] # validate the narrowed untyped field + team_trace: Final = ( + access_policy.decide( + access_policy.TraceGate( + required=True, + method=request.method, + llm_route=RouteChecks.is_llm_api_route(route=route), + passthrough=pass_through_route, + telemetry=RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.trace_telemetry_routes.value + ), + ) ) - if team_trace_id_gate_applies and team_object is not None - else None + if trace_request.next_trace is not None + else trace_request ) if ( - team_trace_id_gate_applies + team_trace.next_trace is not None and team_object is not None - and team_metadata is not None - and not _request_has_trace_id( - request_body=_typed_request_body(request_body), - request=request, - headers_only=headers_only, - ) + and not _request_has_trace_id(_typed_request_body(request_body), request, headers_only=headers_only) ): raise ProxyException( message=( @@ -1342,7 +1106,8 @@ async def common_checks( ) ## 2.1 If user can call model (if personal key) - if _model and team_object is None and user_object is not None: + if flow.user_models: + assert _model is not None and user_object is not None with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): await can_user_call_model( model=_model, @@ -1351,7 +1116,15 @@ async def common_checks( key_model_aliases=key_model_aliases_for_auth_check(valid_token), ) - if end_user_object is not None and end_user_object.models: + customer_plan: Final = common_flow.plan( + replace( + flow_facts, + customer_present=end_user_object is not None, + customer_models=tuple(end_user_object.models or ()) if end_user_object is not None else (), + ) + ) + if customer_plan.customer_models: + assert end_user_object is not None with tracer.trace("litellm.proxy.auth.common_checks.can_customer_call_model"): if _model: can_customer_access_model( @@ -1427,106 +1200,96 @@ async def common_checks( LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body) async def _user_max_budget_check() -> None: - # 4.1 personal budget - if user_object is None or user_object.max_budget is None: - return - is_team_key: Final = team_object is not None and team_object.team_id is not None - if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True: - return + from litellm.proxy.auth import common_budget_adapter - from litellm.proxy.proxy_server import get_current_spend - - user_budget: Final = user_object.max_budget - user_spend: Final = await get_current_spend( - counter_key=user_counter(user_object.user_id), - fallback_spend=user_object.spend or 0.0, - max_budget=user_budget, + return await common_budget_adapter.user_max_budget_check( + user_object, team_object, general_settings, proxy_logging_obj ) - call_info: Final = CallInfo( - spend=user_spend, - max_budget=user_budget, - user_id=user_object.user_id, - user_email=user_object.user_email, - event_group=Litellm_EntityType.USER, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="user_budget", - user_info=call_info, - ) - ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - entity_type=Litellm_EntityType.USER.value, - entity_id=user_object.user_id, - ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead # of one sequential await per scope. return_exceptions lets every scope # settle, then the first error in scope-priority order propagates exactly # as the sequential path raised. - budget_check_coros: Final = tuple( - coro - for coro in ( - team_max_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ), - _team_multi_budget_check(team_object=team_object), - _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None, - _team_soft_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ), - organization_max_budget_check( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - _tag_max_budget_check( - request_body=request_body, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ), - _model_access_group_max_budget_check( - matched_model_access_groups=matched_model_access_groups, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - if matched_model_access_groups - else None, - _user_max_budget_check(), - _check_team_member_budget( - team_object=team_object, - user_object=user_object, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=loaded_team_membership, - team_membership_loaded=team_membership_loaded, - ), - check_end_user_budget(end_user_obj=end_user_object, route=route) - if end_user_object is not None and end_user_object.litellm_budget_table is not None - else None, + async def _check_common_budget(scope: common_flow.BudgetScope) -> None: + match scope: + case "Team": + await team_max_budget_check( + team_object=team_object, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token + ) + case "TeamWindows": + await _team_multi_budget_check(team_object=team_object) + case "KeyWindows": + assert valid_token is not None + await _virtual_key_multi_budget_check(valid_token=valid_token) + case "TeamSoft": + await _team_soft_budget_check( + team_object=team_object, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token + ) + case "Organization": + await organization_max_budget_check( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + case "Tags": + await _tag_max_budget_check( + request_body=request_body, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) + case "Groups": + await _model_access_group_max_budget_check( + matched_model_access_groups=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + case "User": + await _user_max_budget_check() + case "Member": + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, + ) + case "Customer": + assert end_user_object is not None + await check_end_user_budget(end_user_obj=end_user_object, route=route) + case _: + assert_never(scope) + + budget_scopes: Final = common_flow.budgets( + common_flow.BudgetFacts( + token_present=valid_token is not None, + group_count=len(matched_model_access_groups), + customer_present=end_user_object is not None, + customer_budget_present=end_user_object is not None + and end_user_object.litellm_budget_table is not None, ) - if coro is not None ) + budget_check_coros: Final = tuple(_check_common_budget(scope) for scope in budget_scopes) with tracer.trace("litellm.proxy.auth.common_checks.budget_checks"): budget_results: Final = await asyncio.gather(*budget_check_coros, return_exceptions=True) - budget_error: Final = next((r for r in budget_results if isinstance(r, BaseException)), None) - if budget_error is not None: + failure: Final = access_policy.decide( + access_policy.ConcurrentFailures( + tuple(isinstance(result, BaseException) for result in budget_results), + tuple(isinstance(result, NativeAuthFailure) for result in budget_results), + ) + ).failure_index + if failure is not None: + budget_error: Final = budget_results[failure] + if not isinstance(budget_error, BaseException): + raise NativeAuthFailure("native selected an absent budget failure") raise budget_error _enforce_user_param_check(general_settings, request, request_body, route) @@ -1587,18 +1350,17 @@ async def common_checks( def effective_user_role(user_role: str | None) -> LitellmUserRoles: - try: - return LitellmUserRoles(user_role) - except ValueError: - return LitellmUserRoles.INTERNAL_USER + from litellm.proxy.auth import route_role_adapter + + return route_role_adapter.effective_user_role(user_role) def get_user_role( user_obj: LiteLLM_UserTable | None, ) -> LitellmUserRoles | None: - if user_obj is None: - return None - return effective_user_role(user_obj.user_role) + from litellm.proxy.auth import route_role_adapter + + return route_role_adapter.get_user_role(user_obj) _get_user_role: Final = get_user_role @@ -1611,59 +1373,24 @@ def _is_api_route_allowed( valid_token: UserAPIKeyAuth | None, user_obj: LiteLLM_UserTable | None = None, ) -> bool: - """ - - Route b/w api token check and normal token check - """ - _user_role: Final = get_user_role(user_obj=user_obj) + from litellm.proxy.auth import route_role_adapter - if valid_token is None: - raise Exception("Invalid proxy server token passed. valid_token=None.") - - if not is_user_proxy_admin(user_obj=user_obj): # if non-admin - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=_user_role, - route=route, - request=request, - request_data=request_data, - valid_token=valid_token, - ) - return True + return route_role_adapter.is_api_route_allowed(route, request, request_data, valid_token, user_obj) def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: - if user_obj is None: - return False + from litellm.proxy.auth import route_role_adapter - if user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return True - - return False + return route_role_adapter.is_user_proxy_admin(user_obj) _is_user_proxy_admin: Final = is_user_proxy_admin -def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: - """ - Return if a user is allowed to access route. Helper function for `allowed_routes_check`. +def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint + from litellm.proxy.auth import route_role_adapter - Parameters: - - user_route: str - the route the user is trying to call - - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name - (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*"). - """ - from starlette.routing import compile_path - - for allowed_route in allowed_routes: - if allowed_route in LiteLLMRoutes.__members__: - for template in LiteLLMRoutes[allowed_route].value: - regex, _, _ = compile_path(template) - if regex.match(user_route): - return True - elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route): - return True - return False + return route_role_adapter.allowed_routes_list_check(user_route, allowed_routes) def allowed_routes_check( @@ -1671,62 +1398,24 @@ def allowed_routes_check( user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth, ) -> bool: - """ - Check if user -> not admin - allowed to access these routes - """ + from litellm.proxy.auth import route_role_adapter - if user_role == LitellmUserRoles.PROXY_ADMIN: - is_allowed = _allowed_routes_check( - user_route=user_route, - allowed_routes=litellm_proxy_roles.admin_allowed_routes, - ) - return is_allowed - - elif user_role == LitellmUserRoles.TEAM: - if litellm_proxy_roles.team_allowed_routes is None: - """ - By default allow a team to call openai + info routes - """ - is_allowed = _allowed_routes_check(user_route=user_route, allowed_routes=["openai_routes", "info_routes"]) - return is_allowed - elif litellm_proxy_roles.team_allowed_routes is not None: - is_allowed = _allowed_routes_check( - user_route=user_route, - allowed_routes=litellm_proxy_roles.team_allowed_routes, - ) - return is_allowed - return False + return route_role_adapter.allowed_routes_check(user_role, user_route, litellm_proxy_roles) def allowed_route_check_inside_route( user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None, ) -> bool: - ret_val = True - if ( - user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ): - ret_val = False - if requested_user_id is not None and user_api_key_dict.user_id is not None: - if user_api_key_dict.user_id == requested_user_id: - ret_val = True - return ret_val + from litellm.proxy.auth import route_role_adapter + + return route_role_adapter.allowed_route_check_inside_route(user_api_key_dict, requested_user_id) def get_actual_routes(allowed_routes: list) -> list: - actual_routes: Final[list] = [] - for route_name in allowed_routes: - try: - route_value = LiteLLMRoutes[route_name].value - if isinstance(route_value, set): - actual_routes.extend(list(route_value)) - else: - actual_routes.extend(route_value) + from litellm.proxy.auth import route_role_adapter - except KeyError: - actual_routes.append(route_name) - return actual_routes + return route_role_adapter.get_actual_routes(allowed_routes) KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" @@ -1735,11 +1424,9 @@ _OBJECTS_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[obj def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: - """The default budget a key assigns to end users that carry no budget of their own.""" - if key_metadata is None: - return None - budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD) - return budget_id if isinstance(budget_id, str) and budget_id != "" else None + from litellm.proxy.auth import end_user_context_adapter + + return end_user_context_adapter.get_key_end_user_budget_id(key_metadata) @with_service_target(AUTH_OBJECTS_TARGET) @@ -1749,63 +1436,11 @@ async def get_default_end_user_budget( parent_otel_span: Span | None = None, budget_id: str | None = None, ) -> LiteLLM_BudgetTable | None: - """ - Fetches the default end user budget from the database. + from litellm.proxy.auth import end_user_context_adapter - ``budget_id`` selects the budget row; when omitted the proxy-wide - ``litellm.max_end_user_budget_id`` is used. This budget is applied to end - users who don't have an explicit budget_id set. Results are cached for performance. - - Args: - prisma_client: Database client instance - user_api_key_cache: Cache for storing/retrieving budget data - parent_otel_span: Optional OpenTelemetry span for tracing - budget_id: Budget row to load instead of the proxy-wide default - - Returns: - LiteLLM_BudgetTable if configured and found, None otherwise - """ - default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id - if prisma_client is None or default_budget_id is None: - return None - - cache_key: Final = f"default_end_user_budget:{default_budget_id}" - - # Check cache first - cached_budget: Final = await user_api_key_cache.async_get_cache( - key=cache_key, - model_type=LiteLLM_BudgetTable, + return await end_user_context_adapter.get_default_end_user_budget( + prisma_client, user_api_key_cache, parent_otel_span, budget_id ) - if cached_budget is not None: - return cached_budget - - # Fetch from database - try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( - where={"budget_id": default_budget_id} - ) - - if budget_record is None: - verbose_proxy_logger.warning( - "Default end user budget not found in database: %s", - default_budget_id.replace("\r", "").replace("\n", ""), - ) - return None - - _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) - # Cache the budget for 60 seconds - await user_api_key_cache.async_set_cache( - key=cache_key, - value=_budget_obj, - model_type=LiteLLM_BudgetTable, - ttl=get_management_object_ttl(user_api_key_cache), - ) - - return _budget_obj - - except Exception as e: - verbose_proxy_logger.error("Error fetching default end user budget: %s", e) - return None @log_db_metrics @@ -1847,6 +1482,8 @@ async def get_team_member_default_budget( budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( where={"budget_id": budget_id} ) + except NativeAuthFailure: + raise except Exception: verbose_proxy_logger.exception("Error fetching team-default member budget %s", budget_id) if raise_on_lookup_error: @@ -1873,13 +1510,9 @@ def resolve_effective_team_member_budget( team_default_budget: LiteLLM_BudgetTable | None, now: datetime, ) -> float | None: - if member_budget is not None and member_budget.max_budget is not None: - return member_budget.effective_max_budget(now) - if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: - return None - return team_default_budget.max_budget + ( - member_budget.active_temp_budget_increase(now) if member_budget is not None else 0.0 - ) + from litellm.proxy.auth import common_budget_adapter + + return common_budget_adapter.resolve_effective_team_member_budget(member_budget, team_default_budget, now) async def resolve_default_end_user_budget( @@ -1888,30 +1521,10 @@ async def resolve_default_end_user_budget( key_end_user_budget_id: str | None, parent_otel_span: Span | None = None, ) -> LiteLLM_BudgetTable | None: - """ - The default budget for an end user with no budget of its own. + from litellm.proxy.auth import budget_context_adapter - The key's ``end_user_budget_id`` takes precedence over the proxy-wide - ``litellm.max_end_user_budget_id``; the proxy-wide default is the fallback when the key - names no budget or its budget row is missing. - """ - if key_end_user_budget_id is not None: - key_budget: Final = await get_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - budget_id=key_end_user_budget_id, - ) - if key_budget is not None: - return key_budget - - if litellm.max_end_user_budget_id is None: - return None - - return await get_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, + return await budget_context_adapter.resolve_default_end_user_budget( + prisma_client, user_api_key_cache, key_end_user_budget_id, parent_otel_span ) @@ -1922,83 +1535,20 @@ async def _apply_default_budget_to_end_user( parent_otel_span: Span | None = None, key_end_user_budget_id: str | None = None, ) -> LiteLLM_EndUserTable: - """ - Returns the end user with the resolved default budget when it has no budget of its own. + from litellm.proxy.auth import end_user_context_adapter - A row whose own ``budget_id`` resolved to a budget is returned unchanged. Otherwise the - default is resolved on every call and set on a copy: the cached row carries at most the - proxy-wide default (readers such as the Prometheus customer gauges rely on that), never a - key's, so requests through keys with different defaults never observe each other's budget. - - Args: - end_user_obj: The end user object to potentially apply default budget to - prisma_client: Database client instance - user_api_key_cache: Cache for storing/retrieving data - parent_otel_span: Optional OpenTelemetry span for tracing - key_end_user_budget_id: The requesting key's ``end_user_budget_id``, if any - """ - if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None: - return end_user_obj - - if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None: - return end_user_obj - - default_budget: Final = await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, + return await end_user_context_adapter.apply_default_budget_to_end_user( + end_user_obj, prisma_client, user_api_key_cache, parent_otel_span, key_end_user_budget_id ) - if default_budget is None: - return end_user_obj - - verbose_proxy_logger.debug( - "Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id - ) - return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget})) - async def check_end_user_budget( end_user_obj: LiteLLM_EndUserTable, route: str, ) -> None: - """ - Check if end user is within their budget limit. + from litellm.proxy.auth import common_budget_adapter - Args: - end_user_obj: The end user object to check - route: The request route - - Raises: - litellm.BudgetExceededError: If end user has exceeded their budget - """ - if RouteChecks.is_info_route(route): - return - - if end_user_obj.litellm_budget_table is None: - return - - end_user_budget: Final = end_user_obj.litellm_budget_table.max_budget - if end_user_budget is None: - return - - from litellm.proxy.proxy_server import get_current_spend - - end_user_spend: Final = await get_current_spend( - counter_key=end_user_counter(end_user_obj.user_id), - fallback_spend=end_user_obj.spend or 0.0, - max_budget=end_user_budget, - fallback_authoritative=True, - ) - if end_user_spend > end_user_budget: - raise litellm.BudgetExceededError( - current_cost=end_user_spend, - max_budget=end_user_budget, - message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", - entity_type=Litellm_EntityType.END_USER.value, - entity_id=end_user_obj.user_id, - ) + return await common_budget_adapter.check_end_user_budget(end_user_obj, route) _check_end_user_budget: Final = check_end_user_budget @@ -2063,6 +1613,8 @@ async def _cache_registry_answer( """Best-effort: a cache backend failure must not turn a registry load into a failed request.""" try: await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error verbose_proxy_logger.warning("Failed to cache registry %s: %s", cache_key, e) @@ -2077,6 +1629,8 @@ async def _fetch_and_cache_registry( """The registry as the database has it, cached whole, or ``None`` when it is unusable.""" try: registry_ids: Final = await fetch_ids() + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # fail-safe: any registry load error must degrade to per-id lookups, never break auth verbose_proxy_logger.warning( "Registry %s could not be loaded from the database, so per-id lookups will run and the " @@ -2145,7 +1699,7 @@ async def _load_bounded_registry( ) -async def _load_end_user_restricted_registry( +async def load_end_user_restricted_registry( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, ) -> frozenset[str] | None: @@ -2168,6 +1722,9 @@ async def _load_end_user_restricted_registry( ) +_load_end_user_restricted_registry: Final = load_end_user_restricted_registry + + async def _end_user_is_known_unrestricted( end_user_id: str, prisma_client: PrismaClient, @@ -2175,30 +1732,11 @@ async def _end_user_is_known_unrestricted( token_end_user_max_budget: float | None, key_end_user_budget_id: str | None = None, ) -> bool: - """ - True when the cached registry proves the id restricts nothing, so its row need not be read. + from litellm.proxy.auth import end_user_context_adapter - Every field ``get_end_user_object`` callers consume (budget, spend under that budget, region, - default model, models, object permission, blocked) is part of the registry predicate, so an id - outside it is indistinguishable from one with no row at all. The skip is off whenever mere existence of - the row is meaningful: ``max_end_user_budget_id`` or the key's ``end_user_budget_id`` grafts a - default budget onto any row that exists, ``validate_end_user_id_in_db`` rejects ids that resolve - to no row, and a token-supplied ``end_user_max_budget`` (a ``user_custom_auth`` callable can set - one against an otherwise unrestricted row) is enforced against the row's recorded spend. - """ - if ( - litellm.max_end_user_budget_id is not None - or key_end_user_budget_id is not None - or litellm.validate_end_user_id_in_db - or token_end_user_max_budget is not None - ): - return False - - registry: Final = await _load_end_user_restricted_registry( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + return await end_user_context_adapter.end_user_is_known_unrestricted( + end_user_id, prisma_client, user_api_key_cache, token_end_user_max_budget, key_end_user_budget_id ) - return registry is not None and end_user_id not in registry @log_db_metrics @@ -2300,6 +1838,8 @@ async def get_end_user_object( key_end_user_budget_id=key_end_user_budget_id, ) + except NativeAuthFailure: + raise except Exception: return None @@ -2318,61 +1858,20 @@ async def resolve_and_validate_end_user_id( route: str = "", key_end_user_budget_id: str | None = None, ) -> str | None: - """Optionally drop end-user ids that don't resolve to a known DB row. + from litellm.proxy.auth import end_user_context_adapter - Default: pass-through. LiteLLM's documented pattern is that the `user` - field is an arbitrary caller-supplied identifier, so validation is - opt-in behind ``litellm.validate_end_user_id_in_db`` to preserve - backwards compatibility. - - When the flag is set: accept the id when it matches any of - - LiteLLM_EndUserTable.user_id - - LiteLLM_UserTable.user_id - - LiteLLM_UserTable.user_email (case-insensitive) - - If the id doesn't match but a default end-user budget is configured - (``litellm.max_end_user_budget_id`` or the key's ``end_user_budget_id``), - we still preserve the id so that budget is applied downstream; otherwise - we return None. - - DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they - share the same cache as the rest of the auth path instead of adding new - raw Prisma queries. - """ - if raw_end_user_id is None: - return None - if not litellm.validate_end_user_id_in_db: - return raw_end_user_id - if prisma_client is None: - return raw_end_user_id - - has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None - cache_key: Final = f"end_user_validation:{raw_end_user_id}" - cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) - if cached == "valid": - return raw_end_user_id - if cached == "invalid": - return raw_end_user_id if has_default_budget else None - - is_valid: Final = await _end_user_id_exists_in_db( - end_user_id=raw_end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, + return await end_user_context_adapter.resolve_and_validate_end_user_id( + raw_end_user_id, + prisma_client, + user_api_key_cache, + parent_otel_span, + proxy_logging_obj, + route, + key_end_user_budget_id, ) - await user_api_key_cache.async_set_cache( - key=cache_key, - value="valid" if is_valid else "invalid", - ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL), - ) - return raw_end_user_id if is_valid or has_default_budget else None - - -async def _end_user_id_exists_in_db( +async def end_user_id_exists_in_db( end_user_id: str, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, @@ -2380,38 +1879,14 @@ async def _end_user_id_exists_in_db( proxy_logging_obj: ProxyLogging | None = None, route: str = "", ) -> bool: - """True when the id matches an EndUser, User, or user_email row.""" - try: - end_user_obj: Final = await get_end_user_object( - end_user_id=end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - if end_user_obj is not None: - return True - except Exception as e: - verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", e) + from litellm.proxy.auth import end_user_context_adapter - try: - user_obj: Final = await get_user_object( - user_id=end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_db_only=False, - user_email=end_user_id if "@" in end_user_id else None, - ) - if user_obj is not None: - return True - except Exception as e: - verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", e) + return await end_user_context_adapter.end_user_id_exists_in_db( + end_user_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route + ) - return False + +_end_user_id_exists_in_db: Final = end_user_id_exists_in_db async def _load_tag_registry( @@ -2436,7 +1911,7 @@ async def _load_tag_registry( ) -async def _load_model_access_group_registry( +async def load_model_access_group_registry( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, ) -> frozenset[str] | None: @@ -2458,6 +1933,9 @@ async def _load_model_access_group_registry( ) +_load_model_access_group_registry: Final = load_model_access_group_registry + + @with_service_target(AUTH_OBJECTS_TARGET) async def _fetch_uncached_model_access_group_budgets( uncached_groups: Sequence[str], @@ -2488,6 +1966,8 @@ async def _fetch_uncached_model_access_group_budgets( model_type=ModelAccessGroupBudget, ttl=get_management_object_ttl(user_api_key_cache), ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e) return () @@ -2570,6 +2050,8 @@ async def _fetch_uncached_tags( model_type=LiteLLM_TagTable, ttl=get_management_object_ttl(user_api_key_cache), ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # returned as a value: budget checks fail open, the baggage mapper fails closed return e else: @@ -2828,31 +2310,9 @@ async def get_team_membership( def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: - from collections import defaultdict + from litellm.proxy.auth import model_helper_adapter - if team_models is None: - return True - if model in team_models: - return True - - access_groups: dict[str, list[str]] = defaultdict(list) - if llm_router: - access_groups = llm_router.registry().access_groups(model_name=model) - - if len(access_groups) > 0: # check if token contains any model access groups - for idx, m in enumerate( - team_models - ): # loop token models, if any of them are an access group add the access group - if m in access_groups: - return True - - # Filter out models that are access_groups - filtered_models: Final = [m for m in team_models if m not in access_groups] - - if model in filtered_models: - return True - - return False + return model_helper_adapter.model_in_access_group(model, team_models, llm_router) def _should_check_db(key: str, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int) -> bool: @@ -3150,6 +2610,8 @@ async def get_user_object( return _response except UserNotFoundError: raise + except NativeAuthFailure: + raise except Exception as e: _log_budget_lookup_failure("user", e) raise _user_read_failure(user_id=user_id, error=e) @@ -3233,6 +2695,8 @@ async def cache_team_object( await user_api_key_cache.async_delete_cache(key=alias_key) else: await pipelined_delete + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the mutation verbose_proxy_logger.warning( "Failed to invalidate cached team alias entry %s; " @@ -3261,6 +2725,8 @@ async def _invalidate_usage_cache_entry( usage_cache.in_memory_cache.delete_cache(key) else: await usage_cache.async_delete_cache(key=key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the write verbose_proxy_logger.warning( "Failed to invalidate internal usage cache entry %s; a stale %s may be served until its TTL expires: %s", @@ -3360,6 +2826,8 @@ async def invalidate_team_member_spend_state( ) try: await spend_counter_store.delete((spend_counter_key,)) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success verbose_proxy_logger.warning( "Failed to delete stale spend counter %s in Redis after a failed reset write", @@ -3425,6 +2893,8 @@ async def delete_cache_team_object( ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete verbose_proxy_logger.warning( "Failed to invalidate cached team entry %s on delete; " @@ -3487,6 +2957,8 @@ async def delete_cache_key_object( ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # best-effort: a cache error must not fail a committed write verbose_proxy_logger.warning( "Failed to invalidate cached key entry %s; a stale key object may be served until its TTL expires: %s", @@ -3608,10 +3080,6 @@ async def _get_team_db_check( return response -async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient) -> "_PrismaTeamRow | None": - return await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) - - async def _get_team_object_from_user_api_key_cache( team_id: str, prisma_client: PrismaClient, @@ -3656,6 +3124,8 @@ async def _get_team_object_from_user_api_key_cache( proxy_logging_obj=proxy_logging_obj, check_db_only=use_writer, ) + except NativeAuthFailure: + raise except Exception as e: if use_writer: raise @@ -3756,6 +3226,8 @@ async def get_team_object( ) except TeamNotFoundError: raise + except NativeAuthFailure: + raise except Exception: raise HTTPException( status_code=404, @@ -3854,6 +3326,8 @@ async def get_access_object( return _response except HTTPException: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.exception( "Error getting access group for access_group_id: %s", @@ -3944,6 +3418,8 @@ async def get_team_object_by_alias( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug( "Failed to load object_permission for team %s with object_permission_id=%s: %s", @@ -3972,6 +3448,8 @@ async def get_team_object_by_alias( except HTTPException: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias) raise HTTPException( @@ -4061,6 +3539,8 @@ async def get_org_object_by_alias( except HTTPException: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.exception("Error looking up organization by alias: %s", org_alias) raise HTTPException( @@ -4195,6 +3675,8 @@ class ExperimentalUIJWTToken: return None try: return UserAPIKeyAuth.model_validate(json.loads(decrypted_token)) + except NativeAuthFailure: + raise except Exception as e: raise Exception(f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}") @@ -4247,6 +3729,8 @@ async def _fetch_key_object_from_db_unbounded( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: if PrismaDBExceptionHandler.is_database_transport_error(e): did_reconnect = False @@ -4407,6 +3891,8 @@ async def get_key_object( proxy_logging_obj=proxy_logging_obj, check_db_only=check_db_only, ) + except NativeAuthFailure: + raise except Exception as e: if check_db_only: raise @@ -4495,6 +3981,8 @@ async def get_object_permission( ) return _perm_obj + except NativeAuthFailure: + raise except Exception: if check_db_only: raise @@ -4612,21 +4100,13 @@ async def get_org_object( if deserialized_org is not None: return deserialized_org # else, check db - try: - query_kwargs: Final[dict[str, Mapping[str, object]]] = {"where": {"organization_id": org_id}} - if include_budget_table: - query_kwargs["include"] = {"litellm_budget_table": True} + query_kwargs: Final[dict[str, Mapping[str, object]]] = {"where": {"organization_id": org_id}} + if include_budget_table: + query_kwargs["include"] = {"litellm_budget_table": True} - response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique( - **query_kwargs - ) - except Exception: - # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed - # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them - # apart — a caller that treats absence as "this org places no restriction" then drops a real - # org ceiling during an outage. Propagate the real error; callers that already catch - # Exception are unaffected. - raise + response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique( + **query_kwargs + ) if response is None: raise OrganizationNotFoundError( @@ -4686,6 +4166,8 @@ async def get_org_object_for_request( ) except OrganizationNotFoundError: return None + except NativeAuthFailure: + raise except Exception as e: if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) @@ -4749,6 +4231,8 @@ async def get_access_group_tables( raise verbose_proxy_logger.debug("One of the requested access groups does not exist, it grants nothing") return None + except NativeAuthFailure: + raise except Exception: if check_db_only or unreadable == "raise": raise @@ -4895,20 +4379,6 @@ async def get_skill_names_from_access_groups( ) -def _resolve_all_team_model_sentinel_for_auth_check( - models: list[str], - llm_router: RegistryReader | None, - team_id: str | None, -) -> list[str]: - if SpecialModelNames.all_team_models.value not in models or team_id is None or llm_router is None: - return models - proxy_models: Final = llm_router.registry().names() - non_sentinel_models: Final = [model for model in models if model != SpecialModelNames.all_team_models.value] - if not proxy_models: - return non_sentinel_models or models - return list(dict.fromkeys(non_sentinel_models + proxy_models)) - - def check_model_access_helper( model: str, llm_router: RegistryReader | None, @@ -4917,65 +4387,14 @@ def check_model_access_helper( team_id: str | None = None, ) -> bool: ## check if model in allowed model names - from collections import defaultdict + from litellm.proxy.auth import model_helper_adapter - access_groups: dict[str, list[str]] = defaultdict(list) - - if llm_router: - access_groups = llm_router.registry().access_groups(model_name=model, team_id=team_id) - - models = _resolve_all_team_model_sentinel_for_auth_check( - models=models, - llm_router=llm_router, - team_id=team_id, - ) - - if len(access_groups) > 0 and llm_router is not None: # check if token contains any model access groups - for idx, m in enumerate(models): # loop token models, if any of them are an access group add the access group - if m in access_groups: - return True - - # Filter out models that are access_groups - filtered_models: Final = [m for m in models if m not in access_groups] - - if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): - return True - - if model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): - return True - - all_model_access: bool = False - - if (len(filtered_models) == 0 and len(models) == 0) or "*" in filtered_models: - all_model_access = True - - if SpecialModelNames.all_proxy_models.value in filtered_models: - all_model_access = True - - if model is not None and model not in filtered_models and all_model_access is False: - return False - return True + return model_helper_adapter.check_model_access_helper(model, llm_router, models, team_model_aliases, team_id) _check_model_access_helper: Final = check_model_access_helper -def _owning_team_names_for_team_deployment( - model: str, - llm_router: RegistryReader | None, - team_id: str | None, - team_model_aliases: Mapping[str, str] | None, -) -> tuple[str, ...]: - if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): - return () - deployment: Final = llm_router.registry().deployments(model, source="first_in_group") - if deployment is None or deployment.model_info.team_id != team_id: - return () - public_name: Final = deployment.model_info.team_public_model_name - legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) - return (*((public_name,) if public_name else ()), *legacy_aliases) - - def _can_object_call_model( model: str | list[str], llm_router: RegistryReader | None, @@ -4986,136 +4405,32 @@ def _can_object_call_model( object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"] = "user", fallback_depth: int = 0, ) -> Literal[True]: - """ - Checks if token can call a given model + from litellm.proxy.auth import model_access_adapter - Args: - - model: str - - llm_router: Optional[Router] - - models: List[str] - - team_model_aliases: Optional[Dict[str, str]] - - object_type: Literal["user", "team", "key", "org"]. We use the object type to raise the correct exception type - - Returns: - - True: if token allowed to call model - - Raises: - - Exception: If token not allowed to call model - """ - if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: - raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") - if isinstance(model, list): - for m in model: - can_object_call_model( - model=m, - llm_router=llm_router, - models=models, - team_model_aliases=team_model_aliases, - team_id=team_id, - key_model_aliases=key_model_aliases, - object_type=object_type, - fallback_depth=fallback_depth + 1, - ) - return True - - from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent - - compaction_parent: Final = native_compaction_parent(model) - global_or_router_alias_target: Final = ( - litellm.model_alias_map[model] - if model in litellm.model_alias_map - else ( - llm_router.registry().resolve(model, kind="alias") - if llm_router is not None and model in llm_router.registry().aliases() - else None - ) - ) - after_team_alias: Final = team_model_aliases.get(model, model) if team_model_aliases else model - after_key_alias: Final = ( - key_model_aliases.get(after_team_alias, after_team_alias) if key_model_aliases else after_team_alias - ) - after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias) - dispatched_model: Final = ( - key_model_aliases.get(after_global_alias, after_global_alias) if key_model_aliases else after_global_alias - ) - key_alias_applied: Final = after_key_alias != after_team_alias or dispatched_model != after_global_alias - potential_models: Final = ( - (dispatched_model,) - if key_alias_applied - else ( - *((model, compaction_parent) if compaction_parent is not None else (model,)), - *((global_or_router_alias_target,) if global_or_router_alias_target else ()), - ) - ) - owning_team_names: Final = ( - () - if key_alias_applied - else _owning_team_names_for_team_deployment( - model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases - ) - ) - - ## check model access for alias + underlying model - allow if either is in allowed models - for m in potential_models: - if check_model_access_helper( - model=m, - llm_router=llm_router, - models=models, - team_model_aliases=team_model_aliases, - team_id=team_id, - ): - return True - if any( - check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id) - for name in owning_team_names - ): - return True - - internal_message: Final = ( - f"{object_type} not allowed to access model. This {object_type} can only access models={models}. " - f"Tried to access {model}" - ) - raise ModelAccessDeniedProxyException( - message=( - customer_model_access_denied_client_message(model=model) - if object_type == "customer" - else model_access_denied_client_message(model=model) - ), - internal_message=internal_message, - type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type=object_type), - param="model", - code=status.HTTP_403_FORBIDDEN, + return model_access_adapter.can_object_call_model( + model=model, + llm_router=llm_router, + models=models, + team_model_aliases=team_model_aliases, + team_id=team_id, + key_model_aliases=key_model_aliases, + object_type=object_type, + fallback_depth=fallback_depth, ) can_object_call_model: Final = _can_object_call_model -def _resolve_team_alias( +def _resolve_team_alias( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint model: str | list[str], team_model_aliases: Mapping[str, str] | None, team_id: str | None, llm_router: Router | None, ) -> str | list[str]: - if not team_model_aliases: - return model - if isinstance(model, str): - return _live_team_alias_target(model, team_model_aliases, team_id, llm_router) - return [_live_team_alias_target(name, team_model_aliases, team_id, llm_router) for name in model] + from litellm.proxy.auth import team_alias_adapter - -def _live_team_alias_target( - model: str, team_model_aliases: Mapping[str, str], team_id: str | None, llm_router: RegistryReader | None -) -> str: - target: Final = team_model_aliases.get(model) - if target is None: - return model - deleted_team_deployment: Final = ( - llm_router is not None - and target.startswith(f"model_name_{team_id}_") - and target not in llm_router.registry().index("model_name") - ) - return model if deleted_team_deployment else target + return team_alias_adapter.resolve_team_alias(model, team_model_aliases, team_id, llm_router) async def can_agent_call_model( @@ -5126,93 +4441,46 @@ async def can_agent_call_model( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: - managed_policy: Final = managed_agent_policy(valid_token) - if model and valid_token is not None and managed_policy is not None: - managed_models: Final = (managed_policy.object_permission or {}).get("models", ()) - if not isinstance(managed_models, (list, tuple)) or not managed_models: - raise HTTPException(403, "This agent has no model grants") - can_object_call_model( - model=_resolve_team_alias( - model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router - ), - llm_router=llm_router, - models=list(managed_models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - await _check_agent_access_group_model_access(model=model, valid_token=valid_token, llm_router=llm_router) - await _check_agent_caller_model_access( - model=model, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.can_agent_call_model( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj ) -async def _agent_access_group_ceilings( +async def _agent_access_group_ceilings( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver ) -> tuple[AgentAccessGroupCeiling, ...]: - if not valid_token.agent_id: - return () + from litellm.proxy.auth import managed_model_adapter - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - managed: Final = managed_agent_policy(valid_token) - if managed is not None: - return await resolve_managed_agent_ceilings(managed) - unmanaged: Final = await resolve_ceiling(valid_token.agent_id) - return (unmanaged,) if unmanaged is not None else () + return await managed_model_adapter.agent_access_group_ceilings(valid_token, resolve_ceiling) -def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: - managed: Final = managed_agent_policy(valid_token) - granted: Final = (managed.object_permission or {}).get("models") if managed is not None else None - if not isinstance(granted, (list, tuple)): - return frozenset() - return frozenset(name for name in _OBJECTS_ADAPTER.validate_python(granted) if isinstance(name, str)) +def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint + from litellm.proxy.auth import managed_model_adapter + + return managed_model_adapter.managed_agent_model_grants(valid_token) async def agent_model_grants( valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling ) -> frozenset[str]: - ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) - return _managed_agent_model_grants(valid_token).union(*(ceiling.models for ceiling in ceilings)) + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.agent_model_grants(valid_token, resolve_ceiling) -async def _check_agent_access_group_model_access( +async def _check_agent_access_group_model_access( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, llm_router: Router | None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: - """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" - if not model or valid_token is None: - return True - ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) - dispatched: Final = _resolve_team_alias( - model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.check_agent_access_group_model_access( + model, valid_token, llm_router, resolve_ceiling ) - for ceiling in ceilings: - if not ceiling.models: - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=f"agent {valid_token.agent_id} access groups grant no models", - type=ProxyErrorTypes.agent_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - can_object_call_model( - model=dispatched, - llm_router=llm_router, - models=sorted(ceiling.models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - return True LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None @@ -5221,7 +4489,7 @@ CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerT CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerUser]] -async def _check_agent_caller_model_access( +async def _check_agent_caller_model_access( # pyright: ignore[reportUnusedFunction] # backwards-compatible auth helper entrypoint model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str] valid_token: UserAPIKeyAuth | None, llm_router: Router | None, @@ -5231,63 +4499,13 @@ async def _check_agent_caller_model_access( load_team: CallerTeamLoader = load_agent_caller_team, load_user: CallerUserLoader = load_agent_caller_user, ) -> None: - """An agent key acting for an invoking user may call only what that user's own key could: the - invoking team's models (and per-member scope) when a team was echoed, else the user's models.""" - if not model or valid_token is None: - return - caller_auth: Final = agent_caller_auth(valid_token) - if caller_auth is None: - return - caller_team: Final = await load_team(valid_token) - caller_key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) - if caller_team is not None: - await can_team_access_model( - model=model, - team_object=caller_team, - llm_router=llm_router, - prisma_client=prisma_client, - key_model_aliases=caller_key_model_aliases, - ) - await check_team_member_model_access( - model=model, - team_object=caller_team, - valid_token=caller_auth, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - key_model_aliases=caller_key_model_aliases, - ) - return - caller_user: Final = await load_user(valid_token) - if caller_user is None: - return - await can_user_call_model( - model=model, - llm_router=llm_router, - user_object=caller_user, - key_model_aliases=caller_key_model_aliases, + from litellm.proxy.auth import managed_model_adapter + + return await managed_model_adapter.check_agent_caller_model_access( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj, load_team, load_user ) -def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: - """ - Returns True if `model` being accessed is an alias of a team model - - - `model=gpt-4o` - - `team_model_aliases={"gpt-4o": "gpt-4o-team-1"}` - - returns True - - - `model=gp-4o` - - `team_model_aliases={"o-3": "o3-preview"}` - - returns False - """ - if team_model_aliases: - if model in team_model_aliases: - return True - return False - - def key_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth | None) -> Mapping[str, str] | None: return alias_map(valid_token.aliases) if valid_token is not None and valid_token.aliases else None @@ -5296,7 +4514,7 @@ def team_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth) -> Mapping[st return alias_map(valid_token.team_model_aliases) if valid_token.team_model_aliases else None -def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]: +def resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]: """ Expand key model sentinels before auth checks. @@ -5314,101 +4532,7 @@ def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str] return models -def _model_access_groups_serving_model( - model: str | Sequence[str], - llm_router: Router, - team_id: str | None, -) -> frozenset[str]: - """Every model access group whose deployments serve the requested model(s).""" - requested: Final = (model,) if isinstance(model, str) else tuple(model) - return frozenset( - group - for requested_model in requested - for group in llm_router.registry().access_groups(model_name=requested_model, team_id=team_id) - ) - - -async def _team_member_granted_models( - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - team_membership: LiteLLM_TeamMembership | None = None, - team_membership_loaded: bool = False, -) -> Sequence[str]: - """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" - if team_object is None or valid_token.user_id is None: - return () - - if not team_membership_loaded: - team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - return () if team_membership is None else _member_allowed_models(team_membership) - - -async def _org_granted_models( - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> Sequence[str]: - """The org allowlist reached through the key, or through its team when the key names no org.""" - org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) - if org_id is None: - return () - - try: - org_object: Final = await get_org_object( - org_id=org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: # noqa: BLE001 # fail-safe: attribution degrades to "no org grant", it must never break auth - verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", e) - return () - return org_object.models if org_object is not None else () - - -async def _granted_model_lists( - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, - project_object: LiteLLM_ProjectTableCachedObj | None, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - team_membership: LiteLLM_TeamMembership | None = None, - team_membership_loaded: bool = False, -) -> tuple[Sequence[str], ...]: - """One model allowlist per level that participates in authorizing the request.""" - return ( - _resolve_key_models_for_auth_check(valid_token=valid_token), - team_object.models if team_object is not None else (), - await _team_member_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=team_membership_loaded, - ), - project_object.models if project_object is not None else (), - await _org_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - ) +_resolve_key_models_for_auth_check: Final = resolve_key_models_for_auth_check async def collect_matched_model_access_groups( @@ -5423,56 +4547,20 @@ async def collect_matched_model_access_groups( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ) -> tuple[str, ...]: - """ - The budgeted model access groups that authorized this request, sorted and deduplicated. + from litellm.proxy.auth import model_attribution_adapter - A group is charged only when its name appears on an allowlist the caller was granted -- key, - team, team-member scope, project or org -- *and* that group serves the requested model. Asking - for a model that merely belongs to a group attributes nothing, because nothing about the caller - named the group. - - Levels are unioned, never ranked: a team granted ``*`` whose member is scoped to one group is - still a caller gated by that group. An unrestricted allowlist (empty, ``*``) names no group and - so contributes nothing. - - The whole walk is gated on the budget registry, because collecting every match costs a full scan - of each allowlist where the plain access check stops at the first hit. An empty registry means no - group carries a budget, so there is nothing to attribute and no work worth doing. - """ - if model is None or valid_token is None or llm_router is None or prisma_client is None: - return () - - registry: Final = await _load_model_access_group_registry( + return await model_attribution_adapter.collect_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) - if registry is not None and not registry: - return () - - covering_groups: Final = _model_access_groups_serving_model( - model=model, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - budgeted_groups: Final = covering_groups if registry is None else covering_groups & registry - if not budgeted_groups: - return () - - granted: Final = frozenset( - granted_model - for granted_models in await _granted_model_lists( - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=team_membership_loaded, - ) - for granted_model in granted_models - ) - return tuple(sorted(budgeted_groups & granted)) async def stamp_matched_model_access_groups( @@ -5487,32 +4575,20 @@ async def stamp_matched_model_access_groups( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ) -> tuple[str, ...]: - """Record the groups that authorized this request on its auth object, for the post-call spend - writer and the reservation counters, and hand them back for the budget check.""" - if valid_token is None: - return () + from litellm.proxy.auth import model_attribution_adapter - try: - matched: Final = await collect_matched_model_access_groups( - model=model, - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=team_membership_loaded, - ) - except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth - verbose_proxy_logger.debug("model access group attribution failed: %s", e) - return () - if not matched: - return () - matched_groups: Final = list(matched) - valid_token.matched_model_access_groups = matched_groups # rebind-ok: request-scoped carrier for the writer - return matched + return await model_attribution_adapter.stamp_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, + ) async def can_key_call_model( @@ -5522,76 +4598,15 @@ async def can_key_call_model( llm_router: RegistryReader | None, prisma_client: DatabaseClient | None = None, ) -> Literal[True]: - """ - Checks if token can call a given model + from litellm.proxy.auth import model_access_adapter - 1. First checks native key-level model permissions (current implementation) - 2. If not allowed natively, falls back to access_group_ids on the key - - Returns: - - True: if token allowed to call model - - Raises: - - Exception: If token not allowed to call model - """ - key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) - try: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=key_models, - team_model_aliases=valid_token.team_model_aliases, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="key", - ) - except ProxyException: - # Fallback: check key's access_group_ids - key_access_group_ids: Final = valid_token.access_group_ids or [] - if key_access_group_ids: - models_from_groups: Final = await get_models_from_access_groups( - access_group_ids=key_access_group_ids, - prisma_client=prisma_client, - ) - if models_from_groups: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=models_from_groups, - team_model_aliases=valid_token.team_model_aliases, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="key", - ) - raise - - -async def _personal_key_owner( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> LiteLLM_UserTable | None: - if valid_token.user_id is None: - return None - try: - return await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except Exception as e: # noqa: BLE001 # fail closed: an owner that cannot be loaded grants nothing - verbose_proxy_logger.warning( - "Personal key team model access: owner lookup failed, granting nothing: %s", - e, - ) - return None + return await model_access_adapter.can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + ) async def can_key_call_resolved_model( @@ -5600,128 +4615,9 @@ async def can_key_call_resolved_model( valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None, ) -> None: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.auth import resolved_model_adapter - skip_key_model_check: Final = valid_token.config or ( - isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models - ) - if not skip_key_model_check: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - - team_object: LiteLLM_TeamTableCachedObj | None = None - team_object_from_lookup = False - if valid_token.team_id is not None: - try: - team_object = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - team_object_from_lookup = True - except Exception: - team_object = LiteLLM_TeamTableCachedObj( - team_id=valid_token.team_id, - models=valid_token.team_models, - blocked=valid_token.team_blocked, - team_alias=valid_token.team_alias, - metadata=valid_token.team_metadata, - object_permission_id=valid_token.team_object_permission_id, - object_permission=valid_token.team_object_permission, - ) - - if team_object is not None: - try: - await can_team_access_model( - model=model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=valid_token.team_model_aliases, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - except ProxyException as team_denial: - if team_denial.type != ProxyErrorTypes.team_model_access_denied: - raise - if not await key_access_group_grants_model( - model=model, - valid_token=valid_token, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ): - raise - - if valid_token.user_id is not None and team_object_from_lookup: - await check_team_member_model_access( - model=model, - team_object=team_object, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - if personal_key_team_model_access_applies(valid_token) and valid_token.user_id is not None: - owner: Final = await _personal_key_owner( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - await can_personal_key_call_model_via_teams( - model=model, - user_object=owner, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - if valid_token.project_id is not None: - project_object: Final = await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if project_object is not None and len(project_object.models) > 0: - can_project_access_model( - model=model, - project_object=project_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - if valid_token.end_user_id is not None and prisma_client is not None: - key_metadata: Final = _KEY_METADATA_ADAPTER.validate_python(valid_token.metadata) - end_user_object: Final = await get_end_user_object( - end_user_id=valid_token.end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - token_end_user_max_budget=valid_token.end_user_max_budget, - key_end_user_budget_id=get_key_end_user_budget_id(key_metadata), - ) - if end_user_object is not None and end_user_object.models: - can_customer_access_model( - model=model, - end_user_object=end_user_object, - llm_router=llm_router, - valid_token=valid_token, - ) + return await resolved_model_adapter.can_key_call_resolved_model(model, llm_model_list, valid_token, llm_router) def can_org_access_model( @@ -5730,16 +4626,10 @@ def can_org_access_model( llm_router: Router | None, team_model_aliases: dict[str, str] | None = None, ) -> Literal[True]: - """ - Returns True if the team can access a specific model. + from litellm.proxy.auth import model_access_adapter - """ - return can_object_call_model( - model=model, - llm_router=llm_router, - models=org_object.models if org_object else [], - team_model_aliases=team_model_aliases, - object_type="org", + return model_access_adapter.can_org_access_model( + model=model, org_object=org_object, llm_router=llm_router, team_model_aliases=team_model_aliases ) @@ -5751,41 +4641,16 @@ async def can_team_access_model( key_model_aliases: Mapping[str, str] | None = None, prisma_client: DatabaseClient | None = None, ) -> Literal[True]: - """ - Returns True if the team can access a specific model. + from litellm.proxy.auth import model_access_adapter - 1. First checks native team-level model permissions (current implementation) - 2. If not allowed natively, falls back to access_group_ids on the team - """ - try: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=team_object.models if team_object else [], - team_model_aliases=team_model_aliases, - team_id=team_object.team_id if team_object else None, - key_model_aliases=key_model_aliases, - object_type="team", - ) - except ProxyException: - # Fallback: check team's access_group_ids - team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else [] - if team_access_group_ids: - models_from_groups: Final = await get_models_from_access_groups( - access_group_ids=team_access_group_ids, - prisma_client=prisma_client, - ) - if models_from_groups: - return can_object_call_model( - model=model, - llm_router=llm_router, - models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])), - team_model_aliases=team_model_aliases, - team_id=team_object.team_id if team_object else None, - key_model_aliases=key_model_aliases, - object_type="team", - ) - raise + return await model_access_adapter.can_team_access_model( + model=model, + team_object=team_object, + llm_router=llm_router, + team_model_aliases=team_model_aliases, + key_model_aliases=key_model_aliases, + prisma_client=prisma_client, + ) async def get_authorized_resources_from_key_access_groups( @@ -5795,54 +4660,15 @@ async def get_authorized_resources_from_key_access_groups( check_db_only: bool = False, strict_errors: bool = False, ) -> list[str]: - """ - For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row - and contribute its `resource_field` only if the group authorizes the caller - as an owner — that is, the group's `assigned_team_ids` includes the key's - `team_id`, or the group's `assigned_key_ids` includes the key's token. This - preserves the team-as-owner boundary while still letting a group reach the - key without first being added to the team's `access_group_ids` list. - """ - if valid_token is None: - return [] - key_access_group_ids: Final = list(valid_token.access_group_ids or []) - if not key_access_group_ids: - return [] + from litellm.proxy.auth import model_group_adapter - from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache - - if _prisma_client is None or _user_api_key_cache is None: - return [] - - key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None) - key_token: Final = valid_token.token - - authorized_resources: Final[list[str]] = [] - for ag_id in key_access_group_ids: - try: - ag = await get_access_object( - access_group_id=ag_id, - prisma_client=_prisma_client, - user_api_key_cache=_user_api_key_cache, - proxy_logging_obj=_proxy_logging_obj, - check_db_only=check_db_only, - ) - except AccessGroupNotFoundError: - if check_db_only or strict_errors: - raise - continue - except Exception: - if check_db_only or strict_errors: - raise - continue - team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) - key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) - if team_authorized or key_authorized: - authorized_resources.extend(getattr(ag, resource_field, []) or []) - - return list(set(authorized_resources)) + return await model_group_adapter.get_authorized_resources_from_key_access_groups( + valid_token=valid_token, + team_object=team_object, + resource_field=resource_field, + check_db_only=check_db_only, + strict_errors=strict_errors, + ) async def key_access_group_grants_model( @@ -5888,18 +4714,10 @@ def can_project_access_model( llm_router: RegistryReader | None, key_model_aliases: Mapping[str, str] | None = None, ) -> Literal[True]: - """ - Returns True if the project can access a specific model. + from litellm.proxy.auth import model_access_adapter - Raises ProxyException if access is denied. - """ - return can_object_call_model( - model=model, - llm_router=llm_router, - models=project_object.models if project_object else [], - team_id=project_object.team_id if project_object else None, - key_model_aliases=key_model_aliases, - object_type="project", + return model_access_adapter.can_project_access_model( + model=model, project_object=project_object, llm_router=llm_router, key_model_aliases=key_model_aliases ) @@ -5909,27 +4727,11 @@ def can_customer_access_model( llm_router: RegistryReader | None, valid_token: UserAPIKeyAuth | None, ) -> Literal[True]: - team_model_aliases: Final = team_model_aliases_for_auth_check(valid_token) if valid_token is not None else None - team_id: Final = valid_token.team_id if valid_token is not None else None - key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) + from litellm.proxy.auth import model_access_adapter - def check(name: str) -> None: - team_target: Final = ( - _live_team_alias_target(name, team_model_aliases, team_id, llm_router) if team_model_aliases else name - ) - if team_target != name and name in (end_user_object.models or ()): - return - can_object_call_model( - model=team_target, - llm_router=llm_router, - models=end_user_object.models, - key_model_aliases=key_model_aliases, - object_type="customer", - ) - - for name in (model,) if isinstance(model, str) else model: - check(name) - return True + return model_access_adapter.can_customer_access_model( + model=model, end_user_object=end_user_object, llm_router=llm_router, valid_token=valid_token + ) async def can_user_call_model( @@ -5938,235 +4740,26 @@ async def can_user_call_model( user_object: LiteLLM_UserTable | None, key_model_aliases: Mapping[str, str] | None = None, ) -> Literal[True]: - if user_object is None: - return True + from litellm.proxy.auth import model_access_adapter - if SpecialModelNames.no_default_models.value in user_object.models: - internal_message: Final = ( - f"User not allowed to access model. No default model access, only team models allowed. " - f"Tried to access {model}" - ) - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=internal_message, - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - - return can_object_call_model( - model=model, - llm_router=llm_router, - models=user_object.models, - key_model_aliases=key_model_aliases, - object_type="user", + return await model_access_adapter.can_user_call_model( + model=model, llm_router=llm_router, user_object=user_object, key_model_aliases=key_model_aliases ) def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -> bool: - return ( - bool(litellm.personal_key_model_access_from_teams) - and valid_token is not None - and valid_token.team_id is None - and valid_token.user_id is not None - and valid_token.api_key != LITELLM_PROXY_MASTER_KEY_ALIAS - and valid_token.user_role != LitellmUserRoles.PROXY_ADMIN - ) + from litellm.proxy.auth.personal_model_adapter import personal_key_team_model_access_applies as applies + + return applies(valid_token) def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool: - """MCP server, MCP tool and agent inheritance is for a virtual key only. A keyless auth with no - team (one grant source of an admitted subject, a JWT user) resolves as the plain caller it is.""" - return ( - personal_key_team_model_access_applies(valid_token) - and valid_token is not None - and valid_token.api_key is not None - ) + from litellm.proxy.auth.personal_model_adapter import personal_key_inherits_team_grants as inherits + + return inherits(valid_token) -@dataclass(frozen=True, slots=True) -class _PersonalKeyTeamLoad: - team: LiteLLM_TeamTableCachedObj | None - unreadable: bool - withheld_team: LiteLLM_TeamTableCachedObj | None - - -async def _load_personal_key_team_with_status( - team_id: str, - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> _PersonalKeyTeamLoad: - try: - team_object: Final = await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except Exception as e: # noqa: BLE001 # fail closed: a team that cannot be loaded grants nothing - verbose_proxy_logger.warning( - "Personal key team model access: team=%s lookup failed, granting nothing: %s", team_id, e - ) - return _PersonalKeyTeamLoad(team=None, unreadable=True, withheld_team=None) - if team_object.blocked is True: - return _PersonalKeyTeamLoad(team=None, unreadable=False, withheld_team=team_object) - return _PersonalKeyTeamLoad(team=team_object, unreadable=False, withheld_team=None) - - -async def _load_personal_key_team( - team_id: str, - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> LiteLLM_TeamTableCachedObj | None: - loaded: Final = await _load_personal_key_team_with_status( - team_id=team_id, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - return loaded.team - - -async def _load_personal_key_teams( - user_object: LiteLLM_UserTable, - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - *, - check_db_only: bool = False, -) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: - return tuple( - await asyncio.gather( - *( - _load_personal_key_team( - team_id=team_id, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - for team_id in dict.fromkeys(user_object.teams) - ) - ) - ) - - -async def _team_grants_personal_key_model( - model: str, - team_object: LiteLLM_TeamTableCachedObj, - valid_token: UserAPIKeyAuth, - llm_router: RegistryReader | None, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> bool: - key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) - try: - await can_team_access_model( - model=model, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases, - prisma_client=prisma_client, - ) - await _check_team_member_model_access( - model=model, - team_object=team_object, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - key_model_aliases=key_model_aliases, - ) - except ProxyException: - return False - except Exception as e: # noqa: BLE001 # fail closed: a team check that errors grants nothing - verbose_proxy_logger.warning( - "Personal key team model access: team=%s check failed, granting nothing: %s", team_object.team_id, e - ) - return False - return True - - -async def _loaded_teams_allow_model( - model: str, - teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], - valid_token: UserAPIKeyAuth, - llm_router: RegistryReader | None, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> bool: - loaded: Final = tuple(team for team in teams if team is not None) - if not loaded: - return False - grants: Final = await asyncio.gather( - *( - _team_grants_personal_key_model( - model=model, - team_object=team, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - for team in loaded - ) - ) - return any(grants) - - -async def _personal_key_listing_teams( - user_api_key_dict: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache | None, - proxy_logging_obj: ProxyLogging | None, - *, - check_db_only: bool = False, -) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: - if ( - user_api_key_dict.user_id is None - or prisma_client is None - or user_api_key_cache is None - or proxy_logging_obj is None - ): - return () - owner: Final = await _personal_key_owner( - valid_token=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - if owner is None: - return () - return await _load_personal_key_teams( - user_object=owner, - valid_token=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - - -async def _personal_key_team_within_budget( +async def personal_key_team_within_budget( team: LiteLLM_TeamTableCachedObj, valid_token: UserAPIKeyAuth, prisma_client: PrismaClient, @@ -6175,60 +4768,14 @@ async def _personal_key_team_within_budget( *, check_db_only: bool = False, ) -> bool: - source: Final = valid_token.model_copy( - update={"team_id": team.team_id, "org_id": valid_token.org_id or team.organization_id} + from litellm.proxy.auth.personal_budget_adapter import personal_key_team_within_budget + + return await personal_key_team_within_budget( + team, valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only ) - try: - await team_max_budget_check( - team_object=team, - valid_token=source, - proxy_logging_obj=proxy_logging_obj, - ) - await _team_multi_budget_check(team_object=team) - if check_db_only: - team_member_user_id: Final = valid_token.user_id - if team_member_user_id is None: - return False - team_membership: Final = await get_team_membership( - user_id=team_member_user_id, - team_id=team.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=True, - ) - await _check_team_member_budget( - team_object=team, - user_object=None, - valid_token=source, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=True, - ) - else: - await _check_team_member_budget( - team_object=team, - user_object=None, - valid_token=source, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - await organization_max_budget_check( - valid_token=source, - team_object=team, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - except litellm.BudgetExceededError: - return False - except Exception as e: # noqa: BLE001 # fail closed: a team whose budget cannot be checked grants nothing - verbose_proxy_logger.warning("Failed to check budget for personal key team %s, skipping: %s", team.team_id, e) - return False - return True + + +_personal_key_team_within_budget: Final = personal_key_team_within_budget @dataclass(frozen=True, slots=True) @@ -6246,71 +4793,10 @@ async def personal_key_owner_teams_with_status( *, check_db_only: bool = False, ) -> PersonalKeyOwnerTeams: - """The owner's current, loadable, unblocked, within-budget teams: the grant sources a personal key inherits - MCP server, MCP tool and agent access from under ``personal_key_model_access_from_teams``.""" - if valid_token.user_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: - return PersonalKeyOwnerTeams(teams=(), unreadable=False, withheld=()) - owner: Final = await _personal_key_owner( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - if owner is None: - return PersonalKeyOwnerTeams(teams=(), unreadable=False, withheld=()) - team_ids: Final = tuple(team_id for team_id in dict.fromkeys(owner.teams) if team_id and team_id != UI_TEAM_ID) - team_loads: Final = await asyncio.gather( - *( - _load_personal_key_team_with_status( - team_id=team_id, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - for team_id in team_ids - ) - ) - blocked_teams: Final = tuple( - load.withheld_team - for load in team_loads - if load.withheld_team is not None - and (valid_token.org_id is None or load.withheld_team.organization_id in (None, valid_token.org_id)) - ) - loaded_teams: Final = tuple(load.team for load in team_loads if load.team is not None) - eligible: Final = tuple( - team - for team in loaded_teams - if valid_token.org_id is None or team.organization_id in (None, valid_token.org_id) - ) - if not eligible: - return PersonalKeyOwnerTeams( - teams=(), - unreadable=any(load.unreadable for load in team_loads), - withheld=blocked_teams, - ) - within_budget: Final = await asyncio.gather( - *( - _personal_key_team_within_budget( - team=team, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - for team in eligible - ) - ) - return PersonalKeyOwnerTeams( - teams=tuple(team for team, allowed in zip(eligible, within_budget, strict=True) if allowed), - unreadable=any(load.unreadable for load in team_loads), - withheld=( - *blocked_teams, - *(team for team, allowed in zip(eligible, within_budget, strict=True) if not allowed), - ), + from litellm.proxy.auth import personal_owner_adapter + + return await personal_owner_adapter.personal_key_owner_teams_with_status( + valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only ) @@ -6341,31 +4827,10 @@ async def personal_key_team_visible_models( user_api_key_cache: UserApiKeyCache | None, proxy_logging_obj: ProxyLogging | None, ) -> tuple[str, ...]: - teams: Final = await _personal_key_listing_teams( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - return tuple( - [ - model - for model in (models if teams else ()) - if ( - prisma_client is not None - and user_api_key_cache is not None - and proxy_logging_obj is not None - and await _loaded_teams_allow_model( - model=model, - teams=teams, - valid_token=user_api_key_dict, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - ] + from litellm.proxy.auth import personal_model_listing_adapter + + return await personal_model_listing_adapter.personal_key_team_visible_models( + models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj ) @@ -6378,40 +4843,11 @@ async def can_personal_key_call_model_via_teams( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> Literal[True]: - requested_models: Final = (model,) if isinstance(model, str) else tuple(model) - teams: Final = ( - await _load_personal_key_teams( - user_object=user_object, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if user_object is not None - else () + from litellm.proxy.auth import personal_model_request_adapter + + return await personal_model_request_adapter.can_personal_key_call_model_via_teams( + model, user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj ) - for requested_model in requested_models: - if user_object is None or not await _loaded_teams_allow_model( - model=requested_model, - teams=teams, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ): - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=requested_model), - internal_message=( - f"Personal key model access is derived from the owner's teams. User={valid_token.user_id}, " - f"teams={user_object.teams if user_object is not None else None}, " - f"model={requested_model}" - ), - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - return True def _search_tool_names_from_object_permission( @@ -6448,11 +4884,7 @@ def _can_object_call_search_tool( ProxyException if access is denied """ # Empty list means all search tools are allowed - if not allowed_search_tools: - return True - - # Check if the search tool is in the allowlist - if search_tool_name in allowed_search_tools: + if access_policy.decide(access_policy.Search(search_tool_name, tuple(allowed_search_tools), False)).denial is None: return True # Access denied @@ -6531,6 +4963,8 @@ def _search_tool_denied(object_type: GrantLayer, message: str) -> ProxyException async def _strict_team_object(load_team_object: TeamObjectLoader) -> LiteLLM_TeamTable | None: try: return await load_team_object() + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # an unresolved team grants nothing under deny-by-default verbose_proxy_logger.debug("Team lookup failed under search_tool_deny_by_default: %s", e) return None @@ -6549,6 +4983,8 @@ async def _strict_user_object( user_id_upsert=False, parent_otel_span=valid_token.parent_otel_span, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # an unresolved user grants nothing under deny-by-default verbose_proxy_logger.debug("User lookup failed under search_tool_deny_by_default: %s", e) return None @@ -6599,7 +5035,14 @@ def check_search_tool_grants(search_tool_name: str, grants: SearchToolGrants) -> for layer in ("key", "team", "user"): grant = grants.grants.get(layer) if layer in grants.strict_layers: - if grant is None or search_tool_name not in (grant.search_tools or ()): + if ( + access_policy.decide( + access_policy.Search( + search_tool_name, tuple(grant.search_tools or ()) if grant is not None else (), True + ) + ).denial + is not None + ): raise _search_tool_denied( layer, f"{layer.capitalize()} not allowed to access search tool: {search_tool_name}. " @@ -6702,24 +5145,6 @@ async def is_valid_fallback_model( _MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?") -def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: - """ - Throttle an over-budget key instead of blocking it, when the key opted in - via `throttle_on_budget_exceeded` and a global percentage is configured. - - Records the percentage on the request-scoped `budget_throttle_pct` so the - rate limiter scales the key's TPM/RPM down to it; the persistent limits are - left untouched so the throttle never compounds across requests. Returns True - when the key was throttled (caller skips raising), False when it should still - be hard-blocked. - """ - pct: Final = budget_throttle_percentage() - if pct is None or not should_throttle_budget_exceeded(valid_token): - return False - valid_token.budget_throttle_pct = pct - return True - - def send_token_budget_alert( valid_token: UserAPIKeyAuth, spend: float, @@ -6772,30 +5197,9 @@ async def virtual_key_max_budget_check( proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, ) -> None: - """ - Raises: - BudgetExceededError if the token is over it's max budget. - Triggers a budget alert if the token is over it's max budget. + from litellm.proxy.auth import common_budget_adapter - """ - if valid_token.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - spend: Final = await get_current_spend( - counter_key=key_counter(str(valid_token.token)), - fallback_spend=valid_token.spend or 0.0, - max_budget=valid_token.max_budget, - ) - send_token_budget_alert(valid_token, spend, proxy_logging_obj, user_obj) - - # Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False, - # so a NaN max_budget would silently disable enforcement. Treat a - # non-finite max_budget as "no configured limit" rather than as a bypass. - if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: - if _apply_budget_exceeded_throttle(valid_token): - return - raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget) + return await common_budget_adapter.virtual_key_max_budget_check(valid_token, proxy_logging_obj, user_obj) _virtual_key_max_budget_check: Final = virtual_key_max_budget_check @@ -6804,45 +5208,9 @@ _virtual_key_max_budget_check: Final = virtual_key_max_budget_check async def _virtual_key_multi_budget_check( valid_token: UserAPIKeyAuth, ): - """ - Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded. + from litellm.proxy.auth import budget_windows_adapter - Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}. - Using budget_duration (not list index) keeps counters stable when windows are reordered - or removed during a key update. - - Note: counters are not seeded from DB on Redis cold-start. After a Redis flush, - per-window spend resets to zero within the current window period. This is an acceptable - trade-off: the DB stores reset_at timestamps but not per-window accumulated spend. - """ - if not valid_token.budget_limits: - return - - from litellm.proxy.proxy_server import get_current_spend - - for window in valid_token.budget_limits: - w: dict = window if isinstance(window, dict) else window.model_dump() - counter_key = key_window_counter(str(valid_token.token), str(w["budget_duration"])) - window_spend = await get_current_spend( - counter_key=counter_key, - fallback_spend=0.0, - max_budget=w["max_budget"], - window_entity_type="Key", - window_entity_id=valid_token.token, - window_duration=str(w["budget_duration"]), - window_start=get_budget_window_start(w), - ) - if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: - raise litellm.BudgetExceededError( - current_cost=window_spend, - max_budget=w["max_budget"], - message=( - f"ExceededBudget: Key over {w['budget_duration']} budget. " - f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" - ), - entity_type=Litellm_EntityType.KEY.value, - entity_id=valid_token.token, - ) + return await budget_windows_adapter.virtual_key_multi_budget_check(valid_token) async def virtual_key_soft_budget_check( @@ -6850,13 +5218,9 @@ async def virtual_key_soft_budget_check( proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, ) -> None: - """ - Triggers a budget alert if the token is over it's soft budget. + from litellm.proxy.auth import common_budget_adapter - """ - - if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: - send_soft_budget_alert(valid_token, proxy_logging_obj, user_obj) + return await common_budget_adapter.virtual_key_soft_budget_check(valid_token, proxy_logging_obj, user_obj) def send_soft_budget_alert( @@ -6894,127 +5258,18 @@ def send_soft_budget_alert( _virtual_key_soft_budget_check: Final = virtual_key_soft_budget_check -def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: - """Parse emails from a list or comma-separated string.""" - if isinstance(raw, list): - return [e.strip() for e in raw if isinstance(e, str) and e.strip()] - elif isinstance(raw, str): - return [e.strip() for e in raw.split(",") if e.strip()] - return [] - - -def _normalize_alert_emails( - cfg: Mapping[str, str | Sequence[object] | None] | None, -) -> dict[str, list[str]]: - """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. - - Values may legitimately arrive as list, comma-separated string, or None - from YAML/metadata; _parse_email_list tolerates all three. - """ - if not cfg: - return {} - return {k: _parse_email_list(v) for k, v in cfg.items()} - - -def _merge_budget_alert_email_configs( - global_cfg: Mapping[str, str | Sequence[object] | None] | None, - per_key_cfg: Mapping[str, str | Sequence[object] | None] | None, -) -> dict[str, list[str]] | None: - """ - Per-threshold additive merge: each threshold's recipient list is the union - of global + per-key entries (deduped, global-first ordering). Missing - thresholds on one side are inherited from the other. - """ - global_cfg_normalized: Final = _normalize_alert_emails(global_cfg) - per_key_cfg_normalized: Final = _normalize_alert_emails(per_key_cfg) - if not global_cfg_normalized and not per_key_cfg_normalized: - return None - thresholds: Final = set(global_cfg_normalized) | set(per_key_cfg_normalized) - return { - t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds - } - - async def virtual_key_max_budget_alert_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, ) -> None: - """ - Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - (default 80%) of its max budget. - This is a warning alert before the token actually exceeds the max budget. + from litellm.proxy.auth import budget_alert_adapter - """ - - if valid_token.max_budget is not None and valid_token.spend is not None and valid_token.spend > 0: - owner_email: Final = user_obj.user_email if user_obj else None - alert_email_config: Final[dict[str, list[str]] | None] = _merge_budget_alert_email_configs( - global_cfg=litellm.default_key_max_budget_alert_emails, - per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), - ) - - if isinstance(alert_email_config, dict) and alert_email_config: - # New path: only create task if spend has crossed the lowest threshold - min_pct: Final = min( - (int(k) for k in alert_email_config if k.isdigit()), - default=None, - ) - if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0): - return - - call_info = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=owner_email, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - max_budget_alert_emails=alert_email_config, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, - ) - ) - else: - # Old path: existing single 80% threshold — completely unchanged - alert_threshold: Final = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - - if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget: - verbose_proxy_logger.debug( - "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", - valid_token.token, - valid_token.spend, - valid_token.max_budget, - alert_threshold, - ) - call_info = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=owner_email, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) - - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, - ) - ) + return await budget_alert_adapter.virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, + ) _virtual_key_max_budget_alert_check: Final = virtual_key_max_budget_alert_check @@ -7024,26 +5279,6 @@ TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_e _TEAM_MEMBER_ALERT_CONFIG_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -def _is_valid_alert_threshold_pct(pct: str) -> bool: - return pct.isdigit() and len(pct) <= 3 and 1 <= int(pct) <= 100 - - -def _alert_recipients(raw: object) -> Sequence[str] | None: - if isinstance(raw, (str, Sequence)): - return _parse_email_list(raw) - return None - - -def _valid_alert_threshold_config(raw_config: object) -> Mapping[str, str | Sequence[object] | None] | None: - try: - config: Final = _TEAM_MEMBER_ALERT_CONFIG_ADAPTER.validate_python(raw_config) - except ValidationError: - return None - return MappingProxyType( - {pct: _alert_recipients(emails) for pct, emails in config.items() if _is_valid_alert_threshold_pct(pct)} - ) - - def team_member_max_budget_alert_check( team_id: str, team_alias: str | None, @@ -7055,27 +5290,19 @@ def team_member_max_budget_alert_check( spend: float, max_budget: float, ) -> None: - raw_config: Final = team_metadata.get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY) if team_metadata else None - alert_email_config: Final = _merge_budget_alert_email_configs( - global_cfg=None, per_key_cfg=_valid_alert_threshold_config(raw_config) - ) - if not alert_email_config or spend <= 0: - return - min_pct: Final = min(int(pct) for pct in alert_email_config) - if spend < max_budget * (min_pct / 100.0): - return - call_info: Final = CallInfo( - spend=spend, - max_budget=max_budget, - user_id=user_id, + from litellm.proxy.auth import budget_alert_adapter + + return budget_alert_adapter.team_member_max_budget_alert_check( team_id=team_id, team_alias=team_alias, + team_metadata=team_metadata, organization_id=organization_id, + user_id=user_id, user_email=user_email, - event_group=Litellm_EntityType.TEAM_MEMBER, - max_budget_alert_emails=alert_email_config, + proxy_logging_obj=proxy_logging_obj, + spend=spend, + max_budget=max_budget, ) - asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info)) _team_member_max_budget_alert_check: Final = team_member_max_budget_alert_check @@ -7091,79 +5318,18 @@ async def _check_team_member_budget( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ): - """Check if team member is over their max budget within the team.""" - if ( - team_object is not None - and team_object.team_id is not None - and valid_token is not None - and valid_token.user_id is not None - ): - if not team_membership_loaded: - team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - loaded_membership = team_membership + from litellm.proxy.auth import common_budget_lookups - # Per-member override wins; otherwise fall back to the team-level - # default configured via team.metadata["team_member_budget_id"]. - member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None - now: Final = get_utc_datetime() - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - team_default_budget: Final = ( - await get_team_member_default_budget( - budget_id=default_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - if (member_budget_row is None or member_budget_row.max_budget is None) - and isinstance(default_budget_id, str) - else None - ) - team_member_budget: Final = resolve_effective_team_member_budget( - member_budget=member_budget_row, - team_default_budget=team_default_budget, - now=now, - ) - - if team_member_budget is not None: - team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 - - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = await get_current_spend( - counter_key=team_member_counter(valid_token.user_id, team_object.team_id), - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - - if not math.isfinite(team_member_budget): - return - - team_member_max_budget_alert_check( - team_id=team_object.team_id, - team_alias=team_object.team_alias, - team_metadata=team_object.metadata, - organization_id=team_object.organization_id, - user_id=valid_token.user_id, - user_email=user_object.user_email if user_object is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) - - if team_member_spend >= team_member_budget: - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{team_object.team_id}", - ) + return await common_budget_lookups.check_team_member_budget( + team_object, + user_object, + valid_token, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + ) async def check_team_member_model_access( @@ -7178,54 +5344,20 @@ async def check_team_member_model_access( team_membership_loaded: bool = False, key_model_aliases: Mapping[str, str] | None = None, ) -> None: - """ - Check if a team member's per-member model scope allows access to the requested model. + from litellm.proxy.auth import member_model_adapter - Only enforced when the member's budget table has a non-empty allowed_models list. - If allowed_models is empty or absent, the team-level models list applies (no extra restriction). - """ - if valid_token.user_id is None or team_object.team_id is None: - return - - if not team_membership_loaded: - team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - loaded_membership = team_membership - - if ( - loaded_membership is None - or loaded_membership.litellm_budget_table is None - or not loaded_membership.litellm_budget_table.allowed_models - ): - return # no per-member restriction — inherit team-level check - - member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models - try: - can_object_call_model( - model=model, - llm_router=llm_router, - models=member_allowed_models, - object_type="team", - team_id=team_object.team_id, - key_model_aliases=key_model_aliases, - ) - except ProxyException: - internal_message: Final = ( - f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, " - f"Model={model}. Allowed member models = {member_allowed_models}" - ) - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=internal_message, - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) + return await member_model_adapter.check_team_member_model_access( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + key_model_aliases, + ) _check_team_member_model_access: Final = check_team_member_model_access @@ -7236,49 +5368,9 @@ async def team_max_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ) -> None: - """ - Check if the team is over it's max budget. + from litellm.proxy.auth import common_budget_adapter - Raises: - BudgetExceededError if the team is over it's max budget. - Triggers a budget alert if the team is over it's max budget. - """ - if team_object is not None and team_object.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - spend: Final = await get_current_spend( - counter_key=team_counter(team_object.team_id), - fallback_spend=team_object.spend or 0.0, - max_budget=team_object.max_budget, - ) - - if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: - if valid_token: - call_info: Final = CallInfo( - token=valid_token.token, - spend=spend, - max_budget=team_object.max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.TEAM, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="team_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=spend, - max_budget=team_object.max_budget, - message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", - entity_type=Litellm_EntityType.TEAM.value, - entity_id=team_object.team_id, - ) + return await common_budget_adapter.team_max_budget_check(team_object, valid_token, proxy_logging_obj) _team_max_budget_check: Final = team_max_budget_check @@ -7287,41 +5379,9 @@ _team_max_budget_check: Final = team_max_budget_check async def _team_multi_budget_check( team_object: LiteLLM_TeamTable | None, ): - """ - Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded. + from litellm.proxy.auth import budget_windows_adapter - Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}. - Using budget_duration (not list index) keeps counters stable when windows are reordered - or removed during a team update. - """ - if team_object is None or not team_object.budget_limits: - return - - from litellm.proxy.proxy_server import get_current_spend - - for window in team_object.budget_limits: - w: dict = window if isinstance(window, dict) else window.model_dump() - counter_key = team_window_counter(team_object.team_id, str(w["budget_duration"])) - window_spend = await get_current_spend( - counter_key=counter_key, - fallback_spend=0.0, - max_budget=w["max_budget"], - window_entity_type="Team", - window_entity_id=team_object.team_id, - window_duration=str(w["budget_duration"]), - window_start=get_budget_window_start(w), - ) - if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: - raise litellm.BudgetExceededError( - current_cost=window_spend, - max_budget=w["max_budget"], - message=( - f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " - f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" - ), - entity_type=Litellm_EntityType.TEAM.value, - entity_id=team_object.team_id, - ) + return await budget_windows_adapter.team_multi_budget_check(team_object) async def _team_soft_budget_check( @@ -7329,70 +5389,9 @@ async def _team_soft_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ): - """ - Triggers a budget alert if the team is over it's soft budget. - """ - if ( - team_object is not None - and team_object.soft_budget is not None - and team_object.spend is not None - and team_object.spend >= team_object.soft_budget - ): - verbose_proxy_logger.debug( - "Crossed Soft Budget for team %s, spend %s, soft_budget %s", - team_object.team_id, - team_object.spend, - team_object.soft_budget, - ) - if valid_token: - # Extract alert emails from team metadata - alert_emails: list[str] | None = None - if team_object.metadata is not None and isinstance(team_object.metadata, dict): - soft_budget_alert_emails: Final = team_object.metadata.get("soft_budget_alerting_emails") - if soft_budget_alert_emails is not None: - if isinstance(soft_budget_alert_emails, list): - alert_emails = [ - email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() - ] - elif isinstance(soft_budget_alert_emails, str): - # Handle comma-separated string - alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] - # Filter out empty strings - if alert_emails: - alert_emails = [email for email in alert_emails if email] - else: - alert_emails = None + from litellm.proxy.auth import common_budget_adapter - # Only send team soft budget alerts if alert_emails are configured - # Team soft budget alerts are sent via metadata.soft_budget_alerting_emails, not global alerting - if alert_emails is None or len(alert_emails) == 0: - verbose_proxy_logger.debug( - "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails", - team_object.team_id, - ) - return - - call_info: Final = CallInfo( - token=valid_token.token, - spend=team_object.spend, - max_budget=team_object.max_budget, - soft_budget=team_object.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=None, # Team-level alert, no specific user email - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.TEAM, - alert_emails=alert_emails, - ) - - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) - ) + return await common_budget_adapter.team_soft_budget_check(team_object, valid_token, proxy_logging_obj) async def _project_max_budget_check( @@ -7400,57 +5399,9 @@ async def _project_max_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ): - """ - Check if the project is over its max budget. + from litellm.proxy.auth import common_budget_adapter - Raises: - BudgetExceededError if the project is over its max budget. - Triggers a budget alert if the project is over its max budget. - """ - if project_object is None: - return - - max_budget = None - if project_object.litellm_budget_table is not None: - max_budget = project_object.litellm_budget_table.max_budget - - if max_budget is None or not math.isfinite(max_budget): - return - - from litellm.proxy.proxy_server import get_current_spend - - project_spend: Final = await get_current_spend( - counter_key=project_counter(project_object.project_id), - fallback_spend=project_object.spend or 0.0, - max_budget=max_budget, - ) - - if project_spend >= max_budget: - if valid_token: - call_info: Final = CallInfo( - token=valid_token.token, - spend=project_spend, - max_budget=max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.PROJECT, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="project_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=project_spend, - max_budget=max_budget, - message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}", - entity_type=Litellm_EntityType.PROJECT.value, - entity_id=project_object.project_id, - ) + return await common_budget_adapter.project_max_budget_check(project_object, valid_token, proxy_logging_obj) async def _project_soft_budget_check( @@ -7458,43 +5409,9 @@ async def _project_soft_budget_check( valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, ): - """ - Triggers a budget alert if the project is over its soft budget. + from litellm.proxy.auth import common_budget_adapter - Mirrors _team_soft_budget_check() pattern. - """ - if project_object is None: - return - - soft_budget = None - if project_object.litellm_budget_table is not None: - soft_budget = project_object.litellm_budget_table.soft_budget - - if soft_budget is not None and project_object.spend is not None and project_object.spend >= soft_budget: - verbose_proxy_logger.debug( - "Crossed Soft Budget for project %s, spend %s, soft_budget %s", - project_object.project_id, - project_object.spend, - soft_budget, - ) - if valid_token: - call_info: Final = CallInfo( - token=valid_token.token, - spend=project_object.spend, - max_budget=None, - soft_budget=soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.PROJECT, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) - ) + return await common_budget_adapter.project_soft_budget_check(project_object, valid_token, proxy_logging_obj) @with_service_target(AUTH_OBJECTS_TARGET) @@ -7569,97 +5486,12 @@ async def organization_max_budget_check( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: - """ - Check if the organization is over its max budget. + from litellm.proxy.auth import common_budget_lookups - This function checks the organization budget using: - 1. First, tries to use valid_token.org_id (if key has organization_id set) - 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) - - This ensures organization budget checks work even when keys don't have organization_id - set directly, as long as their team belongs to an organization. - - Raises: - BudgetExceededError if the organization is over its max budget. - Triggers a budget alert if the organization is over its max budget. - """ - if valid_token is None or prisma_client is None: - return - - # Determine organization_id: first try from token, then fallback to team - org_id: str | None = None - if valid_token.org_id is not None: - org_id = valid_token.org_id - elif team_object is not None and team_object.organization_id is not None: - org_id = team_object.organization_id - - # If no organization_id found, skip the check - if org_id is None: - return - - # Get organization object with budget table - use get_org_object so it can be mocked in tests - try: - org_table: Final = await get_org_object( - org_id=org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - include_budget_table=True, - ) - except Exception: - # If organization lookup fails, skip the check - return - - if org_table is None: - return - - carry_organization_budget_state(valid_token=valid_token, org_table=org_table) - - # Get max_budget from organization's budget table - org_max_budget: float | None = None - if org_table.litellm_budget_table is not None: - org_max_budget = org_table.litellm_budget_table.max_budget - - if org_max_budget is None: - return - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - from litellm.proxy.proxy_server import get_current_spend - - org_spend: Final = await get_current_spend( - counter_key=org_counter(org_id), - fallback_spend=org_table.spend or 0.0, - max_budget=org_max_budget, + return await common_budget_lookups.organization_max_budget_check( + valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj ) - # Check if organization spend exceeds max budget - if math.isfinite(org_max_budget) and org_spend >= org_max_budget: - # Trigger budget alert - call_info: Final = CallInfo( - token=valid_token.token, - spend=org_spend, - max_budget=org_max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=org_id, - event_group=Litellm_EntityType.ORGANIZATION, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="organization_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=org_spend, - max_budget=org_max_budget, - message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", - entity_type=Litellm_EntityType.ORGANIZATION.value, - entity_id=org_id, - ) - _organization_max_budget_check: Final = organization_max_budget_check @@ -7696,97 +5528,24 @@ async def tag_max_budget_check_for_tags( proxy_logging_obj: ProxyLogging, valid_token: UserAPIKeyAuth | None, ) -> None: - if prisma_client is None or not tags: - return + from litellm.proxy.auth import common_budget_lookups - tag_objects: Final = await get_tag_objects_batch( - tag_names=tags, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + return await common_budget_lookups.tag_max_budget_check_for_tags( + tags, prisma_client, user_api_key_cache, proxy_logging_obj, valid_token ) - # Check budget for each tag - for tag_name in tags: - tag_object = tag_objects.get(tag_name) - if tag_object is None: - continue - - # Check if tag has budget limits - if tag_object.litellm_budget_table is not None and tag_object.litellm_budget_table.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - tag_spend = await get_current_spend( - counter_key=tag_counter(tag_name), - fallback_spend=tag_object.spend or 0.0, - max_budget=tag_object.litellm_budget_table.max_budget, - fallback_authoritative=True, - ) - if tag_spend <= tag_object.litellm_budget_table.max_budget: - continue - raise litellm.BudgetExceededError( - current_cost=tag_spend, - max_budget=tag_object.litellm_budget_table.max_budget, - message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", - entity_type=Litellm_EntityType.TAG.value, - entity_id=tag_name, - ) - async def _model_access_group_max_budget_check( matched_model_access_groups: Sequence[str], prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, ) -> None: - """Block the request when a model access group that authorized it is over its max budget. + from litellm.proxy.auth import common_budget_lookups - Only the groups auth already matched are charged and therefore only they are checked, so a - request that no budgeted group authorized costs nothing here. - - Like the tag check this is a plain read with no reservation, so concurrent requests can - overshoot the ceiling slightly. The reservation counters are the precise path; this one covers - the ``disable_budget_reservation`` case. - - The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded - spend has reached ``max_budget`` has nothing left to give, so the next request is refused. - Keys and organizations already draw the line there. A non-positive budget means no budget, - matching what the reservation path treats as unbudgeted. - - Raises: - BudgetExceededError if a matched group is over its max budget. - """ - if prisma_client is None or not matched_model_access_groups: - return - - budgets: Final = await get_model_access_group_budgets_batch( - access_group_names=matched_model_access_groups, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + return await common_budget_lookups.model_access_group_max_budget_check( + matched_model_access_groups, prisma_client, user_api_key_cache ) - from litellm.proxy.proxy_server import get_current_spend - - for group in matched_model_access_groups: - budget = budgets.get(group) - if budget is None or budget.max_budget is None or budget.max_budget <= 0: - continue - - group_spend = await get_current_spend( - counter_key=model_access_group_counter(group), - fallback_spend=budget.spend, - max_budget=budget.max_budget, - fallback_authoritative=True, - ) - if group_spend < budget.max_budget: - continue - raise litellm.BudgetExceededError( - current_cost=group_spend, - max_budget=budget.max_budget, - message=f"Budget has been exceeded! Model access group={group} Current cost: {group_spend}, Max budget: {budget.max_budget}", - entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, - entity_id=group, - ) - def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: """ @@ -7854,6 +5613,8 @@ def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_mode """ try: stripped_model, custom_llm_provider, _, _ = get_llm_provider(model=model) + except NativeAuthFailure: + raise except Exception: return False @@ -7903,13 +5664,18 @@ def _strict_grant_layers( The identities that must each grant an object under a deny-by-default policy: a virtual key and its team, a keyless team member's team, or a keyless user's own grant. The master key and dashboard sessions need none """ - if not deny_by_default or valid_token is None or not _is_strict_grant_identity(valid_token): - return frozenset() - virtual_key: Final = valid_token.via_virtual_key and not valid_token.is_session_token - has_team: Final = team_object is not None or valid_token.team_id is not None - if not virtual_key and not has_team: - return frozenset(("user",)) - return frozenset(layer for layer, required in (("key", virtual_key), ("team", has_team)) if required) + decision: Final = access_policy.decide( + access_policy.GrantLayers( + deny_by_default=deny_by_default, + identity_present=valid_token is not None, + master_key=valid_token is not None and valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS, + dashboard=valid_token is not None and valid_token.team_id == UI_TEAM_ID, + virtual_key=valid_token is not None and valid_token.via_virtual_key, + session_token=valid_token is not None and valid_token.is_session_token, + team_present=team_object is not None or (valid_token is not None and valid_token.team_id is not None), + ) + ) + return frozenset(decision.layers) async def _identity_grants( @@ -8074,7 +5840,16 @@ def _require_vector_store_grant( vector_store_ids_to_run: Sequence[str], object_permission: _VectorStorePermissionsRow | None, ) -> None: - if object_permission is None or not object_permission.vector_stores: + decision: Final = access_policy.decide( + access_policy.VectorStores( + requested=tuple(map(_provider_vector_store_id, vector_store_ids_to_run)), + allowed=None + if object_permission is None or object_permission.vector_stores is None + else tuple(map(_provider_vector_store_id, object_permission.vector_stores)), + strict=True, + ) + ) + if decision.denial == "MissingVectorGrant": raise ProxyException( message=f"{object_type.capitalize()} not allowed to access vector store. Tried to access {vector_store_ids_to_run[0]}. vector_store_deny_by_default is enabled and the {object_type} has no vector store grants", type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), @@ -8218,24 +5993,20 @@ def _can_object_call_vector_stores( """ Raises ProxyException if the object (key, team, org) cannot access the specific vector store. """ - if object_permissions is None: - return True - - if object_permissions.vector_stores is None: - return True - - # If length is 0, then the object has access to all vector stores. - if len(object_permissions.vector_stores) == 0: - return True - - allowed_stores: Final = frozenset(map(_provider_vector_store_id, object_permissions.vector_stores)) - for vector_store_id in vector_store_ids_to_run: - if _provider_vector_store_id(vector_store_id) not in allowed_stores: - raise ProxyException( - message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}", - type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), - param="vector_store", - code=status.HTTP_401_UNAUTHORIZED, - ) + grants: Final = object_permissions.vector_stores if object_permissions is not None else None + decision: Final = access_policy.decide( + access_policy.VectorStores( + requested=tuple(map(_provider_vector_store_id, vector_store_ids_to_run)), + allowed=tuple(map(_provider_vector_store_id, grants)) if grants is not None else None, + strict=False, + ) + ) + if decision.rejected_index is not None: + raise ProxyException( + message=f"User not allowed to access vector store. Tried to access {vector_store_ids_to_run[decision.rejected_index]}. Only allowed to access {grants}", + type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), + param="vector_store", + code=status.HTTP_401_UNAUTHORIZED, + ) return True diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index 6820ded0eba..47e2ef6efea 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/litellm/proxy/auth/auth_checks_organization.py @@ -5,8 +5,6 @@ Auth Checks for Organizations from collections.abc import Awaitable, Callable from typing import Final -from fastapi import status - from litellm.proxy._types import * @@ -15,96 +13,11 @@ def organization_role_based_access_check( user_object: LiteLLM_UserTable | None, route: str, ): - """ - Role based access control checks only run if a user is part of an Organization + from litellm.proxy.auth import org_route_adapter - Organization Checks: - ONLY RUN IF user_object.organization_memberships is not None - - 1. Only Proxy Admins can access /organization/new - 2. IF route is a LiteLLMRoutes.org_admin_only_routes, then check if user is an Org Admin for that organization - - """ - - if user_object is None: - return - - passed_organization_id: Final[str | None] = request_body.get("organization_id", None) - - if route == "/organization/new": - if user_object.user_role != LitellmUserRoles.PROXY_ADMIN.value: - raise ProxyException( - message=f"Only proxy admins can create new organizations. You are {user_object.user_role}", - type=ProxyErrorTypes.auth_error.value, - param="user_role", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if user_object.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return - - # Checks if route is an Org Admin Only Route - if route in LiteLLMRoutes.org_admin_only_routes.value: - ( - _user_organizations, - _user_organization_role_mapping, - ) = get_user_organization_info(user_object) - - if user_object.organization_memberships is None: - raise ProxyException( - message=f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if passed_organization_id is None: - raise ProxyException( - message="Passed organization_id is None, please pass an organization_id in your request", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - user_role: Final[LitellmUserRoles | None] = _user_organization_role_mapping.get(passed_organization_id) - if user_role is None: - raise ProxyException( - message=f"You do not have a role within the selected organization. Passed organization_id: {passed_organization_id}. Please contact the organization admin to request access.", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if user_role != LitellmUserRoles.ORG_ADMIN.value: - raise ProxyException( - message=f"You do not have the required role to perform {route} in Organization {passed_organization_id}. Your role is {user_role} in Organization {passed_organization_id}", - type=ProxyErrorTypes.auth_error.value, - param="user_role", - code=status.HTTP_401_UNAUTHORIZED, - ) - elif route == "/team/new": - # if user is part of multiple teams, then they need to specify the organization_id - ( - _user_organizations, - _user_organization_role_mapping, - ) = get_user_organization_info(user_object) - if user_object.organization_memberships is not None and len(user_object.organization_memberships) > 0: - if passed_organization_id is None: - raise ProxyException( - message=f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {_user_organizations}", - type=ProxyErrorTypes.auth_error.value, - param="organization_id", - code=status.HTTP_401_UNAUTHORIZED, - ) - - _user_role_in_passed_org: Final = _user_organization_role_mapping.get(passed_organization_id) - if _user_role_in_passed_org != LitellmUserRoles.ORG_ADMIN.value: - raise ProxyException( - message=f"You do not have the required role to call {route}. Your role is {_user_role_in_passed_org} in Organization {passed_organization_id}", - type=ProxyErrorTypes.auth_error.value, - param="user_role", - code=status.HTTP_401_UNAUTHORIZED, - ) + return org_route_adapter.organization_role_based_access_check( + request_body=request_body, user_object=user_object, route=route + ) def get_user_organization_info( @@ -137,40 +50,9 @@ def user_is_org_admin( request_data: dict, user_object: LiteLLM_UserTable | None = None, ) -> bool: - """ - Helper function to check if user is an org admin for all of the passed organizations. + from litellm.proxy.auth import org_route_adapter - Checks both: - - `organization_id` (singular string) — legacy callers - - `organizations` (list of strings) — used by /user/new - """ - if user_object is None: - return False - - if user_object.organization_memberships is None: - return False - - # Collect candidate org IDs from both fields - candidate_org_ids: Final[list[str]] = [] - singular: Final = request_data.get("organization_id", None) - if singular is not None: - candidate_org_ids.append(singular) - orgs_list: Final = request_data.get("organizations", None) - if isinstance(orgs_list, list): - candidate_org_ids.extend(orgs_list) - - if not candidate_org_ids: - return False - - # Build set of orgs where user is admin - admin_org_ids: Final = { - _membership.organization_id - for _membership in user_object.organization_memberships - if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value and _membership.organization_id is not None - } - - # User must be admin of ALL requested orgs, not just any one - return all(org_id in admin_org_ids for org_id in candidate_org_ids) + return org_route_adapter.user_is_org_admin(request_data=request_data, user_object=user_object) _user_is_org_admin: Final = user_is_org_admin diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index 5b872561e6f..957844269d0 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -12,6 +12,7 @@ from litellm._auth.decision import ( MalformedKey, MasterKeyOnlyRoute, MissingKey, + ModelAccessDenied, ModelBudgetExceeded, NoDatabase, NotAVirtualKey, @@ -19,10 +20,12 @@ from litellm._auth.decision import ( TeamMemberBudgetExceeded, UserAuthDisabled, ) +from litellm._auth.failure import NativeAuthFailure from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_KEY_ERROR_MESSAGE from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import key_budget_exceeded_error from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.proxy.auth.model_access_adapter import raise_model_denial from litellm.types.proxy.auth.user_api_key_auth import is_jwt _JWT_AUTH_DISABLED_HINT: Final = ( @@ -52,6 +55,9 @@ def denial_exception( key: UserAPIKeyAuth | None, ) -> ProxyException | HTTPException | litellm.BudgetExceededError: match denial: + case ModelAccessDenied(decision=model_decision): + raise_model_denial(model_decision, model_decision.model or "") + raise NativeAuthFailure("native model denial contained an allow") case MissingKey(): return ProxyException( message="Authentication Error, No api key passed in.", diff --git a/litellm/proxy/auth/budget_alert_adapter.py b/litellm/proxy/auth/budget_alert_adapter.py new file mode 100644 index 00000000000..2537af0f8b8 --- /dev/null +++ b/litellm/proxy/auth/budget_alert_adapter.py @@ -0,0 +1,178 @@ +import asyncio +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._auth import budget_alerts +from litellm._auth.budget_alerts import ConfigFacts, Decision, Facts, Failure, ThresholdFacts +from litellm._logging import verbose_proxy_logger +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.proxy._types import CallInfo, Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.common_budget_adapter import metadata +from litellm.proxy.utils import ProxyLogging + +_CONFIG: Final = TypeAdapter(Mapping[str, object]) +_OBJECT: Final = TypeAdapter(object) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def _threshold(key: str, raw: object) -> ThresholdFacts: + text: Final = raw if isinstance(raw, str) else None + items: Final = ( + tuple(item if isinstance(item, str) else None for item in _ITEMS.validate_python(raw)) + if isinstance(raw, list) + else None + ) + initial: Final = ThresholdFacts(key=key, digit=key.isdigit(), text=text, items=items) + try: + integer: Final = int(key) + except ValueError as error: + return replace(initial, integer_failure=Failure("ValueError", str(error))) + parsed: Final = replace(initial, digits=str(integer)) + try: + return replace(parsed, ratio=integer / 100.0) + except OverflowError as error: + return replace(parsed, ratio_failure=Failure("OverflowError", str(error))) + + +def _config(raw: object) -> tuple[tuple[ThresholdFacts, ...], Failure | None]: + if not raw: + return (), None + try: + value: Final = _CONFIG.validate_python(raw) + except ValidationError: + return (), Failure("AttributeError", f"'{type(raw).__name__}' object has no attribute 'items'") + return tuple(_threshold(key, value) for key, value in value.items()), None + + +def read_alert_facts(facts: Facts, global_config: object, local_config: object) -> Facts: + decision: Final = budget_alerts.plan(facts) + if not decision.read_config: + return facts + global_entries, global_failure = _config(global_config) + local_entries, local_failure = _config(local_config) + loaded: Final = replace(facts, config=ConfigFacts(global_entries, local_entries, global_failure, local_failure)) + order: Final = budget_alerts.plan(loaded) + if not order.read_order: + return loaded + assert loaded.config is not None + return replace(loaded, config=replace(loaded.config, order=tuple(set(order.global_keys) | set(order.local_keys)))) + + +def read_key_budget_alert_facts(valid_token: UserAPIKeyAuth) -> Facts: + facts: Final = Facts( + False, valid_token.spend, valid_token.max_budget, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + ) + if not budget_alerts.plan(facts).read_config: + return facts + return read_alert_facts( + facts, + _OBJECT.validate_python(litellm.default_key_max_budget_alert_emails), # pyright: ignore[reportUnknownMemberType] # validate the legacy global config snapshot + metadata(valid_token.metadata).get("max_budget_alert_emails"), # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # legacy metadata is validated at the boundary + ) + + +def read_member_budget_alert_facts( + spend: float, + max_budget: float, + team_metadata: Mapping[str, object] | None, +) -> Facts: + facts: Final = Facts(True, spend, max_budget, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) + return read_alert_facts( + facts, None, team_metadata.get("team_member_max_budget_alert_emails") if team_metadata else None + ) + + +def raise_alert_failure(decision: Decision) -> None: + failure: Final = decision.failure + if failure is None: + return + match failure.kind: + case "AttributeError": + raise AttributeError(failure.message) + case "OverflowError": + raise OverflowError(failure.message) + case _: + raise ValueError(failure.message) + + +def schedule_budget_alert(decision: Decision, call_info: CallInfo, proxy_logging_obj: ProxyLogging) -> None: + raise_alert_failure(decision) + if decision.emit: + projected: Final = ( + call_info + if decision.thresholds is None + else call_info.model_copy(update={"max_budget_alert_emails": alert_emails(decision)}) + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=projected)) + + +def alert_emails( + decision: Decision, +) -> dict[str, list[str]] | None: # mutable-ok: CallInfo requires the legacy dict of lists projection + return None if decision.thresholds is None else {entry.key: list(entry.recipients) for entry in decision.thresholds} + + +async def virtual_key_max_budget_alert_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + decision: Final = budget_alerts.plan(read_key_budget_alert_facts(valid_token)) + raise_alert_failure(decision) + if not decision.emit: + return + if decision.thresholds is None: + assert valid_token.max_budget is not None + verbose_proxy_logger.debug( + "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", + valid_token.token, + valid_token.spend, + valid_token.max_budget, + valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, + ) + call_info: Final = CallInfo( + token=valid_token.token, + spend=valid_token.spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + schedule_budget_alert(decision, call_info, proxy_logging_obj) + + +def team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + decision: Final = budget_alerts.plan(read_member_budget_alert_facts(spend, max_budget, team_metadata)) + raise_alert_failure(decision) + if not decision.emit: + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + ) + schedule_budget_alert(decision, call_info, proxy_logging_obj) diff --git a/litellm/proxy/auth/budget_context_adapter.py b/litellm/proxy/auth/budget_context_adapter.py new file mode 100644 index 00000000000..fb9413cbf66 --- /dev/null +++ b/litellm/proxy/auth/budget_context_adapter.py @@ -0,0 +1,385 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import ( + Final, + Literal, + Protocol, + assert_never, + cast, # noqa: TID251 # preserve identity of validated legacy JSON mappings +) + +from opentelemetry.trace import Span + +import litellm +from litellm._auth import budget_context as core +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class TeamModelLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + +class EndUserModelLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + +def mapping_value(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return cast(Mapping[str, object], value) # cast-ok: preserve the original JSON map's identity + return None + + +def mapping_count(value: object) -> int | None: + mapping: Final = mapping_value(value) + return len(mapping) if mapping is not None else None + + +def budget_fields(row: LiteLLM_BudgetTable | None) -> core.BudgetFields | None: + if row is None: + return None + return core.BudgetFields( + row.max_budget, + row.tpm_limit, + row.rpm_limit, + row.tpd_limit, + mapping_count(getattr(row, "model_max_budget", None)), + ) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy budget table declares an unparameterized dict + + +async def check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: TeamModelLimiter, + models: Sequence[str], +) -> None: + budgets: Final = valid_token.team_model_max_budget + selected: Final = core.model_checks( + core.ModelCheckFacts("Team", valid_token.team_id, len(budgets) if budgets is not None else 0, len(models)) + ) + for index in selected: + assert valid_token.team_id is not None and budgets is not None + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=budgets, + key_model_max_budget=mapping_value(getattr(valid_token, "model_max_budget", None)), + model=models[index], + ) + + +async def check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: EndUserModelLimiter, + models: Sequence[str], + team_id: str | None, + *, + scope: Literal["EndUser", "CustomEndUser"] = "EndUser", +) -> None: + selected: Final = core.model_checks( + core.ModelCheckFacts( + scope, + end_user_id, + len(end_user_model_max_budget) if end_user_model_max_budget is not None else 0, + len(models), + ) + ) + for index in selected: + assert end_user_id is not None and end_user_model_max_budget is not None + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=models[index], + team_id=team_id, + ) + + +@dataclass(frozen=True, slots=True) +class EndUserIO: + prisma_client: PrismaClient | None + cache: UserApiKeyCache + span: Span | None + + +async def _end_user_row( + facts: core.EndUserFacts, + io: EndUserIO, + own: LiteLLM_EndUserTable | None, + key: LiteLLM_BudgetTable | None = None, + proxy: LiteLLM_BudgetTable | None = None, +) -> LiteLLM_BudgetTable | None: + from litellm.proxy.auth.auth_checks import get_default_end_user_budget + + decision: Final = core.end_user_source(facts) + match decision: + case "Own": + assert own is not None + return own.litellm_budget_table + case "Key": + return key + case "Global": + return proxy + case "None": + return None + case "ReadKey": + row: Final = await get_default_end_user_budget( + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + budget_id=facts.key_budget_id, + ) + return await _end_user_row( + replace(facts, key_read=True, key_budget=budget_fields(row)), io, own, row, proxy + ) + case "ReadGlobalConfig": + return await _end_user_row( + replace(facts, global_config_read=True, global_budget_id=litellm.max_end_user_budget_id), + io, + own, + key, + proxy, + ) + case "ReadGlobal": + row_global: Final = await get_default_end_user_budget( + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + budget_id=facts.global_budget_id, + ) + return await _end_user_row( + replace(facts, global_read=True, global_budget=budget_fields(row_global)), io, own, key, row_global + ) + case _: + assert_never(decision) + + +async def resolve_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + key_end_user_budget_id: str | None, + parent_otel_span: Span | None = None, +) -> LiteLLM_BudgetTable | None: + return await _end_user_row( + core.EndUserFacts("Default", False, None, prisma_client is not None, key_end_user_budget_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + None, + ) + + +async def end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + row: Final = await _end_user_row( + core.EndUserFacts( + "Models", end_user_object is not None, end_user_id, prisma_client is not None, key_end_user_budget_id + ), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + end_user_object, + ) + return getattr(row, "model_max_budget", None) if row is not None else None # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy row field + + +async def apply_key_end_user_default_budget_to_token( + valid_token: UserAPIKeyAuth, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + keep_token_limits: bool, +) -> None: + row: Final = await _end_user_row( + core.EndUserFacts("Token", end_user_object is not None, valid_token.end_user_id, True, key_end_user_budget_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + end_user_object, + ) + current: Final = core.BudgetFields( + valid_token.end_user_max_budget, + valid_token.end_user_tpm_limit, + valid_token.end_user_rpm_limit, + valid_token.end_user_tpd_limit, + mapping_count(valid_token.end_user_model_max_budget), + ) + for field in core.token_assignments(core.TokenDefaultsFacts(current, keep_token_limits, row is not None)): + assert row is not None + match field: + case "MaxBudget": + valid_token.end_user_max_budget = ( # rebind-ok: request-local native projection + row.max_budget + ) + case "Tpm": + valid_token.end_user_tpm_limit = ( # rebind-ok: request-local native projection + row.tpm_limit + ) + case "Rpm": + valid_token.end_user_rpm_limit = ( # rebind-ok: request-local native projection + row.rpm_limit + ) + case "Tpd": + valid_token.end_user_tpd_limit = ( # rebind-ok: request-local native projection + row.tpd_limit + ) + case "Models": + valid_token.end_user_model_max_budget = row.model_max_budget # pyright: ignore[reportUnknownMemberType] # preserve the validated row dict and its identity # rebind-ok: native-selected request-local budget assignment + case _: + assert_never(field) + + +@dataclass(frozen=True, slots=True) +class MemberIO: + prisma_client: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging | None + + +@dataclass(frozen=True, slots=True) +class MemberRead: + member: LiteLLM_TeamMembership | None = None + default: LiteLLM_BudgetTable | None = None + error: Exception | None = None + + +async def _read_member(facts: core.MemberFacts, io: MemberIO) -> MemberRead: + from litellm.proxy.auth.auth_checks import get_team_membership + + assert facts.user_id is not None and facts.team_id is not None + try: + row: Final = await get_team_membership( + user_id=facts.user_id, + team_id=facts.team_id, + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # only the database/cache read is caught; native failures rethrow + return MemberRead(error=error) + return MemberRead(member=row) + + +async def _read_member_default(facts: core.MemberFacts, io: MemberIO) -> MemberRead: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget + + assert facts.default_id is not None + try: + row: Final = await get_team_member_default_budget( + budget_id=facts.default_id, + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + raise_on_lookup_error=True, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # only the database/cache read is caught; native failures rethrow + return MemberRead(error=error) + return MemberRead(default=row) + + +async def _member_model_budget( + token: UserAPIKeyAuth, + facts: core.MemberFacts, + io: MemberIO, + reading: MemberRead, +) -> None: + decision: Final = core.member_source(facts) + match decision: + case "Done": + return + case "RaiseLookup": + assert reading.error is not None + raise reading.error + case "SkipLookup": + verbose_proxy_logger.warning( + "Skipping team-member model budget lookup after a lookup error", exc_info=reading.error + ) + return + case "ReadMembership": + member_read: Final = await _read_member(facts, io) + member: Final = member_read.member + member_row: Final = member.litellm_budget_table if member is not None else None + snapshot: Final = ( + core.MembershipSnapshot(member.budget_id, mapping_count(getattr(member_row, "model_max_budget", None))) + if member is not None + else None + ) + await _member_model_budget( + token, + replace( + facts, + membership_read=True, + membership=snapshot, + lookup_failure="Membership" if member_read.error is not None else None, + ), + io, + member_read, + ) + case "ReadDefault": + default_read: Final = await _read_member_default(facts, io) + await _member_model_budget( + token, + replace( + facts, + default_read=True, + default_budget=budget_fields(default_read.default), + lookup_failure="Default" if default_read.error is not None else None, + ), + io, + default_read, + ) + case "Member": + assert reading.member is not None + token.team_member_model_max_budget = mapping_value( # rebind-ok: request-local native projection + getattr(reading.member.litellm_budget_table, "model_max_budget", None) + ) + case "Default": + assert reading.default is not None + token.team_member_model_max_budget = mapping_value( # rebind-ok: request-local native projection + getattr(reading.default, "model_max_budget", None) + ) + case "Clear": + token.team_member_model_max_budget = None # rebind-ok: request-local native projection + case _: + assert_never(decision) + + +async def load_team_member_default_model_budget( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool, + proxy_logging_obj: ProxyLogging | None, +) -> None: + metadata: Final = mapping_value(getattr(team_object, "metadata", None)) + default_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None + facts: Final = core.MemberFacts( + user_id=user_api_key_auth_obj.user_id if isinstance(user_api_key_auth_obj.user_id, str) else None, + team_id=user_api_key_auth_obj.team_id if isinstance(user_api_key_auth_obj.team_id, str) else None, + prisma_present=prisma_client is not None, + default_id=default_id if isinstance(default_id, str) else None, + skip_budget_checks=skip_budget_checks, + ) + await _member_model_budget( + user_api_key_auth_obj, facts, MemberIO(prisma_client, user_api_key_cache, proxy_logging_obj), MemberRead() + ) diff --git a/litellm/proxy/auth/budget_lookup_adapter.py b/litellm/proxy/auth/budget_lookup_adapter.py new file mode 100644 index 00000000000..1f12b0eb053 --- /dev/null +++ b/litellm/proxy/auth/budget_lookup_adapter.py @@ -0,0 +1,144 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import datetime +from typing import Final, assert_never + +from litellm._auth import budget_lookup as core +from litellm._auth.failure import NativeAuthFailure +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import ( + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.common_budget_adapter import metadata +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@dataclass(frozen=True, slots=True) +class LookupIO: + token: UserAPIKeyAuth | None + team: LiteLLM_TeamTable | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + + +@dataclass(frozen=True, slots=True) +class OrganizationReading: + org_id: str + row: LiteLLM_OrganizationTable + + +async def organization( + facts: core.OrganizationFacts, io: LookupIO, org_id: str | None = None +) -> OrganizationReading | None: + from litellm.proxy.auth.auth_checks import get_org_object + + action: Final = core.organization(facts) + match action: + case "Skip": + return None + case "ReadTokenId": + assert io.token is not None + token_id: Final = io.token.org_id + return await organization(replace(facts, token_id=token_id is not None), io, token_id) + case "ReadTeamId": + assert io.team is not None + team_id: Final = io.team.organization_id + return await organization(replace(facts, team_id=team_id is not None), io, team_id) + case "ReadOrganization": + assert org_id is not None and io.prisma is not None + try: + row: Final = await get_org_object( + org_id=org_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + include_budget_table=True, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # pass the raw lookup failure to native Phase 1 policy + return await organization(replace(facts, lookup_failed=True), io, org_id) + result: Final = core.organization(replace(facts, row_present=row is not None)) + assert result in ("Skip", "UseOrganization") + if result == "Skip": + return None + assert row is not None + return OrganizationReading(org_id, row) + case "UseOrganization": + raise AssertionError("organization row was not retained") + case _: + assert_never(action) + + +@dataclass(frozen=True, slots=True) +class MemberReading: + membership: LiteLLM_TeamMembership | None + row: LiteLLM_BudgetTable | None + now: datetime + metadata: Mapping[str, object] + default_id: object + default: LiteLLM_BudgetTable | None = None + + +async def member( + facts: core.MemberFacts, + io: LookupIO, + membership: LiteLLM_TeamMembership | None, + reading: MemberReading | None = None, +) -> MemberReading | None: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget, get_team_membership, get_utc_datetime + + action: Final = core.member(facts) + match action: + case "Skip": + return None + case "ReadTeamId": + assert io.team is not None + return await member(replace(facts, team_id=io.team.team_id is not None), io, membership) # pyright: ignore[reportUnnecessaryComparison] # Phase 1 accepts constructed cache rows with missing IDs + case "ReadUserId": + assert io.token is not None + return await member(replace(facts, user_id=io.token.user_id is not None), io, membership) + case "ReadMembership": + assert io.token is not None and io.token.user_id is not None and io.team is not None + loaded: Final = await get_team_membership( + user_id=io.token.user_id, + team_id=io.team.team_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + ) + return await member(replace(facts, membership_loaded=True), io, loaded) + case "ReadLimits": + assert io.team is not None + row: Final = membership.litellm_budget_table if membership is not None else None + now: Final = get_utc_datetime() + values: Final = metadata(io.team.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw legacy metadata boundary + default_id: Final = values.get("team_member_budget_id") + snapshot: Final = MemberReading(membership, row, now, values, default_id) + return await member( + replace( + facts, + limits_read=True, + maximum_present=row is not None and row.max_budget is not None, + default_id_string=isinstance(default_id, str), + ), + io, + membership, + snapshot, + ) + case "ReadDefault": + assert reading is not None and isinstance(reading.default_id, str) + default: Final = await get_team_member_default_budget( + budget_id=reading.default_id, prisma_client=io.prisma, user_api_key_cache=io.cache + ) + return await member(replace(facts, default_read=True), io, membership, replace(reading, default=default)) + case "UseLimits": + assert reading is not None + return reading + case _: + assert_never(action) diff --git a/litellm/proxy/auth/budget_pricing_adapter.py b/litellm/proxy/auth/budget_pricing_adapter.py new file mode 100644 index 00000000000..863294731d0 --- /dev/null +++ b/litellm/proxy/auth/budget_pricing_adapter.py @@ -0,0 +1,180 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter + +import litellm +from litellm._auth import budget_pricing +from litellm._auth.budget_pricing import DeploymentPrice, ModelPriceFacts, RouteBudgetFacts +from litellm._auth.failure import NativeAuthFailure +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.router_api import RegistryReader +from litellm.types.router import DeploymentTypedDict + +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT: Final = TypeAdapter(object) + + +def _number(raw: object) -> float | None: + return float(raw) if isinstance(raw, (float, int)) and not isinstance(raw, bool) else None + + +def _deployment_price(deployment: DeploymentTypedDict) -> DeploymentPrice: + info: Final = _MAPPING.validate_python(deployment.get("model_info") or {}) # pyright: ignore[reportUnknownMemberType] # deployment TypedDict contains legacy metadata + model_id: Final = info.get("id") + raw: Final = _OBJECT.validate_python(litellm.model_cost.get(model_id)) if model_id is not None else None # pyright: ignore[reportUnknownMemberType] # project one legacy global price entry + entry: Final = as_str_mapping(raw) or {} + params: Final = _MAPPING.validate_python(deployment["litellm_params"]) + provider: Final = params.get("custom_llm_provider") or str(params.get("model")).partition("/")[0] + query_price: Final = decisions_request_price( + {**entry, "litellm_provider": provider}, decisions_pricing_fields(entry) + ) + return DeploymentPrice( + has_input_price="input_cost_per_token" in entry, + has_output_price="output_cost_per_token" in entry, + input_price=_number(entry.get("input_cost_per_token")), + output_price=_number(entry.get("output_cost_per_token")), + query_price=query_price, + ptu_count=0.0 if info.get("ptu_count") is not None else None, + cost_per_ptu_per_hour=0.0 if info.get("cost_per_ptu_per_hour") is not None else None, + ) + + +def read_model_price_facts(model: str, router: RegistryReader, cached: bool | None = None) -> ModelPriceFacts: + from litellm.proxy.auth.auth_checks import ( + _deployments_served_for, # pyright: ignore[reportPrivateUsage] # shared registry projection + ) + + if cached is not None: + return ModelPriceFacts(cached=cached) + try: + group: Final = router.registry().group(model, include_hidden=True) + facts: Final = ModelPriceFacts( + group_found=group is not None, + input_price=group.input_cost_per_token if group is not None else None, + output_price=group.output_cost_per_token if group is not None else None, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching + return ModelPriceFacts(lookup_failed=True) + if not budget_pricing.price(facts).read_deployments: + return facts + try: + deployments: Final = tuple( + _deployment_price(deployment) for deployment in _deployments_served_for(model, router) + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching + return ModelPriceFacts(lookup_failed=True) + return replace(facts, deployments=deployments) + + +def _read_cached_model_price(model: str, router: RegistryReader, cache: Mapping[str, bool] | None) -> ModelPriceFacts: + cached: Final = cache.get(model) if isinstance(cache, dict) else None + facts: Final = read_model_price_facts(model, router, cached) + decision: Final = budget_pricing.price(facts) + if isinstance(cache, dict) and decision.cache_value is not None: + cache[model] = decision.cache_value # rebind-ok: native-selected registry cache write + return facts + + +def read_model_cost_facts(model: str, router: RegistryReader) -> ModelPriceFacts: + return _read_cached_model_price(model, router, router.registry().zero_cost()) + + +def is_model_cost_zero(model: str | Sequence[str] | None, llm_router: RegistryReader | None) -> bool: + if model is None or llm_router is None: + return budget_pricing.price(ModelPriceFacts()).zero_cost + names: Final = (model,) if isinstance(model, str) else tuple(model) + cache: Final = llm_router.registry().zero_cost() + return all(budget_pricing.price(_read_cached_model_price(name, llm_router, cache)).zero_cost for name in names) + + +def read_route_budget_facts(route: str, body: Mapping[str, object] | None = None) -> RouteBudgetFacts: + from litellm.proxy.auth.auth_checks import BUDGET_ENFORCED_SIDE_EFFECT_ROUTES, MODEL_DISCOVERY_ROUTES + + method: Final = body.get("method") if body is not None else None + return RouteBudgetFacts( + side_effect_route=route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES, + discovery_route=route in MODEL_DISCOVERY_ROUTES, + llm_route=RouteChecks.is_llm_api_route(route=route), + route=route, + method=method if isinstance(method, str) else None, + body_present=body is not None, + ) + + +def _route_decision(route: str, body: Mapping[str, object] | None = None) -> budget_pricing.RouteBudgetDecision: + facts: Final = read_route_budget_facts(route, body) + decision: Final = budget_pricing.route(facts) + if decision != "ReadPassThrough": + return decision + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered: Final = InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + return budget_pricing.route(replace(facts, registered_pass_through=registered)) + + +def route_skips_budget_checks(route: str) -> bool: + return _route_decision(route) == "Skip" + + +def read_budget_exemption_facts( + route: str, + model: str | Sequence[str] | None, + llm_router: RegistryReader | None, + *, + request_body: Mapping[str, object] | None = None, +) -> budget_pricing.ExemptionFacts: + names: Final = (model,) if isinstance(model, str) else tuple(model or ()) + initial: Final = budget_pricing.ExemptionFacts( + read_route_budget_facts(route, request_body), + model is not None, + llm_router is not None, + len(names), + ) + return _read_exemption(initial, names, llm_router, None) + + +def _read_exemption( + facts: budget_pricing.ExemptionFacts, + models: tuple[str, ...], + router: RegistryReader | None, + cache: Mapping[str, bool] | None, +) -> budget_pricing.ExemptionFacts: + current = facts # rebind-ok: native-selected IO driver state + selected_cache = cache # rebind-ok: cache handle acquired at the first model read + while True: + decision = budget_pricing.exemption(current) + if decision.read_pass_through: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered = InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=current.route.route) + current = replace(current, route=replace(current.route, registered_pass_through=registered)) + continue + if decision.read_price_cache: + assert router is not None + selected_cache = router.registry().zero_cost() + current = replace(current, cache_read=True) + continue + if decision.read_model is None: + return current + assert router is not None + price = _read_cached_model_price(models[decision.read_model], router, selected_cache) + current = replace(current, prices=(*current.prices, price)) + + +def request_skips_budget_checks( + route: str, + model: str | Sequence[str] | None, + llm_router: RegistryReader | None, + *, + request_body: Mapping[str, object] | None = None, +) -> bool: + facts: Final = read_budget_exemption_facts(route, model, llm_router, request_body=request_body) + return budget_pricing.exemption(facts).exempt diff --git a/litellm/proxy/auth/budget_throttle.py b/litellm/proxy/auth/budget_throttle.py index 7122c8edd05..8c06454e3b2 100644 --- a/litellm/proxy/auth/budget_throttle.py +++ b/litellm/proxy/auth/budget_throttle.py @@ -10,39 +10,20 @@ so it never compounds across requests. """ import math -from typing import Final -import litellm from litellm.proxy._types import UserAPIKeyAuth def budget_throttle_percentage() -> float | None: - """ - The global throttle percentage, or None when throttling is disabled / - misconfigured (in which case an over-budget key is hard-blocked, the safe - default). - """ - pct: Final = litellm.budget_exceeded_throttle_percentage - if not isinstance(pct, (int, float)) or isinstance(pct, bool): - return None - if not 0 < pct <= 1: - return None - return float(pct) + from litellm.proxy.auth import key_budget_context_adapter + + return key_budget_context_adapter.budget_throttle_percentage() def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: - """ - True when a key that exceeded its own ``max_budget`` should be throttled - rather than blocked: it opted in, a valid global percentage is set, and the - key has a TPM or RPM limit to scale down. A key with neither limit has - nothing to throttle, so it stays hard-blocked (the safe default) rather than - serving unlimited requests past its budget. - """ - if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True: - return False - if valid_token.tpm_limit is None and valid_token.rpm_limit is None: - return False - return budget_throttle_percentage() is not None + from litellm.proxy.auth import key_budget_context_adapter + + return key_budget_context_adapter.should_throttle_budget_exceeded(valid_token) def throttled_limit(limit: int | None, pct: float | None) -> int | None: diff --git a/litellm/proxy/auth/budget_windows_adapter.py b/litellm/proxy/auth/budget_windows_adapter.py new file mode 100644 index 00000000000..647d5bdb75f --- /dev/null +++ b/litellm/proxy/auth/budget_windows_adapter.py @@ -0,0 +1,101 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Final + +from pydantic import BaseModel, TypeAdapter + +from litellm._auth.common_budgets import BudgetFacts, WindowFacts, window_start +from litellm._auth.failure import NativeAuthFailure +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.common_budget_adapter import CounterRead, raise_budget, read_budget + + +@dataclass(frozen=True, slots=True) +class BudgetWindow: + budget_duration: str + max_budget: float + reset_at: object = None + + +_WINDOW: Final = TypeAdapter(BudgetWindow) + + +def _window(raw: object) -> BudgetWindow: + return _WINDOW.validate_python(raw.model_dump() if isinstance(raw, BaseModel) else raw) + + +def _seconds(duration: str) -> int | None: + try: + return duration_in_seconds(duration) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve invalid-duration skip behavior + return None + + +def _reset(value: object) -> datetime | None: + if isinstance(value, datetime): + return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value + if isinstance(value, str): + try: + parsed: Final = datetime.fromisoformat(value.replace("Z", "+00:00")) + except NativeAuthFailure: + raise + except ValueError: + return None + return _reset(parsed) + return None + + +async def check_windows(windows: Sequence[object], *, entity: str, identity: str | None) -> None: + for raw in windows: + await check_window(raw, entity=entity, identity=identity) + + +async def check_window(raw: object, *, entity: str, identity: str | None) -> None: + window: Final = _window(raw) + seconds: Final = _seconds(window.budget_duration) + reset: Final = _reset(window.reset_at) + start_timestamp: Final = window_start( + WindowFacts(seconds, reset.timestamp() if reset is not None else None, datetime.now(timezone.utc).timestamp()) + ) + start: Final = datetime.fromtimestamp(start_timestamp, timezone.utc) if start_timestamp is not None else None + facts, result = await read_budget( + BudgetFacts("KeyWindow" if entity == "Key" else "TeamWindow", window.max_budget), + CounterRead( + f"spend:{entity.lower()}:{identity}:window:{window.budget_duration}", + 0.0, + window.max_budget, + window_entity_type=entity, + window_entity_id=identity, + window_duration=window.budget_duration, + window_start=start, + ), + ) + if result.denied and facts.spend is not None: + label: Final = "Key" if entity == "Key" else f"Team={identity}" + raise_budget( + facts, + Litellm_EntityType.KEY if entity == "Key" else Litellm_EntityType.TEAM, + identity, + f"ExceededBudget: {label} over {window.budget_duration} budget. Spend=${facts.spend:.4f}, Limit=${window.max_budget:.2f}", + ) + + +async def virtual_key_multi_budget_check(valid_token: UserAPIKeyAuth) -> None: + await check_windows( + TypeAdapter(tuple[object, ...]).validate_python(valid_token.budget_limits or ()), # pyright: ignore[reportUnknownMemberType] # legacy input + entity="Key", + identity=valid_token.token, + ) + + +async def team_multi_budget_check(team_object: LiteLLM_TeamTable | None) -> None: + if team_object is not None: + await check_windows( + TypeAdapter(tuple[object, ...]).validate_python(team_object.budget_limits or ()), # pyright: ignore[reportUnknownMemberType] # validate legacy window dictionaries + entity="Team", + identity=team_object.team_id, + ) diff --git a/litellm/proxy/auth/common_budget_adapter.py b/litellm/proxy/auth/common_budget_adapter.py new file mode 100644 index 00000000000..e1894685d02 --- /dev/null +++ b/litellm/proxy/auth/common_budget_adapter.py @@ -0,0 +1,431 @@ +import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm._auth import budget_lookup, common_budgets +from litellm._auth.common_budgets import AlertFacts, BudgetDecision, BudgetFacts, MemberLimitFacts +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import ( + CallInfo, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.utils import ProxyLogging + +_METADATA: Final = TypeAdapter(Mapping[str, object]) +AlertKind: TypeAlias = Literal[ + "token_budget", + "user_budget", + "soft_budget", + "max_budget_alert", + "team_budget", + "organization_budget", + "proxy_budget", + "projected_limit_exceeded", + "project_budget", +] +_ALERT: Final = TypeAdapter[AlertKind](AlertKind) + + +def metadata(raw: object) -> Mapping[str, object]: + return _METADATA.validate_python(raw or {}) + + +@dataclass(frozen=True, slots=True) +class CounterRead: + key: str + fallback: float + max_budget: float | None + authoritative: bool = False + window_entity_type: str | None = None + window_entity_id: str | None = None + window_duration: str | None = None + window_start: datetime | None = None + + +class _CounterOptions(TypedDict, total=False): + fallback_authoritative: ReadOnly[bool] + window_entity_type: ReadOnly[str | None] + window_entity_id: ReadOnly[str | None] + window_duration: ReadOnly[str | None] + window_start: ReadOnly[datetime | None] + + +async def read_counter(counter: CounterRead) -> float: + from litellm.proxy.proxy_server import get_current_spend + + authoritative: Final[_CounterOptions] = {"fallback_authoritative": True} if counter.authoritative else {} + window: Final[_CounterOptions] = ( + { + "window_entity_type": counter.window_entity_type, + "window_entity_id": counter.window_entity_id, + "window_duration": counter.window_duration, + "window_start": counter.window_start, + } + if counter.window_entity_type is not None + else {} + ) + options: Final[_CounterOptions] = {**authoritative, **window} + return await get_current_spend( + counter_key=counter.key, + fallback_spend=counter.fallback, + max_budget=counter.max_budget, + **options, + ) + + +async def read_budget( + facts: BudgetFacts, + counter: CounterRead, + reader: Callable[[CounterRead], Awaitable[float]] = read_counter, +) -> tuple[BudgetFacts, BudgetDecision]: + initial: Final = common_budgets.check(facts) + if not initial.read_spend: + return facts, initial + snapshot: Final = replace(facts, spend=await reader(counter)) + return snapshot, common_budgets.check(snapshot) + + +def emit_alert(logging: ProxyLogging, kind: str, info: CallInfo) -> None: + asyncio.create_task(logging.budget_alerts(type=_ALERT.validate_python(kind), user_info=info)) + + +def key_alert_info( + key: UserAPIKeyAuth, spend: float | None, limit: float | None, entity: Litellm_EntityType +) -> CallInfo: + return CallInfo( + token=key.token, + spend=spend if spend is not None else 0.0, + max_budget=limit, + user_id=key.user_id, + team_id=key.team_id, + team_alias=key.team_alias, + organization_id=key.org_id, + event_group=entity, + ) + + +def raise_budget(facts: BudgetFacts, entity: Litellm_EntityType, identity: str | None, message: str | None) -> None: + assert facts.spend is not None and facts.limit is not None + raise litellm.BudgetExceededError( + current_cost=facts.spend, + max_budget=facts.limit, + message=message, + entity_type=entity.value, + entity_id=identity, + ) + + +async def team_max_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if team_object is None: + return + facts, result = await read_budget( + BudgetFacts("Team", team_object.max_budget, token_present=valid_token is not None), + CounterRead(f"spend:team:{team_object.team_id}", team_object.spend or 0.0, team_object.max_budget), + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + key_alert_info(valid_token, facts.spend, facts.limit, Litellm_EntityType.TEAM), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.TEAM, + team_object.team_id, + f"Budget has been exceeded! Team={team_object.team_id} Current cost: {facts.spend}, Max budget: {facts.limit}", + ) + + +async def virtual_key_max_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + from litellm.proxy.auth.auth_checks import key_budget_exceeded_error, send_token_budget_alert + + raw_percentage: Final[object] = litellm.budget_exceeded_throttle_percentage + percentage: Final = ( + float(raw_percentage) + if isinstance(raw_percentage, (int, float)) and not isinstance(raw_percentage, bool) + else None + ) + facts, result = await read_budget( + BudgetFacts( + "Key", + valid_token.max_budget, + throttle_percentage=percentage, + throttle_opt_in=metadata(valid_token.metadata).get("throttle_on_budget_exceeded") is True, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy metadata is validated by metadata() + tpm_limit=valid_token.tpm_limit, + rpm_limit=valid_token.rpm_limit, + ), + CounterRead(f"spend:key:{valid_token.token}", valid_token.spend or 0.0, valid_token.max_budget), + ) + if result.alert is not None and facts.spend is not None: + send_token_budget_alert(valid_token, facts.spend, proxy_logging_obj, user_obj) + if result.throttle_percentage is not None: + valid_token.budget_throttle_pct = result.throttle_percentage # rebind-ok: native request-local throttle + if result.denied and facts.spend is not None and facts.limit is not None: + raise key_budget_exceeded_error(valid_token, facts.spend, facts.limit) + + +async def project_max_budget_check( + project_object: LiteLLM_ProjectTableCachedObj | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + from litellm.redis.stores.spend_counters import project_counter + + if project_object is None: + return + limit: Final = ( + project_object.litellm_budget_table.max_budget if project_object.litellm_budget_table is not None else None + ) + facts, result = await read_budget( + BudgetFacts("Project", limit, token_present=valid_token is not None), + CounterRead(project_counter(project_object.project_id), project_object.spend or 0.0, limit), + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + key_alert_info(valid_token, facts.spend, limit, Litellm_EntityType.PROJECT), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.PROJECT, + project_object.project_id, + f"Budget has been exceeded! Project={project_object.project_id} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def user_max_budget_check( + user_object: LiteLLM_UserTable | None, + team_object: LiteLLM_TeamTable | None, + general_settings: Mapping[str, object], + proxy_logging_obj: ProxyLogging, +) -> None: + if user_object is None: + return + facts, result = await read_budget( + BudgetFacts( + "User", + user_object.max_budget, + team_key=team_object is not None and team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # retain the legacy None guard for constructed cache rows + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + ), + CounterRead(f"spend:user:{user_object.user_id}", user_object.spend or 0.0, user_object.max_budget), + ) + if result.alert is not None: + emit_alert( + proxy_logging_obj, + result.alert, + CallInfo( + spend=facts.spend if facts.spend is not None else 0.0, + max_budget=facts.limit, + user_id=user_object.user_id, + user_email=user_object.user_email, + event_group=Litellm_EntityType.USER, + ), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.USER, + user_object.user_id, + f"ExceededBudget: User={user_object.user_id} over budget. Spend={facts.spend}, Budget={facts.limit}", + ) + + +async def check_end_user_budget(end_user_obj: LiteLLM_EndUserTable, route: str) -> None: + from litellm.proxy.auth.route_checks import RouteChecks + + limit: Final = ( + end_user_obj.litellm_budget_table.max_budget if end_user_obj.litellm_budget_table is not None else None + ) + facts, result = await read_budget( + BudgetFacts("EndUser", limit, info_route=RouteChecks.is_info_route(route)), + CounterRead(f"spend:end_user:{end_user_obj.user_id}", end_user_obj.spend or 0.0, limit, authoritative=True), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.END_USER, + end_user_obj.user_id, + f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={facts.spend}, Budget={facts.limit}", + ) + + +def global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_checks: bool, route: str) -> None: + from litellm.proxy._types import MODEL_LIST_ROUTES + from litellm.proxy.auth.route_checks import RouteChecks + + facts: Final = BudgetFacts( + "Proxy", + litellm.max_budget, + global_proxy_spend, + skip=skip_budget_checks, + llm_route=RouteChecks.is_llm_api_route(route=route), + model_list_route=route in MODEL_LIST_ROUTES, + ) + if common_budgets.check(facts).denied: + raise_budget(facts, Litellm_EntityType.PROXY, None, None) + + +async def virtual_key_soft_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + from litellm.proxy.auth.auth_checks import send_soft_budget_alert + + if common_budgets.check(BudgetFacts("KeySoft", valid_token.soft_budget, valid_token.spend)).alert is not None: + send_soft_budget_alert(valid_token, proxy_logging_obj, user_obj) + + +def _soft_emails(raw: object) -> tuple[str, ...]: + return budget_lookup.emails(raw) + + +async def team_soft_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if team_object is None: + return + if not budget_lookup.read_soft_emails( + budget_lookup.SoftFacts(team_object.soft_budget, team_object.spend, valid_token is not None) + ): + return + emails: Final = _soft_emails(budget_lookup.soft_email_value(team_object.metadata)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw legacy metadata projection + result: Final = common_budgets.check( + BudgetFacts( + "TeamSoft", + team_object.soft_budget, + team_object.spend, + token_present=valid_token is not None, + alert_recipients=emails, + ) + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + CallInfo( + token=valid_token.token, + spend=team_object.spend if team_object.spend is not None else 0.0, + max_budget=team_object.max_budget, + soft_budget=team_object.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.TEAM, + alert_emails=list(emails), + ), + ) + + +async def project_soft_budget_check( + project_object: LiteLLM_ProjectTableCachedObj | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if project_object is None: + return + limit: Final = ( + project_object.litellm_budget_table.soft_budget if project_object.litellm_budget_table is not None else None + ) + result: Final = common_budgets.check( + BudgetFacts("ProjectSoft", limit, project_object.spend, token_present=valid_token is not None) + ) + if result.alert is not None and valid_token is not None: + emit_alert( + proxy_logging_obj, + result.alert, + CallInfo( + token=valid_token.token, + spend=project_object.spend, + max_budget=None, + soft_budget=limit, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.PROJECT, + ), + ) + + +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + return common_budgets.member_limit( + read_member_limit_facts( + member_budget, team_default_budget.max_budget if team_default_budget is not None else None, now + ) + ) + + +def read_member_limit_facts( + member_budget: LiteLLM_BudgetTable | None, default_limit: float | None, now: datetime +) -> MemberLimitFacts: + expiry: Final = member_budget.temp_budget_expiry if member_budget is not None else None + aware_expiry: Final = ( + expiry.replace(tzinfo=timezone.utc) if expiry is not None and expiry.tzinfo is None else expiry + ) + expiry_number: Final = ( + aware_expiry.replace(tzinfo=timezone.utc).timestamp() + if aware_expiry is not None and aware_expiry.tzinfo is now.tzinfo + else aware_expiry.timestamp() + if aware_expiry is not None + else None + ) + now_number: Final = ( + now.replace(tzinfo=timezone.utc).timestamp() + if aware_expiry is not None and aware_expiry.tzinfo is now.tzinfo + else now.timestamp() + ) + return MemberLimitFacts( + member_limit=member_budget.max_budget if member_budget is not None else None, + default_limit=default_limit, + temporary_increase=member_budget.temp_budget_increase if member_budget is not None else None, + expires_at=expiry_number, + now=now_number, + now_naive=now.utcoffset() is None, + ) + + +def threshold_alert_decision( + spend: float | None, + max_budget: float | None, + config: Mapping[str, Sequence[str]] | None, + *, + team_member: bool = False, +) -> bool: + minimum: Final = min((int(key) for key in (config or {}) if team_member or key.isdigit()), default=None) + return common_budgets.alert( + AlertFacts(spend, max_budget, bool(config), minimum, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, team_member) + ) diff --git a/litellm/proxy/auth/common_budget_lookups.py b/litellm/proxy/auth/common_budget_lookups.py new file mode 100644 index 00000000000..071892cfa17 --- /dev/null +++ b/litellm/proxy/auth/common_budget_lookups.py @@ -0,0 +1,237 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Final + +from litellm._auth import budget_lookup, common_budgets +from litellm._auth.common_budgets import BudgetFacts +from litellm.proxy._types import ( + Litellm_EntityType, + LiteLLM_TagTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth import budget_lookup_adapter +from litellm.proxy.auth.common_budget_adapter import ( + CounterRead, + emit_alert, + key_alert_info, + raise_budget, + read_budget, + resolve_effective_team_member_budget, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.redis.stores.spend_counters import model_access_group_counter +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + +async def organization_max_budget_check( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + reading: Final = await budget_lookup_adapter.organization( + budget_lookup.OrganizationFacts(valid_token is not None, prisma_client is not None, team_object is not None), + budget_lookup_adapter.LookupIO(valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + if reading is None: + return + assert valid_token is not None + org_id: Final = reading.org_id + org: Final = reading.row + carry_organization_budget_state(valid_token=valid_token, org_table=org) + limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None + facts, result = await read_budget( + BudgetFacts("Organization", limit), CounterRead(f"spend:org:{org_id}", org.spend or 0.0, limit) + ) + if result.alert is not None: + info: Final = key_alert_info(valid_token, facts.spend, limit, Litellm_EntityType.ORGANIZATION) + info.organization_id = org_id + emit_alert(proxy_logging_obj, result.alert, info) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.ORGANIZATION, + org_id, + f"Budget has been exceeded! Organization={org_id} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def tag_max_budget_check_for_tags( + tags: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + valid_token: UserAPIKeyAuth | None, +) -> None: + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + if prisma_client is None or not tags: + return + rows: Final = await get_tag_objects_batch( + tag_names=tags, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for name in tags: + await check_tag_row(name, rows.get(name)) + + +async def check_tag_row(name: str, row: LiteLLM_TagTable | None) -> None: + limit: Final = ( + row.litellm_budget_table.max_budget if row is not None and row.litellm_budget_table is not None else None + ) + facts, result = await read_budget( + BudgetFacts("Tag", limit), + CounterRead(f"spend:tag:{name}", (row.spend or 0.0) if row is not None else 0.0, limit, authoritative=True), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.TAG, + name, + f"Budget has been exceeded! Tag={name} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def model_access_group_max_budget_check( + matched_model_access_groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + if prisma_client is None or not matched_model_access_groups: + return + rows: Final = await get_model_access_group_budgets_batch( + access_group_names=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + for name in matched_model_access_groups: + await check_access_group_row(name, rows.get(name)) + + +async def check_access_group_row(name: str, row: ModelAccessGroupBudget | None) -> None: + limit: Final = row.max_budget if row is not None else None + facts, result = await read_budget( + BudgetFacts("AccessGroup", limit), + CounterRead( + model_access_group_counter(name), + row.spend if row is not None else 0.0, + limit, + authoritative=True, + ), + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.MODEL_ACCESS_GROUP, + name, + f"Budget has been exceeded! Model access group={name} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +async def check_team_member_budget( + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> None: + from litellm.proxy.auth.budget_alert_adapter import team_member_max_budget_alert_check + + reading: Final = await budget_lookup_adapter.member( + budget_lookup.MemberFacts(team_object is not None, valid_token is not None, team_membership_loaded), + budget_lookup_adapter.LookupIO(valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj), + team_membership, + ) + if reading is None: + return + assert team_object is not None and valid_token is not None and valid_token.user_id is not None + membership: Final = reading.membership + team_metadata: Final = reading.metadata + limit: Final = resolve_effective_team_member_budget(reading.row, reading.default, reading.now) + facts, result = await read_budget( + BudgetFacts("TeamMember", limit), + CounterRead( + f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + (membership.spend or 0.0) if membership is not None else 0.0, + limit, + ), + ) + if result.alert is not None and facts.spend is not None and limit is not None: + team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=dict(team_metadata), + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=facts.spend, + max_budget=limit, + ) + if result.denied: + raise_budget( + facts, + Litellm_EntityType.TEAM_MEMBER, + f"{valid_token.user_id}:{team_object.team_id}", + f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {facts.spend}, Max budget: {limit}", + ) + + +@dataclass(frozen=True, slots=True) +class AccessGroupBudgetReading: + name: str + facts: BudgetFacts + + +async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget | None) -> AccessGroupBudgetReading: + limit: Final = row.max_budget if row is not None else None + facts, _decision = await read_budget( + BudgetFacts("AccessGroup", limit), + CounterRead( + model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True + ), + ) + return AccessGroupBudgetReading(name, facts) + + +async def read_access_group_budget_facts( + groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> tuple[AccessGroupBudgetReading, ...]: + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + if prisma_client is None or not groups: + return () + rows: Final = await get_model_access_group_budgets_batch( + access_group_names=groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return await _read_access_group_rows(tuple(groups), rows) + + +async def _read_access_group_rows( + groups: tuple[str, ...], + rows: Mapping[str, ModelAccessGroupBudget], +) -> tuple[AccessGroupBudgetReading, ...]: + readings: tuple[AccessGroupBudgetReading, ...] = () # rebind-ok: native-selected IO driver accumulator + for name in groups: + reading = await read_access_group_budget_row(name, rows.get(name)) + readings = (*readings, reading) + if common_budgets.check(reading.facts).denied: + return readings + return readings diff --git a/litellm/proxy/auth/end_user_context_adapter.py b/litellm/proxy/auth/end_user_context_adapter.py new file mode 100644 index 00000000000..aae5d73f770 --- /dev/null +++ b/litellm/proxy/auth/end_user_context_adapter.py @@ -0,0 +1,406 @@ +from collections.abc import Mapping, MutableMapping +from dataclasses import replace +from types import MappingProxyType +from typing import Final, assert_never + +from opentelemetry.trace import Span + +import litellm +from litellm._auth import end_user_context as core +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_EndUserTable, UserAPIKeyAuth +from litellm.proxy.auth.budget_context_adapter import mapping_value, resolve_default_end_user_budget +from litellm.proxy.auth.end_user_context_io import EndUserIO, read_default_budget +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +_PARAM_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType( + { + "EndUserId": "end_user_id", + "Tpm": "end_user_tpm_limit", + "Rpm": "end_user_rpm_limit", + "Tpd": "end_user_tpd_limit", + "MaxBudget": "end_user_max_budget", + "Models": "end_user_model_max_budget", + "Region": "allowed_model_region", + } +) +_BUDGET_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType( + { + "Tpm": "tpm_limit", + "Rpm": "rpm_limit", + "Tpd": "tpd_limit", + "MaxBudget": "max_budget", + "Models": "model_max_budget", + } +) + + +def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: + raw: Final = key_metadata.get("end_user_budget_id") if key_metadata is not None else None + return core.normalize_id(core.IdFacts(raw if isinstance(raw, str) else None)) + + +def select_default_id(facts: core.DefaultFacts) -> str | None: + decision: Final = core.default_id(facts) + if decision.read_global: + return select_default_id(replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id)) + return decision.budget_id + + +async def get_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + budget_id: str | None = None, +) -> LiteLLM_BudgetTable | None: + selected: Final = select_default_id(core.DefaultFacts(prisma_client is not None, budget_id, None)) + if selected is None: + return None + return await read_default_budget(selected, EndUserIO(prisma_client, user_api_key_cache, parent_otel_span)) + + +async def _apply_default( + facts: core.ApplyFacts, row: LiteLLM_EndUserTable, io: EndUserIO, default: LiteLLM_BudgetTable | None = None +) -> LiteLLM_EndUserTable: + decision: Final = core.apply_default(facts) + match decision: + case "Keep": + return row + case "ReadGlobal": + return await _apply_default( + replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), row, io + ) + case "ReadDefault": + budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span) + return await _apply_default( + replace(facts, default_read=True, default_present=budget is not None), row, io, budget + ) + case "Copy": + assert default is not None + verbose_proxy_logger.debug("Applied default budget %s to end user %s", default.budget_id, row.user_id) + return row.model_copy(update=MappingProxyType({"litellm_budget_table": default})) + case _: + assert_never(decision) + + +async def apply_default_budget_to_end_user( + end_user_obj: LiteLLM_EndUserTable, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable: + return await _apply_default( + core.ApplyFacts(end_user_obj.budget_id, end_user_obj.litellm_budget_table is not None, key_end_user_budget_id), + end_user_obj, + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + ) + + +async def _registry(facts: core.RegistryFacts, io: EndUserIO) -> bool: + from litellm.proxy.auth import auth_checks + + decision: Final = core.registry(facts) + match decision: + case "Restricted": + return False + case "Unrestricted": + return True + case "Read": + assert io.prisma is not None + ids: Final = await auth_checks._load_end_user_restricted_registry( # pyright: ignore[reportPrivateUsage] # raw cached registry IO + prisma_client=io.prisma, user_api_key_cache=io.cache + ) + return await _registry( + replace(facts, registry_read=True, registry=tuple(ids) if ids is not None else None), io + ) + case _: + assert_never(decision) + + +async def end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, + key_end_user_budget_id: str | None = None, +) -> bool: + return await _registry( + core.RegistryFacts( + end_user_id, + litellm.max_end_user_budget_id, + key_end_user_budget_id, + litellm.validate_end_user_id_in_db, + token_end_user_max_budget is not None, + ), + EndUserIO(prisma_client, user_api_key_cache), + ) + + +def _budget_fields(row: LiteLLM_BudgetTable | None) -> core.Fields: + return core.Fields( + tpm=getattr(row, "tpm_limit", None) is not None, + rpm=getattr(row, "rpm_limit", None) is not None, + tpd=getattr(row, "tpd_limit", None) is not None, + max_budget=getattr(row, "max_budget", None) is not None, + models=getattr(row, "model_max_budget", None) is not None, + ) + + +def _budget_params(row: LiteLLM_BudgetTable | None) -> Mapping[str, object]: + selected: Final = core.assignments(core.AssignmentFacts("Params", _budget_fields(row))) + return MappingProxyType({_PARAM_NAMES[field]: getattr(row, _BUDGET_NAMES[field]) for field in selected}) + + +def apply_budget_limits_to_end_user_params( + end_user_params: MutableMapping[str, object], # mutable-ok: legacy request-local output parameter + budget_info: LiteLLM_BudgetTable, + end_user_id: str | None, +) -> None: + for name, value in _budget_params(budget_info).items(): + end_user_params[name] = value # rebind-ok: native-selected legacy request parameter projection + verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + + +def _token_fields(end_user_params: Mapping[str, object]) -> tuple[core.Field, ...]: + fields: Final = core.Fields( + tpm=end_user_params.get("end_user_tpm_limit") is not None, + rpm=end_user_params.get("end_user_rpm_limit") is not None, + tpd=end_user_params.get("end_user_tpd_limit") is not None, + models=end_user_params.get("end_user_model_max_budget") is not None, + region=end_user_params.get("allowed_model_region") is not None, + ) + return core.assignments(core.AssignmentFacts("Token", fields)) + + +def _project_token( + valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object], fields: tuple[core.Field, ...] +) -> UserAPIKeyAuth: + for field in fields: + setattr(valid_token, _PARAM_NAMES[field], end_user_params.get(_PARAM_NAMES[field])) + return valid_token + + +def update_valid_token_with_end_user_params( + valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object] +) -> UserAPIKeyAuth: + return _project_token(valid_token, end_user_params, _token_fields(end_user_params)) + + +async def _validation(facts: core.ValidationFacts, io: EndUserIO) -> core.ValidationDecision: + from litellm.proxy.auth import auth_checks + + decision: Final = core.validation(facts) + match decision.read: + case None: + return decision + case "Cache": + cached: Final[object] = await auth_checks._raw_cache(io.cache).async_get_cache( # pyright: ignore[reportPrivateUsage] # shared raw cache IO + key=f"end_user_validation:{facts.end_user_id}" + ) + return await _validation( + replace(facts, cache_read=True, cached=cached if isinstance(cached, str) else None), io + ) + case "EndUser": + try: + row: Final = await auth_checks.get_end_user_object( + end_user_id=facts.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + route=io.route, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve end-user validation lookup failure handling + verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", error) + return await _validation(replace(facts, end_user_read=True), io) + return await _validation(replace(facts, end_user_read=True, end_user_present=row is not None), io) + case "User": + assert facts.end_user_id is not None and io.prisma is not None + try: + user: Final = await auth_checks.get_user_object( + user_id=facts.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + user_id_upsert=False, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + check_db_only=False, + user_email=decision.user_email, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve secondary identity lookup failure handling + verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", error) + return await _validation(replace(facts, user_read=True), io) + return await _validation(replace(facts, user_read=True, user_present=user is not None), io) + case "WriteCache": + await io.cache.async_set_cache( + key=f"end_user_validation:{facts.end_user_id}", value=decision.cache_value, ttl=decision.cache_ttl + ) + return await _validation(replace(facts, cache_written=True), io) + case _: + assert_never(decision.read) + + +async def resolve_and_validate_end_user_id( + raw_end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", + key_end_user_budget_id: str | None = None, +) -> str | None: + decision: Final = await _validation( + core.ValidationFacts( + "Resolve", + raw_end_user_id, + litellm.validate_end_user_id_in_db, + prisma_client is not None, + litellm.max_end_user_budget_id, + key_end_user_budget_id, + ), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) + return decision.end_user_id + + +async def end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", +) -> bool: + decision: Final = await _validation( + core.ValidationFacts("Exists", end_user_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) + return decision.valid + + +def _complete_lookup( + decision: core.LookupDecision, + token: UserAPIKeyAuth, + row: LiteLLM_EndUserTable | None, + default: LiteLLM_BudgetTable | None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + if decision.set_max_budget: + assert default is not None + try: + token.end_user_max_budget = default.max_budget # rebind-ok: request-local native projection + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve request-local setter failure handling + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return token, row + + +async def _lookup( + facts: core.LookupFacts, + token: UserAPIKeyAuth, + io: EndUserIO, + row: LiteLLM_EndUserTable | None = None, + default: LiteLLM_BudgetTable | None = None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + from litellm.proxy.auth.auth_checks import get_end_user_object + + decision: Final = core.lookup(facts) + match decision.action: + case "Done": + return _complete_lookup(decision, token, row, default) + case "ReadGlobal": + return await _lookup( + replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), token, io, row + ) + case "ReadEndUser": + try: + end_user: Final = await get_end_user_object( + end_user_id=token.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + route=io.route, + token_end_user_max_budget=token.end_user_max_budget, + key_end_user_budget_id=facts.key_id, + ) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve the legacy lookup IO failure path + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace( + facts, + end_user_read=True, + end_user_present=end_user is not None, + own_budget_present=end_user is not None and end_user.litellm_budget_table is not None, + ), + token, + io, + end_user, + ) + case "ReadDefault": + try: + budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve the legacy default lookup IO failure path + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace(facts, default_read=True, default_present=budget is not None), token, io, row, budget + ) + case "ApplyOwn" | "ApplyDefault": + budget_row: Final = ( + row.litellm_budget_table if decision.action == "ApplyOwn" and row is not None else default + ) + params: Final = MappingProxyType( + { + "end_user_id": token.end_user_id, + "allowed_model_region": getattr(row, "allowed_model_region", None), + **_budget_params(budget_row), + } + ) + fields: Final = _token_fields(params) + try: + if decision.project_budget: + verbose_proxy_logger.debug("Applied budget limits to end user %s", token.end_user_id or "") + _project_token(token, params, fields) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve request-local setter failure handling + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace(facts, applied=True, token_max_present=token.end_user_max_budget is not None), + token, + io, + row, + default, + ) + case _: + assert_never(decision.action) + + +async def lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Span | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + key_id: Final = get_key_end_user_budget_id(mapping_value(getattr(valid_token, "metadata", None))) + return await _lookup( + core.LookupFacts(key_id), + valid_token, + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) diff --git a/litellm/proxy/auth/end_user_context_io.py b/litellm/proxy/auth/end_user_context_io.py new file mode 100644 index 00000000000..83cf9784886 --- /dev/null +++ b/litellm/proxy/auth/end_user_context_io.py @@ -0,0 +1,50 @@ +from dataclasses import dataclass +from typing import Final + +from opentelemetry.trace import Span + +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_BudgetTable +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.budget_repository import BudgetRepository + + +@dataclass(frozen=True, slots=True) +class EndUserIO: + prisma: PrismaClient | None + cache: UserApiKeyCache + span: Span | None = None + logging: ProxyLogging | None = None + route: str = "" + + +async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTable | None: + from litellm.proxy.auth import auth_checks + + assert io.prisma is not None + key: Final = f"default_end_user_budget:{budget_id}" + cached: Final = await io.cache.async_get_cache(key=key, model_type=LiteLLM_BudgetTable) + if cached is not None: + return cached + try: + table: Final = auth_checks._dictable_table(BudgetRepository(io.prisma), "budget") # pyright: ignore[reportPrivateUsage] # shared raw repository projection + record: Final = await table.find_unique(where={"budget_id": budget_id}) + if record is None: + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", budget_id.replace("\r", "").replace("\n", "") + ) + return None + row: Final = LiteLLM_BudgetTable.model_validate(record.dict()) + await io.cache.async_set_cache( + key=key, value=row, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(io.cache) + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve the default budget DB/cache IO failure contract + verbose_proxy_logger.error("Error fetching default end user budget: %s", error) + return None + + else: + return row diff --git a/litellm/proxy/auth/fallback_access_adapter.py b/litellm/proxy/auth/fallback_access_adapter.py new file mode 100644 index 00000000000..26e7530c924 --- /dev/null +++ b/litellm/proxy/auth/fallback_access_adapter.py @@ -0,0 +1,154 @@ +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope, fallback_authorization +from litellm._auth.model_access import ModelAccessFacts, key_model_grants +from litellm.proxy._types import LiteLLM_ProjectTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Sources: + model: str + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: DatabaseClient | None + team: LiteLLM_TeamTable | None + project: LiteLLM_ProjectTable | None + + +async def load_fallback_authorization_facts( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: DatabaseClient | None = None, + team_membership: LiteLLM_TeamMembership | None = None, + project_object: LiteLLM_ProjectTable | None = None, + refuse_agent_keys: bool = False, + team_object: LiteLLM_TeamTable | None = None, +) -> FallbackAuthorizationFacts: + facts: Final = FallbackAuthorizationFacts( + refuse_agent_keys=refuse_agent_keys, + agent_id=valid_token.agent_id, + team_present=team_object is not None, + embedded_team_grants=tuple(valid_token.team_models or ()), + member_grants=tuple(team_membership.litellm_budget_table.allowed_models or ()) + if team_membership is not None and team_membership.litellm_budget_table is not None + else (), + project_grants=tuple(project_object.models) if project_object is not None else (), + ) + return await _read(facts, _Sources(model, valid_token, llm_router, prisma_client, team_object, project_object)) + + +async def _read(facts: FallbackAuthorizationFacts, source: _Sources) -> FallbackAuthorizationFacts: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + decision: Final = fallback_authorization(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadKey": + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + key: Final = GroupedModelScope( + project_model_facts( + source.model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ), + tuple(source.token.access_group_ids or ()), + ) + return await _read(replace(facts, key=key), source) + case "ReadTeam": + if source.team is None: + raise NativeAuthRuntimeError("missing fallback team source") + team: Final = GroupedModelScope( + project_model_facts( + source.model, + source.router, + source.team.models, + source.token.team_model_aliases, + source.team.team_id, + key_model_aliases_for_auth_check(source.token), + "team", + ), + tuple(source.team.access_group_ids or ()), + ) + return await _read(replace(facts, team=team), source) + case "ReadEmbeddedTeam": + embedded: Final = project_model_facts( + source.model, + source.router, + facts.embedded_team_grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "team", + ) + return await _read(replace(facts, embedded_team=embedded), source) + case "ReadMember": + member: Final = project_model_facts( + source.model, + source.router, + facts.member_grants, + team_id=source.token.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="team", + ) + return await _read(replace(facts, member=member), source) + case "ReadProject": + project: Final = _project_scope(facts, source) + return await _read(replace(facts, project=project), source) + case "ReadKeyGroups": + if facts.key is None: + raise NativeAuthRuntimeError("missing fallback key scope") + key_scope: Final = await read_grouped_scope(facts.key, source.router, source.prisma, merge=False) + return await _read(replace(facts, key=key_scope), source) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing fallback team scope") + team_scope: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=team_scope), source) + case "ReadAssignedGroups": + assigned: Final = await get_authorized_resources_from_key_access_groups( + source.token, source.team, "access_model_names" + ) + return await _read(replace(facts, assigned_models=tuple(assigned)), source) + case "ReadAssignedScope": + if facts.assigned_models is None: + raise NativeAuthRuntimeError("missing fallback assigned models") + scope: Final = project_model_facts( + source.model, + source.router, + facts.assigned_models, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ) + return await _read(replace(facts, assigned_scope=scope), source) + + +def _project_scope(facts: FallbackAuthorizationFacts, source: _Sources) -> ModelAccessFacts: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + if source.project is None: + raise NativeAuthRuntimeError("missing fallback project source") + return project_model_facts( + source.model, + source.router, + facts.project_grants, + team_id=source.project.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="project", + ) diff --git a/litellm/proxy/auth/fallback_candidate_adapter.py b/litellm/proxy/auth/fallback_candidate_adapter.py new file mode 100644 index 00000000000..32b0870798b --- /dev/null +++ b/litellm/proxy/auth/fallback_candidate_adapter.py @@ -0,0 +1,320 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Final, Protocol + +from pydantic import BaseModel + +import litellm +from litellm._auth.decision import FallbackTarget +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_candidate import ( + CandidateBudgetRead, + CandidateBudgetScope, + FallbackCandidateDecision, + FallbackCandidateFacts, + fallback_candidate, +) +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.model_budget_data import ModelBudgetEntity +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.budget_pricing_adapter import read_model_cost_facts +from litellm.proxy.auth.common_budget_lookups import read_access_group_budget_facts +from litellm.proxy.auth.fallback_access_adapter import load_fallback_authorization_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_attribution_adapter import read_model_group_attribution_facts +from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading +from litellm.proxy.auth.route_policy_adapter import route_context +from litellm.proxy.auth.unpriced_model_adapter import read_unpriced_model_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.hooks.model_max_budget_limiter import team_member_budget_entity_id +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router import Router + + +class ModelBudgetReader(Protocol): + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> ModelBudgetReading: ... + + +class _BudgetFields(BaseModel, frozen=True): + model_max_budget: object = None + user_model_max_budget: object = None + end_user_model_max_budget: object = None + team_model_max_budget: object = None + team_member_model_max_budget: object = None + + +@dataclass(frozen=True, slots=True) +class _BudgetSource: + entity: ModelBudgetEntity + identity: str | None + budgets: Mapping[str, object] | None + + +@dataclass(frozen=True, slots=True) +class _Dependencies: + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + + +@dataclass(frozen=True, slots=True) +class _Sources: + target: FallbackTarget + token: UserAPIKeyAuth + limiter: ModelBudgetReader + router: Router | None + route: str + project: LiteLLM_ProjectTable | None + team: LiteLLM_TeamTable | None + customer: LiteLLM_EndUserTable | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + dependencies: _Dependencies | None = None + membership: LiteLLM_TeamMembership | None = None + budgets: tuple[_BudgetSource, ...] = () + + +def _mapping(raw: object) -> Mapping[str, object] | None: + from litellm.proxy.auth.budget_context_adapter import mapping_value + + return mapping_value(raw) + + +def _budget_sources(token: UserAPIKeyAuth) -> tuple[_BudgetSource, ...]: + fields: Final = _BudgetFields.model_validate(token, from_attributes=True) + member_id: Final = ( + team_member_budget_entity_id(token.user_id, token.team_id) + if token.user_id is not None and token.team_id is not None + else None + ) + return ( + _BudgetSource("Key", token.token, _mapping(fields.model_max_budget)), + _BudgetSource("User", token.user_id, _mapping(fields.user_model_max_budget)), + _BudgetSource("EndUser", token.end_user_id, _mapping(fields.end_user_model_max_budget)), + _BudgetSource("Team", token.team_id, _mapping(fields.team_model_max_budget)), + _BudgetSource("TeamMember", member_id, _mapping(fields.team_member_model_max_budget)), + ) + + +async def read_fallback_candidate_facts( + *, + target: FallbackTarget, + valid_token: UserAPIKeyAuth, + limiter: ModelBudgetReader, + llm_router: Router | None, + route: str, + project_object: LiteLLM_ProjectTable | None = None, + team_object: LiteLLM_TeamTable | None = None, + end_user_object: LiteLLM_EndUserTable | None = None, + prisma_client: PrismaClient | None = None, + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> FallbackCandidateFacts: + facts: Final = FallbackCandidateFacts( + router_present=llm_router is not None, + user_present=valid_token.user_id is not None, + team_present=valid_token.team_id is not None, + block_unpriced=litellm.block_requests_for_models_without_pricing, + customer_present=end_user_object is not None, + ) + return await _read( + facts, + _Sources( + target, + valid_token, + limiter, + llm_router, + route, + project_object, + team_object, + end_user_object, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + ), + ) + + +def _warn(source: _Sources, *, exception: bool) -> None: + verbose_proxy_logger.warning( + "Skipping fallback candidate %s after a budget or authorization check error", + source.target.model.replace("\r", "").replace("\n", ""), + exc_info=exception, + ) + + +async def _read(facts: FallbackCandidateFacts, source: _Sources) -> FallbackCandidateFacts: + decision: Final = fallback_candidate(facts) + if decision.kind == "Complete": + if decision.log_failure: + _warn(source, exception=False) + return facts + if decision.kind == "ReadMembership": + return await _read_membership(facts, source) + try: + updated, updated_source = await _effect(decision, facts, source) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy candidate IO errors become failure facts, never authorization verdicts + _warn(source, exception=True) + return replace(facts, io_failed=True) + return await _read(updated, updated_source) + + +async def _read_membership(facts: FallbackCandidateFacts, source: _Sources) -> FallbackCandidateFacts: + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_checks import get_team_membership + + if source.token.user_id is None or source.token.team_id is None: + raise NativeAuthRuntimeError("candidate membership read requested without identity") + dependencies: Final = _Dependencies( + source.prisma if source.prisma is not None else proxy_server.prisma_client, + source.cache if source.cache is not None else proxy_server.user_api_key_cache, + source.logging if source.logging is not None else proxy_server.proxy_logging_obj, + ) + try: + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.token.team_id, + prisma_client=dependencies.prisma, + user_api_key_cache=dependencies.cache, + proxy_logging_obj=dependencies.logging, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve silent membership IO failure as a native candidate fact + return replace(facts, membership="Failed") + return await _read( + replace(facts, membership="Loaded"), replace(source, dependencies=dependencies, membership=membership) + ) + + +async def _effect( + decision: FallbackCandidateDecision, facts: FallbackCandidateFacts, source: _Sources +) -> tuple[FallbackCandidateFacts, _Sources]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check, team_model_aliases_for_auth_check + + router: Final = source.router + dependencies: Final = source.dependencies + if router is None or dependencies is None: + raise NativeAuthRuntimeError("candidate read requested without router or dependencies") + match decision.kind: + case "ReadRoute": + return replace(facts, route=route_context(source.route)), source + case "ReadUnpriced": + return replace(facts, unpriced=read_unpriced_model_facts(source.target.served_model, router)), source + case "ReadCustomerGrants": + if source.customer is None: + raise NativeAuthRuntimeError("candidate customer grants requested without source") + return replace(facts, customer_grants=object_model_grants(source.customer)), source + case "ReadCustomer": + if facts.customer_grants is None: + raise NativeAuthRuntimeError("candidate customer scope requested before grants") + customer: Final = project_model_facts( + source.target.model, + router, + facts.customer_grants, + team_model_aliases_for_auth_check(source.token), + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "customer", + ) + return replace(facts, customer=customer), source + case "ReadPrice": + return replace(facts, price=read_model_cost_facts(source.target.served_model, router)), source + case "ReadBudgetModels": + models: Final = ( + source.target.requested, + source.target.model, + source.target.served_model, + router.registry().resolve(source.target.served_model), + ) + return replace(facts, budget_models=models), source + case "ReadBudgetScopes": + budgets: Final = _budget_sources(source.token) + scopes: Final = tuple( + CandidateBudgetScope(budget.entity, budget.identity is not None, bool(budget.budgets)) + for budget in budgets + ) + return replace(facts, budget_scopes=scopes), replace(source, budgets=budgets) + case "ReadBudget": + if facts.budget_models is None: + raise NativeAuthRuntimeError("candidate budget read requested without model names") + budget: Final = next(budget for budget in source.budgets if budget.entity == decision.entity) + key_budget: Final = next(budget.budgets for budget in source.budgets if budget.entity == "Key") + reading: Final = await source.limiter.read_model_budget_raw( + decision.entity, + budget.identity, + source.token.team_id, + facts.budget_models[decision.model_index], + budget.budgets, + key_budget if decision.entity == "Team" else None, + ) + answer: Final = CandidateBudgetRead(decision.model_index, decision.entity, reading.snapshot) + return replace(facts, budgets=(*facts.budgets, answer)), source + case "ReadAuthorization": + authorization: Final = await load_fallback_authorization_facts( + source.target.model, + source.token, + router, + source.prisma, + source.membership, + source.project, + True, + source.team, + ) + return replace(facts, authorization=authorization), source + case "ResetAttribution": + source.token.matched_model_access_groups = [] # rebind-ok: preserve native-selected request-local reset before attribution IO + return replace(facts, attribution_reset=True), source + case "ReadAttribution": + return replace(facts, attribution=await _attribution(source, dependencies, router)), source + case "StampAttribution": + source.token.matched_model_access_groups = list(decision.groups) # rebind-ok: native attribution effect + return replace(facts, attribution_stamped=True), source + case "ReadGroupBudgets": + readings: Final = await read_access_group_budget_facts( + decision.groups, dependencies.prisma, dependencies.cache + ) + return replace( + facts, group_budgets_read=True, group_budgets=tuple(reading.facts for reading in readings) + ), source + case "Complete" | "ReadMembership": + raise NativeAuthRuntimeError("unexpected candidate data read") + + +async def _attribution(source: _Sources, dependencies: _Dependencies, router: Router) -> ModelGroupAttributionFacts: + try: + return await read_model_group_attribution_facts( + source.target.model, + source.token, + source.team, + source.project, + router, + dependencies.prisma, + dependencies.cache, + dependencies.logging, + source.membership, + True, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve attribution telemetry IO failure separately from candidate IO failure + verbose_proxy_logger.debug("model access group attribution failed: %s", error) + return ModelGroupAttributionFacts(io_failed=True) diff --git a/litellm/proxy/auth/key_budget_context_adapter.py b/litellm/proxy/auth/key_budget_context_adapter.py new file mode 100644 index 00000000000..3f7d245704b --- /dev/null +++ b/litellm/proxy/auth/key_budget_context_adapter.py @@ -0,0 +1,95 @@ +from collections.abc import Mapping +from dataclasses import replace +from datetime import datetime, timezone +from typing import Final, assert_never + +import litellm +from litellm._auth import budget_context as core +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.budget_context_adapter import mapping_value + + +def numeric_facts(value: object) -> core.NumericFacts: + return core.NumericFacts(float(value) if isinstance(value, (int, float)) else None, isinstance(value, bool)) + + +def budget_throttle_percentage() -> float | None: + return core.percentage(numeric_facts(litellm.budget_exceeded_throttle_percentage)) + + +def read_key_throttle_facts(valid_token: UserAPIKeyAuth) -> core.ThrottleFacts: + metadata: Final = mapping_value(getattr(valid_token, "metadata", None)) + facts: Final = core.ThrottleFacts( + (metadata or {}).get("throttle_on_budget_exceeded") is True, + valid_token.tpm_limit, + valid_token.rpm_limit, + ) + decision: Final = core.throttle(facts) + return ( + replace(facts, percentage=numeric_facts(litellm.budget_exceeded_throttle_percentage)) + if decision.read_percentage + else facts + ) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + return core.throttle(read_key_throttle_facts(valid_token)).percentage is not None + + +def inspect_key_budget_throttle(valid_token: UserAPIKeyAuth) -> float | None: + decision: Final = core.throttle(read_key_throttle_facts(valid_token)) + return budget_throttle_percentage() if decision.percentage is not None else None + + +def _microseconds(value: datetime) -> int: + delta: Final = value - datetime(1970, 1, 1, tzinfo=timezone.utc) + return (delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds + + +def _temporary( + token: UserAPIKeyAuth, + facts: core.TemporaryFacts, + metadata: Mapping[str, object] | None, +) -> UserAPIKeyAuth: + decision: Final = core.temporary(facts) + if decision.invalid_increase_type is not None: + raise TypeError( + f"unsupported operand type(s) for +: '{type(token.max_budget).__name__}' and '{decision.invalid_increase_type}'" + ) + if decision.copy: + return token.model_copy(update={"max_budget": decision.max_budget}) + match decision.read: + case None: + return token + case "Metadata": + raw: Final = mapping_value(getattr(token, "metadata", None)) + return _temporary( + token, + replace( + facts, + metadata_read=True, + has_increase=raw is not None and "temp_budget_increase" in raw, + has_expiry=raw is not None and "temp_budget_expiry" in raw, + ), + raw, + ) + case "Expiry": + assert metadata is not None + expiry: Final = datetime.fromisoformat(metadata["temp_budget_expiry"]) # pyright: ignore[reportArgumentType] # preserve datetime's public TypeError for malformed metadata + aware: Final = expiry.replace(tzinfo=timezone.utc) if expiry.tzinfo is None else expiry + return _temporary(token, replace(facts, expires_at=_microseconds(aware)), metadata) + case "Now": + return _temporary(token, replace(facts, now=_microseconds(datetime.now(timezone.utc))), metadata) + case "Increase": + assert metadata is not None + value: Final = metadata["temp_budget_increase"] + raw_increase: Final = core.IncreaseFacts( + float(value) if isinstance(value, (int, float)) else None, bool(value), type(value).__name__ + ) + return _temporary(token, replace(facts, increase=raw_increase), metadata) + case _: + assert_never(decision.read) + + +def update_key_budget_with_temp_budget_increase(valid_token: UserAPIKeyAuth) -> UserAPIKeyAuth: + return _temporary(valid_token, core.TemporaryFacts(valid_token.max_budget), None) diff --git a/litellm/proxy/auth/key_model_access_adapter.py b/litellm/proxy/auth/key_model_access_adapter.py new file mode 100644 index 00000000000..eb79c6d93eb --- /dev/null +++ b/litellm/proxy/auth/key_model_access_adapter.py @@ -0,0 +1,159 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from fastapi import Request +from pydantic import BaseModel, Field + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +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.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 + + +class _Config(BaseModel, frozen=True): + config: Mapping[str, object] = Field(default_factory=dict[str, object]) + + +@dataclass(frozen=True, slots=True) +class _Sources: + token: UserAPIKeyAuth + request_data: dict[str, object] # mutable-ok: alias resolution mutates the caller's request + route: str + request: Request | None + router: Router | None + + +async def read_key_and_fallback_model_access_facts( + *, + valid_token: UserAPIKeyAuth, + request_data: dict[str, object], # mutable-ok: preserve request-local alias mutations + route: str, + request: Request | None, + llm_model_list: Sequence[object] | None, + llm_router: Router | None, +) -> 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) + facts: Final = KeyAndFallbackModelAccessFacts( + config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config), + key_models=object_model_grants(valid_token), + ) + return await _read(facts, _Sources(valid_token, request_data, route, request, llm_router), ()) + + +def _scope(model: str | Sequence[str], source: _Sources) -> GroupedModelScope: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + return GroupedModelScope( + project_model_facts( + model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ), + tuple(source.token.access_group_ids or ()), + ) + + +async def _groups(scope: GroupedModelScope, source: _Sources) -> GroupedModelScope: + return await read_grouped_scope(scope, source.router, merge=False) + + +async def _route(name: str, source: _Sources) -> Exception | None: + from litellm.proxy.route_llm_request import ( + route_request, # pyright: ignore[reportUnknownVariableType] # preserve the existing routing probe and exception identity + ) + + try: + await route_request( + data={"model": name, "messages": [{"role": "user", "content": "Who was Alexander?"}]}, + llm_router=source.router, + user_model=None, + route_type="acompletion", + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # native policy sequences a legacy routing IO failure before rethrowing + return error + return None + + +async def _read( + facts: KeyAndFallbackModelAccessFacts, source: _Sources, route_errors: tuple[Exception | None, ...] +) -> KeyAndFallbackModelAccessFacts: + from litellm.proxy.auth.auth_utils import fallback_target_model_name, iter_request_fallback_targets + from litellm.proxy.auth.user_api_key_auth import ( + _get_model_from_request_context, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # read the existing mutable request model + ) + + decision: Final = key_and_fallback_model_access(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadPrimary": + model: Final = _get_model_from_request_context( + request_data=source.request_data, + route=source.route, + request=source.request, + llm_router=source.router, + team_id=source.token.team_id, + ) + return await _read( + replace(facts, primary_loaded=True, primary=_scope(model, source) if model is not None else None), + source, + route_errors, + ) + case "ReadPrimaryGroups": + if facts.primary is None: + raise NativeAuthRuntimeError("missing primary model scope") + return await _read(replace(facts, primary=await _groups(facts.primary, source)), source, route_errors) + case "ReadFallbackNames": + names: Final = tuple( + dict.fromkeys( + name + for target in iter_request_fallback_targets(source.request_data) + if (name := fallback_target_model_name(target)) is not None + ) + ) + return await _read(replace(facts, fallback_names=names), source, route_errors) + case "ReadFallbackScope": + if facts.fallback_names is None: + raise NativeAuthRuntimeError("missing fallback names") + scope: Final = _scope(facts.fallback_names[decision.index], source) + return await _read(replace(facts, fallback_scopes=(*facts.fallback_scopes, scope)), source, route_errors) + case "ReadFallbackGroups": + grouped: Final = await _groups(facts.fallback_scopes[decision.index], source) + scopes: Final = tuple( + grouped if index == decision.index else scope for index, scope in enumerate(facts.fallback_scopes) + ) + return await _read(replace(facts, fallback_scopes=scopes), source, route_errors) + case "RouteFallback": + if facts.fallback_names is None: + raise NativeAuthRuntimeError("missing fallback names") + error: Final = await _route(facts.fallback_names[decision.index], source) + return await _read( + replace(facts, routing=(*facts.routing, "Routed" if error is None else "Failed")), + source, + (*route_errors, error), + ) + case "RoutingFailure": + failure: Final = route_errors[decision.index] + if failure is None: + raise NativeAuthRuntimeError("missing fallback routing failure") + raise failure diff --git a/litellm/proxy/auth/managed_model_adapter.py b/litellm/proxy/auth/managed_model_adapter.py new file mode 100644 index 00000000000..069e739615a --- /dev/null +++ b/litellm/proxy/auth/managed_model_adapter.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.managed_models import ( + AgentCeilingSourceFacts, + ManagedModelFacts, + agent_ceiling_source, + managed_model_grants, + managed_models, +) +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, + resolve_agent_access_group_ceiling, + resolve_managed_agent_ceilings, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import ( + agent_caller_auth, + load_agent_caller_team, + load_agent_caller_user, +) +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy +from litellm.proxy.auth.member_model_adapter import read_member_model_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.proxy.auth.team_alias_adapter import resolve_team_alias +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader +from litellm.types.agents import AgentResponse + +CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] +CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] +CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_UserTable | None]] +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] | None + token: UserAPIKeyAuth | None + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + ceiling: CeilingResolver + team_loader: CallerTeamLoader + user_loader: CallerUserLoader + managed_grants: tuple[object, ...] = () + dispatched: str | Sequence[str] = () + caller: UserAPIKeyAuth | None = None + team: LiteLLM_TeamTable | None = None + user: LiteLLM_UserTable | None = None + key_aliases: Mapping[str, str] | None = None + + +async def agent_access_group_ceilings( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> tuple[AgentAccessGroupCeiling, ...]: + return await _ceilings(AgentCeilingSourceFacts(bool(valid_token.agent_id)), valid_token, resolve_ceiling) + + +async def _ceilings( + facts: AgentCeilingSourceFacts, + token: UserAPIKeyAuth, + resolver: CeilingResolver, + managed: AgentResponse | None = None, + rows: tuple[AgentAccessGroupCeiling, ...] = (), +) -> tuple[AgentAccessGroupCeiling, ...]: + decision: Final = agent_ceiling_source(facts) + match decision.kind: + case "Complete": + return tuple(rows[index] for index in decision.indices) + case "ReadManaged": + policy: Final = managed_agent_policy(token) + return await _ceilings( + replace(facts, managed_loaded=True, managed_present=policy is not None), token, resolver, policy + ) + case "ReadManagedCeilings": + if managed is None: + raise NativeAuthRuntimeError("missing managed agent ceiling source") + loaded: Final = await resolve_managed_agent_ceilings(managed) + return await _ceilings(replace(facts, loaded=True, count=len(loaded)), token, resolver, managed, loaded) + case "ReadUnmanagedCeiling": + if token.agent_id is None: + raise NativeAuthRuntimeError("missing agent ceiling identity") + row: Final = await resolver(token.agent_id) + projected: Final = (row,) if row is not None else () + return await _ceilings( + replace(facts, loaded=True, count=len(projected)), token, resolver, managed, projected + ) + + +async def read_managed_model_facts( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + purpose: Literal["All", "Ceilings", "Caller"] = "All", + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> ManagedModelFacts: + facts, _source = await _read( + ManagedModelFacts(purpose, bool(model), valid_token is not None), + _Source( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + resolve_ceiling, + load_team, + load_user, + ), + ) + return facts + + +async def can_agent_call_model( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + purpose: Literal["All", "Ceilings", "Caller"] = "All", + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + facts, source = await _read( + ManagedModelFacts(purpose, bool(model), valid_token is not None), + _Source( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + resolve_ceiling, + load_team, + load_user, + ), + ) + decision: Final = managed_models(facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None or model is None: + raise NativeAuthRuntimeError("incomplete managed model authorization") + raise_model_denial( + decision.denial, + model, + user_id=source.caller.user_id + if decision.member_models is not None and source.caller is not None + else valid_token.agent_id + if valid_token is not None + else None, + team_id=source.caller.team_id if source.caller is not None else None, + member_models=decision.member_models, + ) + + +@native_boundary +def _managed_grants(raw: tuple[object, ...]) -> tuple[str, ...]: + return _NAMES.validate_python(raw) + + +def _declared_grants(token: UserAPIKeyAuth) -> tuple[str | None, ...] | None: + managed: Final = managed_agent_policy(token) + raw: Final = (managed.object_permission or {}).get("models") if managed is not None else None + if not isinstance(raw, (list, tuple)): + return None + return tuple(value if isinstance(value, str) else None for value in _ITEMS.validate_python(raw)) + + +def managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: + return managed_model_grants(_declared_grants(valid_token), ()) + + +async def agent_model_grants( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> frozenset[str]: + ceilings: Final = await agent_access_group_ceilings(valid_token, resolve_ceiling) + return managed_model_grants(_declared_grants(valid_token), tuple(tuple(ceiling.models) for ceiling in ceilings)) + + +async def check_agent_access_group_model_access( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> Literal[True]: + await can_agent_call_model( + model, valid_token, llm_router, None, None, None, purpose="Ceilings", resolve_ceiling=resolve_ceiling + ) + return True + + +async def check_agent_caller_model_access( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + await can_agent_call_model( + model, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + purpose="Caller", + load_team=load_team, + load_user=load_user, + ) + + +def _model(source: _Source) -> str | Sequence[str]: + if source.model is None: + raise NativeAuthRuntimeError("managed model read without a request model") + return source.model + + +def _token(source: _Source) -> UserAPIKeyAuth: + if source.token is None: + raise NativeAuthRuntimeError("managed model read without a token") + return source.token + + +def _dispatch(source: _Source) -> str | Sequence[str]: + from litellm.proxy.auth.auth_checks import team_model_aliases_for_auth_check + + token: Final = _token(source) + return resolve_team_alias(_model(source), team_model_aliases_for_auth_check(token), token.team_id, source.router) + + +async def _read(facts: ManagedModelFacts, source: _Source) -> tuple[ManagedModelFacts, _Source]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + decision: Final = managed_models(facts) + match decision.kind: + case "Allow" | "Deny": + return facts, source + case "ReadManaged": + policy: Final = managed_agent_policy(source.token) + raw: Final = (policy.object_permission or {}).get("models", ()) if policy is not None else () + items: Final = _ITEMS.validate_python(raw) if isinstance(raw, (list, tuple)) else () + return await _read( + replace( + facts, + managed_loaded=True, + managed_present=policy is not None, + managed_models_valid=isinstance(raw, (list, tuple)), + managed_models_count=len(items), + ), + replace(source, managed_grants=items), + ) + case "ReadManagedScope": + dispatched: Final = _dispatch(source) + managed: Final = project_model_facts( + dispatched, + source.router, + _managed_grants(source.managed_grants), + team_id=_token(source).team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="agent", + ) + return await _read(replace(facts, managed=managed), source) + case "ReadCeilings": + ceilings: Final = await agent_access_group_ceilings(_token(source), source.ceiling) + return await _read( + replace(facts, ceilings=tuple(tuple(sorted(ceiling.models)) for ceiling in ceilings)), source + ) + case "ReadAliases": + return await _read(replace(facts, aliases_loaded=True), replace(source, dispatched=_dispatch(source))) + case "ReadCeilingScope": + if facts.ceilings is None: + raise NativeAuthRuntimeError("managed model scope without ceilings") + ceiling_scope: Final = project_model_facts( + source.dispatched, + source.router, + facts.ceilings[decision.index], + team_id=_token(source).team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="agent", + ) + return await _read(replace(facts, scopes=(*facts.scopes, ceiling_scope)), source) + case "ReadCaller": + caller: Final = agent_caller_auth(_token(source)) + return await _read( + replace(facts, caller_loaded=True, caller_present=caller is not None), replace(source, caller=caller) + ) + case "ReadTeam": + team: Final = await source.team_loader(_token(source)) + aliases: Final = key_model_aliases_for_auth_check(source.token) + return await _read( + replace(facts, team_loaded=True, team_present=team is not None), + replace(source, team=team, key_aliases=aliases), + ) + case "ReadTeamScope": + if source.team is None: + raise NativeAuthRuntimeError("missing caller team for model scope") + scope: Final = project_model_facts( + _model(source), + source.router, + object_model_grants(source.team), + team_id=source.team.team_id, + key_model_aliases=source.key_aliases, + object_type="team", + ) + return await _read( + replace(facts, team=GroupedModelScope(scope, tuple(source.team.access_group_ids or ()))), source + ) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing caller team group scope") + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=grouped), source) + case "ReadMember": + if source.team is None or source.caller is None or source.cache is None or source.logging is None: + raise NativeAuthRuntimeError("missing caller member model context") + member: Final = await read_member_model_facts( + _model(source), + source.team, + source.caller, + source.router, + source.prisma, + source.cache, + source.logging, + key_model_aliases=source.key_aliases, + ) + return await _read(replace(facts, member=member), source) + case "ReadUser": + user: Final = await source.user_loader(_token(source)) + return await _read( + replace( + facts, + user_loaded=True, + user_present=user is not None, + user_grants=object_model_grants(user) if user is not None else (), + ), + replace(source, user=user), + ) + case "ReadUserScope": + user_scope: Final = project_model_facts( + _model(source), source.router, facts.user_grants, key_model_aliases=source.key_aliases + ) + return await _read(replace(facts, user=user_scope), source) diff --git a/litellm/proxy/auth/member_model_adapter.py b/litellm/proxy/auth/member_model_adapter.py new file mode 100644 index 00000000000..08c2af532d6 --- /dev/null +++ b/litellm/proxy/auth/member_model_adapter.py @@ -0,0 +1,132 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.member_models import MemberModelFacts, member_models +from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import project_model_facts, raise_model_denial +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + team: LiteLLM_TeamTable + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + membership: LiteLLM_TeamMembership | None + aliases: Mapping[str, str] | None + + +async def read_member_model_facts( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, + key_model_aliases: Mapping[str, str] | None = None, +) -> MemberModelFacts: + facts: Final = MemberModelFacts( + user_present=valid_token.user_id is not None, + team_present=team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # preserve unvalidated cache row identity guard + membership_loaded=team_membership_loaded, + ) + return await _read( + facts, + _Source( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + key_model_aliases, + ), + ) + + +async def check_team_member_model_access( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, + key_model_aliases: Mapping[str, str] | None = None, +) -> None: + facts: Final = await read_member_model_facts( + model, + team_object, + valid_token, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + key_model_aliases, + ) + decision: Final = member_models(facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete member model access facts") + raise_model_denial( + decision.denial, + model, + user_id=valid_token.user_id, + team_id=team_object.team_id, + member_models=decision.member_models, + ) + + +async def _read(facts: MemberModelFacts, source: _Source) -> MemberModelFacts: + from litellm.proxy.auth.auth_checks import get_team_membership + + decision: Final = member_models(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("member model lookup without user identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + return await _read(replace(facts, membership_loaded=True), replace(source, membership=membership)) + case "ReadGrants": + row: Final = source.membership.litellm_budget_table if source.membership is not None else None + return await _read( + replace(facts, grants=tuple(row.allowed_models or ()) if row is not None else ()), source + ) + case "ReadScope": + if facts.grants is None: + raise NativeAuthRuntimeError("member model scope read without grants") + scope: Final = project_model_facts( + source.model, + source.router, + facts.grants, + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return await _read(replace(facts, scope=scope), source) diff --git a/litellm/proxy/auth/model_access_adapter.py b/litellm/proxy/auth/model_access_adapter.py new file mode 100644 index 00000000000..f5322635870 --- /dev/null +++ b/litellm/proxy/auth/model_access_adapter.py @@ -0,0 +1,310 @@ +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from fastapi import HTTPException +from pydantic import BaseModel + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ( + ModelAccessDecision, + ModelAccessFacts, + ModelAliases, + ModelObject, + ModelRegistry, + grouped_model_access, + key_model_grants, + model_access, + user_model_precheck, +) +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + ModelAccessDeniedProxyException, + ProxyErrorTypes, + UserAPIKeyAuth, +) +from litellm.proxy.auth.model_access_denied import ( + customer_model_access_denied_client_message, + model_access_denied_client_message, +) +from litellm.proxy.auth.model_projection_adapter import project_scope, read_grouped_scope +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +class ModelLists(BaseModel, frozen=True): + models: tuple[str, ...] = () + + +@native_boundary +def object_model_grants(source: object) -> tuple[str, ...]: + return ModelLists.model_validate(source, from_attributes=True).models + + +@native_boundary +def _initial_model_facts( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> ModelAccessFacts: + requested: Final = (model,) if isinstance(model, str) else tuple(model) + aliases: Final = ModelAliases( + team=dict(team_model_aliases or {}), + key=dict(key_model_aliases or {}), + global_=dict(litellm.model_alias_map), + team_order=tuple(team_model_aliases or ()), + ) + registry: Final = ModelRegistry( + present=llm_router is not None, + helper_present=bool(llm_router), + evaluations=(), + ) + initial: Final = ModelAccessFacts( + requested=requested, + grants=tuple(models), + team_id=team_id, + aliases=aliases, + registry=registry, + object_type=object_type, + fallback_depth=fallback_depth, + max_depth=DEFAULT_MAX_RECURSE_DEPTH, + is_list=not isinstance(model, str), + ) + return initial + + +def project_model_facts( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> ModelAccessFacts: + initial: Final = _initial_model_facts( + model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type, fallback_depth + ) + return project_scope(initial, llm_router) + + +def raise_model_denial( + decision: ModelAccessDecision, + requested: str | Sequence[str], + *, + user_id: str | None = None, + team_id: str | None = None, + teams: Sequence[str] | None = None, + member_models: Sequence[str] | None = None, +) -> Literal[True]: + model: Final = requested if isinstance(requested, str) else list(requested) + match decision.kind: + case "Allow": + return True + case "Depth": + raise Exception( # noqa: TRY002 # preserve the legacy depth failure type and envelope + f"Unable to parse model, max fallback depth exceeded - received model: {decision.model if decision.model is not None else model}" + ) + case "NoAgentModels": + raise HTTPException(403, "This agent has no model grants") + case "ReadGroups" | "ReadScope" | "ReadGroupScope": + raise NativeAuthRuntimeError("model grant lookup was not answered") + case "Deny" | "NoDefaultModels" | "EmptyAgentCeiling" | "PersonalTeam": + error_model: Final = decision.model if decision.model is not None else model + object_type: Final = decision.object_type + message: Final = ( + customer_model_access_denied_client_message(error_model) + if object_type == "customer" + else model_access_denied_client_message(error_model) + ) + internal: Final = ( + f"Team member not allowed to access model. User={user_id}, Team={team_id}, Model={model}. Allowed member models = {list(member_models)}" + if member_models is not None + else ( + f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}" + if decision.kind == "NoDefaultModels" + else ( + f"agent {user_id} access groups grant no models" + if decision.kind == "EmptyAgentCeiling" + else ( + f"Personal key model access is derived from the owner's teams. User={user_id}, teams={list(teams) if teams is not None else None}, model={error_model}" + if decision.kind == "PersonalTeam" + else ( + f"{object_type} not allowed to access model. This {object_type} can only access models={list(decision.grants)}. Tried to access {error_model}" + ) + ) + ) + ) + ) + error_type: Final = ( + ProxyErrorTypes.team_model_access_denied + if member_models is not None + else ( + ProxyErrorTypes.key_model_access_denied + if decision.kind in ("NoDefaultModels", "PersonalTeam") + else ( + ProxyErrorTypes.agent_model_access_denied + if decision.kind == "EmptyAgentCeiling" + else ProxyErrorTypes.get_model_access_error_type_for_object(object_type) + ) + ) + ) + raise ModelAccessDeniedProxyException( + message=message, internal_message=internal, type=error_type, param="model", code=403 + ) + + +def can_object_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: ModelObject = "user", + fallback_depth: int = 0, +) -> Literal[True]: + facts: Final = project_model_facts( + model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type, fallback_depth + ) + return raise_model_denial(model_access(facts), model) + + +async def _grouped_model_access( + facts: ModelAccessFacts, + group_ids: Sequence[str], + prisma_client: DatabaseClient | None, + *, + merge: bool, + router: RegistryReader | None = None, +) -> ModelAccessDecision: + scope: Final = await read_grouped_scope( + GroupedModelScope(facts, tuple(group_ids)), router, prisma_client, merge=merge + ) + return grouped_model_access(scope.facts, scope.group_ids, scope.group_models, merge=merge) + + +async def can_key_call_model( + model: str | Sequence[str], + llm_model_list: Sequence[object] | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + grants: Final = key_model_grants( + object_model_grants(valid_token), valid_token.team_id, tuple(valid_token.team_models or ()) + ) + facts: Final = project_model_facts( + model, + llm_router, + grants, + valid_token.team_model_aliases, + valid_token.team_id, + key_model_aliases_for_auth_check(valid_token), + "key", + ) + decision: Final = await _grouped_model_access( + facts, valid_token.access_group_ids or (), prisma_client, merge=False, router=llm_router + ) + return raise_model_denial(decision, model) + + +async def can_team_access_model( + model: str | Sequence[str], + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + facts: Final = project_model_facts( + model, + llm_router, + team_object.models if team_object is not None else (), + team_model_aliases, + team_object.team_id if team_object is not None else None, + key_model_aliases, + "team", + ) + decision: Final = await _grouped_model_access( + facts, + (team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client, + merge=True, + router=llm_router, + ) + return raise_model_denial(decision, model) + + +def can_org_access_model( + model: str, + org_object: LiteLLM_OrganizationTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, +) -> Literal[True]: + return can_object_call_model( + model, llm_router, org_object.models if org_object is not None else (), team_model_aliases, object_type="org" + ) + + +def can_project_access_model( + model: str | Sequence[str], + project_object: LiteLLM_ProjectTable, + llm_router: RegistryReader | None, + key_model_aliases: Mapping[str, str] | None = None, +) -> Literal[True]: + return can_object_call_model( + model, + llm_router, + project_object.models, + team_id=project_object.team_id, + key_model_aliases=key_model_aliases, + object_type="project", + ) + + +def can_customer_access_model( + model: str | Sequence[str], + end_user_object: LiteLLM_EndUserTable, + llm_router: RegistryReader | None, + valid_token: UserAPIKeyAuth | None, +) -> Literal[True]: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check, team_model_aliases_for_auth_check + + facts: Final = project_model_facts( + model, + llm_router, + end_user_object.models or (), + team_model_aliases_for_auth_check(valid_token) if valid_token is not None else None, + valid_token.team_id if valid_token is not None else None, + key_model_aliases_for_auth_check(valid_token), + "customer", + ) + return raise_model_denial(model_access(facts, kind="Customer"), model) + + +async def can_user_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + user_object: LiteLLM_UserTable | None, + key_model_aliases: Mapping[str, str] | None = None, +) -> Literal[True]: + grants: Final = object_model_grants(user_object) if user_object is not None else () + precheck: Final = user_model_precheck(present=user_object is not None, grants=grants) + if precheck.kind != "ReadScope": + return raise_model_denial(precheck, model) + return can_object_call_model(model, llm_router, grants, key_model_aliases=key_model_aliases) diff --git a/litellm/proxy/auth/model_attribution_adapter.py b/litellm/proxy/auth/model_attribution_adapter.py new file mode 100644 index 00000000000..f7ea8e588e5 --- /dev/null +++ b/litellm/proxy/auth/model_attribution_adapter.py @@ -0,0 +1,213 @@ +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.model_access import key_model_grants +from litellm._auth.model_attribution import ModelGroupAttributionFacts, model_group_attribution +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_ProjectTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Sources: + model: tuple[str, ...] + token: UserAPIKeyAuth + team: LiteLLM_TeamTable | None + project: LiteLLM_ProjectTable | None + router: RegistryReader + prisma: PrismaClient + cache: UserApiKeyCache + logging: ProxyLogging + membership: LiteLLM_TeamMembership | None + + +async def read_model_group_attribution_facts( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTable | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> ModelGroupAttributionFacts: + facts: Final = ModelGroupAttributionFacts( + model_present=model is not None, + token_present=valid_token is not None, + router_present=llm_router is not None, + database_present=prisma_client is not None, + team_present=team_object is not None, + user_present=valid_token is not None and valid_token.user_id is not None, + membership_loaded=team_membership_loaded, + token_org=valid_token.org_id if valid_token is not None else None, + team_org=team_object.organization_id if team_object is not None else None, + ) + if model_group_attribution(facts).kind == "Complete": + return facts + if model is None or valid_token is None or llm_router is None or prisma_client is None: + raise NativeAuthRuntimeError("attribution read requested without its source") + return await _read( + facts, + _Sources( + (model,) if isinstance(model, str) else tuple(model), + valid_token, + team_object, + project_object, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + ), + ) + + +def _covering_groups(source: _Sources) -> Iterator[str]: + for model in source.model: + yield from source.router.registry().access_groups(model_name=model, team_id=source.token.team_id) + + +async def _read(facts: ModelGroupAttributionFacts, source: _Sources) -> ModelGroupAttributionFacts: + from litellm.proxy.auth.auth_checks import ( + _load_model_access_group_registry, # pyright: ignore[reportPrivateUsage] # reuse the existing registry IO without its attribution policy + ) + + decision: Final = model_group_attribution(facts) + match decision.kind: + case "Complete": + return facts + case "ReadRegistry": + registry: Final = await _load_model_access_group_registry(source.prisma, source.cache) + return await _read( + replace(facts, registry_loaded=True, registry=tuple(registry) if registry is not None else None), source + ) + case "ReadCovering": + return await _read(replace(facts, covering=tuple(_covering_groups(source))), source) + case "ReadKey": + key: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + return await _read(replace(facts, grants=(*facts.grants, key)), source) + case "ReadTeam": + team: Final = object_model_grants(source.team) if source.team is not None else () + return await _read(replace(facts, grants=(*facts.grants, team)), source) + case "ReadMember": + membership: Final = await _read_membership(source) if decision.lookup else source.membership + member: Final = ( + tuple(membership.litellm_budget_table.allowed_models or ()) + if decision.include and membership is not None and membership.litellm_budget_table is not None + else () + ) + return await _read(replace(facts, grants=(*facts.grants, member)), source) + case "ReadProject": + project: Final = object_model_grants(source.project) if source.project is not None else () + return await _read(replace(facts, grants=(*facts.grants, project)), source) + case "ReadOrg": + org, failed = await _org_grants(decision.id, source) + return await _read(replace(facts, grants=(*facts.grants, org), org_failed=failed), source) + + +async def _read_membership(source: _Sources) -> LiteLLM_TeamMembership | None: + from litellm.proxy.auth.auth_checks import get_team_membership + + if source.team is None or source.token.user_id is None: + raise NativeAuthRuntimeError("member attribution lookup requested without identity") + return await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + + +async def _org_grants(org_id: str | None, source: _Sources) -> tuple[tuple[str, ...], bool]: + from litellm.proxy.auth.auth_checks import get_org_object + + if org_id is None: + return (), False + try: + org: Final = await get_org_object( + org_id=org_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve attribution's org IO failure behavior, never swallow native failure + verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", error) + return (), True + return (tuple(org.models) if org is not None else ()), False + + +async def collect_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTable | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> tuple[str, ...]: + facts: Final = await read_model_group_attribution_facts( + model, + valid_token, + team_object, + project_object, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + ) + return model_group_attribution(facts).groups + + +async def stamp_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTable | None, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +) -> tuple[str, ...]: + try: + facts: Final = await read_model_group_attribution_facts( + model, + valid_token, + team_object, + project_object, + llm_router, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + team_membership, + team_membership_loaded, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # native attribution decides the legacy telemetry IO failure result + verbose_proxy_logger.debug("model access group attribution failed: %s", error) + return model_group_attribution(ModelGroupAttributionFacts(io_failed=True)).groups + decision: Final = model_group_attribution(facts) + if decision.stamp: + if valid_token is None: + raise NativeAuthRuntimeError("attribution stamping requested without token") + valid_token.matched_model_access_groups = list(decision.groups) # rebind-ok: native attribution effect + return decision.groups diff --git a/litellm/proxy/auth/model_budget_adapter.py b/litellm/proxy/auth/model_budget_adapter.py new file mode 100644 index 00000000000..50d804bb960 --- /dev/null +++ b/litellm/proxy/auth/model_budget_adapter.py @@ -0,0 +1,196 @@ +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from typing import Final + +from pydantic import TypeAdapter + +import litellm +from litellm._auth import model_budget_data +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.model_budget_data import ( + ModelBudgetEntity, + ModelBudgetEntry, + ModelBudgetQueryFacts, + ModelBudgetResolution, + ModelBudgetResolveFacts, + ModelBudgetSnapshot, +) +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows +from litellm.proxy.hooks.model_max_budget_limiter import ModelAccessGroups, model_budget_spend_cache_key +from litellm.router_strategy.budget_limiter import budget_window_seconds +from litellm.types.utils import BudgetConfig + +_FLOAT: Final = TypeAdapter(float) +_OBJECT: Final = TypeAdapter(object) +_ENTITIES: Final[Mapping[ModelBudgetEntity, Litellm_EntityType]] = { + "Key": Litellm_EntityType.KEY, + "User": Litellm_EntityType.USER, + "Team": Litellm_EntityType.TEAM, + "TeamMember": Litellm_EntityType.TEAM_MEMBER, + "EndUser": Litellm_EntityType.END_USER, +} + + +def _entry(name: str, raw: object) -> ModelBudgetEntry: + try: + config: Final = BudgetConfig.model_validate(raw) + seconds: Final = duration_in_seconds(config.budget_duration) if config.budget_duration is not None else None + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # preserve malformed entry fallthrough + return ModelBudgetEntry(name, None, None, None) + return ModelBudgetEntry(name, config.max_budget, config.budget_duration, seconds) + + +def _spend(raw: object) -> float: + try: + return _FLOAT.validate_python(raw or 0.0) + except NativeAuthFailure: + raise + except ValueError: + return 0.0 + + +@dataclass(frozen=True, slots=True) +class ModelBudgetReading: + snapshot: ModelBudgetSnapshot + budget_model: str | None + named_model: bool + + +@dataclass(frozen=True, slots=True) +class ModelBudgetIO: + cache: DualCache + model_access_groups: ModelAccessGroups + shared_windows: SharedWindows + process_started_at: float + clock: Callable[[], float] = time.time + monotonic: Callable[[], float] = time.monotonic + + async def counter(self, key: str) -> float: + redis: Final = self.cache.redis_cache + raw: Final = _OBJECT.validate_python( + await self.cache.async_get_cache(key=key) if redis is None else await redis.async_get_cache(key=key) # pyright: ignore[reportUnknownMemberType] # cache APIs remain legacy IO + ) + return _spend(raw) + + +def resolve_model_budget( + model: str, + budgets: Mapping[str, object], + io: ModelBudgetIO, + team_id: str | None, + key_budgets: Mapping[str, object] | None = None, + *, + team: bool = False, +) -> ModelBudgetResolution: + base: Final = get_bedrock_base_model(model) + raw_price: Final = as_str_mapping(_OBJECT.validate_python(litellm.model_cost.get(base))) or {} # pyright: ignore[reportUnknownMemberType] # project the legacy pricing map value + facts: Final = ModelBudgetResolveFacts( + model=model, + bedrock_base=base, + bedrock_provider=str(raw_price.get("litellm_provider") or ""), + entries=tuple(_entry(name, config) for name, config in budgets.items()), + key_entries=tuple(_entry(name, config) for name, config in (key_budgets or {}).items()), + team=team, + ) + initial: Final = model_budget_data.resolve(facts) + if not initial.read_groups: + return initial + groups: Final = tuple(io.model_access_groups(model, team_id)) + return model_budget_data.resolve(replace(facts, access_groups=groups)) + + +async def read_model_budget( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, +) -> ModelBudgetReading: + resolution: Final = resolve_model_budget(model, budgets or {}, io, team_id, key_budgets, team=entity == "Team") + entry: Final = resolution.entry + if entry is None: + return ModelBudgetReading(ModelBudgetSnapshot(), None, False) + spend_key: Final = model_budget_spend_cache_key(_ENTITIES[entity], entity_id, entry.name, entry.duration) + legacy_key: Final = model_budget_spend_cache_key(_ENTITIES[entity], entity_id, model, entry.duration) + facts: Final = ModelBudgetQueryFacts( + entity=entity, + snapshot=ModelBudgetSnapshot(max_budget=entry.max_budget), + duration_seconds=entry.duration_seconds, + process_uptime=io.monotonic() - io.process_started_at, + same_counter=spend_key == legacy_key, + redis_present=io.cache.redis_cache is not None, + ) + snapshot: Final = await _read_snapshot(io, facts, spend_key, legacy_key, entry.duration) + return ModelBudgetReading(snapshot, entry.name, resolution.named_model) + + +async def _read_snapshot( + io: ModelBudgetIO, + facts: ModelBudgetQueryFacts, + key: str, + legacy_key: str, + duration: str | None, +) -> ModelBudgetSnapshot: + decision: Final = model_budget_data.query(facts) + match decision.read: + case None: + return facts.snapshot + case "Local": + snapshot: Final = replace(facts.snapshot, local_spend=await io.counter(key)) + return await _read_snapshot(io, replace(facts, snapshot=snapshot), key, legacy_key, duration) + case "Legacy": + legacy_snapshot: Final = replace(facts.snapshot, legacy_spend=await io.counter(legacy_key)) + return await _read_snapshot(io, replace(facts, snapshot=legacy_snapshot), key, legacy_key, duration) + case "Shared": + now: Final = io.clock() + shared: Final = ( + await io.shared_windows.read(spend_key=key, budget_duration=duration) if duration is not None else None + ) + shared_snapshot: Final = replace( + facts.snapshot, + shared_spend=shared.spend if shared is not None else None, + shared_started_at=shared.started_at if shared is not None else None, + shared_window_seconds=budget_window_seconds(shared.started_at, duration) + if shared is not None and duration is not None + else None, + now=now, + ) + return await _read_snapshot( + io, replace(facts, snapshot=shared_snapshot, shared_read=True), key, legacy_key, duration + ) + + +async def enforce_model_budget( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + exceeded_message: str, + key_budgets: Mapping[str, object] | None = None, +) -> bool: + reading: Final = await read_model_budget(io, entity, entity_id, team_id, model, budgets, key_budgets) + decision: Final = model_budget_data.check(reading.snapshot) + if decision.denied: + assert reading.snapshot.max_budget is not None + raise litellm.BudgetExceededError( + message=exceeded_message + if reading.named_model + else f"{exceeded_message}, budget_model={reading.budget_model}", + current_cost=decision.current_spend, + max_budget=reading.snapshot.max_budget, + entity_type=_ENTITIES[entity].value, + entity_id=entity_id, + ) + return True diff --git a/litellm/proxy/auth/model_group_adapter.py b/litellm/proxy/auth/model_group_adapter.py new file mode 100644 index 00000000000..c6e0ed4f060 --- /dev/null +++ b/litellm/proxy/auth/model_group_adapter.py @@ -0,0 +1,94 @@ +from dataclasses import dataclass, replace +from typing import Final, Literal, TypeAlias + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.model_groups import GroupLookup, GroupResources, ResourceGroupFacts, resource_group_decision +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth + +ResourceField: TypeAlias = Literal[ + "access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names" +] + + +@dataclass(frozen=True, slots=True) +class GroupRead: + lookup: GroupLookup + error: Exception | None = None + + +async def get_authorized_resources_from_key_access_groups( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + resource_field: ResourceField, + check_db_only: bool = False, + strict_errors: bool = False, +) -> list[str]: # mutable-ok: retain the legacy public return contract + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + facts: Final = ResourceGroupFacts( + subject_present=valid_token is not None, + database_present=prisma_client is not None, + cache_present=_present(user_api_key_cache), + group_ids=tuple(valid_token.access_group_ids or ()) if valid_token is not None else (), + team_id=(valid_token.team_id if valid_token is not None else None) + or (team_object.team_id if team_object is not None else None), + key_id=valid_token.token if valid_token is not None else None, + check_db_only=check_db_only, + strict_errors=strict_errors, + ) + return list(set(await _drive(facts, resource_field, ()))) + + +async def _read(group_id: str, resource_field: ResourceField, check_db_only: bool) -> GroupRead: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + group: Final = await get_access_object( + access_group_id=group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + resources: Final = { + "access_model_names": group.access_model_names, + "access_mcp_server_ids": group.access_mcp_server_ids, + "access_agent_ids": group.access_agent_ids, + "access_skill_names": group.access_skill_names, + }[resource_field] + return GroupRead( + GroupLookup( + "Found", + GroupResources( + tuple(group.assigned_team_ids or ()), tuple(group.assigned_key_ids or ()), tuple(resources or ()) + ), + ) + ) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError as error: + return GroupRead(GroupLookup("Missing"), error) + except Exception as error: # noqa: BLE001 # native strict-error policy decides whether to rethrow this lookup failure + return GroupRead(GroupLookup("Failed"), error) + + +async def _drive( + facts: ResourceGroupFacts, resource_field: ResourceField, reads: tuple[GroupRead, ...] +) -> tuple[str, ...]: + decision: Final = resource_group_decision(facts) + match decision.kind: + case "Complete": + return decision.resources + case "Raise": + error: Final = reads[decision.index].error + if error is None: + raise NativeAuthRuntimeError("missing access group lookup failure") + raise error + case "Read": + result: Final = await _read(facts.group_ids[decision.index], resource_field, facts.check_db_only) + return await _drive(replace(facts, loaded=(*facts.loaded, result.lookup)), resource_field, (*reads, result)) + + +def _present(value: object) -> bool: + return value is not None diff --git a/litellm/proxy/auth/model_helper_adapter.py b/litellm/proxy/auth/model_helper_adapter.py new file mode 100644 index 00000000000..b6b02d27281 --- /dev/null +++ b/litellm/proxy/auth/model_helper_adapter.py @@ -0,0 +1,55 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.model_helpers import ModelHelperFacts, model_helper +from litellm.proxy.auth.model_projection_adapter import provider_facts +from litellm.router_api import RegistryReader + + +def read_model_helper_facts(facts: ModelHelperFacts, router: RegistryReader | None) -> ModelHelperFacts: + if facts.provider_reads is None and not facts.provider_loaded: + return read_model_helper_facts(replace(facts, provider_reads=()), router) + decision: Final = model_helper(facts) + match decision.kind: + case "Complete": + return facts + case "ReadGroups": + if router is None: + raise NativeAuthRuntimeError("model helper registry read without router") + groups: Final = tuple(router.registry().access_groups(model_name=facts.model, team_id=facts.team_id)) + return read_model_helper_facts(replace(facts, groups=groups), router) + case "ReadNames": + if router is None: + raise NativeAuthRuntimeError("model helper names read without router") + return read_model_helper_facts(replace(facts, names=tuple(router.registry().names())), router) + case "ReadProvider": + return read_model_helper_facts( + replace( + facts, + provider_loaded=True, + provider_reads=(*(facts.provider_reads or ()), provider_facts(facts.model)), + ), + router, + ) + + +def check_model_access_helper( + model: str, + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, +) -> bool: + facts: Final = ModelHelperFacts( + "Helper", model, tuple(models), team_id, tuple(team_model_aliases or ()), bool(llm_router) + ) + return model_helper(read_model_helper_facts(facts, llm_router)).allowed + + +def model_in_access_group(model: str, team_models: Sequence[str] | None, llm_router: RegistryReader | None) -> bool: + facts: Final = ModelHelperFacts( + "Group", model, tuple(team_models) if team_models is not None else None, router_present=bool(llm_router) + ) + return model_helper(read_model_helper_facts(facts, llm_router)).allowed diff --git a/litellm/proxy/auth/model_projection_adapter.py b/litellm/proxy/auth/model_projection_adapter.py new file mode 100644 index 00000000000..50141d7fb3c --- /dev/null +++ b/litellm/proxy/auth/model_projection_adapter.py @@ -0,0 +1,121 @@ +from collections.abc import Iterable +from dataclasses import dataclass, replace +from itertools import accumulate, repeat +from typing import Final, assert_never + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ( + ModelAccessFacts, + ModelEvaluation, + ModelProvider, + ModelReading, + ModelReadKind, + grouped_model_access, +) +from litellm._auth.model_projection import ModelProjectionDecision, model_projection +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader +from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent + + +def provider_facts(model: str) -> ModelProvider | None: + try: + stripped, provider, _, _ = litellm.get_llm_provider(model=model) + return ModelProvider(stripped, provider) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # provider inference failure is an absent provider, as in the legacy helper + return None + + +def project_scope( + facts: ModelAccessFacts, router: RegistryReader | None, reads: tuple[ModelReading, ...] = () +) -> ModelAccessFacts: + initial: Final = _Projection(reads, model_projection(facts, reads)) + + def advance(state: _Projection, _tick: None) -> _Projection: + decision: Final = state.decision + if decision.read is None or decision.model is None: + raise NativeAuthRuntimeError("incomplete model registry read") + reading: Final = read_registry(decision.read, decision.model, facts.team_id, router) + updated: Final = (*state.reads, reading) + return _Projection(updated, model_projection(facts, updated)) + + complete: Final = next( + state for state in accumulate(repeat(None), advance, initial=initial) if state.decision.kind == "Complete" + ) + return replace( + facts, + registry=replace( + facts.registry, + evaluations=( + *(facts.registry.evaluations or ()), + ModelEvaluation(facts.grants, complete.reads, facts.aliases), + ), + ), + ) + + +@dataclass(frozen=True, slots=True) +class _Projection: + reads: tuple[ModelReading, ...] + decision: ModelProjectionDecision + + +@native_boundary +def _strings(values: Iterable[str]) -> tuple[str, ...]: + return tuple(values) + + +def read_registry(kind: ModelReadKind, model: str, team_id: str | None, router: RegistryReader | None) -> ModelReading: + if kind == "Compaction": + return ModelReading(kind, model, value=native_compaction_parent(model)) + if kind == "Provider": + return ModelReading(kind, model, provider=provider_facts(model)) + if router is None: + raise NativeAuthRuntimeError("model registry read without router") + match kind: + case "Aliases": + return ModelReading(kind, model, values=_strings(router.registry().aliases())) + case "AliasTarget": + return ModelReading(kind, model, value=router.registry().resolve(model, kind="alias")) + case "Names": + return ModelReading(kind, model, values=_strings(router.registry().names())) + case "Groups": + return ModelReading( + kind, model, values=_strings(router.registry().access_groups(model_name=model, team_id=team_id)) + ) + case "Deployment": + deployment: Final = router.registry().deployments(model, source="first_in_group") + return ModelReading( + kind, + model, + value=deployment.model_info.team_public_model_name if deployment is not None else None, + team=deployment.model_info.team_id if deployment is not None else None, + ) + case _: + assert_never(kind) + + +async def read_grouped_scope( + scope: GroupedModelScope, router: RegistryReader | None, prisma: DatabaseClient | None = None, *, merge: bool +) -> GroupedModelScope: + from litellm.proxy.auth.auth_checks import get_models_from_access_groups + + decision: Final = grouped_model_access(scope.facts, scope.group_ids, scope.group_models, merge=merge) + if decision.kind == "ReadGroups": + models: Final = await get_models_from_access_groups(scope.group_ids, prisma_client=prisma) + return await read_grouped_scope(replace(scope, group_models=tuple(models)), router, prisma, merge=merge) + if decision.kind == "ReadGroupScope": + scoped: Final = project_scope( + replace( + scope.facts, + grants=decision.grants, + aliases=replace(scope.facts.aliases, global_=dict(litellm.model_alias_map)), + ), + router, + ) + return replace(scope, facts=replace(scoped, grants=scope.facts.grants)) + return scope diff --git a/litellm/proxy/auth/org_route_adapter.py b/litellm/proxy/auth/org_route_adapter.py new file mode 100644 index 00000000000..ff008f66552 --- /dev/null +++ b/litellm/proxy/auth/org_route_adapter.py @@ -0,0 +1,130 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.org_routes import ( + AdminFacts, + Identifier, + OrganizationMembership, + OrganizationRouteFacts, + admin, + organization_route, +) +from litellm.proxy._types import LiteLLM_UserTable, LiteLLMRoutes, ProxyErrorTypes, ProxyException + +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) + + +def memberships(user: LiteLLM_UserTable | None) -> tuple[OrganizationMembership, ...] | None: + return ( + tuple(OrganizationMembership(item.organization_id, item.user_role) for item in user.organization_memberships) + if user is not None and user.organization_memberships is not None + else None + ) + + +def user_is_org_admin(request_data: Mapping[str, object], user_object: LiteLLM_UserTable | None = None) -> bool: + facts = AdminFacts(memberships(user_object)) # rebind-ok: native-directed read stage + singular: object = None # rebind-ok: raw request field loaded only after the native directive + plural: tuple[object, ...] = () # rebind-ok: raw request field loaded only after the native directive + while True: + decision = admin(facts) + match decision.action: + case "ReadRequest": + singular = request_data.get("organization_id") + raw_plural = request_data.get("organizations") + plural = _OBJECTS.validate_python(raw_plural) if isinstance(raw_plural, list) else () + facts = replace( + facts, request_read=True, singular_present=singular is not None, plural_count=len(plural) + ) + case "ReadSingular" | "ReadPlural": + raw = singular if decision.action == "ReadSingular" else plural[decision.plural_index] + facts = replace(facts, candidate_index=decision.candidate_index, candidate=_identifier(raw)) + case "Allow": + return True + case "Deny": + return False + case "HashFailure": + raise TypeError(decision.hash_failure) + + +def _identifier(value: object) -> Identifier: + try: + hash(value) + except NativeAuthFailure: + raise + except TypeError as error: + return Identifier(None, str(error)) + return Identifier(value if isinstance(value, str) else None) + + +def organization_role_based_access_check( + request_body: Mapping[str, object], user_object: LiteLLM_UserTable | None, route: str +) -> None: + facts = OrganizationRouteFacts( # rebind-ok: native-directed read driver + route, + user_object is not None, + None, + None, + None, + tuple(LiteLLMRoutes.org_admin_only_routes.value), + request_pending=True, + ) + passed: object = None # rebind-ok: raw field is read only at the native-selected stage + while True: + decision = organization_route(facts) + if decision.read_request: + assert user_object is not None + passed = request_body.get("organization_id") + facts = replace( + facts, + user_role=user_object.user_role, + organization_id=passed if isinstance(passed, str) else None, + organization_id_present=passed is not None, + request_pending=False, + memberships_pending=True, + identifier_pending=True, + ) + elif decision.read_memberships: + facts = replace(facts, memberships=memberships(user_object), memberships_pending=False) + elif decision.read_identifier: + facts = replace(facts, identifier_pending=False, hash_failure=_identifier(passed).hash_failure) + else: + break + if decision.hash_failure is not None: + raise TypeError(decision.hash_failure) + organization_id: Final = passed + role: Final = facts.user_role + messages: Final = { + "CreateOrganization": (f"Only proxy admins can create new organizations. You are {role}", "user_role"), + "NoMemberships": ( + f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.", + "organization_id", + ), + "MissingOrganization": ( + "Passed organization_id is None, please pass an organization_id in your request", + "organization_id", + ), + "MissingRole": ( + f"You do not have a role within the selected organization. Passed organization_id: {organization_id}. Please contact the organization admin to request access.", + "organization_id", + ), + "WrongRole": ( + f"You do not have the required role to perform {route} in Organization {organization_id}. Your role is {decision.organization_role} in Organization {organization_id}", + "user_role", + ), + "TeamMissingOrganization": ( + f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {list(decision.organization_ids)}", + "organization_id", + ), + "TeamWrongRole": ( + f"You do not have the required role to call {route}. Your role is {decision.organization_role} in Organization {organization_id}", + "user_role", + ), + } + if decision.denial is not None: + message, param = messages[decision.denial] + raise ProxyException(message=message, type=ProxyErrorTypes.auth_error.value, param=param, code=401) diff --git a/litellm/proxy/auth/owner_team_adapter.py b/litellm/proxy/auth/owner_team_adapter.py new file mode 100644 index 00000000000..e0ea089269b --- /dev/null +++ b/litellm/proxy/auth/owner_team_adapter.py @@ -0,0 +1,104 @@ +import asyncio +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.owner_teams import OwnerTeamBudget, OwnerTeamFacts, OwnerTeamLoad, owner_teams +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.personal_budget_adapter import read_personal_team_budget_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@dataclass(frozen=True, slots=True) +class OwnerTeamReading: + facts: OwnerTeamFacts + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + token: UserAPIKeyAuth + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + check_db_only: bool + + +async def read_owner_team_facts( + user_object: LiteLLM_UserTable, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + purpose: Literal["Models", "Grants"], + check_db_only: bool = False, +) -> OwnerTeamReading: + facts: Final = OwnerTeamFacts( + purpose, + tuple(user_object.teams), + UI_TEAM_ID, + user_present=valid_token.user_id is not None, + database_present=prisma_client is not None, + org=valid_token.org_id, + ) + return await _read( + OwnerTeamReading(facts), + _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only), + ) + + +async def _load(team_id: str, source: _Source) -> tuple[OwnerTeamLoad, LiteLLM_TeamTableCachedObj | None]: + from litellm.proxy.auth.auth_checks import get_team_object + + try: + team: Final = await get_team_object( + team_id=team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + check_db_only=source.check_db_only, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # raw team lookup failures grant nothing; native failure is never converted + verbose_proxy_logger.warning( + "Personal key team model access: team=%s lookup failed, granting nothing: %s", team_id, error + ) + return OwnerTeamLoad(False, failed=True), None + return OwnerTeamLoad(True, blocked=team.blocked is True, org=team.organization_id), team + + +async def _budget(index: int, reading: OwnerTeamReading, source: _Source) -> OwnerTeamBudget: + team: Final = reading.teams[index] + if team is None or source.prisma is None: + raise NativeAuthRuntimeError("owner team budget read requested without its source") + return OwnerTeamBudget( + index, + await read_personal_team_budget_facts( + team, source.token, source.prisma, source.cache, source.logging, check_db_only=source.check_db_only + ), + ) + + +async def _read(reading: OwnerTeamReading, source: _Source) -> OwnerTeamReading: + decision: Final = owner_teams(reading.facts) + match decision.kind: + case "Complete": + return reading + case "ReadTeams": + rows: Final = await asyncio.gather(*(_load(team_id, source) for team_id in decision.ids)) + updated: Final = OwnerTeamReading( + replace(reading.facts, loaded=True, loads=tuple(row[0] for row in rows)), tuple(row[1] for row in rows) + ) + return await _read(updated, source) + case "ReadBudgets": + budgets: Final = tuple( + await asyncio.gather(*(_budget(index, reading, source) for index in decision.indices)) + ) + return await _read( + replace(reading, facts=replace(reading.facts, budgets=(*reading.facts.budgets, *budgets))), source + ) diff --git a/litellm/proxy/auth/personal_budget_adapter.py b/litellm/proxy/auth/personal_budget_adapter.py new file mode 100644 index 00000000000..a251c162fbc --- /dev/null +++ b/litellm/proxy/auth/personal_budget_adapter.py @@ -0,0 +1,280 @@ +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from typing import Final + +from pydantic import BaseModel + +import litellm +from litellm._auth.common_budgets import BudgetFacts, MemberLimitFacts, WindowFacts, window_start +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.personal_budget import PersonalTeamBudgetDecision, PersonalTeamBudgetFacts, personal_team_budget +from litellm._logging import verbose_proxy_logger +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import Litellm_EntityType, LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj, UserAPIKeyAuth +from litellm.proxy.auth.budget_alert_adapter import team_member_max_budget_alert_check +from litellm.proxy.auth.budget_windows_adapter import ( + _reset, # pyright: ignore[reportPrivateUsage] # shared raw window timestamp projection + _seconds, # pyright: ignore[reportPrivateUsage] # shared raw duration projection + _window, # pyright: ignore[reportPrivateUsage] # shared typed window decoding +) +from litellm.proxy.auth.common_budget_adapter import CounterRead, emit_alert, key_alert_info, metadata, read_budget +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class _Windows(BaseModel, frozen=True): + budget_limits: tuple[object, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class _Sources: + team: LiteLLM_TeamTableCachedObj + token: UserAPIKeyAuth + prisma: PrismaClient + cache: UserApiKeyCache + logging: ProxyLogging + windows: tuple[object, ...] = () + membership: LiteLLM_TeamMembership | None = None + + +async def read_personal_team_budget_facts( + team: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, +) -> PersonalTeamBudgetFacts: + facts: Final = PersonalTeamBudgetFacts( + user_present=valid_token.user_id is not None, + team_id_present=team.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # unvalidated cache rows + check_db_only=check_db_only, + token_org=valid_token.org_id, + team_org=team.organization_id, + ) + decision: Final = personal_team_budget(facts) + if decision.kind != "Source": + raise NativeAuthRuntimeError("personal budget source was not selected") + source: Final = valid_token.model_copy(update={"team_id": team.team_id, "org_id": decision.org_id}) + return await _read( + replace(facts, source_ready=True), _Sources(team, source, prisma_client, user_api_key_cache, proxy_logging_obj) + ) + + +async def personal_key_team_within_budget( + team: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, +) -> bool: + facts: Final = await read_personal_team_budget_facts( + team, valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only + ) + return personal_team_budget(facts).allowed + + +async def _read(facts: PersonalTeamBudgetFacts, source: _Sources) -> PersonalTeamBudgetFacts: + decision: Final = personal_team_budget(facts) + if decision.kind == "Complete": + if decision.log_failure: + verbose_proxy_logger.warning( + "Failed to check budget for personal key team %s, skipping: incompatible datetimes", source.team.team_id + ) + return facts + try: + updated, updated_source = await _effect(decision, facts, source) + except NativeAuthFailure: + raise + except litellm.BudgetExceededError: + return replace(facts, io_failed=True) + except Exception as error: # noqa: BLE001 # preserve fail-closed personal budget IO behavior without catching native failures + verbose_proxy_logger.warning( + "Failed to check budget for personal key team %s, skipping: %s", source.team.team_id, error + ) + return replace(facts, io_failed=True) + return await _read(updated, updated_source) + + +def _member_limit(row: LiteLLM_BudgetTable | None, now: datetime) -> MemberLimitFacts: + expiry: Final = row.temp_budget_expiry if row is not None else None + aware: Final = expiry.replace(tzinfo=timezone.utc) if expiry is not None and expiry.tzinfo is None else expiry + expiry_number: Final = ( + aware.replace(tzinfo=timezone.utc).timestamp() + if aware is not None and aware.tzinfo is now.tzinfo + else aware.timestamp() + if aware is not None + else None + ) + now_number: Final = ( + now.replace(tzinfo=timezone.utc).timestamp() + if aware is not None and aware.tzinfo is now.tzinfo + else now.timestamp() + ) + return MemberLimitFacts( + row.max_budget if row is not None else None, + None, + row.temp_budget_increase if row is not None else None, + expiry_number, + now_number, + now.utcoffset() is None, + ) + + +async def _effect( + decision: PersonalTeamBudgetDecision, facts: PersonalTeamBudgetFacts, source: _Sources +) -> tuple[PersonalTeamBudgetFacts, _Sources]: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget, get_team_membership, get_utc_datetime + + team: Final = source.team + match decision.kind: + case "ReadTeam": + budget, _decision = await read_budget( + BudgetFacts("Team", team.max_budget, token_present=True), + CounterRead(f"spend:team:{team.team_id}", team.spend or 0.0, team.max_budget), + ) + return replace(facts, team=budget), source + case "AlertTeam": + if facts.team is None: + raise NativeAuthRuntimeError("missing personal team budget alert facts") + emit_alert( + source.logging, + "team_budget", + key_alert_info(source.token, facts.team.spend, facts.team.limit, Litellm_EntityType.TEAM), + ) + return replace(facts, team_alerted=True), source + case "ReadWindows": + windows: Final = _Windows.model_validate(team, from_attributes=True).budget_limits or () + return replace(facts, window_count=len(windows)), replace(source, windows=windows) + case "ReadWindow": + window: Final = _window(source.windows[decision.index]) + seconds: Final = _seconds(window.budget_duration) + reset: Final = _reset(window.reset_at) + start: Final = window_start( + WindowFacts( + seconds, reset.timestamp() if reset is not None else None, datetime.now(timezone.utc).timestamp() + ) + ) + window_facts, _window_decision = await read_budget( + BudgetFacts("TeamWindow", window.max_budget), + CounterRead( + f"spend:team:{team.team_id}:window:{window.budget_duration}", + 0.0, + window.max_budget, + window_entity_type="Team", + window_entity_id=team.team_id, + window_duration=window.budget_duration, + window_start=datetime.fromtimestamp(start, timezone.utc) if start is not None else None, + ), + ) + return replace(facts, windows=(*facts.windows, window_facts)), source + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("missing personal team membership identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + check_db_only=facts.check_db_only, + ) + return replace(facts, membership_loaded=True), replace(source, membership=membership) + case "ReadMemberLimit": + row: Final = source.membership.litellm_budget_table if source.membership is not None else None + now: Final = get_utc_datetime() + default_id: Final = metadata(team.metadata).get("team_member_budget_id") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy metadata projected as an opaque mapping + return replace( + facts, + member_limit=_member_limit(row, now), + default_id=default_id if isinstance(default_id, str) else None, + ), source + case "ReadDefault": + if facts.member_limit is None: + raise NativeAuthRuntimeError("missing personal member limit facts") + default: Final = await get_team_member_default_budget( + budget_id=decision.id, prisma_client=source.prisma, user_api_key_cache=source.cache + ) + return replace( + facts, + default_loaded=True, + member_limit=replace( + facts.member_limit, default_limit=default.max_budget if default is not None else None + ), + ), source + case "ReadMember": + member, _member_decision = await read_budget( + BudgetFacts("TeamMember", decision.limit), + CounterRead( + f"spend:team_member:{source.token.user_id}:{team.team_id}", + (source.membership.spend or 0.0) if source.membership is not None else 0.0, + decision.limit, + ), + ) + return replace(facts, member=member), source + case "AlertMember": + alert_member: Final = facts.member + if ( + alert_member is None + or alert_member.spend is None + or alert_member.limit is None + or source.token.user_id is None + ): + raise NativeAuthRuntimeError("missing personal member alert facts") + team_member_max_budget_alert_check( + team_id=team.team_id, + team_alias=team.team_alias, + team_metadata=team.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # helper validates legacy alert metadata + organization_id=team.organization_id, + user_id=source.token.user_id, + user_email=None, + proxy_logging_obj=source.logging, + spend=alert_member.spend, + max_budget=alert_member.limit, + ) + return replace(facts, member_alerted=True), source + case "ReadOrg": + org_budget, org_failed = await _read_org(decision.id, source) + return replace(facts, org_loaded=True, org=org_budget, org_failed=org_failed), source + case "AlertOrg": + if facts.org is None: + raise NativeAuthRuntimeError("missing personal org budget alert facts") + emit_alert( + source.logging, + "organization_budget", + key_alert_info( + source.token, facts.org.spend, facts.org.limit, Litellm_EntityType.ORGANIZATION + ).model_copy(update={"organization_id": decision.id}), + ) + return replace(facts, org_alerted=True), source + case "Source" | "Complete": + raise NativeAuthRuntimeError("unexpected personal budget read") + + +async def _read_org(org_id: str, source: _Sources) -> tuple[BudgetFacts | None, bool]: + from litellm.proxy.auth.auth_checks import get_org_object + + try: + org: Final = await get_org_object( + org_id=org_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + include_budget_table=True, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy org lookup failure skips only the org budget, not the other personal-team ceilings + return None, True + if org is None: + return None, False + carry_organization_budget_state(valid_token=source.token, org_table=org) + limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None + budget, _decision = await read_budget( + BudgetFacts("Organization", limit), CounterRead(f"spend:org:{org_id}", org.spend or 0.0, limit) + ) + return budget, False diff --git a/litellm/proxy/auth/personal_model_adapter.py b/litellm/proxy/auth/personal_model_adapter.py new file mode 100644 index 00000000000..9fa87d43502 --- /dev/null +++ b/litellm/proxy/auth/personal_model_adapter.py @@ -0,0 +1,192 @@ +import asyncio +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Final + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import personal_model_access_applies +from litellm._auth.personal_models import ( + LoadedTeamModelFacts, + PersonalModelDecision, + PersonalModelScope, + loaded_team_models, + personal_model_scope, +) +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +def _eligibility(token: UserAPIKeyAuth | None, *, inheritance: bool) -> bool: + return personal_model_access_applies( + enabled=bool(litellm.personal_key_model_access_from_teams), + token_present=token is not None, + team_id=token.team_id if token is not None else None, + user_id=token.user_id if token is not None else None, + master=token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS if token is not None else False, + role=token.user_role if token is not None else None, + api_key_present=token.api_key is not None if token is not None else False, + inheritance=inheritance, + ) + + +def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) -> bool: + return _eligibility(valid_token, inheritance=False) + + +def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool: + return _eligibility(valid_token, inheritance=True) + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str + team: LiteLLM_TeamTableCachedObj + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + aliases: Mapping[str, str] | None + + +async def read_personal_model_scope( + model: str, + team_object: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> PersonalModelScope: + from litellm.proxy.auth.auth_checks import key_model_aliases_for_auth_check + + aliases: Final = key_model_aliases_for_auth_check(valid_token) + facts: Final = PersonalModelScope( + user_present=valid_token.user_id is not None, + team_id_present=team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # unvalidated cache rows + ) + return await _read_scope( + facts, + _Source( + model, team_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj, aliases + ), + ) + + +async def _read_scope(facts: PersonalModelScope, source: _Source) -> PersonalModelScope: + decision: Final = personal_model_scope(facts) + if decision.kind == "Complete": + return facts + try: + updated: Final = await _scope_effect(decision, facts, source) + except NativeAuthFailure: + raise + except ProxyException: + return replace(facts, io_failed=True) + except Exception as error: # noqa: BLE001 # preserve personal team lookup failures as raw failure facts, not policy callbacks + verbose_proxy_logger.warning( + "Personal key team model access: team=%s check failed, granting nothing: %s", source.team.team_id, error + ) + return replace(facts, io_failed=True) + return await _read_scope(updated, source) + + +async def _scope_effect( + decision: PersonalModelDecision, facts: PersonalModelScope, source: _Source +) -> PersonalModelScope: + from litellm.proxy.auth.auth_checks import get_team_membership + + match decision.kind: + case "ReadTeam": + team: Final = project_model_facts( + source.model, + source.router, + object_model_grants(source.team), + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return replace(facts, team=GroupedModelScope(team, tuple(source.team.access_group_ids or ()))) + case "ReadGroups": + if facts.team is None: + raise NativeAuthRuntimeError("personal model group read without team facts") + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return replace(facts, team=grouped) + case "ReadMembership": + if source.token.user_id is None: + raise NativeAuthRuntimeError("personal model membership read without identity") + membership: Final = await get_team_membership( + user_id=source.token.user_id, + team_id=source.team.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + row: Final = membership.litellm_budget_table if membership is not None else None + grants: Final = tuple(row.allowed_models or ()) if row is not None else () + return replace(facts, member_grants=grants) + case "ReadMember": + if facts.member_grants is None: + raise NativeAuthRuntimeError("personal model member scope read without grants") + member: Final = project_model_facts( + source.model, + source.router, + facts.member_grants, + team_id=source.team.team_id, + key_model_aliases=source.aliases, + object_type="team", + ) + return replace(facts, member=member) + case "ReadScopes" | "Complete": + raise NativeAuthRuntimeError("unexpected personal model scope read") + + +async def read_loaded_team_model_facts( + model: str, + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LoadedTeamModelFacts: + facts: Final = LoadedTeamModelFacts(tuple(team is not None for team in teams)) + decision: Final = loaded_team_models(facts) + if decision.kind == "Complete": + return facts + if decision.kind != "ReadScopes": + raise NativeAuthRuntimeError("unexpected loaded team model read") + + async def read(index: int) -> PersonalModelScope: + team: Final = teams[index] + if team is None: + raise NativeAuthRuntimeError("missing selected personal model team") + return await read_personal_model_scope( + model, team, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + + scopes: Final = tuple(await asyncio.gather(*(read(index) for index in decision.indices))) + return replace(facts, scopes=scopes) + + +async def loaded_teams_allow_model( + model: str, + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> bool: + facts: Final = await read_loaded_team_model_facts( + model, teams, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + return loaded_team_models(facts).allowed diff --git a/litellm/proxy/auth/personal_model_listing_adapter.py b/litellm/proxy/auth/personal_model_listing_adapter.py new file mode 100644 index 00000000000..5e79f062535 --- /dev/null +++ b/litellm/proxy/auth/personal_model_listing_adapter.py @@ -0,0 +1,111 @@ +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.personal_model_listing import PersonalModelListingFacts, personal_model_listing +from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth +from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts +from litellm.proxy.auth.personal_owner_adapter import read_personal_owner_facts, selected_team +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + names: Sequence[str] + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +async def personal_key_listing_teams( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: + reading: Final = await read_personal_owner_facts( + user_api_key_dict, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + check_db_only=check_db_only, + purpose="Models", + ) + decision: Final = personal_owner(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete listing team facts") + return tuple(team if index in decision.selected else None for index, team in enumerate(reading.teams)) + + +async def read_personal_model_listing_facts( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> PersonalModelListingFacts: + owner: Final = PersonalOwnerFacts( + user_api_key_dict.user_id is not None, + prisma_client is not None, + user_api_key_cache is not None, + proxy_logging_obj is not None, + ) + return await _read( + PersonalModelListingFacts(owner), + _Source(models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + + +async def personal_key_team_visible_models( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[str, ...]: + facts: Final = await read_personal_model_listing_facts( + models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = personal_model_listing(facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete personal model listing facts") + return tuple((facts.names or ())[index] for index in decision.selected) + + +async def _read(facts: PersonalModelListingFacts, source: _Source) -> PersonalModelListingFacts: + decision: Final = personal_model_listing(facts) + match decision.kind: + case "Complete": + return facts + case "ReadOwner": + owner: Final = await read_personal_owner_facts( + source.token, source.prisma, source.cache, source.logging, purpose="Models" + ) + return await _read(replace(facts, owner=owner.facts), replace(source, teams=owner.teams)) + case "ReadNames": + return await _read(replace(facts, names=tuple(source.names)), source) + case "ReadModel": + if source.prisma is None or source.cache is None or source.logging is None or facts.names is None: + raise NativeAuthRuntimeError("personal listing scope read without dependencies") + teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams) + model: Final = await read_loaded_team_model_facts( + facts.names[decision.index], + teams, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + ) + return await _read(replace(facts, models=(*facts.models, model)), source) diff --git a/litellm/proxy/auth/personal_model_request_adapter.py b/litellm/proxy/auth/personal_model_request_adapter.py new file mode 100644 index 00000000000..f96da4834b1 --- /dev/null +++ b/litellm/proxy/auth/personal_model_request_adapter.py @@ -0,0 +1,93 @@ +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.model_access import ModelAccessDecision +from litellm._auth.personal_model_request import PersonalModelRequestFacts, personal_model_request +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import raise_model_denial +from litellm.proxy.auth.owner_team_adapter import read_owner_team_facts +from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts +from litellm.proxy.auth.personal_owner_adapter import selected_team +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + owner: LiteLLM_UserTable | None + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +async def read_personal_model_request_facts( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> PersonalModelRequestFacts: + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return await _read( + PersonalModelRequestFacts(requested, user_object is not None), + _Source(user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + + +async def can_personal_key_call_model_via_teams( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Literal[True]: + facts: Final = await read_personal_model_request_facts( + model, user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = personal_model_request(facts) + if decision.kind == "Allow": + return True + if decision.kind != "Deny": + raise NativeAuthRuntimeError("incomplete personal model request facts") + return raise_model_denial( + ModelAccessDecision("PersonalTeam", model=decision.model), + model, + user_id=valid_token.user_id, + teams=user_object.teams if user_object is not None else None, + ) + + +async def _read(facts: PersonalModelRequestFacts, source: _Source) -> PersonalModelRequestFacts: + decision: Final = personal_model_request(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadTeams": + if source.owner is None: + raise NativeAuthRuntimeError("personal model team read without owner") + reading: Final = await read_owner_team_facts( + source.owner, source.token, source.prisma, source.cache, source.logging, purpose="Models" + ) + return await _read(replace(facts, teams=reading.facts), replace(source, teams=reading.teams)) + case "ReadModel": + teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams) + model: Final = await read_loaded_team_model_facts( + facts.requested[decision.index], + teams, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + ) + return await _read(replace(facts, models=(*facts.models, model)), source) diff --git a/litellm/proxy/auth/personal_owner_adapter.py b/litellm/proxy/auth/personal_owner_adapter.py new file mode 100644 index 00000000000..01c3ea5ddd1 --- /dev/null +++ b/litellm/proxy/auth/personal_owner_adapter.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Final, Literal + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.owner_team_adapter import read_owner_team_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +if TYPE_CHECKING: + from litellm.proxy.auth.auth_checks import PersonalKeyOwnerTeams + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerReading: + facts: PersonalOwnerFacts + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + token: UserAPIKeyAuth + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + check_db_only: bool + owner: LiteLLM_UserTable | None = None + purpose: Literal["Models", "Grants"] = "Grants" + + +async def read_personal_owner_facts( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, + purpose: Literal["Models", "Grants"] = "Grants", +) -> PersonalOwnerReading: + facts: Final = PersonalOwnerFacts( + valid_token.user_id is not None, + prisma_client is not None, + user_api_key_cache is not None, + proxy_logging_obj is not None, + ) + return await _read( + PersonalOwnerReading(facts), + _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only, purpose=purpose), + ) + + +async def personal_key_owner_teams_with_status( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> PersonalKeyOwnerTeams: + from litellm.proxy.auth.auth_checks import PersonalKeyOwnerTeams + + reading: Final = await read_personal_owner_facts( + valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only + ) + decision: Final = personal_owner(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete personal owner team facts") + return PersonalKeyOwnerTeams( + teams=tuple(selected_team(reading.teams, index) for index in decision.selected), + withheld=tuple(selected_team(reading.teams, index) for index in decision.withheld), + unreadable=decision.unreadable, + ) + + +def selected_team(teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], index: int) -> LiteLLM_TeamTableCachedObj: + team: Final = teams[index] + if team is None: + raise NativeAuthRuntimeError("missing selected owner team row") + return team + + +async def _read(reading: PersonalOwnerReading, source: _Source) -> PersonalOwnerReading: + from litellm.proxy.auth.auth_checks import get_user_object + + decision: Final = personal_owner(reading.facts) + if decision.kind == "Complete": + return reading + if source.token.user_id is None or source.prisma is None or source.cache is None or source.logging is None: + raise NativeAuthRuntimeError("owner read requested without required dependencies") + match decision.kind: + case "ReadOwner": + try: + owner: Final = await get_user_object( + user_id=source.token.user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + user_id_upsert=False, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + check_db_only=source.check_db_only, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve legacy owner lookup failure as raw data + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", error + ) + return PersonalOwnerReading(replace(reading.facts, owner_loaded=True, owner_failed=True)) + return await _read( + PersonalOwnerReading(replace(reading.facts, owner_loaded=True, owner_present=owner is not None)), + replace(source, owner=owner), + ) + case "ReadTeams": + if source.owner is None: + raise NativeAuthRuntimeError("personal owner team read without owner row") + teams: Final = await read_owner_team_facts( + source.owner, + source.token, + source.prisma, + source.cache, + source.logging, + purpose=source.purpose, + check_db_only=source.check_db_only, + ) + return PersonalOwnerReading(replace(reading.facts, teams=teams.facts), teams.teams) diff --git a/litellm/proxy/auth/resolved_model_adapter.py b/litellm/proxy/auth/resolved_model_adapter.py new file mode 100644 index 00000000000..6316a43a885 --- /dev/null +++ b/litellm/proxy/auth/resolved_model_adapter.py @@ -0,0 +1,283 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final + +from pydantic import BaseModel, Field, TypeAdapter + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import key_model_grants +from litellm._auth.resolved_models import ResolvedModelFacts, ResolvedTeamRead, resolved_models +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.member_model_adapter import read_member_model_facts +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.proxy.auth.personal_model_request_adapter import read_personal_model_request_facts +from litellm.proxy.auth.team_override_adapter import read_team_override_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +class _Config(BaseModel, frozen=True): + config: Mapping[str, object] = Field(default_factory=dict[str, object]) + + +_METADATA: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + +@dataclass(frozen=True, slots=True) +class ResolvedModelReading: + facts: ResolvedModelFacts + team: LiteLLM_TeamTableCachedObj | None = None + owner: LiteLLM_UserTable | None = None + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + team: LiteLLM_TeamTableCachedObj | None = None + owner: LiteLLM_UserTable | None = None + project: LiteLLM_ProjectTable | None = None + customer: LiteLLM_EndUserTable | None = None + + +async def read_resolved_model_facts( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> ResolvedModelReading: + facts: Final = ResolvedModelFacts( + model=model, + config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config), + key_models=object_model_grants(valid_token), + team_id=valid_token.team_id, + user_id=valid_token.user_id, + project_present=valid_token.project_id is not None, + customer_present=valid_token.end_user_id is not None, + database_present=prisma_client is not None, + personal_enabled=bool(litellm.personal_key_model_access_from_teams), + master=valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS, + role=valid_token.user_role, + ) + return await _read( + facts, _Source(model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj) + ) + + +async def can_key_call_resolved_model( + model: str, llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None +) -> None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + reading: Final = await read_resolved_model_facts( + model, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = resolved_models(reading.facts) + if decision.kind == "Allow": + return + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete resolved model access facts") + raise_model_denial( + decision.denial, + model, + user_id=valid_token.user_id, + team_id=reading.team.team_id if reading.team is not None else None, + teams=reading.owner.teams if reading.owner is not None else None, + member_models=decision.member_models, + ) + + +async def _read(facts: ResolvedModelFacts, source: _Source) -> ResolvedModelReading: + from litellm.proxy.auth.auth_checks import ( + get_end_user_object, + get_key_end_user_budget_id, + get_project_object, + get_team_object, + get_user_object, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, + ) + + decision: Final = resolved_models(facts) + match decision.kind: + case "Allow" | "Deny": + return ResolvedModelReading(facts, source.team, source.owner) + case "ReadKey": + grants: Final = key_model_grants( + object_model_grants(source.token), source.token.team_id, tuple(source.token.team_models or ()) + ) + key: Final = project_model_facts( + source.model, + source.router, + grants, + source.token.team_model_aliases, + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "key", + ) + return await _read( + replace(facts, key=GroupedModelScope(key, tuple(source.token.access_group_ids or ()))), source + ) + case "ReadKeyGroups": + if facts.key is None: + raise NativeAuthRuntimeError("resolved key group read without scope") + scope: Final = await read_grouped_scope(facts.key, source.router, merge=False) + return await _read(replace(facts, key=scope), source) + case "ReadTeamObject": + if source.token.team_id is None: + raise NativeAuthRuntimeError("resolved team read without identity") + try: + team: Final = await get_team_object( + team_id=source.token.team_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # legacy resolved-model team lookup falls back to the embedded token row + return await _read(replace(facts, team_read=ResolvedTeamRead(failed=True)), source) + return await _read( + replace(facts, team_read=ResolvedTeamRead(found=team is not None)), # pyright: ignore[reportUnnecessaryComparison] # preserve the legacy missing-row result + replace(source, team=team), + ) + case "ReadEmbeddedTeam": + embedded: Final = LiteLLM_TeamTableCachedObj( + team_id=source.token.team_id, # pyright: ignore[reportArgumentType] # native selects this read only for a present team ID + models=source.token.team_models, + blocked=source.token.team_blocked, + team_alias=source.token.team_alias, + metadata=source.token.team_metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy typed row validates embedded metadata + object_permission_id=source.token.team_object_permission_id, + object_permission=source.token.team_object_permission, + ) + return await _read(replace(facts, embedded_ready=True), replace(source, team=embedded)) + case "ReadTeamScope": + team_scope: Final = await read_team_override_facts( + source.model, + source.token, + source.team, + source.router, + source.token.team_model_aliases, + key_model_aliases_for_auth_check(source.token), + ) + return await _read(replace(facts, team=team_scope), source) + case "ReadMember": + if source.team is None: + raise NativeAuthRuntimeError("resolved member model read without team") + member: Final = await read_member_model_facts( + source.model, + source.team, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + ) + return await _read(replace(facts, member=member), source) + case "ReadOwner": + if source.token.user_id is None: + raise NativeAuthRuntimeError("resolved personal owner read without identity") + try: + owner: Final = await get_user_object( + user_id=source.token.user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + user_id_upsert=False, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve raw owner lookup failure and native fail-closed model selection + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", error + ) + return await _read(replace(facts, owner_loaded=True, owner_failed=True), source) + return await _read( + replace(facts, owner_loaded=True, owner_present=owner is not None), replace(source, owner=owner) + ) + case "ReadPersonal": + personal: Final = await read_personal_model_request_facts( + source.model, source.owner, source.token, source.router, source.prisma, source.cache, source.logging + ) + return await _read(replace(facts, personal=personal), source) + case "ReadProjectObject": + if source.token.project_id is None: + raise NativeAuthRuntimeError("resolved project read without identity") + project: Final = await get_project_object( + project_id=source.token.project_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + ) + return await _read( + replace( + facts, + project_loaded=True, + project_grants=object_model_grants(project) if project is not None else (), + ), + replace(source, project=project), + ) + case "ReadProjectScope": + if source.project is None: + raise NativeAuthRuntimeError("resolved project scope read without row") + project_scope: Final = project_model_facts( + source.model, + source.router, + facts.project_grants, + team_id=source.project.team_id, + key_model_aliases=key_model_aliases_for_auth_check(source.token), + object_type="project", + ) + return await _read(replace(facts, project=project_scope), source) + case "ReadCustomerObject": + if source.token.end_user_id is None or source.prisma is None: + raise NativeAuthRuntimeError("resolved customer read without identity or database") + customer: Final = await get_end_user_object( + end_user_id=source.token.end_user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + proxy_logging_obj=source.logging, + token_end_user_max_budget=source.token.end_user_max_budget, + key_end_user_budget_id=get_key_end_user_budget_id(_METADATA.validate_python(source.token.metadata)), # pyright: ignore[reportUnknownMemberType] # validate the legacy metadata field at the IO boundary + ) + return await _read( + replace( + facts, + customer_loaded=True, + customer_grants=tuple(customer.models or ()) if customer is not None else (), + ), + replace(source, customer=customer), + ) + case "ReadCustomerScope": + customer_scope: Final = project_model_facts( + source.model, + source.router, + facts.customer_grants, + team_model_aliases_for_auth_check(source.token), + source.token.team_id, + key_model_aliases_for_auth_check(source.token), + "customer", + ) + return await _read(replace(facts, customer=customer_scope), source) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 0aa2528d46b..4f1e5df0505 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -1,10 +1,9 @@ -import itertools -import re from collections.abc import Collection, Iterable, Mapping 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, @@ -18,10 +17,13 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.auth import route_policy_adapter from .auth_checks_organization import ( # noqa: F401 # legacy module exports _user_is_org_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - user_is_org_admin, +) +from .auth_checks_organization import ( + user_is_org_admin as user_is_org_admin, # noqa: PLC0414 # preserve the public compatibility re-export ) # Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write @@ -112,14 +114,9 @@ _TOKEN_COUNTING_ACTION: Final = "countTokens" class RouteChecks: @staticmethod def check_route_enabled(route: str) -> None: - try: - from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + from litellm.proxy.auth.route_enable_adapter import check_route_enabled - EnterpriseRouteChecks.should_call_route(route=route) - except HTTPException as e: - raise e - except Exception: - pass + check_route_enabled(route) @staticmethod def should_call_route( @@ -141,120 +138,7 @@ class RouteChecks: valid_token: UserAPIKeyAuth, request: Request | None = None, ) -> bool: - """ - Raises Exception if Virtual Key is not allowed to call the route - """ - - # Only check if valid_token.allowed_routes is set and is a list with at least one item - if valid_token.allowed_routes is None: - return True - if not isinstance(valid_token.allowed_routes, list): - return True - if len(valid_token.allowed_routes) == 0: - return True - - denied_auth_enforced_pass_through_route = False - - # explicit check for allowed routes (exact match or prefix match) - for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): - return True - - ## check if 'allowed_route' is a field name in LiteLLMRoutes - if any(allowed_route in LiteLLMRoutes._member_names_ for allowed_route in valid_token.allowed_routes): - for allowed_route in valid_token.allowed_routes: - if allowed_route in LiteLLMRoutes._member_names_: - if RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes[allowed_route].value, - ): - if ( - allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS - and RouteChecks.is_auth_enforced_pass_through_route( - route=route, - method=RouteChecks.get_request_method(request=request), - ) - ): - if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): - return True - denied_auth_enforced_pass_through_route = True - else: - return True - - ################################################ - # For llm_api_routes, also check registered pass-through endpoints - ################################################ - if allowed_route in _LLM_API_ROUTE_GROUPS: - if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET": - return True - - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - InitPassThroughEndpointHelpers, - ) - - if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): - if RouteChecks.is_auth_enforced_pass_through_route( - route=route, - method=RouteChecks.get_request_method(request=request), - ): - if RouteChecks.check_passthrough_route_access( - route=route, user_api_key_dict=valid_token - ): - return True - denied_auth_enforced_pass_through_route = True - else: - return True - - # Method-aware carve-out: allow GET on the two - # read-only MCP-server discovery endpoints - # (`/v1/mcp/server` and `/v1/mcp/server/{server_id}`) - # so virtual keys with allowed_routes=["llm_api_routes"] - # can list/inspect MCP servers. The GET handlers in - # mcp_management_endpoints.py sanitize the response - # for restricted virtual keys (stripping url, - # headers, env, credentials). POST/PUT/DELETE on - # these paths are admin-only management writes and - # are intentionally not covered. - if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route( - route=route, request=request - ): - return True - - # Agent registry CRUD moved from llm_api_routes into - # management_routes so DISABLE_LLM_API_ENDPOINTS stops - # blocking it. Keys configured with - # allowed_routes=["llm_api_routes"] before that split - # could reach these paths, so keep them reachable here; - # the handlers in agent_endpoints/endpoints.py still - # enforce proxy-admin on writes and scope reads by role. - if allowed_route == "llm_api_routes" and RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.agent_management_routes.value, - ): - return True - - # check if wildcard pattern is allowed - for allowed_route in valid_token.allowed_routes: - if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): - return True - - if denied_auth_enforced_pass_through_route: - raise RouteChecks._auth_pass_through_denied_exception(route=route) - - if valid_token.metadata.get("password_reset_required") is True: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=( - "This account's password must be changed before the session can be used: " - "it was either found in a known data breach or set by an admin. " - "Change it via POST /user/password/change (UI: /ui/change-password), then log in again." - ), - ) - - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}", - ) + return route_policy_adapter.check_virtual_route(route, valid_token, request) @staticmethod def _mask_user_id(user_id: str) -> str: @@ -318,110 +202,10 @@ class RouteChecks: valid_token: UserAPIKeyAuth, request_data: dict, ): - """ - Checks if Non Proxy Admin User is allowed to access the route - """ - - # Check user has defined custom admin routes - RouteChecks.custom_admin_only_route_check( - route=route, + return route_policy_adapter.check_non_admin_route( + user_obj, _user_role, route, request, valid_token, request_data ) - if RouteChecks.is_auth_enforced_pass_through_route( - route=route, - method=RouteChecks.get_request_method(request=request), - ): - RouteChecks._require_auth_pass_through_access( - route=route, - valid_token=valid_token, - jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token), - ) - elif RouteChecks.is_llm_api_route(route=route): - pass - elif RouteChecks.is_info_route(route=route): - # check if user allowed to call an info route - if route == "/key/info": - # handled by function itself - pass - elif route == "/user/info": - # check if user can access this route - query_params: Final = request.query_params - user_id: Final = query_params.get("user_id") - verbose_proxy_logger.debug("user_id: %s & valid_token.user_id: %s", user_id, valid_token.user_id) - if ( - user_id - and user_id != valid_token.user_id - and _user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"key not allowed to access this user's info. user_id={user_id}, key's user_id={valid_token.user_id}", - ) - elif route == "/v2/user/info": - # handled by the endpoint itself (full RBAC in handler) - pass - elif route == "/model/info": - # /model/info just shows models user has access to - pass - elif route == "/team/info": - pass # handled by function itself - elif ( - route in LiteLLMRoutes.global_spend_tracking_routes.value - and getattr(valid_token, "permissions", None) is not None - and "get_spend_routes" in getattr(valid_token, "permissions", []) - ): - pass - elif route == "/lens" or route.startswith("/lens/"): - pass # Lens authorizes the delegated role, including its POST-based read endpoints. - elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: - RouteChecks._check_proxy_admin_viewer_access( - route=route, - _user_role=_user_role, - request_data=request_data, - request=request, - ) - elif ( - _user_role == LitellmUserRoles.INTERNAL_USER.value - and ( - RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value) - or RouteChecks.is_self_scoped_credential_route(route=route, request=request) - ) - or user_is_org_admin(request_data=request_data, user_object=user_obj) - and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value) - or _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - and RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value, - ) - or RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.self_managed_routes.value) - ): - pass - elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"): - pass # authN/authZ handled by api itself - elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token) or ( - valid_token.is_team_service_account - and RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.team_service_account_key_routes.value - ) - ): - pass - elif valid_token.allowed_routes is not None: - # check if route is in allowed_routes (exact match or prefix match) - route_allowed = False - for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): - route_allowed = True - break - - if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): - route_allowed = True - break - - if not route_allowed: - RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) - else: - RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) - @staticmethod def custom_admin_only_route_check(route: str): from litellm.proxy.proxy_server import general_settings, premium_user @@ -454,73 +238,11 @@ class RouteChecks: @staticmethod def is_unbilled_route(route: str) -> bool: - resource, _, action = route.rsplit("/", 1)[-1].partition(":") - return route in _UNBILLED_ROUTES or resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION + return route_policy_adapter.classify(route, "Unbilled") @staticmethod def is_llm_api_route(route: str) -> bool: - """ - Helper to checks if provided route is an OpenAI route - - - Returns: - - True: if route is an OpenAI route - - False: if route is not an OpenAI route - """ - # Ensure route is a string before performing checks - if not isinstance(route, str): - return False - - if route in LiteLLMRoutes.openai_routes.value: - return True - - if route in LiteLLMRoutes.anthropic_routes.value: - return True - - if route in LiteLLMRoutes.google_routes.value: - return True - - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): - return True - - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): - return True - - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): - return True - - # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - # Check for routes with placeholders or wildcard patterns - for openai_route in LiteLLMRoutes.openai_routes.value: - # Replace placeholders with regex pattern - # placeholders are written as "/threads/{thread_id}" - if "{" in openai_route: - if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): - return True - # Check for wildcard patterns like "/containers/*" - if RouteChecks.is_wildcard_pattern(pattern=openai_route): - if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): - return True - - # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" - for google_route in LiteLLMRoutes.google_routes.value: - if "{" in google_route: - if RouteChecks._route_matches_pattern(route=route, pattern=google_route): - return True - - # Check for Anthropic routes with placeholders - for anthropic_route in LiteLLMRoutes.anthropic_routes.value: - if "{" in anthropic_route: - if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): - return True - - if RouteChecks._is_azure_openai_route(route=route): - return True - - for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: - if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): - return True - return RouteChecks.is_builtin_eu_jev_decision_route(route=route) + return isinstance(route, str) and route_policy_adapter.classify(route, "Llm") @staticmethod def is_builtin_eu_jev_decision_route(route: str) -> bool: @@ -536,26 +258,7 @@ class RouteChecks: @staticmethod def _is_get_mcp_server_discovery_route(route: str, request: Request | None) -> bool: - """ - Returns True if `request` is a GET against one of the two read-only - MCP-server discovery paths: - - - GET `/v1/mcp/server` (list) - - GET `/v1/mcp/server/{server_id}` (single server, single segment) - - Multi-segment paths (`/v1/mcp/server/{id}/approve`, etc.) and any - non-GET method return False, so admin-only management writes on the - same path prefix are not reachable through this carve-out. - """ - if request is None or request.method.upper() != "GET": - return False - if route == "/v1/mcp/server": - return True - prefix: Final = "/v1/mcp/server/" - if not route.startswith(prefix): - return False - remainder: Final = route[len(prefix) :] - return bool(remainder) and "/" not in remainder + return route_policy_adapter.classify(route, "Discovery", RouteChecks.get_request_method(request)) @staticmethod def is_management_route(route: str) -> bool: @@ -580,67 +283,11 @@ class RouteChecks: @staticmethod def _is_azure_openai_route(route: str) -> bool: - """ - Check if route is a route from AzureOpenAI SDK client - - eg. - route='/openai/deployments/vertex_ai/gemini-1.5-flash/chat/completions' - """ - # Ensure route is a string before attempting regex matching - if not isinstance(route, str): - return False - # Add support for deployment and engine model paths - deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" - engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" - - if re.match(deployment_pattern, route) or re.match(engine_pattern, route): - return True - return False + return isinstance(route, str) and route_policy_adapter.classify(route, "Azure") @staticmethod def _route_matches_pattern(route: str, pattern: str) -> bool: - """ - Check if route matches the pattern placed in proxy/_types.py - - Example: - - pattern: "/threads/{thread_id}" - - route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - - returns: True - - - - pattern: "/key/{token_id}/regenerate" - - route: "/key/regenerate/82akk800000000jjsk" - - returns: False, pattern is "/key/{token_id}/regenerate" - """ - # Ensure route is a string before attempting regex matching - if not isinstance(route, str): - return False - - def _placeholder_to_regex(match: re.Match) -> str: - placeholder: Final = match.group(0).strip("{}") - if not placeholder.endswith(":path"): - return r"[^/]+" - # A ":path" placeholder takes whatever the router's own path - # converter takes, slashes and colons alike, so an id spelled with - # either (or both) still matches the template it was mounted under. - # - # Unless the template puts a ":" literal of its own after the - # placeholder: the Google routes end in ":generateContent" and - # friends, and there the value has to stop before that suffix - # rather than swallow it and match a different verb. - # - # "[\s\S]" rather than ".", because "." stops at a newline and the - # path converter does not: a %0A anywhere in the value would leave - # the route unmatched here while still reaching the handler, which - # turns this gate into a bypass for the lists built on it. - return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" - - pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) - # Anchor the pattern to match the entire string - pattern = f"^{pattern}$" - if re.match(pattern, route): - return True - return False + return isinstance(route, str) and routes.match_route(route, (pattern,), "Template") @staticmethod def is_wildcard_pattern(pattern: str) -> bool: @@ -653,100 +300,15 @@ class RouteChecks: @staticmethod def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: - """ - Check if route matches the wildcard pattern - - eg. - - pattern: "/scim/v2/*" - route: "/scim/v2/Users" - - returns: True - - pattern: "/scim/v2/*" - route: "/chat/completions" - - returns: False - - - pattern: "/scim/v2/*" - route: "/scim/v2/Users/123" - - returns: True - - """ - if pattern.endswith("*"): - # Get the prefix (everything before the wildcard) - prefix: Final = pattern[:-1] - return route.startswith(prefix) - else: - # If there's no wildcard, the pattern and route should match exactly - return route == pattern + return routes.match_route(route, (pattern,), "Wildcard") @staticmethod def _route_matches_allowed_route(route: str, allowed_route: str) -> bool: - """ - Check if route matches the allowed_route pattern. - Supports both exact match and prefix match. - - Examples: - - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6" -> True (exact match) - - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6/v1/chat/completions" -> True (prefix match) - - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-600" -> False (not a valid prefix) - - Args: - route: The actual route being accessed - allowed_route: The allowed route pattern - - Returns: - bool: True if route matches (exact or prefix), False otherwise - """ - # Exact match - if route == allowed_route: - return True - # Prefix match - ensure we add "/" to prevent false matches like /fake-openai-proxy-600 - if route.startswith(allowed_route + "/"): - return True - return False + return routes.match_route(route, (allowed_route,), "Prefix") @staticmethod def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: - """ - Check if a route has access by checking both exact matches and patterns - - Args: - route (str): The route to check - allowed_routes (Sequence): Allowed routes/patterns - - Returns: - bool: True if route is allowed, False otherwise - """ - ######################################################### - # exact match route is in allowed_routes - ######################################################### - if route in allowed_routes: - return True - - ######################################################### - # wildcard match route is in allowed_routes - # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* - ######################################################### - if any( - RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) - for allowed_route in allowed_routes - if RouteChecks.is_wildcard_pattern(pattern=allowed_route) - ): - return True - - ######################################################### - # pattern match route is in allowed_routes - # pattern: "/threads/{thread_id}" - # route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - # returns: True - ######################################################### - if any( # Check pattern match - RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes - ): - return True - - return False + return routes.match_route(route, tuple(allowed_routes)) @staticmethod def get_request_method(request: Request | None) -> str | None: @@ -768,13 +330,7 @@ class RouteChecks: @staticmethod def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: - method: Final = RouteChecks.get_request_method(request) - if method is None: - return False - return any( - method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,)) - for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES - ) + return route_policy_adapter.classify(route, "SelfCredential", RouteChecks.get_request_method(request)) @staticmethod def is_auth_enforced_pass_through_route(route: str, method: str | None = None) -> bool: @@ -806,48 +362,13 @@ class RouteChecks: @staticmethod def _route_matches_denied_route(route: str, denied_route: str) -> bool: - """A `/` entry denies every route, since every route sits under the root.""" - normalized_denied_route: Final = denied_route.rstrip("/") or "/" - return ( - normalized_denied_route == "/" - or RouteChecks._route_matches_allowed_route(route=route, allowed_route=normalized_denied_route) - or RouteChecks.route_matches_wildcard_pattern(route=route, pattern=denied_route) - ) + return routes.match_route(route, (denied_route,), "Denied") @staticmethod def matching_denied_passthrough_route( route: str, metadata_sources: Iterable[Mapping[str, object] | None] ) -> str | None: - """ - First ``denied_passthrough_routes`` entry across ``metadata_sources`` that matches ``route``. - Unlike the allowlist (key list, else team list), every source's deny list applies. - """ - denied_routes: Final = tuple( - itertools.chain.from_iterable( - cast( # cast-ok: management endpoints validate this metadata key as a list of route strings on write - "list[str]", (metadata or {}).get("denied_passthrough_routes") or [] - ) - for metadata in metadata_sources - ) - ) - if not denied_routes: - return None - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - InitPassThroughEndpointHelpers, - ) - - forwarded_routes: Final = InitPassThroughEndpointHelpers.forwarded_routes(route) - return next( - ( - denied_route - for denied_route in denied_routes - if any( - RouteChecks._route_matches_denied_route(route=candidate, denied_route=denied_route) - for candidate in forwarded_routes - ) - ), - None, - ) + return route_policy_adapter.denied_passthrough(route, metadata_sources) @staticmethod def passthrough_route_denied_exception(route: str, denied_route: str) -> HTTPException: @@ -870,16 +391,7 @@ class RouteChecks: @staticmethod def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool: - """ - Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do: - a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names - no path segment (``*``, ``/*``) is skipped. - """ - return any( - RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) - for allowed_route in team_allowed_routes - if allowed_route.rstrip("*").strip("/") - ) + return routes.match_route(route, tuple(team_allowed_routes), "Jwt") @staticmethod def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]: @@ -897,46 +409,18 @@ class RouteChecks: valid_token: UserAPIKeyAuth, jwt_team_allowed_routes: Collection[str] = (), ) -> None: - """ - Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the - key or team, or an explicit JWT ``team_allowed_routes`` entry. A key or team - ``denied_passthrough_routes`` match blocks the route even when one of those grants it. - """ - RouteChecks._raise_if_passthrough_route_denied(route=route, valid_token=valid_token) - if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): - return - if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes): - return - raise RouteChecks._auth_pass_through_denied_exception(route=route) + decision: Final = routes.passthrough_route( + routes.RouteContext(route), + route_policy_adapter.passthrough_facts(route, valid_token, jwt_team_allowed_routes, include_denied=True), + required=True, + ) + route_policy_adapter.raise_route_denial(decision, route, token=valid_token) @staticmethod def check_passthrough_route_access(route: str, user_api_key_dict: UserAPIKeyAuth) -> bool: - """ - Check if route is a passthrough route. - Supports both exact match and prefix match. - """ - metadata: Final = user_api_key_dict.metadata - team_metadata: Final = user_api_key_dict.team_metadata or {} - if metadata is None and team_metadata is None: - return False - if "allowed_passthrough_routes" not in metadata and "allowed_passthrough_routes" not in team_metadata: - return False - if ( - metadata.get("allowed_passthrough_routes") is None - and team_metadata.get("allowed_passthrough_routes") is None - ): - return False - - allowed_passthrough_routes: Final = ( - metadata.get("allowed_passthrough_routes") or team_metadata.get("allowed_passthrough_routes") or [] - ) - - # Check if route matches any allowed passthrough route (exact or prefix match) - for allowed_route in allowed_passthrough_routes: - if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): - return True - - return False + return routes.passthrough_route( + routes.RouteContext(route), route_policy_adapter.passthrough_facts(route, user_api_key_dict) + ).allowed @staticmethod def is_assistants_api_request(request: Request) -> bool: @@ -1017,117 +501,4 @@ class RouteChecks: request_data: dict, request: Request | None = None, ) -> None: - """ - Check access for PROXY_ADMIN_VIEW_ONLY role. - - Admin Viewer follows a read-parity-with-Proxy-Admin rule: anything Proxy - Admin can read/list/get, Admin Viewer can read/list/get. The only - exclusions are cost-incurring inference routes (Playground, /chat/ - completions, etc.) and any state-mutating request. - - Implementation: - 1. LLM/inference routes → 403 (cost-incurring). - 2. Safe HTTP method (GET/HEAD/OPTIONS) → allow by default. This is - the read-parity guarantee — every new GET endpoint added anywhere - in the codebase is automatically readable by Admin Viewer - without needing to remember to add it to an allowlist. - 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): - - Allow `/user/update` only when restricted to user_email. - - Allow `/user/password/change` (endpoint only writes the caller's own row). - - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. - - Otherwise allow only if the route is in admin_viewer_routes / - global_spend_tracking_routes (legacy explicit-allow set). - - Else 403. - """ - if RouteChecks.is_llm_api_route(route=route): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", - ) - - # Check if this is a write operation on management routes - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value): - # For management routes, only allow read operations or specific allowed updates - if route == "/user/update": - # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY - if request_data is not None and isinstance(request_data, dict): - _params_updated: Final = request_data.keys() - for param in _params_updated: - if param != "user_email": - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated", - ) - elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or ( - route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) - ): - # Block write operations for PROXY_ADMIN_VIEW_ONLY - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) - # Allow read operations on management routes (like /user/info, /team/info, /model/info) - method: Final = request.method.upper() if request is not None else "GET" - is_safe_method: Final = method in RouteChecks._SAFE_HTTP_METHODS - - # ── Safe HTTP method: default-allow ────────────────────────────── - if is_safe_method: - return - - # ── Unsafe HTTP method: explicit checks ────────────────────────── - # Allow `/user/update` for self-service email change. - if route == "/user/update": - if request_data is not None and isinstance(request_data, dict): - for param in request_data: - if param != "user_email": - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=( - f"user not allowed to access this route, role= {_user_role}. " - f"Trying to access: {route} and updating invalid param: {param}. " - "only user_email can be updated" - ), - ) - return - - # Self-service password change; the endpoint only writes the caller's own row. - if route == "/user/password/change": - return - - # Self-service logout; the endpoint only revokes the caller's own session key. - if route == "/session/logout": - return - - if RouteChecks.is_self_scoped_credential_route(route=route, request=request): - return - - # Hard-block known write routes regardless of HTTP method (defensive - # — these are POSTs in practice, but pinning them here protects - # against future GET-shaped writes). - if RouteChecks.check_route_access( - route=route, allowed_routes=RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES - ) or (route.startswith("/key/") and route.endswith("/regenerate")): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) - - # Legacy explicit-allow sets (kept for routes that are POST but - # semantically read-only, e.g. /spend/calculate). Both admin_viewer_routes - # and global_spend_tracking_routes are reads/listings. - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value): - return - if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value): - return - - # NOTE: We intentionally do NOT fall back to allowing all - # `management_routes`. That set is a mix of reads (info/list — handled - # via the safe-method branch above) and writes (`/team/block`, - # `/team/permissions_update`, `/jwt/key/mapping/{new,update,delete}`, - # `/key/bulk_update`, `/key/{id}/reset_spend`). A blanket allow would - # let Admin Viewer POST these write endpoints — violating the - # "no writes, ever" rule. Default-deny instead. - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) + return route_policy_adapter.check_viewer_route(route, _user_role, request_data, request) diff --git a/litellm/proxy/auth/route_enable_adapter.py b/litellm/proxy/auth/route_enable_adapter.py new file mode 100644 index 00000000000..4cb119c0381 --- /dev/null +++ b/litellm/proxy/auth/route_enable_adapter.py @@ -0,0 +1,125 @@ +import os +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Final, assert_never + +from fastapi import HTTPException +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting, route_enable +from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.auth.route_policy_adapter import read_eu_jev_configuration, route_context + +_KEYS: Final = {"Admin": "DISABLE_ADMIN_ENDPOINTS", "Llm": "DISABLE_LLM_API_ENDPOINTS"} + + +_EXEMPTIONS: Final = TypeAdapter(tuple[str, ...]) + + +def read_enterprise() -> object: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + return EnterpriseRouteChecks + + +def read_setting(feature: RouteFeature) -> RouteSetting: + from litellm.proxy.proxy_server import premium_user + from litellm.secret_managers.main import ( + get_secret_bool as _get_secret_bool, # noqa: F401 # pyright: ignore[reportUnusedImport] # preserve the Phase 1 import before environment and license reads + ) + + return RouteSetting(_KEYS[feature] in os.environ, bool(premium_user)) + + +def read_secret(feature: RouteFeature) -> bool: + from litellm.secret_managers.main import get_secret_bool + + return get_secret_bool(_KEYS[feature]) is True + + +def read_exemptions(enterprise: object) -> tuple[str, ...]: + return _EXEMPTIONS.validate_python( + getattr(enterprise, "LLM_API_EXEMPT_ROUTES") # noqa: B009 # optional class attributes are validated at this boundary + ) + + +@dataclass(frozen=True, slots=True) +class RouteEnableReaders: + enterprise: Callable[[], object] = read_enterprise + setting: Callable[[RouteFeature], RouteSetting] = read_setting + secret: Callable[[RouteFeature], bool] = read_secret + exemptions: Callable[[object], tuple[str, ...]] = read_exemptions + eu_jev: Callable[[], bool] = read_eu_jev_configuration + + +_READERS: Final = RouteEnableReaders() + + +def _answer( + facts: RouteEnableFacts, decision: RouteEnableDecision, readers: RouteEnableReaders, enterprise: object +) -> tuple[RouteEnableFacts, object]: + match decision.kind: + case "ReadImport": + return replace(facts, imported=True), readers.enterprise() + case "ReadExemptions": + return replace(facts, exemptions=readers.exemptions(enterprise)), enterprise + case "ReadEuJev": + return replace(facts, eu_jev_configured=readers.eu_jev()), enterprise + case "ReadSettings" | "ReadSecret": + if decision.feature is None: + raise NativeAuthRuntimeError("route setting read without feature") + previous: Final = facts.admin if decision.feature == "Admin" else facts.llm + if decision.kind == "ReadSecret" and previous is None: + raise NativeAuthRuntimeError("route secret read without settings") + setting: Final = ( + readers.setting(decision.feature) + if decision.kind == "ReadSettings" or previous is None + else replace(previous, secret=readers.secret(decision.feature)) + ) + return ( + replace(facts, admin=setting) if decision.feature == "Admin" else replace(facts, llm=setting) + ), enterprise + case "Allow" | "Deny" | "PropagateHttp": + raise NativeAuthRuntimeError("completed route enable decision requested IO") + case _: + assert_never(decision.kind) + + +def _read( + facts: RouteEnableFacts, + readers: RouteEnableReaders, + enterprise: object = None, + http_error: HTTPException | None = None, +) -> tuple[RouteEnableFacts, HTTPException | None]: + decision: Final = route_enable(facts) + if decision.kind in ("Allow", "Deny", "PropagateHttp"): + return facts, http_error + try: + updated, imported = _answer(facts, decision, readers, enterprise) + except NativeAuthFailure: + raise + except HTTPException as error: + return _read(replace(facts, failed="Http"), readers, enterprise, error) + except Exception: # noqa: BLE001 # optional enterprise import and configuration IO failures were ignored in Phase 1 + return _read(replace(facts, failed="Other"), readers, enterprise) + return _read(updated, readers, imported) + + +def check_route_enabled(route: str, *, readers: RouteEnableReaders = _READERS) -> None: + facts, error = _read(RouteEnableFacts(route_context(route)), readers) + decision: Final = route_enable(facts) + if decision.kind == "Allow": + return + if decision.kind == "PropagateHttp" and error is not None: + raise error + if decision.kind != "Deny" or decision.feature is None or decision.reason is None: + raise NativeAuthRuntimeError("incomplete route enable decision") + if decision.reason == "License": + label: Final = "ADMIN" if decision.feature == "Admin" else "LLM API" + raise HTTPException( + 500, + f"\U0001f6a8\U0001f6a8\U0001f6a8 DISABLING {label} ENDPOINTS is an Enterprise feature\n\U0001f6a8 {CommonProxyErrors.not_premium_user.value}", + ) + disabled_label: Final = "Management" if decision.feature == "Admin" else "LLM API" + raise HTTPException(403, f"{disabled_label} routes are disabled for this instance.") diff --git a/litellm/proxy/auth/route_policy_adapter.py b/litellm/proxy/auth/route_policy_adapter.py new file mode 100644 index 00000000000..459798b5568 --- /dev/null +++ b/litellm/proxy/auth/route_policy_adapter.py @@ -0,0 +1,253 @@ +from collections.abc import Collection, Iterable, Mapping +from dataclasses import replace +from typing import Final, Literal + +from fastapi import HTTPException, Request +from pydantic import BaseModel, Field, TypeAdapter + +from litellm._auth import routes +from litellm.proxy._types import LiteLLM_UserTable, LiteLLMRoutes, UserAPIKeyAuth + +_STRINGS: Final = TypeAdapter(tuple[str, ...]) + + +class _TokenFields(BaseModel, frozen=True): + allowed_routes: object = None + metadata: Mapping[str, object] | None = Field(default_factory=dict[str, object]) + permissions: Mapping[str, object] | None = None + + +class _ProxyFields(BaseModel, frozen=True): + general_settings: Mapping[str, object] + premium_user: bool + + +def _allowed_routes(value: object) -> tuple[str, ...] | None: + return _STRINGS.validate_python(value) if isinstance(value, list) else None + + +def route_catalog() -> routes.RouteCatalog: + from litellm.proxy.agent_endpoints.auth.managed_authorization import ( + _MANAGED_LIVE_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_MCP_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_MODEL_PATHS, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_MODEL_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _MANAGED_REALTIME_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + ) + from litellm.proxy.auth.route_checks import ( + _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _SELF_SCOPED_CREDENTIAL_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + _UNBILLED_ROUTES, # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + RouteChecks, + ) + + return routes.RouteCatalog( + groups={name: tuple(member.value) for name, member in LiteLLMRoutes.__members__.items()}, + viewer_blocked=tuple(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES), + viewer_write_blocked=tuple(RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES), # pyright: ignore[reportPrivateUsage] # read-only legacy catalog projection + self_credentials=_SELF_SCOPED_CREDENTIAL_ROUTES, + managed_models=tuple(_MANAGED_MODEL_ROUTES), + managed_paths=_MANAGED_MODEL_PATHS, + managed_realtime=tuple(_MANAGED_REALTIME_ROUTES), + managed_live=tuple(_MANAGED_LIVE_ROUTES), + managed_mcp=_MANAGED_MCP_ROUTES, + unbilled=tuple(_UNBILLED_ROUTES), + ) + + +def route_context(route: str, method: str | None = None, *, registration: bool = False) -> routes.RouteContext: + if not registration: + return routes.RouteContext(route=route, method=method, catalog=route_catalog()) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered: Final = InitPassThroughEndpointHelpers.get_registered_pass_through_route(route=route, method=method) + return routes.RouteContext( + route=route, + method=method, + catalog=route_catalog(), + registered=InitPassThroughEndpointHelpers.is_registered_pass_through_route(route), + registered_auth=registered.get("auth") is True if registered is not None else None, + configured_eu_jev=InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + "/eu_jev/v1/systemone" + ), + ) + + +def passthrough_facts( + route: str, token: UserAPIKeyAuth, jwt_team_routes: Collection[str] = (), *, include_denied: bool = False +) -> routes.PassThroughFacts: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + metadata: Final = _TokenFields.model_validate(token, from_attributes=True).metadata or {} + return routes.PassThroughFacts( + key_allowed=_STRINGS.validate_python(metadata.get("allowed_passthrough_routes") or ()), + team_allowed=_STRINGS.validate_python((token.team_metadata or {}).get("allowed_passthrough_routes") or ()), + denied=( + *_STRINGS.validate_python(metadata.get("denied_passthrough_routes") or ()), + *_STRINGS.validate_python((token.team_metadata or {}).get("denied_passthrough_routes") or ()), + ) + if include_denied + else (), + forwarded=InitPassThroughEndpointHelpers.forwarded_routes(route) if include_denied else (), + jwt_team_routes=tuple(jwt_team_routes), + ) + + +def raise_route_denial( + decision: routes.RouteDecision, + route: str, + *, + token: UserAPIKeyAuth | None = None, + role: str | None = None, + user: LiteLLM_UserTable | None = None, + query_user_id: str | None = None, +) -> None: + from litellm.proxy.auth.route_checks import RouteChecks + + match decision.denial: + case None: + return + case "AdminOnly": + RouteChecks._raise_admin_only_route_exception(user_obj=user, route=route) # pyright: ignore[reportPrivateUsage] # preserve the public exception envelope + case "PassThrough": + raise RouteChecks._auth_pass_through_denied_exception(route) # pyright: ignore[reportPrivateUsage] # preserve the public exception envelope + case "DeniedPassThrough": + raise RouteChecks.passthrough_route_denied_exception(route, decision.detail or "") + case "VirtualRoute": + raise HTTPException( + 403, + f"Virtual key is not allowed to call this route. Only allowed to call routes: {_TokenFields.model_validate(token, from_attributes=True).allowed_routes if token is not None else None}. Tried to call route: {route}", + ) + case "PasswordReset": + raise HTTPException( + 403, + "This account's password must be changed before the session can be used: " + "it was either found in a known data breach or set by an admin. " + "Change it via POST /user/password/change (UI: /ui/change-password), then log in again.", + ) + case "CustomAdmin": + raise HTTPException(403, f"user not allowed to access this route. Route={route} is an admin only route") + case "UserInfo": + raise HTTPException( + 403, + f"key not allowed to access this user's info. user_id={query_user_id}, key's user_id={token.user_id if token is not None else None}", + ) + case "ViewerInference": + raise HTTPException(403, f"user not allowed to access this OpenAI routes, role= {role}") + case "ViewerWrite": + raise HTTPException(403, f"user not allowed to access this route, role= {role}. Trying to access: {route}") + case "ViewerParameter": + raise HTTPException( + 403, + f"user not allowed to access this route, role= {role}. Trying to access: {route} and updating invalid param: {decision.detail}. only user_email can be updated", + ) + + +def check_virtual_route(route: str, token: UserAPIKeyAuth, request: Request | None) -> bool: + from litellm.proxy.auth.route_checks import RouteChecks + + fields: Final = _TokenFields.model_validate(token, from_attributes=True) + facts: Final = 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) + return True + + +def check_viewer_route(route: str, role: str, request_data: Mapping[str, object], request: Request | None) -> None: + from litellm.proxy.auth.route_checks import RouteChecks + + context: Final = route_context(route, RouteChecks.get_request_method(request), registration=True) + facts: Final = routes.ViewerFacts( + replace(context, method="") if request is not None and context.method is None else context, tuple(request_data) + ) + raise_route_denial(routes.viewer_route(facts), route, role=role) + + +def check_non_admin_route( + user: LiteLLM_UserTable | None, + role: str | None, + route: str, + request: Request, + token: UserAPIKeyAuth, + request_data: Mapping[str, object], +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.route_checks import RouteChecks + + proxy: Final = _ProxyFields.model_validate(vars(proxy_server)) + fields: Final = _TokenFields.model_validate(token, from_attributes=True) + singular: Final = request_data.get("organization_id") + plural: Final = request_data.get("organizations") + organization_ids: Final = ( + *((singular,) if isinstance(singular, str) else ()), + *(_STRINGS.validate_python(plural) if isinstance(plural, list) else ()), + ) + context: Final = route_context(route, RouteChecks.get_request_method(request), registration=True) + facts: Final = routes.NonAdminFacts( + viewer=routes.ViewerFacts( + replace(context, method="") if context.method is None else context, tuple(request_data) + ), + role=role, + custom_admin_routes=_STRINGS.validate_python(proxy.general_settings.get("admin_only_routes") or ()), + premium=proxy.premium_user is True, + query_user_id=request.query_params.get("user_id") if route == "/user/info" else None, + token_user_id=token.user_id, + permissions=tuple(fields.permissions or ()), + organization_ids=organization_ids, + admin_organization_ids=tuple( + membership.organization_id + for membership in (user.organization_memberships or ()) + if membership.user_role == "org_admin" + ) + if user is not None + else (), + team_service_account=bool(token.is_team_service_account), + allowed_routes=_allowed_routes(fields.allowed_routes), + passthrough=passthrough_facts(route, token, RouteChecks._jwt_team_allowed_routes(token), include_denied=True), # pyright: ignore[reportPrivateUsage] # project legacy JWT team route metadata + ) + raise_route_denial( + routes.non_admin_route(facts), route, token=token, role=role, user=user, query_user_id=facts.query_user_id + ) + + +def classify( + route: str, + kind: Literal["Llm", "Unbilled", "Discovery", "SelfCredential", "Azure", "Managed"], + method: str | None = None, +) -> bool: + from litellm._auth.route_enable import llm_route + + context: Final = route_context(route, method) + if kind != "Llm": + return routes.classify_route(context, kind) + decision: Final = llm_route(context) + if decision == "ReadEuJev": + return llm_route(context, read_eu_jev_configuration()) == "Allow" + return decision == "Allow" + + +def read_eu_jev_configuration() -> bool: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + return InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash("/eu_jev/v1/systemone") + + +def denied_passthrough(route: str, metadata_sources: Iterable[Mapping[str, object] | None]) -> str | None: + from itertools import chain + + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + facts: Final = routes.PassThroughFacts( + denied=tuple( + chain.from_iterable( + _STRINGS.validate_python((source or {}).get("denied_passthrough_routes") or ()) + for source in metadata_sources + ) + ), + forwarded=InitPassThroughEndpointHelpers.forwarded_routes(route), + ) + return routes.denied_passthrough_route(route, facts).detail diff --git a/litellm/proxy/auth/route_role_adapter.py b/litellm/proxy/auth/route_role_adapter.py new file mode 100644 index 00000000000..238a5a53d1c --- /dev/null +++ b/litellm/proxy/auth/route_role_adapter.py @@ -0,0 +1,163 @@ +from collections.abc import Iterable +from dataclasses import replace +from itertools import chain +from typing import Final, NoReturn, assert_never + +from fastapi import Request +from pydantic import TypeAdapter +from starlette.routing import compile_path + +from litellm._auth import route_roles as core +from litellm._auth.failure import NativeAuthFailure +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + UserAPIKeyAuth, +) + +_STRINGS: Final = TypeAdapter(tuple[str, ...]) +_END: Final = object() + + +def _role(user: LiteLLM_UserTable | None) -> core.RoleFacts: + return core.RoleFacts(user is not None, user.user_role if user is not None else None) + + +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + return LitellmUserRoles(core.role(core.RoleFacts(True, user_role)).effective) + + +def get_user_role(user_obj: LiteLLM_UserTable | None) -> LitellmUserRoles | None: + role: Final = core.role(_role(user_obj)).user_role + return LitellmUserRoles(role) if role is not None else None + + +def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: + return core.role(_role(user_obj)).admin + + +def is_api_route_allowed( + route: str, + request: Request, + request_data: dict[str, object], # mutable-ok: pass the original request dictionary to the shared route wrapper + valid_token: UserAPIKeyAuth | None, + user_obj: LiteLLM_UserTable | None = None, +) -> bool: + from litellm.proxy.auth.route_checks import RouteChecks + + role: Final = get_user_role(user_obj) + initial: Final = core.ApiFacts(valid_token is not None, False, core.RoleFacts(False, None)) + decision: Final = core.api(initial) + if decision == "InvalidToken": + raise Exception("Invalid proxy server token passed. valid_token=None.") # noqa: TRY002 # preserve the public legacy exception class + result: Final = core.api(replace(initial, admin_read=True, user=_role(user_obj))) + if result == "NonAdmin": + assert valid_token is not None + RouteChecks.non_proxy_admin_allowed_routes_check( # pyright: ignore[reportUnknownMemberType] # the shared wrapper still declares an unparameterized request dictionary + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + request_data=request_data, + valid_token=valid_token, + ) + return True + + +def _entry(raw: object) -> core.Entry: + try: + hash(raw) + except NativeAuthFailure: + raise + except TypeError as error: + return core.Entry(None, type(raw).__name__, hash_failure=str(error)) + text: Final = raw if isinstance(raw, str) else None + group: Final = LiteLLMRoutes.__members__.get(text) if text is not None else None + return core.Entry( + text, + type(raw).__name__, + templates=_STRINGS.validate_python(group.value) if group is not None else None, # pyright: ignore[reportUnknownMemberType] # enum groups are legacy route collections + ) + + +def _raise_failure(failure: core.Failure | None) -> NoReturn: + assert failure is not None + match failure.kind: + case "TypeError": + raise TypeError(failure.message) + case "AttributeError": + raise AttributeError(failure.message) + case _: + assert_never(failure.kind) + + +def allowed_routes_list_check(user_route: str, allowed_routes: Iterable[object]) -> bool: + entries: Final = iter(allowed_routes) + facts = core.ListedFacts(user_route) # rebind-ok: native-selected data-read driver state + while True: + decision = core.listed(facts) + match decision.action: + case "ReadEntry": + raw = next(entries, _END) + facts = core.ListedFacts(user_route, None if raw is _END else _entry(raw), raw is _END) + case "ReadTemplate": + assert facts.entry is not None and facts.entry.templates is not None + compiled, _, _ = compile_path(facts.entry.templates[decision.template_index]) + facts = replace(facts, template_index=decision.template_index, compiled=compiled.pattern) + case "Allow": + return True + case "Deny": + return False + case "Failure": + _raise_failure(decision.failure) + case _: + raise AssertionError(f"Unexpected listed-route directive: {decision.action}") + + +def allowed_routes_check( + user_role: LitellmUserRoles, + user_route: str, + litellm_proxy_roles: LiteLLM_JWTAuth, +) -> bool: + facts: Final = core.JwtFacts(user_role) + decision: Final = core.jwt(facts) + match decision: + case "Deny": + return False + case "ReadAdmin": + return allowed_routes_list_check(user_route, litellm_proxy_roles.admin_allowed_routes) + case "ReadTeam": + configured: Final = litellm_proxy_roles.team_allowed_routes + selected: Final = core.jwt(replace(facts, config_read=True, config_missing=configured is None)) # pyright: ignore[reportUnnecessaryComparison] # constructed legacy JWT configs retain the None default branch + return allowed_routes_list_check( + user_route, ("openai_routes", "info_routes") if selected == "DefaultTeam" else configured + ) + case _: + raise AssertionError(f"Unexpected JWT route directive: {decision}") + + +def allowed_route_check_inside_route(user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None) -> bool: + return core.inside(core.InsideFacts(user_api_key_dict.user_role, user_api_key_dict.user_id, requested_user_id)) + + +def _expanded(raw: object) -> tuple[object, ...]: + facts: Final = _entry(raw) + decision: Final = core.expand(facts) + match decision.action: + case "Expand": + assert facts.templates is not None + return facts.templates + case "Keep": + return (raw,) + case "Failure": + _raise_failure(decision.failure) + case _: + raise AssertionError(f"Unexpected route expansion directive: {decision.action}") + + +def get_actual_routes( + allowed_routes: Iterable[object], +) -> list[object]: # mutable-ok: preserve the public list projection + return list(chain.from_iterable(_expanded(raw) for raw in allowed_routes)) diff --git a/litellm/proxy/auth/spend_counter_adapter.py b/litellm/proxy/auth/spend_counter_adapter.py new file mode 100644 index 00000000000..9985d5c5ab0 --- /dev/null +++ b/litellm/proxy/auth/spend_counter_adapter.py @@ -0,0 +1,34 @@ +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, replace +from typing import Final, NoReturn + +from litellm._auth import spend_counter_policy as core +from litellm._auth.failure import NativeAuthRuntimeError + + +@dataclass(frozen=True, slots=True) +class SpendIO: + read_strict: Callable[[], bool] + read_floor: Callable[[], Awaitable[float | None]] + repair: Callable[[float], Awaitable[None]] + unavailable: Callable[[], NoReturn] + + +async def resolve_spend(facts: core.Spend, io: SpendIO) -> float: + decision: Final = core.decide(facts) + match decision: + case core.SpendValue(kind="Value", spend=spend): + return spend + case core.SpendValue(kind="Repair", spend=spend): + await io.repair(spend) + return spend + case core.Action(kind="ReadStrictBefore"): + return await resolve_spend(replace(facts, strict_before=io.read_strict()), io) + case core.Action(kind="ReadStrictAfter"): + return await resolve_spend(replace(facts, strict_after=io.read_strict()), io) + case core.Action(kind="ReadFloor"): + return await resolve_spend(replace(facts, authoritative_read=True, authoritative=await io.read_floor()), io) + case core.Action(kind="Unverifiable"): + io.unavailable() + case _: + raise NativeAuthRuntimeError("Unexpected spend-counter policy decision") diff --git a/litellm/proxy/auth/team_alias_adapter.py b/litellm/proxy/auth/team_alias_adapter.py new file mode 100644 index 00000000000..1c83caf1dba --- /dev/null +++ b/litellm/proxy/auth/team_alias_adapter.py @@ -0,0 +1,33 @@ +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError, native_boundary +from litellm._auth.team_aliases import TeamAliasFacts, team_aliases +from litellm.router_api import RegistryReader + + +def resolve_team_alias( + model: str | Sequence[str], + team_model_aliases: Mapping[str, str] | None, + team_id: str | None, + llm_router: RegistryReader | None, +) -> str | list[str]: # mutable-ok: the existing public wrapper preserves the input scalar/list shape + initial: Final = TeamAliasFacts( + (model,) if isinstance(model, str) else tuple(model), team_model_aliases or {}, team_id, llm_router is not None + ) + decision: Final = team_aliases(initial) + if decision.kind == "Complete": + return decision.models[0] if isinstance(model, str) else list(decision.models) + if llm_router is None: + raise NativeAuthRuntimeError("team aliases requested names without a router") + names: Final = tuple(_names(llm_router.registry().index("model_name")) for _ in range(decision.count)) + completed: Final = team_aliases(replace(initial, names=names)) + if completed.kind != "Complete": + raise NativeAuthRuntimeError("incomplete team alias resolution") + return completed.models[0] if isinstance(model, str) else list(completed.models) + + +@native_boundary +def _names(names: Iterable[str]) -> tuple[str, ...]: + return tuple(names) diff --git a/litellm/proxy/auth/team_override_adapter.py b/litellm/proxy/auth/team_override_adapter.py new file mode 100644 index 00000000000..ee0ed0c1d3b --- /dev/null +++ b/litellm/proxy/auth/team_override_adapter.py @@ -0,0 +1,102 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.team_override import TeamOverrideFacts, team_override +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.proxy.auth.model_projection_adapter import read_grouped_scope +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + token: UserAPIKeyAuth | None + team: LiteLLM_TeamTable | None + router: RegistryReader | None + team_aliases: Mapping[str, str] | None + key_aliases: Mapping[str, str] | None + prisma: DatabaseClient | None + + +async def read_team_override_facts( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> TeamOverrideFacts: + return await _read( + TeamOverrideFacts(), + _Source(model, valid_token, team_object, llm_router, team_model_aliases, key_model_aliases, prisma_client), + ) + + +async def can_team_call_model_with_override( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + facts: Final = await read_team_override_facts( + model, valid_token, team_object, llm_router, team_model_aliases, key_model_aliases, prisma_client + ) + decision: Final = team_override(facts) + if decision.kind == "Allow": + return True + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete team override authorization") + return raise_model_denial(decision.denial, model) + + +async def _read(facts: TeamOverrideFacts, source: _Source) -> TeamOverrideFacts: + + decision: Final = team_override(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadTeam": + scope: Final = project_model_facts( + source.model, + source.router, + object_model_grants(source.team) if source.team is not None else (), + source.team_aliases, + source.team.team_id if source.team is not None else None, + source.key_aliases, + "team", + ) + groups: Final = tuple(source.team.access_group_ids or ()) if source.team is not None else () + return await _read(replace(facts, team=GroupedModelScope(scope, groups)), source) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing team override scope") + grouped: Final = await read_grouped_scope(facts.team, source.router, source.prisma, merge=True) + return await _read(replace(facts, team=grouped), source) + case "ReadAssignedGroups": + assigned: Final = await get_authorized_resources_from_key_access_groups( + source.token, source.team, "access_model_names" + ) + return await _read(replace(facts, assigned_models=tuple(assigned)), source) + case "ReadAssignedScope": + if facts.assigned_models is None: + raise NativeAuthRuntimeError("missing assigned model override grants") + assigned_scope: Final = project_model_facts( + source.model, + source.router, + facts.assigned_models, + source.token.team_model_aliases if source.token is not None else None, + source.token.team_id if source.token is not None else None, + source.key_aliases, + "key", + ) + return await _read(replace(facts, assigned_scope=assigned_scope), source) diff --git a/litellm/proxy/auth/unpriced_model_adapter.py b/litellm/proxy/auth/unpriced_model_adapter.py new file mode 100644 index 00000000000..7ed766df168 --- /dev/null +++ b/litellm/proxy/auth/unpriced_model_adapter.py @@ -0,0 +1,128 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from pydantic import BaseModel, Field, TypeAdapter + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.unpriced_models import DeploymentPricing, PriceEntry, PriceMetric, UnpricedModelFacts, unpriced_model +from litellm.router_api import RegistryReader + +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OPTIONAL_STRING: Final = TypeAdapter[str | None](str | None) + + +class _ModelInfo(BaseModel, frozen=True): + id: str | None = None + + +class _Deployment(BaseModel, frozen=True): + litellm_params: Mapping[str, object] = Field(default_factory=dict[str, object]) + model_info: _ModelInfo = Field(default_factory=_ModelInfo) + + +class _CostMap(BaseModel, frozen=True): + model_cost: Mapping[str, object] + + +def _number(value: object) -> float | None: + return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None + + +def _metric(name: str, value: object) -> PriceMetric: + return PriceMetric( + name, + _number(value), + tuple(_number(item) for item in _MAPPING.validate_python(value).values()) if isinstance(value, dict) else (), + ) + + +def price_entry(entry: Mapping[str, object]) -> PriceEntry: + return PriceEntry( + entry.get("tiered_pricing") is not None, tuple(_metric(key, value) for key, value in entry.items()) + ) + + +def read_unpriced_model_facts(model: str | None, llm_router: RegistryReader | None) -> UnpricedModelFacts: + return _read(UnpricedModelFacts(model, llm_router is not None), llm_router, ()) + + +def _read(facts: UnpricedModelFacts, router: RegistryReader | None, rows: Sequence[_Deployment]) -> UnpricedModelFacts: + decision: Final = unpriced_model(facts) + if decision.kind == "Complete": + return facts + if router is None or facts.model is None: + raise NativeAuthRuntimeError("unpriced model read without a router or model") + match decision.kind: + case "ReadGroup": + found: Final = router.registry().group(facts.model) is not None + return _read(replace(facts, group_found=found), router, rows) + case "ReadDeployments" | "ReadCostDeployments": + deployments: Final = tuple( + _Deployment.model_validate(row) for row in router.registry().deployments(facts.model) or () + ) + projected: Final = tuple( + DeploymentPricing(tuple(row.litellm_params), row.model_info.id is not None) for row in deployments + ) + snapshot: Final = ( + replace(facts, deployments=projected) + if decision.kind == "ReadDeployments" + else replace(facts, cost_deployments=projected) + ) + return _read(snapshot, router, deployments) + case "ReadInfo" | "ReadCost": + if facts.deployments is None: + raise NativeAuthRuntimeError("missing unpriced deployment facts") + row: Final = rows[decision.index] + model_id: Final = row.model_info.id + if model_id is None: + raise NativeAuthRuntimeError("missing unpriced deployment identifier") + current: Final = ( + facts.deployments[decision.index] + if decision.kind == "ReadInfo" + else (facts.cost_deployments or ())[decision.index] + ) + if decision.kind == "ReadInfo": + info: Final = router.registry().deployment_model_info( + model_id, _OPTIONAL_STRING.validate_python(row.litellm_params.get("model")) or "" + ) + updated: Final = replace( + current, + info_loaded=True, + info=price_entry(_MAPPING.validate_python(info)) if info is not None else None, + ) + return _read( + replace( + facts, + deployments=tuple( + updated if index == decision.index else item for index, item in enumerate(facts.deployments) + ), + ), + router, + rows, + ) + keys: Final = tuple( + _MAPPING.validate_python(_CostMap.model_validate(vars(litellm)).model_cost.get(model_id, {})) + ) + cost: Final = replace(current, cost_keys=keys) + return _read( + replace( + facts, + cost_deployments=tuple( + cost if index == decision.index else item + for index, item in enumerate(facts.cost_deployments or ()) + ), + ), + router, + rows, + ) + + +def model_has_no_cost_mapping(model: str | None, llm_router: RegistryReader | None) -> bool: + return unpriced_model(read_unpriced_model_facts(model, llm_router)).unpriced + + +def unpriced_models_in_request(model: str | Sequence[str] | None, llm_router: RegistryReader | None) -> tuple[str, ...]: + names: Final = (model,) if isinstance(model, str) else tuple(model or ()) + return tuple(name for name in names if model_has_no_cost_mapping(name, llm_router)) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a97335f5061..50f09a7373b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -9,11 +9,12 @@ Returns a UserAPIKeyAuth object if the API key is valid import asyncio import fnmatch -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Mapping, Sequence +from dataclasses import replace from datetime import datetime, timezone from enum import Enum from types import MappingProxyType -from typing import Final, NamedTuple, Protocol, Union, assert_never, cast +from typing import Final, Literal, NamedTuple, Protocol, Union, assert_never, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -22,6 +23,9 @@ from pydantic import TypeAdapter from starlette.exceptions import WebSocketException import litellm +from litellm._auth import access as access_policy +from litellm._auth import admission, auth_context, common_flow +from litellm._auth.budget_pricing import ExemptionFacts from litellm._auth.credential import ( PresentedKeys, WebSocketDenial, @@ -37,7 +41,7 @@ from litellm._auth.decision import ( Grant, KeyFacts, MasterKey, - ModelBudgetExceeded, + ModelBudgetRead, ModelBudgetScope, NoMasterKey, NonAdminKey, @@ -54,7 +58,13 @@ from litellm._auth.decision import ( WithinModelBudgets, ) from litellm._auth.driver import drive, drive_model_budgets, drive_team_member_model_budgets +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_and_fallback_model_access +from litellm._auth.model_budget_data import ModelBudgetEntity, ModelBudgetSnapshot from litellm._auth.runtime import get_auth_core, scheme_credential, select_api_key, select_websocket_api_key +from litellm._auth.rust_core import expired from litellm._internal_context import service_target, set_request_parent_otel_span from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging @@ -83,18 +93,11 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo _get_user_role, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _is_user_proxy_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - _model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # reuse common-check group budget enforcement _team_member_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - _unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # reuse common-check pricing enforcement _virtual_key_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _virtual_key_max_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _virtual_key_soft_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export cache_key_object, - can_customer_access_model, - can_key_call_model, - can_object_call_model, - can_project_access_model, - can_team_access_model, check_end_user_budget, common_checks, delete_cache_key_object, @@ -105,25 +108,18 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo get_object_permission, get_org_object_for_request, get_project_object, - get_team_member_default_budget, get_team_membership, get_team_object, get_user_object, get_user_role, is_model_cost_zero, is_user_proxy_admin, - is_valid_fallback_model, jwt_key_mapping_cache_key, - key_access_group_grants_model, - key_model_aliases_for_auth_check, - model_max_budget_of, request_skips_budget_checks, resolve_and_validate_end_user_id, - resolve_default_end_user_budget, route_skips_budget_checks, send_soft_budget_alert, send_token_budget_alert, - stamp_matched_model_access_groups, team_member_max_budget_alert_check, virtual_key_max_budget_alert_check, virtual_key_max_budget_check, @@ -140,13 +136,11 @@ from litellm.proxy.auth.auth_object_prefetch import ( ) from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, - fallback_target_model_name, get_end_user_id_from_request_body, get_model_from_request, get_request_route, get_request_route_template, is_invalid_virtual_key_error, - iter_request_fallback_targets, log_model_access_denial, normalize_request_route, pre_db_read_auth_checks, @@ -155,7 +149,11 @@ from litellm.proxy.auth.auth_utils import ( route_in_additonal_public_routes, ) from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, should_throttle_budget_exceeded +from litellm.proxy.auth.fallback_access_adapter import load_fallback_authorization_facts from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.auth.key_model_access_adapter import read_key_and_fallback_model_access_facts +from litellm.proxy.auth.model_access_adapter import raise_model_denial +from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request @@ -446,7 +444,19 @@ def _budget_fallback_target_model(target: object) -> str | None: return model if isinstance(model, str) else None -class _KeyAndUserModelBudgetLimiter(Protocol): +class _RawModelBudgetLimiter(Protocol): + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> ModelBudgetReading: ... + + +class _KeyAndUserModelBudgetLimiter(_RawModelBudgetLimiter, Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... async def is_user_within_model_budget( @@ -474,7 +484,7 @@ class _EndUserModelBudgetLimiter(Protocol): ) -> bool: ... -class _EveryModelBudgetLimiter(Protocol): +class _EveryModelBudgetLimiter(_RawModelBudgetLimiter, Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... async def is_user_within_model_budget( @@ -526,21 +536,6 @@ def _token_team_models(valid_token: _TokenTeamModels) -> list[str]: return valid_token.team_models -class _TokenModelBudgetFields(Protocol): - @property - def model_max_budget(self) -> Mapping[str, object] | None: ... - - @property - def end_user_model_max_budget(self) -> Mapping[str, object] | None: ... - - @property - def team_model_aliases(self) -> dict[str, str] | None: ... # mutable-ok: mirrors _can_object_call_model's contract - - -def _token_model_budget_fields(valid_token: _TokenModelBudgetFields) -> _TokenModelBudgetFields: - return valid_token - - def _router_fallbacks(router: litellm.Router) -> Sequence[object] | None: fallbacks: Final[Sequence[object] | None] = router.settings().fallbacks return fallbacks @@ -569,6 +564,8 @@ async def _read_user_model_max_budget( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # mirrors the main path's tolerance verbose_logger.debug("Unable to read user for the per-model budget check: %s", e) return None @@ -580,18 +577,9 @@ async def _check_team_model_budget( model_max_budget_limiter: _TeamModelBudgetLimiter, models: list[str], ) -> None: - """Enforce the team's `model_max_budget` for every requested model the key does not override.""" - team_model_max_budget: Final = valid_token.team_model_max_budget - if valid_token.team_id is None or not team_model_max_budget: - return - key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget - for model_name in models: - await model_max_budget_limiter.is_team_within_model_budget( - team_id=valid_token.team_id, - team_model_max_budget=team_model_max_budget, - key_model_max_budget=key_model_max_budget, - model=model_name, - ) + from litellm.proxy.auth import budget_context_adapter + + return await budget_context_adapter.check_team_model_budget(valid_token, model_max_budget_limiter, models) async def _end_user_model_max_budget( @@ -602,21 +590,11 @@ async def _end_user_model_max_budget( user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, ) -> Mapping[str, object] | None: - budget_table: Final = ( - end_user_object.litellm_budget_table - if end_user_object is not None - else ( - await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, - ) - if end_user_id and prisma_client is not None - else None - ) + from litellm.proxy.auth import budget_context_adapter + + return await budget_context_adapter.end_user_model_max_budget( + end_user_id, end_user_object, key_end_user_budget_id, prisma_client, user_api_key_cache, parent_otel_span ) - return model_max_budget_of(budget_table) if budget_table is not None else None async def _check_end_user_model_budget( @@ -626,184 +604,54 @@ async def _check_end_user_model_budget( models: Sequence[str], team_id: str | None, ) -> None: - if not end_user_id or not end_user_model_max_budget: - return - for model_name in models: - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=end_user_id, - end_user_model_max_budget=end_user_model_max_budget, - model=model_name, - team_id=team_id, - ) + from litellm.proxy.auth import budget_context_adapter + + return await budget_context_adapter.check_end_user_model_budget( + end_user_id, end_user_model_max_budget, model_max_budget_limiter, models, team_id + ) -async def _is_within_every_model_budget( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - model: str, -) -> bool: - budget_fields: Final = _token_model_budget_fields(valid_token) - key_model_max_budget: Final = budget_fields.model_max_budget - user_model_max_budget: Final = valid_token.user_model_max_budget - end_user_model_max_budget: Final = budget_fields.end_user_model_max_budget - team_model_max_budget: Final = valid_token.team_model_max_budget - team_member_model_max_budget: Final = valid_token.team_member_model_max_budget - try: - if valid_token.token is not None and isinstance(key_model_max_budget, Mapping) and key_model_max_budget: - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=model, - ) - if valid_token.user_id is not None and isinstance(user_model_max_budget, Mapping) and user_model_max_budget: - await model_max_budget_limiter.is_user_within_model_budget( - user_id=valid_token.user_id, - user_model_max_budget=user_model_max_budget, - model=model, - team_id=valid_token.team_id, - ) - if ( - valid_token.end_user_id is not None - and isinstance(end_user_model_max_budget, Mapping) - and end_user_model_max_budget - ): - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_model_max_budget, - model=model, - team_id=valid_token.team_id, - ) - if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget: - await model_max_budget_limiter.is_team_within_model_budget( - team_id=valid_token.team_id, - team_model_max_budget=team_model_max_budget, - key_model_max_budget=key_model_max_budget, - model=model, - ) - if ( - valid_token.user_id is not None - and valid_token.team_id is not None - and isinstance(team_member_model_max_budget, Mapping) - and team_member_model_max_budget - ): - await model_max_budget_limiter.is_team_member_within_model_budget( - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_member_model_max_budget=team_member_model_max_budget, - model=model, - ) - except litellm.BudgetExceededError: - return False - return True +def _model_budget_read( + reading: ModelBudgetReading, message: str, entity_type: str, entity_id: str | None +) -> ModelBudgetRead: + return ModelBudgetRead( + reading.snapshot, + message if reading.named_model else f"{message}, budget_model={reading.budget_model}", + entity_type, + entity_id, + ) -async def _is_fallback_model_authorized( - model: str, - valid_token: UserAPIKeyAuth, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router | None, - prisma_client: PrismaClient | None = None, - team_membership: LiteLLM_TeamMembership | None = None, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - refuse_agent_keys: bool = False, - team_object: LiteLLM_TeamTable | None = None, -) -> bool: - if refuse_agent_keys and valid_token.agent_id is not None: - return False - try: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - prisma_client=prisma_client, - ) - if team_object is not None: - try: - await can_team_access_model( - model=model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - prisma_client=prisma_client, - ) - except ProxyException as team_denial: - if team_denial.type != ProxyErrorTypes.team_model_access_denied: - raise - if not await key_access_group_grants_model( - model=model, - valid_token=valid_token, - team_object=team_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ): - raise - else: - team_models: Final = _token_team_models(valid_token) - if team_models: - can_object_call_model( - model=model, - llm_router=llm_router, - models=team_models, - team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="team", - ) - member_allowed_models: Final = ( - team_membership.litellm_budget_table.allowed_models - if team_membership is not None and team_membership.litellm_budget_table is not None - else None - ) - if member_allowed_models: - can_object_call_model( - model=model, - llm_router=llm_router, - models=member_allowed_models, - team_id=valid_token.team_id, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - object_type="team", - ) - if project_object is not None and len(project_object.models) > 0: - can_project_access_model( - model=model, - project_object=project_object, - llm_router=llm_router, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - except ProxyException: - return False - return True +def _model_budget_scope_enabled(scope: Literal["Key", "User", "Member"], token: UserAPIKeyAuth, budget: object) -> bool: + facts = common_flow.ModelBudgets(scope) # rebind-ok: native-selected raw field read stages + while True: + match common_flow.read(facts): + case "ReadModelBudgetPresence": + facts = replace(facts, budget_present=bool(budget)) + case "ReadModelBudgetShape": + facts = replace(facts, mapping=isinstance(budget, Mapping)) + case "ReadModelBudgetToken": + facts = replace(facts, entity_present=token.token is not None) + case "ReadModelBudgetUser": + facts = replace(facts, entity_present=token.user_id is not None) + case "ReadModelBudgetTeam": + facts = replace(facts, team_present=token.team_id is not None) + case "CheckModelBudgets": + return True + case "Continue": + return False + case _: + raise NativeAuthRuntimeError("Unexpected model-budget scope read") -async def _is_within_all_fallback_model_budgets( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - models: Sequence[str], -) -> bool: - for model in models: - if not await _is_within_every_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model=model, - ): - return False - return True - - -async def _model_budget_verdict(check: Awaitable[bool]) -> ModelBudgetExceeded | None: - try: - with tracer.trace("litellm.proxy.auth.budget_checks"): - _ = await check - except litellm.BudgetExceededError as e: - return ModelBudgetExceeded( - message=e.message, - current_cost=e.current_cost, - max_budget=e.max_budget, - entity_type=e.entity_type, - entity_id=e.entity_id, - ) - return None +def _model_budget_fallback_enabled(router_present: bool, data: Mapping[str, object]) -> bool: + initial: Final = common_flow.FallbackAllowed(router_present) + facts: Final = ( + replace(initial, disabled=fallbacks_disabled(data)) + if common_flow.read(initial) == "ReadFallbackDisabled" + else initial + ) + return common_flow.read(facts) == "AllowFallback" class _ModelBudgetEffects: @@ -851,41 +699,45 @@ class _ModelBudgetEffects: ) ), rewritable_model=requested if isinstance(requested, str) else None, - key_has_model_budgets=bool(self._valid_token.model_max_budget) and self._valid_token.token is not None, + key_has_model_budgets=_model_budget_scope_enabled( + "Key", self._valid_token, self._valid_token.model_max_budget + ), key_budget_fallbacks=MappingProxyType({model: tuple(chain) for model, chain in budget_fallbacks.items()}), - user_has_model_budgets=self._valid_token.user_id is not None - and isinstance(user_model_max_budget, Mapping) - and bool(user_model_max_budget), + user_has_model_budgets=_model_budget_scope_enabled("User", self._valid_token, user_model_max_budget), ) - async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: - return await _model_budget_verdict( - self._limiter.is_key_within_model_budget(user_api_key_dict=self._valid_token, model=model) + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: + token: Final = self._valid_token + reading: Final = await self._limiter.read_model_budget_raw( + "Key", token.token, token.team_id, model, token.model_max_budget + ) + return _model_budget_read( + reading, + f"LiteLLM Virtual Key: {token.token}, key_alias: {token.key_alias}, exceeded budget for model={model}", + "key", + token.token, ) - async def authorize_budget_fallback(self, model: str) -> bool: - return await _is_fallback_model_authorized( + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: + return await load_fallback_authorization_facts( model=model, valid_token=self._valid_token, - llm_model_list=self._llm_model_list, llm_router=self._llm_router, ) async def rewrite_request_model(self, model: str) -> None: _reroute_request(self._request_data, self._request, model) - async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: user_id: Final = self._valid_token.user_id user_model_max_budget: Final = self._valid_token.user_model_max_budget if user_id is None or not isinstance(user_model_max_budget, Mapping): - return None - return await _model_budget_verdict( - self._limiter.is_user_within_model_budget( - user_id=user_id, - user_model_max_budget=user_model_max_budget, - model=model, - team_id=self._valid_token.team_id, - ) + return ModelBudgetRead(ModelBudgetSnapshot()) + reading: Final = await self._limiter.read_model_budget_raw( + "User", user_id, self._valid_token.team_id, model, user_model_max_budget + ) + return _model_budget_read( + reading, f"LiteLLM User: {user_id}, exceeded budget for model={model}", "user", user_id ) @@ -966,33 +818,46 @@ class _TeamMemberModelBudgetEffects: return TeamMemberModelBudgetScope( models=self._models, rewritable_model=requested if isinstance(requested, str) else None, - has_model_budgets=self._valid_token.user_id is not None - and self._valid_token.team_id is not None - and isinstance(budget, Mapping) - and bool(budget), - can_fall_back=self._llm_router is not None and not fallbacks_disabled(self._request_data), + has_model_budgets=_model_budget_scope_enabled("Member", self._valid_token, budget), + can_fall_back=_model_budget_fallback_enabled(self._llm_router is not None, self._request_data), ) - async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: user_id: Final = self._valid_token.user_id team_id: Final = self._valid_token.team_id budget: Final = self._valid_token.team_member_model_max_budget if user_id is None or team_id is None or not isinstance(budget, Mapping): - return None - return await _model_budget_verdict( - self._limiter.is_team_member_within_model_budget( - user_id=user_id, team_id=team_id, team_member_model_max_budget=budget, model=model - ) + return ModelBudgetRead(ModelBudgetSnapshot()) + from litellm.proxy.hooks.model_max_budget_limiter import team_member_budget_entity_id + + identity: Final = team_member_budget_entity_id(user_id, team_id) + reading: Final = await self._limiter.read_model_budget_raw("TeamMember", identity, team_id, model, budget) + return _model_budget_read( + reading, + f"LiteLLM Team Member: user={user_id}, team={team_id}, exceeded budget for model={model}", + "team_member", + identity, ) async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: llm_router: Final = self._llm_router - if ( - llm_router is None - or self._request.headers.get("content-type", "").lower().startswith("multipart/form-data") - or references_provider_scoped_resource(self._request_data) - ): + initial: Final = common_flow.MemberFallback(llm_router is not None) + content: Final = ( + replace( + initial, + multipart=self._request.headers.get("content-type", "").lower().startswith("multipart/form-data"), + ) + if common_flow.read(initial) == "ReadContentType" + else initial + ) + resource: Final = ( + replace(content, resource_bound=references_provider_scoped_resource(self._request_data)) + if common_flow.read(content) == "ReadResourceBinding" + else content + ) + if common_flow.read(resource) != "ReadFallbackSettings": return () + assert llm_router is not None from litellm.proxy.proxy_server import proxy_config prisma_client, _, proxy_logging_obj = self._resolved_dependencies() @@ -1007,6 +872,8 @@ class _TeamMemberModelBudgetEffects: ), ) ) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. return () @@ -1045,107 +912,22 @@ class _TeamMemberModelBudgetEffects: for candidate, fallback_model in resolved ) - async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: - llm_router: Final = self._llm_router - user_id: Final = self._valid_token.user_id - team_id: Final = self._valid_token.team_id - if llm_router is None or user_id is None or team_id is None: - return False - prisma_client, user_api_key_cache, proxy_logging_obj = self._resolved_dependencies() - try: - team_membership: Final = await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback - return False - try: - return await self._is_usable_fallback( - target, llm_router, team_membership, prisma_client, user_api_key_cache, proxy_logging_obj - ) - except Exception: # noqa: BLE001 # A candidate check failure must not escape the fail-closed walk. - verbose_proxy_logger.warning( - "Skipping fallback candidate %s after a budget or authorization check error", - target.model.replace("\r", "").replace("\n", ""), - exc_info=True, - ) - return False + async def vet_team_member_fallback(self, target: FallbackTarget) -> FallbackCandidateFacts: + from litellm.proxy.auth.fallback_candidate_adapter import read_fallback_candidate_facts - async def _is_usable_fallback( - self, - target: FallbackTarget, - llm_router: litellm.Router, - team_membership: LiteLLM_TeamMembership | None, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - ) -> bool: - if ( - litellm.block_requests_for_models_without_pricing - and RouteChecks.is_llm_api_route(route=self._route) - and _unpriced_models_in_request(model=target.served_model, llm_router=llm_router) - ): - return False - if self._end_user_object is not None and self._end_user_object.models: - can_customer_access_model( - model=target.model, - end_user_object=self._end_user_object, - llm_router=llm_router, - valid_token=self._valid_token, - ) - if not is_model_cost_zero( - model=target.served_model, llm_router=llm_router - ) and not await _is_within_all_fallback_model_budgets( + return await read_fallback_candidate_facts( + target=target, valid_token=self._valid_token, - model_max_budget_limiter=self._limiter, - models=tuple( - dict.fromkeys( - ( - target.requested, - target.model, - target.served_model, - llm_router.registry().resolve(target.served_model), - ) - ) - ), - ): - return False - if not await _is_fallback_model_authorized( - model=target.model, - valid_token=self._valid_token, - llm_model_list=self._llm_model_list, - llm_router=llm_router, + limiter=self._limiter, + llm_router=self._llm_router, + route=self._route, + project_object=self._project_object, + team_object=self._team_object, + end_user_object=self._end_user_object, prisma_client=self._prisma_client, - team_object=self._team_object, - team_membership=team_membership, - project_object=self._project_object, - refuse_agent_keys=True, - ): - return False - self._valid_token.matched_model_access_groups = [] - matched_groups: Final = await stamp_matched_model_access_groups( - model=target.model, - valid_token=self._valid_token, - team_object=self._team_object, - project_object=self._project_object, - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=True, + user_api_key_cache=self._user_api_key_cache, + proxy_logging_obj=self._proxy_logging_obj, ) - self._valid_token.matched_model_access_groups = list(matched_groups) - if matched_groups: - await _model_access_group_max_budget_check( - matched_model_access_groups=list(matched_groups), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - return True async def reroute_team_member_request(self, model: str) -> None: _reroute_request(self._request_data, self._request, model) @@ -1295,30 +1077,9 @@ def _apply_budget_limits_to_end_user_params( budget_info: LiteLLM_BudgetTable, end_user_id: str | None, ) -> None: - """ - Helper function to apply budget limits to end user parameters. + from litellm.proxy.auth import end_user_context_adapter - Args: - end_user_params: Dictionary to update with budget parameters - budget_info: Budget table object containing limits - end_user_id: ID of the end user for logging - """ - if budget_info.tpm_limit is not None: - end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit - - if budget_info.rpm_limit is not None: - end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit - - if budget_info.tpd_limit is not None: - end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit - - if budget_info.max_budget is not None: - end_user_params["end_user_max_budget"] = budget_info.max_budget - - if budget_info.model_max_budget is not None: - end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget - - verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + return end_user_context_adapter.apply_budget_limits_to_end_user_params(end_user_params, budget_info, end_user_id) _WEBSOCKET_DENIAL_DETAIL: Final = MappingProxyType( @@ -1409,6 +1170,8 @@ async def user_api_key_auth_websocket_for_model( raise HTTPException( 403, "Required baggage tag mappings need HTTP inference, WebSocket inference is disabled" ) + except NativeAuthFailure: + raise except Exception as e: if not close_on_auth_error: raise @@ -1432,22 +1195,9 @@ async def user_api_key_auth_accepted_websocket_for_model(websocket: WebSocket, m def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: - valid_token.end_user_id = end_user_params.get("end_user_id") - # Only overwrite token fields when the DB-derived value is not None. - # This prevents DB lookups (where the budget table has no value set) - # from silently clearing values that a custom auth function may have - # already set on the token. - if end_user_params.get("end_user_tpm_limit") is not None: - valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] - if end_user_params.get("end_user_rpm_limit") is not None: - valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] - if end_user_params.get("end_user_tpd_limit") is not None: - valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] - if end_user_params.get("allowed_model_region") is not None: - valid_token.allowed_model_region = end_user_params["allowed_model_region"] - if end_user_params.get("end_user_model_max_budget") is not None: - valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] - return valid_token + from litellm.proxy.auth import end_user_context_adapter + + return end_user_context_adapter.update_valid_token_with_end_user_params(valid_token, end_user_params) # Reusable coordinator for global spend to prevent cache stampede @@ -1786,6 +1536,8 @@ async def _auto_register_jwt_mapping( "updated_by": behavior.value, } ) + except NativeAuthFailure: + raise except Exception as e: error_str: Final = str(e).lower() if "unique" in error_str or "p2002" in error_str: @@ -1804,6 +1556,8 @@ async def _auto_register_jwt_mapping( try: async with db_span("delete_orphaned_jwt_key", "LiteLLM_VerificationToken"): await VerificationTokenRepository(prisma_client).table.delete(where={"token": token_hash}) + except NativeAuthFailure: + raise except Exception as delete_err: # Don't fail the request if cleanup fails — the orphan is # unmapped and inert. Log so an operator can prune it later. @@ -2107,6 +1861,8 @@ def _ensure_litellm_received_at_on_request_state(request: Request) -> datetime: received_at: Final = datetime.now(timezone.utc) try: request.state.litellm_received_at = received_at + except NativeAuthFailure: + raise except Exception: pass return received_at @@ -2189,6 +1945,8 @@ async def _record_unparsable_body_failure( error_type=ProxyErrorTypes.bad_request_error, route=route, ) + except NativeAuthFailure: + raise except Exception as e: # noqa: BLE001 # any logging failure must leave the caller's 400 untouched verbose_proxy_logger.exception("Failed to log the request rejected for an unparsable body: %s", e) @@ -2341,6 +2099,8 @@ async def _resolve_end_user( return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=default_budget) except litellm.BudgetExceededError: raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) @@ -2428,6 +2188,8 @@ async def _refresh_team_fields_from_cache( field_name = f"team_{k}" if field_name in valid_token.__fields__: setattr(valid_token, field_name, v) + except NativeAuthFailure: + raise except Exception as e: verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. @@ -2510,6 +2272,8 @@ class _BuilderEffects: cached: Final = IdentityStore.key_from_principal( await self._identity_store(check_cache_only=True).resolve(hashed_token=hashed_token) ) + except NativeAuthFailure: + raise except Exception: verbose_logger.debug("api key not found in cache.") return None @@ -2582,10 +2346,10 @@ class _BuilderEffects: budget_throttle_pct=budget_throttle_percentage() if should_throttle_budget_exceeded(key) else None, ) - async def enforce_model_access(self) -> None: + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: from litellm.proxy.proxy_server import llm_router - await enforce_key_and_fallback_model_access( + return await read_key_and_fallback_model_access_facts( valid_token=self.key, request_data=self._request_data, route=self._route, @@ -2593,6 +2357,8 @@ class _BuilderEffects: llm_model_list=_proxy_model_list(), llm_router=llm_router, ) + + async def prefetch_auth_objects(self) -> None: await _prefetch_referenced_auth_objects( self.key, end_user_id=self.end_user.end_user_id, @@ -2611,6 +2377,8 @@ class _BuilderEffects: parent_otel_span=self._parent_otel_span, proxy_logging_obj=self._proxy_logging_obj, ) + except NativeAuthFailure: + raise except Exception as e: verbose_logger.debug("Unable to get user %s from db/cache, continuing without it: %s", user_id, e) return None @@ -2623,7 +2391,8 @@ class _BuilderEffects: key.user_model_max_budget = owner.model_max_budget return OwnerFacts(deactivated=isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False) - async def check_budget_exemption(self) -> bool: + async def check_budget_exemption(self) -> ExemptionFacts: + from litellm.proxy.auth.budget_pricing_adapter import read_budget_exemption_facts from litellm.proxy.proxy_server import llm_router model: Final = _get_model_from_request_context( @@ -2633,11 +2402,13 @@ class _BuilderEffects: llm_router=llm_router, team_id=self.key.team_id, ) - return request_skips_budget_checks( + return read_budget_exemption_facts( route=self._route, model=model, llm_router=llm_router, request_body=self._request_data ) async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: + from litellm._auth import reservation + from litellm.proxy.auth.common_budget_adapter import resolve_effective_team_member_budget from litellm.proxy.proxy_server import get_current_spend if self._prisma_client is None: @@ -2647,9 +2418,12 @@ class _BuilderEffects: membership: Final = cached if cached is not None else await self._load_team_membership(user_id, team_id) if membership is None or membership.litellm_budget_table is None: return None - max_budget: Final = membership.litellm_budget_table.effective_max_budget(now=datetime.now(timezone.utc)) - if max_budget is None or max_budget <= 0: + max_budget: Final = resolve_effective_team_member_budget( + membership.litellm_budget_table, None, datetime.now(timezone.utc) + ) + if not reservation.counter("Member", max_budget): return None + assert max_budget is not None spend: Final = await get_current_spend( counter_key=team_member_counter(user_id, team_id), fallback_spend=self.key.team_member_spend or 0.0, @@ -2726,16 +2500,16 @@ class _BuilderEffects: async def read_model_budget_scope(self) -> ModelBudgetScope: return await self._model_budgets().read_model_budget_scope() - async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: return await self._model_budgets().check_key_model_budget(model) - async def authorize_budget_fallback(self, model: str) -> bool: + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: return await self._model_budgets().authorize_budget_fallback(model) async def rewrite_request_model(self, model: str) -> None: await self._model_budgets().rewrite_request_model(model) - async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: return await self._model_budgets().check_user_model_budget(model) async def load_key_context(self) -> None: @@ -2802,6 +2576,8 @@ async def user_api_key_auth_builder( start_time: Final = getattr(request.state, "litellm_received_at", None) or datetime.now(timezone.utc) try: request.state.litellm_received_at = start_time + except NativeAuthFailure: + raise except Exception: pass route: Final[str] = get_request_route(request=request) @@ -2992,6 +2768,8 @@ async def user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, user_email=mapped_user_email, ) + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug("JWT mapped-key user_email backfill skipped: %s", e) else: @@ -3216,6 +2994,8 @@ async def user_api_key_auth_builder( start_time=start_time, jwt_auth_enabled=enable_jwt_auth, ) + except NativeAuthFailure: + raise except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=e, @@ -3440,6 +3220,8 @@ async def _safe_fetch(label: str, awaitable): e, ) raise + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.debug( "centralized auth: %s fetch swallowed (%s: %s)", @@ -3472,36 +3254,6 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) -def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseException) -> bool: - """Whether the token's own team fields may stand in for a team that failed to - resolve, without widening access. - - The UI dashboard mints every session key against the ``UI_TEAM_ID`` sentinel, - which by design never has a team row, so a failed lookup for it is not a - degraded read to be treated with suspicion; it always vouches, exactly as it - always safely has (these keys are restricted elsewhere to UI-only routes). - - For every other team, a team that is provably gone is a definitive answer, - not a degraded read, so nothing may stand in for it and no setting may - override that. - - Otherwise the team's grant is merely unknown. A token carrying one may vouch, - since replaying a recorded grant cannot widen it and denying every team key - while the row is briefly unreadable would trade the widening for an outage. A - token carrying none may not: ``team_models=[]`` reads as every model and - ``team_blocked=False`` as unblocked. ``allow_requests_on_db_unavailable`` opts - back out, and is only consulted here because the failure is known by this - point to be a degraded read. - """ - if valid_token.team_id == UI_TEAM_ID: - return True - if isinstance(lookup_error, TeamNotFoundError): - return False - if valid_token.team_models: - return True - return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - - async def _inherit_org_identity( user_api_key_auth_obj: UserAPIKeyAuth, team_object: LiteLLM_TeamTableCachedObj | None, @@ -3558,68 +3310,10 @@ async def _load_team_member_default_model_budget( skip_budget_checks: bool, proxy_logging_obj: ProxyLogging | None, ) -> None: - team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None - team_metadata: Final[Mapping[str, object] | None] = ( - cast( # cast-ok: team metadata is a string-keyed JSON object - Mapping[str, object], team_metadata_value - ) - if isinstance(team_metadata_value, Mapping) - else None - ) - team_member_budget_id: Final[object | None] = ( - team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None - ) - if ( - not isinstance(user_api_key_auth_obj.user_id, str) - or not isinstance(user_api_key_auth_obj.team_id, str) - or prisma_client is None - ): - return + from litellm.proxy.auth import budget_context_adapter - try: - membership: Final = await get_team_membership( - user_id=user_api_key_auth_obj.user_id, - team_id=user_api_key_auth_obj.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - member_budget_row: Final = membership.litellm_budget_table if membership is not None else None - member_budget_value: Final[object] = ( - getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None - ) - if ( - membership is not None - and membership.budget_id is not None - and membership.budget_id != team_member_budget_id - and isinstance(member_budget_value, Mapping) - and member_budget_value - ): - user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream - member_budget_value - ) - return - if not isinstance(team_member_budget_id, str): - return - team_member_default_budget: Final = await get_team_member_default_budget( - budget_id=team_member_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - raise_on_lookup_error=True, - ) - except Exception: - if not skip_budget_checks: - raise - verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True) - return - - default_budget_value: Final[object] = ( - getattr(team_member_default_budget, "model_max_budget", None) - if team_member_default_budget is not None - else None - ) - user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream - default_budget_value if isinstance(default_budget_value, Mapping) else None + return await budget_context_adapter.load_team_member_default_model_budget( + user_api_key_auth_obj, team_object, prisma_client, user_api_key_cache, skip_budget_checks, proxy_logging_obj ) @@ -3667,41 +3361,34 @@ async def run_centralized_common_checks( user_custom_auth, ) - # Public routes (e.g. /health/liveness) are exempt from - # auth in the builder — the wrapper must not retroactively apply - # authz on top, or k8s readiness probes and other unauthenticated - # callers get 401. - if not force_virtual_key_checks and ( - route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route) - ): - return - - # User-configured pass-through endpoints with ``auth: false`` are - # explicitly unauthenticated — the builder returns an empty - # UserAPIKeyAuth() and the request is forwarded as-is. Running - # common_checks on the empty token would reject the request as - # admin-only. The "auth" flag on the endpoint config is the - # contract; honor it. pass_through_endpoints: Final = general_settings.get("pass_through_endpoints", None) - if not force_virtual_key_checks and pass_through_endpoints is not None: - for endpoint in pass_through_endpoints: - if isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True: - return - - # No-auth dev mode: master_key unset AND no JWT/OAuth2 auth - # configured. The builder returns an INTERNAL_USER token for any - # api_key; the proxy is unauthenticated by configuration. - # Running common_checks would block every admin route on these - # deployments where that was previously not the contract. If any - # authn is enabled (JWT, OAuth2, OAuth2-proxy), authz must run. - if not force_virtual_key_checks and is_no_auth_dev_mode(master_key, general_settings): - return - - if ( - not force_virtual_key_checks - and user_custom_auth is not None - and not general_settings.get("custom_auth_run_common_checks", False) - ): + plan: Final = admission.common_plan( + admission.CommonGate( + force_virtual_key_checks=force_virtual_key_checks, + public_route=route in LiteLLMRoutes.public_routes.value, + additional_public_route=route_in_additonal_public_routes(current_route=route), + unauthenticated_passthrough=any( + isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True + for endpoint in (pass_through_endpoints or ()) + ), + master_key_present=master_key is not None, + jwt_enabled=bool(general_settings.get("enable_jwt_auth", False)), + oauth_enabled=bool(general_settings.get("enable_oauth2_auth", False)), + oauth_proxy_enabled=bool(general_settings.get("enable_oauth2_proxy_auth", False)), + custom_auth=user_custom_auth is not None, + custom_auth_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), + ) + ) + request.state.auth_ownership_plan = admission.ownership_plan( + admission.Ownership( + virtual_key=user_api_key_auth_obj.via_virtual_key, + custom_auth=user_api_key_auth_obj.authenticated_by_custom_auth, + jwt=bool(user_api_key_auth_obj.jwt_claims), + oauth=bool(general_settings.get("enable_oauth2_auth", False)) + or bool(general_settings.get("enable_oauth2_proxy_auth", False)), + ) + ) + if plan.skip is not None: return parent_otel_span: Final = user_api_key_auth_obj.parent_otel_span @@ -3726,7 +3413,17 @@ async def run_centralized_common_checks( user_api_key_auth_obj.end_user_id = end_user_id fetch_coros: Final = [] - if user_api_key_auth_obj.team_id is not None and user_api_key_auth_obj.team_id != UI_TEAM_ID: + fetch_plan: Final = auth_context.fetch( + auth_context.Fetch( + user_api_key_auth_obj.team_id, + UI_TEAM_ID, + user_api_key_auth_obj.user_id, + user_api_key_auth_obj.project_id, + end_user_id, + ) + ) + if fetch_plan.team: + assert user_api_key_auth_obj.team_id is not None fetch_coros.append( _safe_fetch( "team", @@ -3742,7 +3439,8 @@ async def run_centralized_common_checks( else: fetch_coros.append(_safe_fetch("team", _noop_none())) - if user_api_key_auth_obj.user_id is not None: + if fetch_plan.user: + assert user_api_key_auth_obj.user_id is not None fetch_coros.append( _safe_fetch( "user", @@ -3759,7 +3457,8 @@ async def run_centralized_common_checks( else: fetch_coros.append(_safe_fetch("user", _noop_none())) - if user_api_key_auth_obj.project_id is not None: + if fetch_plan.project: + assert user_api_key_auth_obj.project_id is not None fetch_coros.append( _safe_fetch( "project", @@ -3774,7 +3473,8 @@ async def run_centralized_common_checks( else: fetch_coros.append(_safe_fetch("project", _noop_none())) - if end_user_id: + if fetch_plan.end_user: + assert end_user_id is not None fetch_coros.append( _safe_fetch( "end_user", @@ -3822,38 +3522,46 @@ async def run_centralized_common_checks( global_spend_result, ) = await asyncio.gather(*fetch_coros, return_exceptions=True) - # ProxyException / BudgetExceededError are authorization failures — - # propagate so the wrapper renders them. HTTPException is fallback - # material (404 from get_team_object is the only known producer). - for r in ( - team_result, - user_result, - project_result, - end_user_result, - global_spend_result, - ): - if isinstance(r, (ProxyException, litellm.BudgetExceededError)): - raise r - - # Use BaseException (not HTTPException) in the narrowing checks so - # mypy can narrow ``Any | BaseException`` to the typed object in the - # else branch. After the for-loop above, the only BaseException that - # can still appear here is HTTPException (other listed re-raises were - # propagated; non-listed exceptions were already swallowed to None). - team_object: LiteLLM_TeamTableCachedObj | None - if isinstance(team_result, BaseException): - # Token-derived fallback only valid when a team_id is set; - # _team_obj_from_token asserts that precondition. - if user_api_key_auth_obj.team_id is None: - team_object = None - elif _token_can_vouch_for_team(user_api_key_auth_obj, team_result): - team_object = _team_obj_from_token(user_api_key_auth_obj) - else: - raise team_result - else: - team_object = ( - _team_obj_from_token(user_api_key_auth_obj) if user_api_key_auth_obj.team_id == UI_TEAM_ID else team_result + results: Final = (team_result, user_result, project_result, end_user_result, global_spend_result) + failed: Final = access_policy.decide( + access_policy.ConcurrentFailures( + failures=tuple(isinstance(value, (ProxyException, litellm.BudgetExceededError)) for value in results), + terminal=tuple(isinstance(value, NativeAuthFailure) for value in results), ) + ).failure_index + if failed is not None: + failure: Final = results[failed] + assert isinstance(failure, BaseException) + raise failure + team_facts: Final = auth_context.Team( + user_api_key_auth_obj.team_id, + UI_TEAM_ID, + failed=isinstance(team_result, BaseException), + missing=isinstance(team_result, TeamNotFoundError), + recorded_models=tuple(user_api_key_auth_obj.team_models or ()), + ) + initial_team: Final = auth_context.team(team_facts) + team_plan: Final = ( + auth_context.team( + replace(team_facts, allow_unavailable=PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()) + ) + if initial_team.source == "ReadUnavailablePolicy" + else initial_team + ) + team_object: LiteLLM_TeamTableCachedObj | None + match team_plan.source: + case "None": + team_object = None + case "Token": + team_object = _team_obj_from_token(user_api_key_auth_obj) + case "Loaded": + assert not isinstance(team_result, BaseException) + team_object = team_result + case "Failure": + assert isinstance(team_result, BaseException) + raise team_result + case _: + raise NativeAuthFailure("incomplete team context policy") user_object: LiteLLM_UserTable | None = user_result if isinstance(user_result, LiteLLM_UserTable) else None project_object: Final[LiteLLM_ProjectTableCachedObj | None] = ( @@ -3888,12 +3596,16 @@ async def run_centralized_common_checks( # username/password login) keeps the row's personal budget (max_budget, # and user_email for budget alerts) so common_checks enforces it; every # other field keeps its default. - if user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN: - budget_row: Final = ( - user_object - if isinstance(user_object, LiteLLM_UserTable) and user_api_key_auth_obj.user_id != litellm_proxy_admin_name - else None + admin_plan: Final = auth_context.admin( + auth_context.Admin( + user_api_key_auth_obj.user_role, + user_api_key_auth_obj.user_id, + litellm_proxy_admin_name, + isinstance(user_object, LiteLLM_UserTable), ) + ) + if admin_plan.force_admin: + budget_row: Final = user_object if admin_plan.preserve_budget else None user_object = LiteLLM_UserTable( user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name, user_role=LitellmUserRoles.PROXY_ADMIN, @@ -4089,35 +3801,17 @@ async def _apply_key_end_user_default_budget_to_token( parent_otel_span: Span | None, keep_token_limits: bool, ) -> None: - """The builder's end-user pass runs before the key is resolved, so only here can the key's - ``end_user_budget_id`` win over the proxy-wide default on the token that reservation reads. - On the virtual-key path the token's end-user limits are the builder's proxy-wide defaults and - the key budget replaces them wholesale. With ``keep_token_limits`` (custom auth) the token's - limits are caps the custom auth callable set, so the key budget only fills the ones it left - unset.""" - default_budget: Final = ( - end_user_object.litellm_budget_table - if end_user_object is not None - else await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, - ) - ) - if default_budget is None: - return + from litellm.proxy.auth import budget_context_adapter - if not keep_token_limits or valid_token.end_user_max_budget is None: - valid_token.end_user_max_budget = default_budget.max_budget - if not keep_token_limits or valid_token.end_user_tpm_limit is None: - valid_token.end_user_tpm_limit = default_budget.tpm_limit - if not keep_token_limits or valid_token.end_user_rpm_limit is None: - valid_token.end_user_rpm_limit = default_budget.rpm_limit - if not keep_token_limits or valid_token.end_user_tpd_limit is None: - valid_token.end_user_tpd_limit = default_budget.tpd_limit - if not keep_token_limits or valid_token.end_user_model_max_budget is None: - valid_token.end_user_model_max_budget = default_budget.model_max_budget + return await budget_context_adapter.apply_key_end_user_default_budget_to_token( + valid_token, + end_user_object, + key_end_user_budget_id, + prisma_client, + user_api_key_cache, + parent_otel_span, + keep_token_limits, + ) async def _reserve_budget_after_common_checks( @@ -4137,7 +3831,16 @@ async def _reserve_budget_after_common_checks( request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None - if not skip_budget_checks and general_settings.get("disable_budget_reservation") is not True: + plan: Final = admission.reservation_plan( + admission.Reservation( + skip_budget_checks=skip_budget_checks, + disabled=general_settings.get("disable_budget_reservation") is True, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + fail_closed=general_settings.get("fail_closed_budget_enforcement") is True, + billing_agent=user_api_key_auth_obj.billing_agent_policy is not None, + ) + ) + if plan.reserve: from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, ) @@ -4154,11 +3857,8 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, - apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, - fail_closed_budget_enforcement=( - general_settings.get("fail_closed_budget_enforcement") is True - or user_api_key_auth_obj.billing_agent_policy is not None - ), + apply_user_budget_to_team_keys=plan.apply_user_budget_to_team_keys, + fail_closed_budget_enforcement=plan.fail_closed, raw_body=await read_raw_json_body(request=request), ) if request is not None: @@ -4291,6 +3991,8 @@ async def _authorize_authenticated_request( route=route, force_virtual_key_checks=force_virtual_key_checks, ) + except NativeAuthFailure: + raise except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=e, @@ -4412,6 +4114,8 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque fan_out_provider(), ) ) + except NativeAuthFailure: + raise except Exception as exc: # noqa: BLE001 # telemetry routing is best-effort and must never break authentication verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc) @@ -4461,6 +4165,8 @@ async def user_api_key_auth( request_data=request_data, custom_litellm_key_header=custom_litellm_key_header, ) + except NativeAuthFailure: + raise except Exception: # The body was read first, so a caller who sent both a malformed body and # a rejected key used to get the 400; the response is unchanged, and the @@ -4516,6 +4222,8 @@ async def user_api_key_auth( # any future consumer must treat a missing principal as deny, not allow. try: request.state.principal = _resolve_request_principal(request, user_api_key_auth_obj) + except NativeAuthFailure: + raise except Exception as e: verbose_proxy_logger.warning("Principal projection at auth seam failed (non-fatal): %s", e) @@ -4570,26 +4278,12 @@ async def _return_user_api_key_auth_obj( return UserAPIKeyAuth.model_validate(user_api_key_kwargs) -def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): - valid_token_metadata: Final = valid_token.metadata - if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: - expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) - if expiry.tzinfo is None: - expiry = expiry.replace(tzinfo=timezone.utc) - if expiry > datetime.now(timezone.utc): - return valid_token_metadata["temp_budget_increase"] - return None - - def _update_key_budget_with_temp_budget_increase( valid_token: UserAPIKeyAuth, ) -> UserAPIKeyAuth: - if valid_token.max_budget is None: - return valid_token - temp_budget_increase: Final = _get_temp_budget_increase(valid_token) - if not temp_budget_increase: - return valid_token - return valid_token.model_copy(update={"max_budget": valid_token.max_budget + temp_budget_increase}) + from litellm.proxy.auth import key_budget_context_adapter + + return key_budget_context_adapter.update_key_budget_with_temp_budget_increase(valid_token) async def _lookup_end_user_and_apply_budget( @@ -4600,119 +4294,40 @@ async def _lookup_end_user_and_apply_budget( user_api_key_cache, proxy_logging_obj, ): - """Look up end_user from DB and apply budget limits to valid_token.""" - end_user_object = None - key_end_user_budget_id: Final = get_key_end_user_budget_id(valid_token.metadata) - try: - end_user_object = await get_end_user_object( - end_user_id=valid_token.end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - token_end_user_max_budget=valid_token.end_user_max_budget, - key_end_user_budget_id=key_end_user_budget_id, - ) - if end_user_object is not None: - end_user_params = { - "end_user_id": valid_token.end_user_id, - "allowed_model_region": end_user_object.allowed_model_region, - } - if end_user_object.litellm_budget_table is not None: - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=end_user_object.litellm_budget_table, - end_user_id=valid_token.end_user_id or "", - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params - ) - elif key_end_user_budget_id is not None or litellm.max_end_user_budget_id is not None: - default_budget: Final = await resolve_default_end_user_budget( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_end_user_budget_id=key_end_user_budget_id, - parent_otel_span=parent_otel_span, - ) - if default_budget is not None: - end_user_params = {"end_user_id": valid_token.end_user_id} - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=default_budget, - end_user_id=valid_token.end_user_id or "", - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params - ) - if valid_token.end_user_max_budget is None: - valid_token.end_user_max_budget = default_budget.max_budget - except Exception as e: - if isinstance(e, litellm.BudgetExceededError): - raise e - verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) - return valid_token, end_user_object + from litellm.proxy.auth import end_user_context_adapter + + return await end_user_context_adapter.lookup_end_user_and_apply_budget( + valid_token, route, parent_otel_span, prisma_client, user_api_key_cache, proxy_logging_obj + ) async def enforce_key_and_fallback_model_access( *, valid_token: UserAPIKeyAuth, - request_data: dict, + request_data: dict[str, object], route: str, request: Request | None, - llm_model_list: list | None, + llm_model_list: Sequence[object] | None, llm_router: litellm.Router | None, ) -> None: - """ - Key-level model allowlist and client fallbacks (same as standard auth). - Not included in common_checks — common_checks enforces team/user/project model access only. - """ - await _normalize_claude_model(request_data, valid_token, request, route) - await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) - _resolve_codex_native_chatgpt_model(request_data, valid_token, request, route) - config: Final = valid_token.config - - if config != {}: - model_list: Final = config.get("model_list", []) - new_model_list: Final = model_list - verbose_proxy_logger.debug("\n new llm router model list %s", new_model_list) - elif isinstance(valid_token.models, list) and "all-team-models" in valid_token.models: - pass - else: - model: Final = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - - if model is not None: - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - - fallback_names: Final = tuple( - name - for target in iter_request_fallback_targets(request_data) - if (name := fallback_target_model_name(target)) is not None - ) - - for _name in dict.fromkeys(fallback_names): # dedupe, preserve order - await can_key_call_model( - model=_name, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - await is_valid_fallback_model( - model=_name, - llm_router=llm_router, - user_model=None, - ) + facts: Final = await read_key_and_fallback_model_access_facts( + valid_token=valid_token, + request_data=request_data, + route=route, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + ) + decision: Final = key_and_fallback_model_access(facts) + match decision.kind: + case "Allow": + return + case "Deny": + if decision.denial is None: + raise NativeAuthFailure("missing native model denial") + raise_model_denial(decision.denial, decision.denial.model or "") + case _: + raise NativeAuthFailure("unanswered native model read") _enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias @@ -4736,8 +4351,14 @@ async def _run_post_custom_auth_checks( ) # 1. Look up end_user object from DB if end_user_id is set + custom_facts: Final = common_flow.CustomFacts( + valid_token.end_user_id, + False, + bool(general_settings.get("custom_auth_run_common_checks", False)), + ) + custom_plan: Final = common_flow.custom(custom_facts) end_user_object = None - if valid_token.end_user_id is not None: + if custom_plan.read_end_user: valid_token, end_user_object = await _lookup_end_user_and_apply_budget( valid_token=valid_token, route=route, @@ -4750,19 +4371,15 @@ async def _run_post_custom_auth_checks( # gate skips it for custom-auth deployments unless # custom_auth_run_common_checks is set. Enforce it here on that path # so an over-budget end user can't keep making requests. - if end_user_object is not None and not general_settings.get("custom_auth_run_common_checks", False): + if common_flow.custom(replace(custom_facts, end_user_present=end_user_object is not None)).check_end_user: + assert end_user_object is not None await check_end_user_budget(end_user_obj=end_user_object, route=route) # 2. Check token expiry if valid_token.expires is not None: current_time: Final = datetime.now(timezone.utc) - if isinstance(valid_token.expires, datetime): - expiry_time = valid_token.expires - else: - expiry_time = datetime.fromisoformat(valid_token.expires) - if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: - expiry_time = expiry_time.replace(tzinfo=timezone.utc) - if expiry_time < current_time: + expiry_time: Final = _aware_expiry(valid_token.expires) + if expired(expiry_time, current_time): raise ProxyException( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, @@ -4770,7 +4387,7 @@ async def _run_post_custom_auth_checks( param=(abbreviate_api_key(api_key=valid_token.token) if valid_token.token else ""), ) - if general_settings.get("custom_auth_run_common_checks", False): + if custom_plan.check_model_access: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, @@ -4794,7 +4411,9 @@ async def _run_post_custom_auth_checks( # be refused under custom auth and served under the other two. skip_budget_checks: Final = ( is_model_cost_zero(model=current_model, llm_router=llm_router) - if current_model is not None and llm_router is not None + if common_flow.custom( + replace(custom_facts, model_present=current_model is not None, router_present=llm_router is not None) + ).read_pricing else False ) @@ -4811,7 +4430,7 @@ async def _run_post_custom_auth_checks( proxy_logging_obj=proxy_logging_obj, ) valid_token.user_model_max_budget = user_budget # rebind-ok: the spend hook reads it off this token - if not skip_budget_checks: + if common_flow.read(common_flow.BudgetSkip(skip_budget_checks, False)) == "CheckBudgets": await _enforce_model_budgets( _AUTH_CORE.check_model_budgets(), _ModelBudgetEffects( @@ -4839,21 +4458,17 @@ async def _run_post_custom_auth_checks( ) end_user_mmb: Final = valid_token.end_user_model_max_budget - if ( - not skip_budget_checks - and end_user_mmb is not None - and isinstance(end_user_mmb, dict) - and len(end_user_mmb) > 0 - and current_models - and valid_token.end_user_id is not None - ): - for model_name in current_models: - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_mmb, - model=model_name, - team_id=valid_token.team_id, - ) + if common_flow.read(common_flow.BudgetSkip(skip_budget_checks, False)) == "CheckBudgets": + from litellm.proxy.auth import budget_context_adapter + + await budget_context_adapter.check_end_user_model_budget( + valid_token.end_user_id, + end_user_mmb if isinstance(end_user_mmb, dict) else None, + model_max_budget_limiter, + current_models, + valid_token.team_id, + scope="CustomEndUser", + ) # team / user / end_user / project context objects are fetched by # the centralized common_checks gate in user_api_key_auth after diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 0b389e3dfca..e030a85b0bf 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -15,6 +15,8 @@ from typing import ( ) import litellm +from litellm._auth import model_budget_data +from litellm._auth.model_budget_data import ModelBudgetEntity from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -41,6 +43,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload if TYPE_CHECKING: + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading from litellm.router import Router VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" @@ -988,9 +991,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): entity_type=Litellm_EntityType.TEAM, entity_id=team_id, team_id=team_id, - model_max_budget=team_model_budget_entries( - model=model, team_model_max_budget=team_model_max_budget, key_model_max_budget=key_model_max_budget - ), + model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, model=model, exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", ) @@ -1014,6 +1016,27 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> "ModelBudgetReading": + from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_model_budget + + io: Final = ModelBudgetIO( + self.dual_cache, + self._model_access_groups, + self._shared_windows, + _PROCESS_STARTED_AT, + clock=self._clock, + monotonic=time.monotonic, + ) + return await read_model_budget(io, entity, entity_id, team_id, model, budgets, key_budgets) + async def _is_entity_within_model_budget( self, entity_type: Litellm_EntityType, @@ -1022,6 +1045,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model_max_budget: Mapping[str, object] | None, model: str, exceeded_message: str, + key_model_max_budget: Mapping[str, object] | None = None, ) -> bool: for candidate in expand_model_latest_alias(model): await self._is_single_entity_within_model_budget( @@ -1031,6 +1055,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model_max_budget=model_max_budget, model=candidate, exceeded_message=exceeded_message, + key_model_max_budget=key_model_max_budget, ) return True @@ -1042,45 +1067,32 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model_max_budget: Mapping[str, object] | None, model: str, exceeded_message: str, + key_model_max_budget: Mapping[str, object] | None = None, ) -> bool: - if not model_max_budget: - return True - resolved: Final = resolve_model_budget( - model=model, - model_max_budget=model_max_budget, - model_access_groups=self._model_access_groups, - team_id=team_id, + entities: Final[Mapping[Litellm_EntityType, ModelBudgetEntity]] = { + Litellm_EntityType.KEY: "Key", + Litellm_EntityType.USER: "User", + Litellm_EntityType.TEAM: "Team", + Litellm_EntityType.TEAM_MEMBER: "TeamMember", + Litellm_EntityType.END_USER: "EndUser", + } + reading: Final = await self.read_model_budget_raw( + entities[entity_type], + entity_id, + team_id, + model, + model_max_budget, + key_model_max_budget, ) - if resolved is None: - verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value) - return True - - max_budget: Final = resolved.budget_config.max_budget - if max_budget is None or max_budget < 0: - return True - - local_spend: Final = await self._get_spend_for_model_budget( - entity_type=entity_type, - entity_id=entity_id, - model=model, - resolved=resolved, - ) - current_spend: Final = ( - local_spend - if local_spend >= max_budget - else await self._with_shared_spend( - entity_type=entity_type, entity_id=entity_id, resolved=resolved, local_spend=local_spend - ) - ) - if current_spend >= max_budget: + decision: Final = model_budget_data.check(reading.snapshot) + if decision.denied: + assert reading.snapshot.max_budget is not None raise litellm.BudgetExceededError( - message=( - exceeded_message - if resolved.budget_model in _budget_model_candidates(model) - else f"{exceeded_message}, budget_model={resolved.budget_model}" - ), - current_cost=current_spend, - max_budget=max_budget, + message=exceeded_message + if reading.named_model + else f"{exceeded_message}, budget_model={reading.budget_model}", + current_cost=decision.current_spend, + max_budget=reading.snapshot.max_budget, entity_type=entity_type.value, entity_id=entity_id, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d21b6169133..3ad323f8c1f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -46,6 +46,7 @@ from typing import ( Final, Literal, NamedTuple, + NoReturn, Optional, Protocol, TypeAlias, @@ -3378,7 +3379,7 @@ def _fail_closed_budget_enforcement() -> bool: return general_settings.get("fail_closed_budget_enforcement") is True -def _raise_budget_unverifiable(counter_key: str) -> None: +def _raise_budget_unverifiable(counter_key: str) -> NoReturn: verbose_proxy_logger.warning( "fail_closed_budget_enforcement: rejecting request — spend for %s could " "not be verified against Redis or the database", @@ -3436,46 +3437,28 @@ async def get_current_spend( and cached in-process for a few seconds, so a persistently stale counter drives at most one read per counter per window rather than one per request. """ + from functools import partial + + from litellm._auth.spend_counter_policy import Spend + from litellm.proxy.auth.spend_counter_adapter import SpendIO, resolve_spend + current, verified = await _read_spend_counter_estimate(counter_key=counter_key, fallback_spend=fallback_spend) - if fallback_authoritative: - verified = True - - if max_budget is None or current >= max_budget: - return current - - # Cheap staleness signal for primary counters: the counter reads below the - # spend this caller already knows about. Window counters have no such signal - # (fallback is 0), so they always re-check, bounded by the cache. Strict mode - # (fail_closed_budget_enforcement) always re-checks against the authoritative - # source too, so a counter that is stale-low at the same time as the caller's - # cached spend cannot slip through; the 5s cache keeps that bounded. - is_window: Final = window_start is not None - if fallback_spend > current or is_window or _fail_closed_budget_enforcement(): - authoritative: Final = await _authoritative_floor_spend( - counter_key=counter_key, - window_entity_type=window_entity_type, - window_entity_id=window_entity_id, - window_duration=window_duration, - window_start=window_start, - ) - if authoritative is not None: - verified = True - if authoritative > current: - await _repair_stale_spend_counter(counter_key=counter_key, db_spend=authoritative) - return authoritative - elif fallback_spend > current: - # nothing to read (tag counters, an end user without a row or a DB client, a - # failed read); fallback_spend is the authoritative recorded value loaded in auth. - return fallback_spend - - # Opt-in hard guarantee: when the spend backing this admit decision came - # only from a per-pod cache (Redis and DB both unreadable), reject rather - # than admit on an unverifiable budget. No-op unless the flag is set, so - # default behavior is unchanged. - if not verified and _fail_closed_budget_enforcement(): - _raise_budget_unverifiable(counter_key) - - return current + return await resolve_spend( + Spend(current, verified, fallback_spend, max_budget, window_start is not None, fallback_authoritative), + SpendIO( + read_strict=_fail_closed_budget_enforcement, + read_floor=partial( + _authoritative_floor_spend, + counter_key=counter_key, + window_entity_type=window_entity_type, + window_entity_id=window_entity_id, + window_duration=window_duration, + window_start=window_start, + ), + repair=lambda spend: _repair_stale_spend_counter(counter_key=counter_key, db_spend=spend), + unavailable=partial(_raise_budget_unverifiable, counter_key), + ), + ) @with_service_target(SPEND_COUNTERS_TARGET) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8f0584bd0d0..9855ff77225 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,10 +2,9 @@ from __future__ import annotations import asyncio import json -import math import time from collections.abc import AsyncIterator, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType @@ -15,6 +14,9 @@ from fastapi import HTTPException, status from pydantic import ConfigDict, TypeAdapter import litellm +from litellm._auth import common_budgets, spend_counter_policy +from litellm._auth import reservation as reservation_policy +from litellm._auth.failure import NativeAuthRuntimeError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route @@ -29,6 +31,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded +from litellm.proxy.auth.common_budget_adapter import read_member_limit_facts from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, @@ -182,9 +185,16 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: requests never reach this branch so their concurrent-overspend protection is untouched. """ - if valid_token is None: + decision: Final = spend_counter_policy.decide( + spend_counter_policy.ThrottleKey( + counter_key, valid_token is not None, valid_token.token if valid_token is not None else None + ) + ) + if decision.kind == "Skip": return False - return counter_key == key_counter(str(valid_token.token)) and should_throttle_budget_exceeded(valid_token) + if decision.kind != "ReadThrottle" or valid_token is None: + raise NativeAuthRuntimeError("Unexpected throttle-counter decision") + return should_throttle_budget_exceeded(valid_token) async def _apply_over_budget_reservation_policy( @@ -204,26 +214,29 @@ async def _apply_over_budget_reservation_policy( down to what is left, unless strict enforcement is on, because the known estimate already does not fit; anything else hard-blocks by raising. """ - if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): - await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) - applied_entries.remove(entry) - return reservation_cost - - remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation <= 1e-12: - _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) - if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: - _raise_counter_budget_exceeded( - counter=counter, - current_cost=current_spend - reservation_cost, - estimated_cost=reservation_cost, + decision: Final = spend_counter_policy.decide( + spend_counter_policy.OverBudget( + _key_reservation_should_release_for_throttle(counter.counter_key, valid_token), + current_spend, + counter.max_budget, + reservation_cost, + fail_closed_budget_enforcement, ) - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, ) - return remaining_before_reservation + match decision: + case spend_counter_policy.Action(kind="Release"): + await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) + applied_entries.remove(entry) + return reservation_cost + case spend_counter_policy.Rejected(current=current, estimate=estimate): + _raise_counter_budget_exceeded(counter=counter, current_cost=current, estimated_cost=estimate) + case spend_counter_policy.Resize(cost=cost): + await _resize_applied_reservation( + entries=applied_entries, current_reserved_cost=reservation_cost, new_reserved_cost=cost + ) + return cost + case _: + raise NativeAuthRuntimeError("Unexpected reservation-overflow decision") def _raise_counter_budget_exceeded( @@ -263,11 +276,20 @@ async def reserve_budget_for_request( fail_closed_budget_enforcement: bool = False, raw_body: bytes | None = None, ) -> dict | None: - if valid_token is None or not RouteChecks.is_llm_api_route(route=route): + start: Final = reservation_policy.Gate(token_present=valid_token is not None) + if reservation_policy.action(start).kind == "Skip": return None - if RouteChecks.is_unbilled_route(route): + assert valid_token is not None + route_facts: Final = replace(start, llm_route=RouteChecks.is_llm_api_route(route=route)) + if reservation_policy.action(route_facts).kind == "Skip": return None - if get_model_from_request(request_body, route, llm_router=llm_router) is None: + billing_facts: Final = replace(route_facts, unbilled=RouteChecks.is_unbilled_route(route)) + if reservation_policy.action(billing_facts).kind == "Skip": + return None + model_facts: Final = replace( + billing_facts, model_present=get_model_from_request(request_body, route, llm_router=llm_router) is not None + ) + if reservation_policy.action(model_facts).kind == "Skip": return None counters: Final = await _get_budget_counters( @@ -282,7 +304,8 @@ async def reserve_budget_for_request( end_user_object=end_user_object, apply_user_budget_to_team_keys=apply_user_budget_to_team_keys, ) - if not counters: + counter_facts: Final = replace(model_facts, counter_count=len(counters)) + if reservation_policy.action(counter_facts).kind == "Skip": return None input_token_counts: Final = await count_request_input_tokens( @@ -302,8 +325,9 @@ async def reserve_budget_for_request( # estimate_request_max_cost still returns None when the model is unknown # to the cost map (no token-priced cost fields, e.g. image/audio routes). # In that case we fall back to read-time enforcement only. - if reservation_cost is None or reservation_cost <= 0: + if reservation_policy.action(replace(counter_facts, cost_read=True, cost=reservation_cost)).kind == "Skip": return None + assert reservation_cost is not None applied_entries: Final[list[dict[str, float | str]]] = [] try: @@ -494,9 +518,21 @@ async def _get_budget_counters( apply_user_budget_to_team_keys: bool = False, ) -> list[_BudgetCounter]: counters: Final[list[_BudgetCounter]] = [] + sources: Final = reservation_policy.sources( + reservation_policy.Sources( + key_present=valid_token.token is not None, + team_present=team_object is not None, + team_id=team_object.team_id if team_object is not None else None, + user_present=user_object is not None, + user_id=user_object.user_id if user_object is not None else None, + apply_user_budget_to_team_keys=apply_user_budget_to_team_keys, + ) + ) - if valid_token.token is not None: - if valid_token.max_budget is not None and valid_token.max_budget > 0: + if sources.key: + assert valid_token.token is not None + if reservation_policy.counter("Key", valid_token.max_budget): + assert valid_token.max_budget is not None counters.append( _BudgetCounter( counter_key=key_counter(valid_token.token), @@ -517,9 +553,11 @@ async def _get_budget_counters( ) ) - if team_object is not None and team_object.team_id is not None: + if sources.team: + assert team_object is not None and team_object.team_id is not None team_id: Final = team_object.team_id - if team_object.max_budget is not None and team_object.max_budget > 0: + if reservation_policy.counter("Team", team_object.max_budget): + assert team_object.max_budget is not None counters.append( _BudgetCounter( counter_key=team_counter(team_id), @@ -540,14 +578,8 @@ async def _get_budget_counters( ) ) - is_team_key: Final = team_object is not None and team_object.team_id is not None - if ( - (not is_team_key or apply_user_budget_to_team_keys) - and user_object is not None - and user_object.user_id is not None - and user_object.max_budget is not None - and user_object.max_budget > 0 - ): + if sources.user and reservation_policy.counter("User", user_object.max_budget if user_object is not None else None): + assert user_object is not None and user_object.user_id is not None and user_object.max_budget is not None counters.append( _BudgetCounter( counter_key=user_counter(user_object.user_id), @@ -616,7 +648,9 @@ async def _get_end_user_budget_counter( end_user_id: str | None, end_user_object: object, ) -> _BudgetCounter | None: - end_user_id = end_user_id or valid_token.end_user_id + end_user_id = reservation_policy.identity( + reservation_policy.Identity("EndUser", end_user_id, valid_token.end_user_id) + ) if end_user_id is None: return None @@ -629,8 +663,9 @@ async def _get_end_user_budget_counter( budget_table: Final = _get_value(end_user_object, "litellm_budget_table") max_budget = _to_float(_get_value(budget_table, "max_budget")) - if max_budget is None or max_budget <= 0: + if not reservation_policy.counter("EndUser", max_budget): return None + assert max_budget is not None return _BudgetCounter( counter_key=end_user_counter(end_user_id), @@ -669,8 +704,9 @@ async def _get_tag_budget_counters( continue budget_table = _get_value(tag_object, "litellm_budget_table") max_budget = _to_float(_get_value(budget_table, "max_budget")) - if max_budget is None or max_budget <= 0: + if not reservation_policy.counter("Tag", max_budget): continue + assert max_budget is not None counters.append( _BudgetCounter( counter_key=tag_counter(tag_name), @@ -712,8 +748,9 @@ async def _get_model_access_group_budget_counters( def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | None) -> _BudgetCounter | None: """A counter for one group, or nothing when the group carries no budget to reserve against.""" - if budget is None or budget.max_budget is None or budget.max_budget <= 0: + if not reservation_policy.counter("AccessGroup", budget.max_budget if budget is not None else None): return None + assert budget is not None and budget.max_budget is not None return _BudgetCounter( counter_key=model_access_group_counter(group), source_cache_key=model_access_group_cache_key(group), @@ -742,8 +779,19 @@ async def _get_team_member_budget_counter( user_object: LiteLLM_UserTable | None, user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: - if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: + if ( + reservation_policy.action( + reservation_policy.MemberGate( + team_object is not None, + team_object.team_id if team_object is not None else None, + user_object is not None, + valid_token.user_id, + ) + ).kind + == "Skip" + ): return None + assert team_object is not None and valid_token.user_id is not None membership_cache_key: Final = team_membership_reservation_cache_key( user_id=valid_token.user_id, team_id=team_object.team_id @@ -759,25 +807,33 @@ async def _get_team_member_budget_counter( member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None now: Final = datetime.now(timezone.utc) - team_member_budget: float | None = None - if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) - else: + default_cap: float | None = None + if ( + reservation_policy.action( + reservation_policy.MemberCap(member_budget_row.max_budget if member_budget_row is not None else None) + ).kind + == "ReadDefault" + ): default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): default_budget: Final = _CACHED_VALUE.validate_python( await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}") ) - default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) - if default_cap is not None and default_cap > 0: - team_member_budget = default_cap + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) + default_cap = _to_float(_get_value(default_budget, "max_budget")) - if team_member_budget is None or team_member_budget <= 0: + team_member_budget: Final = common_budgets.member_limit( + read_member_limit_facts( + member_budget_row, + default_cap if reservation_policy.counter("Key", default_cap) else None, + now, + ) + ) + + if not reservation_policy.counter("Member", team_member_budget): return None + assert team_member_budget is not None - team_member_spend = cast(LiteLLM_TeamMembership, team_membership).spend if team_membership is not None else 0.0 + team_member_spend = team_membership.spend if team_membership is not None else 0.0 return _BudgetCounter( counter_key=team_member_counter(valid_token.user_id, team_object.team_id), source_cache_key=membership_cache_key, @@ -794,11 +850,11 @@ async def _get_org_budget_counter( team_object: LiteLLM_TeamTable | None, user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: - org_id: str | None = None - if valid_token.org_id is not None: - org_id = valid_token.org_id - elif team_object is not None and team_object.organization_id is not None: - org_id = team_object.organization_id + org_id: Final = reservation_policy.identity( + reservation_policy.Identity( + "Organization", valid_token.org_id, team_object.organization_id if team_object is not None else None + ) + ) if org_id is None: return None @@ -813,8 +869,9 @@ async def _get_org_budget_counter( return None org_max_budget: Final = _to_float(_get_value(org_budget_table, "max_budget")) - if org_max_budget is None or org_max_budget <= 0: + if not reservation_policy.counter("Organization", org_max_budget): return None + assert org_max_budget is not None org_spend: Final = _to_float(_get_value(org_table, "spend")) or 0.0 return _BudgetCounter( @@ -847,8 +904,9 @@ async def _get_project_budget_counter( return None project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget")) - if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget): + if not reservation_policy.counter("Project", project_max_budget): return None + assert project_max_budget is not None return _BudgetCounter( counter_key=project_counter(valid_token.project_id), @@ -875,10 +933,11 @@ def _get_budget_limit_counters( window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") max_budget = _to_float(window_dict.get("max_budget")) - if not budget_duration or max_budget is None or max_budget <= 0: + facts = reservation_policy.Window(bool(budget_duration), max_budget) + if reservation_policy.action(facts).kind == "Skip": continue window_start = get_budget_window_start(window_dict) - if window_start is None: + if reservation_policy.action(replace(facts, start_present=window_start is not None)).kind == "WarnWindow": verbose_proxy_logger.warning( "Skipping budget window with invalid duration for %s=%s: %s", entity_type, @@ -886,6 +945,7 @@ def _get_budget_limit_counters( budget_duration, ) continue + assert max_budget is not None and window_start is not None counters.append( _BudgetCounter( counter_key=f"{entity_prefix}:window:{budget_duration}", @@ -937,7 +997,10 @@ async def _reservation_counter_loaded(counter: _BudgetCounter, fail_closed_budge try: await _initialize_reservation_counter(counter=counter) except _CounterReservationUnavailable: - if fail_closed_budget_enforcement: + if ( + spend_counter_policy.decide(spend_counter_policy.ReservationFailure(fail_closed_budget_enforcement)).kind + == "Unavailable" + ): _raise_reservation_unavailable(counter_key=counter.counter_key) return False return True @@ -1013,7 +1076,12 @@ async def _reserve_reservable_counters( if reserved_values is None: for entry in entries: applied_entries.remove(entry) - if fail_closed_budget_enforcement: + if ( + spend_counter_policy.decide( + spend_counter_policy.ReservationFailure(fail_closed_budget_enforcement) + ).kind + == "Unavailable" + ): _raise_reservation_unavailable(counter_key=group[0].counter_key) continue for counter, entry, reserved_value in zip(group, entries, reserved_values): @@ -1023,7 +1091,10 @@ async def _reserve_reservable_counters( current_spend = reserved_value - (charged_cost - reservation_cost) else: current_spend = current_spend_by_counter_key[counter.counter_key] + reservation_cost - if current_spend > counter.max_budget: + if ( + spend_counter_policy.decide(spend_counter_policy.Overflow(current_spend, counter.max_budget)).kind + == "CheckOverflow" + ): reservation_cost = await _apply_over_budget_reservation_policy( counter=counter, valid_token=valid_token, @@ -1044,14 +1115,16 @@ def _reservation_groups( """Every counter the batch read says still has room for the estimate is charged in one pipeline. As soon as one does not, the counters are charged one at a time so the over-budget policy settles each before the next is touched, and a rejection charges nothing after it.""" - if not counters: - return () - if all( - current_spend_by_counter_key[counter.counter_key] + reservation_cost <= counter.max_budget - for counter in counters - ): - return (tuple(counters),) - return tuple((counter,) for counter in counters) + decision: Final = spend_counter_policy.decide( + spend_counter_policy.ReservationGroups( + tuple(current_spend_by_counter_key[counter.counter_key] for counter in counters), + tuple(counter.max_budget for counter in counters), + reservation_cost, + ) + ) + if not isinstance(decision, spend_counter_policy.Groups): + raise NativeAuthRuntimeError("Unexpected reservation grouping decision") + return tuple(tuple(counters[index] for index in group) for group in decision.groups) async def _reserve_counters( @@ -1163,7 +1236,12 @@ async def _set_reserved_entries_actual_cost( ) ) inconsistent: Final = tuple(item for item, ok in zip(adjustments, consistent) if not ok) - if inconsistent and not reseed_on_inconsistent: + if ( + spend_counter_policy.decide( + spend_counter_policy.ResizeConsistency(bool(inconsistent), reseed_on_inconsistent) + ).kind + == "RejectResize" + ): # Pre-call admission resize: the in-flight reservation cost is not yet # persisted, so the DB floor would discard it. Keep the original # fail-closed behavior (raise -> reserve_budget_for_request releases and @@ -1206,10 +1284,7 @@ async def _counter_can_apply_adjustment( current_value, _ = await read_spend_counter_cache_value(counter_key=counter_key) except (TypeError, ValueError): return False - if current_value is None: - return False - - return not (adjustment < 0 and current_value + adjustment < -1e-12) + return spend_counter_policy.decide(spend_counter_policy.Adjustment(current_value, adjustment)).kind == "Compatible" async def _release_applied_entries_best_effort( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8a7c43372d8..e264363e463 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -28,10 +28,62 @@ class NativeAuthContinuation: def request_auth_start(facts: object, settings: object, now: object, resolved: bool) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... def request_auth_budgets(entry: str) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... +def request_auth_access(facts: object) -> dict[str, object]: ... +def request_auth_agent_access(facts: object) -> str: ... +def request_auth_agent_grants(facts: object) -> object: ... +def request_auth_context(facts: object) -> object: ... +def request_auth_reservation(facts: object) -> object: ... +def request_auth_spend_counter_policy(facts: object) -> object: ... +def request_auth_common_flow(facts: object) -> object: ... +def request_auth_common_read(facts: object) -> object: ... +def request_auth_common_budget_scopes(facts: object) -> object: ... +def request_auth_custom_flow(facts: object) -> object: ... +def request_auth_expired(expires: object, now: object) -> bool: ... +def request_auth_agent_invocation(facts: object) -> dict[str, object]: ... +def request_auth_admission(facts: object) -> dict[str, object]: ... +def request_auth_managed(facts: object) -> dict[str, object]: ... +def request_auth_common_budget(facts: object) -> dict[str, object]: ... +def request_auth_budget_alert(facts: object) -> bool: ... +def request_auth_budget_window_start(facts: object) -> float | None: ... +def request_auth_member_limit(facts: object) -> float | None: ... +def request_auth_budget_price(facts: object) -> dict[str, object]: ... +def request_auth_budget_route(facts: object) -> str: ... +def request_auth_model_budget_resolve(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_query(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_check(facts: object) -> dict[str, object]: ... def request_auth_select_api_key(keys: object) -> dict[str, str] | None: ... def request_auth_select_websocket_api_key(keys: object) -> dict[str, str]: ... def request_auth_scheme_credential(value: str) -> str | None: ... def request_auth_litellm_header_token(value: str) -> str: ... +def request_auth_route_access(request: object) -> object: ... +def request_auth_org_route(facts: object) -> object: ... +def request_auth_org_admin(memberships: object, organizations: tuple[str, ...]) -> bool: ... +def request_auth_model_access(request: object) -> object: ... +def request_auth_model_groups(facts: object) -> object: ... +def request_auth_fallback_access(facts: object) -> object: ... +def request_auth_key_model_access(facts: object) -> object: ... +def request_auth_unpriced_model(facts: object) -> object: ... +def request_auth_model_attribution(facts: object) -> object: ... +def request_auth_fallback_candidate(facts: object) -> object: ... +def request_auth_personal_budget(facts: object) -> object: ... +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_personal_owner(facts: object) -> object: ... +def request_auth_personal_model_request(facts: object) -> object: ... +def request_auth_member_models(facts: object) -> object: ... +def request_auth_resolved_models(facts: object) -> object: ... +def request_auth_personal_model_listing(facts: object) -> object: ... +def request_auth_agent_ceiling(facts: object) -> object: ... +def request_auth_agent_group_load(facts: object) -> object: ... +def request_auth_team_aliases(facts: object) -> object: ... +def request_auth_managed_models(facts: object) -> object: ... +def request_auth_agent_ceiling_source(facts: object) -> object: ... +def request_auth_managed_model_grants(facts: object) -> object: ... +def request_auth_model_projection(facts: object, reads: object) -> object: ... +def request_auth_route_enable(facts: object) -> object: ... +def request_auth_llm_route(context: object, configured: bool | None = None) -> object: ... @final class NativeClickHouseSpendConfig: @@ -1182,3 +1234,54 @@ class RouterCall(Generic[_S]): # candidate ids, pick first; pre_route: see _PreRouteStep; cache: see _CacheStep; secrets, build and model_info: # see those steps; first call: None. def next(self, outcome: object) -> _S: ... + +def request_auth_budget_exemption(facts: object) -> dict[str, object]: ... + +def request_auth_budget_alert_plan(facts: object) -> dict[str, object]: ... + +def request_auth_context_models(facts: object) -> object: ... + +def request_auth_context_end_user(facts: object) -> object: ... + +def request_auth_context_token(facts: object) -> object: ... + +def request_auth_context_member(facts: object) -> object: ... + +def request_auth_context_temporary(facts: object) -> object: ... + +def request_auth_context_throttle(facts: object) -> object: ... + +def request_auth_context_percentage(facts: object) -> float | None: ... + +def request_auth_end_user_normalize_id(facts: object) -> object: ... + +def request_auth_end_user_default_id(facts: object) -> object: ... + +def request_auth_end_user_apply(facts: object) -> object: ... + +def request_auth_end_user_registry(facts: object) -> object: ... + +def request_auth_end_user_assignments(facts: object) -> object: ... + +def request_auth_end_user_lookup(facts: object) -> object: ... + +def request_auth_end_user_validation(facts: object) -> object: ... + +def request_auth_budget_lookup_organization(facts: object) -> object: ... +def request_auth_budget_lookup_member(facts: object) -> object: ... +def request_auth_budget_lookup_emails(facts: object) -> object: ... +def request_auth_budget_lookup_soft(facts: object) -> bool: ... + +def request_auth_route_role(facts: object) -> object: ... + +def request_auth_route_role_api(facts: object) -> object: ... + +def request_auth_route_role_jwt(facts: object) -> object: ... + +def request_auth_route_role_inside(facts: object) -> object: ... + +def request_auth_route_role_expand(facts: object) -> object: ... + +def request_auth_route_role_listed(facts: object) -> object: ... + +def request_auth_org_admin_raw(facts: object) -> object: ... diff --git a/tests/unit/_auth/budget_reference.py b/tests/unit/_auth/budget_reference.py new file mode 100644 index 00000000000..a0219ea8259 --- /dev/null +++ b/tests/unit/_auth/budget_reference.py @@ -0,0 +1,117 @@ +import math +from dataclasses import replace +from typing import Final + +from litellm._auth.common_budgets import BudgetDecision, BudgetFacts +from litellm._auth.budget_pricing import ExemptionFacts, ModelPriceFacts +from litellm._auth.model_budget_data import ModelBudgetDecision, ModelBudgetSnapshot + + +def price(f: ModelPriceFacts) -> bool: + if f.cached is not None: + return f.cached + if f.lookup_failed or not f.group_found or f.input_price is None or f.output_price is None: + return False + if f.input_price > 0 or f.output_price > 0: + return False + deployments: Final = f.deployments or () + return any(d.has_input_price or d.has_output_price for d in deployments) and not any( + (d.input_price is not None and d.input_price > 0) + or (d.output_price is not None and d.output_price > 0) + or (d.query_price is not None and d.query_price > 0) + or (d.ptu_count is not None and d.cost_per_ptu_per_hour is not None) + for d in deployments + ) + + +def exemption(f: ExemptionFacts) -> bool: + route: Final = f.route + if not route.side_effect_route and ( + route.discovery_route or (not route.llm_route and not route.registered_pass_through) + ): + return True + if route.body_present and ( + route.route in ("/mcp-rest/tools/list", "/v1/mcp/tools") + or route.route not in ("/mcp/tools/call", "/mcp-rest/tools/call") + and (route.route == "/mcp" or route.route.startswith("/mcp/")) + and route.method in ("initialize", "notifications/initialized", "ping", "tools/list") + ): + return True + return f.model_input_present and f.router_present and all(price(p) for p in f.prices) + + +def common(f: BudgetFacts) -> BudgetDecision: + limit: Final = f.limit + if limit is None: + return BudgetDecision() + if f.scope == "User" and f.team_key and not f.apply_user_budget_to_team_keys: + return BudgetDecision() + if f.scope == "EndUser" and f.info_route: + return BudgetDecision() + if f.scope == "Project" and not math.isfinite(limit): + return BudgetDecision() + if f.scope == "AccessGroup" and limit <= 0: + return BudgetDecision() + if f.scope == "Proxy" and (not limit > 0 or f.skip or not f.llm_route or f.model_list_route): + return BudgetDecision() + if f.spend is None: + return BudgetDecision(read_spend=f.scope not in ("KeySoft", "TeamSoft", "ProjectSoft", "Proxy")) + if f.scope.endswith("Soft"): + enabled: Final = { + "KeySoft": bool(limit), + "TeamSoft": f.token_present and bool(f.alert_recipients), + "ProjectSoft": f.token_present, + }[f.scope] + return BudgetDecision(alert="soft_budget" if enabled and f.spend >= limit else None) + denied: Final = ( + not f.spend <= limit + if f.scope == "Tag" + else not f.spend < limit + if f.scope == "AccessGroup" + else f.spend > limit + if f.scope == "EndUser" + else math.isfinite(limit) and (f.spend > limit if f.scope in ("Team", "Proxy") else f.spend >= limit) + ) + result: Final = BudgetDecision( + denied=denied, + alert=( + "token_budget" + if f.scope == "Key" + else "user_budget" + if f.scope == "User" + else "team_budget" + if f.scope == "Team" and denied and f.token_present + else "project_budget" + if f.scope == "Project" and denied and f.token_present + else "organization_budget" + if f.scope == "Organization" and denied + else "team_member_threshold" + if f.scope == "TeamMember" and math.isfinite(limit) + else None + ), + ) + if ( + f.scope == "Key" + and denied + and f.throttle_opt_in + and (f.tpm_limit is not None or f.rpm_limit is not None) + and f.throttle_percentage is not None + and 0 < f.throttle_percentage <= 1 + ): + return replace(result, denied=False, throttle_percentage=f.throttle_percentage) + return result + + +def model(f: ModelBudgetSnapshot) -> ModelBudgetDecision: + local: Final = (f.local_spend or 0.0) + (f.legacy_spend or 0.0) + current: Final = ( + max(local, f.shared_spend) + if f.shared_spend is not None + and f.shared_started_at is not None + and f.shared_window_seconds is not None + and f.now < f.shared_started_at + f.shared_window_seconds + else local + ) + return ModelBudgetDecision( + None, f.max_budget is not None and not f.max_budget < 0 and current >= f.max_budget, current + ) diff --git a/tests/unit/_auth/conftest.py b/tests/unit/_auth/conftest.py index 1e171ca0db5..e114732e771 100644 --- a/tests/unit/_auth/conftest.py +++ b/tests/unit/_auth/conftest.py @@ -3,10 +3,10 @@ from typing import Final, Protocol, cast import pytest -from litellm._auth import runtime, rust_core +from litellm._auth import credential, rust_core from litellm._auth.decision import AuthCore -from litellm._auth.python_core import PythonCore -from litellm._auth.shadow import ShadowCore +from tests.unit._auth.reference_core import PythonCore +from tests.unit._auth.shadow import ShadowCore from litellm._auth.credential import PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys @@ -23,7 +23,6 @@ def core(request: pytest.FixtureRequest) -> AuthCore: return factory() -@pytest.fixture(params=("python", "rust", "shadow")) -def credentials(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> CredentialSelector: - monkeypatch.setenv("LITELLM_AUTH_CORE", cast(str, request.param)) - return runtime +@pytest.fixture(params=("python", "rust")) +def credentials(request: pytest.FixtureRequest) -> CredentialSelector: + return credential if cast(str, request.param) == "python" else rust_core diff --git a/tests/unit/_auth/raw_answers.py b/tests/unit/_auth/raw_answers.py new file mode 100644 index 00000000000..3ef95c7c294 --- /dev/null +++ b/tests/unit/_auth/raw_answers.py @@ -0,0 +1,40 @@ +from litellm._auth.decision import ModelBudgetExceeded, ModelBudgetRead +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.model_budget_data import ModelBudgetSnapshot + + +def fallback_read(model: str, authorized: bool = True) -> FallbackAuthorizationFacts: + return FallbackAuthorizationFacts( + key=GroupedModelScope( + ModelAccessFacts((model,), (model,) if authorized else ("not-allowed",), object_type="key") + ) + ) + + +def budget_read(exceeded: ModelBudgetExceeded | None = None) -> ModelBudgetRead: + if exceeded is None: + return ModelBudgetRead(ModelBudgetSnapshot()) + return ModelBudgetRead( + ModelBudgetSnapshot(max_budget=exceeded.max_budget, local_spend=exceeded.current_cost), + exceeded.message, + exceeded.entity_type, + exceeded.entity_id, + ) + + +def candidate_read(model: str, usable: bool = True) -> FallbackCandidateFacts: + return FallbackCandidateFacts( + router_present=True, + user_present=True, + team_present=True, + membership="Loaded", + price=ModelPriceFacts(cached=True), + authorization=fallback_read(model, authorized=usable), + attribution_reset=True, + attribution=ModelGroupAttributionFacts(), + attribution_stamped=True, + ) diff --git a/tests/unit/_auth/raw_model_reference.py b/tests/unit/_auth/raw_model_reference.py new file mode 100644 index 00000000000..bc484145f0e --- /dev/null +++ b/tests/unit/_auth/raw_model_reference.py @@ -0,0 +1,52 @@ +from dataclasses import replace +from typing import Final + +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from litellm._auth.model_access import ModelAccessDecision +from tests.unit._auth.support.model_reference import model_access + + +def grouped(scope: GroupedModelScope) -> ModelAccessDecision: + direct: Final = model_access(scope.facts) + if direct.kind == "Allow" or not scope.group_ids: + return direct + if scope.group_models is None: + raise AssertionError("unread resource groups") + return model_access(replace(scope.facts, grants=scope.group_models)) if scope.group_models else direct + + +def initial_access(facts: KeyAndFallbackModelAccessFacts) -> ModelAccessDecision | None: + if facts.config_nonempty or "all-team-models" in facts.key_models: + return None + if not facts.primary_loaded: + raise AssertionError("unread primary") + if facts.primary is not None: + primary: Final = grouped(facts.primary) + if primary.kind != "Allow": + return primary + if facts.fallback_names is None: + raise AssertionError("unread fallbacks") + for index, _ in enumerate(facts.fallback_names): + if (result := grouped(facts.fallback_scopes[index])).kind != "Allow": + return result + if facts.routing[index] != "Routed": + raise AssertionError("routing failure should be projected by the effect") + return None + + +def fallback_access(facts: FallbackAuthorizationFacts) -> bool: + if facts.refuse_agent_keys and facts.agent_id is not None: + return False + if facts.key is None: + raise AssertionError("unread key scope") + if grouped(facts.key).kind != "Allow": + return False + if facts.team is not None and grouped(facts.team).kind != "Allow": + assigned: Final = replace(facts.team.facts, grants=facts.assigned_models or ()) + if not facts.assigned_models or model_access(assigned).kind != "Allow": + return False + return all( + scope is None or model_access(scope).kind == "Allow" + for scope in (facts.embedded_team, facts.member, facts.project) + ) diff --git a/litellm/_auth/python_core.py b/tests/unit/_auth/reference_core.py similarity index 93% rename from litellm/_auth/python_core.py rename to tests/unit/_auth/reference_core.py index c7d6210bdc6..6cb7a030233 100644 --- a/litellm/_auth/python_core.py +++ b/tests/unit/_auth/reference_core.py @@ -28,12 +28,14 @@ from litellm._auth.decision import ( MasterKey, MasterKeyOnlyRoute, MissingKey, + ModelAccessDenied, NoDatabase, NoMasterKey, NonAdminKey, NotAVirtualKey, OwnerDeactivated, OwnerFacts, + PrefetchAuthObjects, PrepareKey, PresentedKey, ReadKeySpend, @@ -54,7 +56,8 @@ from litellm._auth.decision import ( WithinModelBudgets, ) from litellm._auth.key import hash_token, is_master_key -from litellm._auth.model_budgets import ( +from tests.unit._auth.raw_model_reference import initial_access +from tests.unit._auth.reference_model_budgets import ( check_model_budgets, check_team_member_model_budgets, check_user_model_budgets, @@ -148,7 +151,13 @@ def _check_key(check: _KeyCheck) -> Step: def _after_inspect(check: _KeyCheck, key: KeyFacts) -> Step: if key.blocked: return Deny(KeyBlocked()) - return EnforceModelAccess(resume=lambda: _check_owner(check, key)) + return EnforceModelAccess( + resume=lambda facts: ( + Deny(ModelAccessDenied(denial)) + if (denial := initial_access(facts)) is not None + else PrefetchAuthObjects(resume=lambda: _check_owner(check, key)) + ) + ) def _check_owner(check: _KeyCheck, key: KeyFacts) -> Step: @@ -165,7 +174,9 @@ def _after_owner(check: _KeyCheck, key: KeyFacts, user_id: str, owner: OwnerFact def _check_budget_exemption(check: _KeyCheck, key: KeyFacts) -> Step: - return CheckBudgetExemption(resume=lambda exempt: _check_team_member_budget(check, key, exempt)) + from tests.unit._auth.budget_reference import exemption + + return CheckBudgetExemption(resume=lambda facts: _check_team_member_budget(check, key, exemption(facts))) def _check_team_member_budget(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: diff --git a/litellm/_auth/model_budgets.py b/tests/unit/_auth/reference_model_budgets.py similarity index 74% rename from litellm/_auth/model_budgets.py rename to tests/unit/_auth/reference_model_budgets.py index 7bf2e109287..07bc4ae2315 100644 --- a/litellm/_auth/model_budgets.py +++ b/tests/unit/_auth/reference_model_budgets.py @@ -1,3 +1,4 @@ +from tests.unit._auth.test_fallback_candidate import reference as candidate_reference from collections.abc import Callable from typing import Final @@ -10,6 +11,7 @@ from litellm._auth.decision import ( Deny, FallbackTarget, ModelBudgetExceeded, + ModelBudgetRead, ModelBudgetScope, ReadModelBudgetScope, ReadTeamMemberFallbacks, @@ -22,6 +24,18 @@ from litellm._auth.decision import ( VetTeamMemberFallback, WithinModelBudgets, ) +from tests.unit._auth.budget_reference import model +from tests.unit._auth.raw_model_reference import fallback_access + + +def exceeded(reading: ModelBudgetRead) -> ModelBudgetExceeded | None: + result: Final = model(reading.snapshot) + if not result.denied: + return None + assert reading.snapshot.max_budget is not None + return ModelBudgetExceeded( + reading.message, result.current_spend, reading.snapshot.max_budget, reading.entity_type, reading.entity_id + ) def check_model_budgets(then: Callable[[], Step]) -> Step: @@ -58,11 +72,13 @@ def _check_key_budget_of( return Deny(exceeded) return _first_fallback_within_budget( scope.key_budget_fallbacks.get(model, ()), - found=lambda fallback: _reroute(fallback, exceeded, then=lambda: _check_key_budget_of(scope, rest, fallback, then)), + found=lambda fallback: _reroute( + fallback, exceeded, then=lambda: _check_key_budget_of(scope, rest, fallback, then) + ), exhausted=lambda: Deny(exceeded), ) - return CheckKeyModelBudget(model=model, resume=after) + return CheckKeyModelBudget(model=model, resume=lambda reading: after(exceeded(reading))) def _first_fallback_within_budget( @@ -73,9 +89,9 @@ def _first_fallback_within_budget( candidate: Final = candidates[0] return CheckKeyModelBudget( model=candidate, - resume=lambda exceeded: ( + resume=lambda reading: ( found(candidate) - if exceeded is None + if exceeded(reading) is None else _first_fallback_within_budget(candidates[1:], found, exhausted) ), ) @@ -84,7 +100,9 @@ def _first_fallback_within_budget( def _reroute(fallback: str, exceeded: ModelBudgetExceeded, then: Callable[[], Step]) -> Step: return AuthorizeBudgetFallback( model=fallback, - resume=lambda authorized: RewriteRequestModel(model=fallback, resume=then) if authorized else Deny(exceeded), + resume=lambda facts: ( + RewriteRequestModel(model=fallback, resume=then) if fallback_access(facts) else Deny(exceeded) + ), ) @@ -93,7 +111,9 @@ def _check_user_budgets(scope: ModelBudgetScope, models: tuple[str, ...], then: return then() return CheckUserModelBudget( model=models[0], - resume=lambda exceeded: _check_user_budgets(scope, models[1:], then) if exceeded is None else Deny(exceeded), + resume=lambda reading: ( + _check_user_budgets(scope, models[1:], then) if (failure := exceeded(reading)) is None else Deny(failure) + ), ) @@ -134,7 +154,7 @@ def _check_team_member_budget_of( ), ) - return CheckTeamMemberModelBudget(model=model, resume=after) + return CheckTeamMemberModelBudget(model=model, resume=lambda reading: after(exceeded(reading))) def _first_usable_fallback( @@ -147,5 +167,9 @@ def _first_usable_fallback( target: Final = targets[0] return VetTeamMemberFallback( target=target, - resume=lambda usable: found(target) if usable else _first_usable_fallback(targets[1:], found, exhausted), + resume=lambda facts: ( + found(target) + if candidate_reference(facts).usable + else _first_usable_fallback(targets[1:], found, exhausted) + ), ) diff --git a/tests/unit/_auth/reference_policy.py b/tests/unit/_auth/reference_policy.py new file mode 100644 index 00000000000..27a0951b551 --- /dev/null +++ b/tests/unit/_auth/reference_policy.py @@ -0,0 +1,186 @@ +from collections.abc import Mapping +from typing import Final + +from litellm._auth import access, admission, managed + + +def access_decision(facts: access.AccessFacts) -> access.AccessDecision: + match facts: + case access.ConcurrentFailures(failures=failures, terminal=terminal): + return access.AccessDecision( + failure_index=next( + (index for index, failed in enumerate(terminal) if failed), + next((index for index, failed in enumerate(failures) if failed), None), + ) + ) + case access.UserParamRequest(enabled=enabled): + return access.AccessDecision(read="UserParam" if enabled else None) + case access.ClientTagRequest(enabled=enabled, llm_route=llm_route): + return access.AccessDecision( + read=None + if not enabled or llm_route is False + else "ClientTagRoute" + if llm_route is None + else "ClientTags" + ) + case access.GuardrailRequest(requested_keys=keys): + modifies: Final = any( + key in keys + for key in ( + "guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", + ) + ) + return access.AccessDecision(read="GuardrailPermission" if modifies else None) + case access.TraceRequest(required=required): + return access.AccessDecision(next_trace="Gate" if required else None) + case access.Search(name=name, allowed=allowed, strict=strict): + return access.AccessDecision(denial="Search" if (strict or allowed) and name not in allowed else None) + case access.Blocked(blocked=blocked): + return access.AccessDecision(denial="Blocked" if blocked else None) + case access.UserParam( + enabled=enabled, method=method, llm_route=llm_route, mcp_route=mcp_route, has_user=has_user + ): + return access.AccessDecision( + denial="UserParam" + if enabled and method and method.upper() == "POST" and llm_route and not mcp_route and not has_user + else None + ) + case access.ClientTags(enabled=enabled, llm_route=llm_route, has_tags=has_tags): + return access.AccessDecision(denial="ClientTags" if enabled and llm_route and has_tags else None) + case access.Guardrails(requested_keys=keys, modify_guardrails=modify): + return access.AccessDecision( + denial="Guardrails" + if modify is False + and any( + key in keys + for key in ( + "guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", + ) + ) + else None + ) + case access.AgentTrace(required=required, header=header): + return access.AccessDecision(denial="AgentTrace" if required and not header else None) + case access.Tools(requested=requested, key_allowed=key, team_allowed=team): + tool_allowed: Final = key or team + rejected: Final = tuple(tool for tool in requested if tool_allowed and tool not in tool_allowed) + return access.AccessDecision(denial="Tools" if rejected else None, rejected=rejected) + case access.GrantLayers(): + if not facts.deny_by_default or not facts.identity_present or facts.master_key or facts.dashboard: + return access.AccessDecision() + layers: Final[tuple[access.GrantLayer, ...]] = ( + *(("key",) if facts.virtual_key and not facts.session_token else ()), + *(("team",) if facts.team_present else ()), + ) + return access.AccessDecision(layers=layers or ("user",)) + case access.VectorStores(requested=requested, allowed=allowed, strict=strict): + if not allowed: + return access.AccessDecision(denial="MissingVectorGrant" if strict else None) + rejected_index: Final = next((index for index, store in enumerate(requested) if store not in allowed), None) + return access.AccessDecision( + denial="VectorStores" if rejected_index is not None else None, rejected_index=rejected_index + ) + case access.TeamLookup(dashboard=dashboard, missing=missing, recorded_models=models, allow_unavailable=allow): + if not dashboard and not missing and not models and allow is None: + return access.AccessDecision(read="UnavailablePolicy") + return access.AccessDecision( + denial="TeamLookup" if not dashboard and (missing or (not models and not allow)) else None + ) + case access.TraceGate( + required=required, method=method, llm_route=llm, passthrough=passthrough, telemetry=telemetry + ): + return access.AccessDecision( + next_trace="Header" + if required and method not in ("GET", "HEAD", "OPTIONS") and (llm or passthrough) and not telemetry + else None + ) + case access.TraceStep(source=source, value=value, present=present, headers_only=headers_only): + if value: + return access.AccessDecision() + if source == "Header" and not headers_only: + return access.AccessDecision(next_trace="Body") + if source == "Body": + return access.AccessDecision(next_trace="Selected" if present else "Traceparent") + if source == "Traceparent": + return access.AccessDecision(next_trace="Selected") + if source == "Selected" and not present: + return access.AccessDecision(next_trace="Other") + return access.AccessDecision(denial="TeamTrace") + case access.TeamTrace(): + gate: Final = access_decision( + access.TraceGate(facts.required, facts.method, facts.llm_route, facts.passthrough, facts.telemetry) + ) + trace_present: Final = facts.header or ( + not facts.headers_only + and ( + facts.body_trace + or (not facts.body_trace_present and facts.traceparent_id) + or (facts.selected_trace if facts.selected_trace_present else facts.other_trace) + ) + ) + return access.AccessDecision(denial="TeamTrace" if gate.next_trace and not trace_present else None) + + +def common_plan(facts: admission.CommonGate) -> admission.CommonPlan: + if facts.force_virtual_key_checks: + return admission.CommonPlan(None) + if not facts.force_virtual_key_checks and (facts.public_route or facts.additional_public_route): + return admission.CommonPlan("PublicRoute") + if facts.unauthenticated_passthrough: + return admission.CommonPlan("UnauthenticatedPassthrough") + if ( + not facts.master_key_present + and not facts.jwt_enabled + and not facts.oauth_enabled + and not facts.oauth_proxy_enabled + ): + return admission.CommonPlan("NoAuthDevelopment") + if facts.custom_auth and not facts.custom_auth_common_checks: + return admission.CommonPlan("CustomAuth") + return admission.CommonPlan(None) + + +def reservation_plan(facts: admission.Reservation) -> admission.ReservationPlan: + return admission.ReservationPlan( + not facts.skip_budget_checks and not facts.disabled, + facts.apply_user_budget_to_team_keys, + facts.fail_closed or facts.billing_agent, + ) + + +def actor(facts: managed.Actor) -> managed.ManagedDecision: + agent, context = facts.agent, facts.context + if not agent.enabled or agent.identity is None or not agent.identity.active: + return managed.ManagedDecision("Deny", "Disabled", None) + if context is None: + return managed.ManagedDecision("Deny", "MissingContext", None) + if context.agent_id != agent.agent_id or context.binding_revision != agent.identity.revision: + return managed.ManagedDecision("Deny", "ChangedIdentity", None) + if agent.execution_mode not in (context.mode, "both"): + return managed.ManagedDecision("Deny", "ExecutionMode", None) + if context.mode == "delegated" and not context.user_id: + return managed.ManagedDecision("Deny", "MissingHuman", None) + return managed.ManagedDecision("Allow", None, None) + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + if entry.get("tiered_pricing") is not None: + return True + for key, value in entry.items(): + if "cost_per" not in key: + continue + if _is_positive_cost(value): + return True + if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()): + return True + return False diff --git a/litellm/_auth/shadow.py b/tests/unit/_auth/shadow.py similarity index 98% rename from litellm/_auth/shadow.py rename to tests/unit/_auth/shadow.py index fc3bd4073e2..af753d85178 100644 --- a/litellm/_auth/shadow.py +++ b/tests/unit/_auth/shadow.py @@ -6,7 +6,7 @@ from functools import reduce from typing import Final from litellm._auth import decision -from litellm._auth.python_core import PythonCore +from tests.unit._auth.reference_core import PythonCore from litellm._auth.rust_core import ( AUTH_STEP, MEMBER_STEP, diff --git a/tests/unit/_auth/support/fallback_reference.py b/tests/unit/_auth/support/fallback_reference.py new file mode 100644 index 00000000000..d3c47e288a7 --- /dev/null +++ b/tests/unit/_auth/support/fallback_reference.py @@ -0,0 +1,68 @@ +from dataclasses import replace +from typing import Final + +from litellm._auth.fallback_access import FallbackAuthorizationDecision as Decision +from litellm._auth.fallback_access import FallbackAuthorizationFacts as Facts +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision +from tests.unit._auth.support.model_reference import model_access + + +def grouped(scope: GroupedModelScope, *, merge: bool = False) -> ModelAccessDecision: + primary: Final = model_access(scope.facts) + if primary.kind != "Deny" or not scope.group_ids: + return primary + if scope.group_models is None: + return ModelAccessDecision("ReadGroups") + if not scope.group_models: + return primary + grants: Final = tuple(dict.fromkeys((*scope.facts.grants, *scope.group_models))) if merge else scope.group_models + return model_access(replace(scope.facts, grants=grants)) + + +def authorization(facts: Facts) -> Decision: + if facts.refuse_agent_keys and facts.agent_id is not None: + return Decision("Deny") + if facts.key is None: + return Decision("ReadKey") + key: Final = grouped(facts.key) + if key.kind == "ReadGroups": + return Decision("ReadKeyGroups") + if key.kind != "Allow": + return Decision("Deny", key) + if facts.team_present and facts.team is None: + return Decision("ReadTeam") + if facts.team is not None: + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if team.kind != "Allow": + if team.kind != "Deny": + return Decision("Deny", team) + if facts.assigned_models is None: + return Decision("ReadAssignedGroups") + if not facts.assigned_models: + return Decision("Deny", team) + if facts.assigned_scope is None: + return Decision("ReadAssignedScope") + if model_access(facts.assigned_scope).kind != "Allow": + return Decision("Deny", team) + elif facts.embedded_team is None and facts.embedded_team_grants: + return Decision("ReadEmbeddedTeam") + elif facts.embedded_team is not None: + embedded: Final = model_access(facts.embedded_team) + if embedded.kind != "Allow": + return Decision("Deny", embedded) + if facts.member is None and facts.member_grants: + return Decision("ReadMember") + if facts.member is not None: + member: Final = model_access(facts.member) + if member.kind != "Allow": + return Decision("Deny", member) + if facts.project is None and facts.project_grants: + return Decision("ReadProject") + if facts.project is not None: + project: Final = model_access(facts.project) + if project.kind != "Allow": + return Decision("Deny", project) + return Decision("Allow") diff --git a/tests/unit/_auth/support/model_reference.py b/tests/unit/_auth/support/model_reference.py new file mode 100644 index 00000000000..e9155f8d256 --- /dev/null +++ b/tests/unit/_auth/support/model_reference.py @@ -0,0 +1,89 @@ +import re +from dataclasses import replace +from typing import Final + +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts + + +def _pattern(model: str, pattern: str) -> bool: + return "*" in pattern and re.match("^" + pattern.replace("*", ".*") + "$", model) is not None + + +def helper(model: str, facts: ModelAccessFacts, *, team_aliases: bool = True) -> bool: + ordinary: Final = tuple(grant for grant in facts.grants if grant != "all-team-models") + expanded: Final = tuple(dict.fromkeys((*ordinary, *facts.registry.names))) + grants: Final = ( + (expanded or facts.grants) + if "all-team-models" in facts.grants and facts.team_id is not None and facts.registry.present + else facts.grants + ) + record: Final = facts.registry.records.get(model) + groups: Final = record.access_groups if record is not None else () + if facts.registry.present and any(grant in groups for grant in grants): + return True + filtered: Final = tuple(grant for grant in grants if grant not in groups) + if team_aliases and model in facts.aliases.team: + return True + if any(_pattern(model, grant) for grant in filtered): + return True + provider: Final = record.provider if record is not None else None + if ( + provider is not None + and (provider.stripped_model != model or "/" not in model) + and any(_pattern(f"{provider.provider}/{provider.stripped_model}", grant) for grant in filtered) + ): + return True + return not grants or "*" in filtered or "all-proxy-models" in filtered or model in filtered + + +def _single(model: str, facts: ModelAccessFacts) -> ModelAccessDecision: + after_team: Final = facts.aliases.team.get(model, model) + after_key: Final = facts.aliases.key.get(after_team, after_team) + after_global: Final = facts.aliases.global_.get(after_key, after_key) + dispatched: Final = facts.aliases.key.get(after_global, after_global) + key_applied: Final = after_key != after_team or dispatched != after_global + alias: Final = ( + facts.aliases.global_.get(model) if model in facts.aliases.global_ else facts.registry.aliases.get(model) + ) + parent: Final = facts.registry.compaction_parents.get(model) + potential: Final = ( + (dispatched,) + if key_applied + else (model, *((parent,) if parent is not None else ()), *((alias,) if alias else ())) + ) + if any(helper(candidate, facts) for candidate in potential): + return ModelAccessDecision("Allow") + record: Final = facts.registry.records.get(model) + public: Final = ( + record.deployment_public_name if record is not None and record.deployment_team == facts.team_id else None + ) + owning: Final = ( + ( + *((public,) if public else ()), + *(alias for alias, target in facts.aliases.team.items() if target == model), + ) + if not key_applied + and facts.registry.present + and facts.team_id is not None + and model.startswith(f"model_name_{facts.team_id}_") + else () + ) + if any(helper(candidate, facts, team_aliases=False) for candidate in owning): + return ModelAccessDecision("Allow") + return ModelAccessDecision("Deny", model, facts.grants, facts.object_type) + + +def _list_item(model: str, facts: ModelAccessFacts) -> ModelAccessDecision: + decision: Final = model_access( + replace(facts, requested=(model,), fallback_depth=facts.fallback_depth + 1, is_list=False) + ) + return replace(decision, model=model) if decision.kind == "Depth" else decision + + +def model_access(facts: ModelAccessFacts) -> ModelAccessDecision: + if facts.fallback_depth >= facts.max_depth: + return ModelAccessDecision("Depth") + decisions: Final = ( + _list_item(model, facts) if facts.is_list else _single(model, facts) for model in facts.requested + ) + return next((decision for decision in decisions if decision.kind != "Allow"), ModelAccessDecision("Allow")) diff --git a/tests/unit/_auth/support/model_registry_reference.py b/tests/unit/_auth/support/model_registry_reference.py new file mode 100644 index 00000000000..d499bab565e --- /dev/null +++ b/tests/unit/_auth/support/model_registry_reference.py @@ -0,0 +1,187 @@ +"""Model authorization and registry ordering frozen from Phase 1, 8bd7f5089a75.""" + +import re +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy._types import LiteLLM_UserTable, ModelAccessDeniedProxyException, ProxyErrorTypes +from litellm.proxy.auth.model_access_denied import ( + customer_model_access_denied_client_message, + model_access_denied_client_message, +) +from litellm.router_api import RegistryReader +from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent + + +def _pattern(model: str, pattern: str) -> bool: + return "*" in pattern and re.match("^" + pattern.replace("*", ".*") + "$", model) is not None + + +def _provider_pattern(model: str, pattern: str) -> bool: + try: + stripped, provider, _, _ = litellm.get_llm_provider(model=model) + except Exception: # noqa: BLE001 # the Phase 1 provider inference helper swallowed lookup failures + return False + if stripped == model and "/" in model: + return False + return _pattern(f"{provider}/{stripped}", pattern) + + +def _helper( + model: str, + router: RegistryReader | None, + grants: Sequence[str], + team: str | None, + aliases: Mapping[str, str] | None, +) -> bool: + groups: Final = router.registry().access_groups(model_name=model, team_id=team) if router else {} + names: Final = ( + router.registry().names() if "all-team-models" in grants and team is not None and router is not None else None + ) + ordinary: Final = tuple(grant for grant in grants if grant != "all-team-models") + expanded: Final = ( + (tuple(dict.fromkeys((*ordinary, *names))) or tuple(grants)) if names is not None else tuple(grants) + ) + if groups and any(grant in groups for grant in expanded): + return True + filtered: Final = tuple(grant for grant in expanded if grant not in groups) + if aliases and model in aliases: + return True + if any(_pattern(model, grant) for grant in filtered): + return True + if any("*" in grant and _provider_pattern(model, grant) for grant in filtered): + return True + return (not filtered and not expanded) or "*" in filtered or "all-proxy-models" in filtered or model in filtered + + +def _owning( + model: str, router: RegistryReader | None, team: str | None, aliases: Mapping[str, str] | None +) -> tuple[str, ...]: + if router is None or team is None or not model.startswith(f"model_name_{team}_"): + return () + deployment: Final = router.registry().deployments(model, source="first_in_group") + public: Final = ( + deployment.model_info.team_public_model_name + if deployment is not None and deployment.model_info.team_id == team + else None + ) + return (*((public,) if public else ()), *(alias for alias, target in (aliases or {}).items() if target == model)) + + +def can_object_call_model( + model: str | Sequence[str], + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, + key_model_aliases: Mapping[str, str] | None = None, + object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"] = "user", + fallback_depth: int = 0, +) -> Literal[True]: + if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: + raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") # noqa: TRY002 # frozen Phase 1 exception contract + if not isinstance(model, str): + for item in model: + can_object_call_model( + item, + llm_router, + models, + team_model_aliases, + team_id, + key_model_aliases, + object_type, + fallback_depth + 1, + ) + return True + return _single(model, llm_router, models, team_model_aliases, team_id, key_model_aliases, object_type) + + +def _single( + model: str, + router: RegistryReader | None, + grants: Sequence[str], + team_aliases: Mapping[str, str] | None, + team: str | None, + key_aliases: Mapping[str, str] | None, + object_type: Literal["user", "customer", "team", "key", "org", "project", "agent"], +) -> Literal[True]: + parent: Final = native_compaction_parent(model) + target: Final = ( + litellm.model_alias_map[model] + if model in litellm.model_alias_map + else router.registry().resolve(model, kind="alias") + if router is not None and model in router.registry().aliases() + else None + ) + after_team: Final = (team_aliases or {}).get(model, model) + after_key: Final = (key_aliases or {}).get(after_team, after_team) + after_global: Final = litellm.model_alias_map.get(after_key, after_key) + dispatched: Final = (key_aliases or {}).get(after_global, after_global) + key_applied: Final = after_key != after_team or dispatched != after_global + candidates: Final = ( + (dispatched,) + if key_applied + else (model, *((parent,) if parent is not None else ()), *((target,) if target else ())) + ) + owning: Final = () if key_applied else _owning(model, router, team, team_aliases) + if any(_helper(candidate, router, grants, team, team_aliases) for candidate in candidates): + return True + if any(_helper(candidate, router, grants, team, None) for candidate in owning): + return True + raise ModelAccessDeniedProxyException( + message=customer_model_access_denied_client_message(model) + if object_type == "customer" + else model_access_denied_client_message(model), + internal_message=f"{object_type} not allowed to access model. This {object_type} can only access models={list(grants)}. Tried to access {model}", + type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type), + param="model", + code=403, + ) + + +class _Grants(BaseModel, frozen=True): + models: tuple[str, ...] = () + + +async def can_user_call_model( + model: str, llm_router: RegistryReader | None, user_object: LiteLLM_UserTable | None +) -> Literal[True]: + if user_object is None: + return True + grants: Final = _Grants.model_validate(user_object, from_attributes=True).models + if "no-default-models" in grants: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model), + internal_message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + if not grants: + return True + return can_object_call_model(model, llm_router, grants) + + +def entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + if "tiered_pricing" in entry and entry["tiered_pricing"] is not None: + return True + return any(_priced(key, value) for key, value in entry.items()) + + +def _priced(key: str, value: object) -> bool: + if "cost_per" not in key: + return False + if isinstance(value, bool): + return False + if isinstance(value, (int, float)): + return value > 0 + if isinstance(value, dict): + return any( + isinstance(number, (int, float)) and not isinstance(number, bool) and number > 0 + for number in TypeAdapter(Mapping[object, object]).validate_python(value).values() + ) + return False diff --git a/tests/unit/_auth/support/route_enable_reference.py b/tests/unit/_auth/support/route_enable_reference.py new file mode 100644 index 00000000000..63c15a2a4f6 --- /dev/null +++ b/tests/unit/_auth/support/route_enable_reference.py @@ -0,0 +1,39 @@ +from typing import Final + +from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting +from tests.unit._auth.support.route_reference import RouteChecks + + +def _setting(facts: RouteSetting | None, feature: RouteFeature) -> RouteEnableDecision | None: + if facts is None: + return RouteEnableDecision("ReadSettings", feature) + if facts.environment_present and not facts.premium: + return RouteEnableDecision("Deny", feature, "License") + if facts.secret is None: + return RouteEnableDecision("ReadSecret", feature) + if facts.secret: + return RouteEnableDecision("Deny", feature, "Disabled") + return None + + +def route_enable(facts: RouteEnableFacts) -> RouteEnableDecision: + if facts.failed is not None: + return RouteEnableDecision("PropagateHttp" if facts.failed == "Http" else "Allow") + if not facts.imported: + return RouteEnableDecision("ReadImport") + if RouteChecks.check_route_access(facts.context.route, facts.context.catalog.groups.get("management_routes", ())): + admin: Final = _setting(facts.admin, "Admin") + if admin is not None: + return admin + if not RouteChecks.is_llm_api_route(facts.context.route, configured_eu_jev=True): + if facts.context.route.rstrip("/") != "/eu_jev/v1/systemone": + return RouteEnableDecision("Allow") + if facts.eu_jev_configured is None: + return RouteEnableDecision("ReadEuJev") + if facts.eu_jev_configured: + return RouteEnableDecision("Allow") + if facts.exemptions is None: + return RouteEnableDecision("ReadExemptions") + if facts.context.route in facts.exemptions: + return RouteEnableDecision("Allow") + return _setting(facts.llm, "Llm") or RouteEnableDecision("Allow") diff --git a/tests/unit/_auth/support/route_reference.py b/tests/unit/_auth/support/route_reference.py new file mode 100644 index 00000000000..3ce40db3218 --- /dev/null +++ b/tests/unit/_auth/support/route_reference.py @@ -0,0 +1,1233 @@ +import itertools +import re +from collections.abc import Collection, Iterable, Mapping, Sequence +from typing import Final, Protocol, cast + +from fastapi import HTTPException, Request, status +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + MODEL_LIST_ROUTES, + UI_TEAM_ID, + CommonProxyErrors, + KeyManagementRoutes, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, +) + + +def user_is_org_admin(request_data: Mapping[str, object], user_object: LiteLLM_UserTable | None = None) -> bool: + if user_object is None or user_object.organization_memberships is None: + return False + singular: Final = request_data.get("organization_id") + multiple: Final = request_data.get("organizations") + candidates: Final = ( + *((singular,) if singular is not None else ()), + *(TypeAdapter(tuple[object, ...]).validate_python(multiple) if isinstance(multiple, list) else ()), + ) + admins: Final = frozenset( + membership.organization_id + for membership in user_object.organization_memberships + if membership.user_role == LitellmUserRoles.ORG_ADMIN.value and membership.organization_id is not None # pyright: ignore[reportUnnecessaryComparison] # Phase 1 rejected absent legacy membership IDs + ) + return bool(candidates) and all(candidate in admins for candidate in candidates) + + +def managed_agent_route_allowed(route: str, method: str | None) -> bool: + if route in ("/agents", "/v1/agents"): + return method in (None, "GET", "HEAD") + if route in ( + "/realtime", + "/v1/realtime", + "/openai/v1/realtime", + "/realtime/translations", + "/v1/realtime/translations", + "/live/sessions", + "/v1/live/sessions", + "/openai/v1/live/sessions", + ): + return method in (None, "GET") + model_routes: Final = frozenset( + f"{prefix}/{operation}" + for prefix, operation in itertools.product( + ("", "/v1"), + ( + "chat/completions", + "completions", + "embeddings", + "responses", + "messages", + "messages/count_tokens", + "images/generations", + "images/edits", + "audio/transcriptions", + "audio/speech", + "moderations", + "rerank", + "decisions", + "systemone", + "extraction", + "ocr", + ), + ) + ) | frozenset( + ( + "/openai/v1/responses", + "/v2/rerank", + "/claude_code_gateway/v1/messages", + "/claude_code_gateway/v1/messages/count_tokens", + "/cursor/chat/completions", + ) + ) + paths: Final = ( + "/engines/{model:path}/chat/completions", + "/engines/{model:path}/completions", + "/engines/{model:path}/embeddings", + "/openai/deployments/{model:path}/chat/completions", + "/openai/deployments/{model:path}/completions", + "/openai/deployments/{model:path}/embeddings", + "/openai/deployments/{model:path}/images/generations", + "/openai/deployments/{model:path}/images/edits", + "/v1beta/models/{model_name:path}:countTokens", + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/{model_name:path}:streamGenerateContent", + "/models/{model_name:path}:countTokens", + "/models/{model_name:path}:generateContent", + "/models/{model_name:path}:streamGenerateContent", + ) + if route in model_routes or RouteChecks.check_route_access(route, paths): + return method in (None, "POST") + mcp: Final = tuple( + candidate + for candidate in LiteLLMRoutes.mcp_inference_routes.value + if candidate not in ("/token", "/introspect") + ) + return RouteChecks.check_route_access(route, mcp) or RouteChecks.check_route_access( + route, LiteLLMRoutes.agent_inference_routes.value + ) + + +class _ReferenceToken(Protocol): + @property + def allowed_routes(self) -> Sequence[str] | None: ... + + @property + def metadata(self) -> Mapping[str, object]: ... + + @property + def team_metadata(self) -> Mapping[str, object] | None: ... + + +def _token(value: UserAPIKeyAuth) -> _ReferenceToken: + return cast(_ReferenceToken, value) + + +# Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write +# endpoint to a management router REQUIRES adding it here too — the surrounding +# check falls through to "allow" if the route is not matched, which previously +# let view-only admins call /team/block, /team/unblock, /key/bulk_update, etc. +_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( + [ + # user + "/user/new", + "/management/v1/users/bulk", + "/user/delete", + "/management/v1/users/bulk_delete", + "/user/bulk_update", + # team + "/team/new", + "/management/v1/teams/{team_id}/members/bulk_delete", + "/management/v1/teams/{team_id}/members/bulk_update", + "/team/update", + "/team/delete", + "/team/block", + "/team/unblock", + "/team/permissions_update", + "/team/permissions_bulk_update", + # model + "/model/new", + "/model/update", + "/model/delete", + # JWT key mapping + "/jwt/key/mapping/new", + "/jwt/key/mapping/update", + "/jwt/key/mapping/delete", + # key management — keep in sync with KeyManagementRoutes write entries + KeyManagementRoutes.KEY_GENERATE.value, + KeyManagementRoutes.KEY_UPDATE.value, + KeyManagementRoutes.KEY_DELETE.value, + KeyManagementRoutes.KEY_REGENERATE.value, + KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT.value, + KeyManagementRoutes.KEY_BLOCK.value, + KeyManagementRoutes.KEY_UNBLOCK.value, + KeyManagementRoutes.KEY_BULK_UPDATE.value, + KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, + ] +) + +# Suffixes for `/key/{key_id}/...` path-parameterized write routes that the +# enum templates with `{key_id}`. The blocklist above can't match templated +# paths directly because the request route carries the resolved key id. +_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spend") + +_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp")) +_EU_JEV_DECISION_ROUTE: Final = "/eu_jev/v1/systemone" +_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS)) + +# Writes that only touch the caller's own credential row (BYOK connections and +# per-user MCP credentials). Method-scoped because the generic credential CRUD +# handlers share these paths: the route check cannot admit PATCH/DELETE on a +# name like "user_connections" without also opening the admin CRUD handlers. +_SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( + ("GET", "/credentials/user_connections"), + ("POST", "/credentials/{credential_name:path}/user_connection/start"), + ("POST", "/credentials/{credential_name:path}/user_connection/poll"), + ("DELETE", "/credentials/{credential_name:path}/user_connection"), + ("POST", "/v1/mcp/server/oauth/{server_id}/register"), + ("POST", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-env-vars"), + ("DELETE", "/v1/mcp/server/{server_id}/user-env-vars"), +) + + +_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( + { + *MODEL_LIST_ROUTES, + "/utils/token_counter", + "/utils/transform_request", + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", + } +) +_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) +_TOKEN_COUNTING_ACTION: Final = "countTokens" + + +class RouteChecks: + @staticmethod + def check_route_enabled(route: str) -> None: + try: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + EnterpriseRouteChecks.should_call_route(route=route) + except HTTPException as e: + raise e + except Exception: + pass + + @staticmethod + def should_call_route( + route: str, + valid_token: UserAPIKeyAuth, + request: Request | None = None, + ): + """ + Check if management route is disabled and raise exception + """ + RouteChecks.check_route_enabled(route) + # Check if Virtual Key is allowed to call the route - Applies to all Roles + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token, request=request) + return True + + @staticmethod + def is_virtual_key_allowed_to_call_route( + route: str, + valid_token: UserAPIKeyAuth, + request: Request | None = None, + ) -> bool: + """ + Raises Exception if Virtual Key is not allowed to call the route + """ + token: Final = _token(valid_token) + + # Only check if token.allowed_routes is set and is a list with at least one item + if token.allowed_routes is None: + return True + if not isinstance(token.allowed_routes, list): + return True + if len(token.allowed_routes) == 0: + return True + + denied_auth_enforced_pass_through_route = False + + # explicit check for allowed routes (exact match or prefix match) + for allowed_route in token.allowed_routes: + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): + return True + + ## check if 'allowed_route' is a field name in LiteLLMRoutes + if any(allowed_route in LiteLLMRoutes._member_names_ for allowed_route in token.allowed_routes): + for allowed_route in token.allowed_routes: + if allowed_route in LiteLLMRoutes._member_names_: + if RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes[allowed_route].value, + ): + if ( + allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS + and RouteChecks.is_auth_enforced_pass_through_route( + route=route, + method=RouteChecks.get_request_method(request=request), + ) + ): + if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): + return True + denied_auth_enforced_pass_through_route = True + else: + return True + + ################################################ + # For llm_api_routes, also check registered pass-through endpoints + ################################################ + if allowed_route in _LLM_API_ROUTE_GROUPS: + if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET": + return True + + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): + if RouteChecks.is_auth_enforced_pass_through_route( + route=route, + method=RouteChecks.get_request_method(request=request), + ): + if RouteChecks.check_passthrough_route_access( + route=route, user_api_key_dict=valid_token + ): + return True + denied_auth_enforced_pass_through_route = True + else: + return True + + # Method-aware carve-out: allow GET on the two + # read-only MCP-server discovery endpoints + # (`/v1/mcp/server` and `/v1/mcp/server/{server_id}`) + # so virtual keys with allowed_routes=["llm_api_routes"] + # can list/inspect MCP servers. The GET handlers in + # mcp_management_endpoints.py sanitize the response + # for restricted virtual keys (stripping url, + # headers, env, credentials). POST/PUT/DELETE on + # these paths are admin-only management writes and + # are intentionally not covered. + if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route( + route=route, request=request + ): + return True + + # Agent registry CRUD moved from llm_api_routes into + # management_routes so DISABLE_LLM_API_ENDPOINTS stops + # blocking it. Keys configured with + # allowed_routes=["llm_api_routes"] before that split + # could reach these paths, so keep them reachable here; + # the handlers in agent_endpoints/endpoints.py still + # enforce proxy-admin on writes and scope reads by role. + if allowed_route == "llm_api_routes" and RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes.agent_management_routes.value, + ): + return True + + # check if wildcard pattern is allowed + for allowed_route in token.allowed_routes: + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): + return True + + if denied_auth_enforced_pass_through_route: + raise RouteChecks._auth_pass_through_denied_exception(route=route) + + if token.metadata.get("password_reset_required") is True: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + "This account's password must be changed before the session can be used: " + "it was either found in a known data breach or set by an admin. " + "Change it via POST /user/password/change (UI: /ui/change-password), then log in again." + ), + ) + + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {token.allowed_routes}. Tried to call route: {route}", + ) + + @staticmethod + def _mask_user_id(user_id: str) -> str: + """ + Mask user_id to prevent leaking sensitive information in error messages + + Args: + user_id (str): The user_id to mask + + Returns: + str: Masked user_id showing only first 2 and last 2 characters + """ + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + + if not user_id or len(user_id) <= 4: + return "***" + + # Use SensitiveDataMasker with custom configuration for user_id + masker: Final = SensitiveDataMasker(visible_prefix=6, visible_suffix=2, mask_char="*") + + return masker.mask_value(user_id) + + @staticmethod + def _raise_admin_only_route_exception( + user_obj: LiteLLM_UserTable | None, + route: str, + ) -> None: + """ + Raise exception for routes that require proxy admin access + + Args: + user_obj (Optional[LiteLLM_UserTable]): The user object + route (str): The route being accessed + + Raises: + ProxyException: 401 with user role and masked user_id information + """ + user_role = "unknown" + user_id = "unknown" + if user_obj is not None: + user_role = user_obj.user_role or "unknown" + user_id = user_obj.user_id or "unknown" + + masked_user_id: Final = RouteChecks._mask_user_id(user_id) + raise ProxyException( + message=( + "Authentication Error, Only proxy admin can be used to generate, delete, update info for new " + f"keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" + ), + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + + @staticmethod + def non_proxy_admin_allowed_routes_check( + user_obj: LiteLLM_UserTable | None, + _user_role: LitellmUserRoles | None, + route: str, + request: Request, + valid_token: UserAPIKeyAuth, + request_data: Mapping[str, object] | None, + ): + """ + Checks if Non Proxy Admin User is allowed to access the route + """ + token: Final = _token(valid_token) + + # Check user has defined custom admin routes + RouteChecks.custom_admin_only_route_check( + route=route, + ) + + if RouteChecks.is_auth_enforced_pass_through_route( + route=route, + method=RouteChecks.get_request_method(request=request), + ): + RouteChecks._require_auth_pass_through_access( + route=route, + valid_token=valid_token, + jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token), + ) + elif RouteChecks.is_llm_api_route(route=route): + pass + elif RouteChecks.is_info_route(route=route): + # check if user allowed to call an info route + if route == "/key/info": + # handled by function itself + pass + elif route == "/user/info": + # check if user can access this route + query_params: Final = request.query_params + user_id: Final = query_params.get("user_id") + verbose_proxy_logger.debug("user_id: %s & valid_token.user_id: %s", user_id, valid_token.user_id) + if ( + user_id + and user_id != valid_token.user_id + and _user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"key not allowed to access this user's info. user_id={user_id}, key's user_id={valid_token.user_id}", + ) + elif route == "/v2/user/info": + # handled by the endpoint itself (full RBAC in handler) + pass + elif route == "/model/info": + # /model/info just shows models user has access to + pass + elif route == "/team/info": + pass # handled by function itself + elif ( + route in LiteLLMRoutes.global_spend_tracking_routes.value + and getattr(valid_token, "permissions", None) is not None + and "get_spend_routes" in getattr(valid_token, "permissions", []) + ): + pass + elif route == "/lens" or route.startswith("/lens/"): + pass # Lens authorizes the delegated role, including its POST-based read endpoints. + elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: + RouteChecks._check_proxy_admin_viewer_access( + route=route, + _user_role=_user_role, + request_data=request_data, + request=request, + ) + elif ( + _user_role == LitellmUserRoles.INTERNAL_USER.value + and ( + RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value) + or RouteChecks.is_self_scoped_credential_route(route=route, request=request) + ) + or user_is_org_admin(request_data=request_data, user_object=user_obj) # pyright: ignore[reportArgumentType] # preserve the legacy request dictionary contract in this reference + and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value) + or _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + and RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value, + ) + or RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.self_managed_routes.value) + ): + pass + elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"): + pass # authN/authZ handled by api itself + elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token) or ( + valid_token.is_team_service_account + and RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.team_service_account_key_routes.value + ) + ): + pass + elif token.allowed_routes is not None: + # check if route is in allowed_routes (exact match or prefix match) + route_allowed = False + for allowed_route in token.allowed_routes: + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): + route_allowed = True + break + + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): + route_allowed = True + break + + if not route_allowed: + RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) + else: + RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) + + @staticmethod + def custom_admin_only_route_check(route: str): + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # preserve the legacy settings import in this reference + premium_user, + ) + + if "admin_only_routes" in general_settings: + if premium_user is not True: + 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", + ) + + @staticmethod + def is_unbilled_route(route: str) -> bool: + resource, _, action = route.rsplit("/", 1)[-1].partition(":") + return route in _UNBILLED_ROUTES or resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION + + @staticmethod + def is_llm_api_route(route: str, configured_eu_jev: bool | None = None) -> bool: + """ + Helper to checks if provided route is an OpenAI route + + + Returns: + - True: if route is an OpenAI route + - False: if route is not an OpenAI route + """ + # Ensure route is a string before performing checks + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # reference retains the legacy runtime validation + return False + + if route in LiteLLMRoutes.openai_routes.value: + return True + + if route in LiteLLMRoutes.anthropic_routes.value: + return True + + if route in LiteLLMRoutes.google_routes.value: + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): + return True + + # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + # Check for routes with placeholders or wildcard patterns + for openai_route in LiteLLMRoutes.openai_routes.value: + # Replace placeholders with regex pattern + # placeholders are written as "/threads/{thread_id}" + if "{" in openai_route: + if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): + return True + # Check for wildcard patterns like "/containers/*" + if RouteChecks.is_wildcard_pattern(pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): + return True + + # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" + for google_route in LiteLLMRoutes.google_routes.value: + if "{" in google_route: + if RouteChecks._route_matches_pattern(route=route, pattern=google_route): + return True + + # Check for Anthropic routes with placeholders + for anthropic_route in LiteLLMRoutes.anthropic_routes.value: + if "{" in anthropic_route: + if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): + return True + + if RouteChecks._is_azure_openai_route(route=route): + return True + + for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): + return True + return RouteChecks.is_builtin_eu_jev_decision_route(route=route, configured_eu_jev=configured_eu_jev) + + @staticmethod + def is_builtin_eu_jev_decision_route(route: str, configured_eu_jev: bool | None = None) -> bool: + if route.rstrip("/") != _EU_JEV_DECISION_ROUTE: + return False + if configured_eu_jev is not None: + return not configured_eu_jev + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass_through_endpoints imports the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + route=_EU_JEV_DECISION_ROUTE + ) + + @staticmethod + def _is_get_mcp_server_discovery_route(route: str, request: Request | None) -> bool: + """ + Returns True if `request` is a GET against one of the two read-only + MCP-server discovery paths: + + - GET `/v1/mcp/server` (list) + - GET `/v1/mcp/server/{server_id}` (single server, single segment) + + Multi-segment paths (`/v1/mcp/server/{id}/approve`, etc.) and any + non-GET method return False, so admin-only management writes on the + same path prefix are not reachable through this carve-out. + """ + if request is None or request.method.upper() != "GET": + return False + if route == "/v1/mcp/server": + return True + prefix: Final = "/v1/mcp/server/" + if not route.startswith(prefix): + return False + remainder: Final = route[len(prefix) :] + return bool(remainder) and "/" not in remainder + + @staticmethod + def is_management_route(route: str) -> bool: + """ + Check if route is a management route + """ + return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value) + + @staticmethod + def is_info_route(route: str) -> bool: + """ + Check if route is an info route + """ + return route in LiteLLMRoutes.info_routes.value + + @staticmethod + def is_dashboard_session_info_route(route: str, team_id: str | None) -> bool: + """ + Check if an Admin UI session key is calling an info route the dashboard polls in the background + """ + return team_id == UI_TEAM_ID and RouteChecks.is_info_route(route) + + @staticmethod + def _is_azure_openai_route(route: str) -> bool: + """ + Check if route is a route from AzureOpenAI SDK client + + eg. + route='/openai/deployments/vertex_ai/gemini-1.5-flash/chat/completions' + """ + # Ensure route is a string before attempting regex matching + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # reference retains the legacy runtime validation + return False + # Add support for deployment and engine model paths + deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" + engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" + + if re.match(deployment_pattern, route) or re.match(engine_pattern, route): + return True + return False + + @staticmethod + def _route_matches_pattern(route: str, pattern: str) -> bool: + """ + Check if route matches the pattern placed in proxy/_types.py + + Example: + - pattern: "/threads/{thread_id}" + - route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + - returns: True + + + - pattern: "/key/{token_id}/regenerate" + - route: "/key/regenerate/82akk800000000jjsk" + - returns: False, pattern is "/key/{token_id}/regenerate" + """ + # Ensure route is a string before attempting regex matching + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # reference retains the legacy runtime validation + return False + + def _placeholder_to_regex(match: re.Match[str]) -> str: + placeholder: Final = match.group(0).strip("{}") + if not placeholder.endswith(":path"): + return r"[^/]+" + # A ":path" placeholder takes whatever the router's own path + # converter takes, slashes and colons alike, so an id spelled with + # either (or both) still matches the template it was mounted under. + # + # Unless the template puts a ":" literal of its own after the + # placeholder: the Google routes end in ":generateContent" and + # friends, and there the value has to stop before that suffix + # rather than swallow it and match a different verb. + # + # "[\s\S]" rather than ".", because "." stops at a newline and the + # path converter does not: a %0A anywhere in the value would leave + # the route unmatched here while still reaching the handler, which + # turns this gate into a bypass for the lists built on it. + return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" + + pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) + # Anchor the pattern to match the entire string + pattern = f"^{pattern}$" + if re.match(pattern, route): + return True + return False + + @staticmethod + def is_wildcard_pattern(pattern: str) -> bool: + """ + Check if pattern is a wildcard pattern + """ + return pattern.endswith("*") + + _is_wildcard_pattern = is_wildcard_pattern + + @staticmethod + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + """ + Check if route matches the wildcard pattern + + eg. + + pattern: "/scim/v2/*" + route: "/scim/v2/Users" + - returns: True + + pattern: "/scim/v2/*" + route: "/chat/completions" + - returns: False + + + pattern: "/scim/v2/*" + route: "/scim/v2/Users/123" + - returns: True + + """ + if pattern.endswith("*"): + # Get the prefix (everything before the wildcard) + prefix: Final = pattern[:-1] + return route.startswith(prefix) + else: + # If there's no wildcard, the pattern and route should match exactly + return route == pattern + + @staticmethod + def _route_matches_allowed_route(route: str, allowed_route: str) -> bool: + """ + Check if route matches the allowed_route pattern. + Supports both exact match and prefix match. + + Examples: + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6" -> True (exact match) + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6/v1/chat/completions" -> True (prefix match) + - allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-600" -> False (not a valid prefix) + + Args: + route: The actual route being accessed + allowed_route: The allowed route pattern + + Returns: + bool: True if route matches (exact or prefix), False otherwise + """ + # Exact match + if route == allowed_route: + return True + # Prefix match - ensure we add "/" to prevent false matches like /fake-openai-proxy-600 + if route.startswith(allowed_route + "/"): + return True + return False + + @staticmethod + def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: + """ + Check if a route has access by checking both exact matches and patterns + + Args: + route (str): The route to check + allowed_routes (Sequence): Allowed routes/patterns + + Returns: + bool: True if route is allowed, False otherwise + """ + ######################################################### + # exact match route is in allowed_routes + ######################################################### + if route in allowed_routes: + return True + + ######################################################### + # wildcard match route is in allowed_routes + # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* + ######################################################### + if any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in allowed_routes + if RouteChecks.is_wildcard_pattern(pattern=allowed_route) + ): + return True + + ######################################################### + # pattern match route is in allowed_routes + # pattern: "/threads/{thread_id}" + # route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + # returns: True + ######################################################### + if any( # Check pattern match + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes + ): + return True + + return False + + @staticmethod + def get_request_method(request: Request | None) -> str | None: + if request is None: + return None + + try: + method: Final = cast( # cast-ok: request.method is str at runtime; tests hand a MagicMock + object, request.method + ) + except (AttributeError, KeyError): + return None + if not isinstance(method, str): + return None + + return method.upper() + + _get_request_method = get_request_method + + @staticmethod + def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: + method: Final = RouteChecks.get_request_method(request) + if method is None: + return False + return any( + method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,)) + for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES + ) + + @staticmethod + def is_auth_enforced_pass_through_route(route: str, method: str | None = None) -> bool: + """ + True for config/DB pass-through endpoints registered with auth=true. + + These routes are injected into ``openai_routes`` for spend/budget hooks but + must not inherit blanket ``openai_routes`` RBAC; access is gated by + ``allowed_passthrough_routes`` on the key or team. + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + route_info: Final = InitPassThroughEndpointHelpers.get_registered_pass_through_route(route=route, method=method) + if route_info is None: + return False + return route_info.get("auth") is True + + @staticmethod + def _auth_pass_through_denied_exception(route: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"Key/team not allowed to access passthrough route {route}. " + "Configure `allowed_passthrough_routes` on the team or key." + ), + ) + + @staticmethod + def _route_matches_denied_route(route: str, denied_route: str) -> bool: + """A `/` entry denies every route, since every route sits under the root.""" + normalized_denied_route: Final = denied_route.rstrip("/") or "/" + return ( + normalized_denied_route == "/" + or RouteChecks._route_matches_allowed_route(route=route, allowed_route=normalized_denied_route) + or RouteChecks.route_matches_wildcard_pattern(route=route, pattern=denied_route) + ) + + @staticmethod + def matching_denied_passthrough_route( + route: str, metadata_sources: Iterable[Mapping[str, object] | None] + ) -> str | None: + """ + First ``denied_passthrough_routes`` entry across ``metadata_sources`` that matches ``route``. + Unlike the allowlist (key list, else team list), every source's deny list applies. + """ + denied_routes: Final = tuple( + itertools.chain.from_iterable( + cast( # cast-ok: management endpoints validate this metadata key as a list of route strings on write + "list[str]", (metadata or {}).get("denied_passthrough_routes") or [] + ) + for metadata in metadata_sources + ) + ) + if not denied_routes: + return None + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + forwarded_routes: Final = InitPassThroughEndpointHelpers.forwarded_routes(route) + return next( + ( + denied_route + for denied_route in denied_routes + if any( + RouteChecks._route_matches_denied_route(route=candidate, denied_route=denied_route) + for candidate in forwarded_routes + ) + ), + None, + ) + + @staticmethod + def passthrough_route_denied_exception(route: str, denied_route: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"Key/team denied access to passthrough route {route}. " + f"Matched `{denied_route}` in `denied_passthrough_routes`." + ), + ) + + @staticmethod + def _raise_if_passthrough_route_denied(route: str, valid_token: UserAPIKeyAuth) -> None: + token: Final = _token(valid_token) + denied_route: Final = RouteChecks.matching_denied_passthrough_route( + route=route, + metadata_sources=(token.metadata, token.team_metadata), + ) + if denied_route is not None: + raise RouteChecks.passthrough_route_denied_exception(route=route, denied_route=denied_route) + + @staticmethod + def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool: + """ + Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do: + a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names + no path segment (``*``, ``/*``) is skipped. + """ + return any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in team_allowed_routes + if allowed_route.rstrip("*").strip("/") + ) + + @staticmethod + def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]: + """``team_allowed_routes`` for team tokens built by JWT auth; JWT-mapped virtual keys stay key-scoped.""" + if valid_token.jwt_claims is None or valid_token.token is not None or valid_token.team_id is None: + return () + + from litellm.proxy.proxy_server import jwt_handler + + return jwt_handler.litellm_jwtauth.team_allowed_routes + + @staticmethod + def _require_auth_pass_through_access( + route: str, + valid_token: UserAPIKeyAuth, + jwt_team_allowed_routes: Collection[str] = (), + ) -> None: + """ + Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the + key or team, or an explicit JWT ``team_allowed_routes`` entry. A key or team + ``denied_passthrough_routes`` match blocks the route even when one of those grants it. + """ + RouteChecks._raise_if_passthrough_route_denied(route=route, valid_token=valid_token) + if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): + return + if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes): + return + raise RouteChecks._auth_pass_through_denied_exception(route=route) + + @staticmethod + def check_passthrough_route_access(route: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Check if route is a passthrough route. + Supports both exact match and prefix match. + """ + token: Final = _token(user_api_key_dict) + metadata: Final = token.metadata + team_metadata: Final = token.team_metadata or {} + if metadata is None and team_metadata is None: # pyright: ignore[reportUnnecessaryComparison] # reference retains the legacy runtime validation + return False + if "allowed_passthrough_routes" not in metadata and "allowed_passthrough_routes" not in team_metadata: + return False + if ( + metadata.get("allowed_passthrough_routes") is None + and team_metadata.get("allowed_passthrough_routes") is None + ): + return False + + allowed_passthrough_routes: Final = cast( + Sequence[str], + metadata.get("allowed_passthrough_routes") or team_metadata.get("allowed_passthrough_routes") or [], + ) + + # Check if route matches any allowed passthrough route (exact or prefix match) + for allowed_route in allowed_passthrough_routes: + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): + return True + + return False + + @staticmethod + def is_assistants_api_request(request: Request) -> bool: + """ + Returns True if `thread` or `assistant` is in the request path + + Args: + request (Request): The request object + + Returns: + bool: True if `thread` or `assistant` is in the request path, False otherwise + """ + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route + + route: Final = get_request_route(request) + if "thread" in route or "assistant" in route: + return True + return False + + _is_assistants_api_request = is_assistants_api_request + + @staticmethod + def is_generate_content_route(route: str) -> bool: + """ + Returns True if this is a google generateContent or streamGenerateContent route + + These routes from google allow passing key=api_key in the query params + """ + if "generateContent" in route: + return True + if "streamGenerateContent" in route: + return True + return False + + # HTTP methods that are intrinsically read-only and therefore safe to + # default-allow for PROXY_ADMIN_VIEW_ONLY. Anything else (POST/PUT/PATCH/ + # DELETE) is treated as a write attempt and goes through the explicit + # write-allowlist below. + _SAFE_HTTP_METHODS = frozenset({"GET", "HEAD", "OPTIONS"}) + + # Explicit write routes that PROXY_ADMIN_VIEW_ONLY must NEVER call. The + # role-principle is "no writes, ever" — the management_routes list is the + # authoritative source for which non-llm routes are writes; we just need + # to filter out the read endpoints (info / list) that share the prefix. + # A cleaner approach is to denylist by HTTP verb (POST/PUT/PATCH/DELETE); + # this block stays as a backstop in case a write is implemented as GET. + _ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset( + [ + "/user/new", + "/management/v1/users/bulk", + "/user/delete", + "/management/v1/users/bulk_delete", + "/user/bulk_update", + "/team/new", + "/management/v1/teams/{team_id}/members/bulk_delete", + "/management/v1/teams/{team_id}/members/bulk_update", + "/team/update", + "/team/delete", + "/model/new", + "/model/update", + "/model/delete", + "/key/generate", + "/key/delete", + "/key/update", + "/key/regenerate", + "/key/service-account/generate", + "/key/block", + "/key/unblock", + "/team/key/bulk_update", + ] + ) + + @staticmethod + def _check_proxy_admin_viewer_access( + route: str, + _user_role: str, + request_data: Mapping[str, object] | None, + request: Request | None = None, + ) -> None: + """ + Check access for PROXY_ADMIN_VIEW_ONLY role. + + Admin Viewer follows a read-parity-with-Proxy-Admin rule: anything Proxy + Admin can read/list/get, Admin Viewer can read/list/get. The only + exclusions are cost-incurring inference routes (Playground, /chat/ + completions, etc.) and any state-mutating request. + + Implementation: + 1. LLM/inference routes → 403 (cost-incurring). + 2. Safe HTTP method (GET/HEAD/OPTIONS) → allow by default. This is + the read-parity guarantee — every new GET endpoint added anywhere + in the codebase is automatically readable by Admin Viewer + without needing to remember to add it to an allowlist. + 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): + - Allow `/user/update` only when restricted to user_email. + - Allow `/user/password/change` (endpoint only writes the caller's own row). + - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. + - Otherwise allow only if the route is in admin_viewer_routes / + global_spend_tracking_routes (legacy explicit-allow set). + - Else 403. + """ + if RouteChecks.is_llm_api_route(route=route): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", + ) + + # Check if this is a write operation on management routes + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value): + # For management routes, only allow read operations or specific allowed updates + if route == "/user/update": + # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY + if request_data is not None and isinstance(request_data, dict): + _params_updated: Final = request_data.keys() + for param in _params_updated: + if param != "user_email": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated", + ) + elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or ( + route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) + ): + # Block write operations for PROXY_ADMIN_VIEW_ONLY + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) + # Allow read operations on management routes (like /user/info, /team/info, /model/info) + method: Final = request.method.upper() if request is not None else "GET" + is_safe_method: Final = method in RouteChecks._SAFE_HTTP_METHODS + + # ── Safe HTTP method: default-allow ────────────────────────────── + if is_safe_method: + return + + # ── Unsafe HTTP method: explicit checks ────────────────────────── + # Allow `/user/update` for self-service email change. + if route == "/user/update": + if request_data is not None and isinstance(request_data, dict): + for param in request_data: + if param != "user_email": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"user not allowed to access this route, role= {_user_role}. " + f"Trying to access: {route} and updating invalid param: {param}. " + "only user_email can be updated" + ), + ) + return + + # Self-service password change; the endpoint only writes the caller's own row. + if route == "/user/password/change": + return + + # Self-service logout; the endpoint only revokes the caller's own session key. + if route == "/session/logout": + return + + if RouteChecks.is_self_scoped_credential_route(route=route, request=request): + return + + # Hard-block known write routes regardless of HTTP method (defensive + # — these are POSTs in practice, but pinning them here protects + # against future GET-shaped writes). + if RouteChecks.check_route_access( + route=route, allowed_routes=RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES + ) or (route.startswith("/key/") and route.endswith("/regenerate")): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) + + # Legacy explicit-allow sets (kept for routes that are POST but + # semantically read-only, e.g. /spend/calculate). Both admin_viewer_routes + # and global_spend_tracking_routes are reads/listings. + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value): + return + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value): + return + + # NOTE: We intentionally do NOT fall back to allowing all + # `management_routes`. That set is a mix of reads (info/list — handled + # via the safe-method branch above) and writes (`/team/block`, + # `/team/permissions_update`, `/jwt/key/mapping/{new,update,delete}`, + # `/key/bulk_update`, `/key/{id}/reset_spend`). A blanket allow would + # let Admin Viewer POST these write endpoints — violating the + # "no writes, ever" rule. Default-deny instead. + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) diff --git a/tests/unit/_auth/test_access.py b/tests/unit/_auth/test_access.py new file mode 100644 index 00000000000..17d161dec1e --- /dev/null +++ b/tests/unit/_auth/test_access.py @@ -0,0 +1,110 @@ +from collections.abc import Callable +from itertools import product +from typing import Final, cast + +import pytest + +from litellm._auth import access +from tests.unit._auth.reference_policy import access_decision + + +@pytest.fixture(params=("reference", "native")) +def decide(request: pytest.FixtureRequest) -> Callable[[access.AccessFacts], access.AccessDecision]: + return access_decision if cast(str, request.param) == "reference" else access.decide + + +@pytest.mark.parametrize( + "facts", + tuple( + access.UserParam(*values) + for values in product((False, True), (None, "POST", "post", "GET"), (False, True), (False, True), (False, True)) + ), +) +def test_user_param_gate_preserves_method_route_and_presence( + decide: Callable[[access.AccessFacts], access.AccessDecision], facts: access.UserParam +) -> None: + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize("facts", tuple(access.GrantLayers(*values) for values in product((False, True), repeat=7))) +def test_default_deny_layers_preserve_master_dashboard_and_session_exceptions( + decide: Callable[[access.AccessFacts], access.AccessDecision], facts: access.GrantLayers +) -> None: + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize( + "key,team,expected", + ((None, None, ()), ((), ("a",), ("b", "b")), (("b",), ("a",), ("a",)), (("a", "b"), (), ()), ((), (), ())), +) +def test_key_tools_override_team_and_keep_duplicate_denials( + decide: Callable[[access.AccessFacts], access.AccessDecision], + key: tuple[str, ...] | None, + team: tuple[str, ...] | None, + expected: tuple[str, ...], +) -> None: + assert decide(access.Tools(("a", "b", "b"), key, team)) == access.AccessDecision( + denial="Tools" if expected else None, rejected=expected + ) + + +@pytest.mark.parametrize( + "strict,allowed,expected", + ( + (False, None, None), + (False, (), None), + (True, None, "MissingVectorGrant"), + (True, (), "MissingVectorGrant"), + (True, ("a",), "VectorStores"), + (False, ("a",), "VectorStores"), + ), +) +def test_vector_grants_distinguish_missing_from_denied_and_keep_first_index( + decide: Callable[[access.AccessFacts], access.AccessDecision], + strict: bool, + allowed: tuple[str, ...] | None, + expected: access.AccessDenial | None, +) -> None: + assert decide(access.VectorStores(("a", "b", "a"), allowed, strict)) == access.AccessDecision( + denial=expected, rejected_index=1 if expected == "VectorStores" else None + ) + + +@pytest.mark.parametrize( + "facts", + tuple( + access.TraceStep(source, value, present, headers_only) + for source, value, present, headers_only in product( + ("Header", "Body", "Traceparent", "Selected", "Other"), + (None, "", " ", "trace"), + (False, True), + (False, True), + ) + ), +) +def test_trace_source_precedence_and_read_short_circuits( + decide: Callable[[access.AccessFacts], access.AccessDecision], facts: access.TraceStep +) -> None: + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize( + "key", + ("guardrails", "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", "unrelated"), +) +@pytest.mark.parametrize("modify", (None, False, True)) +def test_guardrail_permission_requires_literal_false( + decide: Callable[[access.AccessFacts], access.AccessDecision], key: str, modify: bool | None +) -> None: + facts: Final = access.Guardrails((key,), modify) + assert decide(facts) == access_decision(facts) + + +@pytest.mark.parametrize( + "enabled,route,read", + ((False, None, None), (True, None, "ClientTagRoute"), (True, False, None), (True, True, "ClientTags")), +) +def test_client_tags_only_read_route_and_request_state_when_enabled( + decide: Callable[[access.AccessFacts], access.AccessDecision], enabled: bool, route: bool | None, read: str | None +) -> None: + assert decide(access.ClientTagRequest(enabled, route)).read == read diff --git a/tests/unit/_auth/test_admission.py b/tests/unit/_auth/test_admission.py new file mode 100644 index 00000000000..3a7c77dd8db --- /dev/null +++ b/tests/unit/_auth/test_admission.py @@ -0,0 +1,21 @@ +from itertools import product + +import pytest + +from litellm._auth import admission +from tests.unit._auth import reference_policy + + +@pytest.mark.parametrize("facts", tuple(admission.CommonGate(*values) for values in product((False, True), repeat=10))) +def test_common_gate_has_the_reference_skip_priority(facts: admission.CommonGate) -> None: + assert admission.common_plan(facts) == reference_policy.common_plan(facts) + + +@pytest.mark.parametrize("facts", tuple(admission.Reservation(*values) for values in product((False, True), repeat=5))) +def test_reservation_plan_preserves_billing_agent_fail_closed(facts: admission.Reservation) -> None: + assert admission.reservation_plan(facts) == reference_policy.reservation_plan(facts) + + +@pytest.mark.parametrize("facts", tuple(admission.Ownership(*values) for values in product((False, True), repeat=4))) +def test_observation_never_transfers_request_ownership(facts: admission.Ownership) -> None: + assert admission.ownership_plan(facts).owner == "Python" diff --git a/tests/unit/_auth/test_agent_ceilings.py b/tests/unit/_auth/test_agent_ceilings.py new file mode 100644 index 00000000000..a32fb1f6488 --- /dev/null +++ b/tests/unit/_auth/test_agent_ceilings.py @@ -0,0 +1,91 @@ +from collections.abc import Callable +from itertools import chain +from typing import Final + +import pytest + +from litellm._auth.agent_ceilings import AgentCeilingDecision as Decision +from litellm._auth.agent_ceilings import AgentCeilingFacts as Facts +from litellm._auth.agent_ceilings import AgentCeilingGroup as Group +from litellm._auth.agent_ceilings import AgentGroupLoadDecision as LoadDecision +from litellm._auth.agent_ceilings import AgentGroupLoadFacts as LoadFacts +from litellm._auth.agent_ceilings import agent_ceiling, agent_group_load +from litellm._auth.failure import NativeAuthFailure + + +def reference(facts: Facts) -> Decision: + if facts.ids is None: + return Decision("ReadIds") + if not facts.ids: + return Decision("Complete") + if facts.groups is None: + return Decision("ReadGroups", authoritative=facts.check_db_only) + indices: Final = tuple(index for index, group in enumerate(facts.groups) if group.status == "Found") + return Decision( + "Complete", + present=True, + ids=tuple(name for name, group in zip(facts.ids, facts.groups, strict=True) if group.status != "Missing"), + models=frozenset(chain.from_iterable(facts.groups[index].models for index in indices)), + mcp_server_ids=frozenset(chain.from_iterable(facts.groups[index].mcp_server_ids for index in indices)), + agent_ids=frozenset(chain.from_iterable(facts.groups[index].agent_ids for index in indices)), + indices=indices, + ) + + +def load_reference(facts: LoadFacts) -> LoadDecision: + if not facts.database_present: + return LoadDecision("DatabaseUnavailable") + if facts.status is None: + return LoadDecision("Read") + if facts.status == "Success": + return LoadDecision("Return") + return LoadDecision("Raise" if facts.status == "NotFound" or facts.check_db_only else "HttpUnavailable") + + +@pytest.mark.parametrize("policy", (reference, agent_ceiling)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (Facts(), Decision("ReadIds")), + (Facts(()), Decision("Complete")), + (Facts(("group",)), Decision("ReadGroups")), + (Facts(("group",), check_db_only=True), Decision("ReadGroups", authoritative=True)), + (Facts(("group",), (Group("Missing"),)), Decision("Complete", present=True)), + (Facts(("group",), (Group("Unavailable"),)), Decision("Complete", present=True, ids=("group",))), + ( + Facts( + ("one", "two", "one"), + (Group("Found", ("model",), ("server",), ("agent",)), Group("Missing"), Group("Found", ("model",))), + ), + Decision( + "Complete", + present=True, + ids=("one", "one"), + models=frozenset(("model",)), + mcp_server_ids=frozenset(("server",)), + agent_ids=frozenset(("agent",)), + indices=(0, 2), + ), + ), + ), +) +def test_ceiling_union_and_missing_vs_unavailable( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("policy", (load_reference, agent_group_load)) +@pytest.mark.parametrize("database", (False, True)) +@pytest.mark.parametrize("strict", (False, True)) +@pytest.mark.parametrize("status", (None, "Success", "NotFound", "HttpError")) +def test_group_read_error_policy( + policy: Callable[[LoadFacts], LoadDecision], database: bool, strict: bool, status: str | None +) -> None: + facts: Final = LoadFacts(database, strict, status) # pyright: ignore[reportArgumentType] # parameter table covers the complete typed status domain + assert policy(facts) == load_reference(facts) + + +def test_mismatched_group_answers_are_terminal() -> None: + with pytest.raises(NativeAuthFailure): + agent_ceiling(Facts(("group",), ())) diff --git a/tests/unit/_auth/test_agent_grants.py b/tests/unit/_auth/test_agent_grants.py new file mode 100644 index 00000000000..946d1a351e3 --- /dev/null +++ b/tests/unit/_auth/test_agent_grants.py @@ -0,0 +1,193 @@ +from dataclasses import replace +from itertools import product +from typing import Final, assert_never + +import pytest + +from litellm._auth import agent_grants as policy + + +def _decision( + action: policy.Action = "Complete", grants: tuple[str, ...] | None = None, team_id: str | None = None +) -> policy.Decision: + return policy.Decision(action, grants, team_id) + + +def reference(facts: policy.Facts) -> policy.Decision: + match facts: + case policy.Resolve(): + if facts.managed: + return _decision("ReadManaged") + if not facts.key_team_loaded: + return _decision("ReadKeyTeam") + if not facts.strict and facts.role != "proxy_admin" and facts.default_deny is None: + return _decision("ReadDefault") + if (facts.strict or (facts.role != "proxy_admin" and facts.default_deny)) and facts.key_team is None: + return _decision(grants=()) + if not facts.caller_loaded: + return _decision("ReadCaller") + if not facts.ceiling_loaded: + return _decision("ReadCeiling") + return _decision( + grants=facts.ceiling + if facts.own is None + else facts.own + if facts.ceiling is None + else tuple(item for item in facts.own if item in facts.ceiling) + ) + case policy.Failure(): + return _decision("Unavailable") if facts.strict else _decision(grants=() if facts.closed else None) + case policy.OwnerTeam(): + if facts.permission_id_present and not facts.permission_present: + return _decision("Unavailable") if facts.strict else _decision(grants=()) + return _decision("ReadDeclared") + case policy.OwnerUnion(): + if any(source.unavailable for source in facts.teams): + return _decision("Unavailable") + restrictions: Final = ( + *(source.grants for source in facts.teams), + *(() for source in facts.withheld if source.unavailable or source.grants is not None), + *(((),) if facts.unreadable else ()), + ) + return _decision( + grants=None + if restrictions and all(grants is None for grants in restrictions) + else tuple(sorted({item for grants in restrictions for item in (grants or ())})) + ) + case policy.BillingTeam(): + restricted: Final = tuple( + source for source in facts.sources if source.grants is not None and facts.agent_id in source.grants + ) + selected: Final = restricted or ( + facts.sources + if all(source.grants is None and not source.unavailable for source in facts.sources) + else () + ) + return _decision( + team_id=min((source.team_id for source in selected if source.team_id is not None), default=None) + ) + case policy.Human(): + return _decision( + grants=tuple( + sorted( + { + item + for source in facts.sources + if source.team_id is None or source.team_id == facts.team_id + for item in (source.grants or ()) + } + ) + ) + ) + case policy.Delegated(): + if any(source.team_id is None and facts.agent_id in (source.grants or ()) for source in facts.sources): + return _decision(team_id=facts.team_id) + teams: Final = frozenset( + source.team_id + for source in facts.sources + if source.team_id is not None + and facts.agent_id in (source.grants or ()) + and (facts.allowed is None or source.team_id in facts.allowed) + ) + if facts.team_id in teams: + return _decision(team_id=facts.team_id) + if not facts.explicit and teams: + return _decision(team_id=min(teams)) + return _decision("SelectTeam") + case policy.TeamLookup(): + if not facts.subject: + return _decision() + if facts.team_id is None: + return _decision( + "ReadPersonalEligibility" + if facts.personal is None + else "ReadOwnerTeams" + if facts.personal + else "Complete" + ) + if not facts.database or (facts.loaded and not facts.row_present): + return _decision(grants=() if facts.strict else None) + return _decision("ReadDeclared" if facts.loaded else "ReadTeam") + case policy.GroupMatches(): + return _decision( + grants=tuple( + agent.agent_id for agent in facts.agents if any(group in agent.groups for group in facts.groups) + ) + ) + case _: + assert_never(facts) + + +_SOURCES: Final = ( + (), + (policy.Source("b", None),), + (policy.Source("b", ("agent",)), policy.Source("a", ("agent", "other"))), + (policy.Source("b", None, True),), + (policy.Source(None, ("agent",)),), + (policy.Source("b", ()),), +) +_RESOLVE: Final = tuple( + policy.Resolve( + managed, + strict, + role, + loaded, + key, + default, + caller_loaded=caller, + own=("agent", "other"), + ceiling_loaded=ceiling_loaded, + ceiling=ceiling, + ) + for managed, strict, role, loaded, key, default, caller, ceiling_loaded, ceiling in product( + (False, True), + (False, True), + (None, "proxy_admin", "internal_user"), + (False, True), + (None, (), ("agent",)), + (None, False, True), + (False, True), + (False, True), + (None, (), ("agent",)), + ) +) +_CASES: Final[tuple[policy.Facts, ...]] = ( + *_RESOLVE, + *( + policy.OwnerUnion(teams, withheld, unreadable) + for teams, withheld, unreadable in product(_SOURCES, _SOURCES, (False, True)) + ), + *(policy.Failure(strict, closed) for strict, closed in product((False, True), repeat=2)), + *(policy.OwnerTeam(strict, declared, loaded) for strict, declared, loaded in product((False, True), repeat=3)), + *(policy.BillingTeam(sources, "agent") for sources in _SOURCES), + *(policy.Human(sources, team) for sources, team in product(_SOURCES, (None, "a", "b"))), + *( + policy.Delegated(sources, "agent", team, explicit, allowed) + for sources, team, explicit, allowed in product(_SOURCES, (None, "a", "b"), (False, True), (None, (), ("b",))) + ), + *( + policy.TeamLookup(subject, team, database, strict, personal, loaded, row) + for subject, team, database, strict, personal, loaded, row in product( + (False, True), + (None, "", "team"), + (False, True), + (False, True), + (None, False, True), + (False, True), + (False, True), + ) + ), + policy.GroupMatches((policy.ConfiguredAgent("agent", ("group",)), policy.ConfiguredAgent("other", ())), ("group",)), +) + + +@pytest.mark.parametrize("facts", _CASES) +def test_agent_grants_match_reference(facts: policy.Facts) -> None: + assert policy.decide(facts) == reference(facts) + + +def test_default_deny_precedes_caller_and_ceiling_reads() -> None: + facts: Final = policy.Resolve(False, False, "internal_user", key_team_loaded=True) + assert policy.decide(facts).action == "ReadDefault" + assert policy.decide(replace(facts, default_deny=True)) == _decision(grants=()) + assert policy.decide(replace(facts, default_deny=False)).action == "ReadCaller" diff --git a/tests/unit/_auth/test_auth_context.py b/tests/unit/_auth/test_auth_context.py new file mode 100644 index 00000000000..8585e627b18 --- /dev/null +++ b/tests/unit/_auth/test_auth_context.py @@ -0,0 +1,64 @@ +from itertools import product +from typing import Final + +import pytest + +from litellm._auth import auth_context + +_FETCH: Final = tuple( + auth_context.Fetch(team, "ui", user, project, end_user) + for team, user, project, end_user in product( + (None, "ui", "", "team"), (None, "", "user"), (None, "", "project"), (None, "", "customer") + ) +) +_TEAM: Final = tuple( + auth_context.Team(team, "ui", failed, missing, models, allow) + for team, failed, missing, models, allow in product( + (None, "ui", "team"), (False, True), (False, True), ((), ("model",)), (None, False, True) + ) +) +_ADMIN: Final = tuple( + auth_context.Admin(role, user, "admin", row) + for role, user, row in product((None, "proxy_admin", "internal_user"), (None, "admin", "user"), (False, True)) +) + + +@pytest.mark.parametrize("facts", _FETCH) +def test_context_fetches_preserve_missing_and_empty_ids(facts: auth_context.Fetch) -> None: + assert auth_context.fetch(facts) == auth_context.FetchPlan( + facts.team_id is not None and facts.team_id != facts.ui_team_id, + facts.user_id is not None, + facts.project_id is not None, + bool(facts.end_user_id), + ) + + +@pytest.mark.parametrize("facts", _TEAM) +def test_team_context_preserves_fallback_rules(facts: auth_context.Team) -> None: + expected: Final = ( + ("Token" if facts.team_id == facts.ui_team_id else "Loaded") + if not facts.failed + else ( + "None" + if facts.team_id is None + else "Token" + if facts.team_id == facts.ui_team_id + else "Failure" + if facts.missing + else "Token" + if facts.recorded_models + else "ReadUnavailablePolicy" + if facts.allow_unavailable is None + else "Token" + if facts.allow_unavailable + else "Failure" + ) + ) + assert auth_context.team(facts).source == expected + + +@pytest.mark.parametrize("facts", _ADMIN) +def test_admin_projection_preserves_personal_budgets(facts: auth_context.Admin) -> None: + assert auth_context.admin(facts) == auth_context.AdminPlan( + facts.role == "proxy_admin", facts.user_row and facts.user_id != facts.admin_id + ) diff --git a/tests/unit/_auth/test_budget_pricing.py b/tests/unit/_auth/test_budget_pricing.py new file mode 100644 index 00000000000..b193bfe90c5 --- /dev/null +++ b/tests/unit/_auth/test_budget_pricing.py @@ -0,0 +1,78 @@ +from math import inf, nan +from typing import Final + +import pytest + +from litellm._auth import budget_pricing as native +from litellm._auth.budget_pricing import DeploymentPrice, ModelPriceFacts, PriceDecision, RouteBudgetFacts + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (ModelPriceFacts(cached=True, lookup_failed=True), PriceDecision(True, None)), + (ModelPriceFacts(cached=False), PriceDecision(False, None)), + (ModelPriceFacts(lookup_failed=True), PriceDecision(False, None)), + (ModelPriceFacts(), PriceDecision(False, False)), + (ModelPriceFacts(group_found=True), PriceDecision(False, False)), + (ModelPriceFacts(group_found=True, input_price=0, output_price=0), PriceDecision(False, None, True)), + (ModelPriceFacts(group_found=True, input_price=0, output_price=0, deployments=()), PriceDecision(False, False)), + ), +) +def test_pricing_fail_closed_and_read_plan(facts: ModelPriceFacts, expected: PriceDecision) -> None: + assert native.price(facts) == expected + + +@pytest.mark.parametrize( + ("price", "expected"), ((None, False), (-1, True), (0, True), (1, False), (inf, False), (nan, True)) +) +def test_zero_price_keeps_reference_nonfinite_semantics(price: float | None, expected: bool) -> None: + facts: Final = ModelPriceFacts( + group_found=True, input_price=price, output_price=0, deployments=(DeploymentPrice(has_input_price=True),) + ) + assert native.price(facts) == PriceDecision(expected, expected) + + +@pytest.mark.parametrize( + ("deployment", "expected"), + ( + (DeploymentPrice(), False), + (DeploymentPrice(has_input_price=True), True), + (DeploymentPrice(has_output_price=True), True), + (DeploymentPrice(has_input_price=True, input_price=1), False), + (DeploymentPrice(has_input_price=True, output_price=1), False), + (DeploymentPrice(has_input_price=True, query_price=1), False), + (DeploymentPrice(has_input_price=True, ptu_count=0, cost_per_ptu_per_hour=0), False), + (DeploymentPrice(has_input_price=True, ptu_count=0), True), + (DeploymentPrice(has_input_price=True, input_price=nan), True), + ), +) +def test_explicit_deployment_cost_and_ptu_override_zero_group_price( + deployment: DeploymentPrice, expected: bool +) -> None: + facts: Final = ModelPriceFacts(group_found=True, input_price=0, output_price=0, deployments=(deployment,)) + assert native.price(facts) == PriceDecision(expected, expected) + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (RouteBudgetFacts(side_effect_route=True, discovery_route=True), "ReadPricing"), + (RouteBudgetFacts(discovery_route=True, llm_route=True), "Skip"), + (RouteBudgetFacts(llm_route=True), "ReadPricing"), + (RouteBudgetFacts(), "ReadPassThrough"), + (RouteBudgetFacts(registered_pass_through=False), "Skip"), + (RouteBudgetFacts(registered_pass_through=True), "ReadPricing"), + (RouteBudgetFacts(llm_route=True, route="/mcp", method="tools/list", body_present=True), "Skip"), + (RouteBudgetFacts(llm_route=True, route="/mcp", method="tools/list", body_present=False), "ReadPricing"), + ( + RouteBudgetFacts(llm_route=True, route="/mcp/tools/call", method="tools/list", body_present=True), + "ReadPricing", + ), + (RouteBudgetFacts(llm_route=True, route="/mcp-rest/tools/list", body_present=True), "Skip"), + (RouteBudgetFacts(llm_route=True, route="/mcp/team", method="initialize", body_present=True), "Skip"), + (RouteBudgetFacts(llm_route=True, route="/mcp", method="tools/call", body_present=True), "ReadPricing"), + ), +) +def test_route_and_mcp_exemption_policy(facts: RouteBudgetFacts, expected: str) -> None: + assert native.route(facts) == expected diff --git a/tests/unit/_auth/test_common_budgets.py b/tests/unit/_auth/test_common_budgets.py new file mode 100644 index 00000000000..3d3efb4fdcd --- /dev/null +++ b/tests/unit/_auth/test_common_budgets.py @@ -0,0 +1,134 @@ +from collections.abc import Callable +from dataclasses import replace +from math import inf, nan +from typing import Final, get_args + +import pytest + +from litellm._auth import common_budgets as native +from litellm._auth.common_budgets import ( + AlertFacts, + BudgetDecision, + BudgetFacts, + BudgetScope, + MemberLimitFacts, + WindowFacts, +) + +from .budget_reference import common as reference + + +@pytest.mark.parametrize("scope", get_args(BudgetScope)) +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, nan, inf, -inf)) +@pytest.mark.parametrize("spend", (None, -1.0, 0.0, 1.0, 2.0, nan, inf, -inf)) +def test_numeric_budget_case_table_matches_reference( + scope: BudgetScope, limit: float | None, spend: float | None +) -> None: + facts: Final = BudgetFacts(scope, limit, spend, token_present=True, alert_recipients=("owner",)) + assert native.check(facts) == reference(facts) + + +@pytest.mark.parametrize("policy", (reference, native.check), ids=("reference", "native")) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (BudgetFacts("Team", 1, 1), BudgetDecision()), + (BudgetFacts("Key", 1, 1), BudgetDecision(denied=True, alert="token_budget")), + (BudgetFacts("Team", 1, 2, token_present=True), BudgetDecision(denied=True, alert="team_budget")), + (BudgetFacts("User", 1, 2, team_key=True), BudgetDecision()), + ( + BudgetFacts("User", 1, 2, team_key=True, apply_user_budget_to_team_keys=True), + BudgetDecision(denied=True, alert="user_budget"), + ), + (BudgetFacts("EndUser", 1, 2, info_route=True), BudgetDecision()), + (BudgetFacts("EndUser", 1, 1), BudgetDecision()), + (BudgetFacts("Project", inf), BudgetDecision()), + (BudgetFacts("Team", inf), BudgetDecision(read_spend=True)), + (BudgetFacts("Proxy", 1, 2, skip=True), BudgetDecision()), + (BudgetFacts("Proxy", 1, 2, llm_route=False), BudgetDecision()), + (BudgetFacts("Proxy", 1, 2, model_list_route=True), BudgetDecision()), + (BudgetFacts("Proxy", 1, 2), BudgetDecision(denied=True)), + ( + BudgetFacts("TeamSoft", 0, 0, token_present=True, alert_recipients=("a",)), + BudgetDecision(alert="soft_budget"), + ), + (BudgetFacts("TeamSoft", 1, 2, token_present=True), BudgetDecision()), + (BudgetFacts("ProjectSoft", 1, 2, token_present=True), BudgetDecision(alert="soft_budget")), + (BudgetFacts("KeySoft", 0, 1), BudgetDecision()), + (BudgetFacts("KeySoft", -1, 0), BudgetDecision(alert="soft_budget")), + (BudgetFacts("Tag", 0, nan), BudgetDecision(denied=True)), + (BudgetFacts("AccessGroup", nan, 0), BudgetDecision(denied=True)), + ), +) +def test_reference_and_native_obey_the_same_policy_table( + policy: Callable[[BudgetFacts], BudgetDecision], + facts: BudgetFacts, + expected: BudgetDecision, +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("percentage", (None, -1.0, 0.0, 0.5, 1.0, 1.1, nan, inf)) +@pytest.mark.parametrize("opt_in", (False, True)) +@pytest.mark.parametrize(("tpm", "rpm"), ((None, None), (0, None), (None, 0), (100, 100))) +def test_throttles_only_an_exceeded_opted_in_key( + percentage: float | None, opt_in: bool, tpm: int | None, rpm: int | None +) -> None: + facts: Final = BudgetFacts( + "Key", 1, 2, throttle_percentage=percentage, throttle_opt_in=opt_in, tpm_limit=tpm, rpm_limit=rpm + ) + assert native.check(facts) == reference(facts) + assert native.check(replace(facts, spend=0)) == BudgetDecision(alert="token_budget") + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (AlertFacts(8, 10), True), + (AlertFacts(10, 10), False), + (AlertFacts(0, 0), False), + (AlertFacts(nan, 10), False), + (AlertFacts(nan, 10, True, 80, team_member=True), True), + (AlertFacts(0, 0, True, 0, team_member=True), False), + (AlertFacts(10, 10, True, 80), True), + (AlertFacts(7, 10, True, 80), False), + (AlertFacts(10, 10, True, None), False), + (AlertFacts(10, 10, False, team_member=True), False), + ), +) +def test_threshold_alert_policy(facts: AlertFacts, expected: bool) -> None: + assert native.alert(facts) is expected + + +@pytest.mark.parametrize( + ("member", "default", "increase", "expiry", "expected"), + ( + (1, 5, 2, 101, 3), + (1, 5, 2, 100, 1), + (1, 5, 2, 99, 1), + (0, 5, 2, 101, 2), + (None, 5, 2, 101, 7), + (None, 0, 2, 101, None), + (None, -1, 2, 101, None), + (None, None, 2, 101, None), + (1, 5, None, 101, 1), + (1, 5, 2, None, 1), + ), +) +def test_temporary_member_budget_expiry_is_exclusive( + member: float | None, default: float | None, increase: float | None, expiry: float | None, expected: float | None +) -> None: + assert native.member_limit(MemberLimitFacts(member, default, increase, expiry, 100)) == expected + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (WindowFacts(60, 100, 200), 40), + (WindowFacts(60, None, 200), 140), + (WindowFacts(None, 100, 200), None), + (WindowFacts(0, 100, 200), 100), + ), +) +def test_counter_window_start(facts: WindowFacts, expected: float | None) -> None: + assert native.window_start(facts) == expected diff --git a/tests/unit/_auth/test_common_flow.py b/tests/unit/_auth/test_common_flow.py new file mode 100644 index 00000000000..365a10e654c --- /dev/null +++ b/tests/unit/_auth/test_common_flow.py @@ -0,0 +1,147 @@ +from itertools import product +from typing import Final + +import pytest + +from litellm._auth import common_flow + +_FACTS: Final = tuple( + common_flow.Facts(model, skip, team, user, token, user_id, agent_id, customer, models) + for model, skip, team, user, token, user_id, agent_id, customer, models in product( + (False, True), + (False, True), + (False, True), + (False, True), + (False, True), + (None, "", "user"), + (None, "", "agent"), + (False, True), + ((), ("model",)), + ) +) +_BUDGETS: Final = tuple( + common_flow.BudgetFacts(token, groups, customer, table) + for token, groups, customer, table in product((False, True), (0, 1), (False, True), (False, True)) +) +_CUSTOM: Final = tuple( + common_flow.CustomFacts(end_user_id, end_user, common, model, router) + for end_user_id, end_user, common, model, router in product( + (None, "", "customer"), (False, True), (False, True), (False, True), (False, True) + ) +) + + +@pytest.mark.parametrize("facts", _FACTS) +def test_common_model_plan_matches_reference(facts: common_flow.Facts) -> None: + assert common_flow.plan(facts) == common_flow.Plan( + membership=facts.token_present + and (facts.model_present or not facts.skip_budgets) + and facts.team_present + and facts.user_id is not None, + team_models=facts.model_present and facts.team_present, + member_models=facts.model_present and facts.team_present and facts.token_present and bool(facts.user_id), + user_models=facts.model_present and not facts.team_present and facts.user_present, + customer_models=facts.customer_present and bool(facts.customer_models), + agent_trace=facts.token_present and bool(facts.agent_id), + ) + + +@pytest.mark.parametrize("entity", (False, True)) +@pytest.mark.parametrize("budget", (False, True)) +@pytest.mark.parametrize("mapping", (False, True)) +@pytest.mark.parametrize("team", (False, True)) +def test_model_budget_scope_matches_reference(entity: bool, budget: bool, mapping: bool, team: bool) -> None: + expected: Final = "CheckModelBudgets" if entity and budget and mapping and team else "Continue" + assert common_flow.read(common_flow.ModelBudgets("Member", entity, budget, mapping, team)) == expected + + +@pytest.mark.parametrize("router", (False, True)) +@pytest.mark.parametrize("disabled", (False, True)) +def test_fallback_eligibility_matches_reference(router: bool, disabled: bool) -> None: + expected: Final = "AllowFallback" if router and not disabled else "Continue" + assert common_flow.read(common_flow.FallbackAllowed(router, disabled)) == expected + + +@pytest.mark.parametrize( + "facts, expected", + ( + (common_flow.ModelBudgets("Key"), "ReadModelBudgetPresence"), + (common_flow.ModelBudgets("Key", budget_present=False), "Continue"), + (common_flow.ModelBudgets("Key", budget_present=True), "ReadModelBudgetToken"), + (common_flow.ModelBudgets("User"), "ReadModelBudgetUser"), + (common_flow.ModelBudgets("User", entity_present=False), "Continue"), + (common_flow.ModelBudgets("User", entity_present=True), "ReadModelBudgetShape"), + (common_flow.ModelBudgets("User", entity_present=True, mapping=False), "Continue"), + (common_flow.ModelBudgets("Member", entity_present=True), "ReadModelBudgetTeam"), + (common_flow.ModelBudgets("Member", entity_present=True, team_present=False), "Continue"), + (common_flow.FallbackAllowed(False), "Continue"), + (common_flow.FallbackAllowed(True), "ReadFallbackDisabled"), + ), +) +def test_scope_reads_preserve_reference_short_circuit( + facts: common_flow.ModelBudgets | common_flow.FallbackAllowed, expected: common_flow.ReadPlan +) -> None: + assert common_flow.read(facts) == expected + + +@pytest.mark.parametrize("facts", _BUDGETS) +def test_budget_scopes_preserve_reference_priority(facts: common_flow.BudgetFacts) -> None: + assert common_flow.budgets(facts) == ( + "Team", + "TeamWindows", + *(("KeyWindows",) if facts.token_present else ()), + "TeamSoft", + "Organization", + "Tags", + *(("Groups",) if facts.group_count else ()), + "User", + "Member", + *(("Customer",) if facts.customer_present and facts.customer_budget_present else ()), + ) + + +@pytest.mark.parametrize("facts", _CUSTOM) +def test_custom_flow_matches_reference(facts: common_flow.CustomFacts) -> None: + assert common_flow.custom(facts) == common_flow.CustomPlan( + facts.end_user_id is not None, + facts.end_user_present and not facts.common_checks, + facts.common_checks, + facts.model_present and facts.router_present, + ) + + +@pytest.mark.parametrize("skip", (False, True)) +@pytest.mark.parametrize("exempt", (None, False, True)) +def test_budget_reads_preserve_short_circuit(skip: bool, exempt: bool | None) -> None: + expected: Final = "SkipBudgets" if skip or exempt else "ReadExemption" if exempt is None else "CheckBudgets" + assert common_flow.read(common_flow.BudgetSkip(skip, exempt)) == expected + + +@pytest.mark.parametrize("enabled", (False, True)) +@pytest.mark.parametrize("route", (None, False, True)) +def test_pricing_read_order(enabled: bool, route: bool | None) -> None: + expected: Final = "Continue" if not enabled or route is False else "ReadRoute" if route is None else "ReadUnpriced" + assert common_flow.read(common_flow.Pricing(enabled, route)) == expected + + +@pytest.mark.parametrize("count", (0, 1, 2)) +def test_unpriced_denial(count: int) -> None: + assert common_flow.read(common_flow.Unpriced(count)) == ("RejectUnpriced" if count else "Continue") + + +@pytest.mark.parametrize("router", (False, True)) +@pytest.mark.parametrize("multipart", (None, False, True)) +@pytest.mark.parametrize("resource", (None, False, True)) +def test_fallback_reads_preserve_body_binding_guards( + router: bool, multipart: bool | None, resource: bool | None +) -> None: + expected: Final = ( + "Continue" + if not router or multipart or resource + else "ReadContentType" + if multipart is None + else "ReadResourceBinding" + if resource is None + else "ReadFallbackSettings" + ) + assert common_flow.read(common_flow.MemberFallback(router, multipart, resource)) == expected diff --git a/tests/unit/_auth/test_failure.py b/tests/unit/_auth/test_failure.py new file mode 100644 index 00000000000..d2e42ac3455 --- /dev/null +++ b/tests/unit/_auth/test_failure.py @@ -0,0 +1,36 @@ +import pytest + +from litellm._auth.failure import NativeAuthFailure, native_boundary +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentAccess, AgentRequestHandler + + +@pytest.mark.parametrize("error", (TypeError("bad fact"), ValueError("bad wire"), RuntimeError("native error"))) +def test_native_errors_keep_public_exception_type_and_are_marked_terminal(error: Exception) -> None: + @native_boundary + def entry(value: int) -> int: + raise error + + with pytest.raises(type(error), match=str(error)) as raised: + entry(1) + assert isinstance(raised.value, NativeAuthFailure) + assert raised.value.__cause__ is error + + +@pytest.mark.asyncio +async def test_native_failure_never_becomes_unrestricted_agent_access() -> None: + @native_boundary + def broken_entry() -> None: + raise ValueError("invalid native policy answer") + + async def key_lookup(auth: UserAPIKeyAuth | None, *, strict: bool) -> AgentAccess: + broken_entry() + raise AssertionError("native failure returned") + + async def team_lookup(auth: UserAPIKeyAuth | None, *, strict: bool) -> AgentAccess: + raise AssertionError("team lookup followed native failure") + + with pytest.raises(NativeAuthFailure, match="invalid native policy answer"): + await AgentRequestHandler.resolve_key_team_agent_access( + UserAPIKeyAuth(), lookup_key_access=key_lookup, lookup_team_access=team_lookup + ) diff --git a/tests/unit/_auth/test_fallback_access.py b/tests/unit/_auth/test_fallback_access.py new file mode 100644 index 00000000000..ba7c7036c02 --- /dev/null +++ b/tests/unit/_auth/test_fallback_access.py @@ -0,0 +1,106 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import ( + FallbackAuthorizationDecision, + FallbackAuthorizationFacts, + GroupedModelScope, + fallback_authorization, +) +from litellm._auth.model_access import ModelAccessFacts +from tests.unit._auth.support.fallback_reference import authorization + +_ALLOW: Final = ModelAccessFacts(("model",), ("model",), object_type="key") +_DENY: Final = replace(_ALLOW, grants=("other",)) +_BASE: Final = FallbackAuthorizationFacts(GroupedModelScope(_ALLOW)) + + +@pytest.mark.parametrize("policy", (authorization, fallback_authorization)) +@pytest.mark.parametrize( + ("facts", "kind"), + ( + pytest.param(_BASE, "Allow", id="no-extra-ceilings"), + pytest.param(replace(_BASE, refuse_agent_keys=True, agent_id=""), "Deny", id="agent-presence-not-truthiness"), + pytest.param( + replace(_BASE, key=GroupedModelScope(_DENY, ("group",))), "ReadKeyGroups", id="key-groups-before-team" + ), + pytest.param( + replace(_BASE, key=GroupedModelScope(_DENY, ("group",), ())), "Deny", id="empty-key-group-retains-denial" + ), + pytest.param( + replace(_BASE, key=GroupedModelScope(_DENY, ("group",), ("model",))), "Allow", id="key-group-fallback" + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY, ("group",))), "ReadTeamGroups", id="team-group-fallback" + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY)), "ReadAssignedGroups", id="team-denial-asks-owned-key-grants" + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY), assigned_models=()), + "Deny", + id="empty-assigned-does-not-widen", + ), + pytest.param( + replace(_BASE, team=GroupedModelScope(_DENY), assigned_models=("model",)), + "ReadAssignedScope", + id="assigned-reads-scope-before-override", + ), + pytest.param( + replace( + _BASE, + team=GroupedModelScope(_DENY), + assigned_models=("model",), + assigned_scope=replace(_DENY, grants=("model",)), + ), + "Allow", + id="assigned-overrides-team", + ), + pytest.param( + replace(_BASE, embedded_team=_DENY, assigned_models=("model",)), + "Deny", + id="embedded-team-no-assigned-override", + ), + pytest.param(replace(_BASE, member=_DENY), "Deny", id="member-ceiling"), + pytest.param(replace(_BASE, project=_DENY), "Deny", id="project-ceiling"), + ), +) +def test_fallback_authorization_stage_order( + policy: Callable[[FallbackAuthorizationFacts], FallbackAuthorizationDecision], + facts: FallbackAuthorizationFacts, + kind: str, +) -> None: + assert policy(facts).kind == kind + + +def test_key_denial_wins_before_later_scopes() -> None: + decision: Final = fallback_authorization( + replace(_BASE, key=GroupedModelScope(_DENY), project=replace(_DENY, object_type="project")) + ) + assert decision.denial is not None + assert decision.denial.object_type == "key" + + +@pytest.mark.parametrize("policy", (authorization, fallback_authorization)) +@pytest.mark.parametrize( + ("facts", "kind"), + ( + (FallbackAuthorizationFacts(), "ReadKey"), + (FallbackAuthorizationFacts(refuse_agent_keys=True, agent_id="agent"), "Deny"), + (replace(_BASE, team_present=True), "ReadTeam"), + (replace(_BASE, embedded_team_grants=("model",)), "ReadEmbeddedTeam"), + (replace(_BASE, member_grants=("model",)), "ReadMember"), + (replace(_BASE, project_grants=("model",)), "ReadProject"), + (replace(_BASE, key=GroupedModelScope(_DENY), team_present=True, member_grants=("model",)), "Deny"), + (replace(_BASE, member=_DENY, project_grants=("model",)), "Deny"), + ), +) +def test_later_scope_projection_is_requested_only_after_prior_allow( + policy: Callable[[FallbackAuthorizationFacts], FallbackAuthorizationDecision], + facts: FallbackAuthorizationFacts, + kind: str, +) -> None: + assert policy(facts).kind == kind diff --git a/tests/unit/_auth/test_fallback_candidate.py b/tests/unit/_auth/test_fallback_candidate.py new file mode 100644 index 00000000000..f7ccd135b12 --- /dev/null +++ b/tests/unit/_auth/test_fallback_candidate.py @@ -0,0 +1,211 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.common_budgets import BudgetFacts +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import FallbackAuthorizationFacts, GroupedModelScope +from litellm._auth.fallback_candidate import CandidateBudgetRead, CandidateBudgetScope, fallback_candidate +from litellm._auth.fallback_candidate import FallbackCandidateDecision as Decision +from litellm._auth.fallback_candidate import FallbackCandidateFacts as Facts +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.model_attribution import ModelGroupAttributionFacts +from litellm._auth.model_budget_data import ModelBudgetEntity, ModelBudgetSnapshot +from litellm._auth.routes import RouteCatalog, RouteContext +from litellm._auth.unpriced_models import UnpricedModelFacts +from tests.unit._auth.support.fallback_reference import authorization +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_model_attribution import reference as attribution + +_ENTITIES: Final[tuple[ModelBudgetEntity, ...]] = ("Key", "User", "EndUser", "Team", "TeamMember") + + +def reference(facts: Facts) -> Decision: + if not all((facts.router_present, facts.user_present, facts.team_present)) or facts.membership == "Failed": + return Decision("Complete") + if facts.membership == "Unread": + return Decision("ReadMembership") + if facts.io_failed: + return Decision("Complete") + if facts.block_unpriced: + if facts.route is None: + return Decision("ReadRoute") + if facts.route.route in facts.route.catalog.groups.get("openai_routes", ()): + if facts.unpriced is None: + return Decision("ReadUnpriced") + if ( + facts.unpriced.group_found + and facts.unpriced.deployments == () + and facts.unpriced.cost_deployments == () + ): + return Decision("Complete") + if facts.customer_present: + if facts.customer_grants is None: + return Decision("ReadCustomerGrants") + if facts.customer_grants: + if facts.customer is None: + return Decision("ReadCustomer") + if model_access(facts.customer).kind != "Allow": + return Decision("Complete", log_failure=True) + if facts.price is None: + return Decision("ReadPrice") + if not facts.price.cached: + if facts.budget_models is None: + return Decision("ReadBudgetModels") + if facts.budget_scopes is None: + return Decision("ReadBudgetScopes") + for index, model in enumerate(facts.budget_models): + if model in facts.budget_models[:index]: + continue + for entity in _ENTITIES: + if not any( + scope.entity == entity and scope.identity_present and scope.configured + for scope in facts.budget_scopes + ): + continue + reading = next( + (reading for reading in facts.budgets if reading.model_index == index and reading.entity == entity), + None, + ) + if reading is None: + return Decision("ReadBudget", model_index=index, entity=entity) + snapshot = reading.snapshot + spend = (snapshot.local_spend or 0.0) + (snapshot.legacy_spend or 0.0) + if snapshot.max_budget is not None and snapshot.max_budget >= 0 and spend >= snapshot.max_budget: + return Decision("Complete") + if facts.authorization is None: + return Decision("ReadAuthorization") + if authorization(facts.authorization).kind == "Deny": + return Decision("Complete") + if not facts.attribution_reset: + return Decision("ResetAttribution") + if facts.attribution is None: + return Decision("ReadAttribution") + groups: Final = attribution(facts.attribution).groups + if not facts.attribution_stamped: + return Decision("StampAttribution", groups=groups) + if groups: + if not facts.group_budgets_read: + return Decision("ReadGroupBudgets", groups=groups) + if any( + budget.limit is not None and budget.limit > 0 and budget.spend is not None and budget.spend >= budget.limit + for budget in facts.group_budgets + ): + return Decision("Complete", log_failure=True) + return Decision("Complete", usable=True) + + +_INITIAL: Final = Facts(True, True, True) +_MEMBER: Final = replace(_INITIAL, membership="Loaded") +_FREE: Final = replace(_MEMBER, price=ModelPriceFacts(cached=True)) +_PAID: Final = replace(_MEMBER, price=ModelPriceFacts(cached=False)) +_SCOPE: Final = ModelAccessFacts(("model",), ("model",), object_type="key") +_AUTH: Final = replace(_FREE, authorization=FallbackAuthorizationFacts(key=GroupedModelScope(_SCOPE))) +_RESET: Final = replace(_AUTH, attribution_reset=True) +_GROUPS: Final = ModelGroupAttributionFacts( + model_present=True, + token_present=True, + router_present=True, + database_present=True, + registry_loaded=True, + registry=("group",), + covering=("group",), + grants=(("group",), (), (), (), ()), +) +_ATTRIBUTED: Final = replace(_RESET, attribution=_GROUPS) +_STAMPED: Final = replace(_ATTRIBUTED, attribution_stamped=True) +_SCOPES: Final = tuple(CandidateBudgetScope(entity, True, True) for entity in _ENTITIES) +_BUDGET: Final = replace(_PAID, budget_models=("model",), budget_scopes=_SCOPES) +_ROUTE: Final = RouteContext( + "/chat/completions", catalog=RouteCatalog(groups={"openai_routes": ("/chat/completions",)}) +) + + +@pytest.mark.parametrize("policy", (reference, fallback_candidate)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_INITIAL, Decision("ReadMembership")), + (replace(_INITIAL, router_present=False), Decision("Complete")), + (replace(_INITIAL, user_present=False), Decision("Complete")), + (replace(_INITIAL, team_present=False), Decision("Complete")), + (replace(_INITIAL, membership="Failed"), Decision("Complete")), + (replace(_MEMBER, io_failed=True), Decision("Complete")), + (replace(_MEMBER, block_unpriced=True), Decision("ReadRoute")), + (replace(_MEMBER, block_unpriced=True, route=_ROUTE), Decision("ReadUnpriced")), + ( + replace( + _MEMBER, + block_unpriced=True, + route=_ROUTE, + unpriced=UnpricedModelFacts("model", True, group_found=True, deployments=(), cost_deployments=()), + customer_present=True, + ), + Decision("Complete"), + ), + (replace(_MEMBER, block_unpriced=True, route=RouteContext("/health")), Decision("ReadPrice")), + (replace(_MEMBER, customer_present=True), Decision("ReadCustomerGrants")), + (replace(_MEMBER, customer_present=True, customer_grants=()), Decision("ReadPrice")), + (replace(_MEMBER, customer_present=True, customer_grants=("model",)), Decision("ReadCustomer")), + ( + replace( + _MEMBER, customer_present=True, customer_grants=("other",), customer=replace(_SCOPE, grants=("other",)) + ), + Decision("Complete", log_failure=True), + ), + (_MEMBER, Decision("ReadPrice")), + (_FREE, Decision("ReadAuthorization")), + (_PAID, Decision("ReadBudgetModels")), + (replace(_PAID, budget_models=("model",)), Decision("ReadBudgetScopes")), + (_BUDGET, Decision("ReadBudget", entity="Key")), + ( + replace( + _BUDGET, budgets=(CandidateBudgetRead(0, "Key", ModelBudgetSnapshot(max_budget=1, local_spend=1)),) + ), + Decision("Complete"), + ), + (replace(_BUDGET, budget_scopes=()), Decision("ReadAuthorization")), + ( + replace(_FREE, authorization=FallbackAuthorizationFacts(refuse_agent_keys=True, agent_id="")), + Decision("Complete"), + ), + (_AUTH, Decision("ResetAttribution")), + (_RESET, Decision("ReadAttribution")), + (_ATTRIBUTED, Decision("StampAttribution", groups=("group",))), + (_STAMPED, Decision("ReadGroupBudgets", groups=("group",))), + ( + replace(_STAMPED, group_budgets_read=True, group_budgets=(BudgetFacts("AccessGroup", 1, 1),)), + Decision("Complete", log_failure=True), + ), + ( + replace(_STAMPED, group_budgets_read=True, group_budgets=(BudgetFacts("AccessGroup", 1, 0),)), + Decision("Complete", usable=True), + ), + (replace(_STAMPED, attribution=ModelGroupAttributionFacts(io_failed=True)), Decision("Complete", usable=True)), + ), +) +def test_candidate_staged_case_table(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("policy", (reference, fallback_candidate)) +@pytest.mark.parametrize("index", range(5)) +def test_candidate_model_budget_order(policy: Callable[[Facts], Decision], index: int) -> None: + reads: Final = tuple(CandidateBudgetRead(0, entity, ModelBudgetSnapshot()) for entity in _ENTITIES[:index]) + assert policy(replace(_BUDGET, budgets=reads)) == Decision("ReadBudget", entity=_ENTITIES[index]) + + +@pytest.mark.parametrize("policy", (reference, fallback_candidate)) +def test_candidate_deduplicates_model_names_without_reordering(policy: Callable[[Facts], Decision]) -> None: + reads: Final = tuple(CandidateBudgetRead(0, entity, ModelBudgetSnapshot()) for entity in _ENTITIES) + assert policy(replace(_BUDGET, budget_models=("model", "model", "other"), budgets=reads)) == Decision( + "ReadBudget", model_index=2, entity="Key" + ) + + +def test_candidate_native_decode_failure_is_not_a_failed_candidate() -> None: + with pytest.raises(NativeAuthFailure): + fallback_candidate(replace(_BUDGET, budgets=(CandidateBudgetRead(-1, "Key", ModelBudgetSnapshot()),))) diff --git a/tests/unit/_auth/test_key_model_access.py b/tests/unit/_auth/test_key_model_access.py new file mode 100644 index 00000000000..9456f307518 --- /dev/null +++ b/tests/unit/_auth/test_key_model_access.py @@ -0,0 +1,80 @@ +from dataclasses import replace +from typing import Final + +import pytest + +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 ModelAccessFacts + +_ALLOW: Final = GroupedModelScope(ModelAccessFacts(("model",), ("model",), object_type="key")) +_DENY: Final = GroupedModelScope(replace(_ALLOW.facts, grants=("other",))) +_BASE: Final = KeyAndFallbackModelAccessFacts(False, ("model",)) + + +@pytest.mark.parametrize( + ("facts", "kind", "index"), + ( + pytest.param(_BASE, "ReadPrimary", 0, id="first-read"), + pytest.param(replace(_BASE, config_nonempty=True), "Allow", 0, id="config-before-primary"), + pytest.param(replace(_BASE, key_models=("all-team-models",)), "Allow", 0, id="team-sentinel-before-primary"), + pytest.param( + replace(_BASE, primary_loaded=True, primary=_DENY), "Deny", 0, id="primary-denial-before-fallback-parsing" + ), + pytest.param( + replace(_BASE, primary_loaded=True, primary=replace(_DENY, group_ids=("group",))), + "ReadPrimaryGroups", + 0, + id="key-group-before-fallback-parsing", + ), + pytest.param( + replace(_BASE, primary_loaded=True), "ReadFallbackNames", 0, id="no-primary-still-checks-fallbacks" + ), + pytest.param( + replace(_BASE, primary_loaded=True, fallback_names=("model",)), + "ReadFallbackScope", + 0, + id="fallback-needs-scope", + ), + pytest.param( + replace(_BASE, primary_loaded=True, fallback_names=("model",), fallback_scopes=(_ALLOW,)), + "RouteFallback", + 0, + id="allowed-fallback-then-probe", + ), + pytest.param( + replace(_BASE, primary_loaded=True, fallback_names=("model",), fallback_scopes=(_DENY,)), + "Deny", + 0, + id="denied-fallback-before-probe", + ), + pytest.param( + replace( + _BASE, + primary_loaded=True, + fallback_names=("first", "second"), + fallback_scopes=(_ALLOW,), + routing=("Failed",), + ), + "RoutingFailure", + 0, + id="routing-failure-before-next-scope", + ), + pytest.param( + replace( + _BASE, + primary_loaded=True, + fallback_names=("first", "second"), + fallback_scopes=(_ALLOW,), + routing=("Routed",), + ), + "ReadFallbackScope", + 1, + id="next-scope-after-routing", + ), + pytest.param(replace(_BASE, primary_loaded=True, fallback_names=()), "Allow", 0, id="empty-fallbacks"), + ), +) +def test_key_and_fallback_access_order(facts: KeyAndFallbackModelAccessFacts, kind: str, index: int) -> None: + decision: Final = key_and_fallback_model_access(facts) + assert (decision.kind, decision.index) == (kind, index) diff --git a/tests/unit/_auth/test_managed.py b/tests/unit/_auth/test_managed.py new file mode 100644 index 00000000000..9f33aab692e --- /dev/null +++ b/tests/unit/_auth/test_managed.py @@ -0,0 +1,84 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth import managed +from tests.unit._auth import reference_policy + +AGENT: Final = managed.Agent("agent", True, True, managed.BoundIdentity(True, "revision"), "both") +CONTEXT: Final = managed.Context("agent", "revision", "delegated", "user") + + +@pytest.mark.parametrize( + "agent,context", + ( + (AGENT, CONTEXT), + (replace(AGENT, enabled=False), None), + (replace(AGENT, identity=None), CONTEXT), + (replace(AGENT, identity=managed.BoundIdentity(False, "old")), CONTEXT), + (AGENT, None), + (AGENT, replace(CONTEXT, agent_id="other")), + (AGENT, replace(CONTEXT, binding_revision=None)), + (AGENT, replace(CONTEXT, binding_revision="old")), + (replace(AGENT, execution_mode="autonomous"), CONTEXT), + (AGENT, replace(CONTEXT, user_id=None)), + (AGENT, replace(CONTEXT, user_id="")), + ), +) +def test_actor_denial_order_preserves_verified_context_contract( + agent: managed.Agent, context: managed.Context | None +) -> None: + facts: Final = managed.Actor(agent, context) + assert managed.decide(facts) == reference_policy.actor(facts) + + +def test_actor_read_plan_assigns_policy_before_delegation_grant_read() -> None: + facts: Final = managed.ActorLookup("agent", True, context=CONTEXT) + assert managed.decide(facts).action == "ReadAgent" + loaded: Final = replace(facts, agent_loaded=True, agent=AGENT) + assert managed.decide(loaded).action == "AssignPolicy" + assigned: Final = replace(loaded, policy_assigned=True) + assert managed.decide(assigned).action == "ReadHumanGrants" + assert managed.decide(replace(assigned, human_grants=())).denial == "HumanNotPermitted" + assert managed.decide(replace(assigned, human_grants=("agent",))).action == "Allow" + assert managed.decide(replace(assigned, delegation_verified=True)).action == "Allow" + + +def test_missing_actor_still_reads_retirement_before_denial() -> None: + facts: Final = managed.ActorLookup("agent", True, agent_loaded=True, context=CONTEXT) + assert managed.decide(facts).action == "ReadRetired" + assert managed.decide(replace(facts, retired=False)).denial == "Removed" + assert managed.decide(replace(facts, retired=False, context=None)).action == "Allow" + + +def test_declared_grants_read_groups_in_order_and_keep_empty_declarations_restricted() -> None: + facts: Final = managed.DeclaredGrants((), ("group",), ("unified",)) + assert managed.decide(facts).action == "ReadDeclaredGroups" + grouped: Final = replace(facts, group_agents=()) + assert managed.decide(grouped).action == "ReadUnifiedGroups" + loaded: Final = replace(grouped, unified_agents=()) + assert managed.decide(loaded).grants == () + assert managed.decide(managed.DeclaredGrants((), (), ())).grants is None + + +@pytest.mark.parametrize("mode", (None, "autonomous", "delegated")) +@pytest.mark.parametrize("caller", (None, (), ("a",), ("a", "b"))) +@pytest.mark.parametrize("human", ((), ("b",), ("a", "b"))) +def test_managed_grants_intersect_each_authority_without_inheriting_unrestricted( + mode: str | None, caller: tuple[str, ...] | None, human: tuple[str, ...] +) -> None: + facts: Final = managed.ManagedAgentGrants( + True, + mode, + True, + own=("a", "b", "c"), + ceilings=(("a", "b"),), + caller_loaded=True, + caller=caller, + human=human, + ) + caller_ids: Final = frozenset(("a", "b", "c")) if caller is None else frozenset(caller) + human_ids: Final = frozenset(human) if mode == "delegated" else frozenset(("a", "b", "c")) + expected: Final = frozenset(("a", "b")) & caller_ids & human_ids + assert frozenset(managed.decide(facts).grants or ()) == expected diff --git a/tests/unit/_auth/test_managed_models.py b/tests/unit/_auth/test_managed_models.py new file mode 100644 index 00000000000..167ab0da702 --- /dev/null +++ b/tests/unit/_auth/test_managed_models.py @@ -0,0 +1,116 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.managed_models import ManagedModelDecision as Decision +from litellm._auth.managed_models import ManagedModelFacts as Facts +from litellm._auth.managed_models import managed_models +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_member_models import reference as member_reference + + +def reference(facts: Facts) -> Decision: + if facts.purpose == "All" and not facts.managed_loaded: + return Decision("ReadManaged") + if not facts.model_present or not facts.token_present: + return Decision("Allow") + if facts.purpose == "All" and facts.managed_present: + if not facts.managed_models_valid or not facts.managed_models_count: + return Decision("Deny", ModelAccessDecision("NoAgentModels")) + if facts.managed is None: + return Decision("ReadManagedScope") + managed: Final = model_access(facts.managed) + if managed.kind != "Allow": + return Decision("Deny", managed) + if facts.purpose != "Caller": + if facts.ceilings is None: + return Decision("ReadCeilings") + if not facts.aliases_loaded: + return Decision("ReadAliases") + for index, ceiling in enumerate(facts.ceilings): + if not ceiling: + return Decision("Deny", ModelAccessDecision("EmptyAgentCeiling")) + if index >= len(facts.scopes): + return Decision("ReadCeilingScope", index=index) + if (scope := model_access(facts.scopes[index])).kind != "Allow": + return Decision("Deny", scope) + if facts.purpose == "Ceilings": + return Decision("Allow") + if not facts.caller_loaded: + return Decision("ReadCaller") + if not facts.caller_present: + return Decision("Allow") + if not facts.team_loaded: + return Decision("ReadTeam") + if facts.team_present: + if facts.team is None: + return Decision("ReadTeamScope") + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if team.kind != "Allow": + return Decision("Deny", team) + if facts.member is None: + return Decision("ReadMember") + member: Final = member_reference(facts.member) + return Decision("Allow") if member.kind == "Allow" else Decision("Deny", member.denial, member.member_models) + if not facts.user_loaded: + return Decision("ReadUser") + if not facts.user_present or not facts.user_grants: + return Decision("Allow") + if "no-default-models" in facts.user_grants: + return Decision("Deny", ModelAccessDecision("NoDefaultModels")) + if facts.user is None: + return Decision("ReadUserScope") + user: Final = model_access(facts.user) + return Decision("Allow") if user.kind == "Allow" else Decision("Deny", user) + + +_ALL: Final = Facts("All", True, True) +_GROUPS: Final = Facts("Ceilings", True, True) +_CALLER: Final = Facts("Caller", True, True, caller_loaded=True, caller_present=True, team_loaded=True) +_ALLOW: Final = ModelAccessFacts(("model",), ("model",), object_type="agent") +_DENY: Final = replace(_ALLOW, grants=("other",)) +_TEAM: Final = GroupedModelScope(replace(_ALLOW, object_type="team")) + + +@pytest.mark.parametrize("policy", (reference, managed_models)) +@pytest.mark.parametrize( + "facts", + ( + _ALL, + replace(_ALL, model_present=False), + replace(_ALL, managed_loaded=True, model_present=False), + replace(_ALL, managed_loaded=True, managed_present=True), + replace(_ALL, managed_loaded=True, managed_present=True, managed_models_valid=True, managed_models_count=1), + replace( + _ALL, + managed_loaded=True, + managed_present=True, + managed_models_valid=True, + managed_models_count=1, + managed=_DENY, + ), + _GROUPS, + replace(_GROUPS, ceilings=()), + replace(_GROUPS, ceilings=(), aliases_loaded=True), + replace(_GROUPS, ceilings=((),), aliases_loaded=True), + replace(_GROUPS, ceilings=(("model",), ("other",)), aliases_loaded=True, scopes=(_ALLOW,)), + replace(_GROUPS, ceilings=(("model",), ("other",)), aliases_loaded=True, scopes=(_ALLOW, _DENY)), + _CALLER, + replace(_CALLER, team_present=True), + replace(_CALLER, team_present=True, team=replace(_TEAM, group_ids=("g",))), + replace(_CALLER, team_present=True, team=_TEAM), + replace(_CALLER, team_present=True, team=_TEAM, member=MemberModelFacts(True, True, True, ())), + replace(_CALLER, user_loaded=True, user_present=True, user_grants=("no-default-models",)), + replace(_CALLER, user_loaded=True, user_present=True, user_grants=("model",)), + ), +) +def test_managed_caller_and_ceiling_read_denial_order(policy: Callable[[Facts], Decision], facts: Facts) -> None: + assert policy(facts) == reference(facts) diff --git a/tests/unit/_auth/test_member_models.py b/tests/unit/_auth/test_member_models.py new file mode 100644 index 00000000000..b9901fd8706 --- /dev/null +++ b/tests/unit/_auth/test_member_models.py @@ -0,0 +1,56 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.member_models import MemberModelDecision as Decision +from litellm._auth.member_models import MemberModelFacts as Facts +from litellm._auth.member_models import member_models +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Facts) -> Decision: + if not facts.user_present or not facts.team_present: + return Decision("Allow") + if not facts.membership_loaded: + return Decision("ReadMembership") + if facts.grants is None: + return Decision("ReadGrants") + if not facts.grants: + return Decision("Allow") + if facts.scope is None: + return Decision("ReadScope") + denied: Final = model_access(facts.scope) + if denied.kind == "Allow": + return Decision("Allow") + return Decision("Deny", denied, facts.grants if denied.kind == "Deny" else None) + + +_BASE: Final = Facts(True, True) +_GRANTS: Final = replace(_BASE, membership_loaded=True, grants=("allowed",)) +_SCOPE: Final = ModelAccessFacts(("requested",), ("allowed",), object_type="team") + + +@pytest.mark.parametrize("policy", (reference, member_models)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadMembership")), + (replace(_BASE, user_present=False), Decision("Allow")), + (replace(_BASE, team_present=False), Decision("Allow")), + (replace(_BASE, membership_loaded=True), Decision("ReadGrants")), + (replace(_GRANTS, grants=()), Decision("Allow")), + (_GRANTS, Decision("ReadScope")), + (replace(_GRANTS, scope=replace(_SCOPE, grants=())), Decision("Allow")), + ( + replace(_GRANTS, scope=_SCOPE), + Decision("Deny", ModelAccessDecision("Deny", "requested", ("allowed",), "team"), ("allowed",)), + ), + ), +) +def test_member_scope_load_order_and_denial_envelope( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_model_access.py b/tests/unit/_auth/test_model_access.py new file mode 100644 index 00000000000..205e10ee2ba --- /dev/null +++ b/tests/unit/_auth/test_model_access.py @@ -0,0 +1,258 @@ +from collections.abc import Awaitable, Callable +from dataclasses import replace +from typing import Final, Literal +from unittest.mock import MagicMock + +import pytest + +from litellm._auth import model_access as native +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.model_access import ( + ModelAccessFacts as Facts, +) +from litellm._auth.model_access import ( + ModelAliases as Aliases, +) +from litellm._auth.model_access import ( + ModelProvider as Provider, +) +from litellm._auth.model_access import ( + ModelRecord as Record, +) +from litellm._auth.model_access import ( + ModelRegistry as Registry, +) +from litellm.proxy._types import LiteLLM_UserTable, ModelAccessDeniedProxyException +from litellm.proxy.auth import model_access_adapter +from tests.unit._auth.support import model_reference as reference +from tests.unit._auth.support import model_registry_reference + +_BASE: Final = Facts(("model",), object_type="key") +_CASES: Final = ( + pytest.param(_BASE, "Allow", id="empty-is-unrestricted"), + pytest.param(replace(_BASE, grants=("model",)), "Allow", id="exact-grant"), + pytest.param(replace(_BASE, grants=("other",)), "Deny", id="denied-model"), + pytest.param(replace(_BASE, grants=("*",)), "Allow", id="global-wildcard"), + pytest.param(replace(_BASE, grants=("all-proxy-models",)), "Allow", id="proxy-sentinel"), + pytest.param(replace(_BASE, grants=("mo*",)), "Allow", id="prefix-wildcard"), + pytest.param(replace(_BASE, requested=("model\n",), grants=("mo*",)), "Allow", id="python-final-newline"), + pytest.param(replace(_BASE, requested=("mod\nel",), grants=("mo*l",)), "Deny", id="embedded-newline"), + pytest.param( + replace(_BASE, grants=("p/*",), registry=Registry(records={"model": Record(provider=Provider("model", "p"))})), + "Allow", + id="inferred-provider", + ), + pytest.param( + replace( + _BASE, + requested=("p/model",), + grants=("q/*",), + registry=Registry(records={"p/model": Record(provider=Provider("model", "p"))}), + ), + "Deny", + id="different-provider", + ), + pytest.param( + replace(_BASE, grants=("other",), aliases=Aliases(team={"model": "target"})), "Allow", id="team-alias-exists" + ), + pytest.param( + replace(_BASE, grants=("model",), aliases=Aliases(key={"model": "target"})), + "Deny", + id="key-alias-replaces-source", + ), + pytest.param( + replace(_BASE, grants=("target",), aliases=Aliases(key={"model": "target"})), "Allow", id="key-target" + ), + pytest.param( + replace(_BASE, grants=("target",), aliases=Aliases(global_={"model": "target"})), "Allow", id="global-target" + ), + pytest.param( + replace( + _BASE, + grants=("target",), + aliases=Aliases( + team={"model": "middle"}, key={"middle": "third", "global": "target"}, global_={"third": "global"} + ), + ), + "Allow", + id="alias-order", + ), + pytest.param( + replace(_BASE, grants=("target",), registry=Registry(present=True, aliases={"model": "target"})), + "Allow", + id="router-alias", + ), + pytest.param( + replace( + _BASE, + grants=("target",), + registry=Registry(present=True, aliases={"model": "target"}), + aliases=Aliases(global_={"model": "wrong"}), + ), + "Deny", + id="global-over-router", + ), + pytest.param( + replace(_BASE, grants=("parent",), registry=Registry(compaction_parents={"model": "parent"})), + "Allow", + id="compaction-parent", + ), + pytest.param( + replace( + _BASE, + grants=("group",), + registry=Registry(present=True, records={"model": Record(access_groups=("group",))}), + ), + "Allow", + id="registry-group", + ), + pytest.param( + replace( + _BASE, + grants=("other-group",), + registry=Registry(present=True, records={"model": Record(access_groups=("group",))}), + ), + "Deny", + id="different-group", + ), + pytest.param( + replace(_BASE, grants=("all-team-models",), team_id="team", registry=Registry(present=True, names=("model",))), + "Allow", + id="team-sentinel-expands", + ), + pytest.param( + replace(_BASE, grants=("all-team-models",), team_id="team", registry=Registry(present=True)), + "Deny", + id="empty-registry-sentinel", + ), + pytest.param( + replace(_BASE, grants=("all-team-models",), registry=Registry(present=True, names=("model",))), + "Deny", + id="team-sentinel-no-team", + ), + pytest.param( + replace( + _BASE, + requested=("model_name_team_id",), + grants=("public",), + team_id="team", + registry=Registry( + present=True, + records={"model_name_team_id": Record(deployment_team="team", deployment_public_name="public")}, + ), + ), + "Allow", + id="owning-team-public-name", + ), + pytest.param( + replace( + _BASE, + requested=("model_name_team_id",), + grants=("public",), + team_id="team", + registry=Registry( + present=True, + records={"model_name_team_id": Record(deployment_team="other", deployment_public_name="public")}, + ), + ), + "Deny", + id="public-name-wrong-team", + ), + pytest.param( + replace( + _BASE, + requested=("model_name_team_id",), + grants=("other",), + team_id="team", + aliases=Aliases(team={"public": "model_name_team_id"}), + registry=Registry(present=True), + ), + "Deny", + id="owning-alias-not-unconditional", + ), + pytest.param( + replace(_BASE, requested=("model", "other"), grants=("model",), is_list=True), + "Deny", + id="all-list-items-required", + ), + pytest.param(replace(_BASE, requested=(), grants=("other",), is_list=True), "Allow", id="empty-list"), + pytest.param(replace(_BASE, fallback_depth=5), "Depth", id="depth"), + pytest.param(replace(_BASE, fallback_depth=4, is_list=True), "Depth", id="list-depth"), +) + + +@pytest.mark.parametrize("policy", (reference.model_access, native.model_access), ids=("reference", "native")) +@pytest.mark.parametrize(("facts", "kind"), _CASES) +def test_model_policy_case_table( + policy: Callable[[Facts], native.ModelAccessDecision], facts: Facts, kind: str +) -> None: + decision: Final = policy(facts) + assert decision.kind == kind + assert decision == reference.model_access(facts) + + +@pytest.mark.parametrize( + ("groups", "models", "merge", "kind"), + ( + ((), None, False, "Deny"), + (("g",), None, False, "ReadGroups"), + (("g",), (), False, "Deny"), + (("g",), ("model",), False, "Allow"), + (("g",), ("model",), True, "Allow"), + (("g",), ("other",), False, "Deny"), + ), +) +def test_group_fallback_waits_for_data( + groups: tuple[str, ...], models: tuple[str, ...] | None, merge: bool, kind: str +) -> None: + assert native.grouped_model_access(replace(_BASE, grants=("no",)), groups, models, merge=merge).kind == kind + + +@pytest.mark.parametrize( + ("present", "grants", "kind"), + ( + (False, ("no-default-models",), "Allow"), + (True, ("no-default-models", "model"), "NoDefaultModels"), + (True, (), "Allow"), + (True, ("other",), "Deny"), + ), +) +def test_user_default_scope(present: bool, grants: tuple[str, ...], kind: str) -> None: + assert native.user_model_access(replace(_BASE, grants=grants), present=present).kind == kind + + +@pytest.mark.parametrize( + "policy", (model_registry_reference.can_user_call_model, model_access_adapter.can_user_call_model) +) +@pytest.mark.asyncio +async def test_user_precheck_does_not_read_registry( + policy: Callable[[str, MagicMock, LiteLLM_UserTable | None], Awaitable[Literal[True]]], +) -> None: + router: Final = MagicMock() + user: Final = LiteLLM_UserTable(user_id="user", models=["no-default-models"]) + with pytest.raises(ModelAccessDeniedProxyException, match="model"): + await policy("model", router, user) + assert not router.mock_calls + assert await policy("model", router, None) is True + assert not router.mock_calls + + +def test_native_model_decode_failure_is_terminal() -> None: + with pytest.raises(NativeAuthFailure): + native.model_access(replace(_BASE, fallback_depth=-1)) + + +@pytest.mark.parametrize( + ("managed", "ceilings", "kind"), + ( + (None, (), "Allow"), + ((), (), "NoAgentModels"), + (("model",), ((),), "EmptyAgentCeiling"), + (("model",), (("model",), ("other",)), "Deny"), + (("model",), (("model",),), "Allow"), + ), +) +def test_agent_grants_are_intersected( + managed: tuple[str, ...] | None, ceilings: tuple[tuple[str, ...], ...], kind: str +) -> None: + assert native.agent_model_access(_BASE, managed, ceilings).kind == kind diff --git a/tests/unit/_auth/test_model_attribution.py b/tests/unit/_auth/test_model_attribution.py new file mode 100644 index 00000000000..09f85e0c9e2 --- /dev/null +++ b/tests/unit/_auth/test_model_attribution.py @@ -0,0 +1,87 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.model_attribution import ModelGroupAttributionDecision as Decision +from litellm._auth.model_attribution import ModelGroupAttributionFacts as Facts +from litellm._auth.model_attribution import model_group_attribution + + +def reference(facts: Facts) -> Decision: + if facts.io_failed or not all( + (facts.model_present, facts.token_present, facts.router_present, facts.database_present) + ): + return Decision("Complete") + if not facts.registry_loaded: + return Decision("ReadRegistry") + if facts.registry is not None and not facts.registry: + return Decision("Complete") + if facts.covering is None: + return Decision("ReadCovering") + budgeted: Final = set(facts.covering) if facts.registry is None else set(facts.covering) & set(facts.registry) + if not budgeted: + return Decision("Complete") + match len(facts.grants): + case 0: + return Decision("ReadKey") + case 1: + return Decision("ReadTeam") + case 2: + return Decision( + "ReadMember", + include=facts.team_present and facts.user_present, + lookup=facts.team_present and facts.user_present and not facts.membership_loaded, + ) + case 3: + return Decision("ReadProject") + case 4: + return Decision("ReadOrg", id=facts.token_org or facts.team_org) + case _: + scopes: Final = facts.grants[:4] if facts.org_failed else facts.grants + granted: Final = set[str]().union(*scopes) + groups: Final = tuple(sorted(budgeted & granted)) + return Decision("Complete", groups, stamp=bool(groups)) + + +_BASE: Final = Facts(model_present=True, token_present=True, router_present=True, database_present=True) +_COVERED: Final = replace(_BASE, registry_loaded=True, registry=("a", "b"), covering=("b", "a", "c", "b")) +_FINAL: Final = replace(_COVERED, grants=(("a",), (), ("b",), (), ("c",))) + + +@pytest.mark.parametrize("policy", (reference, model_group_attribution)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadRegistry")), + (replace(_BASE, model_present=False), Decision("Complete")), + (replace(_BASE, token_present=False), Decision("Complete")), + (replace(_BASE, router_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, io_failed=True), Decision("Complete")), + (replace(_BASE, registry_loaded=True, registry=()), Decision("Complete")), + (replace(_BASE, registry_loaded=True), Decision("ReadCovering")), + (replace(_COVERED, covering=()), Decision("Complete")), + (_COVERED, Decision("ReadKey")), + (replace(_COVERED, grants=((),)), Decision("ReadTeam")), + (replace(_COVERED, grants=((), ())), Decision("ReadMember")), + ( + replace(_COVERED, grants=((), ()), team_present=True, user_present=True), + Decision("ReadMember", include=True, lookup=True), + ), + ( + replace(_COVERED, grants=((), ()), team_present=True, user_present=True, membership_loaded=True), + Decision("ReadMember", include=True), + ), + (replace(_COVERED, grants=((), (), ())), Decision("ReadProject")), + (replace(_COVERED, grants=((), (), (), ()), token_org="key", team_org="team"), Decision("ReadOrg", id="key")), + (replace(_COVERED, grants=((), (), (), ()), token_org="", team_org="team"), Decision("ReadOrg", id="team")), + (_FINAL, Decision("Complete", ("a", "b"), stamp=True)), + (replace(_FINAL, registry=None), Decision("Complete", ("a", "b", "c"), stamp=True)), + (replace(_FINAL, registry=None, org_failed=True), Decision("Complete", ("a", "b"), stamp=True)), + (replace(_FINAL, grants=(("*",), (), (), (), ())), Decision("Complete")), + ), +) +def test_attribution_case_table(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_model_budget_data.py b/tests/unit/_auth/test_model_budget_data.py new file mode 100644 index 00000000000..6fb053d2d4b --- /dev/null +++ b/tests/unit/_auth/test_model_budget_data.py @@ -0,0 +1,90 @@ +from math import inf, isnan, nan +from typing import Final + +import pytest + +from litellm._auth import model_budget_data as native +from litellm._auth.model_budget_data import ( + ModelBudgetEntry, + ModelBudgetQueryFacts, + ModelBudgetResolveFacts, + ModelBudgetSnapshot, +) + +from .budget_reference import model as reference + + +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("local", (None, -1.0, 0.0, 1.0, 2.0, inf, nan)) +@pytest.mark.parametrize("shared", (None, 0.0, 2.0, inf, nan)) +@pytest.mark.parametrize("now", (99, 100)) +def test_raw_counter_and_shared_window_case_table( + limit: float | None, local: float | None, shared: float | None, now: float +) -> None: + facts: Final = ModelBudgetSnapshot(limit, local, 0, shared, 40, 60, now) + result: Final = native.check(facts) + expected: Final = reference(facts) + assert result.denied == expected.denied + assert result.read is None + assert ( + isnan(result.current_spend) and isnan(expected.current_spend) + ) or result.current_spend == expected.current_spend + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1)), "Local"), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(-1)), None), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(None)), None), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1, 0), 60, 59), "Legacy"), + (ModelBudgetQueryFacts("EndUser", ModelBudgetSnapshot(1, 0), 60, 59), "Legacy"), + (ModelBudgetQueryFacts("EndUser", ModelBudgetSnapshot(1, 0), 60, 60), None), + (ModelBudgetQueryFacts("User", ModelBudgetSnapshot(1, 0), 60, 1), "Shared"), + (ModelBudgetQueryFacts("Team", ModelBudgetSnapshot(1, 0), 60, 1, redis_present=True), None), + (ModelBudgetQueryFacts("TeamMember", ModelBudgetSnapshot(1, 1), 60), None), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1, 0), 60, same_counter=True), "Shared"), + (ModelBudgetQueryFacts("Key", ModelBudgetSnapshot(1, 0, 0), 60, shared_read=True), None), + ), +) +def test_native_selects_the_next_io_read(facts: ModelBudgetQueryFacts, expected: str | None) -> None: + assert native.query(facts).read == expected + + +def entry(name: str, limit: float | None = 1, duration: str | None = "1h") -> ModelBudgetEntry: + return ModelBudgetEntry(name, limit, duration, 3600 if duration is not None else None) + + +@pytest.mark.parametrize( + ("entries", "groups", "expected"), + ( + ((entry("*"), entry("m*"), entry("model")), ("group",), "model"), + ((entry("*"), entry("m*"), entry("group")), ("group",), "group"), + ((entry("*"), entry("m*"), entry("model", duration=None)), (), "m*"), + ((entry("provider/model"), entry("model")), (), "provider/model"), + ((entry("b"), entry("a"), entry("*")), ("b", "a"), "a"), + ((entry("model", None), entry("*")), (), "model"), + ((entry("m?del"),), (), None), + ((entry("*a*"), entry("mo*"), entry("mod*l")), (), "mod*l"), + ), +) +def test_resolution_prioritizes_exact_then_groups_then_specific_globs( + entries: tuple[ModelBudgetEntry, ...], + groups: tuple[str, ...], + expected: str | None, +) -> None: + result: Final = native.resolve(ModelBudgetResolveFacts("provider/model", "", "", entries, access_groups=groups)) + assert (result.entry.name if result.entry is not None else None) == expected + assert result.read_groups is False + + +@pytest.mark.parametrize( + ("key_limit", "expected"), ((0, "*"), (1, "*"), (-1, "model"), (None, "model"), (nan, "model")) +) +def test_only_enforced_named_key_entry_replaces_the_team_entry(key_limit: float | None, expected: str) -> None: + facts: Final = ModelBudgetResolveFacts( + "model", "", "", (entry("model"), entry("*")), (entry("model", key_limit),), team=True, access_groups=() + ) + result: Final = native.resolve(facts) + assert result.entry is not None + assert result.entry.name == expected diff --git a/tests/unit/_auth/test_model_budgets.py b/tests/unit/_auth/test_model_budgets.py index 5c704f8b465..d091e3d8c66 100644 --- a/tests/unit/_auth/test_model_budgets.py +++ b/tests/unit/_auth/test_model_budgets.py @@ -1,3 +1,5 @@ +from litellm._auth.fallback_candidate import FallbackCandidateFacts +from tests.unit._auth.raw_answers import candidate_read from dataclasses import dataclass, replace from typing import Final @@ -5,9 +7,11 @@ import pytest from litellm._auth.decision import ( Allow, + AuthCore, Deny, FallbackTarget, ModelBudgetExceeded, + ModelBudgetRead, ModelBudgetScope, Outcome, Step, @@ -15,7 +19,8 @@ from litellm._auth.decision import ( WithinModelBudgets, ) from litellm._auth.driver import drive_model_budgets, drive_team_member_model_budgets -from litellm._auth.decision import AuthCore +from litellm._auth.fallback_access import FallbackAuthorizationFacts +from tests.unit._auth.raw_answers import budget_read, fallback_read Call = tuple[str, str | None] @@ -54,13 +59,13 @@ class FakeRequest: self._record("read") return self.scope - async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: self._record("key", model) - return _exceeded("key", model) if model in self.key_over else None + return budget_read(_exceeded("key", model) if model in self.key_over else None) - async def authorize_budget_fallback(self, model: str) -> bool: + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: self._record("authorize", model) - return model not in self.unauthorized + return fallback_read(model, model not in self.unauthorized) async def rewrite_request_model(self, model: str) -> None: self._record("rewrite", model) @@ -70,9 +75,9 @@ class FakeRequest: rewritable_model=model, ) - async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: self._record("user", model) - return _exceeded("user", model) if model in self.user_over else None + return budget_read(_exceeded("user", model) if model in self.user_over else None) async def _run(step: Step, request: FakeRequest) -> tuple[Outcome, tuple[Call, ...]]: @@ -211,18 +216,18 @@ class FakeTeamMemberRequest: self._record("read") return self.scope - async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: self._record("member", model) - return _exceeded("member", model) if model in self.over else None + return budget_read(_exceeded("member", model) if model in self.over else None) async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: self._record("fallbacks", model) return self.targets if self.targets is not None else tuple(_target(fallback) for fallback in self.fallbacks) - async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: + async def vet_team_member_fallback(self, target: FallbackTarget) -> FallbackCandidateFacts: self._record("vet", target.model) self.vetted = (*self.vetted, target) - return target.model not in self.unusable + return candidate_read(target.model, target.model not in self.unusable) async def reroute_team_member_request(self, model: str) -> None: self._record("reroute", model) diff --git a/tests/unit/_auth/test_model_groups.py b/tests/unit/_auth/test_model_groups.py new file mode 100644 index 00000000000..49afbf26491 --- /dev/null +++ b/tests/unit/_auth/test_model_groups.py @@ -0,0 +1,51 @@ +from dataclasses import replace +from typing import Final, Literal + +import pytest + +from litellm._auth.model_groups import GroupLookup, GroupResources, ResourceGroupFacts, resource_group_decision + +_BASE: Final = ResourceGroupFacts(True, True, True, ("group",), "team", "key") + + +@pytest.mark.parametrize( + ("team_ids", "key_ids", "expected"), + ( + (("team",), (), ("model",)), + ((), ("key",), ("model",)), + (("other",), (), ()), + ((), ("other",), ()), + ((), (), ()), + ), +) +def test_group_ownership_boundary( + team_ids: tuple[str, ...], key_ids: tuple[str, ...], expected: tuple[str, ...] +) -> None: + decision: Final = resource_group_decision( + replace(_BASE, loaded=(GroupLookup("Found", GroupResources(team_ids, key_ids, ("model", "model"))),)) + ) + assert (decision.kind, decision.resources) == ("Complete", expected) + + +@pytest.mark.parametrize("kind", ("Missing", "Failed")) +@pytest.mark.parametrize("check_db_only", (False, True)) +@pytest.mark.parametrize("strict_errors", (False, True)) +def test_strict_lookup_errors_stop_before_next_group( + kind: Literal["Missing", "Failed"], check_db_only: bool, strict_errors: bool +) -> None: + decision: Final = resource_group_decision( + replace( + _BASE, + group_ids=("first", "second"), + check_db_only=check_db_only, + strict_errors=strict_errors, + loaded=(GroupLookup(kind),), + ) + ) + assert (decision.kind, decision.index) == (("Raise", 0) if check_db_only or strict_errors else ("Read", 1)) + + +@pytest.mark.parametrize("field", ("subject_present", "database_present", "cache_present")) +def test_unavailable_dependency_grants_nothing(field: str) -> None: + decision: Final = resource_group_decision(replace(_BASE, **{field: False})) + assert (decision.kind, decision.resources) == ("Complete", ()) diff --git a/tests/unit/_auth/test_model_helpers.py b/tests/unit/_auth/test_model_helpers.py new file mode 100644 index 00000000000..c038776f218 --- /dev/null +++ b/tests/unit/_auth/test_model_helpers.py @@ -0,0 +1,105 @@ +import re +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.model_access import ModelProvider +from litellm._auth.model_helpers import ModelHelperDecision as Decision +from litellm._auth.model_helpers import ModelHelperFacts as Facts +from litellm._auth.model_helpers import model_helper + + +def reference(facts: Facts) -> Decision: + if facts.kind == "Group" and (facts.grants is None or facts.model in facts.grants): + return Decision("Complete", True) + if facts.router_present and facts.groups is None: + return Decision("ReadGroups") + groups: Final = facts.groups or () + grants: Final = facts.grants or () + if facts.kind == "Group": + return Decision("Complete", any(grant in groups for grant in grants)) + expand: Final = "all-team-models" in grants and facts.team_id is not None and facts.router_present + if expand and facts.names is None: + return Decision("ReadNames") + expanded: Final = ( + ( + tuple( + dict.fromkeys((*tuple(grant for grant in grants if grant != "all-team-models"), *(facts.names or ()))) + ) + or grants + ) + if expand + else grants + ) + if any(grant in groups for grant in expanded) or facts.model in facts.team_aliases: + return Decision("Complete", True) + filtered: Final = tuple(grant for grant in expanded if grant not in groups) + patterns: Final = tuple(grant.replace("*", ".*") + "$" for grant in filtered if "*" in grant) + if any(re.match(pattern, facts.model) is not None for pattern in patterns): + return Decision("Complete", True) + if patterns: + if not facts.provider_loaded: + return Decision("ReadProvider") + provider: Final = facts.provider + if ( + provider is not None + and (provider.stripped_model != facts.model or "/" not in facts.model) + and any( + re.match(pattern, f"{provider.provider}/{provider.stripped_model}") is not None for pattern in patterns + ) + ): + return Decision("Complete", True) + return Decision( + "Complete", + not filtered and not expanded or "*" in filtered or "all-proxy-models" in filtered or facts.model in filtered, + ) + + +_GROUP: Final = Facts("Group", "model", ("model",), router_present=True) +_HELPER: Final = replace(_GROUP, kind="Helper") +_READ: Final = replace(_HELPER, groups=()) + + +@pytest.mark.parametrize("policy", (reference, model_helper)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_GROUP, Decision("Complete", True)), + (replace(_GROUP, grants=None), Decision("Complete", True)), + (replace(_GROUP, grants=()), Decision("ReadGroups")), + (replace(_GROUP, grants=(), groups=()), Decision("Complete")), + (replace(_GROUP, grants=("group",), groups=("group",)), Decision("Complete", True)), + (_HELPER, Decision("ReadGroups")), + (_READ, Decision("Complete", True)), + (replace(_READ, grants=()), Decision("Complete", True)), + (replace(_READ, grants=("other",)), Decision("Complete")), + (replace(_READ, grants=("other",), team_aliases=("model",)), Decision("Complete", True)), + (replace(_READ, grants=("group",), groups=("group",)), Decision("Complete", True)), + (replace(_READ, grants=("all-team-models",), team_id="team"), Decision("ReadNames")), + (replace(_READ, grants=("all-team-models",), team_id="team", names=()), Decision("Complete")), + (replace(_READ, grants=("all-team-models",), team_id="team", names=("model",)), Decision("Complete", True)), + (replace(_READ, grants=("mod*",)), Decision("Complete", True)), + (replace(_READ, grants=("provider/*",)), Decision("ReadProvider")), + (replace(_READ, grants=("provider/*",), provider_loaded=True), Decision("Complete")), + ( + replace(_READ, grants=("provider/*",), provider_loaded=True, provider=ModelProvider("model", "provider")), + Decision("Complete", True), + ), + ( + replace( + _READ, + model="unknown/model", + grants=("provider/*",), + provider_loaded=True, + provider=ModelProvider("unknown/model", "provider"), + ), + Decision("Complete"), + ), + ), +) +def test_model_helper_read_order_and_grants( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_owner_teams.py b/tests/unit/_auth/test_owner_teams.py new file mode 100644 index 00000000000..d8d40c79d7a --- /dev/null +++ b/tests/unit/_auth/test_owner_teams.py @@ -0,0 +1,91 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.owner_teams import OwnerTeamBudget, OwnerTeamLoad, owner_teams +from litellm._auth.owner_teams import OwnerTeamDecision as Decision +from litellm._auth.owner_teams import OwnerTeamFacts as Facts +from tests.unit._auth.test_personal_budget import ALLOW, INITIAL +from tests.unit._auth.test_personal_budget import reference as budget + + +def reference(facts: Facts) -> Decision: + if facts.purpose == "Grants" and not (facts.user_present and facts.database_present): + return Decision("Complete") + if not facts.loaded: + ids: Final = tuple(dict.fromkeys(facts.ids)) + return Decision( + "ReadTeams", + ids=tuple(team for team in ids if facts.purpose == "Models" or team and team != facts.ui_team_id), + ) + eligible: Final = tuple( + index + for index, row in enumerate(facts.loads) + if row.present + and not row.blocked + and (facts.purpose == "Models" or facts.org is None or row.org in (None, facts.org)) + ) + unreadable: Final = any(row.failed for row in facts.loads) + if facts.purpose == "Models": + return Decision("Complete", selected=eligible, unreadable=unreadable) + missing: Final = tuple(index for index in eligible if not any(read.index == index for read in facts.budgets)) + if missing: + return Decision("ReadBudgets", indices=missing) + outcomes: Final = tuple((read.index, budget(read.facts).allowed) for read in facts.budgets) + blocked: Final = tuple( + index + for index, row in enumerate(facts.loads) + if row.present and row.blocked and (facts.org is None or row.org in (None, facts.org)) + ) + return Decision( + "Complete", + selected=tuple(index for index in eligible if (index, True) in outcomes), + withheld=(*blocked, *(index for index in eligible if (index, False) in outcomes)), + unreadable=unreadable, + ) + + +_BASE: Final = Facts("Grants", ("", "ui", "a", "a", "b"), "ui", org="org") +_ROWS: Final = replace( + _BASE, + loaded=True, + loads=( + OwnerTeamLoad(True, blocked=True, org="org"), + OwnerTeamLoad(True, org="org"), + OwnerTeamLoad(True, org="other"), + OwnerTeamLoad(False, failed=True), + OwnerTeamLoad(True), + ), +) +_COMPLETE: Final = replace( + _ROWS, budgets=(OwnerTeamBudget(1, ALLOW), OwnerTeamBudget(4, replace(ALLOW, io_failed=True))) +) + + +@pytest.mark.parametrize("policy", (reference, owner_teams)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadTeams", ids=("a", "b"))), + (replace(_BASE, user_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, purpose="Models"), Decision("ReadTeams", ids=("", "ui", "a", "b"))), + (_ROWS, Decision("ReadBudgets", indices=(1, 4))), + (replace(_ROWS, purpose="Models"), Decision("Complete", selected=(1, 2, 4), unreadable=True)), + (_COMPLETE, Decision("Complete", selected=(1,), withheld=(0, 4), unreadable=True)), + (replace(_ROWS, org=None), Decision("ReadBudgets", indices=(1, 2, 4))), + (replace(_ROWS, org=""), Decision("ReadBudgets", indices=(4,))), + ), +) +def test_owner_team_selection_keeps_order_scope_and_withheld_rows( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +def test_later_incomplete_budget_is_not_hidden_by_an_allowed_team() -> None: + with pytest.raises(NativeAuthFailure): + owner_teams(replace(_ROWS, budgets=(OwnerTeamBudget(1, ALLOW), OwnerTeamBudget(4, INITIAL)))) diff --git a/tests/unit/_auth/test_personal_budget.py b/tests/unit/_auth/test_personal_budget.py new file mode 100644 index 00000000000..d22c35e44e5 --- /dev/null +++ b/tests/unit/_auth/test_personal_budget.py @@ -0,0 +1,149 @@ +import math +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.common_budgets import BudgetFacts, MemberLimitFacts +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.personal_budget import PersonalTeamBudgetDecision as Decision +from litellm._auth.personal_budget import PersonalTeamBudgetFacts as Facts +from litellm._auth.personal_budget import personal_team_budget + + +def denied(budget: BudgetFacts) -> bool: + if budget.limit is None or budget.spend is None or not math.isfinite(budget.limit): + return False + return budget.spend > budget.limit if budget.scope == "Team" else budget.spend >= budget.limit + + +def reference(facts: Facts) -> Decision: + org: Final = facts.token_org or facts.team_org + if not facts.source_ready: + return Decision("Source", org_id=org) + if facts.io_failed: + return Decision("Complete") + if facts.team is None: + return Decision("ReadTeam") + if denied(facts.team): + return Decision("Complete") if facts.team_alerted else Decision("AlertTeam") + if facts.window_count is None: + return Decision("ReadWindows") + for index in range(facts.window_count): + if index >= len(facts.windows): + return Decision("ReadWindow", index=index) + if denied(facts.windows[index]): + return Decision("Complete") + if facts.check_db_only and not facts.user_present: + return Decision("Complete") + if facts.user_present and (facts.team_id_present or facts.check_db_only) and not facts.membership_loaded: + return Decision("ReadMembership") + if facts.user_present and facts.team_id_present: + limit: Final = facts.member_limit + if limit is None: + return Decision("ReadMemberLimit") + if limit.member_limit is None and facts.default_id is not None and not facts.default_loaded: + return Decision("ReadDefault", id=facts.default_id) + base: Final = ( + limit.member_limit + if limit.member_limit is not None + else limit.default_limit + if limit.default_limit is not None and not limit.default_limit <= 0 + else None + ) + if ( + base is not None + and limit.temporary_increase is not None + and limit.expires_at is not None + and limit.now_naive + ): + return Decision("Complete", log_failure=True) + increase: Final = ( + limit.temporary_increase + if limit.temporary_increase is not None and limit.expires_at is not None and limit.expires_at > limit.now + else 0.0 + ) + if facts.member is None: + return Decision("ReadMember", limit=base + increase if base is not None else None) + if facts.member.limit is not None and math.isfinite(facts.member.limit) and not facts.member_alerted: + return Decision("AlertMember") + if denied(facts.member): + return Decision("Complete") + if org is not None: + if not facts.org_loaded: + return Decision("ReadOrg", id=org) + if not facts.org_failed and facts.org is not None and denied(facts.org): + return Decision("Complete") if facts.org_alerted else Decision("AlertOrg", id=org) + return Decision("Complete", allowed=True) + + +INITIAL: Final = Facts(user_present=True, token_org="org") +TEAM: Final = replace(INITIAL, source_ready=True, team=BudgetFacts("Team", None, token_present=True)) +WINDOWS: Final = replace(TEAM, window_count=0) +MEMBERSHIP: Final = replace(WINDOWS, membership_loaded=True) +LIMIT: Final = replace(MEMBERSHIP, member_limit=MemberLimitFacts(None, None, None, None, 100)) +MEMBER: Final = replace(LIMIT, member=BudgetFacts("TeamMember", None)) +ALLOW: Final = replace(MEMBER, org_loaded=True) + + +@pytest.mark.parametrize("policy", (reference, personal_team_budget)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (INITIAL, Decision("Source", org_id="org")), + (replace(INITIAL, token_org="", team_org="team-org"), Decision("Source", org_id="team-org")), + (replace(INITIAL, token_org=None, team_org=""), Decision("Source", org_id="")), + (replace(INITIAL, source_ready=True), Decision("ReadTeam")), + (replace(TEAM, io_failed=True), Decision("Complete")), + (replace(TEAM, team=BudgetFacts("Team", 1, 2, token_present=True)), Decision("AlertTeam")), + (replace(TEAM, team=BudgetFacts("Team", 1, 2, token_present=True), team_alerted=True), Decision("Complete")), + (replace(TEAM, team=BudgetFacts("Team", 1, 1, token_present=True)), Decision("ReadWindows")), + (TEAM, Decision("ReadWindows")), + (replace(TEAM, window_count=2), Decision("ReadWindow")), + (replace(TEAM, window_count=2, windows=(BudgetFacts("TeamWindow", 1, 0),)), Decision("ReadWindow", index=1)), + (replace(TEAM, window_count=2, windows=(BudgetFacts("TeamWindow", 1, 1),)), Decision("Complete")), + (replace(WINDOWS, user_present=False, check_db_only=True), Decision("Complete")), + (replace(WINDOWS, user_present=False), Decision("ReadOrg", id="org")), + (WINDOWS, Decision("ReadMembership")), + (replace(WINDOWS, team_id_present=False), Decision("ReadOrg", id="org")), + (replace(WINDOWS, team_id_present=False, check_db_only=True), Decision("ReadMembership")), + (replace(MEMBERSHIP, team_id_present=False, check_db_only=True), Decision("ReadOrg", id="org")), + (MEMBERSHIP, Decision("ReadMemberLimit")), + (replace(LIMIT, default_id=""), Decision("ReadDefault")), + ( + replace(LIMIT, default_id="default", member_limit=MemberLimitFacts(0, None, None, None, 100)), + Decision("ReadMember", limit=0), + ), + ( + replace( + LIMIT, default_id="default", default_loaded=True, member_limit=MemberLimitFacts(None, 10, 2, 101, 100) + ), + Decision("ReadMember", limit=12), + ), + (replace(LIMIT, member_limit=MemberLimitFacts(10, None, 2, 100, 100)), Decision("ReadMember", limit=10)), + ( + replace(LIMIT, member_limit=MemberLimitFacts(10, None, 2, 101, 100, True)), + Decision("Complete", log_failure=True), + ), + (LIMIT, Decision("ReadMember")), + (replace(MEMBER, member=BudgetFacts("TeamMember", 1, 1)), Decision("AlertMember")), + (replace(MEMBER, member=BudgetFacts("TeamMember", 1, 1), member_alerted=True), Decision("Complete")), + (replace(MEMBER, member=BudgetFacts("TeamMember", float("inf"), 100)), Decision("ReadOrg", id="org")), + (MEMBER, Decision("ReadOrg", id="org")), + (replace(MEMBER, token_org=None), Decision("Complete", allowed=True)), + (ALLOW, Decision("Complete", allowed=True)), + (replace(ALLOW, org_failed=True), Decision("Complete", allowed=True)), + (replace(ALLOW, org=BudgetFacts("Organization", 1, 1)), Decision("AlertOrg", id="org")), + (replace(ALLOW, org=BudgetFacts("Organization", 1, 1), org_alerted=True), Decision("Complete")), + ), +) +def test_personal_budget_read_and_denial_order( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +def test_incomplete_budget_read_is_terminal_protocol_failure() -> None: + with pytest.raises(NativeAuthFailure): + personal_team_budget(replace(TEAM, team=BudgetFacts("Team", 1))) diff --git a/tests/unit/_auth/test_personal_model_listing.py b/tests/unit/_auth/test_personal_model_listing.py new file mode 100644 index 00000000000..0d40963c03f --- /dev/null +++ b/tests/unit/_auth/test_personal_model_listing.py @@ -0,0 +1,66 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_model_listing import PersonalModelListingDecision as Decision +from litellm._auth.personal_model_listing import PersonalModelListingFacts as Facts +from litellm._auth.personal_model_listing import personal_model_listing +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm._auth.personal_owner import PersonalOwnerFacts +from tests.unit._auth.test_personal_models import ALLOW, INITIAL +from tests.unit._auth.test_personal_models import reference as scope_reference +from tests.unit._auth.test_personal_owner import reference as owner_reference + + +def reference(facts: Facts) -> Decision: + owner: Final = owner_reference(facts.owner) + if owner.kind != "Complete": + return Decision("ReadOwner") + if facts.owner.teams is None or not facts.owner.teams.loads: + return Decision("Complete") + if facts.names is None: + return Decision("ReadNames") + if len(facts.models) < len(facts.names): + return Decision("ReadModel", index=len(facts.models), teams=owner.selected) + return Decision( + "Complete", + selected=tuple( + index + for index, model in enumerate(facts.models) + if any(tuple(scope_reference(scope).allowed for scope in model.scopes)) + ), + ) + + +_BASE: Final = Facts(PersonalOwnerFacts(True, True, True, True)) +_OWNER: Final = replace( + _BASE.owner, + owner_loaded=True, + owner_present=True, + teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)), +) +_LOADED: Final = replace(_BASE, owner=_OWNER, names=("first", "second", "first")) +_ALLOWED: Final = LoadedTeamModelFacts((True,), (ALLOW,)) +_DENIED: Final = LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)) + + +@pytest.mark.parametrize("policy", (reference, personal_model_listing)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadOwner")), + (replace(_BASE, owner=replace(_BASE.owner, database_present=False)), Decision("Complete")), + (replace(_BASE, owner=replace(_BASE.owner, owner_loaded=True, owner_failed=True)), Decision("Complete")), + (replace(_BASE, owner=_OWNER), Decision("ReadNames")), + (_LOADED, Decision("ReadModel", teams=(0,))), + (replace(_LOADED, models=(_DENIED,)), Decision("ReadModel", index=1, teams=(0,))), + (replace(_LOADED, models=(_ALLOWED, _DENIED, _ALLOWED)), Decision("Complete", selected=(0, 2))), + ), +) +def test_listing_preserves_order_duplicates_and_checks_after_denial( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_personal_model_request.py b/tests/unit/_auth/test_personal_model_request.py new file mode 100644 index 00000000000..40b2edf752f --- /dev/null +++ b/tests/unit/_auth/test_personal_model_request.py @@ -0,0 +1,57 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_model_request import PersonalModelRequestDecision as Decision +from litellm._auth.personal_model_request import PersonalModelRequestFacts as Facts +from litellm._auth.personal_model_request import personal_model_request +from litellm._auth.personal_models import LoadedTeamModelFacts +from tests.unit._auth.test_owner_teams import reference as teams_reference +from tests.unit._auth.test_personal_models import ALLOW, INITIAL, reference as scope_reference + + +def reference(facts: Facts) -> Decision: + if facts.owner_present and facts.teams is None: + return Decision("ReadTeams") + teams: Final = teams_reference(facts.teams).selected if facts.teams is not None else () + for index, model in enumerate(facts.requested): + if not facts.owner_present: + return Decision("Deny", model=model) + if index >= len(facts.models): + return Decision("ReadModel", index=index, teams=teams) + if not any(scope_reference(scope).allowed for scope in facts.models[index].scopes): + return Decision("Deny", model=model) + return Decision("Allow") + + +_BASE: Final = Facts(("first", "second"), True) +_TEAMS: Final = replace( + _BASE, teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)) +) +_ONE: Final = replace(_TEAMS, models=(LoadedTeamModelFacts((True,), (ALLOW,)),)) + + +@pytest.mark.parametrize("policy", (reference, personal_model_request)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadTeams")), + (replace(_BASE, requested=()), Decision("ReadTeams")), + (replace(_BASE, owner_present=False), Decision("Deny", model="first")), + (replace(_BASE, owner_present=False, requested=()), Decision("Allow")), + (_TEAMS, Decision("ReadModel", teams=(0,))), + (_ONE, Decision("ReadModel", index=1, teams=(0,))), + (replace(_ONE, models=(*_ONE.models, LoadedTeamModelFacts((True,), (ALLOW,)))), Decision("Allow")), + ( + replace(_TEAMS, models=(LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)),)), + Decision("Deny", model="first"), + ), + ), +) +def test_personal_requests_read_all_teams_but_stop_at_first_model_denial( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_personal_models.py b/tests/unit/_auth/test_personal_models.py new file mode 100644 index 00000000000..5e0dffaa0bb --- /dev/null +++ b/tests/unit/_auth/test_personal_models.py @@ -0,0 +1,110 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts, personal_model_access_applies +from litellm._auth.personal_models import LoadedTeamModelFacts, loaded_team_models, personal_model_scope +from litellm._auth.personal_models import PersonalModelDecision as Decision +from litellm._auth.personal_models import PersonalModelScope as Scope +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Scope) -> Decision: + if facts.io_failed: + return Decision("Complete") + if facts.team is None: + return Decision("ReadTeam") + team: Final = grouped(facts.team, merge=True) + if team.kind == "ReadGroups": + return Decision("ReadGroups") + if team.kind != "Allow": + return Decision("Complete") + if not facts.user_present or not facts.team_id_present: + return Decision("Complete", allowed=True) + if facts.member_grants is None: + return Decision("ReadMembership") + if not facts.member_grants: + return Decision("Complete", allowed=True) + if facts.member is None: + return Decision("ReadMember") + return Decision("Complete", allowed=model_access(facts.member).kind == "Allow") + + +INITIAL: Final = Scope(True, True) +MODEL: Final = ModelAccessFacts(("model",), ("model",), object_type="team") +TEAM: Final = replace(INITIAL, team=GroupedModelScope(MODEL)) +ALLOW: Final = replace(TEAM, member_grants=()) + + +@pytest.mark.parametrize("policy", (reference, personal_model_scope)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (INITIAL, Decision("ReadTeam")), + (replace(INITIAL, io_failed=True), Decision("Complete")), + (TEAM, Decision("ReadMembership")), + (replace(TEAM, team=GroupedModelScope(replace(MODEL, grants=("other",)))), Decision("Complete")), + (replace(TEAM, team=GroupedModelScope(replace(MODEL, grants=("other",)), ("group",))), Decision("ReadGroups")), + ( + replace(TEAM, team=GroupedModelScope(replace(MODEL, grants=("other",)), ("group",), ("model",))), + Decision("ReadMembership"), + ), + (replace(TEAM, team_id_present=False), Decision("Complete", allowed=True)), + (replace(TEAM, user_present=False), Decision("Complete", allowed=True)), + (ALLOW, Decision("Complete", allowed=True)), + (replace(TEAM, member_grants=("member",)), Decision("ReadMember")), + (replace(TEAM, member_grants=("other",), member=replace(MODEL, grants=("other",))), Decision("Complete")), + (replace(TEAM, member_grants=("model",), member=MODEL), Decision("Complete", allowed=True)), + ), +) +def test_team_model_scope_orders_grants_before_membership( + policy: Callable[[Scope], Decision], facts: Scope, expected: Decision +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (LoadedTeamModelFacts(()), Decision("Complete")), + (LoadedTeamModelFacts((False, True, False, True)), Decision("ReadScopes", indices=(1, 3))), + ( + LoadedTeamModelFacts((True, True), (ALLOW, replace(INITIAL, io_failed=True))), + Decision("Complete", allowed=True), + ), + (LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)), Decision("Complete")), + ), +) +def test_loaded_teams_are_all_read_before_union(facts: LoadedTeamModelFacts, expected: Decision) -> None: + assert loaded_team_models(facts) == expected + + +def test_later_native_failure_is_not_hidden_by_first_team_allow() -> None: + with pytest.raises(NativeAuthFailure): + loaded_team_models(LoadedTeamModelFacts((True, True), (ALLOW, INITIAL))) + + +@pytest.mark.parametrize( + ("api_key_present", "inheritance", "expected"), ((False, False, True), (False, True, False), (True, True, True)) +) +def test_inheritance_requires_virtual_key_not_just_owner( + api_key_present: bool, inheritance: bool, expected: bool +) -> None: + assert ( + personal_model_access_applies( + enabled=True, + token_present=True, + team_id=None, + user_id="", + master=False, + role=None, + api_key_present=api_key_present, + inheritance=inheritance, + ) + is expected + ) diff --git a/tests/unit/_auth/test_personal_owner.py b/tests/unit/_auth/test_personal_owner.py new file mode 100644 index 00000000000..02e75764fcd --- /dev/null +++ b/tests/unit/_auth/test_personal_owner.py @@ -0,0 +1,61 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_owner import PersonalOwnerDecision as Decision +from litellm._auth.personal_owner import PersonalOwnerFacts as Facts +from litellm._auth.personal_owner import personal_owner +from tests.unit._auth.test_owner_teams import reference as select_teams + + +def reference(facts: Facts) -> Decision: + if not all((facts.user_present, facts.database_present, facts.cache_present, facts.logging_present)): + return Decision("Complete") + if not facts.owner_loaded: + return Decision("ReadOwner") + if facts.owner_failed or not facts.owner_present: + return Decision("Complete") + if facts.teams is None: + return Decision("ReadTeams") + teams: Final = select_teams(facts.teams) + return Decision("Complete", teams.selected, teams.withheld, teams.unreadable) + + +_BASE: Final = Facts(True, True, True, True) +_OWNER: Final = replace(_BASE, owner_loaded=True, owner_present=True) + + +@pytest.mark.parametrize("policy", (reference, personal_owner)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadOwner")), + (replace(_BASE, user_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, cache_present=False), Decision("Complete")), + (replace(_BASE, logging_present=False), Decision("Complete")), + (replace(_BASE, owner_loaded=True), Decision("Complete")), + (replace(_OWNER, owner_failed=True), Decision("Complete")), + (_OWNER, Decision("ReadTeams")), + ( + replace( + _OWNER, + teams=OwnerTeamFacts( + "Grants", + (), + "ui", + loaded=True, + loads=(OwnerTeamLoad(True, blocked=True), OwnerTeamLoad(False, failed=True)), + ), + ), + Decision("Complete", withheld=(0,), unreadable=True), + ), + ), +) +def test_owner_lookup_and_team_status_order( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py index 4f1fe014d0f..3d37c7b557d 100644 --- a/tests/unit/_auth/test_python_core.py +++ b/tests/unit/_auth/test_python_core.py @@ -6,22 +6,25 @@ from zoneinfo import ZoneInfo import pytest +from litellm._auth.budget_pricing import ExemptionFacts, RouteBudgetFacts from litellm._auth.decision import ( AdminKey, Allow, + AuthCore, CachedAdminKey, Deny, ExpiredKey, KeyBlocked, KeyBudgetExceeded, KeyFacts, - ModelBudgetExceeded, - ModelBudgetScope, KeyWithoutToken, MalformedKey, MasterKey, MasterKeyOnlyRoute, MissingKey, + ModelBudgetExceeded, + ModelBudgetRead, + ModelBudgetScope, NoDatabase, NoMasterKey, NonAdminKey, @@ -40,8 +43,10 @@ from litellm._auth.decision import ( VirtualKey, ) from litellm._auth.driver import drive +from litellm._auth.fallback_access import FallbackAuthorizationFacts from litellm._auth.key import hash_token -from litellm._auth.decision import AuthCore +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts +from tests.unit._auth.raw_answers import budget_read, fallback_read MASTER: Final = "sk-master" NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc) @@ -121,16 +126,20 @@ class RecordingHandler: self._record("inspect_key") return self.key - async def enforce_model_access(self) -> None: + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: self._record("enforce_model_access") + return KeyAndFallbackModelAccessFacts(True, ()) + + async def prefetch_auth_objects(self) -> None: + return None async def load_key_owner(self, user_id: str) -> OwnerFacts | None: self._record("load_key_owner", user_id) return self.owner - async def check_budget_exemption(self) -> bool: + async def check_budget_exemption(self) -> ExemptionFacts: self._record("check_budget_exemption") - return self.exempt + return ExemptionFacts(RouteBudgetFacts(discovery_route=self.exempt, llm_route=True), False, False, 0) async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: self._record("read_team_member_budget", (user_id, team_id)) @@ -156,20 +165,20 @@ class RecordingHandler: self._record("read_model_budget_scope") return self.model_budget_scope - async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_key_model_budget(self, model: str) -> ModelBudgetRead: self._record("check_key_model_budget", model) - return None + return budget_read() - async def authorize_budget_fallback(self, model: str) -> bool: + async def authorize_budget_fallback(self, model: str) -> FallbackAuthorizationFacts: self._record("authorize_budget_fallback", model) - return True + return fallback_read(model) async def rewrite_request_model(self, model: str) -> None: self._record("rewrite_request_model", model) - async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + async def check_user_model_budget(self, model: str) -> ModelBudgetRead: self._record("check_user_model_budget", model) - return self.user_model_budget_exceeded + return budget_read(self.user_model_budget_exceeded) async def load_key_context(self) -> None: self._record("load_key_context") diff --git a/tests/unit/_auth/test_reservation.py b/tests/unit/_auth/test_reservation.py new file mode 100644 index 00000000000..23398b47c0b --- /dev/null +++ b/tests/unit/_auth/test_reservation.py @@ -0,0 +1,123 @@ +from dataclasses import replace +from itertools import product +from math import inf, nan +from typing import Final + +import pytest + +from litellm._auth import reservation + +_SCOPES: Final[tuple[reservation.Scope, ...]] = ( + "Key", + "Team", + "User", + "EndUser", + "Tag", + "AccessGroup", + "Member", + "Organization", + "Project", +) +_SOURCES: Final = tuple( + reservation.Sources(key, team, team_id, user, user_id, apply) + for key, team, team_id, user, user_id, apply in product( + (False, True), (False, True), (None, "", "team"), (False, True), (None, "", "user"), (False, True) + ) +) + + +@pytest.mark.parametrize("scope", _SCOPES) +@pytest.mark.parametrize("limit", (None, -inf, -1.0, 0.0, 1.0, inf, nan)) +def test_reservation_counter_reference(scope: reservation.Scope, limit: float | None) -> None: + expected: Final = limit is not None and ( + limit > 0 + if scope in ("Key", "Team", "User") + else limit > 0 and limit < inf + if scope == "Project" + else not limit <= 0 + ) + assert reservation.counter(scope, limit) == expected + + +@pytest.mark.parametrize("facts", _SOURCES) +def test_source_selection_preserves_team_key_user_opt_in(facts: reservation.Sources) -> None: + team: Final = facts.team_present and facts.team_id is not None + assert reservation.sources(facts) == reservation.SourcePlan( + facts.key_present, + team, + (not team or facts.apply_user_budget_to_team_keys) and facts.user_present and facts.user_id is not None, + ) + + +@pytest.mark.parametrize("cost", (None, -inf, -1.0, 0.0, 1.0, inf, nan)) +def test_estimate_and_reservation_staging(cost: float | None) -> None: + initial: Final = reservation.Gate(True) + route: Final = replace(initial, llm_route=True) + billable: Final = replace(route, unbilled=False) + model: Final = replace(billable, model_present=True) + counters: Final = replace(model, counter_count=1) + assert tuple(reservation.action(facts).kind for facts in (initial, route, billable, model, counters)) == ( + "ReadRoute", + "ReadUnbilled", + "ReadModel", + "ReadCounters", + "ReadEstimate", + ) + assert reservation.action(replace(counters, cost_read=True, cost=cost)).kind == ( + "Skip" if cost is None or cost <= 0 else "Reserve" + ) + + +@pytest.mark.parametrize("duration", (False, True)) +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("parsed", (None, False, True)) +def test_window_parse_and_warning_timing(duration: bool, limit: float | None, parsed: bool | None) -> None: + expected: Final = ( + "Skip" + if not duration or limit is None or limit <= 0 + else "ReadStart" + if parsed is None + else "AppendWindow" + if parsed + else "WarnWindow" + ) + assert reservation.action(reservation.Window(duration, limit, parsed)).kind == expected + + +@pytest.mark.parametrize("scope", ("EndUser", "Organization", "Project")) +@pytest.mark.parametrize("primary", (None, "", "primary")) +@pytest.mark.parametrize("fallback", (None, "", "fallback")) +def test_reservation_identity_keeps_empty_scope_semantics( + scope: str, primary: str | None, fallback: str | None +) -> None: + from pydantic import TypeAdapter + + facts: Final = TypeAdapter(reservation.Identity).validate_python( + {"scope": scope, "primary": primary, "fallback": fallback} + ) + expected: Final = ( + primary or fallback + if scope == "EndUser" + else primary + if scope == "Project" or primary is not None + else fallback + ) + assert reservation.identity(facts) == expected + + +@pytest.mark.parametrize("team", (False, True)) +@pytest.mark.parametrize("user", (False, True)) +@pytest.mark.parametrize("team_id", (None, "", "team")) +@pytest.mark.parametrize("user_id", (None, "", "user")) +def test_reservation_membership_requires_both_subjects( + team: bool, user: bool, team_id: str | None, user_id: str | None +) -> None: + expected: Final = "ReadMember" if team and team_id is not None and user and user_id is not None else "Skip" + assert reservation.action(reservation.MemberGate(team, team_id, user, user_id)).kind == expected + + +@pytest.mark.parametrize("cap", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("read", (False, True)) +def test_member_default_read_only_when_own_cap_absent(cap: float | None, read: bool) -> None: + expected: Final = "ReadDefault" if cap is None and not read else "ReadCap" + assert reservation.action(reservation.MemberCap(cap, read)).kind == expected diff --git a/tests/unit/_auth/test_resolved_models.py b/tests/unit/_auth/test_resolved_models.py new file mode 100644 index 00000000000..7b301e9b3be --- /dev/null +++ b/tests/unit/_auth/test_resolved_models.py @@ -0,0 +1,156 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.member_models import MemberModelFacts +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm._auth.resolved_models import ResolvedModelDecision as Decision +from litellm._auth.resolved_models import ResolvedModelFacts as Facts +from litellm._auth.resolved_models import ResolvedTeamRead, resolved_models +from litellm._auth.team_override import TeamOverrideFacts +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access +from tests.unit._auth.test_member_models import reference as member_reference +from tests.unit._auth.test_personal_model_request import reference as personal_reference +from tests.unit._auth.test_team_override import reference as team_reference + + +def reference(facts: Facts) -> Decision: + if not facts.config_nonempty and "all-team-models" not in facts.key_models: + if facts.key is None: + return Decision("ReadKey") + key: Final = grouped(facts.key, merge=False) + if key.kind == "ReadGroups": + return Decision("ReadKeyGroups") + if key.kind != "Allow": + return Decision("Deny", key) + if facts.team_id is not None: + if facts.team_read is None: + return Decision("ReadTeamObject") + if facts.team_read.failed and not facts.embedded_ready: + return Decision("ReadEmbeddedTeam") + if facts.team_read.found or facts.team_read.failed: + if facts.team is None: + return Decision("ReadTeamScope") + team: Final = team_reference(facts.team) + if team.kind == "Deny": + return Decision("Deny", team.denial) + if facts.team_read.found and facts.user_id is not None: + if facts.member is None: + return Decision("ReadMember") + member: Final = member_reference(facts.member) + if member.kind == "Deny": + return Decision("Deny", member.denial, member.member_models) + if ( + facts.personal_enabled + and facts.team_id is None + and facts.user_id is not None + and not facts.master + and facts.role != "proxy_admin" + ): + if not facts.owner_loaded: + return Decision("ReadOwner") + if facts.owner_failed or not facts.owner_present: + return Decision("Deny", ModelAccessDecision("PersonalTeam", model=facts.model)) + if facts.personal is None: + return Decision("ReadPersonal") + personal: Final = personal_reference(facts.personal) + if personal.kind == "Deny": + return Decision("Deny", ModelAccessDecision("PersonalTeam", model=personal.model)) + if facts.project_present: + if not facts.project_loaded: + return Decision("ReadProjectObject") + if facts.project_grants: + if facts.project is None: + return Decision("ReadProjectScope") + project: Final = model_access(facts.project) + if project.kind != "Allow": + return Decision("Deny", project) + if facts.customer_present and facts.database_present: + if not facts.customer_loaded: + return Decision("ReadCustomerObject") + if facts.customer_grants: + if facts.customer is None: + return Decision("ReadCustomerScope") + customer: Final = model_access(facts.customer) + if customer.kind != "Allow": + return Decision("Deny", customer) + return Decision("Allow") + + +_MODEL: Final = ModelAccessFacts(("model",), ("model",)) +_PASS: Final = Facts(model="model", config_nonempty=True) +_TEAM: Final = replace( + _PASS, + team_id="team", + team_read=ResolvedTeamRead(found=True), + team=TeamOverrideFacts(GroupedModelScope(replace(_MODEL, object_type="team"))), +) +_PERSONAL: Final = replace(_PASS, user_id="owner", personal_enabled=True) +_PROJECT: Final = replace(_PASS, project_present=True, project_loaded=True, project_grants=("model",)) +_CUSTOMER: Final = replace( + _PASS, customer_present=True, database_present=True, customer_loaded=True, customer_grants=("model",) +) + + +@pytest.mark.parametrize("policy", (reference, resolved_models)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (Facts(), Decision("ReadKey")), + (_PASS, Decision("Allow")), + (Facts(key_models=("all-team-models",)), Decision("Allow")), + (Facts(key=GroupedModelScope(_MODEL)), Decision("Allow")), + ( + Facts(key=GroupedModelScope(replace(_MODEL, grants=("other",)), ("group",)), project_present=True), + Decision("ReadKeyGroups"), + ), + (replace(_TEAM, team_read=None), Decision("ReadTeamObject")), + (replace(_TEAM, team_read=ResolvedTeamRead(failed=True)), Decision("ReadEmbeddedTeam")), + (replace(_TEAM, team=None), Decision("ReadTeamScope")), + (_TEAM, Decision("Allow")), + (replace(_TEAM, user_id="member", project_present=True), Decision("ReadMember")), + ( + replace(_TEAM, user_id="member", member=MemberModelFacts(True, False), project_present=True), + Decision("ReadProjectObject"), + ), + ( + replace( + _TEAM, + user_id="member", + team_read=ResolvedTeamRead(failed=True), + embedded_ready=True, + project_present=True, + ), + Decision("ReadProjectObject"), + ), + (_PERSONAL, Decision("ReadOwner")), + (replace(_PERSONAL, master=True), Decision("Allow")), + (replace(_PERSONAL, role="proxy_admin"), Decision("Allow")), + (replace(_PERSONAL, owner_loaded=True), Decision("Deny", ModelAccessDecision("PersonalTeam", model="model"))), + ( + replace(_PERSONAL, owner_loaded=True, owner_present=True, owner_failed=True), + Decision("Deny", ModelAccessDecision("PersonalTeam", model="model")), + ), + (replace(_PERSONAL, owner_loaded=True, owner_present=True), Decision("ReadPersonal")), + (replace(_PROJECT, project_loaded=False), Decision("ReadProjectObject")), + (_PROJECT, Decision("ReadProjectScope")), + (replace(_PROJECT, project_grants=()), Decision("Allow")), + ( + replace( + _PROJECT, project=replace(_MODEL, object_type="project"), customer_present=True, database_present=True + ), + Decision("ReadCustomerObject"), + ), + (replace(_CUSTOMER, database_present=False), Decision("Allow")), + (replace(_CUSTOMER, customer_loaded=False), Decision("ReadCustomerObject")), + (_CUSTOMER, Decision("ReadCustomerScope")), + (replace(_CUSTOMER, customer_grants=()), Decision("Allow")), + (replace(_CUSTOMER, customer=replace(_MODEL, object_type="customer")), Decision("Allow")), + ), +) +def test_resolved_scope_and_lookup_order(policy: Callable[[Facts], Decision], facts: Facts, expected: Decision) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_route_enable.py b/tests/unit/_auth/test_route_enable.py new file mode 100644 index 00000000000..bdc2ab66335 --- /dev/null +++ b/tests/unit/_auth/test_route_enable.py @@ -0,0 +1,75 @@ +from collections.abc import Callable +from typing import Final, Literal + +import pytest + +from litellm._auth.route_enable import ( + LlmRouteDecision, + RouteEnableDecision, + RouteEnableFacts, + RouteSetting, + llm_route, + route_enable, +) +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_enable_reference import route_enable as reference + + +@pytest.mark.parametrize("policy", (reference, route_enable)) +@pytest.mark.parametrize( + "route", ("/unknown", "/models", "/v1/models", "/chat/completions", "/key/generate", "/user/info") +) +@pytest.mark.parametrize("imported", (False, True)) +@pytest.mark.parametrize("failed", (None, "Other", "Http")) +@pytest.mark.parametrize("exemptions", (None, (), ("/models", "/v1/models"))) +@pytest.mark.parametrize( + "setting", + ( + None, + RouteSetting(True, False), + RouteSetting(False, False), + RouteSetting(True, True), + RouteSetting(True, True, True), + RouteSetting(True, True, False), + ), +) +def test_optional_route_enable_decision_table( + policy: Callable[[RouteEnableFacts], RouteEnableDecision], + route: str, + imported: bool, + failed: Literal["Other", "Http"] | None, + exemptions: tuple[str, ...] | None, + setting: RouteSetting | None, +) -> None: + facts: Final = RouteEnableFacts(route_context(route), imported, failed, setting, setting, exemptions) + assert policy(facts) == reference(facts) + + +@pytest.mark.parametrize( + ("route", "configured", "expected"), + ( + ("/chat/completions", None, "Allow"), + ("/unknown", None, "Deny"), + ("/eu_jev/v1/systemone", None, "ReadEuJev"), + ("/eu_jev/v1/systemone", False, "Allow"), + ("/eu_jev/v1/systemone", True, "Deny"), + ("/eu_jev/v1/systemone/", True, "Deny"), + ), +) +def test_llm_route_configuration_read_is_selected_only_at_builtin_fallback( + route: str, configured: bool | None, expected: LlmRouteDecision +) -> None: + assert llm_route(route_context(route), configured) == expected + + +@pytest.mark.parametrize("configured", (None, False, True)) +@pytest.mark.parametrize("route", ("/eu_jev/v1/systemone", "/eu_jev/v1/systemone/")) +def test_optional_route_builtin_configuration_matches_reference(configured: bool | None, route: str) -> None: + facts: Final = RouteEnableFacts( + route_context(route), + imported=True, + llm=RouteSetting(True, True, True), + exemptions=(), + eu_jev_configured=configured, + ) + assert route_enable(facts) == reference(facts) diff --git a/tests/unit/_auth/test_routes.py b/tests/unit/_auth/test_routes.py new file mode 100644 index 00000000000..31f7d5b6b37 --- /dev/null +++ b/tests/unit/_auth/test_routes.py @@ -0,0 +1,191 @@ +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException, Request + +from litellm._auth import routes +from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.route_policy_adapter import route_context +from tests.unit._auth.support.route_reference import RouteChecks as ReferenceRoutes +from tests.unit._auth.support.route_reference import managed_agent_route_allowed + + +@dataclass(frozen=True, slots=True) +class RouteResult: + allowed: bool + status: int | None = None + detail: object = None + + +def outcome(call: Callable[[], object]) -> RouteResult: + try: + result: Final = call() + return RouteResult(result is not False) + except HTTPException as error: + return RouteResult(False, error.status_code, error.detail) + + +_ROUTES: Final = tuple( + dict.fromkeys( + ( + "/unknown", + "/chat/completions", + "/v1/mcp/server", + "/v1/mcp/server/id", + "/v1/mcp/server/id/more", + "/auto_router/session", + "/models", + "/v1/responses/id", + "/agents", + "/v1/agents/id", + "/key/id/regenerate", + "/user/update", + "/user/password/change", + "/session/logout", + "/user/info", + "/organization/new", + "/team/new", + "/openai/deployments/model/versions/chat/completions", + "/engines/model/chat/completions", + "/foo/count_tokens", + "/v1beta/models/model:countTokens", + "/unknown\n", + "/key/generate", + "/key/list", + "/session/list", + ) + ) +) + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize("method", (None, "GET", "POST", "DELETE", "HEAD", "OPTIONS")) +def test_route_classifications_match_reference(route: str, method: str | None) -> None: + request: Final = Request({"type": "http", "method": method or "GET", "path": route, "headers": []}) + assert RouteChecks.is_llm_api_route(route) == ReferenceRoutes.is_llm_api_route(route) + assert RouteChecks.is_unbilled_route(route) == ReferenceRoutes.is_unbilled_route(route) + assert RouteChecks.is_self_scoped_credential_route( + route, request + ) == ReferenceRoutes.is_self_scoped_credential_route(route, request) + assert RouteChecks._is_get_mcp_server_discovery_route( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route, request + ) == ReferenceRoutes._is_get_mcp_server_discovery_route(route, request) # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize( + "allowed", + ( + None, + (), + ("llm_api_routes",), + ("llm_api_routes_without_mcp",), + ("management_routes",), + ("/v1/*",), + ("/user",), + ("*",), + ("/no",), + ), +) +@pytest.mark.parametrize("method", ("GET", "POST")) +def test_virtual_route_allow_and_denial_envelope(route: str, allowed: tuple[str, ...] | None, method: str) -> None: + request: Final = Request({"type": "http", "method": method, "path": route, "headers": []}) + token: Final = UserAPIKeyAuth(allowed_routes=list(allowed) if allowed is not None else None) + assert outcome(lambda: RouteChecks.is_virtual_key_allowed_to_call_route(route, token, request)) == outcome( + lambda: ReferenceRoutes.is_virtual_key_allowed_to_call_route(route, token, request) + ) + + +@pytest.mark.parametrize("route", _ROUTES) +@pytest.mark.parametrize("method", ("GET", "POST", "PATCH", "HEAD")) +@pytest.mark.parametrize("keys", ((), ("user_email",), ("user_id", "user_email"))) +def test_viewer_role_denial_order_and_envelope(route: str, method: str, keys: tuple[str, ...]) -> None: + request: Final = Request({"type": "http", "method": method, "path": route, "headers": []}) + data: Final = {key: "value" for key in keys} + assert outcome( + lambda: RouteChecks._check_proxy_admin_viewer_access( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # compare the preserved private route behavior against native projection + route=route, _user_role="proxy_admin_viewer", request_data=data, request=request + ) + ) == outcome( + lambda: ReferenceRoutes._check_proxy_admin_viewer_access( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route=route, _user_role="proxy_admin_viewer", request_data=data, request=request + ) + ) + + +@pytest.mark.parametrize("group", tuple(LiteLLMRoutes)) +def test_full_catalog_classification(group: LiteLLMRoutes) -> None: + assert tuple(RouteChecks.is_llm_api_route(route) for route in group.value) == tuple( + ReferenceRoutes.is_llm_api_route(route) for route in group.value + ) + + +@pytest.mark.parametrize( + ("route", "pattern"), + ( + ("/a/id", "/a/{id}"), + ("/a/x/y", "/a/{id:path}"), + ("/a/x/y:action", "/a/{id:path}:action"), + ("/a/x\ny", "/a/{id:path}"), + ("/a/id\n", "/a/{id}"), + ("/a/b", "/a/*"), + ("/ab", "/a"), + ("/a", ""), + ("/a", "/"), + ("/a/id:action", "/a/{id}:action"), + ("/a/id", "/a/{id:path}"), + ), +) +def test_route_pattern_modes(route: str, pattern: str) -> None: + assert RouteChecks._route_matches_pattern(route, pattern) == ReferenceRoutes._route_matches_pattern(route, pattern) # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + assert RouteChecks._route_matches_allowed_route(route, pattern) == ReferenceRoutes._route_matches_allowed_route( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route, pattern + ) + assert RouteChecks.route_matches_wildcard_pattern(route, pattern) == ReferenceRoutes.route_matches_wildcard_pattern( + route, pattern + ) + assert RouteChecks._route_matches_denied_route(route, pattern) == ReferenceRoutes._route_matches_denied_route( # pyright: ignore[reportPrivateUsage] # compare the preserved private route behavior against native projection + route, pattern + ) + + +@pytest.mark.parametrize( + "route", + ( + "/agents", + "/v1/agents", + "/chat/completions", + "/v1/chat/completions", + "/v1/responses/id", + "/realtime", + "/v1/realtime", + "/v1/mcp/server/id", + "/mcp/foo", + "/user/info", + "/v1beta/models/x:generateContent", + ), +) +@pytest.mark.parametrize("method", (None, "GET", "HEAD", "POST", "DELETE", "get")) +def test_managed_actor_route_methods(route: str, method: str | None) -> None: + assert routes.classify_route(route_context(route, method), "Managed") == managed_agent_route_allowed(route, method) + + +@pytest.mark.parametrize( + ("key", "team", "jwt", "denied", "expected"), + ( + (("/p",), (), (), (), True), + ((), ("/p",), (), (), True), + (("/q",), ("/p",), (), (), False), + ((), (), ("/p/*",), (), True), + ((), (), ("*",), (), False), + (("/p",), (), (), ("/p",), False), + ), +) +def test_passthrough_precedence( + key: tuple[str, ...], team: tuple[str, ...], jwt: tuple[str, ...], denied: tuple[str, ...], expected: bool +) -> None: + facts: Final = routes.PassThroughFacts(key, team, denied, ("/p/resource",), jwt) + assert routes.passthrough_route(routes.RouteContext("/p/resource"), facts, required=True).allowed == expected diff --git a/tests/unit/_auth/test_runtime.py b/tests/unit/_auth/test_runtime.py index 830b743ee6e..be773a265c1 100644 --- a/tests/unit/_auth/test_runtime.py +++ b/tests/unit/_auth/test_runtime.py @@ -8,7 +8,7 @@ from litellm._auth.runtime import get_auth_core from tests.unit._auth.test_python_core import NOW, PLAIN_KEY, SETTINGS -@pytest.mark.parametrize("backend", (None, "rust", "shadow")) +@pytest.mark.parametrize("backend", (None, "rust", "python", "shadow")) def test_default_and_native_modes_consume_their_decision_once(backend: str | None, monkeypatch: pytest.MonkeyPatch): if backend is None: monkeypatch.delenv("LITELLM_AUTH_CORE", raising=False) @@ -20,17 +20,3 @@ def test_default_and_native_modes_consume_their_decision_once(backend: str | Non assert step.resume(blocked) == Deny(KeyBlocked()) with pytest.raises(RuntimeError, match="already consumed"): step.resume(blocked) - - -def test_python_reference_can_be_selected_explicitly(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("LITELLM_AUTH_CORE", "python") - step: Final = get_auth_core().start_resolved_key(RequestFacts("/chat/completions", "hash", True), SETTINGS, NOW) - assert isinstance(step, InspectKey) - blocked: Final = replace(PLAIN_KEY, blocked=True) - assert step.resume(blocked) == step.resume(blocked) == Deny(KeyBlocked()) - - -def test_invalid_backend_fails_before_native_entry(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("LITELLM_AUTH_CORE", "invalid") - with pytest.raises(ValueError, match="LITELLM_AUTH_CORE must be python, rust or shadow"): - get_auth_core() diff --git a/tests/unit/_auth/test_rust_core.py b/tests/unit/_auth/test_rust_core.py index 96e39e0ffa2..8db6fda188c 100644 --- a/tests/unit/_auth/test_rust_core.py +++ b/tests/unit/_auth/test_rust_core.py @@ -15,9 +15,10 @@ from litellm._auth.decision import ( WithinModelBudgets, ) from litellm._auth.driver import drive +from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts from litellm._auth.rust_core import RustCore -from litellm._auth.shadow import ShadowCore, ShadowMismatch from litellm.rust_bridge import _native +from tests.unit._auth.shadow import ShadowCore, ShadowMismatch from tests.unit._auth.test_python_core import NOW, PLAIN_KEY, SETTINGS, RecordingHandler @@ -115,7 +116,7 @@ async def test_effect_failure_after_native_entry_is_never_replayed(core: AuthCor failure: Final = RuntimeError("effect failed") class FailingHandler(RecordingHandler): - async def enforce_model_access(self) -> None: + async def enforce_model_access(self) -> KeyAndFallbackModelAccessFacts: self._record("enforce_model_access") raise failure diff --git a/tests/unit/_auth/test_spend_counter_policy.py b/tests/unit/_auth/test_spend_counter_policy.py new file mode 100644 index 00000000000..e62dfc009da --- /dev/null +++ b/tests/unit/_auth/test_spend_counter_policy.py @@ -0,0 +1,81 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth import spend_counter_policy as core + +_SPEND: Final = core.Spend(2.0, False, 0.0, 10.0, False, False) + + +@pytest.mark.parametrize( + "facts, expected", + ( + (_SPEND, core.Action("ReadStrictBefore")), + (replace(_SPEND, limit=None), core.SpendValue("Value", 2.0)), + (replace(_SPEND, current=10.0), core.SpendValue("Value", 10.0)), + (replace(_SPEND, current=11.0), core.SpendValue("Value", 11.0)), + (replace(_SPEND, fallback=3.0), core.Action("ReadFloor")), + (replace(_SPEND, window=True), core.Action("ReadFloor")), + (replace(_SPEND, strict_before=True), core.Action("ReadFloor")), + (replace(_SPEND, strict_before=False), core.Action("ReadStrictAfter")), + (replace(_SPEND, strict_before=False, verified=True), core.SpendValue("Value", 2.0)), + (replace(_SPEND, strict_before=False, fallback_authoritative=True), core.SpendValue("Value", 2.0)), + (replace(_SPEND, strict_before=True, authoritative_read=True), core.Action("ReadStrictAfter")), + ( + replace(_SPEND, strict_before=True, authoritative_read=True, authoritative=3.0), + core.SpendValue("Repair", 3.0), + ), + ( + replace(_SPEND, strict_before=True, authoritative_read=True, authoritative=1.0), + core.SpendValue("Value", 2.0), + ), + (replace(_SPEND, fallback=3.0, authoritative_read=True), core.SpendValue("Value", 3.0)), + (replace(_SPEND, strict_before=False, strict_after=True), core.Action("Unverifiable")), + ( + replace(_SPEND, strict_before=True, authoritative_read=True, strict_after=False), + core.SpendValue("Value", 2.0), + ), + ), +) +def test_spend_selection_preserves_lazy_flag_reads(facts: core.Spend, expected: core.Decision) -> None: + assert core.decide(facts) == expected + + +@pytest.mark.parametrize("current", (10.0 + 1e-13, 10.5, 11.0, 12.0, 13.0)) +@pytest.mark.parametrize("strict", (False, True)) +@pytest.mark.parametrize("throttled", (False, True)) +def test_overflow_matches_reference(current: float, strict: bool, throttled: bool) -> None: + remaining: Final = 10.0 - (current - 2.0) + expected: Final = ( + core.Action("Release") + if throttled + else core.Rejected("Reject", current, None) + if remaining <= 1e-12 + else core.Rejected("Reject", current - 2.0, 2.0) + if strict and current - 10.0 > 1e-12 + else core.Resize("Resize", remaining) + ) + assert core.decide(core.OverBudget(throttled, current, 10.0, 2.0, strict)) == expected + + +@pytest.mark.parametrize("token_present", (False, True)) +@pytest.mark.parametrize("token", (None, "", "hash")) +@pytest.mark.parametrize("counter", ("spend:key:hash", "spend:key:None", "spend:key:", "spend:team:hash")) +def test_throttle_is_limited_to_the_key_counter(token_present: bool, token: str | None, counter: str) -> None: + expected: Final = "ReadThrottle" if token_present and counter == f"spend:key:{token}" else "Skip" + assert core.decide(core.ThrottleKey(counter, token_present, token)) == core.Action(expected) + + +@pytest.mark.parametrize("currents, limits", (((), ()), ((0.0, 1.0), (2.0, 2.0)), ((0.0, 2.0), (2.0, 2.0)))) +def test_batching_preserves_order_and_overflow_boundaries( + currents: tuple[float, ...], limits: tuple[float, ...] +) -> None: + expected: Final = ( + () + if not currents + else (tuple(range(len(currents))),) + if all(current + 1.0 <= limit for current, limit in zip(currents, limits, strict=True)) + else tuple((index,) for index in range(len(currents))) + ) + assert core.decide(core.ReservationGroups(currents, limits, 1.0)) == core.Groups("ReservationGroups", expected) diff --git a/tests/unit/_auth/test_team_aliases.py b/tests/unit/_auth/test_team_aliases.py new file mode 100644 index 00000000000..b61a34204dc --- /dev/null +++ b/tests/unit/_auth/test_team_aliases.py @@ -0,0 +1,49 @@ +from collections.abc import Callable +from typing import Final + +import pytest + +from litellm._auth.team_aliases import TeamAliasDecision as Decision +from litellm._auth.team_aliases import TeamAliasFacts as Facts +from litellm._auth.team_aliases import team_aliases + + +def reference(facts: Facts) -> Decision: + targets: Final = tuple(facts.aliases.get(name) for name in facts.requested) + count: Final = sum( + bool(target is not None and facts.router_present and target.startswith(f"model_name_{facts.team_id}_")) + for target in targets + ) + if count and facts.names is None: + return Decision("ReadNames", count=count) + reads: Final = iter(facts.names or ()) + return Decision( + "Complete", + models=tuple( + name + if target is None + or ( + facts.router_present and target.startswith(f"model_name_{facts.team_id}_") and target not in next(reads) + ) + else target + for name, target in zip(facts.requested, targets, strict=True) + ), + ) + + +@pytest.mark.parametrize("policy", (reference, team_aliases)) +@pytest.mark.parametrize( + "facts", + ( + Facts(("alias",), {}, "team", True), + Facts(("alias",), {"alias": "target"}, "team", True), + Facts(("alias",), {"alias": "model_name_team_target"}, "team", False), + Facts(("alias",), {"alias": "model_name_team_target"}, "team", True), + Facts(("alias", "alias"), {"alias": "model_name_team_target"}, "team", True, (("model_name_team_target",), ())), + Facts(("alias",), {"alias": "model_name_None_target"}, None, True, ((),)), + ), +) +def test_aliases_preserve_deleted_targets_and_repeated_registry_reads( + policy: Callable[[Facts], Decision], facts: Facts +) -> None: + assert policy(facts) == reference(facts) diff --git a/tests/unit/_auth/test_team_override.py b/tests/unit/_auth/test_team_override.py new file mode 100644 index 00000000000..8e956b54e3e --- /dev/null +++ b/tests/unit/_auth/test_team_override.py @@ -0,0 +1,74 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.team_override import TeamOverrideDecision as Decision +from litellm._auth.team_override import TeamOverrideFacts as Facts +from litellm._auth.team_override import team_override +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Facts) -> Decision: + if facts.team is None: + return Decision("ReadTeam") + denied: Final = grouped(facts.team, merge=True) + if denied.kind == "Allow": + return Decision("Allow") + if denied.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if denied.kind != "Deny" or denied.object_type != "team": + return Decision("Deny", denied) + if facts.assigned_models is None: + return Decision("ReadAssignedGroups") + if not facts.assigned_models: + return Decision("Deny", denied) + if facts.assigned_scope is None: + return Decision("ReadAssignedScope") + assigned: Final = model_access(facts.assigned_scope) + if assigned.kind == "Allow": + return Decision("Allow") + return Decision("Deny", denied if assigned.kind == "Deny" else assigned) + + +_MODEL: Final = ModelAccessFacts(("model",), ("other",), object_type="team") +_TEAM: Final = Facts(GroupedModelScope(_MODEL)) +_ASSIGNED: Final = replace(_TEAM, assigned_models=("model",)) + + +@pytest.mark.parametrize("policy", (reference, team_override)) +@pytest.mark.parametrize( + "facts", + ( + Facts(), + _TEAM, + Facts(GroupedModelScope(replace(_MODEL, grants=()))), + Facts(GroupedModelScope(_MODEL, ("team-group",))), + Facts(GroupedModelScope(_MODEL, ("team-group",), ("model",))), + Facts(GroupedModelScope(_MODEL, ("team-group",), ("no-access",))), + Facts(GroupedModelScope(replace(_MODEL, object_type="key"))), + Facts(GroupedModelScope(replace(_MODEL, fallback_depth=10))), + replace(_TEAM, assigned_models=()), + _ASSIGNED, + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key", grants=("model",))), + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key")), + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key", fallback_depth=10)), + ), +) +def test_team_override_reads_groups_only_after_team_denial(policy: Callable[[Facts], Decision], facts: Facts) -> None: + assert policy(facts) == reference(facts) + + +def test_assigned_denial_keeps_team_error_envelope() -> None: + facts: Final = replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key")) + assert team_override(facts) == Decision("Deny", model_access(_MODEL)) + + +def test_bad_assigned_scope_is_terminal_not_team_denial() -> None: + with pytest.raises(NativeAuthFailure): + team_override(replace(_ASSIGNED, assigned_scope=replace(_MODEL, fallback_depth=-1))) diff --git a/tests/unit/_auth/test_unpriced_models.py b/tests/unit/_auth/test_unpriced_models.py new file mode 100644 index 00000000000..fac21104611 --- /dev/null +++ b/tests/unit/_auth/test_unpriced_models.py @@ -0,0 +1,106 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.unpriced_models import DeploymentPricing, PriceEntry, UnpricedModelFacts, unpriced_model +from litellm.proxy.auth.unpriced_model_adapter import price_entry +from tests.unit._auth.support.model_registry_reference import entry_has_priced_metric + +_BASE: Final = UnpricedModelFacts("model", True) +_ROW: Final = DeploymentPricing((), True) + + +@pytest.mark.parametrize( + ("entry", "expected"), + ( + ({}, False), + ({"input_cost_per_token": 0}, False), + ({"input_cost_per_token": -1}, False), + ({"cost_per_image": 0.1}, True), + ({"cost_per_image": True}, False), + ({"cost_per_image": "0.1"}, False), + ({"cost_per_image": {"small": 0, "large": 0.2}}, True), + ({"cost_per_image": {"large": True}}, False), + ({"cost_per_image": {"large": {"nested": 1}}}, False), + ({"tiered_pricing": None}, False), + ({"tiered_pricing": ()}, True), + ({"tiered_pricing": False}, True), + ({"maximum_tokens": 1}, False), + ({"input_cost_per_token": float("nan")}, False), + ({"input_cost_per_token": float("inf")}, True), + ), +) +def test_native_price_metrics_match_reference(entry: Mapping[str, object], expected: bool) -> None: + facts: Final = replace( + _BASE, + group_found=True, + deployments=(replace(_ROW, info_loaded=True, info=price_entry(entry)),), + cost_deployments=(), + ) + assert entry_has_priced_metric(entry) == expected + assert unpriced_model(facts).unpriced == (not expected) + + +@pytest.mark.parametrize( + ("facts", "kind", "index", "unpriced"), + ( + (replace(_BASE, model=None), "Complete", 0, False), + (replace(_BASE, model=""), "Complete", 0, False), + (replace(_BASE, router_present=False), "Complete", 0, False), + (_BASE, "ReadGroup", 0, False), + (replace(_BASE, group_found=False), "Complete", 0, False), + (replace(_BASE, group_found=True), "ReadDeployments", 0, False), + (replace(_BASE, group_found=True, deployments=(_ROW,)), "ReadInfo", 0, False), + ( + replace(_BASE, group_found=True, deployments=(replace(_ROW, parameter_keys=("input_cost_per_token",)),)), + "Complete", + 0, + False, + ), + ( + replace(_BASE, group_found=True, deployments=(replace(_ROW, parameter_keys=("tiered_pricing",)),)), + "Complete", + 0, + False, + ), + (replace(_BASE, group_found=True, deployments=()), "ReadCostDeployments", 0, False), + (replace(_BASE, group_found=True, deployments=(), cost_deployments=(_ROW,)), "ReadCost", 0, False), + ( + replace( + _BASE, + group_found=True, + deployments=(), + cost_deployments=(replace(_ROW, cost_keys=("input_cost_per_token",)),), + ), + "Complete", + 0, + False, + ), + ( + replace( + _BASE, + group_found=True, + deployments=(), + cost_deployments=(replace(_ROW, cost_keys=("cost_per_second",)),), + ), + "Complete", + 0, + True, + ), + ( + replace( + _BASE, group_found=True, deployments=(replace(_ROW, info_loaded=True, info=PriceEntry(False, ())), _ROW) + ), + "ReadInfo", + 1, + False, + ), + ), +) +def test_unpriced_reads_preserve_legacy_short_circuit_order( + facts: UnpricedModelFacts, kind: str, index: int, unpriced: bool +) -> None: + decision: Final = unpriced_model(facts) + assert (decision.kind, decision.index, decision.unpriced) == (kind, index, unpriced) diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 4692088e536..c6406b82c15 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -1369,7 +1369,7 @@ def personal_key_owner_teams(monkeypatch: pytest.MonkeyPatch, request: pytest.Fi monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) - monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) monkeypatch.setattr( auth_checks, "get_user_object", @@ -1428,7 +1428,9 @@ async def test_empty_owner_team_id_does_not_restrict_agent_access( assert owner_teams.unreadable is False assert tuple(team.team_id for team in owner_teams.teams) == ("team-open",) - assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + assert ( + await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + ) assert all(call.kwargs["team_id"] == "team-open" for call in team_loader.await_args_list) @@ -1671,7 +1673,6 @@ async def test_personal_key_agent_billing_keeps_team_controls_when_org_lookup_fa org_loader: Final = AsyncMock(side_effect=RuntimeError("org lookup failed")) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) - monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader) @@ -1717,7 +1718,9 @@ async def test_personal_key_agent_billing_keeps_team_controls_when_membership_lo AsyncMock(side_effect=RuntimeError("membership lookup failed")), ) - billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, "agent-b") + billed: Final = await AgentRequestHandler._personal_key_agent_billing_auth_for_team( # pyright: ignore[reportPrivateUsage] # exercise billing after owner-team selection has completed + personal_key_owner_teams, "agent-b", await _get_team_object("team-b"), strict=False + ) assert billed.team_id == "team-b" assert billed.team_rpm_limit == 5 @@ -1761,9 +1764,7 @@ async def test_personal_key_agent_billing_uses_fresh_membership_for_strict_polic ), ) - async def _get_team_membership( - user_id: str, team_id: str, **kwargs: object - ) -> LiteLLM_TeamMembership: + async def _get_team_membership(user_id: str, team_id: str, **kwargs: object) -> LiteLLM_TeamMembership: database_value: Final = bool(kwargs.get("check_db_only")) return LiteLLM_TeamMembership( user_id=user_id, @@ -1774,7 +1775,6 @@ async def test_personal_key_agent_billing_uses_fresh_membership_for_strict_polic team_loader: Final = AsyncMock(side_effect=_get_team_object) monkeypatch.setattr(auth_checks, "get_team_object", team_loader) monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(side_effect=_get_team_membership)) - monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) auth: Final = personal_key_owner_teams.model_copy(update={"requires_fresh_policy": strict}) @@ -2135,9 +2135,7 @@ async def test_strict_personal_key_agent_access_reads_owner_teams_from_the_datab return LiteLLM_TeamTable( team_id=team_id, object_permission_id=f"perm-{team_id}", - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id=f"perm-{team_id}", agents=["agent-a"] - ), + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id=f"perm-{team_id}", agents=["agent-a"]), ) monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) diff --git a/tests/unit/proxy/auth/budget_alert_reference.py b/tests/unit/proxy/auth/budget_alert_reference.py new file mode 100644 index 00000000000..a47fd364db6 --- /dev/null +++ b/tests/unit/proxy/auth/budget_alert_reference.py @@ -0,0 +1,42 @@ +from collections.abc import Mapping + +from pydantic import TypeAdapter, ValidationError + +from litellm._auth.budget_alerts import Facts + + +def reference_recipients(raw: object) -> list[str]: + if isinstance(raw, list): + return [email.strip() for email in raw if isinstance(email, str) and email.strip()] + if isinstance(raw, str): + return [email.strip() for email in raw.split(",") if email.strip()] + return [] + + +def reference_alert(facts: Facts, global_config: object, local_config: object) -> tuple[bool, object]: + if not facts.member and not (facts.limit is not None and facts.spend is not None and facts.spend > 0): + return False, None + if facts.member: + try: + config = TypeAdapter(Mapping[str, object]).validate_python(local_config) + except ValidationError: + config = {} + local_config = { + key: value for key, value in config.items() if key.isdigit() and len(key) <= 3 and 1 <= int(key) <= 100 + } + global_entries = {key: reference_recipients(value) for key, value in (global_config or {}).items()} + local_entries = {key: reference_recipients(value) for key, value in (local_config or {}).items()} + merged = { + key: list(dict.fromkeys(global_entries.get(key, []) + local_entries.get(key, []))) + for key in set(global_entries) | set(local_entries) + } + if merged: + if facts.member and facts.spend <= 0: + return False, None + minimum = min((int(key) for key in merged if key.isdigit()), default=None) + if minimum is None or facts.spend < facts.limit * (minimum / 100.0): + return False, None + return True, tuple((key, tuple(value)) for key, value in merged.items()) + return ( + not facts.member and facts.spend >= facts.limit * facts.default_fraction and facts.spend < facts.limit + ), None diff --git a/tests/unit/proxy/auth/budget_reference.py b/tests/unit/proxy/auth/budget_reference.py new file mode 100644 index 00000000000..a04a997ab46 --- /dev/null +++ b/tests/unit/proxy/auth/budget_reference.py @@ -0,0 +1,44 @@ +from collections.abc import Mapping +from typing import Final + +import litellm +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + entry_names_model, + resolve_model_budget, +) + + +async def model_allowed( + limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, + entity_type: Litellm_EntityType, + entity_id: str | None, + team_id: str | None, + budgets: Mapping[str, object] | None, + model: str, + message: str, +) -> bool: + if not budgets: + return True + resolved: Final = resolve_model_budget(model, budgets, limiter._model_access_groups, team_id) + if resolved is None: + return True + limit: Final = resolved.budget_config.max_budget + if limit is None or limit < 0: + return True + local: Final = await limiter._get_spend_for_model_budget(entity_type, entity_id, model, resolved) + current: Final = ( + local if local >= limit else await limiter._with_shared_spend(entity_type, entity_id, resolved, local) + ) + if current >= limit: + raise litellm.BudgetExceededError( + message=message + if entry_names_model(resolved.budget_model, model) + else f"{message}, budget_model={resolved.budget_model}", + current_cost=current, + max_budget=limit, + entity_type=entity_type.value, + entity_id=entity_id, + ) + return True diff --git a/tests/unit/proxy/auth/phase1_budget_auth_reference.py b/tests/unit/proxy/auth/phase1_budget_auth_reference.py new file mode 100644 index 00000000000..42c4673d874 --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_auth_reference.py @@ -0,0 +1,915 @@ +"""Policy frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9, with independently supplied raw IO.""" + +from __future__ import annotations + +import asyncio +import json +import math +import re +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final, Optional, Protocol + +from opentelemetry.trace import Span +from pydantic import ConfigDict, TypeAdapter + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import ( + MODEL_LIST_ROUTES, + CallInfo, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _dictable_table, # pyright: ignore[reportPrivateUsage] # data-only Prisma row projection +) +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + get_management_object_ttl, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.budget_repository import BudgetRepository +from litellm.router import Router +from litellm.router_api import RegistryReader +from litellm.types.router import DeploymentTypedDict + +from .phase1_budget_route_reference import RouteChecks +from .phase1_budget_throttle_reference import budget_throttle_percentage, should_throttle_budget_exceeded + + +class _RawCacheRead(Protocol): + async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... + + +def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: + return cache + + +def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None: # mutable-ok: Phase 1 public collection and identity contract + + cache: Final = llm_router.registry().zero_cost() + return cache if isinstance(cache, dict) else None + + +def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool: # mutable-ok: Phase 1 public collection and identity contract + + if model is None or llm_router is None: + return False + + model_list: Final = [model] if isinstance(model, str) else model + + zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) + + for model_name in model_list: + if zero_cost_cache is not None: + cached = zero_cost_cache.get(model_name) + if cached is not None: + if cached is False: + return False + continue + try: + model_group_info = llm_router.registry().group(model_name, include_hidden=True) + + if model_group_info is None: + verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + input_cost = model_group_info.input_cost_per_token + output_cost = model_group_info.output_cost_per_token + + if input_cost is None or output_cost is None: + verbose_proxy_logger.debug( + "Model %s has undefined cost (input: %s, output: %s), assuming it has cost", + model_name, + input_cost, + output_cost, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if input_cost > 0 or output_cost > 0: + verbose_proxy_logger.debug( + "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + safe_name = str(model_name).replace("\n", "").replace("\r", "") + served_deployments = _deployments_served_for(model_name, llm_router) + if not _is_cost_explicitly_configured(served_deployments): + verbose_proxy_logger.debug( + "Model %s has zero cost but no explicit cost " + "configuration in model_cost entry — treating as unknown " + "cost (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if _has_explicit_positive_cost(served_deployments): + verbose_proxy_logger.debug( + "Model %s routes to a deployment with an explicit positive per-token price (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if _has_ptu_flat_cost(served_deployments): + verbose_proxy_logger.debug( + "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " + "rate is not a free model (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + verbose_proxy_logger.debug( + "Model %s has zero cost explicitly configured (input: %s, output: %s)", + safe_name, + input_cost, + output_cost, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = True + + except Exception as e: + verbose_proxy_logger.debug("Error checking cost for model %s: %s, assuming it has cost", model_name, e) + return False + + return True + + +_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: + + registry: Final = llm_router.registry() + routed_group: Final = registry.resolve(model) + if registry.resolve(routed_group) == routed_group: + return registry.deployments(routed_group) or () + named: Final = registry.deployments(model, source="alias") + if named: + return named + return tuple( + DeploymentTypedDict(**deployment) # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + for deployment in registry.patterns().get_deployments_by_pattern(model=routed_group) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + ) + + +def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_info = deployment.get("model_info") or _NO_MODEL_INFO # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return True + return False + + +def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_id is None: + continue + raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: + return True + return False + + +def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_id is None: + continue + raw_entry = as_str_mapping(litellm.model_cost.get(model_id)) or _EMPTY_COST_ENTRY # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + if ( + _is_positive_cost(raw_entry.get("input_cost_per_token")) + or _is_positive_cost(raw_entry.get("output_cost_per_token")) + or _is_positive_cost(_deployment_query_price(deployment["litellm_params"], raw_entry)) + ): + return True + return False + + +def _deployment_query_price(litellm_params: Mapping[str, object], cost_entry: Mapping[str, object]) -> float | None: + provider: Final = litellm_params.get("custom_llm_provider") or str(litellm_params.get("model")).partition("/")[0] + return decisions_request_price({**cost_entry, "litellm_provider": provider}, decisions_pricing_fields(cost_entry)) + + +_EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +MODEL_DISCOVERY_ROUTES: Final = frozenset( + { + *MODEL_LIST_ROUTES, + "/model/info", + "/v1/model/info", + "/v2/model/info", + "/model_group/info", + "/utils/model_info", + } +) + + +BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset( + { + "/health", + "/health/services", + "/health/test_connection", + "/auto_router/test_routing", + } +) + + +MCP_DISCOVERY_ROUTES: Final = frozenset({"/mcp-rest/tools/list", "/v1/mcp/tools"}) + + +MCP_ZERO_SPEND_JSONRPC_METHODS: Final = frozenset({"initialize", "notifications/initialized", "ping", "tools/list"}) + + +MCP_TOOL_CALL_ROUTES: Final = frozenset({"/mcp/tools/call", "/mcp-rest/tools/call"}) + + +def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) -> bool: + if route in MCP_DISCOVERY_ROUTES: + return True + if route in MCP_TOOL_CALL_ROUTES or not (route == "/mcp" or route.startswith("/mcp/")): + return False + method: Final = request_body.get("method") + return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS + + +def route_skips_budget_checks(route: str) -> bool: + if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: + return False + if route in MODEL_DISCOVERY_ROUTES: + return True + if RouteChecks.is_llm_api_route(route=route): + return False + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + + +def request_skips_budget_checks( + route: str, + model: str | list[str] | None, # mutable-ok: Phase 1 public collection and identity contract + llm_router: Router | None, + *, + request_body: Mapping[str, object] | None = None, +) -> bool: + if route_skips_budget_checks(route=route): + return True + if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body): + return True + return is_model_cost_zero(model=model, llm_router=llm_router) + + +KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" + + +def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: + + if key_metadata is None: + return None + budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD) + return budget_id if isinstance(budget_id, str) and budget_id != "" else None + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def get_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + budget_id: str | None = None, +) -> LiteLLM_BudgetTable | None: + + default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id + if prisma_client is None or default_budget_id is None: + return None + + cache_key: Final = f"default_end_user_budget:{default_budget_id}" + + cached_budget: Final = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_BudgetTable, + ) + if cached_budget is not None: + return cached_budget + + try: + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( + where={"budget_id": default_budget_id} + ) + + if budget_record is None: + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", + default_budget_id.replace("\r", "").replace("\n", ""), + ) + return None + + _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value=_budget_obj, + model_type=LiteLLM_BudgetTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + + return _budget_obj + + except Exception as e: + verbose_proxy_logger.error("Error fetching default end user budget: %s", e) + return None + + +async def get_team_member_default_budget( + budget_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + raise_on_lookup_error: bool = False, +) -> LiteLLM_BudgetTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_team_member_default_budget") + + +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + if member_budget is not None and member_budget.max_budget is not None: + return _effective_member_max(member_budget, now) + if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: + return None + return team_default_budget.max_budget + ( + _active_member_increase(member_budget, now) if member_budget is not None else 0.0 + ) + + +async def resolve_default_end_user_budget( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + key_end_user_budget_id: str | None, + parent_otel_span: Span | None = None, +) -> LiteLLM_BudgetTable | None: + + if key_end_user_budget_id is not None: + key_budget: Final = await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + budget_id=key_end_user_budget_id, + ) + if key_budget is not None: + return key_budget + + if litellm.max_end_user_budget_id is None: + return None + + return await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + + +async def _apply_default_budget_to_end_user( + end_user_obj: LiteLLM_EndUserTable, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable: + + if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None: + return end_user_obj + + if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None: + return end_user_obj + + default_budget: Final = await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + + if default_budget is None: + return end_user_obj + + verbose_proxy_logger.debug( + "Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id + ) + return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget})) + + +async def check_end_user_budget( + end_user_obj: LiteLLM_EndUserTable, + route: str, +) -> None: + + if RouteChecks.is_info_route(route): + return + + if end_user_obj.litellm_budget_table is None: + return + + end_user_budget: Final = end_user_obj.litellm_budget_table.max_budget + if end_user_budget is None: + return + + from litellm.proxy.proxy_server import get_current_spend + + end_user_spend: Final = await get_current_spend( + counter_key=f"spend:end_user:{end_user_obj.user_id}", + fallback_spend=end_user_obj.spend or 0.0, + max_budget=end_user_budget, + fallback_authoritative=True, + ) + if end_user_spend > end_user_budget: + raise litellm.BudgetExceededError( + current_cost=end_user_spend, + max_budget=end_user_budget, + message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", + entity_type=Litellm_EntityType.END_USER.value, + entity_id=end_user_obj.user_id, + ) + + +async def _load_end_user_restricted_registry( + prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache +) -> frozenset[str] | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: _load_end_user_restricted_registry") + + +async def _end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, + key_end_user_budget_id: str | None = None, +) -> bool: + + if ( + litellm.max_end_user_budget_id is not None + or key_end_user_budget_id is not None + or litellm.validate_end_user_id_in_db + or token_end_user_max_budget is not None + ): + return False + + registry: Final = await _load_end_user_restricted_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return registry is not None and end_user_id not in registry + + +async def get_end_user_object( + end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + route: str | None = "", + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + token_end_user_max_budget: float | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_end_user_object") + + +_END_USER_VALIDATION_NEGATIVE_TTL: Final = 60 + + +_END_USER_VALIDATION_POSITIVE_TTL: Final = 300 + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def resolve_and_validate_end_user_id( + raw_end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", + key_end_user_budget_id: str | None = None, +) -> str | None: + + if raw_end_user_id is None: + return None + if not litellm.validate_end_user_id_in_db: + return raw_end_user_id + if prisma_client is None: + return raw_end_user_id + + has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None + cache_key: Final = f"end_user_validation:{raw_end_user_id}" + cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) + if cached == "valid": + return raw_end_user_id + if cached == "invalid": + return raw_end_user_id if has_default_budget else None + + is_valid: Final = await _end_user_id_exists_in_db( + end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value="valid" if is_valid else "invalid", + ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL), + ) + + return raw_end_user_id if is_valid or has_default_budget else None + + +async def _end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", +) -> bool: + + try: + end_user_obj: Final = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", e) + + try: + user_obj: Final = await get_user_object( + user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=False, + user_email=end_user_id if "@" in end_user_id else None, + ) + if user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", e) + + return False + + +async def get_team_membership( + user_id: str, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + *, + check_db_only: bool = False, +) -> Optional["LiteLLM_TeamMembership"]: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_team_membership") + + +async def get_user_object( + user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + user_id_upsert: bool, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + sso_user_id: str | None = None, + user_email: str | None = None, + check_db_only: bool | None = None, +) -> LiteLLM_UserTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_user_object") + + +_MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?") + + +def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: + + pct: Final = budget_throttle_percentage() + if pct is None or not should_throttle_budget_exceeded(valid_token): + return False + valid_token.budget_throttle_pct = pct # rebind-ok: preserve Phase 1 mutation order and identity + return True + + +def send_token_budget_alert( + valid_token: UserAPIKeyAuth, + spend: float, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj is not None else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) + ) + + +def key_budget_exceeded_error( + valid_token: UserAPIKeyAuth, spend: float, max_budget: float +) -> litellm.BudgetExceededError: + + key_label: Final = valid_token.key_alias or "key" + key_descriptor: Final = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) + else key_label + ) + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {max_budget}", + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + +async def virtual_key_max_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + + if valid_token.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + max_budget=valid_token.max_budget, + ) + send_token_budget_alert(valid_token, spend, proxy_logging_obj, user_obj) + + if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: + if _apply_budget_exceeded_throttle(valid_token): + return + raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget) + + +async def _virtual_key_multi_budget_check( + valid_token: UserAPIKeyAuth, +): + + if not valid_token.budget_limits: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.budget_limits: # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + w: dict = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportMissingTypeArgument, reportUnknownVariableType, reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_entity_type="Key", + window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_start=get_budget_window_start(w), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + message=( + f"ExceededBudget: Key over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + +async def team_max_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + + if team_object is not None and team_object.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:team:{team_object.team_id}", + fallback_spend=team_object.spend or 0.0, + max_budget=team_object.max_budget, + ) + + if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: + if valid_token: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=team_object.max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.TEAM, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=spend, + max_budget=team_object.max_budget, + message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +async def _team_multi_budget_check( + team_object: LiteLLM_TeamTable | None, +): + + if team_object is None or not team_object.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in team_object.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportMissingTypeArgument] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_entity_type="Team", + window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_start=get_budget_window_start(w), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + message=( + f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +class ModelMaxBudgetHolder(Protocol): + @property + def model_max_budget(self) -> Mapping[str, object] | None: ... + + +def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | None: + return holder.model_max_budget + + +def _active_member_increase(self: LiteLLM_BudgetTable, now: datetime) -> float: + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + return 0.0 + expiry: Final = ( + self.temp_budget_expiry.replace(tzinfo=timezone.utc) + if self.temp_budget_expiry.tzinfo is None + else self.temp_budget_expiry + ) + return 0.0 if expiry <= now else self.temp_budget_increase + + +def _effective_member_max(self: LiteLLM_BudgetTable, now: datetime) -> float | None: + if self.max_budget is None: + return None + return self.max_budget + _active_member_increase(self, now) + + +_WINDOW_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +def _coerce_window(window: object) -> Mapping[str, object]: + if isinstance(window, Mapping): + return window # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if isinstance(window, str): + try: + return _WINDOW_OBJECT.validate_python(json.loads(window)) + except Exception: + return {} + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return {} + dumped: Final[object] = model_dump() + return dumped if isinstance(dumped, Mapping) else {} # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + + +def get_budget_window_start(window: object) -> datetime | None: + window_dict: Final = _coerce_window(window) + budget_duration: Final = window_dict.get("budget_duration") + if budget_duration is None: + return None + try: + duration_seconds: Final = duration_in_seconds(str(budget_duration)) + except Exception: + return None + + reset_at = _coerce_datetime(window_dict.get("reset_at")) # rebind-ok: preserve Phase 1 mutation order and identity + if reset_at is None: + return datetime.now(timezone.utc) - timedelta(seconds=duration_seconds) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) # rebind-ok: preserve Phase 1 mutation order and identity + return reset_at - timedelta(seconds=duration_seconds) + + +def _coerce_datetime(value: object) -> datetime | None: + if value is None: + return None + if isinstance(value, datetime): + return value + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return None + + +__all__ = ( + "_raw_cache", + "_get_router_zero_cost_cache", + "_deployments_served_for", + "_has_ptu_flat_cost", + "_is_cost_explicitly_configured", + "_has_explicit_positive_cost", + "_deployment_query_price", + "_is_positive_cost", + "_apply_default_budget_to_end_user", + "_load_end_user_restricted_registry", + "_end_user_is_known_unrestricted", + "_end_user_id_exists_in_db", + "_apply_budget_exceeded_throttle", + "_virtual_key_multi_budget_check", + "_team_multi_budget_check", + "_active_member_increase", + "_effective_member_max", + "_coerce_window", + "_coerce_datetime", +) diff --git a/tests/unit/proxy/auth/phase1_budget_context_reference.py b/tests/unit/proxy/auth/phase1_budget_context_reference.py new file mode 100644 index 00000000000..46900a23d09 --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_context_reference.py @@ -0,0 +1,352 @@ +"""Policy frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9, with independently supplied raw IO.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from typing import Final, Protocol, cast + +from opentelemetry.trace import Span + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from .phase1_budget_auth_reference import ( + get_end_user_object, + get_key_end_user_budget_id, + get_team_member_default_budget, + get_team_membership, + model_max_budget_of, + resolve_default_end_user_budget, +) + + +class _TeamModelBudgetLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + +class _EndUserModelBudgetLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + +async def _check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _TeamModelBudgetLimiter, + models: list[str], # mutable-ok: Phase 1 public collection and identity contract +) -> None: + + team_model_max_budget: Final = valid_token.team_model_max_budget + if valid_token.team_id is None or not team_model_max_budget: + return + key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + for model_name in models: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + model=model_name, + ) + + +async def _end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + budget_table: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else ( + await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if end_user_id and prisma_client is not None + else None + ) + ) + return model_max_budget_of(budget_table) if budget_table is not None else None + + +async def _check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: _EndUserModelBudgetLimiter, + models: Sequence[str], + team_id: str | None, +) -> None: + if not end_user_id or not end_user_model_max_budget: + return + for model_name in models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=model_name, + team_id=team_id, + ) + + +def _apply_budget_limits_to_end_user_params( + end_user_params: dict, # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + budget_info: LiteLLM_BudgetTable, + end_user_id: str | None, +) -> None: + + if budget_info.tpm_limit is not None: + end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.rpm_limit is not None: + end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.tpd_limit is not None: + end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.max_budget is not None: + end_user_params["end_user_max_budget"] = budget_info.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.model_max_budget is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + + +def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + valid_token.end_user_id = end_user_params.get("end_user_id") # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + if end_user_params.get("end_user_tpm_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_rpm_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_tpd_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("allowed_model_region") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.allowed_model_region = end_user_params["allowed_model_region"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_model_max_budget") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] # rebind-ok: preserve Phase 1 mutation order and identity + return valid_token + + +async def _load_team_member_default_model_budget( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool, + proxy_logging_obj: ProxyLogging | None, +) -> None: + team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None + team_metadata: Final[Mapping[str, object] | None] = ( + cast( # cast-ok: team metadata is a string-keyed JSON object + Mapping[str, object], team_metadata_value + ) + if isinstance(team_metadata_value, Mapping) + else None + ) + team_member_budget_id: Final[object | None] = ( + team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None + ) + if ( + not isinstance(user_api_key_auth_obj.user_id, str) + or not isinstance(user_api_key_auth_obj.team_id, str) + or prisma_client is None + ): + return + + try: + membership: Final = await get_team_membership( + user_id=user_api_key_auth_obj.user_id, + team_id=user_api_key_auth_obj.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + member_budget_row: Final = membership.litellm_budget_table if membership is not None else None + member_budget_value: Final[object] = ( + getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None + ) + if ( + membership is not None + and membership.budget_id is not None + and membership.budget_id != team_member_budget_id + and isinstance(member_budget_value, Mapping) + and member_budget_value + ): + user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream + member_budget_value + ) + return + if not isinstance(team_member_budget_id, str): + return + team_member_default_budget: Final = await get_team_member_default_budget( + budget_id=team_member_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + raise_on_lookup_error=True, + ) + except Exception: + if not skip_budget_checks: + raise + verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True) + return + + default_budget_value: Final[object] = ( + getattr(team_member_default_budget, "model_max_budget", None) + if team_member_default_budget is not None + else None + ) + user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream + default_budget_value if isinstance(default_budget_value, Mapping) else None + ) + + +async def _apply_key_end_user_default_budget_to_token( + valid_token: UserAPIKeyAuth, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + keep_token_limits: bool, +) -> None: + + default_budget: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + ) + if default_budget is None: + return + + if not keep_token_limits or valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_tpm_limit is None: + valid_token.end_user_tpm_limit = default_budget.tpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_rpm_limit is None: + valid_token.end_user_rpm_limit = default_budget.rpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_tpd_limit is None: + valid_token.end_user_tpd_limit = default_budget.tpd_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_model_max_budget is None: + valid_token.end_user_model_max_budget = default_budget.model_max_budget # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + +def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): # pyright: ignore[reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values + valid_token_metadata: Final = valid_token.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: + expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + if expiry.tzinfo is None: + expiry = expiry.replace(tzinfo=timezone.utc) # rebind-ok: preserve Phase 1 mutation order and identity + if expiry > datetime.now(timezone.utc): + return valid_token_metadata["temp_budget_increase"] # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + return None + + +def _update_key_budget_with_temp_budget_increase( + valid_token: UserAPIKeyAuth, +) -> UserAPIKeyAuth: + if valid_token.max_budget is None: + return valid_token + temp_budget_increase: Final = _get_temp_budget_increase(valid_token) # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if not temp_budget_increase: + return valid_token + return valid_token.model_copy(update={"max_budget": valid_token.max_budget + temp_budget_increase}) + + +async def _lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Span | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +): + + end_user_object = None # rebind-ok: preserve Phase 1 mutation order and identity + key_end_user_budget_id: Final = get_key_end_user_budget_id(valid_token.metadata) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + try: + end_user_object = await get_end_user_object( # rebind-ok: preserve Phase 1 mutation order and identity + end_user_id=valid_token.end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + token_end_user_max_budget=valid_token.end_user_max_budget, + key_end_user_budget_id=key_end_user_budget_id, + ) + if end_user_object is not None: + end_user_params = { # rebind-ok: preserve Phase 1 mutation order and identity + "end_user_id": valid_token.end_user_id, + "allowed_model_region": end_user_object.allowed_model_region, + } + if end_user_object.litellm_budget_table is not None: + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=end_user_object.litellm_budget_table, + end_user_id=valid_token.end_user_id or "", + ) + valid_token = update_valid_token_with_end_user_params( # rebind-ok: preserve Phase 1 mutation order and identity + valid_token=valid_token, end_user_params=end_user_params + ) + elif key_end_user_budget_id is not None or litellm.max_end_user_budget_id is not None: + default_budget: Final = await resolve_default_end_user_budget( + prisma_client=prisma_client, # pyright: ignore[reportArgumentType] # frozen Phase 1 forwards the optional raw client + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if default_budget is not None: + end_user_params = {"end_user_id": valid_token.end_user_id} # rebind-ok: preserve Phase 1 mutation order and identity + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=default_budget, + end_user_id=valid_token.end_user_id or "", + ) + valid_token = update_valid_token_with_end_user_params( # rebind-ok: preserve Phase 1 mutation order and identity + valid_token=valid_token, end_user_params=end_user_params + ) + if valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + except Exception as e: + if isinstance(e, litellm.BudgetExceededError): + raise e + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) + return valid_token, end_user_object + + +__all__ = ( + "_check_team_model_budget", + "_end_user_model_max_budget", + "_check_end_user_model_budget", + "_apply_budget_limits_to_end_user_params", + "_load_team_member_default_model_budget", + "_apply_key_end_user_default_budget_to_token", + "_get_temp_budget_increase", + "_update_key_budget_with_temp_budget_increase", + "_lookup_end_user_and_apply_budget", +) diff --git a/tests/unit/proxy/auth/phase1_budget_lookup_reference.py b/tests/unit/proxy/auth/phase1_budget_lookup_reference.py new file mode 100644 index 00000000000..6d11b7f398c --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_lookup_reference.py @@ -0,0 +1,341 @@ +"""Unchanged Phase 1 budget lookup policy with raw IO supplied by tests.""" + +from __future__ import annotations + +import asyncio +import math +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + CallInfo, + Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.utils import get_utc_datetime + +from .phase1_budget_auth_reference import ( + get_team_member_default_budget, + get_team_membership, + resolve_effective_team_member_budget, +) + +TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_emails" +_TEAM_MEMBER_ALERT_CONFIG_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +async def get_org_object( + org_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + include_budget_table: bool, +) -> LiteLLM_OrganizationTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_org_object") + + +def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """Parse emails from a list or comma-separated string.""" + if isinstance(raw, list): + return [e.strip() for e in raw if isinstance(e, str) and e.strip()] + elif isinstance(raw, str): + return [e.strip() for e in raw.split(",") if e.strip()] + return [] + + +def _merge_budget_alert_email_configs( + global_cfg: Mapping[str, str | Sequence[object] | None] | None, + per_key_cfg: Mapping[str, str | Sequence[object] | None] | None, +) -> dict[str, list[str]] | None: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """ + Per-threshold additive merge: each threshold's recipient list is the union + of global + per-key entries (deduped, global-first ordering). Missing + thresholds on one side are inherited from the other. + """ + global_cfg_normalized: Final = _normalize_alert_emails(global_cfg) + per_key_cfg_normalized: Final = _normalize_alert_emails(per_key_cfg) + if not global_cfg_normalized and (not per_key_cfg_normalized): + return None + thresholds: Final = set(global_cfg_normalized) | set(per_key_cfg_normalized) + return { + t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds + } + + +def _is_valid_alert_threshold_pct(pct: str) -> bool: + return pct.isdigit() and len(pct) <= 3 and (1 <= int(pct) <= 100) + + +def _alert_recipients(raw: object) -> Sequence[str] | None: + if isinstance(raw, (str, Sequence)): + return _parse_email_list(raw) # pyright: ignore[reportUnknownArgumentType] # preserve unvalidated Phase 1 oracle values and call shape + return None + + +def _valid_alert_threshold_config(raw_config: object) -> Mapping[str, str | Sequence[object] | None] | None: + try: + config: Final = _TEAM_MEMBER_ALERT_CONFIG_ADAPTER.validate_python(raw_config) + except ValidationError: + return None + return MappingProxyType( + {pct: _alert_recipients(emails) for pct, emails in config.items() if _is_valid_alert_threshold_pct(pct)} + ) + + +def team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + raw_config: Final = team_metadata.get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY) if team_metadata else None + alert_email_config: Final = _merge_budget_alert_email_configs( + global_cfg=None, per_key_cfg=_valid_alert_threshold_config(raw_config) + ) + if not alert_email_config or spend <= 0: + return + min_pct: Final = min(int(pct) for pct in alert_email_config) + if spend < max_budget * (min_pct / 100.0): + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + max_budget_alert_emails=alert_email_config, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info)) + + +async def _check_team_member_budget( # pyright: ignore[reportUnusedFunction] # preserve unvalidated Phase 1 oracle values and call shape + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +): + """Check if team member is over their max budget within the team.""" + if ( + team_object is not None + and team_object.team_id is not None # pyright: ignore[reportUnnecessaryComparison] # preserve unvalidated Phase 1 oracle values and call shape + and (valid_token is not None) + and (valid_token.user_id is not None) + ): + if not team_membership_loaded: + team_membership = await get_team_membership( # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None + now: Final = get_utc_datetime() + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + team_default_budget: Final = ( + await get_team_member_default_budget( + budget_id=default_budget_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache + ) + if (member_budget_row is None or member_budget_row.max_budget is None) + and isinstance(default_budget_id, str) + else None + ) + team_member_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget_row, team_default_budget=team_default_budget, now=now + ) + if team_member_budget is not None: + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + from litellm.proxy.proxy_server import get_current_spend + + team_member_spend = await get_current_spend( # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + fallback_spend=team_member_spend, + max_budget=team_member_budget, + ) + if not math.isfinite(team_member_budget): + return + team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=team_object.metadata, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # preserve unvalidated Phase 1 oracle values and call shape + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=team_member_spend, + max_budget=team_member_budget, + ) + if team_member_spend >= team_member_budget: + raise litellm.BudgetExceededError( + current_cost=team_member_spend, + max_budget=team_member_budget, + message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{team_object.team_id}", + ) + + +async def _team_soft_budget_check( # pyright: ignore[reportUnusedFunction] # preserve unvalidated Phase 1 oracle values and call shape + team_object: LiteLLM_TeamTable | None, valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging +): + """ + Triggers a budget alert if the team is over it's soft budget. + """ + if ( + team_object is not None + and team_object.soft_budget is not None + and (team_object.spend is not None) + and (team_object.spend >= team_object.soft_budget) + ): + verbose_proxy_logger.debug( + "Crossed Soft Budget for team %s, spend %s, soft_budget %s", + team_object.team_id, + team_object.spend, + team_object.soft_budget, + ) + if valid_token: + alert_emails: list[str] | None = None # mutable-ok: unchanged Phase 1 collection contract in independent oracle # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if team_object.metadata is not None and isinstance(team_object.metadata, dict): # pyright: ignore[reportUnknownMemberType, reportUnnecessaryIsInstance] # preserve unvalidated Phase 1 oracle values and call shape + soft_budget_alert_emails: Final = team_object.metadata.get("soft_budget_alerting_emails") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + if soft_budget_alert_emails is not None: + if isinstance(soft_budget_alert_emails, list): + alert_emails = [ # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() # pyright: ignore[reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + ] + elif isinstance(soft_budget_alert_emails, str): + alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if alert_emails: + alert_emails = [email for email in alert_emails if email] # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + else: + alert_emails = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if alert_emails is None or len(alert_emails) == 0: + verbose_proxy_logger.debug( + "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails", + team_object.team_id, + ) + return + call_info: Final = CallInfo( + token=valid_token.token, + spend=team_object.spend, + max_budget=team_object.max_budget, + soft_budget=team_object.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.TEAM, + alert_emails=alert_emails, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="soft_budget", user_info=call_info)) + + +async def organization_max_budget_check( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Check if the organization is over its max budget. + + This function checks the organization budget using: + 1. First, tries to use valid_token.org_id (if key has organization_id set) + 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) + + This ensures organization budget checks work even when keys don't have organization_id + set directly, as long as their team belongs to an organization. + + Raises: + BudgetExceededError if the organization is over its max budget. + Triggers a budget alert if the organization is over its max budget. + """ + if valid_token is None or prisma_client is None: + return + org_id: str | None = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if valid_token.org_id is not None: + org_id = valid_token.org_id # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + elif team_object is not None and team_object.organization_id is not None: + org_id = team_object.organization_id # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_id is None: + return + try: + org_table: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + include_budget_table=True, + ) + except Exception: + return + if org_table is None: + return + carry_organization_budget_state(valid_token=valid_token, org_table=org_table) + org_max_budget: float | None = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_table.litellm_budget_table is not None: + org_max_budget = org_table.litellm_budget_table.max_budget # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_max_budget is None: + return + from litellm.proxy.proxy_server import get_current_spend + + org_spend: Final = await get_current_spend( + counter_key=f"spend:org:{org_id}", fallback_spend=org_table.spend or 0.0, max_budget=org_max_budget + ) + if math.isfinite(org_max_budget) and org_spend >= org_max_budget: + call_info: Final = CallInfo( + token=valid_token.token, + spend=org_spend, + max_budget=org_max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=org_id, + event_group=Litellm_EntityType.ORGANIZATION, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="organization_budget", user_info=call_info)) + raise litellm.BudgetExceededError( + current_cost=org_spend, + max_budget=org_max_budget, + message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", + entity_type=Litellm_EntityType.ORGANIZATION.value, + entity_id=org_id, + ) + + +def _normalize_alert_emails(cfg: Mapping[str, str | Sequence[object] | None] | None) -> dict[str, list[str]]: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. + + Values may legitimately arrive as list, comma-separated string, or None + from YAML/metadata; _parse_email_list tolerates all three. + """ + if not cfg: + return {} + return {k: _parse_email_list(v) for k, v in cfg.items()} diff --git a/tests/unit/proxy/auth/phase1_budget_route_reference.py b/tests/unit/proxy/auth/phase1_budget_route_reference.py new file mode 100644 index 00000000000..f409731dc7e --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_route_reference.py @@ -0,0 +1,151 @@ +"""Route classification frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9 for independent budget references.""" + +from __future__ import annotations + +import re +from collections.abc import Collection +from typing import Final + +from litellm.proxy._types import LiteLLMRoutes + +_EU_JEV_DECISION_ROUTE: Final = "/eu_jev/v1/systemone" + + +class RouteChecks: + @staticmethod + def is_llm_api_route(route: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + if route in LiteLLMRoutes.openai_routes.value: + return True + + if route in LiteLLMRoutes.anthropic_routes.value: + return True + + if route in LiteLLMRoutes.google_routes.value: + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): + return True + + for openai_route in LiteLLMRoutes.openai_routes.value: + if "{" in openai_route: + if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): + return True + + if RouteChecks.is_wildcard_pattern(pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): + return True + + for google_route in LiteLLMRoutes.google_routes.value: + if "{" in google_route: + if RouteChecks._route_matches_pattern(route=route, pattern=google_route): + return True + + for anthropic_route in LiteLLMRoutes.anthropic_routes.value: + if "{" in anthropic_route: + if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): + return True + + if RouteChecks._is_azure_openai_route(route=route): + return True + + for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): + return True + return RouteChecks.is_builtin_eu_jev_decision_route(route=route) + + @staticmethod + def is_builtin_eu_jev_decision_route(route: str) -> bool: + if route.rstrip("/") != _EU_JEV_DECISION_ROUTE: + return False + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass_through_endpoints imports the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + route=_EU_JEV_DECISION_ROUTE + ) + + @staticmethod + def is_info_route(route: str) -> bool: + + return route in LiteLLMRoutes.info_routes.value + + @staticmethod + def _is_azure_openai_route(route: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" + engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" + + if re.match(deployment_pattern, route) or re.match(engine_pattern, route): + return True + return False + + @staticmethod + def _route_matches_pattern(route: str, pattern: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + def _placeholder_to_regex(match: re.Match[str]) -> str: + placeholder: Final = match.group(0).strip("{}") + if not placeholder.endswith(":path"): + return r"[^/]+" + + return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" + + pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) # rebind-ok: preserve Phase 1 mutation order and identity + + pattern = f"^{pattern}$" # rebind-ok: preserve Phase 1 mutation order and identity + if re.match(pattern, route): + return True + return False + + @staticmethod + def is_wildcard_pattern(pattern: str) -> bool: + + return pattern.endswith("*") + + @staticmethod + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + + if pattern.endswith("*"): + prefix: Final = pattern[:-1] + return route.startswith(prefix) + else: + return route == pattern + + @staticmethod + def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: + + if route in allowed_routes: + return True + + if any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in allowed_routes + if RouteChecks.is_wildcard_pattern(pattern=allowed_route) + ): + return True + + if any( + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes + ): + return True + + return False + + +__all__ = () diff --git a/tests/unit/proxy/auth/phase1_budget_throttle_reference.py b/tests/unit/proxy/auth/phase1_budget_throttle_reference.py new file mode 100644 index 00000000000..5b3d8469d5a --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_throttle_reference.py @@ -0,0 +1,45 @@ +""" +Throttle a key after it exceeds its own ``max_budget`` instead of blocking it. + +When a key opts in via ``throttle_on_budget_exceeded`` and a global +``budget_exceeded_throttle_percentage`` is configured, an over-budget key keeps +serving requests but at a reduced TPM/RPM (the configured percentage of its +configured limits). The decision (over budget + opted in) is made once during +auth; the scaling is recomputed from the key's original limits on every request +so it never compounds across requests. +""" + +import math +from typing import Final + +import litellm +from litellm.proxy._types import UserAPIKeyAuth + + +def budget_throttle_percentage() -> float | None: + + pct: Final = litellm.budget_exceeded_throttle_percentage + if not isinstance(pct, (int, float)) or isinstance(pct, bool): + return None + if not 0 < pct <= 1: + return None + return float(pct) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + + if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return False + if valid_token.tpm_limit is None and valid_token.rpm_limit is None: + return False + return budget_throttle_percentage() is not None + + +def throttled_limit(limit: int | None, pct: float | None) -> int | None: + + if limit is None or pct is None: + return limit + return max(1, math.floor(limit * pct)) + + +__all__ = () diff --git a/tests/unit/proxy/auth/route_role_reference.py b/tests/unit/proxy/auth/route_role_reference.py new file mode 100644 index 00000000000..8e0ade7c067 --- /dev/null +++ b/tests/unit/proxy/auth/route_role_reference.py @@ -0,0 +1,81 @@ +from collections.abc import Iterable +from itertools import chain +from typing import Final + +from fastapi import Request +from starlette.routing import compile_path + +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.route_checks import RouteChecks + + +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + try: + return LitellmUserRoles(user_role) + except ValueError: + return LitellmUserRoles.INTERNAL_USER + + +def get_user_role(user_obj: LiteLLM_UserTable | None) -> LitellmUserRoles | None: + return None if user_obj is None else effective_user_role(user_obj.user_role) + + +def is_user_proxy_admin(user_obj: LiteLLM_UserTable | None) -> bool: + return user_obj is not None and user_obj.user_role is not None and user_obj.user_role == "proxy_admin" + + +def is_api_route_allowed( + route: str, request: Request, request_data: dict[str, object], valid_token: UserAPIKeyAuth | None, # mutable-ok: preserve the original request identity at the shared route effect + user_obj: LiteLLM_UserTable | None = None, +) -> bool: + role: Final = get_user_role(user_obj) + if valid_token is None: + raise Exception("Invalid proxy server token passed. valid_token=None.") + if not is_user_proxy_admin(user_obj): + RouteChecks.non_proxy_admin_allowed_routes_check( # pyright: ignore[reportUnknownMemberType] # Phase 1 public effect signature + user_obj=user_obj, _user_role=role, route=route, request=request, + request_data=request_data, valid_token=valid_token, + ) + return True + + +def allowed_routes_list_check(user_route: str, allowed_routes: Iterable[object]) -> bool: + for allowed in allowed_routes: + if allowed in LiteLLMRoutes.__members__: + for template in LiteLLMRoutes[allowed].value: # pyright: ignore[reportArgumentType, reportUnknownVariableType, reportUnknownMemberType] # reference intentionally preserves the untyped Phase 1 input contract + regex, _, _ = compile_path(template) # pyright: ignore[reportUnknownArgumentType] # enum values are legacy route collections + if regex.match(user_route): + return True + elif allowed.endswith("*"): # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # invalid legacy entries must raise their original AttributeError + if user_route.startswith(allowed[:-1]): # pyright: ignore[reportIndexIssue, reportUnknownArgumentType] # preserve the reference's raw entry operations + return True + elif user_route == allowed: + return True + return False + + +def allowed_routes_check(user_role: LitellmUserRoles, user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth) -> bool: + if user_role == LitellmUserRoles.PROXY_ADMIN: + return allowed_routes_list_check(user_route, litellm_proxy_roles.admin_allowed_routes) + if user_role == LitellmUserRoles.TEAM: + configured: Final = litellm_proxy_roles.team_allowed_routes + return allowed_routes_list_check( + user_route, ("openai_routes", "info_routes") if configured is None else configured # pyright: ignore[reportUnnecessaryComparison] # retain Phase 1 constructed None config behavior + ) + return False + + +def allowed_route_check_inside_route(auth: UserAPIKeyAuth, requested_user_id: str | None) -> bool: + return auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) or ( + requested_user_id is not None and auth.user_id is not None and auth.user_id == requested_user_id + ) + + +def get_actual_routes(allowed_routes: Iterable[object]) -> list[object]: # mutable-ok: public Phase 1 projection + def expanded(raw: object) -> tuple[object, ...]: + try: + return tuple(LiteLLMRoutes[raw].value) # pyright: ignore[reportArgumentType, reportUnknownArgumentType, reportUnknownMemberType] # preserve enum indexing and iterable values from Phase 1 + except KeyError: + return (raw,) + + return list(chain.from_iterable(expanded(raw) for raw in allowed_routes)) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 78ce9728f3b..ed01367c0fb 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -5926,11 +5926,12 @@ class TestGuardrailModificationCheck: `_get_admin_metadata` silently ignores at read time. """ - def _call(self, request_body): + def _call(self, request_body, *, allowed=False): from litellm.proxy.auth.auth_checks import guardrail_modification_check - team_object = MagicMock() - team_object.metadata = {} # no permission + team_object = LiteLLM_TeamTable( + team_id="guardrail-team", metadata={"guardrails": {"modify_guardrails": allowed}} + ) return guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): @@ -5940,47 +5941,31 @@ class TestGuardrailModificationCheck: def test_rejects_guardrails_list(self): from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"guardrails": ["custom"]}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"guardrails": ["custom"]}}) + assert exc.value.status_code == 403 def test_rejects_disable_global_guardrails_plural(self): from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"disable_global_guardrails": True}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"disable_global_guardrails": True}}) + assert exc.value.status_code == 403 def test_rejects_disable_global_guardrail_singular(self): """VERIA-28's originally-reported singular-key typo variant.""" from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"disable_global_guardrail": True}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"disable_global_guardrail": True}}) + assert exc.value.status_code == 403 def test_rejects_opted_out_global_guardrails(self): from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}}) + assert exc.value.status_code == 403 @pytest.mark.parametrize( "key", @@ -6001,46 +5986,30 @@ class TestGuardrailModificationCheck: """ from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": {key: empty_value}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {key: empty_value}}) + assert exc.value.status_code == 403 def test_rejects_injection_via_litellm_metadata_key(self): """Caller can populate the OTHER metadata key; that must also 403.""" from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"litellm_metadata": {"disable_global_guardrails": True}}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"litellm_metadata": {"disable_global_guardrails": True}}) + assert exc.value.status_code == 403 def test_rejects_root_level_injection(self): """Top-level injection (`request_body["disable_global_guardrails"]`) was VERIA-28's easiest variant to hit — keep it rejected.""" from fastapi import HTTPException - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"disable_global_guardrails": True}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"disable_global_guardrails": True}) + assert exc.value.status_code == 403 def test_allows_when_team_has_permission(self): - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=True, - ): - # no-op, should not raise - self._call({"metadata": {"disable_global_guardrails": True}}) + # no-op, should not raise + self._call({"metadata": {"disable_global_guardrails": True}}, allowed=True) def test_rejects_string_encoded_metadata_bypass(self): """Regression: attacker sends metadata as JSON string to bypass the @@ -6051,13 +6020,9 @@ class TestGuardrailModificationCheck: from fastapi import HTTPException attacker_payload = {"disable_global_guardrails": True} - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"metadata": _json.dumps(attacker_payload)}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"metadata": _json.dumps(attacker_payload)}) + assert exc.value.status_code == 403 def test_rejects_string_encoded_litellm_metadata_bypass(self): """Same bypass via the litellm_metadata key.""" @@ -6066,13 +6031,9 @@ class TestGuardrailModificationCheck: from fastapi import HTTPException attacker_payload = {"guardrails": ["evaded"]} - with patch( - "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", - return_value=False, - ): - with pytest.raises(HTTPException) as exc: - self._call({"litellm_metadata": _json.dumps(attacker_payload)}) - assert exc.value.status_code == 403 + with pytest.raises(HTTPException) as exc: + self._call({"litellm_metadata": _json.dumps(attacker_payload)}) + assert exc.value.status_code == 403 def test_noop_when_string_is_not_json_object(self): """Unparseable strings should not trigger a 403 — they have no keys.""" @@ -11850,7 +11811,9 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_ unmapped_request = _client_tag_request(client_tags, baggage=False) unmapped = _client_tag_body(client_tags, client_slot) LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(unmapped, route) - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key")) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key") + ) mapped_request = _client_tag_request(client_tags, baggage=baggage) mapped = _client_tag_body(client_tags, client_slot) diff --git a/tests/unit/proxy/auth/test_budget_alert_adapter.py b/tests/unit/proxy/auth/test_budget_alert_adapter.py new file mode 100644 index 00000000000..8582074fc6a --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_alert_adapter.py @@ -0,0 +1,125 @@ +import asyncio +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth import budget_alerts +from litellm._auth.budget_alerts import Facts +from litellm._auth.failure import NativeAuthFailure, NativeAuthValueError +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.proxy._types import CallInfo, Litellm_EntityType, UserAPIKeyAuth +from litellm.proxy.auth import budget_alert_adapter as native +from litellm.rust_bridge import _native +from tests.unit.proxy.auth.budget_alert_reference import reference_alert + + +def observed_reference(facts: Facts, global_config: object, local_config: object) -> object: + try: + return reference_alert(facts, global_config, local_config) + except (ValueError, OverflowError, AttributeError) as error: + return type(error).__name__, str(error) + + +def observed_native(facts: Facts, global_config: object, local_config: object) -> object: + try: + decision: Final = budget_alerts.plan(native.read_alert_facts(facts, global_config, local_config)) + native.raise_alert_failure(decision) + config: Final = native.alert_emails(decision) if decision.emit else None + return decision.emit, None if config is None else tuple((key, tuple(value)) for key, value in config.items()) + except (ValueError, OverflowError, AttributeError) as error: + return type(error).__name__, str(error) + + +@pytest.mark.parametrize("member", (False, True)) +@pytest.mark.parametrize("spend", (-1.0, 0.0, 7.9, 8.0, 9.0, 10.0, float("nan"), float("inf"))) +@pytest.mark.parametrize("limit", (-1.0, 0.0, 10.0, float("nan"), float("inf"))) +@pytest.mark.parametrize( + "configs", + ( + (None, None), + ({"80": " a@x.test,b@x.test "}, {"80": [" b@x.test ", "c@x.test", "", None]}), + (None, {"bad": "a@x.test"}), + (None, {"0": [], "101": None, "080": [" a@x.test "]}), + (None, {"80": ("a@x.test",)}), + (None, {"٨٠": "a@x.test"}), + (None, {"²": []}), + (None, {"9" * 400: [], "80": "a@x.test"}), + (None, {"9" * 400: []}), + (None, "invalid"), + ), +) +def test_native_threshold_merge_and_trigger_match_reference( + member: bool, spend: float, limit: float, configs: tuple[object, object] +) -> None: + facts: Final = Facts(member, spend, limit, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) + global_config, local_config = configs + global_value: Final = None if member else global_config + assert observed_native(facts, global_value, local_config) == observed_reference(facts, global_value, local_config) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("config", (None, {"80": " a@x.test, b@x.test "})) +async def test_key_alert_schedules_once_without_awaiting_and_preserves_projection( + monkeypatch: pytest.MonkeyPatch, config: object +) -> None: + monkeypatch.setattr(litellm, "default_key_max_budget_alert_emails", None) + key: Final = UserAPIKeyAuth( + token="token", + spend=8, + max_budget=10, + metadata={"max_budget_alert_emails": config}, + user_id="owner", + team_id="team", + ) + logger: Final = Mock(budget_alerts=AsyncMock()) + await native.virtual_key_max_budget_alert_check(key, logger) + assert logger.budget_alerts.call_count == 1 + assert logger.budget_alerts.await_count == 0 + call_info: Final = logger.budget_alerts.call_args.kwargs["user_info"] + expected: Final = CallInfo( + token="token", + spend=8, + max_budget=10, + soft_budget=None, + user_id="owner", + team_id="team", + team_alias=None, + organization_id=None, + user_email=None, + key_alias=None, + event_group=Litellm_EntityType.KEY, + ) + projected: Final = ( + expected + if config is None + else expected.model_copy(update={"max_budget_alert_emails": {"80": ["a@x.test", "b@x.test"]}}) + ) + assert call_info.model_dump() == projected.model_dump() + assert call_info.model_fields_set == projected.model_fields_set + await asyncio.sleep(0) + assert logger.budget_alerts.await_count == 1 + + +def test_native_failure_is_terminal_before_alert_config_reads(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = ValueError("native alert decoder failed") + backend: Final = Mock(side_effect=failure) + monkeypatch.setattr(_native, "request_auth_budget_alert_plan", backend) + with pytest.raises(NativeAuthFailure, match="native alert decoder failed") as error: + native.read_alert_facts(Facts(False, 8.0, 10.0, 0.8), None, {"80": []}) + assert isinstance(error.value, ValueError) + assert backend.call_count == 1 + + +def test_pricing_native_failure_never_becomes_lookup_failure() -> None: + from litellm.proxy.auth.budget_pricing_adapter import read_model_price_facts + + failure: Final = NativeAuthValueError("native registry failed") + registry: Final = Mock(group=Mock(side_effect=failure)) + router: Final = Mock(registry=Mock(return_value=registry)) + with pytest.raises(NativeAuthFailure) as error: + read_model_price_facts("model", router) + assert error.value is failure + registry.group.assert_called_once() + registry.deployments.assert_not_called() diff --git a/tests/unit/proxy/auth/test_budget_context_adapter.py b/tests/unit/proxy/auth/test_budget_context_adapter.py new file mode 100644 index 00000000000..662bd369f0d --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_context_adapter.py @@ -0,0 +1,335 @@ +from collections.abc import Mapping +from typing import Final, Literal +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth import budget_context as core +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks +from litellm.proxy.auth import budget_context_adapter as adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.proxy.auth import phase1_budget_auth_reference as auth_reference +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy + + +@pytest.fixture(params=(False, True), ids=("reference", "native")) +def native(request: pytest.FixtureRequest) -> bool: + return request.param is True # pyright: ignore[reportAny] # fixture parameters are the two literal booleans above + + +@pytest.mark.parametrize("scope", ("Team", "EndUser", "CustomEndUser")) +@pytest.mark.parametrize("identity", (None, "", "entity")) +@pytest.mark.parametrize("budget_count", (0, 1, 2)) +@pytest.mark.parametrize("model_count", (0, 1, 3)) +def test_native_model_read_selection( + scope: Literal["Team", "EndUser", "CustomEndUser"], identity: str | None, budget_count: int, model_count: int +) -> None: + facts: Final = core.ModelCheckFacts(scope, identity, budget_count, model_count) + expected: Final = ( + tuple(range(model_count)) if budget_count and identity is not None and (scope != "EndUser" or identity) else () + ) + assert core.model_checks(facts) == expected + + +@pytest.mark.parametrize("keep", (False, True)) +@pytest.mark.parametrize("present", (False, True)) +@pytest.mark.parametrize("limit", (None, 0, -1, 1)) +@pytest.mark.parametrize("models", (None, 0, 2)) +def test_native_token_assignment_table(keep: bool, present: bool, limit: int | None, models: int | None) -> None: + facts: Final = core.TokenDefaultsFacts(core.BudgetFields(limit, limit, limit, limit, models), keep, present) + expected: Final = tuple( + name + for name, value in (("MaxBudget", limit), ("Tpm", limit), ("Rpm", limit), ("Tpd", limit), ("Models", models)) + if present and (not keep or value is None) + ) + assert core.token_assignments(facts) == expected + + +@pytest.mark.parametrize("mode", ("Models", "Token", "Default")) +@pytest.mark.parametrize("object_present", (False, True)) +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("db", (False, True)) +@pytest.mark.parametrize("key_id", (None, "", "key")) +def test_end_user_initial_table( + mode: Literal["Models", "Token", "Default"], + object_present: bool, + identity: str | None, + db: bool, + key_id: str | None, +) -> None: + facts: Final = core.EndUserFacts(mode, object_present, identity, db, key_id) + expected: Final = ( + "Own" + if mode != "Default" and object_present + else "None" + if mode == "Models" and (not identity or not db) + else "ReadKey" + if key_id is not None + else "ReadGlobalConfig" + ) + assert core.end_user_source(facts) == expected + + +@pytest.mark.parametrize("key_row", (False, True)) +@pytest.mark.parametrize("global_id", (None, "", "proxy")) +@pytest.mark.parametrize("global_read", (False, True)) +@pytest.mark.parametrize("global_row", (False, True)) +def test_end_user_fallback_stages(key_row: bool, global_id: str | None, global_read: bool, global_row: bool) -> None: + facts: Final = core.EndUserFacts( + "Default", + False, + None, + True, + "key", + key_read=True, + key_budget=core.BudgetFields() if key_row else None, + global_config_read=True, + global_budget_id=global_id, + global_read=global_read, + global_budget=core.BudgetFields() if global_row else None, + ) + expected: Final = ( + "Key" + if key_row + else "None" + if global_id is None + else "ReadGlobal" + if not global_read + else "Global" + if global_row + else "None" + ) + assert core.end_user_source(facts) == expected + + +@pytest.mark.parametrize("default_id", (None, "", "default")) +@pytest.mark.parametrize("member_id", (None, "", "member", "default")) +@pytest.mark.parametrize("member_models", (None, 0, 2)) +@pytest.mark.parametrize("default_models", (None, 0, 2)) +@pytest.mark.parametrize("skip", (False, True)) +def test_member_snapshot_table( + default_id: str | None, member_id: str | None, member_models: int | None, default_models: int | None, skip: bool +) -> None: + facts: Final = core.MemberFacts( + "", + "", + True, + default_id, + skip, + membership_read=True, + membership=core.MembershipSnapshot(member_id, member_models), + default_read=True, + default_budget=core.BudgetFields(model_budget_count=default_models), + ) + expected: Final = ( + "Member" + if member_id is not None and member_id != default_id and member_models + else "Done" + if default_id is None + else "Default" + if default_models is not None + else "Clear" + ) + assert core.member_source(facts) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (None, "", "entity")) +@pytest.mark.parametrize("scope", ("Team", "EndUser")) +@pytest.mark.parametrize("fail", (False, True)) +async def test_selected_model_reads_keep_order_and_stop_on_failure( + native: bool, identity: str | None, scope: Literal["Team", "EndUser"], fail: bool +) -> None: + failure: Final = RuntimeError("counter failed") + read: Final = AsyncMock(side_effect=(True, failure, True) if fail else None, return_value=True) + limiter: Final = Mock() + limiter.is_team_within_model_budget = read + limiter.is_end_user_within_model_budget = read + caps: Final[Mapping[str, object]] = {"a": {"max_budget": 1}} + token: Final = UserAPIKeyAuth(team_id=identity, team_model_max_budget=dict(caps)) + models: Final = ["a", "a", "b"] + coroutine: Final = ( + (adapter.check_team_model_budget if native else legacy._check_team_model_budget)(token, limiter, models) # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + if scope == "Team" + else (adapter.check_end_user_model_budget if native else legacy._check_end_user_model_budget)( # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + identity, caps, limiter, models, "team" + ) + ) + active: Final = identity is not None and (scope == "Team" or bool(identity)) + if fail and active: + with pytest.raises(RuntimeError, match="counter failed") as raised: + await coroutine + assert raised.value is failure + else: + await coroutine + assert read.await_count == (2 if fail else 3) if active else read.await_count == 0 + assert tuple(item.kwargs["model"] for item in read.await_args_list) == tuple(models[: read.await_count]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own", (False, True)) +@pytest.mark.parametrize("own_row", (False, True)) +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("key_found", (False, True)) +async def test_end_user_source_read_order_and_identity( + native: bool, monkeypatch: pytest.MonkeyPatch, own: bool, own_row: bool, identity: str | None, key_found: bool +) -> None: + budget: Final = LiteLLM_BudgetTable(model_max_budget={"m": {"max_budget": 1}}) + end_user: Final = ( + LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if own_row else None) + if own + else None + ) + read: Final = AsyncMock(side_effect=(budget,) if key_found else (None, budget)) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", read) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", read) + monkeypatch.setattr(litellm, "max_end_user_budget_id", "proxy") + fn: Final = adapter.end_user_model_max_budget if native else legacy._end_user_model_max_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + value: Final = await fn(identity, end_user, "key", Mock(), UserApiKeyCache(), None) + expected: Final = getattr(budget, "model_max_budget", None) if (own and own_row) or (not own and identity) else None + assert value is expected + assert read.await_count == (0 if own or not identity else 1 if key_found else 2) + if read.await_count: + assert read.await_args_list[0].kwargs["budget_id"] == "key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("keep", (False, True)) +@pytest.mark.parametrize("present", (False, True)) +async def test_token_defaults_preserve_caps_and_projection_identity(native: bool, keep: bool, present: bool) -> None: + budget: Final = LiteLLM_BudgetTable( + max_budget=9, tpm_limit=8, rpm_limit=7, tpd_limit=6, model_max_budget={"m": {"max_budget": 5}} + ) + token: Final = UserAPIKeyAuth( + end_user_max_budget=0 if present else None, + end_user_tpm_limit=0 if present else None, + end_user_model_max_budget={} if present else None, + ) + original_models: Final = token.end_user_model_max_budget + own: Final = LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget) + fn: Final = ( + adapter.apply_key_end_user_default_budget_to_token + if native + else legacy._apply_key_end_user_default_budget_to_token # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + await fn(token, own, "key", Mock(), UserApiKeyCache(), None, keep) + assert token.end_user_max_budget == (0 if keep and present else 9) + assert token.end_user_tpm_limit == (0 if keep and present else 8) + assert (token.end_user_rpm_limit, token.end_user_tpd_limit) == (7, 6) + assert token.end_user_model_max_budget is ( + original_models if keep and present else getattr(budget, "model_max_budget", None) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("default_id", (None, "", "default")) +@pytest.mark.parametrize("member_id", (None, "default", "override")) +@pytest.mark.parametrize("member_caps", (None, {}, {"member": {"max_budget": 1}})) +@pytest.mark.parametrize("default_caps", (None, {}, {"default": {"max_budget": 2}})) +async def test_member_default_read_and_mutation_order( + native: bool, + monkeypatch: pytest.MonkeyPatch, + default_id: str | None, + member_id: str | None, + member_caps: Mapping[str, object] | None, + default_caps: Mapping[str, object] | None, +) -> None: + member_budget: Final = LiteLLM_BudgetTable(model_max_budget=dict(member_caps) if member_caps is not None else None) + default_budget: Final = LiteLLM_BudgetTable( + model_max_budget=dict(default_caps) if default_caps is not None else None + ) + membership: Final = LiteLLM_TeamMembership( + user_id="", team_id="", budget_id=member_id, litellm_budget_table=member_budget + ) + member_read: Final = AsyncMock(return_value=membership) + default_read: Final = AsyncMock(return_value=default_budget) + monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(auth_reference, "get_team_membership", member_read) + monkeypatch.setattr(legacy, "get_team_membership", member_read) + monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(auth_reference, "get_team_member_default_budget", default_read) + monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) + token: Final = UserAPIKeyAuth(user_id="", team_id="", team_member_model_max_budget={"old": {}}) + prior: Final = token.team_member_model_max_budget + team: Final = LiteLLM_TeamTable(team_id="", metadata={"team_member_budget_id": default_id}) + fn: Final = ( + adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + await fn(token, team, Mock(), UserApiKeyCache(), False, None) + override: Final = member_id is not None and member_id != default_id and bool(member_caps) + expected: Final = ( + getattr(member_budget, "model_max_budget", None) + if override + else getattr(default_budget, "model_max_budget", None) + if default_id is not None + else prior + ) + assert token.team_member_model_max_budget is expected + member_read.assert_awaited_once() + assert default_read.await_count == (0 if override or default_id is None else 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("skip", (False, True)) +@pytest.mark.parametrize("during_default", (False, True)) +async def test_member_lookup_error_policy( + native: bool, monkeypatch: pytest.MonkeyPatch, skip: bool, during_default: bool +) -> None: + error: Final = RuntimeError("budget read failed") + member_read: Final = AsyncMock(return_value=None, side_effect=None if during_default else error) + default_read: Final = AsyncMock(side_effect=error) + monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(auth_reference, "get_team_membership", member_read) + monkeypatch.setattr(legacy, "get_team_membership", member_read) + monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(auth_reference, "get_team_member_default_budget", default_read) + monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) + token: Final = UserAPIKeyAuth(user_id="u", team_id="t", team_member_model_max_budget={"old": {}}) + prior: Final = token.team_member_model_max_budget + team: Final = LiteLLM_TeamTable(team_id="t", metadata={"team_member_budget_id": "default"}) + fn: Final = ( + adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + pending: Final = fn(token, team, Mock(), UserApiKeyCache(), skip, None) + if skip: + await pending + else: + with pytest.raises(RuntimeError, match="budget read failed") as raised: + await pending + assert raised.value is error + assert token.team_member_model_max_budget is prior + + +@pytest.mark.asyncio +async def test_member_native_failure_never_becomes_skippable_lookup_error(monkeypatch: pytest.MonkeyPatch) -> None: + error: Final = NativeAuthRuntimeError("native planner failed") + read: Final = AsyncMock(side_effect=error) + monkeypatch.setattr(auth_checks, "get_team_membership", read) + monkeypatch.setattr(auth_reference, "get_team_membership", read) + with pytest.raises(NativeAuthRuntimeError, match="native planner failed") as raised: + await adapter.load_team_member_default_model_budget( + UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None + ) + assert raised.value is error + read.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_member_planner_failure_stops_before_io(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge import _native + + native: Final = Mock(side_effect=RuntimeError("decode failed")) + read: Final = AsyncMock() + monkeypatch.setattr(_native, "request_auth_context_member", native) + monkeypatch.setattr(auth_checks, "get_team_membership", read) + monkeypatch.setattr(auth_reference, "get_team_membership", read) + with pytest.raises(NativeAuthRuntimeError, match="decode failed"): + await adapter.load_team_member_default_model_budget( + UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None + ) + read.assert_not_awaited() + native.assert_called_once() diff --git a/tests/unit/proxy/auth/test_budget_lookup_adapter.py b/tests/unit/proxy/auth/test_budget_lookup_adapter.py new file mode 100644 index 00000000000..38707a68983 --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_lookup_adapter.py @@ -0,0 +1,276 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from math import inf, isnan, nan +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_OrganizationTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks, common_budget_adapter, common_budget_lookups +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from . import phase1_budget_lookup_reference as reference +from .test_common_budget_adapter import observe + +_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) +_CALL: Final = TypeAdapter(tuple[str, tuple[object, ...], Mapping[str, object]]) + + +def normalized_calls(events: Mock) -> tuple[tuple[str, tuple[object, ...], Mapping[str, object]], ...]: + return tuple( + ( + name, + args, + MappingProxyType( + {key: "NaN" if isinstance(value, float) and isnan(value) else value for key, value in kwargs.items()} + ), + ) + for name, args, kwargs in (_CALL.validate_python(value) for value in events.mock_calls) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("token_id", "team_id"), ((None, None), (None, "team-org"), ("", "team-org"), ("key-org", "team-org")) +) +@pytest.mark.parametrize(("token_present", "client_present"), ((True, True), (False, True), (True, False))) +@pytest.mark.parametrize("row_state", ("present", "missing", "failed")) +@pytest.mark.parametrize("limit", (None, 1.0, inf)) +async def test_organization_source_and_raw_reads_match_phase1( + monkeypatch: pytest.MonkeyPatch, + token_id: str | None, + team_id: str | None, + token_present: bool, + client_present: bool, + row_state: str, + limit: float | None, +) -> None: + row: Final = LiteLLM_OrganizationTable.model_construct( + organization_id="row", spend=2.0, budget_id="budget", litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit) + ) + reader: Final = AsyncMock( + return_value=row if row_state == "present" else None, + side_effect=RuntimeError("unavailable") if row_state == "failed" else None, + ) + counter: Final = AsyncMock(return_value=2.0) + monkeypatch.setattr(reference, "get_org_object", reader) + monkeypatch.setattr(auth_checks, "get_org_object", reader) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + old: Final = UserAPIKeyAuth(org_id=token_id, token="token", key_alias="alias") + new: Final = old.model_copy(deep=True) + team: Final = LiteLLM_TeamTable(team_id="team", organization_id=team_id) + prisma: Final = Mock(spec=PrismaClient) if client_present else None + cache: Final = UserApiKeyCache() + alerts: Final = AsyncMock() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + new_log: Final = old_log + before: Final = await observe( + lambda: reference.organization_max_budget_check(old if token_present else None, team, prisma, cache, old_log), + old_log, + old, + ) + old_reads: Final = (tuple(reader.await_args_list), tuple(counter.await_args_list)) + reader.reset_mock() + counter.reset_mock() + alerts.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.organization_max_budget_check( + new if token_present else None, team, prisma, cache, new_log + ), + new_log, + new, + ) + assert after == before + assert (tuple(reader.await_args_list), tuple(counter.await_args_list)) == old_reads + assert new.model_dump() == old.model_dump() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loaded", (False, True)) +@pytest.mark.parametrize("member_limit", (None, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("default_id", (None, "", "default", 1)) +@pytest.mark.parametrize("missing_member", (False, True)) +async def test_member_default_and_counter_order_match_phase1( + monkeypatch: pytest.MonkeyPatch, + loaded: bool, + member_limit: float | None, + default_id: object, + missing_member: bool, +) -> None: + row: Final = LiteLLM_TeamMembership( + team_id="team", user_id="user", spend=2.0, litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_limit) + ) + membership: Final = None if missing_member else row + member_read: Final = AsyncMock(return_value=membership) + default_read: Final = AsyncMock(return_value=LiteLLM_BudgetTable(max_budget=1)) + counter: Final = AsyncMock(return_value=2.0) + clock: Final = Mock(return_value=_NOW) + events: Final = Mock() + events.attach_mock(member_read, "membership") + events.attach_mock(default_read, "default") + events.attach_mock(counter, "counter") + events.attach_mock(clock, "clock") + for module in (reference, auth_checks): + monkeypatch.setattr(module, "get_team_membership", member_read) + monkeypatch.setattr(module, "get_team_member_default_budget", default_read) + monkeypatch.setattr(module, "get_utc_datetime", clock) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + team: Final = LiteLLM_TeamTable( + team_id="team", + metadata={ + "team_member_budget_id": default_id, + "team_member_max_budget_alert_emails": {"50": ["owner"]}, + }, + ) + key: Final = UserAPIKeyAuth(user_id="user") + prisma: Final = Mock(spec=PrismaClient) + cache: Final = UserApiKeyCache() + alerts: Final = AsyncMock() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + new_log: Final = old_log + before: Final = await observe( + lambda: reference._check_team_member_budget(team, None, key, prisma, cache, old_log, membership, loaded), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + old_reads: Final = normalized_calls(events) + events.reset_mock() + alerts.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.check_team_member_budget( + team, None, key, prisma, cache, new_log, membership, loaded + ), + new_log, + key, + ) + assert after == before + assert normalized_calls(events) == old_reads + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw", (None, "", " a,b , ,a", "\x1ca\x1f,\u0085b\u00a0", [" a ", "", "\x1c", 4, "a"], ("a",), 4, {"a": 1}) +) +@pytest.mark.parametrize( + ("limit", "spend", "token"), + ((1, 2, True), (1, 0, True), (0, 0, True), (1, 2, False), (None, 2, True), (nan, 2, True)), +) +async def test_soft_email_normalization_matches_phase1( + raw: object, + limit: float | None, + spend: float, + token: bool, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="team", soft_budget=limit, spend=spend, metadata={"soft_budget_alerting_emails": raw} + ) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe( + lambda: reference._team_soft_budget_check(team, key if token else None, old_log), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + after: Final = await observe( + lambda: common_budget_adapter.team_soft_budget_check(team, key if token else None, new_log), + new_log, + key, + ) + assert after == before + + +@pytest.mark.asyncio +async def test_organization_native_failure_from_lookup_is_terminal(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("native failed") + reader: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(auth_checks, "get_org_object", reader) + with pytest.raises(NativeAuthRuntimeError) as raised: + await common_budget_lookups.organization_max_budget_check( + UserAPIKeyAuth(org_id="org"), None, Mock(spec=PrismaClient), UserApiKeyCache(), Mock(spec=ProxyLogging) + ) + assert raised.value is failure + assert reader.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_present", "team_id", "token_present", "user_id"), + ( + (False, "team", True, "user"), + (True, None, True, "user"), + (True, "team", False, "user"), + (True, "team", True, None), + (True, "", True, ""), + ), +) +async def test_member_identity_stages_match_phase1( + monkeypatch: pytest.MonkeyPatch, + team_present: bool, + team_id: str | None, + token_present: bool, + user_id: str | None, +) -> None: + team: Final = LiteLLM_TeamTable.model_construct(team_id=team_id) if team_present else None + key: Final = UserAPIKeyAuth(user_id=user_id) + reader: Final = AsyncMock(return_value=None) + clock: Final = Mock(return_value=_NOW) + for module in (reference, auth_checks): + monkeypatch.setattr(module, "get_team_membership", reader) + monkeypatch.setattr(module, "get_utc_datetime", clock) + logging: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + cache: Final = UserApiKeyCache() + before: Final = await observe( + lambda: reference._check_team_member_budget(team, None, key if token_present else None, None, cache, logging), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + logging, + key, + ) + reads: Final = (tuple(reader.await_args_list), clock.call_count) + reader.reset_mock() + clock.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.check_team_member_budget( + team, None, key if token_present else None, None, cache, logging + ), + logging, + key, + ) + assert after == before + assert (tuple(reader.await_args_list), clock.call_count) == reads + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True)) +async def test_soft_budget_skip_does_not_read_metadata(native: bool) -> None: + team: Final = LiteLLM_TeamTable(team_id="team", soft_budget=10, spend=0) + object.__setattr__(team, "metadata", object()) + alerts: Final = AsyncMock() + logging: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + policy: Final = common_budget_adapter.team_soft_budget_check if native else reference._team_soft_budget_check # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + result: Final = await policy(team, UserAPIKeyAuth(), logging) + assert result is None + alerts.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw", (None, [], 1, "metadata", {1: "ignored", "soft_budget_alerting_emails": " a "})) +async def test_soft_metadata_shape_preserves_phase1_projection(raw: object) -> None: + team: Final = LiteLLM_TeamTable(team_id="team", soft_budget=1, spend=2) + object.__setattr__(team, "metadata", raw) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe( + lambda: reference._team_soft_budget_check(team, key, old_log), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + after: Final = await observe(lambda: common_budget_adapter.team_soft_budget_check(team, key, new_log), new_log, key) + assert after == before diff --git a/tests/unit/proxy/auth/test_budget_pricing_adapter.py b/tests/unit/proxy/auth/test_budget_pricing_adapter.py new file mode 100644 index 00000000000..ff5fdd62a7d --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_pricing_adapter.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm.proxy.auth import budget_pricing_adapter as native +from tests.unit.proxy.auth import phase1_budget_auth_reference as reference + + +def router(input_price: float | None, output_price: float | None, deployments: tuple[object, ...] = ()) -> Mock: + view: Final = Mock() + view.group.return_value = SimpleNamespace(input_cost_per_token=input_price, output_cost_per_token=output_price) + view.zero_cost.return_value = {} + view.deployments.return_value = deployments + return Mock(registry=Mock(return_value=view)) + + +@pytest.mark.parametrize(("input_price", "output_price"), ((None, 0), (0, None), (1, 0), (0, 1))) +def test_priced_or_undefined_group_does_not_read_deployments( + input_price: float | None, output_price: float | None +) -> None: + fake: Final = router(input_price, output_price) + facts: Final = native.read_model_price_facts("model", fake) + assert facts == ModelPriceFacts(group_found=True, input_price=input_price, output_price=output_price) + fake.registry().deployments.assert_not_called() + + +@pytest.mark.parametrize( + "route", + ( + "/v1/models", + "/model/info", + "/health/services", + "/chat/completions", + "/mcp", + "/mcp-rest/tools/list", + "/mcp/tools/call", + "/key/info", + ), +) +@pytest.mark.parametrize("body", (None, {}, {"method": "tools/list"}, {"method": "tools/call"}, {"method": 1})) +def test_route_and_mcp_decision_matches_production_reference(route: str, body: dict[str, object] | None) -> None: + assert native.request_skips_budget_checks( + route, None, None, request_body=body + ) == reference.request_skips_budget_checks(route, None, None, request_body=body) + assert native.route_skips_budget_checks(route) == reference.route_skips_budget_checks(route) + + +def test_missing_models_and_empty_model_list_preserve_exemption_behavior() -> None: + fake: Final = router(0, 0) + assert native.is_model_cost_zero(None, fake) is False + assert native.is_model_cost_zero("model", None) is False + assert native.is_model_cost_zero([], fake) is True + + +@pytest.mark.parametrize( + ("prices", "reads", "exempt"), (((0, 1, 0), 2, False), ((0, 0, 0), 3, True), ((1, 0, 0), 1, False)) +) +def test_raw_exemption_prices_only_native_selected_models_and_preserves_cache_writes( + monkeypatch: pytest.MonkeyPatch, prices: tuple[float, ...], reads: int, exempt: bool +) -> None: + from litellm._auth import budget_pricing + + old: Final = router(0, 0) + new: Final = router(0, 0) + monkeypatch.setitem( + litellm.model_cost, "budget-zero-fixture", {"input_cost_per_token": 0, "output_cost_per_token": 0} + ) + deployment: Final = { + "model_name": "model", + "litellm_params": {"model": "openai/fixture"}, + "model_info": {"id": "budget-zero-fixture"}, + } + old.registry().deployments.return_value = (deployment,) + new.registry().deployments.return_value = (deployment,) + old.registry().resolve.side_effect = lambda model: model + new.registry().resolve.side_effect = lambda model: model + old.registry().group.side_effect = tuple( + SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices + ) + new.registry().group.side_effect = tuple( + SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices + ) + old_cache: Final[dict[str, bool]] = {} + new_cache: Final[dict[str, bool]] = {} + old.registry().zero_cost.return_value = old_cache + new.registry().zero_cost.return_value = new_cache + models: Final = ["first", "second", "third"] + before: Final = reference.request_skips_budget_checks("/chat/completions", models, old) + facts: Final = native.read_budget_exemption_facts("/chat/completions", models, new) + assert budget_pricing.exemption(facts).exempt is before is exempt + assert len(facts.prices) == reads + assert old.registry().group.call_count == new.registry().group.call_count == reads + assert old_cache == new_cache + + +def test_raw_exemption_route_skip_does_not_touch_model_registry() -> None: + fake: Final = router(0, 0) + facts: Final = native.read_budget_exemption_facts("/v1/models", ["first", "second"], fake) + assert facts.prices == () + fake.registry.assert_not_called() + + +def test_raw_cached_price_does_not_read_model_group_or_write_again() -> None: + fake: Final = router(0, 0) + fake.registry().zero_cost.return_value = {"model": False} + facts: Final = native.read_model_cost_facts("model", fake) + assert facts.cached is False + fake.registry().group.assert_not_called() + + +def test_empty_model_list_still_reads_zero_cost_cache() -> None: + fake: Final = router(0, 0) + facts: Final = native.read_budget_exemption_facts("/chat/completions", [], fake) + assert facts.prices == () + assert facts.cache_read is True + fake.registry().zero_cost.assert_called_once_with() + fake.registry().group.assert_not_called() diff --git a/tests/unit/proxy/auth/test_common_budget_adapter.py b/tests/unit/proxy/auth/test_common_budget_adapter.py new file mode 100644 index 00000000000..ace55803f6b --- /dev/null +++ b/tests/unit/proxy/auth/test_common_budget_adapter.py @@ -0,0 +1,297 @@ +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from math import inf +from typing import Final, Literal, Protocol +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth.common_budgets import BudgetFacts +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + BudgetLimitEntry, + LiteLLM_EndUserTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth import budget_windows_adapter, common_budget_lookups +from litellm.proxy.auth import common_budget_adapter as native +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from tests.unit.proxy.auth import phase1_budget_auth_reference as reference + + +@dataclass(frozen=True, slots=True) +class Observed: + error: tuple[object, ...] | None + alerts: tuple[object, ...] + throttle: float | None + + +class AlertLog(Protocol): + @property + def budget_alerts(self) -> AsyncMock: ... + + +async def observe(run: Callable[[], Awaitable[None]], logging: AlertLog, key: UserAPIKeyAuth) -> Observed: + try: + await run() + except litellm.BudgetExceededError as error: + failure: Final = (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id) + await asyncio.sleep(0) + return Observed(failure, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct) + await asyncio.sleep(0) + return Observed(None, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2))) +async def test_team_error_envelope_and_alert_projection_match_reference( + monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + team: Final = LiteLLM_TeamTable(team_id="team", max_budget=limit, spend=0) + key: Final = UserAPIKeyAuth(token="hash", team_id="team") + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.team_max_budget_check(team, key, old_log), old_log, key) + after: Final = await observe(lambda: native.team_max_budget_check(team, key, new_log), new_log, key) + assert after == before + assert reader.await_count == (0 if limit is None else 2) + assert reader.await_args_list[:1] == reader.await_args_list[1:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2))) +@pytest.mark.parametrize("throttle", (None, 0.5, 2)) +async def test_key_error_envelope_alert_and_throttle_mutation_match_reference( + monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float, throttle: float | None +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", throttle) + old: Final = UserAPIKeyAuth( + token="hash", key_alias="alias", max_budget=limit, tpm_limit=10, metadata={"throttle_on_budget_exceeded": True} + ) + new: Final = old.model_copy(deep=True) + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.virtual_key_max_budget_check(old, old_log), old_log, old) + after: Final = await observe(lambda: native.virtual_key_max_budget_check(new, new_log), new_log, new) + assert after == before + assert reader.await_args_list[:1] == reader.await_args_list[1:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("/chat/completions", "/key/info")) +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 1), (1, 2), (0, 0), (-1, 0))) +async def test_end_user_info_skip_and_strict_ceiling_match_reference( + monkeypatch: pytest.MonkeyPatch, route: str, limit: float | None, spend: float +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + user: Final = LiteLLM_EndUserTable( + user_id="end", blocked=False, litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit) + ) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.check_end_user_budget(user, route), old_log, key) + after: Final = await observe(lambda: native.check_end_user_budget(user, route), new_log, key) + assert after == before + assert reader.await_args_list[:1] == reader.await_args_list[1:] + + +@pytest.mark.asyncio +async def test_spend_io_failure_is_terminal_and_does_not_replay() -> None: + failure: Final = RuntimeError("counter unavailable") + reader: Final = AsyncMock(side_effect=failure) + with pytest.raises(RuntimeError, match="counter unavailable") as raised: + await native.read_budget(BudgetFacts("Key", 1), native.CounterRead("spend:key:hash", 0, 1), reader) + assert raised.value is failure + assert reader.await_count == 1 + + +@pytest.mark.parametrize( + ("member", "default", "increase", "expiry"), + ( + (1, 5, 2, datetime(2025, 1, 2)), + (None, 5, 2, datetime(2025, 1, 2)), + (1, 5, 2, datetime(2025, 1, 1, tzinfo=timezone.utc)), + (None, 0, 2, datetime(2025, 1, 2)), + (None, -1, 2, None), + (0, 5, None, None), + ), +) +def test_member_budget_selection_and_temporary_expiry_match_reference( + member: float | None, default: float | None, increase: float | None, expiry: datetime | None +) -> None: + now: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) + row: Final = LiteLLM_BudgetTable(max_budget=member, temp_budget_increase=increase, temp_budget_expiry=expiry) + team: Final = LiteLLM_BudgetTable(max_budget=default) + assert native.resolve_effective_team_member_budget( + row, team, now + ) == reference.resolve_effective_team_member_budget(row, team, now) + + +@pytest.mark.parametrize( + "policy", (reference.resolve_effective_team_member_budget, native.resolve_effective_team_member_budget) +) +def test_naive_now_preserves_datetime_error(policy: Callable[..., float | None]) -> None: + row: Final = LiteLLM_BudgetTable(max_budget=1, temp_budget_increase=2, temp_budget_expiry=datetime(2025, 1, 2)) + with pytest.raises(TypeError, match="offset-naive"): + policy(row, None, datetime(2025, 1, 1)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Key", "Team")) +@pytest.mark.parametrize("as_model", (False, True)) +@pytest.mark.parametrize("spend", (1.0, 10.0)) +async def test_budget_window_objects_preserve_counter_contract_and_denial( + monkeypatch: pytest.MonkeyPatch, + entity: str, + as_model: bool, + spend: float, +) -> None: + entry: Final = BudgetLimitEntry( + budget_duration="24h", max_budget=10.0, reset_at=datetime(2025, 1, 2, tzinfo=timezone.utc) + ) + raw: Final = entry if as_model else entry.model_dump() + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + key: Final = UserAPIKeyAuth(token="key") + team: Final = LiteLLM_TeamTable(team_id="team") + object.__setattr__(key, "budget_limits", [raw]) + object.__setattr__(team, "budget_limits", [raw]) + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + old_call: Final = ( + ( + lambda: reference._virtual_key_multi_budget_check(key) # pyright: ignore[reportPrivateUsage] # Phase 1 reference + ) + if entity == "Key" + else ( + lambda: reference._team_multi_budget_check(team) # pyright: ignore[reportPrivateUsage] # Phase 1 reference + ) + ) + new_call: Final = ( + (lambda: budget_windows_adapter.virtual_key_multi_budget_check(key)) + if entity == "Key" + else (lambda: budget_windows_adapter.team_multi_budget_check(team)) + ) + before: Final = await observe(old_call, old_log, key) + expected_read: Final = reader.await_args + reader.reset_mock() + after: Final = await observe(new_call, new_log, key) + assert after == before + assert (after.error is not None) is (spend == 10.0) + assert reader.await_args_list == [expected_read] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity", ("Key", "Team", "User", "Project", "TeamMember", "Organization", "EndUser", "Tag", "AccessGroup") +) +async def test_non_window_spend_reads_omit_unselected_options( + monkeypatch: pytest.MonkeyPatch, + entity: str, +) -> None: + reader: Final = AsyncMock(return_value=3.0) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + authoritative: Final = entity in {"EndUser", "Tag", "AccessGroup"} + assert ( + await native.read_counter(native.CounterRead(f"spend:{entity}:id", 2.0, 10.0, authoritative=authoritative)) + == 3.0 + ) + reader.assert_awaited_once_with( + counter_key=f"spend:{entity}:id", + fallback_spend=2.0, + max_budget=10.0, + **({"fallback_authoritative": True} if authoritative else {}), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Key", "Team")) +async def test_window_spend_reads_keep_explicit_missing_metadata( + monkeypatch: pytest.MonkeyPatch, + entity: str, +) -> None: + reader: Final = AsyncMock(return_value=3.0) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + assert ( + await native.read_counter( + native.CounterRead( + "window", + 0.0, + 10.0, + window_entity_type=entity, + window_duration="invalid", + ) + ) + == 3.0 + ) + reader.assert_awaited_once_with( + counter_key="window", + fallback_spend=0.0, + max_budget=10.0, + window_entity_type=entity, + window_entity_id=None, + window_duration="invalid", + window_start=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Team", "Project", "TeamMember")) +@pytest.mark.parametrize("spend", (1.0, 11.0)) +async def test_primary_scope_denial_with_baseline_spend_signature( + monkeypatch: pytest.MonkeyPatch, + entity: Literal["Team", "Project", "TeamMember"], + spend: float, +) -> None: + async def current_spend(counter_key: str, fallback_spend: float, max_budget: float | None = None) -> float: + return fallback_spend + + reader: Final = AsyncMock(side_effect=current_spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + team: Final = LiteLLM_TeamTable(team_id="team", spend=spend, max_budget=10.0) + key: Final = UserAPIKeyAuth(user_id="user", team_id="team") + row: Final = LiteLLM_BudgetTable(max_budget=10.0) + logging: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + project: Final = LiteLLM_ProjectTableCachedObj( + project_id="project", team_id="team", spend=spend, litellm_budget_table=row + ) + member: Final = LiteLLM_TeamMembership(user_id="user", team_id="team", spend=spend, litellm_budget_table=row) + invocation: Final = { + "Team": lambda: native.team_max_budget_check(team, None, logging), + "Project": lambda: native.project_max_budget_check(project, None, logging), + "TeamMember": lambda: common_budget_lookups.check_team_member_budget( + team, + None, + key, + None, + UserApiKeyCache(), + logging, + team_membership=member, + team_membership_loaded=True, + ), + }[entity] + result: Final = await observe(invocation, logging, key) + assert (result.error is not None) is (spend > 10.0) + reader.assert_awaited_once_with( + counter_key={ + "Team": "spend:team:team", + "Project": "spend:project:project", + "TeamMember": "spend:team_member:user:team", + }[entity], + fallback_spend=spend, + max_budget=10.0, + ) diff --git a/tests/unit/proxy/auth/test_common_budget_lookups.py b/tests/unit/proxy/auth/test_common_budget_lookups.py new file mode 100644 index 00000000000..b034dafcf47 --- /dev/null +++ b/tests/unit/proxy/auth/test_common_budget_lookups.py @@ -0,0 +1,44 @@ +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm._auth import common_budgets +from litellm.proxy import proxy_server +from litellm.proxy.auth.common_budget_lookups import _read_access_group_rows +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limit", "reads", "denied"), ((None, 0, False), (0, 0, False), (-1, 0, False), (1, 1, True), (2, 2, False)) +) +async def test_raw_access_group_snapshot_stops_reads_at_native_denial( + monkeypatch: pytest.MonkeyPatch, + limit: float | None, + reads: int, + denied: bool, +) -> None: + counter: Final = AsyncMock(return_value=1.0) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + rows: Final = { + "first": ModelAccessGroupBudget(access_group_name="first", max_budget=limit), + "second": ModelAccessGroupBudget(access_group_name="second", max_budget=limit), + } + snapshots: Final = await _read_access_group_rows(("first", "second"), rows) + assert common_budgets.check(snapshots[0].facts).denied is denied + assert len(snapshots) == (1 if denied else 2) + assert snapshots[0].name == "first" + assert counter.await_count == reads + + +@pytest.mark.asyncio +async def test_access_group_counter_failure_is_terminal_before_next_group(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = RuntimeError("counter unavailable") + counter: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + row: Final = ModelAccessGroupBudget(access_group_name="first", max_budget=1) + with pytest.raises(RuntimeError, match="counter unavailable") as raised: + await _read_access_group_rows(("first", "second"), {"first": row, "second": row}) + assert raised.value is failure + assert counter.await_count == 1 diff --git a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py index 0c3487255f6..030a8570913 100644 --- a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py @@ -115,7 +115,7 @@ async def test_custom_auth_enforces_end_user_budget_when_common_checks_skipped() with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + "litellm.proxy.auth.auth_checks.get_end_user_object", new_callable=AsyncMock, return_value=over_budget_end_user, ), diff --git a/tests/unit/proxy/auth/test_end_user_context_adapter.py b/tests/unit/proxy/auth/test_end_user_context_adapter.py new file mode 100644 index 00000000000..8eaee4db5a9 --- /dev/null +++ b/tests/unit/proxy/auth/test_end_user_context_adapter.py @@ -0,0 +1,437 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks +from litellm.proxy.auth import end_user_context_adapter as adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.proxy.auth import phase1_budget_auth_reference as auth_reference +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy + + +@pytest.fixture(params=(False, True), ids=("reference", "native")) +def native(request: pytest.FixtureRequest) -> bool: + return request.param is True # pyright: ignore[reportAny] # boolean fixture parameter + + +@pytest.mark.parametrize( + "value,expected", ((None, None), ("", None), (" ", " "), ("budget", "budget"), (0, None), (False, None), ({}, None)) +) +def test_metadata_id_keeps_only_nonempty_strings(native: bool, value: object, expected: str | None) -> None: + function: Final = adapter.get_key_end_user_budget_id if native else auth_reference.get_key_end_user_budget_id + assert function({"end_user_budget_id": value}) == expected + assert function(None) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit", (None, "", "key")) +@pytest.mark.parametrize("global_id", (None, "", "global")) +@pytest.mark.parametrize("db_present", (False, True)) +async def test_default_budget_id_and_cache_selection( + native: bool, explicit: str | None, global_id: str | None, db_present: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + cache: Final = UserApiKeyCache() + selected: Final = explicit if explicit is not None else global_id + budget: Final = LiteLLM_BudgetTable(budget_id=selected, max_budget=0) + await cache.async_set_cache(key=f"default_end_user_budget:{selected}", value=budget, model_type=LiteLLM_BudgetTable) + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget + result: Final = await function(Mock() if db_present else None, cache, budget_id=explicit) + assert result == (budget if db_present and selected is not None else None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own_id", (None, "", "own")) +@pytest.mark.parametrize("own_present", (False, True)) +@pytest.mark.parametrize("key_id", (None, "", "key")) +@pytest.mark.parametrize("global_id", (None, "", "global")) +@pytest.mark.parametrize("found", (False, True)) +async def test_default_application_preserves_own_budget_and_copies_only_resolved_defaults( + native: bool, + own_id: str | None, + own_present: bool, + key_id: str | None, + global_id: str | None, + found: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + own: Final = LiteLLM_BudgetTable(budget_id="own", max_budget=2) + default: Final = LiteLLM_BudgetTable(budget_id="default", max_budget=3) if found else None + row: Final = LiteLLM_EndUserTable( + user_id="eu", blocked=False, budget_id=own_id, litellm_budget_table=own if own_present else None + ) + load: Final = AsyncMock(return_value=default) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", load) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", load) + function: Final = ( + adapter.apply_default_budget_to_end_user if native else auth_reference._apply_default_budget_to_end_user # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + ) # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + result: Final = await function(row, Mock(), UserApiKeyCache(), key_end_user_budget_id=key_id) + unchanged: Final = (own_id is not None and own_present) or (key_id is None and global_id is None) or not found + assert (result is row) is unchanged + assert result.litellm_budget_table is (row.litellm_budget_table if unchanged else default) + assert row.litellm_budget_table is (own if own_present else None) + expected_ids: Final = ( + () + if own_id is not None and own_present + else ( + ((key_id,) if key_id is not None else ()) + + ((global_id,) if global_id is not None and (key_id is None or not found) else ()) + ) + ) + assert tuple(call.kwargs.get("budget_id", global_id) for call in load.await_args_list) == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("global_id", (None, "", "global")) +@pytest.mark.parametrize("key_id", (None, "", "key")) +@pytest.mark.parametrize("validate", (False, True)) +@pytest.mark.parametrize("max_budget", (None, 0)) +@pytest.mark.parametrize("ids", (None, frozenset[str](), frozenset(("eu",)))) +async def test_unrestricted_registry_requires_absence_of_all_budget_and_validation_constraints( + native: bool, + global_id: str | None, + key_id: str | None, + validate: bool, + max_budget: int | None, + ids: frozenset[str] | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", validate) + load: Final = AsyncMock(return_value=ids) + monkeypatch.setattr(auth_checks, "_load_end_user_restricted_registry", load) + monkeypatch.setattr(auth_reference, "_load_end_user_restricted_registry", load) + function: Final = adapter.end_user_is_known_unrestricted if native else auth_reference._end_user_is_known_unrestricted # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + result: Final = await function("eu", Mock(), UserApiKeyCache(), max_budget, key_id) + blocked: Final = global_id is not None or key_id is not None or validate or max_budget is not None + assert result is (not blocked and ids is not None and "eu" not in ids) + assert load.await_count == (0 if blocked else 1) + + +@pytest.mark.parametrize("value", (None, 0, 1)) +@pytest.mark.parametrize("models", (None, {}, {"model": {"budget_limit": 0}})) +def test_budget_parameter_projection_keeps_zero_and_empty_model_limits( + native: bool, + value: int | None, + models: Mapping[str, object] | None, +) -> None: + budget: Final = LiteLLM_BudgetTable( + tpm_limit=value, + rpm_limit=value, + tpd_limit=value, + max_budget=value, + model_max_budget=dict(models) if models is not None else None, + ) + params: Final[dict[str, object]] = {"existing": True} # mutable-ok: legacy output argument + if native: + adapter.apply_budget_limits_to_end_user_params(params, budget, "eu") + else: + legacy._apply_budget_limits_to_end_user_params(params, budget, "eu") # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # Phase 1 unparameterized reference dictionary + expected: Final = { + "existing": True, + **( + { + name: value + for name in ("end_user_tpm_limit", "end_user_rpm_limit", "end_user_tpd_limit", "end_user_max_budget") + } + if value is not None + else {} + ), + **({"end_user_model_max_budget": models} if models is not None else {}), + } + assert params == expected + assert tuple(params) == tuple(expected) + if models is not None: + assert params["end_user_model_max_budget"] is getattr(budget, "model_max_budget") + + +@pytest.mark.parametrize("value", (None, 0, "opaque")) +def test_token_projection_preserves_unset_limits_and_never_copies_max_budget(native: bool, value: object) -> None: + token: Final = UserAPIKeyAuth( + end_user_id="old", + end_user_tpm_limit=9, + end_user_rpm_limit=9, + end_user_tpd_limit=9, + end_user_max_budget=9, + end_user_model_max_budget={"old": {}}, + allowed_model_region="eu", + ) + params: Final = { + name: value + for name in ( + "end_user_tpm_limit", + "end_user_rpm_limit", + "end_user_tpd_limit", + "end_user_max_budget", + "end_user_model_max_budget", + "allowed_model_region", + ) + } + result: Final = ( + adapter.update_valid_token_with_end_user_params(token, params) + if native + else legacy.update_valid_token_with_end_user_params(token, params) # pyright: ignore[reportUnknownMemberType] # Phase 1 unparameterized reference dictionary + ) + assert result is token + assert token.end_user_id is None + assert token.end_user_tpm_limit == (9 if value is None else value) + assert token.end_user_rpm_limit == (9 if value is None else value) + assert token.end_user_tpd_limit == (9 if value is None else value) + assert token.allowed_model_region == ("eu" if value is None else value) + assert token.end_user_model_max_budget == ({"old": {}} if value is None else value) + assert token.end_user_max_budget == 9 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cached", ("valid", "invalid", None, 0)) +@pytest.mark.parametrize("global_id,key_id", ((None, None), ("", None), ("global", None), (None, ""))) +@pytest.mark.parametrize("end_found,user_found", ((False, False), (True, False), (False, True))) +@pytest.mark.parametrize("identity", ("eu", "email@example.test")) +async def test_validation_uses_raw_row_presence_and_preserves_unknown_ids_with_defaults( + native: bool, + cached: object, + global_id: str | None, + key_id: str | None, + end_found: bool, + user_found: bool, + identity: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=f"end_user_validation:{identity}", value=cached) + end: Final = AsyncMock(return_value=LiteLLM_EndUserTable(user_id=identity, blocked=False) if end_found else None) + user: Final = AsyncMock(return_value=LiteLLM_UserTable(user_id=identity) if user_found else None) + monkeypatch.setattr(auth_checks, "get_end_user_object", end) + monkeypatch.setattr(auth_reference, "get_end_user_object", end) + monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) + function: Final = ( + adapter.resolve_and_validate_end_user_id if native else auth_reference.resolve_and_validate_end_user_id + ) + result: Final = await function(identity, Mock(), cache, key_end_user_budget_id=key_id) + valid: Final = cached == "valid" or (cached != "invalid" and (end_found or user_found)) + assert result == (identity if valid or bool(global_id) or key_id is not None else None) + assert end.await_count == (0 if cached in ("valid", "invalid") else 1) + assert user.await_count == (0 if cached in ("valid", "invalid") or end_found else 1) + if user.await_count: + assert user.await_args is not None + assert user.await_args.kwargs["user_email"] == (identity if "@" in identity else None) + assert await cache.async_get_cache(key=f"end_user_validation:{identity}") == ("valid" if valid else "invalid") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("enabled,db", ((False, False), (False, True), (True, False))) +async def test_validation_disabled_or_without_db_skips_all_lookups( + native: bool, + identity: str | None, + enabled: bool, + db: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", enabled) + cache: Final = Mock(spec=UserApiKeyCache) + read: Final = AsyncMock(side_effect=AssertionError("cache must not be read")) + cache.async_get_cache = read + function: Final = ( + adapter.resolve_and_validate_end_user_id if native else auth_reference.resolve_and_validate_end_user_id + ) + assert await function(identity, Mock() if db else None, cache) == identity + read.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row_mode", ("missing", "empty", "budget")) +@pytest.mark.parametrize("initial_max", (None, 0, 7)) +@pytest.mark.parametrize( + "key_id,global_id,found", + ((None, None, False), ("key", None, True), (None, "global", True), ("key", "global", False)), +) +async def test_lookup_uses_row_before_default_and_only_fills_missing_max_budget( + native: bool, + row_mode: str, + initial_max: int | None, + key_id: str | None, + global_id: str | None, + found: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id) + budget: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=5, tpm_limit=0, model_max_budget={}) + row: Final = ( + LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if row_mode == "budget" else None) + if row_mode != "missing" + else None + ) + end: Final = AsyncMock(return_value=row) + default: Final = AsyncMock(return_value=budget if found else None) + monkeypatch.setattr(auth_checks, "get_end_user_object", end) + monkeypatch.setattr(auth_reference, "get_end_user_object", end) + monkeypatch.setattr(legacy, "get_end_user_object", end) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", default) + token: Final = UserAPIKeyAuth( + end_user_id="eu", end_user_max_budget=initial_max, end_user_tpm_limit=9, metadata={"end_user_budget_id": key_id} + ) + result: Final = ( + await adapter.lookup_end_user_and_apply_budget( + token, "/chat/completions", None, Mock(), UserApiKeyCache(), None + ) + if native + else await legacy._lookup_end_user_and_apply_budget( # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # Phase 1 untyped reference boundary + token, "/chat/completions", None, Mock(), UserApiKeyCache(), None + ) # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # Phase 1 untyped reference boundary + ) + actual_token, actual_row = result + assert actual_token is token and actual_row is row + default_used: Final = row is None and (key_id is not None or global_id is not None) and found + assert token.end_user_max_budget == (5 if default_used and initial_max is None else initial_max) + assert token.end_user_tpm_limit == (0 if row_mode == "budget" or default_used else 9) + assert token.end_user_model_max_budget == ({} if row_mode == "budget" or default_used else None) + assert end.await_count == 1 + assert default.await_count == ( + 0 + if row is not None + else (int(key_id is not None) + int(global_id is not None and (key_id is None or not found))) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_failure", (False, True)) +async def test_lookup_io_failure_does_not_select_defaults_and_native_failure_is_terminal( + native_failure: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure: Final = NativeAuthRuntimeError("native") if native_failure else RuntimeError("io") + load: Final = AsyncMock(side_effect=failure) + default: Final = AsyncMock(side_effect=AssertionError("no defaults after lookup error")) + monkeypatch.setattr(auth_checks, "get_end_user_object", load) + monkeypatch.setattr(auth_reference, "get_end_user_object", load) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", default) + token: Final = UserAPIKeyAuth(end_user_id="eu", metadata={"end_user_budget_id": "key"}) + if native_failure: + with pytest.raises(NativeAuthRuntimeError) as raised: + await adapter.lookup_end_user_and_apply_budget(token, "", None, Mock(), UserApiKeyCache(), None) + assert raised.value is failure + else: + assert await adapter.lookup_end_user_and_apply_budget(token, "", None, Mock(), UserApiKeyCache(), None) == ( + token, + None, + ) + default.assert_not_called() + + +@pytest.mark.asyncio +async def test_native_validation_failure_stops_before_user_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + failure: Final = NativeAuthRuntimeError("native") + monkeypatch.setattr(auth_checks, "get_end_user_object", AsyncMock(side_effect=failure)) + monkeypatch.setattr(auth_reference, "get_end_user_object", AsyncMock(side_effect=failure)) + user: Final = AsyncMock() + monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) + with pytest.raises(NativeAuthRuntimeError) as raised: + await adapter.resolve_and_validate_end_user_id("eu", Mock(), UserApiKeyCache()) + assert raised.value is failure + user.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("found", (False, True)) +async def test_default_budget_cache_miss_reads_database_and_caches_only_found_rows( + native: bool, + found: bool, +) -> None: + cache: Final = UserApiKeyCache() + row: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=0) if found else None + read: Final = AsyncMock(return_value=row) + db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read))) + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget + assert await function(db, cache, budget_id="budget") == row + read.assert_awaited_once_with(where={"budget_id": "budget"}) + assert await cache.async_get_cache(key="default_end_user_budget:budget", model_type=LiteLLM_BudgetTable) == row + if found: + assert await function(db, cache, budget_id="budget") == row + read.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_failure", (False, True)) +async def test_default_budget_database_failure_keeps_native_fault_terminal(native_failure: bool) -> None: + failure: Final = NativeAuthRuntimeError("native") if native_failure else RuntimeError("io") + read: Final = AsyncMock(side_effect=failure) + db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read))) + if native_failure: + with pytest.raises(NativeAuthRuntimeError) as raised: + await adapter.get_default_end_user_budget(db, UserApiKeyCache(), budget_id="budget") + assert raised.value is failure + else: + assert await adapter.get_default_end_user_budget(db, UserApiKeyCache(), budget_id="budget") is None + read.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_default_budget_cache_failure_is_not_swallowed(native: bool) -> None: + failure: Final = RuntimeError("cache unavailable") + cache: Final = Mock(spec=UserApiKeyCache) + read: Final = AsyncMock(side_effect=failure) + cache.async_get_cache = read + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget + with pytest.raises(RuntimeError) as raised: + await function(Mock(), cache, budget_id="budget") + assert raised.value is failure + read.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own_failure", (False, True)) +@pytest.mark.parametrize("user_found", (False, True)) +async def test_validation_lookup_errors_select_secondary_identity_read( + native: bool, + own_failure: bool, + user_found: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + own: Final = AsyncMock(side_effect=RuntimeError("end-user read")) if own_failure else AsyncMock(return_value=None) + user: Final = ( + AsyncMock(return_value=LiteLLM_UserTable(user_id="eu")) + if user_found + else AsyncMock(side_effect=RuntimeError("user read")) + ) + monkeypatch.setattr(auth_checks, "get_end_user_object", own) + monkeypatch.setattr(auth_reference, "get_end_user_object", own) + monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) + function: Final = adapter.end_user_id_exists_in_db if native else auth_reference._end_user_id_exists_in_db # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + assert await function("eu", Mock(), UserApiKeyCache()) is user_found + own.assert_awaited_once() + user.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_lookup_budget_denial_preserves_exception_identity(native: bool, monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = litellm.BudgetExceededError(current_cost=2, max_budget=1) + read: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(auth_checks, "get_end_user_object", read) + monkeypatch.setattr(auth_reference, "get_end_user_object", read) + monkeypatch.setattr(legacy, "get_end_user_object", read) + invocation: Final = ( + adapter.lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None) + if native + else legacy._lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # Phase 1 untyped reference boundary + ) + with pytest.raises(litellm.BudgetExceededError) as raised: + await invocation + assert raised.value is failure + read.assert_awaited_once() diff --git a/tests/unit/proxy/auth/test_key_budget_context_adapter.py b/tests/unit/proxy/auth/test_key_budget_context_adapter.py new file mode 100644 index 00000000000..843f74a30f5 --- /dev/null +++ b/tests/unit/proxy/auth/test_key_budget_context_adapter.py @@ -0,0 +1,108 @@ +import math +from datetime import datetime, timezone +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm._auth import budget_context as core +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth import key_budget_context_adapter as adapter +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy +from tests.unit.proxy.auth import phase1_budget_throttle_reference as budget_throttle + +_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) + + +@pytest.mark.parametrize("native", (False, True), ids=("reference", "native")) +@pytest.mark.parametrize( + "increase", (None, False, True, 0, -3, 3, float("nan"), float("inf"), "", "3", [], [3], {}, {"n": 3}) +) +@pytest.mark.parametrize( + "expiry", ("2024-12-31", "2025-01-01T00:00:00+00:00", "2025-01-01T00:00:00.000001Z", "2026-01-01T09:00:00+09:00") +) +@pytest.mark.parametrize("limit", (None, 0, 10)) +def test_temporary_budget_reference_native( + native: bool, monkeypatch: pytest.MonkeyPatch, increase: object, expiry: str, limit: float | None +) -> None: + now: Final = Mock(return_value=_NOW) + clock: Final = Mock(wraps=datetime, now=now) + monkeypatch.setattr(adapter, "datetime", clock) + monkeypatch.setattr(legacy, "datetime", clock) + token: Final = UserAPIKeyAuth( + max_budget=limit, metadata={"temp_budget_increase": increase, "temp_budget_expiry": expiry} + ) + fn: Final = ( + adapter.update_key_budget_with_temp_budget_increase + if native + else legacy._update_key_budget_with_temp_budget_increase # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + active: Final = expiry.startswith("2026") or expiry == "2025-01-01T00:00:00.000001Z" + changes: Final = limit is not None and active and bool(increase) + if changes and not isinstance(increase, (int, float)): + with pytest.raises(TypeError, match="unsupported operand type"): + fn(token) + return + result: Final = fn(token) + assert (result is token) is not changes + if changes: + assert limit is not None and isinstance(increase, (int, float)) + expected: Final = limit + increase + assert result.max_budget is not None + assert result.max_budget == expected or (math.isnan(result.max_budget) and math.isnan(expected)) + assert token.max_budget == limit + assert now.call_count == (0 if limit is None else 1) + + +@pytest.mark.parametrize("native", (False, True), ids=("reference", "native")) +@pytest.mark.parametrize("expiry", (None, 5, "", "invalid")) +@pytest.mark.parametrize("increase_present", (False, True)) +@pytest.mark.parametrize("limit", (None, 1)) +def test_temporary_metadata_error_timing( + native: bool, monkeypatch: pytest.MonkeyPatch, expiry: object, increase_present: bool, limit: float | None +) -> None: + token: Final = UserAPIKeyAuth( + max_budget=limit, + metadata={"temp_budget_expiry": expiry, **({"temp_budget_increase": 0} if increase_present else {})}, + ) + now: Final = Mock(return_value=_NOW) + clock: Final = Mock(wraps=datetime, now=now) + monkeypatch.setattr(adapter, "datetime", clock) + monkeypatch.setattr(legacy, "datetime", clock) + fn: Final = ( + adapter.update_key_budget_with_temp_budget_increase + if native + else legacy._update_key_budget_with_temp_budget_increase # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + if limit is not None and increase_present: + with pytest.raises(ValueError if isinstance(expiry, str) else TypeError): + fn(token) + else: + assert fn(token) is token + now.assert_not_called() + + +@pytest.mark.parametrize("pct", (None, False, True, "", "0.5", 0, -0.5, 0.25, 1, 2, float("nan"), float("inf"))) +@pytest.mark.parametrize("opt_in", (None, False, True, 1, "true")) +@pytest.mark.parametrize("limit", (None, 0, 1)) +def test_throttle_reference_native( + monkeypatch: pytest.MonkeyPatch, pct: object, opt_in: object, limit: int | None +) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", pct) + token: Final = UserAPIKeyAuth(metadata={"throttle_on_budget_exceeded": opt_in}, tpm_limit=limit) + assert adapter.budget_throttle_percentage() == budget_throttle.budget_throttle_percentage() + assert adapter.should_throttle_budget_exceeded(token) == budget_throttle.should_throttle_budget_exceeded(token) + expected: Final = ( + budget_throttle.budget_throttle_percentage() if budget_throttle.should_throttle_budget_exceeded(token) else None + ) + assert adapter.inspect_key_budget_throttle(token) == expected + + +@pytest.mark.parametrize("tpm", (None, 0, 1)) +@pytest.mark.parametrize("rpm", (None, 0, 1)) +@pytest.mark.parametrize("opt_in", (False, True)) +def test_throttle_native_selects_config_read(tpm: int | None, rpm: int | None, opt_in: bool) -> None: + assert core.throttle(core.ThrottleFacts(opt_in, tpm, rpm)).read_percentage is ( + opt_in and (tpm is not None or rpm is not None) + ) diff --git a/tests/unit/proxy/auth/test_managed_model_adapter.py b/tests/unit/proxy/auth/test_managed_model_adapter.py new file mode 100644 index 00000000000..d96e57ae375 --- /dev/null +++ b/tests/unit/proxy/auth/test_managed_model_adapter.py @@ -0,0 +1,84 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ModelAccessDeniedProxyException, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling +from litellm.proxy.auth import managed_model_adapter +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from litellm.types.agents import AgentResponse + + +@pytest.mark.asyncio +@pytest.mark.parametrize("models", (None, "model", (), ("allowed",))) +async def test_managed_model_grants_preserve_public_denial(models: object) -> None: + token: Final = UserAPIKeyAuth().model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": models} + ) + } + ) + cache: Final = UserApiKeyCache() + authorize: Final = managed_model_adapter.can_agent_call_model + if models == ("allowed",): + with pytest.raises(ModelAccessDeniedProxyException) as denied: + await authorize("denied", token, None, None, cache, ProxyLogging(user_api_key_cache=cache)) + assert ( + denied.value.internal_message + == "agent not allowed to access model. This agent can only access models=['allowed']. Tried to access denied" + ) + else: + with pytest.raises(HTTPException) as denied: + await authorize("denied", token, None, None, cache, ProxyLogging(user_api_key_cache=cache)) + assert (denied.value.status_code, denied.value.detail) == (403, "This agent has no model grants") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("grants", ((), ("allowed",))) +async def test_agent_empty_ceiling_and_model_denial_envelopes(grants: tuple[str, ...]) -> None: + async def resolve(_agent: str) -> AgentAccessGroupCeiling: + return AgentAccessGroupCeiling(("group",), frozenset(grants), frozenset(), frozenset()) + + authorize: Final = managed_model_adapter.check_agent_access_group_model_access + with pytest.raises(ModelAccessDeniedProxyException) as denied: + await authorize("denied", UserAPIKeyAuth(agent_id="agent"), None, resolve) + assert denied.value.internal_message == ( + "agent agent access groups grant no models" + if not grants + else "agent not allowed to access model. This agent can only access models=['allowed']. Tried to access denied" + ) + + +@pytest.mark.asyncio +async def test_managed_alias_dispatch_before_scope_check() -> None: + token: Final = UserAPIKeyAuth(team_model_aliases={"alias": "allowed"}).model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": ("allowed",)} + ), + } + ) + assert await managed_model_adapter.can_agent_call_model("alias", token, None, None, None, None) is None + + +@pytest.mark.parametrize( + ("models", "expected"), + ( + (None, frozenset[str]()), + ("model", frozenset[str]()), + ((), frozenset[str]()), + (("a", 42, "a", "b"), frozenset(("a", "b"))), + ), +) +def test_managed_listing_filters_non_string_grants(models: object, expected: frozenset[str]) -> None: + token: Final = UserAPIKeyAuth().model_copy( + update={ + "managed_agent_policy": AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, object_permission={"models": models} + ) + } + ) + assert managed_model_adapter.managed_agent_model_grants(token) == expected diff --git a/tests/unit/proxy/auth/test_model_budget_adapter.py b/tests/unit/proxy/auth/test_model_budget_adapter.py new file mode 100644 index 00000000000..0393b018e5d --- /dev/null +++ b/tests/unit/proxy/auth/test_model_budget_adapter.py @@ -0,0 +1,138 @@ +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass, field +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +import litellm +from litellm._auth import model_budget_data +from litellm._auth.model_budget_data import ModelBudgetEntity +from litellm.caching.caching import DualCache +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.auth.model_budget_adapter import ( + ModelBudgetIO, + enforce_model_budget, + read_model_budget, + resolve_model_budget, +) +from litellm.proxy.hooks import model_max_budget_limiter as reference +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindow + +from .budget_reference import model_allowed + + +@dataclass(frozen=True, slots=True) +class Windows: + value: SharedWindow | None = None + read: AsyncMock = field(default_factory=AsyncMock) + + async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: + raise AssertionError("request auth must not open a window") + + async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: + raise AssertionError("request auth must not write a counter") + + +def groups(model: str, team_id: str | None) -> frozenset[str]: + return frozenset(("group",)) + + +@pytest.mark.parametrize( + ("model", "budgets"), + ( + ( + "provider/model", + {"model": {"max_budget": 1, "budget_duration": "1h"}, "*": {"max_budget": 2, "budget_duration": "1h"}}, + ), + ("model", {"model": {"budget_duration": "bad"}, "*": {"max_budget": 2, "budget_duration": "1h"}}), + ("model", {"model": {}, "group": {"max_budget": 2, "budget_duration": "1h"}}), + ( + "model", + {"m*": {"max_budget": 2, "budget_duration": "1h"}, "group": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ( + "mødél", + {"m*é*": {"max_budget": 2, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ( + "model", + {"model": {"max_budget": None, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ), +) +def test_resolution_matches_production_reference(model: str, budgets: Mapping[str, object]) -> None: + io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0) + old: Final = reference.resolve_model_budget(model, budgets, groups) + result: Final = resolve_model_budget(model, budgets, io, None) + assert result.read_groups is False + assert (result.entry.name if result.entry else None) == (old.budget_model if old else None) + assert (result.entry.max_budget if result.entry else None) == (old.budget_config.max_budget if old else None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entity", "scope"), + ( + ("Key", Litellm_EntityType.KEY), + ("User", Litellm_EntityType.USER), + ("Team", Litellm_EntityType.TEAM), + ("TeamMember", Litellm_EntityType.TEAM_MEMBER), + ("EndUser", Litellm_EntityType.END_USER), + ), +) +@pytest.mark.parametrize( + ("local", "shared", "now", "limit"), + ((0, 2, 100, 1), (1, 0, 100, 1), (0, 2, 3700, 1), (0, 0, 100, 0), (0, 2, 100, -1), (0, 2, 100, None)), +) +async def test_counter_and_shared_window_match_production_reference( + entity: ModelBudgetEntity, scope: Litellm_EntityType, local: float, shared: float, now: float, limit: float | None +) -> None: + cache: Final = DualCache() + key: Final = reference.model_budget_spend_cache_key(scope, "id", "model", "1h") + cache.set_cache(key, local) + windows: Final = Windows(read=AsyncMock(return_value=SharedWindow(100, shared))) + budgets: Final = {"model": {"max_budget": limit, "budget_duration": "1h"}} + old: Final = reference._PROXY_VirtualKeyModelMaxBudgetLimiter( + cache, clock=lambda: now, model_access_groups=groups, shared_windows=windows + ) + io: Final = ModelBudgetIO(cache, groups, windows, 0, clock=lambda: now, monotonic=lambda: 0) + before: Final = await observed_model_result(model_allowed(old, scope, "id", None, budgets, "model", "exceeded")) + after: Final = await observed_model_result( + enforce_model_budget(io, entity, "id", None, "model", budgets, "exceeded") + ) + assert after == before + + +async def observed_model_result(call: Awaitable[bool]) -> tuple[object, ...]: + try: + return (await call,) + except litellm.BudgetExceededError as error: + return (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id) + + +@pytest.mark.asyncio +async def test_legacy_counter_is_added_to_the_current_counter_before_shared_read() -> None: + cache: Final = DualCache() + cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "model", "1h"), 0.4) + cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "provider/model", "1h"), 0.6) + windows: Final = Windows(read=AsyncMock(side_effect=AssertionError("already over local limit"))) + io: Final = ModelBudgetIO(cache, groups, windows, 0, monotonic=lambda: 1) + reading: Final = await read_model_budget( + io, "Key", "id", None, "provider/model", {"model": {"max_budget": 1, "budget_duration": "1h"}} + ) + assert reading.snapshot.local_spend == 0.4 + assert reading.snapshot.legacy_spend == 0.6 + assert model_budget_data.check(reading.snapshot).denied is True + windows.read.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shared_read_failure_is_terminal_and_not_replayed() -> None: + failure: Final = RuntimeError("database unavailable") + windows: Final = Windows(read=AsyncMock(side_effect=failure)) + io: Final = ModelBudgetIO(DualCache(), groups, windows, 0, monotonic=lambda: 0) + with pytest.raises(RuntimeError, match="database unavailable") as raised: + await read_model_budget(io, "User", "id", None, "model", {"model": {"max_budget": 1, "budget_duration": "1h"}}) + assert raised.value is failure + assert windows.read.await_count == 1 diff --git a/tests/unit/proxy/auth/test_model_projection_adapter.py b/tests/unit/proxy/auth/test_model_projection_adapter.py new file mode 100644 index 00000000000..60817b2fc82 --- /dev/null +++ b/tests/unit/proxy/auth/test_model_projection_adapter.py @@ -0,0 +1,203 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from queue import SimpleQueue +from typing import Final, Literal, cast + +import pytest + +import litellm +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import grouped_model_access +from litellm.proxy._types import ModelAccessDeniedProxyException +from litellm.proxy.auth import model_access_adapter, model_projection_adapter +from litellm.router_api import RegistryReader +from tests.unit._auth.support import model_registry_reference as reference + + +@dataclass(frozen=True, slots=True) +class _Info: + team_id: str | None = "team" + team_public_model_name: str | None = "public" + + +@dataclass(frozen=True, slots=True) +class _Deployment: + model_info: _Info = field(default_factory=_Info) + + +@dataclass(frozen=True, slots=True) +class _Registry: + events: SimpleQueue[tuple[str, str | None]] + entries: tuple[str, ...] = ("registered",) + router_aliases: Mapping[str, str] = field(default_factory=lambda: {"alias": "target"}) + groups: Mapping[str, tuple[str, ...]] = field(default_factory=lambda: {"group": ("m",)}) + deployment: _Deployment | None = field(default_factory=_Deployment) + fail: str | None = None + + def observe(self, kind: str, model: str | None = None) -> None: + self.events.put((kind, model)) + if self.fail == kind: + raise RuntimeError(f"failed {kind}") + + def names(self) -> list[str]: # mutable-ok: the legacy sentinel resolver concatenates the registry's concrete list + self.observe("names") + return list(self.entries) + + def aliases(self) -> Mapping[str, str]: + self.observe("aliases") + return self.router_aliases + + def resolve(self, model: str, *, kind: Literal["alias"]) -> str | None: + self.observe("resolve", model) + return self.router_aliases.get(model) + + def access_groups(self, model_name: str, team_id: str | None) -> Mapping[str, tuple[str, ...]]: + self.observe("groups", model_name) + return self.groups + + def deployments(self, model: str, *, source: Literal["first_in_group"]) -> _Deployment | None: + self.observe("deployment", model) + return self.deployment + + +@dataclass(frozen=True, slots=True) +class _Router: + view: _Registry + + def registry(self) -> _Registry: + self.view.observe("registry") + return self.view + + +def _reader(view: _Registry) -> RegistryReader: + return cast(RegistryReader, _Router(view)) # noqa: TID251 # cast-ok: typed IO double implements only the registry methods consumed by auth + + +def _events(events: SimpleQueue[tuple[str, str | None]]) -> tuple[tuple[str, str | None], ...]: + return tuple(events.get() for _ in range(events.qsize())) + + +def _run(native: bool, model: str | tuple[str, ...], grants: tuple[str, ...], registry: _Registry) -> str: + router: Final = _reader(registry) + aliases: Final = {"z": "model_name_team_deploy", "a": "model_name_team_deploy", "team-alias": "target"} + try: + if native: + model_access_adapter.can_object_call_model(model, router, grants, aliases, "team", {"key": "target"}) + else: + reference.can_object_call_model( + model if isinstance(model, str) else list(model), + router, + list(grants), + aliases, + "team", + {"key": "target"}, + ) + return "allowed" + except ModelAccessDeniedProxyException as error: + return str((error.message, error.internal_message, error.type, error.param, error.code)) + + +@pytest.mark.parametrize( + "model", ("m", "alias", "global", "key", "team-alias", "model_name_team_deploy", ("m", "m"), ("denied", "m"), ()) +) +@pytest.mark.parametrize( + "grants", + ( + (), + ("m",), + ("target",), + ("group",), + ("all-team-models",), + ("all-team-models", "other"), + ("*",), + ("other/*", "provider/*"), + ("a",), + ("public",), + ), +) +def test_registry_observation_order_and_denials_match_legacy( + monkeypatch: pytest.MonkeyPatch, model: str | tuple[str, ...], grants: tuple[str, ...] +) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + registry: Final = _Registry(events) + + def provider(model: str) -> tuple[str, str, None, None]: + events.put(("provider", model)) + return "leaf", "provider", None, None + + monkeypatch.setattr(litellm, "get_llm_provider", provider) + monkeypatch.setattr(litellm, "model_alias_map", {"global": "target"}) + expected: Final = _run(False, model, grants, registry) + expected_events: Final = _events(events) + assert _run(True, model, grants, registry) == expected + assert _events(events) == expected_events + + +@pytest.mark.parametrize("native", (False, True)) +@pytest.mark.parametrize("failed", ("aliases", "groups", "names", "deployment")) +def test_registry_failure_stops_before_later_reads(native: bool, failed: str) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + registry: Final = _Registry(events, fail=failed) + with pytest.raises(RuntimeError, match=f"failed {failed}") as error: + _run(native, "model_name_team_deploy", ("all-team-models",), registry) + assert not isinstance(error.value, NativeAuthFailure) + assert _events(events)[-1][0] == failed + + +@pytest.mark.parametrize("native", (False, True)) +def test_depth_denial_never_reads_registry(native: bool) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + router: Final = _reader(_Registry(events, fail="registry")) + policy: Final = model_access_adapter.can_object_call_model if native else reference.can_object_call_model + with pytest.raises(Exception, match="max fallback depth exceeded"): + policy("m", router, (), fallback_depth=100) + assert _events(events) == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group_models", (("other",), ("other/*",), ("m",))) +async def test_group_retry_gets_fresh_registry_reads_even_for_duplicate_grants(group_models: tuple[str, ...]) -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + router: Final = _reader(_Registry(events)) + facts: Final = model_access_adapter.project_model_facts("m", router, ("other",)) + assert _events(events) == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) + scope: Final = await model_projection_adapter.read_grouped_scope( + GroupedModelScope(facts, ("group",), group_models), router, merge=False + ) + assert _events(events)[:4] == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) + assert grouped_model_access(scope.facts, scope.group_ids, scope.group_models).kind == ( + "Allow" if group_models == ("m",) else "Deny" + ) + + +def test_native_provider_failure_is_not_a_missing_provider(monkeypatch: pytest.MonkeyPatch) -> None: + def provider(model: str) -> tuple[str, str, None, None]: + raise NativeAuthRuntimeError("terminal provider failure") + + monkeypatch.setattr(litellm, "get_llm_provider", provider) + with pytest.raises(NativeAuthRuntimeError, match="terminal provider failure"): + model_access_adapter.can_object_call_model("model", None, ("provider/*",)) + + +@pytest.mark.parametrize("native", (False, True)) +def test_compaction_parent_is_read_only_after_child_denial(native: bool) -> None: + from litellm.router_strategy.complexity_router.context_compaction import native_compaction_call + + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + with native_compaction_call("parent", "child"): + assert _run(native, "child", ("parent",), _Registry(events)) == "allowed" + assert _events(events) == ( + ("registry", None), + ("aliases", None), + ("registry", None), + ("groups", "child"), + ("registry", None), + ("groups", "parent"), + ) + + +def test_large_model_lists_do_not_consume_python_recursion_depth() -> None: + events: Final[SimpleQueue[tuple[str, str | None]]] = SimpleQueue() + assert _run(True, ("m",) * 180, ("m",), _Registry(events)) == "allowed" + assert _events(events) == (("registry", None), ("aliases", None), ("registry", None), ("groups", "m")) * 180 diff --git a/tests/unit/proxy/auth/test_org_route_adapter.py b/tests/unit/proxy/auth/test_org_route_adapter.py new file mode 100644 index 00000000000..ebae5ba615e --- /dev/null +++ b/tests/unit/proxy/auth/test_org_route_adapter.py @@ -0,0 +1,101 @@ +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.proxy._types import LiteLLM_UserTable, ProxyException +from litellm.proxy.auth import org_route_adapter as native + +from .test_route_role_adapter import observed + +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) + + +def reference_admin(body: Mapping[str, object], user: LiteLLM_UserTable | None) -> bool: + if user is None or user.organization_memberships is None: + return False + singular: Final = body.get("organization_id") + plural: Final = body.get("organizations") + candidates: Final = ( + *((singular,) if singular is not None else ()), + *(_OBJECTS.validate_python(plural) if isinstance(plural, list) else ()), + ) + admins: Final = frozenset( + item.organization_id for item in user.organization_memberships + if item.user_role == "org_admin" and item.organization_id is not None # pyright: ignore[reportUnnecessaryComparison] # preserve Phase 1 constructed membership handling + ) + return bool(candidates) and all(candidate in admins for candidate in candidates) + + +def user(kind: str) -> LiteLLM_UserTable | None: + if kind == "absent": + return None + return LiteLLM_UserTable.model_validate({ + "user_id": "u", + "user_role": "proxy_admin" if kind == "proxy" else "internal_user", + "organization_memberships": None if kind == "none" else ( + [] if kind == "empty" else [{ + "organization_id": "org", "user_role": "org_admin", "user_id": "u", + "created_at": "2025-01-01T00:00:00Z", "updated_at": "2025-01-01T00:00:00Z", + }] + ), + }) + + +@pytest.mark.parametrize("kind", ("absent", "none", "empty", "member", "proxy")) +@pytest.mark.parametrize("singular", (None, "org", "other", "", 0, False, 1.5, [], {}, ())) +@pytest.mark.parametrize("plural", (None, [], ["org"], ["other", []], ["org", []], [0], "org")) +def test_raw_singular_and_plural_keep_membership_failure_and_short_circuit( + kind: str, singular: object, plural: object +) -> None: + member: Final = user(kind) + body: Final = {"organization_id": singular, "organizations": plural} + assert observed(lambda: native.user_is_org_admin(body, member)) == observed(lambda: reference_admin(body, member)) + + +@pytest.mark.parametrize("identifier", (0, False, 1.5, (), "")) +@pytest.mark.parametrize("route", ("/organization/delete", "/team/new")) +def test_nonstring_present_identifier_is_not_changed_to_missing(identifier: object, route: str) -> None: + with pytest.raises(ProxyException) as captured: + native.organization_role_based_access_check({"organization_id": identifier}, user("member"), route) + expected: Final = ( + f"You do not have a role within the selected organization. Passed organization_id: {identifier}. Please contact the organization admin to request access." + if route == "/organization/delete" else + f"You do not have the required role to call {route}. Your role is None in Organization {identifier}" + ) + assert captured.value.message == expected + assert captured.value.code == "401" + + +@pytest.mark.parametrize("identifier", ([], {})) +@pytest.mark.parametrize("route", ("/organization/delete", "/team/new")) +def test_unhashable_identifier_raises_only_when_membership_lookup_is_selected(identifier: object, route: str) -> None: + with pytest.raises(TypeError, match="unhashable type"): + native.organization_role_based_access_check({"organization_id": identifier}, user("member"), route) + assert native.organization_role_based_access_check({"organization_id": identifier}, user("proxy"), route) is None + assert native.organization_role_based_access_check({"organization_id": identifier}, user("absent"), route) is None + with pytest.raises(ProxyException, match="not a member of any organization"): + native.organization_role_based_access_check({"organization_id": identifier}, user("none"), "/organization/delete") + + +class UnreadableBody(Mapping[str, object]): + def __getitem__(self, key: str) -> object: + raise AssertionError(f"unexpected body read: {key}") + + def __iter__(self) -> Iterator[str]: + return iter(()) + + def __len__(self) -> int: + return 0 + + +class UnreadableHash: + def __hash__(self) -> int: + raise AssertionError("unexpected candidate hash") + + +def test_no_user_never_reads_request_and_denied_candidate_never_hashes_later_candidate() -> None: + assert native.user_is_org_admin(UnreadableBody(), None) is False + assert native.organization_role_based_access_check(UnreadableBody(), None, "/organization/delete") is None + assert native.user_is_org_admin({"organizations": ["not-allowed", UnreadableHash()]}, user("member")) is False diff --git a/tests/unit/proxy/auth/test_personal_key_team_model_access.py b/tests/unit/proxy/auth/test_personal_key_team_model_access.py index 38f34adb48e..31b5c304bfc 100644 --- a/tests/unit/proxy/auth/test_personal_key_team_model_access.py +++ b/tests/unit/proxy/auth/test_personal_key_team_model_access.py @@ -320,14 +320,12 @@ async def test_model_listing_loads_each_team_once(monkeypatch: pytest.MonkeyPatc @pytest.mark.asyncio async def test_team_check_error_grants_nothing_for_that_team(monkeypatch: pytest.MonkeyPatch, teams: None) -> None: _enable(monkeypatch) - check_team: Final = auth_checks.can_team_access_model - async def _failing_for_team_b(model: str, team_object: LiteLLM_TeamTableCachedObj, **_: object) -> bool: - if team_object.team_id == "team-b": - raise RuntimeError("access group lookup failed") - return await check_team(model=model, team_object=team_object, llm_router=_ROUTER) + async def _failing_for_team_b(*, team_id: str, **_: object) -> None: + if team_id == "team-b": + raise RuntimeError("membership lookup failed") - monkeypatch.setattr(auth_checks, "can_team_access_model", _failing_for_team_b) + monkeypatch.setattr(auth_checks, "get_team_membership", _failing_for_team_b) owner: Final = LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]) token: Final = UserAPIKeyAuth(token="k1", user_id="u1") assert await _common_checks("gpt-a", owner, token) is True diff --git a/tests/unit/proxy/auth/test_route_enable_adapter.py b/tests/unit/proxy/auth/test_route_enable_adapter.py new file mode 100644 index 00000000000..6f3f48d321c --- /dev/null +++ b/tests/unit/proxy/auth/test_route_enable_adapter.py @@ -0,0 +1,140 @@ +from queue import SimpleQueue +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.route_enable import RouteFeature, RouteSetting +from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.auth.route_enable_adapter import RouteEnableReaders, check_route_enabled + + +@pytest.mark.parametrize("route", ("/models", "/chat/completions")) +def test_enterprise_class_exemptions_preserve_enabled_route_policy(route: str) -> None: + class EnterpriseRoutes: + LLM_API_EXEMPT_ROUTES: Final = ("/models",) + + readers: Final = RouteEnableReaders( + enterprise=lambda: EnterpriseRoutes, + setting=lambda _: RouteSetting(True, True), + secret=lambda _: True, + ) + if route == "/models": + assert check_route_enabled(route, readers=readers) is None + else: + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + assert (error.value.status_code, error.value.detail) == ( + 403, + "LLM API routes are disabled for this instance.", + ) + + +def test_missing_optional_class_exemptions_preserve_ignored_read_failure() -> None: + readers: Final = RouteEnableReaders(enterprise=lambda: object()) + assert check_route_enabled("/chat/completions", readers=readers) is None + + +@pytest.mark.parametrize( + ("route", "expected"), + ( + ("/unknown", ("import",)), + ("/models", ("import", "exemptions")), + ("/chat/completions", ("import", "exemptions", "Llm", "secret-Llm")), + ("/key/generate", ("import", "Admin", "secret-Admin")), + ), +) +def test_optional_route_reads_are_lazy(route: str, expected: tuple[str, ...]) -> None: + events: Final[SimpleQueue[str]] = SimpleQueue() + + def enterprise() -> object: + events.put("import") + return object() + + def setting(feature: RouteFeature) -> RouteSetting: + events.put(feature) + return RouteSetting(False, False) + + def secret(feature: RouteFeature) -> bool: + events.put(f"secret-{feature}") + return False + + def exemptions(_enterprise: object) -> tuple[str, ...]: + events.put("exemptions") + return ("/models",) + + check_route_enabled(route, readers=RouteEnableReaders(enterprise, setting, secret, exemptions)) + assert tuple(events.get() for _ in range(events.qsize())) == expected + + +@pytest.mark.parametrize(("route", "label"), (("/key/generate", "ADMIN"), ("/chat/completions", "LLM API"))) +@pytest.mark.parametrize("premium", (False, True)) +def test_optional_route_denial_envelopes(route: str, label: str, premium: bool) -> None: + readers: Final = RouteEnableReaders( + enterprise=lambda: object(), + setting=lambda _: RouteSetting(True, premium), + secret=lambda _: True, + exemptions=lambda _: (), + ) + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + if premium: + disabled: Final = "Management" if label == "ADMIN" else label + assert (error.value.status_code, error.value.detail) == ( + 403, + f"{disabled} routes are disabled for this instance.", + ) + else: + assert (error.value.status_code, error.value.detail) == ( + 500, + f"\U0001f6a8\U0001f6a8\U0001f6a8 DISABLING {label} ENDPOINTS is an Enterprise feature\n\U0001f6a8 {CommonProxyErrors.not_premium_user.value}", + ) + + +@pytest.mark.parametrize( + "error", + ( + ImportError("optional package"), + RuntimeError("settings unavailable"), + HTTPException(418, "read failure"), + NativeAuthRuntimeError("terminal native failure"), + ), +) +def test_optional_import_failure_policy(error: Exception) -> None: + def enterprise() -> object: + raise error + + readers: Final = RouteEnableReaders(enterprise=enterprise) + if isinstance(error, (HTTPException, NativeAuthRuntimeError)): + with pytest.raises(type(error)) as caught: + check_route_enabled("/key/generate", readers=readers) + assert caught.value is error + else: + assert check_route_enabled("/key/generate", readers=readers) is None + + +@pytest.mark.parametrize("configured", (False, True)) +@pytest.mark.parametrize("route", ("/eu_jev/v1/systemone", "/eu_jev/v1/systemone/")) +def test_configured_eu_jev_route_skips_llm_disable_setting(configured: bool, route: str) -> None: + events: Final[SimpleQueue[str]] = SimpleQueue() + + def eu_jev() -> bool: + events.put("configuration") + return configured + + def setting(feature: RouteFeature) -> RouteSetting: + events.put(feature) + return RouteSetting(True, True) + + readers: Final = RouteEnableReaders( + enterprise=lambda: object(), setting=setting, secret=lambda _: True, exemptions=lambda _: (), eu_jev=eu_jev + ) + if configured: + assert check_route_enabled(route, readers=readers) is None + assert tuple(events.get() for _ in range(events.qsize())) == ("configuration",) + else: + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + assert (error.value.status_code, error.value.detail) == (403, "LLM API routes are disabled for this instance.") + assert tuple(events.get() for _ in range(events.qsize())) == ("configuration", "Llm") diff --git a/tests/unit/proxy/auth/test_route_role_adapter.py b/tests/unit/proxy/auth/test_route_role_adapter.py new file mode 100644 index 00000000000..dce460a6b37 --- /dev/null +++ b/tests/unit/proxy/auth/test_route_role_adapter.py @@ -0,0 +1,126 @@ +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException, Request + +from litellm._auth import route_roles as core +from litellm._auth.failure import NativeAuthFailure +from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth import route_role_adapter as native +from litellm.rust_bridge import _native + +from . import route_role_reference as reference + +_ROLES: Final = (None, "", "unknown", "PROXY_ADMIN", *(role.value for role in LitellmUserRoles)) + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: type[Exception] + message: str + status: int | None + detail: object + + +def observed(call: Callable[[], object]) -> object: + try: + return call() + except Exception as error: + return Failure( + type(error), str(error), + error.status_code if isinstance(error, HTTPException) else None, + error.detail if isinstance(error, HTTPException) else None, + ) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("present", (False, True)) +def test_role_and_admin_projection_matches_reference(role: str | None, present: bool) -> None: + user: Final = LiteLLM_UserTable(user_id="u", user_role=role) if present else None + assert native.effective_user_role(role) is reference.effective_user_role(role) + assert native.get_user_role(user) is reference.get_user_role(user) + assert native.is_user_proxy_admin(user) is reference.is_user_proxy_admin(user) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("user_id", (None, "", "u")) +@pytest.mark.parametrize("requested", (None, "", "u", "other")) +def test_inside_route_admin_viewer_or_non_none_self(role: str | None, user_id: str | None, requested: str | None) -> None: + auth: Final = UserAPIKeyAuth.model_construct(user_role=role, user_id=user_id) + assert native.allowed_route_check_inside_route(auth, requested) is reference.allowed_route_check_inside_route(auth, requested) + + +@pytest.mark.parametrize("role", _ROLES) +@pytest.mark.parametrize("token_present", (False, True)) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/team/new")) +def test_api_route_role_precedes_token_and_nonadmin_check(role: str | None, token_present: bool, route: str) -> None: + user: Final = LiteLLM_UserTable(user_id="u", user_role=role) + token: Final = UserAPIKeyAuth(user_id="u") if token_present else None + request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + assert observed(lambda: native.is_api_route_allowed(route, request, {}, token, user)) == observed( + lambda: reference.is_api_route_allowed(route, request, {}, token, user) + ) + + +@pytest.mark.parametrize("role", tuple(LitellmUserRoles)) +@pytest.mark.parametrize("configured", (None, (), ("info_routes",), ("/x*",))) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/x/y", "/not-configured")) +def test_jwt_admin_team_none_default_and_empty_deny( + role: LitellmUserRoles, configured: tuple[str, ...] | None, route: str +) -> None: + config: Final = LiteLLM_JWTAuth.model_construct( + admin_allowed_routes=None if configured is None else list(configured), + team_allowed_routes=None if configured is None else list(configured), + ) + assert observed(lambda: native.allowed_routes_check(role, route, config)) == observed( + lambda: reference.allowed_routes_check(role, route, config) + ) + + +@pytest.mark.parametrize("entries", ( + (), ("/x",), ("/x*",), ("/x/{id}",), ("info_routes",), + ("/x", None), (None, "/x"), (1,), (False,), (list[object](),), (dict[str, object](),), + (set[object](),), ((list[object](),),), (("x",),), ("missing_group", "info_routes"), ("/x", list[object]()), +)) +@pytest.mark.parametrize("route", ("/x", "/x/a", "/x/", "/user/info", "")) +def test_listed_route_invalid_entries_and_short_circuit_match_reference(entries: tuple[object, ...], route: str) -> None: + assert observed(lambda: native.allowed_routes_list_check(route, entries)) == observed( + lambda: reference.allowed_routes_list_check(route, entries) + ) + + +@pytest.mark.parametrize("group", tuple(LiteLLMRoutes.__members__)) +@pytest.mark.parametrize("route", ("/chat/completions", "/user/info", "/key/a/reset_spend", "/credentials/a/b", "/x\n", "")) +def test_every_route_group_uses_starlette_template_semantics(group: str, route: str) -> None: + assert observed(lambda: native.allowed_routes_list_check(route, (group,))) == observed( + lambda: reference.allowed_routes_list_check(route, (group,)) + ) + + +@pytest.mark.parametrize("entries", (("info_routes", "info_routes"), ("/x", None, 1, ("x",)), ([],), ({},))) +def test_expansion_preserves_order_duplicates_nonstring_values_and_errors(entries: tuple[object, ...]) -> None: + assert observed(lambda: native.get_actual_routes(entries)) == observed(lambda: reference.get_actual_routes(entries)) + + +@pytest.mark.parametrize("compiled", ("^/thing/(?P[^/]+)$", "^/thing/(?P.*)$", "^/v1\\.0/$")) +@pytest.mark.parametrize("route", ("/thing/", "/thing/a", "/thing/a/b", "/thing/a\n", "/thing/a\nb", "/v1.0/", "/v1x0/")) +def test_compiled_group_regex_keeps_empty_path_literal_escaping_and_final_newline(compiled: str, route: str) -> None: + import re + + decision: Final = core.listed(core.ListedFacts(route, core.Entry("group", "str", templates=("unused",)), compiled=compiled)) + assert (decision.action == "Allow") is (re.match(compiled, route) is not None) + + +def test_native_fault_is_terminal_before_reading_first_allowed_entry(monkeypatch: pytest.MonkeyPatch) -> None: + def broken(_facts: object) -> object: + raise ValueError("decode failed") + + def entries() -> Iterable[object]: + raise AssertionError("entry read after native failure") + yield "unused" + + monkeypatch.setattr(_native, "request_auth_route_role_listed", broken) + with pytest.raises(NativeAuthFailure, match="decode failed"): + native.allowed_routes_list_check("/x", entries()) diff --git a/tests/unit/proxy/auth/test_router_override_fallback_auth.py b/tests/unit/proxy/auth/test_router_override_fallback_auth.py index d35a5851659..549b7740375 100644 --- a/tests/unit/proxy/auth/test_router_override_fallback_auth.py +++ b/tests/unit/proxy/auth/test_router_override_fallback_auth.py @@ -5,15 +5,23 @@ override is promoted to per-request kwargs after auth and lets a caller execute requests against models their API key cannot call. """ -from typing import List -from unittest.mock import AsyncMock, patch +from typing import Final, List import pytest -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm import Router + +from litellm.proxy._types import LitellmUserRoles, ModelAccessDeniedProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import fallback_target_model_name, iter_request_fallback_targets from litellm.proxy.auth.user_api_key_auth import enforce_key_and_fallback_model_access +_ROUTER: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/unit-model", "api_key": "sk-unit-test"}} + for name in ("gpt-3.5-turbo", "real-fallback", "string-fallback") + ] +) + def _fallback_model_names(fallbacks): """Model names the auth check validates for a top-level ``fallbacks`` value.""" @@ -38,9 +46,11 @@ def _key_with_models(models: List[str]) -> UserAPIKeyAuth: def test_fallback_model_names_router_config_shape(): """Router-config shape: ``[{primary: [fallback_list]}]``.""" - assert _fallback_model_names( - [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] - ) == ["gpt-4", "claude-3", "o1"] + assert _fallback_model_names([{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}]) == [ + "gpt-4", + "claude-3", + "o1", + ] def test_fallback_model_names_simple_string_shape(): @@ -58,9 +68,10 @@ def test_fallback_model_names_client_side_shape(): def test_fallback_model_names_nested_deployment_fallbacks(): """A deployment target's own nested fallback field is unrolled too.""" - assert _fallback_model_names( - [{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}] - ) == ["gpt-4", "deepseek-chat"] + assert _fallback_model_names([{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}]) == [ + "gpt-4", + "deepseek-chat", + ] def test_fallback_model_names_empty_or_none(): @@ -85,34 +96,16 @@ async def test_router_override_fallbacks_validated_against_key_allowlist(): }, } - seen_models: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen_models.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - # Both the primary model and the override-nested fallback must be - # checked against the API key's allowlist. - assert "gpt-3.5-turbo" in seen_models - assert "unauthorized-model" in seen_models + assert denied.value.code == "403" @pytest.mark.asyncio @@ -136,31 +129,16 @@ async def test_router_override_all_fallback_fields_validated(fallback_field): }, } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - assert "smuggled-model" in seen + assert denied.value.code == "403" @pytest.mark.asyncio @@ -183,31 +161,16 @@ async def test_top_level_fallback_fields_validated(fallback_field): fallback_field: [{"gpt-3.5-turbo": ["top-level-smuggled"]}], } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - assert "top-level-smuggled" in seen + assert denied.value.code == "403" @pytest.mark.asyncio @@ -229,31 +192,16 @@ async def test_nested_deployment_fallback_inner_model_validated(): ], } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): + with pytest.raises(ModelAccessDeniedProxyException) as denied: await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, route="/v1/chat/completions", request=None, llm_model_list=None, - llm_router=None, + llm_router=_ROUTER, ) - - assert "deep-smuggled-model" in seen + assert denied.value.code == "403" @pytest.mark.asyncio @@ -274,32 +222,16 @@ async def test_model_less_fallback_dict_is_skipped_never_passed_as_none(): ], } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route="/v1/chat/completions", - request=None, - llm_model_list=None, - llm_router=None, - ) - - assert None not in seen - assert seen == ["gpt-3.5-turbo", "real-fallback", "string-fallback"] + valid_token.models = ["gpt-3.5-turbo", "real-fallback", "string-fallback"] + result = await enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=_ROUTER, + ) + assert result is None @pytest.mark.asyncio @@ -312,28 +244,12 @@ async def test_router_override_without_fallbacks_does_not_break_auth(): "router_settings_override": {"num_retries": 3, "timeout": 30}, } - seen: List[str] = [] - - async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): - seen.append(model) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=fake_can_key_call_model, - ), - patch( - "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", - new=AsyncMock(), - ), - ): - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route="/v1/chat/completions", - request=None, - llm_model_list=None, - llm_router=None, - ) - - assert seen == ["gpt-3.5-turbo"] + result = await enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=_ROUTER, + ) + assert result is None diff --git a/tests/unit/proxy/auth/test_spend_counter_adapter.py b/tests/unit/proxy/auth/test_spend_counter_adapter.py new file mode 100644 index 00000000000..e576a583cb0 --- /dev/null +++ b/tests/unit/proxy/auth/test_spend_counter_adapter.py @@ -0,0 +1,49 @@ +from typing import Final, NoReturn +from unittest.mock import AsyncMock, Mock, call + +import pytest + +from litellm._auth.spend_counter_policy import Spend +from litellm.proxy.auth.spend_counter_adapter import SpendIO, resolve_spend + + +def unavailable() -> NoReturn: + raise RuntimeError("unverified spend") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback, window, floor, strict, expected, calls", + ( + (0.0, False, None, (False, False), 2.0, (call.read_strict(), call.read_strict())), + (0.0, False, 4.0, (True,), 4.0, (call.read_strict(), call.read_floor(), call.repair(4.0))), + (5.0, False, None, (), 5.0, (call.read_floor(),)), + (0.0, True, 1.0, (), 2.0, (call.read_floor(),)), + ), +) +async def test_native_plan_preserves_floor_and_flag_read_order( + fallback: float, + window: bool, + floor: float | None, + strict: tuple[bool, ...], + expected: float, + calls: tuple[object, ...], +) -> None: + effects: Final = Mock() + effects.attach_mock(Mock(side_effect=strict), "read_strict") + effects.attach_mock(AsyncMock(return_value=floor), "read_floor") + effects.attach_mock(AsyncMock(), "repair") + io: Final = SpendIO(effects.read_strict, effects.read_floor, effects.repair, unavailable) + result: Final = await resolve_spend(Spend(2.0, False, fallback, 10.0, window, False), io) + assert (result, tuple(effects.mock_calls)) == (expected, calls) + + +@pytest.mark.asyncio +async def test_unverified_spend_reads_the_latest_fail_closed_flag() -> None: + strict: Final = Mock(side_effect=(False, True)) + floor: Final = AsyncMock() + with pytest.raises(RuntimeError, match="unverified spend"): + await resolve_spend( + Spend(2.0, False, 0.0, 10.0, False, False), SpendIO(strict, floor, AsyncMock(), unavailable) + ) + assert (strict.call_count, floor.await_count) == (2, 0) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index e8ac0e96760..df9a5a1bca8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -15,7 +15,7 @@ from litellm._logging import verbose_proxy_logger import json import logging import litellm -from litellm._auth.python_core import PythonCore +from tests.unit._auth.reference_core import PythonCore from litellm.proxy.auth.user_api_key_auth import ( _enforce_model_budgets, _ModelBudgetEffects, @@ -1404,14 +1404,15 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ calls = [] class _Limiter: - async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: - return True + async def read_model_budget_raw( + self, entity: str, entity_id: str | None, team_id: str | None, model: str, + budgets: Mapping[str, object] | None, key_budgets: Mapping[str, object] | None = None, + ): + from litellm._auth.model_budget_data import ModelBudgetSnapshot + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading - async def is_user_within_model_budget( - self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None - ) -> bool: - calls.append((user_id, user_model_max_budget, model)) - return True + calls.append((entity_id, budgets, model)) + return ModelBudgetReading(ModelBudgetSnapshot(None, None, None, None, None, None, 0), model, model) valid_token = UserAPIKeyAuth( token="hash", @@ -1821,57 +1822,30 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): ) -def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): - """ - The custom-auth helper runs its own key, user and end-user per-model budget - checks. If it does not honour the zero-cost skip that the JWT and - virtual-key paths honour, the same free request is refused under one auth - method and served under the others. - - Asserted structurally, on the same reasoning as the sibling test: the defect - is a check sitting outside a guard, and it must hold for checks added later - rather than only for whichever request a behavioural test happened to drive. - """ - import ast - import inspect - import textwrap - +@pytest.mark.asyncio +async def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): from litellm.proxy.auth import user_api_key_auth as auth_module - src = textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks)) - tree = ast.parse(src) - - assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag" - - budget_calls = ( - "_enforce_model_budgets", - "is_end_user_within_model_budget", + router = litellm.Router(model_list=[{ + "model_name": "free-model", + "litellm_params": { + "model": "openai/free-model", "api_key": "unused", + "input_cost_per_token": 0, "output_cost_per_token": 0, + }, + "model_info": {"id": "free-id"}, + }]) + token = UserAPIKeyAuth( + token="test-key", + model_max_budget={"free-model": {"max_budget": 0}}, ) - - def guarding_ifs(target: ast.AST) -> list[ast.If]: - return [ - node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is target for sub in ast.walk(node)) - ] - - def mentions_skip(node: ast.If) -> bool: - return any(isinstance(sub, ast.Name) and sub.id == "skip_budget_checks" for sub in ast.walk(node.test)) - - for call_name in budget_calls: - calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and ( - (isinstance(node.func, ast.Name) and node.func.id == call_name) - or (isinstance(node.func, ast.Attribute) and node.func.attr == call_name) - ) - ] - assert calls, f"{call_name} is no longer called here; update this invariant" - for call in calls: - assert any(mentions_skip(node) for node in guarding_ifs(call)), ( - f"{call_name} runs even for a zero-cost model, so custom auth refuses " - "requests the JWT and virtual-key paths serve" - ) + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + result = await auth_module._run_post_custom_auth_checks( + token, None, {"model": "free-model"}, "/v1/chat/completions", None + ) + assert result is token def test_custom_auth_attaches_the_user_budget_even_when_it_does_not_enforce(): diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 8b9d682581a..8b87f985da8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -74,7 +74,7 @@ from litellm.proxy.auth.auth_checks import ( cache_key_object, jwt_key_mapping_cache_key, ) -from litellm._auth.python_core import PythonCore +from tests.unit._auth.reference_core import PythonCore from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model from litellm.proxy.auth.user_api_key_auth import ( @@ -451,26 +451,15 @@ async def test_should_not_reuse_cached_key_object_for_request_state(): @pytest.mark.asyncio async def test_custom_auth_does_not_enforce_key_model_access_by_default(): valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch( - "litellm.proxy.proxy_server.general_settings", - {}, - ), - ): - await _run_post_custom_auth_checks( + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = await _run_post_custom_auth_checks( valid_token=valid_token, request=None, - request_data=request_data, + request_data={"model": "gpt-4o"}, route="/v1/chat/completions", parent_otel_span=None, ) - mock_can_key.assert_not_awaited() + assert result is valid_token @pytest.mark.asyncio @@ -499,29 +488,20 @@ async def test_custom_auth_honors_key_level_model_access_restriction_allowed_wit request_data = {"model": "gpt-4o-mini"} with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, ), ): - await _run_post_custom_auth_checks( + result = await _run_post_custom_auth_checks( valid_token=valid_token, request=None, request_data=request_data, route="/v1/chat/completions", parent_otel_span=None, ) - mock_can_key.assert_awaited_once_with( - model="gpt-4o-mini", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) + assert result is valid_token @pytest.mark.asyncio @@ -530,63 +510,37 @@ async def test_custom_auth_enforces_key_model_access_from_file_route_header_with request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"}) with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, ), ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=request, - request_data={}, - route="/v1/files", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="restricted-model", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) + with pytest.raises(ProxyException) as denial: + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=request, + request_data={}, + route="/v1/files", + parent_otel_span=None, + ) + assert denial.value.type == ProxyErrorTypes.key_model_access_denied + assert "restricted-model" in str(denial.value) @pytest.mark.asyncio async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in(): valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - mock_can_key.side_effect = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) + with patch("litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}): with pytest.raises(ProxyException) as exc: await _run_post_custom_auth_checks( valid_token=valid_token, request=None, - request_data=request_data, + request_data={"model": "gpt-4o"}, route="/v1/chat/completions", parent_otel_span=None, ) - - assert exc.value.type == ProxyErrorTypes.key_model_access_denied + assert exc.value.type == ProxyErrorTypes.key_model_access_denied def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): @@ -1896,7 +1850,13 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th from litellm.proxy.proxy_server import hash_token api_key = "sk-prefetch-order-test" - valid_token = UserAPIKeyAuth(api_key=api_key, token=hash_token(api_key), user_id="u1", team_id="t1") + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=hash_token(api_key), + user_id="u1", + team_id="t1", + models=["gpt-4o" if model_allowed else "other-model"], + ) mock_cache = AsyncMock() mock_cache.async_get_cache = AsyncMock(return_value=None) @@ -1912,7 +1872,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "proxy_logging_obj": mock_proxy_logging_obj, "master_key": "sk-master-key", "general_settings": {}, - "llm_model_list": [], + "llm_model_list": [{"model_name": "gpt-4o"}], "llm_router": None, "open_telemetry_logger": None, "model_max_budget_limiter": MagicMock(), @@ -1921,12 +1881,6 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "litellm_proxy_admin_name": "admin", } _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} - denied = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1939,11 +1893,6 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th new_callable=AsyncMock, return_value=valid_token, ), - patch( # test-quality-ok: the observable is whether the prefetch runs before or after this check - "litellm.proxy.auth.user_api_key_auth.enforce_key_and_fallback_model_access", - new_callable=AsyncMock, - side_effect=None if model_allowed else denied, - ), patch( # test-quality-ok: counting prefetch calls on a denied request IS the regression being pinned "litellm.proxy.auth.user_api_key_auth.prefetch_auth_objects", new_callable=AsyncMock ) as mock_prefetch, @@ -9596,13 +9545,23 @@ class TestJWTAuthUserEmail: def _key_over_budget_for(*models: str) -> AsyncMock: - async def is_key_within_model_budget(user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: - if model in models: - raise litellm.BudgetExceededError(current_cost=10, max_budget=5, message=f"over budget for {model}") - return True + from litellm._auth.model_budget_data import ModelBudgetSnapshot + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading + + async def read_model_budget_raw( + entity: str, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> ModelBudgetReading: + return ModelBudgetReading( + ModelBudgetSnapshot(5, 10 if model in models else 0, None, None, None, None, 0), model, model + ) limiter: Final = AsyncMock() - limiter.is_key_within_model_budget.side_effect = is_key_within_model_budget + limiter.read_model_budget_raw.side_effect = read_model_budget_raw return limiter @@ -9696,7 +9655,9 @@ class TestKeyModelBudgetFallback: valid_token, _key_over_budget_for(*over_budget), request_data, self._make_request() ) - assert exc_info.value.message == "over budget for gpt-4o" + assert ( + exc_info.value.message == "LiteLLM Virtual Key: test-key, key_alias: None, exceeded budget for model=gpt-4o" + ) assert (exc_info.value.current_cost, exc_info.value.max_budget) == (10, 5) assert request_data["model"] == "gpt-4o" diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 55beaf5254d..e892fe43e65 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -131,26 +131,29 @@ def test_resolve_model_budget(): assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None -# Test is_key_within_model_budget @pytest.mark.asyncio -async def test_is_key_within_model_budget(budget_limiter): - # Mock user API key dict - user_api_key = UserAPIKeyAuth( +@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0)) +async def test_is_key_within_model_budget( + budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + spend: float, +) -> None: + user_api_key: Final = UserAPIKeyAuth( token="test-key", key_alias="test-alias", model_max_budget={"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, ) - - # Test when model is within budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): - assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True - - # Test when model exceeds budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): + counter: Final = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.KEY, + entity_id="test-key", + budget_model="gpt-4", + budget_duration="1d", + ) + await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend) + if spend >= 100.0: with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") - - # Test model not in budget config + else: + assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True @@ -224,38 +227,26 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim assert call_kwargs["response_cost"] == 0.05 -# Test is_end_user_within_model_budget @pytest.mark.asyncio -async def test_is_end_user_within_model_budget(budget_limiter): - # Test when model is within budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): - assert ( - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "gpt-4", - ) - is True - ) - - # Test when model exceeds budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): - with pytest.raises(litellm.BudgetExceededError): - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "gpt-4", - ) - - # Test model not in budget config - assert ( - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "non-existent", - ) - is True +@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0)) +async def test_is_end_user_within_model_budget( + budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + spend: float, +) -> None: + counter: Final = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.END_USER, + entity_id="test-user", + budget_model="gpt-4", + budget_duration="1d", ) + await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend) + budgets: Final = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + if spend >= 100.0: + with pytest.raises(litellm.BudgetExceededError): + await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") + else: + assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") is True + assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "non-existent") is True # Test _get_spend_for_model_budget for the end-user scope