diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index fb29a92e306..1d56f8b4a29 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -475,6 +475,7 @@ jobs: - shard: rust-bridge-harness test-path: >- + tests/unit/_auth tests/test_litellm_rust/test_dispatch.py tests/test_litellm_rust/test_public_entries.py tests/test_litellm_rust/chat_completions/test_dispatch.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 48b6730b0c6..1d5ca54308e 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4408,6 +4408,7 @@ dependencies = [ "litellm-redis-resolve", "litellm-redis-store", "litellm-redis-types", + "litellm-request-auth-core", "litellm-router-core", "litellm-router-types", "litellm-secret-value", @@ -4562,6 +4563,19 @@ dependencies = [ "url", ] +[[package]] +name = "litellm-request-auth-core" +version = "0.1.0" +dependencies = [ + "litellm-secret-value", + "rstest", + "serde", + "sha2 0.10.9", + "strum", + "thiserror 2.0.19", + "veil", +] + [[package]] name = "litellm-router" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 72d6694ba79..e8a7ade51e8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -9,6 +9,7 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +litellm-request-auth-core = { path = "crates/request-auth-core" } litellm-config = { path = "crates/config" } litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6ecfebf0096..51f02365b59 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +litellm-request-auth-core.workspace = true litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 09b787b2c94..b7500311107 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,6 +13,7 @@ mod model_registry; mod preflight; mod python_settings; mod redis; +mod request_auth; mod router; mod routes; mod secrets; @@ -61,6 +62,12 @@ mod _native { NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns, }; #[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, + }; + #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs new file mode 100644 index 00000000000..d96bf7b5bde --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -0,0 +1,124 @@ +use std::sync::Mutex; + +use litellm_host_python::{from_py, to_py}; +use litellm_request_auth_core::{self as core, credential, protocol::*}; +use pyo3::{ + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, +}; + +type ProjectedStep = (Py, Option); + +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), + } +} + +fn project(py: Python<'_>, step: Result) -> PyResult { + let step = step.map_err(map_error)?; + Ok(( + to_py(py, &step.node)?, + step.resume.map(|resume| NativeAuthContinuation { + resume: Mutex::new(Some(resume)), + }), + )) +} + +#[pyclass(frozen, module = "litellm.rust_bridge._native")] +pub(crate) struct NativeAuthContinuation { + resume: Mutex>, +} + +#[pymethods] +impl NativeAuthContinuation { + fn resume(&self, py: Python<'_>, answer: &Bound<'_, PyAny>) -> PyResult { + let resume = self + .resume + .lock() + .map_err(|_| PyRuntimeError::new_err("auth continuation lock poisoned"))? + .take() + .ok_or_else(|| PyRuntimeError::new_err("auth continuation already consumed"))?; + let answer: Answer = from_py(answer)?; + project(py, py.detach(move || resume(answer))) + } +} + +#[pyfunction] +pub(crate) fn request_auth_start( + py: Python<'_>, + facts: &Bound<'_, PyAny>, + settings: &Bound<'_, PyAny>, + now: &Bound<'_, PyAny>, + resolved: bool, +) -> PyResult { + let facts: RequestFacts = from_py(facts)?; + let settings: Settings = from_py(settings)?; + let now: DateTime = from_py(now)?; + project( + py, + py.detach(move || { + if resolved { + core::start_resolved_key(facts, settings, now) + } else { + core::start(facts, settings, now) + } + }), + ) +} + +#[derive(serde::Deserialize, strum::VariantArray)] +enum BudgetEntry { + Key, + User, + TeamMember, +} + +#[pyfunction] +pub(crate) fn request_auth_budgets( + py: Python<'_>, + entry: &Bound<'_, PyAny>, +) -> PyResult { + let entry: BudgetEntry = from_py(entry)?; + project( + py, + py.detach(move || match entry { + BudgetEntry::Key => core::check_model_budgets(), + BudgetEntry::User => core::check_user_model_budgets(), + BudgetEntry::TeamMember => core::check_team_member_model_budgets(), + }), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_select_api_key( + py: Python<'_>, + keys: &Bound<'_, PyAny>, +) -> PyResult> { + let keys: credential::PresentedKeys = from_py(keys)?; + to_py(py, &py.detach(move || credential::select_api_key(keys))) +} + +#[pyfunction] +pub(crate) fn request_auth_select_websocket_api_key( + py: Python<'_>, + keys: &Bound<'_, PyAny>, +) -> PyResult> { + let keys: credential::WebSocketKeys = from_py(keys)?; + to_py( + py, + &py.detach(move || credential::select_websocket_api_key(keys)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_scheme_credential(value: &str) -> Option { + credential::scheme_credential(value) +} + +#[pyfunction] +pub(crate) fn request_auth_litellm_header_token(value: &str) -> String { + credential::litellm_header_token(value) +} diff --git a/litellm-rust/crates/request-auth-core/Cargo.toml b/litellm-rust/crates/request-auth-core/Cargo.toml new file mode 100644 index 00000000000..6b82bde399e --- /dev/null +++ b/litellm-rust/crates/request-auth-core/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "litellm-request-auth-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true + +[dependencies] +litellm-secret-value.workspace = true +serde.workspace = true +sha2.workspace = true +strum.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/request-auth-core/src/credential.rs b/litellm-rust/crates/request-auth-core/src/credential.rs new file mode 100644 index 00000000000..a246b36e9db --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/credential.rs @@ -0,0 +1,225 @@ +use litellm_secret_value::SecretValue; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Default)] +pub struct PresentedKeys { + pub custom_header: Option, + pub litellm_header: Option, + pub authorization: Option, + pub api_key: Option, + pub x_api_key: Option, + pub google: Option, + pub apim: Option, + pub query: Option, + pub pass_through: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum KeySource { + #[serde(rename = "custom_header")] + CustomHeader, + #[serde(rename = "x-litellm-api-key")] + LitellmHeader, + #[serde(rename = "authorization")] + Authorization, + #[serde(rename = "api-key")] + ApiKey, + #[serde(rename = "x-api-key")] + XApiKey, + #[serde(rename = "x-goog-api-key")] + Google, + #[serde(rename = "ocp-apim-subscription-key")] + Apim, + #[serde(rename = "key")] + Query, + #[serde(rename = "pass_through_endpoint")] + PassThrough, +} + +#[derive(veil::Redact, Serialize)] +pub struct SelectedKey { + pub source: KeySource, + #[redact(with = "[REDACTED]")] + pub presented: String, + #[redact(with = "[REDACTED]")] + pub token: String, +} + +#[derive(Debug, Deserialize, Default)] +pub struct WebSocketKeys { + pub authorization: Option, + pub api_key: Option, + pub subprotocols: Option, +} + +#[derive(Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum WebSocketDenial { + #[serde(rename = "no_api_key")] + NoApiKey, + #[serde(rename = "invalid_authorization")] + InvalidAuthorization, +} + +#[derive(veil::Redact, Serialize)] +#[serde(tag = "kind", content = "value")] +pub enum WebSocketResult { + Token(#[redact(with = "[REDACTED]")] String), + Denial(WebSocketDenial), +} + +fn whitespace(ch: char) -> bool { + ch.is_whitespace() || matches!(ch, '\u{1c}'..='\u{1f}') +} + +fn credential_part(value: &str) -> Option<&str> { + let end = value + .find(|ch| ch == '/' || ch == ',' || whitespace(ch)) + .unwrap_or(value.len()); + (end != 0).then_some(&value[..end]) +} + +pub fn scheme_credential(value: &str) -> Option { + if value.starts_with("AWS4-HMAC-SHA256") { + let bearer = value + .match_indices("Credential=Bearer") + .find_map(|(index, prefix)| { + let tail = &value[index + prefix.len()..]; + if !tail.starts_with(whitespace) { + return None; + } + credential_part(tail.trim_start_matches(whitespace)) + }); + return bearer + .or_else(|| { + value + .match_indices("Credential=") + .find_map(|(index, prefix)| credential_part(&value[index + prefix.len()..])) + }) + .map(str::to_owned); + } + let (scheme, token) = value.split_once(' ')?; + matches!(scheme.to_lowercase().as_str(), "bearer" | "basic").then(|| token.to_owned()) +} + +pub fn litellm_header_token(value: &str) -> String { + scheme_credential(value).unwrap_or_else(|| value.to_owned()) +} + +pub fn select_api_key(keys: PresentedKeys) -> Option { + if let Some(value) = keys.custom_header { + return Some(SelectedKey { + source: KeySource::CustomHeader, + presented: value.expose().to_owned(), + token: scheme_credential(value.expose()).unwrap_or_default(), + }); + } + if let Some(value) = keys.litellm_header { + return Some(SelectedKey { + source: KeySource::LitellmHeader, + presented: value.expose().to_owned(), + token: litellm_header_token(value.expose()), + }); + } + if let Some(value) = keys + .authorization + .filter(|value| !value.expose().is_empty()) + { + return Some(SelectedKey { + source: KeySource::Authorization, + presented: value.expose().to_owned(), + token: scheme_credential(value.expose()).unwrap_or_default(), + }); + } + [ + (KeySource::ApiKey, keys.api_key), + (KeySource::XApiKey, keys.x_api_key), + (KeySource::Google, keys.google), + (KeySource::Apim, keys.apim), + ( + KeySource::Query, + keys.query.filter(|value| !value.expose().is_empty()), + ), + (KeySource::PassThrough, keys.pass_through), + ] + .into_iter() + .find_map(|(source, value)| { + value.map(|value| SelectedKey { + source, + presented: value.expose().to_owned(), + token: value.expose().to_owned(), + }) + }) +} + +pub fn select_websocket_api_key(keys: WebSocketKeys) -> WebSocketResult { + if let Some(value) = keys + .authorization + .filter(|value| !value.expose().is_empty()) + { + let Some((scheme, token)) = value.expose().split_once(' ') else { + return WebSocketResult::Denial(WebSocketDenial::InvalidAuthorization); + }; + if scheme.to_lowercase() != "bearer" { + return WebSocketResult::Denial(WebSocketDenial::InvalidAuthorization); + } + return WebSocketResult::Token(token.trim_matches(whitespace).to_owned()); + } + if let Some(value) = keys.api_key.filter(|value| !value.expose().is_empty()) { + return WebSocketResult::Token(value.expose().to_owned()); + } + let token = keys.subprotocols.as_ref().and_then(|value| { + value.expose().split(',').find_map(|protocol| { + protocol + .trim_matches(whitespace) + .strip_prefix("openai-insecure-api-key.") + }) + }); + match token { + Some(token) if !token.is_empty() => WebSocketResult::Token(token.to_owned()), + _ => WebSocketResult::Denial(WebSocketDenial::NoApiKey), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::bearer("Bearer token", Some(" token"))] + #[case::empty("Bearer ", Some(""))] + #[case::basic("bAsIc token", Some("token"))] + #[case::unsupported("Token token", None)] + #[case::aws("AWS4-HMAC-SHA256 Credential=Bearer\u{1c}token/date", Some("token"))] + #[case::aws_precedence( + "AWS4-HMAC-SHA256 Credential=first/date, Credential=Bearer second/date", + Some("second") + )] + fn schemes(#[case] value: &str, #[case] expected: Option<&str>) { + assert_eq!(scheme_credential(value).as_deref(), expected); + } + + #[rstest] + fn empty_custom_header_still_wins() { + let selected = select_api_key(PresentedKeys { + custom_header: Some("".into()), + api_key: Some("other".into()), + ..Default::default() + }) + .unwrap(); + assert_eq!(selected.source, KeySource::CustomHeader); + assert_eq!(selected.token, ""); + } + + #[rstest] + fn secrets_are_redacted() { + let keys = PresentedKeys { + api_key: Some("credential-value".into()), + ..Default::default() + }; + let keys_debug = format!("{keys:?}"); + let debug = format!("{keys_debug} {:?}", select_api_key(keys).unwrap()); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("REDACTED")); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/error.rs b/litellm-rust/crates/request-auth-core/src/error.rs new file mode 100644 index 00000000000..6656d9f409c --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/error.rs @@ -0,0 +1,7 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid auth effect answer")] + InvalidAnswer, + #[error("can't compare offset-naive and offset-aware datetimes")] + IncompatibleDateTimes, +} diff --git a/litellm-rust/crates/request-auth-core/src/key.rs b/litellm-rust/crates/request-auth-core/src/key.rs new file mode 100644 index 00000000000..b4b20621759 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/key.rs @@ -0,0 +1,418 @@ +use litellm_secret_value::SecretValue; +use sha2::{Digest, Sha256}; + +use crate::{Error, model_budgets, protocol::*}; + +struct KeyCheck { + settings: Settings, + now: DateTime, + is_llm_api_route: bool, + key_label: SecretValue, +} + +pub fn start(facts: RequestFacts, settings: Settings, now: DateTime) -> Result { + let Some(master_key) = &settings.master_key else { + return Step::allow(Grant::NoMasterKey { + token: facts.token.map(|token| token.expose().to_owned()), + }); + }; + let Some(token) = facts.token.clone() else { + return Step::deny(Denial::MissingKey); + }; + if token.expose().is_empty() { + return Step::deny(Denial::MalformedKey); + } + if facts.route == "/user/auth" { + return if settings.allow_user_auth { + Step::allow(Grant::UserAuthRoute) + } else { + Step::deny(Denial::UserAuthDisabled) + }; + } + if &token == master_key { + return Step::effect( + Effect::ResolveEndUser { + prefetch_key: false, + }, + |(): ()| Step::allow(Grant::MasterKey), + ); + } + let hashed_token = format!("{:x}", Sha256::digest(token.expose().as_bytes())); + Step::effect( + Effect::ResolveEndUser { prefetch_key: true }, + move |(): ()| { + Step::effect( + Effect::ResolvePresentedKey { + hashed_token: hashed_token.clone(), + }, + move |presented: Option| { + after_presented_key(presented, token, hashed_token, facts, settings, now) + }, + ) + }, + ) +} + +pub fn start_resolved_key( + facts: RequestFacts, + settings: Settings, + now: DateTime, +) -> Result { + check_key(KeyCheck { + settings, + now, + is_llm_api_route: facts.is_llm_api_route, + key_label: facts.token.unwrap_or_else(|| "".into()), + }) +} + +fn after_presented_key( + presented: Option, + token: SecretValue, + hashed_token: String, + facts: RequestFacts, + settings: Settings, + now: DateTime, +) -> Result { + if let Some(PresentedKey::AdminKey { expires }) = &presented { + if let Some(expires) = expires + && expires.before(&now)? + { + let expires = expires.clone(); + return Step::effect(Effect::EvictPresentedKey { hashed_token }, move |(): ()| { + Step::deny(Denial::ExpiredKey { + expires, + now, + key_label: token.expose().to_owned(), + }) + }); + } + return Step::allow(Grant::CachedAdminKey); + } + if settings.master_key_only_routes.contains(&facts.route) { + return Step::deny(Denial::MasterKeyOnlyRoute); + } + if !settings.has_db { + return Step::deny(Denial::NoDatabase); + } + if presented.is_some() { + let check = KeyCheck { + settings, + now, + is_llm_api_route: facts.is_llm_api_route, + key_label: token, + }; + return Step::effect( + Effect::PrepareKey { + loaded_from_db: false, + }, + move |(): ()| check_key(check), + ); + } + if !token.expose().starts_with("sk-") { + return Step::deny(Denial::NotAVirtualKey); + } + let check = KeyCheck { + settings, + now, + is_llm_api_route: facts.is_llm_api_route, + key_label: hashed_token.clone().into(), + }; + Step::effect(Effect::LoadKey { hashed_token }, move |(): ()| { + Step::effect( + Effect::PrepareKey { + loaded_from_db: true, + }, + move |(): ()| check_key(check), + ) + }) +} + +fn check_key(check: KeyCheck) -> Result { + Step::effect(Effect::InspectKey, move |key: KeyFacts| { + if key.blocked { + return Step::deny(Denial::KeyBlocked); + } + Step::effect(Effect::EnforceModelAccess, move |(): ()| { + check_owner(check, key) + }) + }) +} + +fn check_owner(check: KeyCheck, key: KeyFacts) -> Result { + let Some(user_id) = key.user_id.clone() else { + return check_budget_exemption(check, key); + }; + Step::effect( + Effect::LoadKeyOwner { + user_id: user_id.clone(), + }, + move |owner: Option| { + if owner.is_some_and(|owner| owner.deactivated) { + return Step::deny(Denial::OwnerDeactivated { user_id }); + } + check_budget_exemption(check, key) + }, + ) +} + +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, + }) + }, + ) + }, + ) + }) +} + +fn check_expiry(check: KeyCheck, key: KeyFacts, exempt: bool) -> Result { + if let Some(expires) = &key.expires + && expires.before(&check.now)? + { + return Step::deny(Denial::ExpiredKey { + expires: expires.clone(), + now: check.now, + key_label: check.key_label.expose().to_owned(), + }); + } + if exempt { + return load_context(check, key, None); + } + Step::effect(Effect::SendKeyBudgetAlerts, move |(): ()| { + check_key_budget(check, key) + }) +} + +fn check_key_budget(check: KeyCheck, key: KeyFacts) -> Result { + let Some(max_budget) = key.max_budget.filter(|_| check.is_llm_api_route) else { + return check_soft_budget(check, key, None); + }; + Step::effect(Effect::ReadKeySpend { max_budget }, move |spend: f64| { + Step::effect(Effect::AlertKeyBudget { spend }, move |(): ()| { + if !max_budget.is_finite() || spend < max_budget { + return check_soft_budget(check, key, None); + } + if let Some(throttle) = key.budget_throttle_pct { + return check_soft_budget(check, key, Some(throttle)); + } + Step::deny(Denial::KeyBudgetExceeded { spend, max_budget }) + }) + }) +} + +fn check_soft_budget(check: KeyCheck, key: KeyFacts, throttle: Option) -> Result { + let alert = key + .soft_budget + .is_some_and(|budget| budget != 0.0 && key.spend >= budget); + let then = + move || model_budgets::key_budgets(Box::new(move || load_context(check, key, throttle))); + if alert { + return Step::effect(Effect::AlertKeySoftBudget, move |(): ()| then()); + } + then() +} + +fn load_context(check: KeyCheck, key: KeyFacts, throttle: Option) -> Result { + Step::effect(Effect::LoadKeyContext, move |(): ()| { + if !(check.settings.proxy_max_budget > 0.0 && check.settings.has_db) { + return grant(key, throttle); + } + Step::effect( + Effect::ReadProxySpend, + move |spend: Option| match spend { + None => grant(key, throttle), + Some(spend) => Step::effect(Effect::AlertProxyBudget { spend }, move |(): ()| { + grant(key, throttle) + }), + }, + ) + }) +} + +fn grant(key: KeyFacts, throttle: Option) -> Result { + if !key.has_token { + return Step::deny(Denial::KeyWithoutToken); + } + Step::allow(Grant::VirtualKey { + budget_throttle_pct: throttle, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn check() -> KeyCheck { + KeyCheck { + settings: Settings { + master_key: Some("sk-master".into()), + allow_user_auth: false, + has_db: true, + master_key_only_routes: ["/admin".into()].into(), + proxy_max_budget: 0.0, + }, + now: DateTime { + wall_microseconds: 100, + offset_microseconds: Some(0), + timezone_id: 1, + reference: 0, + }, + is_llm_api_route: true, + key_label: "credential-value".into(), + } + } + + #[fixture] + fn key() -> KeyFacts { + KeyFacts { + has_token: true, + blocked: false, + expires: None, + spend: 0.0, + soft_budget: None, + max_budget: None, + team_member_spend: None, + user_id: None, + team_id: None, + budget_throttle_pct: None, + } + } + + #[rstest] + fn blocked_key_stops_before_model_access(check: KeyCheck, key: KeyFacts) { + let step = check_key(check).unwrap(); + assert!(matches!(step.node, Node::Effect(Effect::InspectKey))); + let result = (step.resume.unwrap())(Answer::Key(KeyFacts { + blocked: true, + ..key + })) + .unwrap(); + assert!(matches!( + result.node, + Node::Outcome(Outcome::Deny { + denial: Denial::KeyBlocked + }) + )); + assert!(result.resume.is_none()); + } + + #[rstest] + #[case::at_limit(1.0, 1.0, true)] + #[case::below_limit(0.9, 1.0, false)] + #[case::nan_limit(1.0, f64::NAN, false)] + #[case::infinite_limit(1.0, f64::INFINITY, false)] + #[case::nan_spend(f64::NAN, 1.0, true)] + fn budget_alert_precedes_the_decision( + check: KeyCheck, + key: KeyFacts, + #[case] spend: f64, + #[case] max_budget: f64, + #[case] denied: bool, + ) { + let read = check_key_budget( + check, + KeyFacts { + max_budget: Some(max_budget), + ..key + }, + ) + .unwrap(); + assert!(matches!( + read.node, + Node::Effect(Effect::ReadKeySpend { .. }) + )); + let alert = (read.resume.unwrap())(Answer::Spend(spend)).unwrap(); + assert!(matches!( + alert.node, + Node::Effect(Effect::AlertKeyBudget { .. }) + )); + let result = (alert.resume.unwrap())(Answer::Unit).unwrap(); + if denied { + assert!(matches!( + result.node, + Node::Outcome(Outcome::Deny { + denial: Denial::KeyBudgetExceeded { .. } + }) + )); + } else { + assert!(matches!( + result.node, + Node::Effect(Effect::ReadModelBudgetScope) + )); + } + } + + #[rstest] + fn cached_admin_skips_route_and_database_checks(check: KeyCheck) { + let settings = Settings { + has_db: false, + ..check.settings + }; + let step = start( + RequestFacts { + route: "/admin".into(), + token: Some("other".into()), + is_llm_api_route: false, + }, + settings, + check.now, + ) + .unwrap(); + let lookup = (step.resume.unwrap())(Answer::Unit).unwrap(); + let result = (lookup.resume.unwrap())(Answer::Presented(Some(PresentedKey::AdminKey { + expires: None, + }))) + .unwrap(); + assert!(matches!( + result.node, + Node::Outcome(Outcome::Allow { + grant: Grant::CachedAdminKey + }) + )); + } + + #[rstest] + fn public_expiry_payload_does_not_leak_into_debug(check: KeyCheck) { + let denial = Denial::ExpiredKey { + expires: check.now.clone(), + now: check.now, + key_label: check.key_label.expose().to_owned(), + }; + let debug = format!("{denial:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("REDACTED")); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs new file mode 100644 index 00000000000..a55e5ed7565 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -0,0 +1,13 @@ +#![forbid(unsafe_code)] + +pub mod credential; +mod error; +mod key; +mod model_budgets; +pub mod protocol; + +pub use error::Error; +pub use key::{start, start_resolved_key}; +pub use model_budgets::{ + check_model_budgets, check_team_member_model_budgets, check_user_model_budgets, +}; diff --git a/litellm-rust/crates/request-auth-core/src/model_budgets.rs b/litellm-rust/crates/request-auth-core/src/model_budgets.rs new file mode 100644 index 00000000000..b02ed6d62c4 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_budgets.rs @@ -0,0 +1,266 @@ +use crate::{Error, protocol::*}; + +fn allowed() -> Result { + Step::allow(Grant::WithinModelBudgets) +} + +pub fn check_model_budgets() -> Result { + key_budgets(Box::new(allowed)) +} + +pub fn check_user_model_budgets() -> Result { + user_budgets(Box::new(allowed)) +} + +pub(crate) fn key_budgets(then: Then) -> Result { + Step::effect( + Effect::ReadModelBudgetScope, + move |scope: ModelBudgetScope| { + if !scope.key_has_model_budgets || scope.models.is_empty() { + return user_models(scope.user_has_model_budgets, scope.models, then); + } + key_models( + scope.key_budget_fallbacks, + scope.models, + scope.rewritable_model, + Box::new(move || user_budgets(then)), + ) + }, + ) +} + +fn user_budgets(then: Then) -> Result { + Step::effect( + Effect::ReadModelBudgetScope, + move |scope: ModelBudgetScope| { + user_models(scope.user_has_model_budgets, scope.models, then) + }, + ) +} + +fn key_models( + fallbacks: std::collections::BTreeMap>, + models: Vec, + rewritable: Option, + then: Then, +) -> Result { + let Some((model, rest)) = models.split_first() else { + return then(); + }; + let model = model.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::CheckKeyModelBudget { + model: model.clone(), + }, + move |exceeded: Option| { + let Some(exceeded) = exceeded else { + return key_models(fallbacks, rest, rewritable, then); + }; + if Some(&model) != rewritable.as_ref() { + return Step::deny(exceeded); + } + let candidates = fallbacks.get(&model).cloned().unwrap_or_default(); + first_affordable( + candidates, + exceeded, + Box::new(move |fallback| key_models(fallbacks, rest, Some(fallback), then)), + ) + }, + ) +} + +type Found = Box Result + Send + Sync>; + +fn first_affordable( + candidates: Vec, + exceeded: ModelBudgetExceeded, + found: Found, +) -> Result { + let Some((candidate, rest)) = candidates.split_first() else { + return Step::deny(exceeded); + }; + let candidate = candidate.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::CheckKeyModelBudget { + model: candidate.clone(), + }, + move |over: Option| { + if over.is_some() { + return first_affordable(rest, exceeded, found); + } + Step::effect( + Effect::AuthorizeBudgetFallback { + model: candidate.clone(), + }, + move |authorized: bool| { + if !authorized { + return Step::deny(exceeded); + } + Step::effect( + Effect::RewriteRequestModel { + model: candidate.clone(), + }, + move |(): ()| found(candidate), + ) + }, + ) + }, + ) +} + +fn user_models(enabled: bool, models: Vec, then: Then) -> Result { + if !enabled { + return then(); + } + let Some((model, rest)) = models.split_first() else { + return then(); + }; + let rest = rest.to_vec(); + Step::effect( + Effect::CheckUserModelBudget { + model: model.clone(), + }, + move |exceeded: Option| match exceeded { + Some(exceeded) => Step::deny(exceeded), + None => user_models(enabled, rest, then), + }, + ) +} + +pub fn check_team_member_model_budgets() -> Result { + Step::effect( + Effect::ReadTeamMemberModelBudgetScope, + |scope: TeamMemberModelBudgetScope| { + if !scope.has_model_budgets { + return allowed(); + } + member_models(scope.models, scope.rewritable_model, scope.can_fall_back) + }, + ) +} + +fn member_models( + models: Vec, + rewritable: Option, + can_fall_back: bool, +) -> Result { + let Some((model, rest)) = models.split_first() else { + return allowed(); + }; + let model = model.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::CheckTeamMemberModelBudget { + model: model.clone(), + }, + move |exceeded: Option| { + let Some(exceeded) = exceeded else { + return member_models(rest, rewritable, can_fall_back); + }; + if Some(&model) != rewritable.as_ref() || !can_fall_back { + return Step::deny(exceeded); + } + Step::effect( + Effect::ReadTeamMemberFallbacks { + model: model.clone(), + }, + move |targets: Vec| { + first_usable( + targets + .into_iter() + .filter(|target| target.model != model) + .collect(), + exceeded, + Box::new(move |fallback| { + member_models(rest, Some(fallback), can_fall_back) + }), + ) + }, + ) + }, + ) +} + +fn first_usable( + targets: Vec, + exceeded: ModelBudgetExceeded, + found: Found, +) -> Result { + let Some((target, rest)) = targets.split_first() else { + return Step::deny(exceeded); + }; + let target = target.clone(); + let rest = rest.to_vec(); + Step::effect( + Effect::VetTeamMemberFallback { + target: target.clone(), + }, + move |usable: bool| { + if !usable { + return first_usable(rest, exceeded, found); + } + Step::effect( + Effect::RerouteTeamMemberRequest { + model: target.model.clone(), + }, + move |(): ()| found(target.model), + ) + }, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn unauthorized_affordable_fallback_denies_without_trying_the_next_one() { + let exceeded = ModelBudgetExceeded { + message: "over budget".into(), + current_cost: 2.0, + max_budget: 1.0, + entity_type: None, + entity_id: None, + }; + let read = first_affordable( + vec!["first".into(), "second".into()], + exceeded, + Box::new(|_| allowed()), + ) + .unwrap(); + assert!( + matches!(read.node, Node::Effect(Effect::CheckKeyModelBudget { ref model }) if model == "first") + ); + let authorize = (read.resume.unwrap())(Answer::Exceeded(None)).unwrap(); + assert!( + matches!(authorize.node, Node::Effect(Effect::AuthorizeBudgetFallback { ref model }) if model == "first") + ); + let result = (authorize.resume.unwrap())(Answer::Bool(false)).unwrap(); + assert!(matches!( + result.node, + Node::Outcome(Outcome::Deny { + denial: Denial::ModelBudgetExceeded { .. } + }) + )); + assert!(result.resume.is_none()); + } + + #[rstest] + fn user_entry_never_checks_key_budgets() { + let step = check_user_model_budgets().unwrap(); + let result = (step.resume.unwrap())(Answer::ModelScope(ModelBudgetScope { + models: vec!["model".into()], + rewritable_model: Some("model".into()), + key_has_model_budgets: true, + key_budget_fallbacks: Default::default(), + user_has_model_budgets: true, + })) + .unwrap(); + assert!( + matches!(result.node, Node::Effect(Effect::CheckUserModelBudget { ref model }) if model == "model") + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/protocol.rs b/litellm-rust/crates/request-auth-core/src/protocol.rs new file mode 100644 index 00000000000..c5a139f4402 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/protocol.rs @@ -0,0 +1,401 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_secret_value::SecretValue; +use serde::{Deserialize, Serialize}; + +use crate::Error; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DateTime { + pub wall_microseconds: i64, + pub offset_microseconds: Option, + pub timezone_id: u64, + pub reference: usize, +} + +impl DateTime { + pub fn before(&self, other: &Self) -> Result { + if self.timezone_id == other.timezone_id { + return Ok(self.wall_microseconds < other.wall_microseconds); + } + match (self.offset_microseconds, other.offset_microseconds) { + (Some(left), Some(right)) => { + Ok(self.wall_microseconds - left < other.wall_microseconds - right) + } + (None, None) => Ok(self.wall_microseconds < other.wall_microseconds), + _ => Err(Error::IncompatibleDateTimes), + } + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct RequestFacts { + pub route: String, + pub token: Option, + pub is_llm_api_route: bool, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Settings { + pub master_key: Option, + pub allow_user_auth: bool, + pub has_db: bool, + pub master_key_only_routes: BTreeSet, + pub proxy_max_budget: f64, +} + +#[derive(Clone, veil::Redact, Serialize)] +#[serde(tag = "kind")] +pub enum Grant { + NoMasterKey { + #[redact(with = "[REDACTED]")] + token: Option, + }, + UserAuthRoute, + MasterKey, + CachedAdminKey, + VirtualKey { + budget_throttle_pct: Option, + }, + WithinModelBudgets, +} + +#[derive(Clone, veil::Redact, Serialize)] +#[serde(tag = "kind")] +pub enum Denial { + MissingKey, + MalformedKey, + UserAuthDisabled, + ExpiredKey { + expires: DateTime, + now: DateTime, + #[redact(with = "[REDACTED]")] + key_label: String, + }, + MasterKeyOnlyRoute, + NoDatabase, + NotAVirtualKey, + KeyBlocked, + OwnerDeactivated { + user_id: String, + }, + TeamMemberBudgetExceeded { + user_id: String, + team_id: String, + spend: f64, + max_budget: f64, + }, + KeyBudgetExceeded { + spend: f64, + max_budget: f64, + }, + KeyWithoutToken, + ModelBudgetExceeded { + message: String, + current_cost: f64, + max_budget: f64, + entity_type: Option, + entity_id: Option, + }, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct ModelBudgetExceeded { + pub message: String, + pub current_cost: f64, + pub max_budget: f64, + pub entity_type: Option, + pub entity_id: Option, +} + +impl From for Denial { + fn from(value: ModelBudgetExceeded) -> Self { + Self::ModelBudgetExceeded { + message: value.message, + current_cost: value.current_cost, + max_budget: value.max_budget, + entity_type: value.entity_type, + entity_id: value.entity_id, + } + } +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(tag = "kind")] +pub enum PresentedKey { + AdminKey { expires: Option }, + NonAdminKey, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct KeyFacts { + pub has_token: bool, + pub blocked: bool, + pub expires: Option, + pub spend: f64, + pub soft_budget: Option, + pub max_budget: Option, + pub team_member_spend: Option, + pub user_id: Option, + pub team_id: Option, + pub budget_throttle_pct: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct OwnerFacts { + pub deactivated: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct TeamMemberBudget { + pub spend: f64, + pub max_budget: f64, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelBudgetScope { + pub models: Vec, + pub rewritable_model: Option, + pub key_has_model_budgets: bool, + pub key_budget_fallbacks: BTreeMap>, + pub user_has_model_budgets: bool, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TeamMemberModelBudgetScope { + pub models: Vec, + pub rewritable_model: Option, + pub has_model_budgets: bool, + pub can_fall_back: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct FallbackTarget { + pub requested: String, + pub model: String, + pub served_model: String, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(tag = "kind")] +pub enum Effect { + ResolveEndUser { + prefetch_key: bool, + }, + ResolvePresentedKey { + hashed_token: String, + }, + EvictPresentedKey { + hashed_token: String, + }, + LoadKey { + hashed_token: String, + }, + PrepareKey { + loaded_from_db: bool, + }, + InspectKey, + EnforceModelAccess, + LoadKeyOwner { + user_id: String, + }, + CheckBudgetExemption, + ReadTeamMemberBudget { + user_id: String, + team_id: String, + }, + AlertTeamMemberBudget { + user_id: String, + team_id: String, + budget: TeamMemberBudget, + }, + SendKeyBudgetAlerts, + ReadKeySpend { + max_budget: f64, + }, + AlertKeyBudget { + spend: f64, + }, + AlertKeySoftBudget, + LoadKeyContext, + ReadProxySpend, + AlertProxyBudget { + spend: f64, + }, + ReadModelBudgetScope, + CheckKeyModelBudget { + model: String, + }, + AuthorizeBudgetFallback { + model: String, + }, + RewriteRequestModel { + model: String, + }, + CheckUserModelBudget { + model: String, + }, + ReadTeamMemberModelBudgetScope, + CheckTeamMemberModelBudget { + model: String, + }, + ReadTeamMemberFallbacks { + model: String, + }, + VetTeamMemberFallback { + target: FallbackTarget, + }, + RerouteTeamMemberRequest { + model: String, + }, +} + +#[derive(Debug, Serialize)] +#[serde(tag = "kind")] +pub enum Outcome { + Allow { grant: Grant }, + Deny { denial: Denial }, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +pub enum Node { + Outcome(Outcome), + Effect(Effect), +} + +#[derive(Deserialize)] +#[serde(tag = "kind", content = "value")] +pub enum Answer { + Unit, + Bool(bool), + Presented(Option), + Key(KeyFacts), + Owner(Option), + TeamBudget(Option), + Spend(f64), + ProxySpend(Option), + ModelScope(ModelBudgetScope), + Exceeded(Option), + TeamScope(TeamMemberModelBudgetScope), + Targets(Vec), +} + +pub trait FromAnswer: Sized { + fn from_answer(answer: Answer) -> Result; +} + +macro_rules! answer { + ($type:ty, $variant:ident) => { + impl FromAnswer for $type { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::$variant(value) => Ok(value), + _ => Err(Error::InvalidAnswer), + } + } + } + }; +} +answer!(bool, Bool); +answer!(Option, Presented); +answer!(KeyFacts, Key); +answer!(Option, Owner); +answer!(Option, TeamBudget); +answer!(f64, Spend); +answer!(Option, ProxySpend); +answer!(ModelBudgetScope, ModelScope); +answer!(Option, Exceeded); +answer!(TeamMemberModelBudgetScope, TeamScope); +answer!(Vec, Targets); + +impl FromAnswer for () { + fn from_answer(answer: Answer) -> Result { + match answer { + Answer::Unit => Ok(()), + _ => Err(Error::InvalidAnswer), + } + } +} + +pub type Continuation = Box Result + Send + Sync>; +pub type Then = Box Result + Send + Sync>; + +pub struct Step { + pub node: Node, + pub resume: Option, +} + +impl Step { + pub fn allow(grant: Grant) -> Result { + Ok(Self { + node: Node::Outcome(Outcome::Allow { grant }), + resume: None, + }) + } + pub fn deny(denial: impl Into) -> Result { + Ok(Self { + node: Node::Outcome(Outcome::Deny { + denial: denial.into(), + }), + resume: None, + }) + } + pub fn effect( + effect: Effect, + resume: impl FnOnce(T) -> Result + Send + Sync + 'static, + ) -> Result { + Ok(Self { + node: Node::Effect(effect), + resume: Some(Box::new(move |answer| resume(T::from_answer(answer)?))), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::same_zone_ignores_offset(10, Some(100), 1, 11, Some(200), 1, Ok(true))] + #[case::different_zones_use_offset(10, Some(100), 1, 11, Some(200), 2, Ok(false))] + #[case::equal_is_not_expired(10, None, 0, 10, None, 0, Ok(false))] + #[case::mixed_naive_aware(10, None, 0, 10, Some(0), 1, Err(()))] + #[allow(clippy::too_many_arguments)] + fn expiry_comparison( + #[case] wall: i64, + #[case] offset: Option, + #[case] zone: u64, + #[case] other_wall: i64, + #[case] other_offset: Option, + #[case] other_zone: u64, + #[case] expected: Result, + ) { + let expires = DateTime { + wall_microseconds: wall, + offset_microseconds: offset, + timezone_id: zone, + reference: 0, + }; + let now = DateTime { + wall_microseconds: other_wall, + offset_microseconds: other_offset, + timezone_id: other_zone, + reference: 1, + }; + assert_eq!(expires.before(&now).map_err(|_| ()), expected); + } + + #[rstest] + fn wrong_answer_never_advances_the_decision() { + let step = Step::effect(Effect::CheckBudgetExemption, |_: bool| { + Step::allow(Grant::MasterKey) + }) + .unwrap(); + assert!(matches!( + (step.resume.unwrap())(Answer::Unit), + Err(Error::InvalidAnswer) + )); + } +} diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md new file mode 100644 index 00000000000..0c78570ae64 --- /dev/null +++ b/litellm/_auth/AGENTS.md @@ -0,0 +1,18 @@ +# litellm/_auth + +## Why this exists + +- Lots of Python code like `litellm/proxy/auth/user_api_key_auth.py` is too complex +- Soon, Rust will implement the same auth stack, and expose common functionality as Python API through pyo3, to avoid maintaining 2 copy of code. +- So the goal of `_auth` here is to define that interface first, not writing really good python code. Should be behavior preserving refactor though + +## The swap point + +- `decision.py` is the interface. `AuthCore.start(RequestFacts, Settings, now) -> Step` is the only thing a core implements, and `python_core.PythonCore` is the Python one. A Rust core will implement the same protocol through pyo3, and the builder picks one without changing anything else +- A core does no IO and reads no globals. The same facts, settings, time and effect answers always give the same outcome +- A `Step` is either an `Outcome` or an effect. Each effect carries its typed `resume` continuation, and `driver.drive` answers effects through an `EffectHandler` until an outcome comes back. The proxy's handler (`_BuilderEffects` in `litellm/proxy/auth/user_api_key_auth.py`) owns all cache and DB IO, so both cores share it +- Projection stays in the proxy. `_authenticated` in the builder module turns an `Allow` grant into `UserAPIKeyAuth`, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract +- The key checks run in the core, in the order the proxy has always applied them: blocked key, model access, SCIM-deactivated owner, budget exemption, team member budget, expiry, budget alerts, the key's own max budget (or its throttle), the soft budget alert, per-model budgets, the team and project context, the proxy budget alert, then the key's token. A budget-exempt request skips everything from the team member budget to per-model budgets except expiry. `start_resolved_key` runs the same checks for keys the proxy resolved itself, such as internal jobs +- `model_budgets.py` decides the key and user per-model budgets: which models to check, when an over-budget key may fall back (only for the top-level request model, to the first `budget_fallbacks` entry it can afford and may call), and that the user check sees the rerouted model. The key chain runs it inline, custom auth runs it through `AuthCore.check_model_budgets`, and JWT runs the user half through `AuthCore.check_user_model_budgets`. Those two entry points end in `Allow(WithinModelBudgets())` and are driven by `driver.drive_model_budgets`, which only needs a `ModelBudgetHandler` +- The team member per-model budget runs through `AuthCore.check_team_member_model_budgets`, which has its own step type (`TeamMemberStep`) and driver (`drive_team_member_model_budgets`). The core decides when an over-budget member may fall back and walks the router's fallback targets in order, skipping the requested model. The proxy handler (`_TeamMemberModelBudgetEffects`) reads those targets and vets each one (pricing, customer access, every per-model budget, model access and access group budgets) +- The team and end-user per-model budgets still run in the proxy after the core, as plain loops with no fallback diff --git a/litellm/_auth/__init__.py b/litellm/_auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/_auth/credential.py b/litellm/_auth/credential.py new file mode 100644 index 00000000000..105eb4afe73 --- /dev/null +++ b/litellm/_auth/credential.py @@ -0,0 +1,141 @@ +import re +from dataclasses import dataclass +from enum import Enum +from typing import Final + +_AWS4_PREFIX: Final = "AWS4-HMAC-SHA256" +_AWS4_BEARER_CREDENTIAL: Final = re.compile(r"Credential=Bearer\s+([^/\s,]+)") +_AWS4_CREDENTIAL: Final = re.compile(r"Credential=([^/\s,]+)") +_TOKEN_SCHEMES: Final = frozenset({"bearer", "basic"}) +_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX: Final = "openai-insecure-api-key." + + +class KeySource(Enum): + CUSTOM_HEADER = "custom_header" + LITELLM_HEADER = "x-litellm-api-key" + AUTHORIZATION = "authorization" + API_KEY = "api-key" + X_API_KEY = "x-api-key" + GOOGLE = "x-goog-api-key" + APIM = "ocp-apim-subscription-key" + QUERY = "key" + PASS_THROUGH = "pass_through_endpoint" + + +@dataclass(frozen=True, slots=True) +class PresentedKeys: + custom_header: str | None = None + litellm_header: str | None = None + authorization: str | None = None + api_key: str | None = None + x_api_key: str | None = None + google: str | None = None + apim: str | None = None + query: str | None = None + pass_through: str | None = None + + +@dataclass(frozen=True, slots=True) +class SelectedKey: + source: KeySource + presented: str + token: str + + +class WebSocketDenial(Enum): + NO_API_KEY = "no_api_key" + INVALID_AUTHORIZATION = "invalid_authorization" + + +@dataclass(frozen=True, slots=True) +class WebSocketKeys: + authorization: str | None = None + api_key: str | None = None + subprotocols: str | None = None + + +def _aws4_access_key(value: str) -> str | None: + match: Final = _AWS4_BEARER_CREDENTIAL.search(value) or _AWS4_CREDENTIAL.search(value) + return match.group(1) if match else None + + +def _split_scheme(value: str) -> tuple[str, str] | None: + scheme, separator, credential = value.partition(" ") + return (scheme.lower(), credential) if separator else None + + +def scheme_credential(value: str) -> str | None: + if value.startswith(_AWS4_PREFIX): + return _aws4_access_key(value) + split: Final = _split_scheme(value) + if split is None: + return None + scheme, credential = split + return credential if scheme in _TOKEN_SCHEMES else None + + +def litellm_header_token(value: str) -> str: + token: Final = scheme_credential(value) + return value if token is None else token + + +def _raw_key(keys: PresentedKeys) -> SelectedKey | None: + raw_sources: Final = ( + (KeySource.API_KEY, keys.api_key), + (KeySource.X_API_KEY, keys.x_api_key), + (KeySource.GOOGLE, keys.google), + (KeySource.APIM, keys.apim), + (KeySource.QUERY, keys.query or None), + (KeySource.PASS_THROUGH, keys.pass_through), + ) + return next( + ( + SelectedKey(source=source, presented=value, token=value) + for source, value in raw_sources + if value is not None + ), + None, + ) + + +def select_api_key(keys: PresentedKeys) -> SelectedKey | None: + if keys.custom_header is not None: + return SelectedKey( + source=KeySource.CUSTOM_HEADER, + presented=keys.custom_header, + token=scheme_credential(keys.custom_header) or "", + ) + if keys.litellm_header is not None: + return SelectedKey( + source=KeySource.LITELLM_HEADER, + presented=keys.litellm_header, + token=litellm_header_token(keys.litellm_header), + ) + if keys.authorization: + return SelectedKey( + source=KeySource.AUTHORIZATION, + presented=keys.authorization, + token=scheme_credential(keys.authorization) or "", + ) + return _raw_key(keys) + + +def _subprotocol_api_key(subprotocols: str | None) -> str | None: + protocols: Final = tuple(protocol.strip() for protocol in (subprotocols or "").split(",")) + return next( + ( + protocol.removeprefix(_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX) + for protocol in protocols + if protocol.startswith(_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX) + ), + None, + ) + + +def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: + if keys.authorization: + split: Final = _split_scheme(keys.authorization) + if split is None or split[0] != "bearer": + return WebSocketDenial.INVALID_AUTHORIZATION + return split[1].strip() + return keys.api_key or _subprotocol_api_key(keys.subprotocols) or WebSocketDenial.NO_API_KEY diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py new file mode 100644 index 00000000000..b002d28b701 --- /dev/null +++ b/litellm/_auth/decision.py @@ -0,0 +1,433 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Protocol, TypeAlias + + +@dataclass(frozen=True, slots=True) +class RequestFacts: + route: str + token: str | None + is_llm_api_route: bool + + +@dataclass(frozen=True, slots=True) +class Settings: + master_key: str | None + allow_user_auth: bool + has_db: bool + master_key_only_routes: frozenset[str] + proxy_max_budget: float + + +@dataclass(frozen=True, slots=True) +class NoMasterKey: + token: str | None + + +@dataclass(frozen=True, slots=True) +class UserAuthRoute: + pass + + +@dataclass(frozen=True, slots=True) +class MasterKey: + pass + + +@dataclass(frozen=True, slots=True) +class CachedAdminKey: + pass + + +@dataclass(frozen=True, slots=True) +class VirtualKey: + budget_throttle_pct: float | None + + +@dataclass(frozen=True, slots=True) +class WithinModelBudgets: + pass + + +Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey | WithinModelBudgets + + +@dataclass(frozen=True, slots=True) +class MissingKey: + pass + + +@dataclass(frozen=True, slots=True) +class MalformedKey: + pass + + +@dataclass(frozen=True, slots=True) +class UserAuthDisabled: + pass + + +@dataclass(frozen=True, slots=True) +class ExpiredKey: + expires: datetime + now: datetime + key_label: str + + +@dataclass(frozen=True, slots=True) +class MasterKeyOnlyRoute: + pass + + +@dataclass(frozen=True, slots=True) +class NoDatabase: + pass + + +@dataclass(frozen=True, slots=True) +class NotAVirtualKey: + pass + + +@dataclass(frozen=True, slots=True) +class KeyBlocked: + pass + + +@dataclass(frozen=True, slots=True) +class OwnerDeactivated: + user_id: str + + +@dataclass(frozen=True, slots=True) +class TeamMemberBudgetExceeded: + user_id: str + team_id: str + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class KeyBudgetExceeded: + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class KeyWithoutToken: + pass + + +@dataclass(frozen=True, slots=True) +class ModelBudgetExceeded: + message: str + current_cost: float + max_budget: float + entity_type: str | None + entity_id: str | None + + +Denial: TypeAlias = ( + MissingKey + | MalformedKey + | UserAuthDisabled + | ExpiredKey + | MasterKeyOnlyRoute + | NoDatabase + | NotAVirtualKey + | KeyBlocked + | OwnerDeactivated + | TeamMemberBudgetExceeded + | KeyBudgetExceeded + | KeyWithoutToken + | ModelBudgetExceeded +) + + +@dataclass(frozen=True, slots=True) +class Allow: + grant: Grant + + +@dataclass(frozen=True, slots=True) +class Deny: + denial: Denial + + +Outcome: TypeAlias = Allow | Deny + + +@dataclass(frozen=True, slots=True) +class AdminKey: + expires: datetime | None + + +@dataclass(frozen=True, slots=True) +class NonAdminKey: + pass + + +PresentedKey: TypeAlias = AdminKey | NonAdminKey + + +@dataclass(frozen=True, slots=True) +class KeyFacts: + has_token: bool + blocked: bool + expires: datetime | None + spend: float + soft_budget: float | None + max_budget: float | None + team_member_spend: float | None + user_id: str | None + team_id: str | None + budget_throttle_pct: float | None + + +@dataclass(frozen=True, slots=True) +class OwnerFacts: + deactivated: bool + + +@dataclass(frozen=True, slots=True) +class TeamMemberBudget: + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class ModelBudgetScope: + models: tuple[str, ...] + rewritable_model: str | None + key_has_model_budgets: bool + key_budget_fallbacks: Mapping[str, tuple[str, ...]] + user_has_model_budgets: bool + + +@dataclass(frozen=True, slots=True) +class TeamMemberModelBudgetScope: + models: tuple[str, ...] + rewritable_model: str | None + has_model_budgets: bool + can_fall_back: bool + + +@dataclass(frozen=True, slots=True) +class FallbackTarget: + requested: str + model: str + served_model: str + + +@dataclass(frozen=True, slots=True) +class ResolveEndUser: + prefetch_key: bool + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ResolvePresentedKey: + hashed_token: str + resume: "Callable[[PresentedKey | None], Step]" + + +@dataclass(frozen=True, slots=True) +class EvictPresentedKey: + hashed_token: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKey: + hashed_token: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class PrepareKey: + loaded_from_db: bool + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class InspectKey: + resume: "Callable[[KeyFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class EnforceModelAccess: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKeyOwner: + user_id: str + resume: "Callable[[OwnerFacts | None], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckBudgetExemption: + resume: "Callable[[bool], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberBudget: + user_id: str + team_id: str + resume: "Callable[[TeamMemberBudget | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertTeamMemberBudget: + user_id: str + team_id: str + budget: TeamMemberBudget + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class SendKeyBudgetAlerts: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadKeySpend: + max_budget: float + resume: "Callable[[float], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertKeyBudget: + spend: float + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertKeySoftBudget: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKeyContext: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadProxySpend: + resume: "Callable[[float | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertProxyBudget: + spend: float + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadModelBudgetScope: + resume: "Callable[[ModelBudgetScope], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckKeyModelBudget: + model: str + resume: "Callable[[ModelBudgetExceeded | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AuthorizeBudgetFallback: + model: str + resume: "Callable[[bool], Step]" + + +@dataclass(frozen=True, slots=True) +class RewriteRequestModel: + model: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckUserModelBudget: + model: str + resume: "Callable[[ModelBudgetExceeded | None], Step]" + + +IdentityEffect: TypeAlias = ( + ResolveEndUser + | ResolvePresentedKey + | EvictPresentedKey + | LoadKey + | PrepareKey + | InspectKey + | EnforceModelAccess + | LoadKeyOwner + | CheckBudgetExemption + | LoadKeyContext +) +BudgetEffect: TypeAlias = ( + ReadTeamMemberBudget + | AlertTeamMemberBudget + | SendKeyBudgetAlerts + | ReadKeySpend + | AlertKeyBudget + | AlertKeySoftBudget + | ReadProxySpend + | AlertProxyBudget +) +ModelBudgetEffect: TypeAlias = ( + ReadModelBudgetScope | CheckKeyModelBudget | AuthorizeBudgetFallback | RewriteRequestModel | CheckUserModelBudget +) +Effect: TypeAlias = IdentityEffect | BudgetEffect | ModelBudgetEffect +Step: TypeAlias = Outcome | Effect + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberModelBudgetScope: + resume: "Callable[[TeamMemberModelBudgetScope], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class CheckTeamMemberModelBudget: + model: str + resume: "Callable[[ModelBudgetExceeded | None], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberFallbacks: + model: str + resume: "Callable[[tuple[FallbackTarget, ...]], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class VetTeamMemberFallback: + target: FallbackTarget + resume: "Callable[[bool], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class RerouteTeamMemberRequest: + model: str + resume: "Callable[[], TeamMemberStep]" + + +TeamMemberEffect: TypeAlias = ( + ReadTeamMemberModelBudgetScope + | CheckTeamMemberModelBudget + | ReadTeamMemberFallbacks + | VetTeamMemberFallback + | RerouteTeamMemberRequest +) +TeamMemberStep: TypeAlias = Outcome | TeamMemberEffect + + +class AuthCore(Protocol): + def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... + + def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... + + def check_model_budgets(self) -> Step: ... + + def check_user_model_budgets(self) -> Step: ... + + def check_team_member_model_budgets(self) -> TeamMemberStep: ... diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py new file mode 100644 index 00000000000..e447d7c0d04 --- /dev/null +++ b/litellm/_auth/driver.py @@ -0,0 +1,264 @@ +from typing import Protocol, assert_never + +from litellm._auth.decision import ( + AlertKeyBudget, + AlertKeySoftBudget, + AlertProxyBudget, + AlertTeamMemberBudget, + Allow, + AuthorizeBudgetFallback, + BudgetEffect, + CheckBudgetExemption, + CheckKeyModelBudget, + CheckTeamMemberModelBudget, + CheckUserModelBudget, + Deny, + Effect, + EnforceModelAccess, + EvictPresentedKey, + FallbackTarget, + IdentityEffect, + InspectKey, + KeyFacts, + LoadKey, + LoadKeyContext, + LoadKeyOwner, + ModelBudgetEffect, + ModelBudgetExceeded, + ModelBudgetScope, + Outcome, + OwnerFacts, + PrepareKey, + PresentedKey, + ReadKeySpend, + ReadModelBudgetScope, + ReadProxySpend, + ReadTeamMemberBudget, + ReadTeamMemberFallbacks, + ReadTeamMemberModelBudgetScope, + RerouteTeamMemberRequest, + ResolveEndUser, + ResolvePresentedKey, + RewriteRequestModel, + SendKeyBudgetAlerts, + Step, + TeamMemberBudget, + TeamMemberEffect, + TeamMemberModelBudgetScope, + TeamMemberStep, + VetTeamMemberFallback, +) + + +class ModelBudgetHandler(Protocol): + async def read_model_budget_scope(self) -> ModelBudgetScope: ... + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + + async def authorize_budget_fallback(self, model: str) -> bool: ... + + async def rewrite_request_model(self, model: str) -> None: ... + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + + +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 read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: ... + + async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: ... + + async def reroute_team_member_request(self, model: str) -> None: ... + + +class EffectHandler(ModelBudgetHandler, Protocol): + async def resolve_end_user(self, prefetch_key: bool) -> None: ... + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ... + + async def evict_presented_key(self, hashed_token: str) -> None: ... + + async def load_key(self, hashed_token: str) -> None: ... + + async def prepare_key(self, loaded_from_db: bool) -> None: ... + + async def inspect_key(self) -> KeyFacts: ... + + async def enforce_model_access(self) -> None: ... + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ... + + async def check_budget_exemption(self) -> bool: ... + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: ... + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: ... + + async def send_key_budget_alerts(self) -> None: ... + + async def read_key_spend(self, max_budget: float) -> float: ... + + async def alert_key_budget(self, spend: float) -> None: ... + + async def alert_key_soft_budget(self) -> None: ... + + async def load_key_context(self) -> None: ... + + async def read_proxy_spend(self) -> float | None: ... + + async def alert_proxy_budget(self, spend: float) -> None: ... + + +async def _answer(step: Effect, handler: EffectHandler) -> Step: + match step: + case ( + ResolveEndUser() + | ResolvePresentedKey() + | EvictPresentedKey() + | LoadKey() + | PrepareKey() + | InspectKey() + | EnforceModelAccess() + | LoadKeyOwner() + | CheckBudgetExemption() + | LoadKeyContext() + ): + return await _answer_identity(step, handler) + case ( + ReadModelBudgetScope() + | CheckKeyModelBudget() + | AuthorizeBudgetFallback() + | RewriteRequestModel() + | CheckUserModelBudget() + ): + return await _answer_model_budget(step, handler) + case _: + return await _answer_budget(step, handler) + + +async def _answer_identity(step: IdentityEffect, handler: EffectHandler) -> Step: + match step: + case ResolveEndUser(prefetch_key=prefetch_key, resume=resume): + await handler.resolve_end_user(prefetch_key) + return resume() + case ResolvePresentedKey(hashed_token=hashed_token, resume=resume): + return resume(await handler.resolve_presented_key(hashed_token)) + case EvictPresentedKey(hashed_token=hashed_token, resume=resume): + await handler.evict_presented_key(hashed_token) + return resume() + case LoadKey(hashed_token=hashed_token, resume=resume): + await handler.load_key(hashed_token) + return resume() + case PrepareKey(loaded_from_db=loaded_from_db, resume=resume): + await handler.prepare_key(loaded_from_db) + return resume() + case InspectKey(resume=resume): + return resume(await handler.inspect_key()) + case EnforceModelAccess(resume=resume): + await handler.enforce_model_access() + return resume() + case LoadKeyOwner(user_id=user_id, resume=resume): + return resume(await handler.load_key_owner(user_id)) + case CheckBudgetExemption(resume=resume): + return resume(await handler.check_budget_exemption()) + case LoadKeyContext(resume=resume): + await handler.load_key_context() + return resume() + case _: + assert_never(step) + + +async def _answer_budget(step: BudgetEffect, handler: EffectHandler) -> Step: + match step: + case ReadTeamMemberBudget(user_id=user_id, team_id=team_id, resume=resume): + return resume(await handler.read_team_member_budget(user_id, team_id)) + case AlertTeamMemberBudget(user_id=user_id, team_id=team_id, budget=budget, resume=resume): + await handler.alert_team_member_budget(user_id, team_id, budget) + return resume() + case SendKeyBudgetAlerts(resume=resume): + await handler.send_key_budget_alerts() + return resume() + case ReadKeySpend(max_budget=max_budget, resume=resume): + return resume(await handler.read_key_spend(max_budget)) + case AlertKeyBudget(spend=spend, resume=resume): + await handler.alert_key_budget(spend) + return resume() + case AlertKeySoftBudget(resume=resume): + await handler.alert_key_soft_budget() + return resume() + case ReadProxySpend(resume=resume): + return resume(await handler.read_proxy_spend()) + case AlertProxyBudget(spend=spend, resume=resume): + await handler.alert_proxy_budget(spend) + return resume() + case _: + assert_never(step) + + +async def _answer_model_budget(step: ModelBudgetEffect, handler: ModelBudgetHandler) -> Step: + match step: + case ReadModelBudgetScope(resume=resume): + return resume(await handler.read_model_budget_scope()) + case CheckKeyModelBudget(model=model, resume=resume): + return resume(await handler.check_key_model_budget(model)) + case AuthorizeBudgetFallback(model=model, resume=resume): + return resume(await handler.authorize_budget_fallback(model)) + case RewriteRequestModel(model=model, resume=resume): + await handler.rewrite_request_model(model) + return resume() + case CheckUserModelBudget(model=model, resume=resume): + return resume(await handler.check_user_model_budget(model)) + case _: + assert_never(step) + + +async def drive_model_budgets(step: Step, handler: ModelBudgetHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case ( + ReadModelBudgetScope() + | CheckKeyModelBudget() + | AuthorizeBudgetFallback() + | RewriteRequestModel() + | CheckUserModelBudget() + ): + return await drive_model_budgets(await _answer_model_budget(step, handler), handler) + case _: + raise RuntimeError(f"{type(step).__name__} is not a model budget effect") + + +async def drive(step: Step, handler: EffectHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case _: + return await drive(await _answer(step, handler), handler) + + +async def _answer_team_member(step: TeamMemberEffect, handler: TeamMemberModelBudgetHandler) -> TeamMemberStep: + match step: + case ReadTeamMemberModelBudgetScope(resume=resume): + return resume(await handler.read_team_member_model_budget_scope()) + case CheckTeamMemberModelBudget(model=model, resume=resume): + return resume(await handler.check_team_member_model_budget(model)) + case ReadTeamMemberFallbacks(model=model, resume=resume): + return resume(await handler.read_team_member_fallbacks(model)) + case VetTeamMemberFallback(target=target, resume=resume): + return resume(await handler.vet_team_member_fallback(target)) + case RerouteTeamMemberRequest(model=model, resume=resume): + await handler.reroute_team_member_request(model) + return resume() + case _: + assert_never(step) + + +async def drive_team_member_model_budgets(step: TeamMemberStep, handler: TeamMemberModelBudgetHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case _: + return await drive_team_member_model_budgets(await _answer_team_member(step, handler), handler) diff --git a/litellm/_auth/key.py b/litellm/_auth/key.py new file mode 100644 index 00000000000..af63778e3cf --- /dev/null +++ b/litellm/_auth/key.py @@ -0,0 +1,12 @@ +import hashlib +import secrets + + +def hash_token(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +def is_master_key(token: str | None, master_key: str | None) -> bool: + if token is None or master_key is None: + return False + return secrets.compare_digest(token.encode(), master_key.encode()) diff --git a/litellm/_auth/model_budgets.py b/litellm/_auth/model_budgets.py new file mode 100644 index 00000000000..7bf2e109287 --- /dev/null +++ b/litellm/_auth/model_budgets.py @@ -0,0 +1,151 @@ +from collections.abc import Callable +from typing import Final + +from litellm._auth.decision import ( + Allow, + AuthorizeBudgetFallback, + CheckKeyModelBudget, + CheckTeamMemberModelBudget, + CheckUserModelBudget, + Deny, + FallbackTarget, + ModelBudgetExceeded, + ModelBudgetScope, + ReadModelBudgetScope, + ReadTeamMemberFallbacks, + ReadTeamMemberModelBudgetScope, + RerouteTeamMemberRequest, + RewriteRequestModel, + Step, + TeamMemberModelBudgetScope, + TeamMemberStep, + VetTeamMemberFallback, + WithinModelBudgets, +) + + +def check_model_budgets(then: Callable[[], Step]) -> Step: + return ReadModelBudgetScope(resume=lambda scope: _check_key_budgets(scope, then)) + + +def check_user_model_budgets(then: Callable[[], Step]) -> Step: + return ReadModelBudgetScope(resume=lambda scope: _check_user_budgets(scope, scope.models, then)) + + +def _check_key_budgets(scope: ModelBudgetScope, then: Callable[[], Step]) -> Step: + if not scope.key_has_model_budgets or not scope.models: + return _check_user_budgets(scope, scope.models, then) + return _check_key_budget_of( + scope, + scope.models, + scope.rewritable_model, + then=lambda: ReadModelBudgetScope(resume=lambda rescoped: _check_user_budgets(rescoped, rescoped.models, then)), + ) + + +def _check_key_budget_of( + scope: ModelBudgetScope, models: tuple[str, ...], rewritable_model: str | None, then: Callable[[], Step] +) -> Step: + if not models: + return then() + model: Final = models[0] + rest: Final = models[1:] + + def after(exceeded: ModelBudgetExceeded | None) -> Step: + if exceeded is None: + return _check_key_budget_of(scope, rest, rewritable_model, then) + if model != rewritable_model: + 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)), + exhausted=lambda: Deny(exceeded), + ) + + return CheckKeyModelBudget(model=model, resume=after) + + +def _first_fallback_within_budget( + candidates: tuple[str, ...], found: Callable[[str], Step], exhausted: Callable[[], Step] +) -> Step: + if not candidates: + return exhausted() + candidate: Final = candidates[0] + return CheckKeyModelBudget( + model=candidate, + resume=lambda exceeded: ( + found(candidate) + if exceeded is None + else _first_fallback_within_budget(candidates[1:], found, exhausted) + ), + ) + + +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), + ) + + +def _check_user_budgets(scope: ModelBudgetScope, models: tuple[str, ...], then: Callable[[], Step]) -> Step: + if not scope.user_has_model_budgets or not models: + return then() + return CheckUserModelBudget( + model=models[0], + resume=lambda exceeded: _check_user_budgets(scope, models[1:], then) if exceeded is None else Deny(exceeded), + ) + + +def check_team_member_model_budgets() -> TeamMemberStep: + return ReadTeamMemberModelBudgetScope( + resume=lambda scope: ( + _check_team_member_budget_of(scope, scope.models, scope.rewritable_model) + if scope.has_model_budgets + else Allow(WithinModelBudgets()) + ) + ) + + +def _check_team_member_budget_of( + scope: TeamMemberModelBudgetScope, models: tuple[str, ...], rewritable_model: str | None +) -> TeamMemberStep: + if not models: + return Allow(WithinModelBudgets()) + model: Final = models[0] + rest: Final = models[1:] + + def reroute(target: FallbackTarget) -> TeamMemberStep: + return RerouteTeamMemberRequest( + model=target.model, resume=lambda: _check_team_member_budget_of(scope, rest, target.model) + ) + + def after(exceeded: ModelBudgetExceeded | None) -> TeamMemberStep: + if exceeded is None: + return _check_team_member_budget_of(scope, rest, rewritable_model) + if model != rewritable_model or not scope.can_fall_back: + return Deny(exceeded) + return ReadTeamMemberFallbacks( + model=model, + resume=lambda targets: _first_usable_fallback( + tuple(target for target in targets if target.model != model), + found=reroute, + exhausted=lambda: Deny(exceeded), + ), + ) + + return CheckTeamMemberModelBudget(model=model, resume=after) + + +def _first_usable_fallback( + targets: tuple[FallbackTarget, ...], + found: Callable[[FallbackTarget], TeamMemberStep], + exhausted: Callable[[], TeamMemberStep], +) -> TeamMemberStep: + if not targets: + return exhausted() + target: Final = targets[0] + return VetTeamMemberFallback( + target=target, + resume=lambda usable: found(target) if usable else _first_usable_fallback(targets[1:], found, exhausted), + ) diff --git a/litellm/_auth/python_core.py b/litellm/_auth/python_core.py new file mode 100644 index 00000000000..c7d6210bdc6 --- /dev/null +++ b/litellm/_auth/python_core.py @@ -0,0 +1,255 @@ +import math +from dataclasses import dataclass +from datetime import datetime +from typing import Final + +from litellm._auth.decision import ( + AdminKey, + AlertKeyBudget, + AlertKeySoftBudget, + AlertProxyBudget, + AlertTeamMemberBudget, + Allow, + CachedAdminKey, + CheckBudgetExemption, + Deny, + EnforceModelAccess, + EvictPresentedKey, + ExpiredKey, + InspectKey, + KeyBlocked, + KeyBudgetExceeded, + KeyFacts, + KeyWithoutToken, + LoadKey, + LoadKeyContext, + LoadKeyOwner, + MalformedKey, + MasterKey, + MasterKeyOnlyRoute, + MissingKey, + NoDatabase, + NoMasterKey, + NonAdminKey, + NotAVirtualKey, + OwnerDeactivated, + OwnerFacts, + PrepareKey, + PresentedKey, + ReadKeySpend, + ReadProxySpend, + ReadTeamMemberBudget, + RequestFacts, + ResolveEndUser, + ResolvePresentedKey, + SendKeyBudgetAlerts, + Settings, + Step, + TeamMemberBudget, + TeamMemberBudgetExceeded, + TeamMemberStep, + UserAuthDisabled, + UserAuthRoute, + VirtualKey, + WithinModelBudgets, +) +from litellm._auth.key import hash_token, is_master_key +from litellm._auth.model_budgets import ( + check_model_budgets, + check_team_member_model_budgets, + check_user_model_budgets, +) + +_USER_AUTH_ROUTE: Final = "/user/auth" +_VIRTUAL_KEY_PREFIX: Final = "sk-" + + +@dataclass(frozen=True, slots=True) +class _KeyCheck: + settings: Settings + now: datetime + is_llm_api_route: bool + key_label: str + + +class PythonCore: + def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: + token: Final = facts.token + if settings.master_key is None: + return Allow(NoMasterKey(token=token)) + if token is None: + return Deny(MissingKey()) + if token == "": + return Deny(MalformedKey()) + if facts.route == _USER_AUTH_ROUTE: + return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(UserAuthDisabled()) + if is_master_key(token, settings.master_key): + return ResolveEndUser(prefetch_key=False, resume=lambda: Allow(MasterKey())) + hashed_token: Final = hash_token(token) + return ResolveEndUser( + prefetch_key=True, + resume=lambda: ResolvePresentedKey( + hashed_token=hashed_token, + resume=lambda presented: _after_presented_key(presented, token, hashed_token, facts, settings, now), + ), + ) + + def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: + return _check_key(_KeyCheck(settings, now, facts.is_llm_api_route, key_label=facts.token or "")) + + def check_model_budgets(self) -> Step: + return check_model_budgets(then=lambda: Allow(WithinModelBudgets())) + + def check_user_model_budgets(self) -> Step: + return check_user_model_budgets(then=lambda: Allow(WithinModelBudgets())) + + def check_team_member_model_budgets(self) -> TeamMemberStep: + return check_team_member_model_budgets() + + +def _after_presented_key( + presented: PresentedKey | None, + token: str, + hashed_token: str, + facts: RequestFacts, + settings: Settings, + now: datetime, +) -> Step: + match presented: + case AdminKey(expires=expires) if expires is not None and expires < now: + return EvictPresentedKey( + hashed_token=hashed_token, + resume=lambda: Deny(ExpiredKey(expires=expires, now=now, key_label=token)), + ) + case AdminKey(): + return Allow(CachedAdminKey()) + case NonAdminKey() | None: + pass + if facts.route in settings.master_key_only_routes: + return Deny(MasterKeyOnlyRoute()) + if not settings.has_db: + return Deny(NoDatabase()) + if presented is not None: + cached: Final = _KeyCheck(settings, now, facts.is_llm_api_route, key_label=token) + return PrepareKey(loaded_from_db=False, resume=lambda: _check_key(cached)) + if not token.startswith(_VIRTUAL_KEY_PREFIX): + return Deny(NotAVirtualKey()) + loaded: Final = _KeyCheck(settings, now, facts.is_llm_api_route, key_label=hashed_token) + return LoadKey( + hashed_token=hashed_token, + resume=lambda: PrepareKey(loaded_from_db=True, resume=lambda: _check_key(loaded)), + ) + + +def _check_key(check: _KeyCheck) -> Step: + return InspectKey(resume=lambda key: _after_inspect(check, key)) + + +def _after_inspect(check: _KeyCheck, key: KeyFacts) -> Step: + if key.blocked: + return Deny(KeyBlocked()) + return EnforceModelAccess(resume=lambda: _check_owner(check, key)) + + +def _check_owner(check: _KeyCheck, key: KeyFacts) -> Step: + user_id: Final = key.user_id + if user_id is None: + return _check_budget_exemption(check, key) + return LoadKeyOwner(user_id=user_id, resume=lambda owner: _after_owner(check, key, user_id, owner)) + + +def _after_owner(check: _KeyCheck, key: KeyFacts, user_id: str, owner: OwnerFacts | None) -> Step: + if owner is not None and owner.deactivated: + return Deny(OwnerDeactivated(user_id=user_id)) + return _check_budget_exemption(check, key) + + +def _check_budget_exemption(check: _KeyCheck, key: KeyFacts) -> Step: + return CheckBudgetExemption(resume=lambda exempt: _check_team_member_budget(check, key, exempt)) + + +def _check_team_member_budget(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: + user_id: Final = key.user_id + team_id: Final = key.team_id + if exempt or key.team_member_spend is None or not check.settings.has_db or user_id is None or team_id is None: + return _check_expiry(check, key, exempt) + return ReadTeamMemberBudget( + user_id=user_id, + team_id=team_id, + resume=lambda budget: _after_team_member_budget(check, key, user_id, team_id, budget), + ) + + +def _after_team_member_budget( + check: _KeyCheck, key: KeyFacts, user_id: str, team_id: str, budget: TeamMemberBudget | None +) -> Step: + if budget is None or budget.spend < budget.max_budget: + return _check_expiry(check, key, exempt=False) + return AlertTeamMemberBudget( + user_id=user_id, + team_id=team_id, + budget=budget, + resume=lambda: Deny( + TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=budget.spend, max_budget=budget.max_budget) + ), + ) + + +def _check_expiry(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: + if key.expires is not None and key.expires < check.now: + return Deny(ExpiredKey(expires=key.expires, now=check.now, key_label=check.key_label)) + if exempt: + return _load_context(check, key, budget_throttle_pct=None) + return SendKeyBudgetAlerts(resume=lambda: _check_key_budget(check, key)) + + +def _check_key_budget(check: _KeyCheck, key: KeyFacts) -> Step: + max_budget: Final = key.max_budget + if not check.is_llm_api_route or max_budget is None: + return _check_soft_budget(check, key, budget_throttle_pct=None) + return ReadKeySpend( + max_budget=max_budget, + resume=lambda spend: AlertKeyBudget( + spend=spend, resume=lambda: _after_key_spend(check, key, spend, max_budget) + ), + ) + + +def _after_key_spend(check: _KeyCheck, key: KeyFacts, spend: float, max_budget: float) -> Step: + # A NaN max_budget compares false against any spend, so a non-finite budget means no limit, never a bypass. + if not math.isfinite(max_budget) or spend < max_budget: + return _check_soft_budget(check, key, budget_throttle_pct=None) + if key.budget_throttle_pct is not None: + return _check_soft_budget(check, key, budget_throttle_pct=key.budget_throttle_pct) + return Deny(KeyBudgetExceeded(spend=spend, max_budget=max_budget)) + + +def _check_soft_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + def model_budgets() -> Step: + return check_model_budgets(then=lambda: _load_context(check, key, budget_throttle_pct)) + + if key.soft_budget and key.spend >= key.soft_budget: + return AlertKeySoftBudget(resume=model_budgets) + return model_budgets() + + +def _load_context(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + return LoadKeyContext(resume=lambda: _check_proxy_budget(check, key, budget_throttle_pct)) + + +def _check_proxy_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + if not (check.settings.proxy_max_budget > 0 and check.settings.has_db): + return _grant(key, budget_throttle_pct) + return ReadProxySpend( + resume=lambda spend: ( + _grant(key, budget_throttle_pct) + if spend is None + else AlertProxyBudget(spend=spend, resume=lambda: _grant(key, budget_throttle_pct)) + ) + ) + + +def _grant(key: KeyFacts, budget_throttle_pct: float | None) -> Step: + if not key.has_token: + return Deny(KeyWithoutToken()) + return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) diff --git a/litellm/_auth/runtime.py b/litellm/_auth/runtime.py new file mode 100644 index 00000000000..63544108511 --- /dev/null +++ b/litellm/_auth/runtime.py @@ -0,0 +1,88 @@ +import os +from typing import Final, Literal, TypeAlias + +from litellm._auth import credential +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 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 + + +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 + + +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 + + +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 diff --git a/litellm/_auth/rust_core.py b/litellm/_auth/rust_core.py new file mode 100644 index 00000000000..1f2a3cf5bd8 --- /dev/null +++ b/litellm/_auth/rust_core.py @@ -0,0 +1,241 @@ +from collections.abc import Callable, Mapping +from dataclasses import fields, is_dataclass +from datetime import datetime +from enum import Enum +from typing import ( + Final, + TypeAlias, + cast, # noqa: TID251 # dataclass and enum reflection erases field types +) + +from pydantic import TypeAdapter + +from litellm._auth import decision +from litellm._auth.credential import KeySource, PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys +from litellm.rust_bridge import _native + +Wire: TypeAlias = None | bool | int | float | str | list["Wire"] | dict[str, "Wire"] +CoreStep: TypeAlias = decision.Step | decision.TeamMemberStep + +OBJECT_MAPPING: Final = TypeAdapter(dict[str, object]) +OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) +AUTH_STEP: Final = TypeAdapter[decision.Step](decision.Step) +MEMBER_STEP: Final = TypeAdapter[decision.TeamMemberStep](decision.TeamMemberStep) +CORE_STEP: Final = TypeAdapter[CoreStep](CoreStep) +TEAM_BUDGET: Final = TypeAdapter(decision.TeamMemberBudget) +FALLBACK_TARGET: Final = TypeAdapter(decision.FallbackTarget) +TEAM_ALERT: Final = TypeAdapter(decision.AlertTeamMemberBudget) +MEMBER_FALLBACK: Final = TypeAdapter(decision.VetTeamMemberFallback) + +_PUBLIC_TYPES: Final[Mapping[str, Callable[..., object]]] = { + cls.__name__: cls + for cls in ( + decision.NoMasterKey, + decision.UserAuthRoute, + decision.MasterKey, + decision.CachedAdminKey, + decision.VirtualKey, + decision.WithinModelBudgets, + decision.MissingKey, + decision.MalformedKey, + decision.UserAuthDisabled, + decision.ExpiredKey, + decision.MasterKeyOnlyRoute, + decision.NoDatabase, + decision.NotAVirtualKey, + decision.KeyBlocked, + decision.OwnerDeactivated, + decision.TeamMemberBudgetExceeded, + decision.KeyBudgetExceeded, + decision.KeyWithoutToken, + decision.ModelBudgetExceeded, + decision.Allow, + decision.Deny, + decision.ResolveEndUser, + decision.ResolvePresentedKey, + decision.EvictPresentedKey, + decision.LoadKey, + decision.PrepareKey, + decision.InspectKey, + decision.EnforceModelAccess, + decision.LoadKeyOwner, + decision.CheckBudgetExemption, + decision.ReadTeamMemberBudget, + decision.AlertTeamMemberBudget, + decision.SendKeyBudgetAlerts, + decision.ReadKeySpend, + decision.AlertKeyBudget, + decision.AlertKeySoftBudget, + decision.LoadKeyContext, + decision.ReadProxySpend, + decision.AlertProxyBudget, + decision.ReadModelBudgetScope, + decision.CheckKeyModelBudget, + decision.AuthorizeBudgetFallback, + decision.RewriteRequestModel, + decision.CheckUserModelBudget, + decision.ReadTeamMemberModelBudgetScope, + decision.CheckTeamMemberModelBudget, + decision.ReadTeamMemberFallbacks, + decision.VetTeamMemberFallback, + decision.RerouteTeamMemberRequest, + ) +} +_ANSWER_TYPES: Final = { + "ResolvePresentedKey": "Presented", + "InspectKey": "Key", + "LoadKeyOwner": "Owner", + "CheckBudgetExemption": "Bool", + "ReadTeamMemberBudget": "TeamBudget", + "ReadKeySpend": "Spend", + "ReadProxySpend": "ProxySpend", + "ReadModelBudgetScope": "ModelScope", + "CheckKeyModelBudget": "Exceeded", + "AuthorizeBudgetFallback": "Bool", + "CheckUserModelBudget": "Exceeded", + "ReadTeamMemberModelBudgetScope": "TeamScope", + "CheckTeamMemberModelBudget": "Exceeded", + "ReadTeamMemberFallbacks": "Targets", + "VetTeamMemberFallback": "Bool", +} + + +class _Projection: + def __init__(self) -> None: + self.datetimes: tuple[datetime, ...] = () + + def encode(self, value: object) -> Wire: + if isinstance(value, datetime): + offset: Final = value.utcoffset() + reference: Final = len(self.datetimes) + self.datetimes = (*self.datetimes, value) + return { + "wall_microseconds": ( + (value.toordinal() * 86400 + value.hour * 3600 + value.minute * 60 + value.second) * 1000000 + + value.microsecond + ), + "offset_microseconds": None + if offset is None + else (offset.days * 86400 + offset.seconds) * 1000000 + offset.microseconds, + "timezone_id": id(value.tzinfo), + "reference": reference, + } + if isinstance(value, Enum): + return self.encode(cast(object, value.value)) # cast-ok: Enum.value is dynamically typed + if value is None or isinstance(value, (str, bool, int, float)): + return value + if is_dataclass(value) and not isinstance(value, type): + return { + "kind": type(value).__name__, + **{field.name: self.encode(field_value(value, field.name)) for field in fields(value)}, + } + if isinstance(value, Mapping): + return {key: self.encode(item) for key, item in OBJECT_MAPPING.validate_python(value).items()} + if isinstance(value, (tuple, list, frozenset)): + return [self.encode(item) for item in OBJECT_TUPLE.validate_python(value)] + raise TypeError("unsupported auth protocol value") + + def decode(self, value: object) -> object: + if not isinstance(value, dict): + return value + data: Final = OBJECT_MAPPING.validate_python(value) + if "wall_microseconds" in data: + reference: Final = data["reference"] + if not isinstance(reference, int): + raise TypeError("invalid auth datetime reference") + return self.datetimes[reference] + kind: Final = data.get("kind") + if kind is None: + return {key: self.decode(item) for key, item in data.items()} + if not isinstance(kind, str): + raise TypeError("invalid auth protocol tag") + constructor: Final = public_type(kind) + return constructor(**{key: self.decode(item) for key, item in data.items() if key != "kind"}) + + def step(self, result: tuple[Mapping[str, object], _native.NativeAuthContinuation | None]) -> CoreStep: + data, continuation = result + kind: Final = data["kind"] + if not isinstance(kind, str): + raise TypeError("invalid auth protocol tag") + if continuation is None: + return CORE_STEP.validate_python(self.decode(data)) + + def resume(*answers: object) -> CoreStep: + answer: Final = { + "kind": _ANSWER_TYPES.get(kind, "Unit"), + "value": self.encode(answers[0]) if answers else None, + } + return self.step(continuation.resume(answer)) + + payload: Final = {key: self.decode(value) for key, value in data.items() if key != "kind"} + if kind == "AlertTeamMemberBudget": + return TEAM_ALERT.validate_python( + {**payload, "budget": TEAM_BUDGET.validate_python(payload["budget"]), "resume": resume} + ) + if kind == "VetTeamMemberFallback": + return MEMBER_FALLBACK.validate_python( + {**payload, "target": FALLBACK_TARGET.validate_python(payload["target"]), "resume": resume} + ) + constructor: Final = public_type(kind) + return CORE_STEP.validate_python(constructor(**payload, resume=resume)) + + +def field_value(value: object, name: str) -> object: + return cast(object, getattr(value, name)) # cast-ok: dataclass fields are discovered dynamically + + +def public_type(kind: str) -> Callable[..., object]: + return _PUBLIC_TYPES[kind] + + +class RustCore: + def start(self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime) -> decision.Step: + projection: Final = _Projection() + return AUTH_STEP.validate_python( + projection.step( + _native.request_auth_start( + projection.encode(facts), projection.encode(settings), projection.encode(now), False + ) + ), + ) + + def start_resolved_key( + self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime + ) -> decision.Step: + projection: Final = _Projection() + return AUTH_STEP.validate_python( + projection.step( + _native.request_auth_start( + projection.encode(facts), projection.encode(settings), projection.encode(now), True + ) + ), + ) + + def check_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python(_Projection().step(_native.request_auth_budgets("Key"))) + + def check_user_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python(_Projection().step(_native.request_auth_budgets("User"))) + + def check_team_member_model_budgets(self) -> decision.TeamMemberStep: + return MEMBER_STEP.validate_python(_Projection().step(_native.request_auth_budgets("TeamMember"))) + + +def select_api_key(keys: PresentedKeys) -> SelectedKey | None: + selected: Final = _native.request_auth_select_api_key(_Projection().encode(keys)) + if selected is None: + return None + return SelectedKey(source=KeySource(selected["source"]), presented=selected["presented"], token=selected["token"]) + + +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"] + + +def scheme_credential(value: str) -> str | None: + return _native.request_auth_scheme_credential(value) + + +def litellm_header_token(value: str) -> str: + return _native.request_auth_litellm_header_token(value) diff --git a/litellm/_auth/shadow.py b/litellm/_auth/shadow.py new file mode 100644 index 00000000000..fc3bd4073e2 --- /dev/null +++ b/litellm/_auth/shadow.py @@ -0,0 +1,115 @@ +import math +from collections.abc import Callable, Mapping +from dataclasses import dataclass, fields, is_dataclass, replace +from datetime import datetime +from functools import reduce +from typing import Final + +from litellm._auth import decision +from litellm._auth.python_core import PythonCore +from litellm._auth.rust_core import ( + AUTH_STEP, + MEMBER_STEP, + OBJECT_MAPPING, + OBJECT_TUPLE, + CoreStep, + RustCore, + field_value, +) + + +class ShadowMismatch(RuntimeError): + pass + + +def snapshot(value: object) -> object: + if isinstance(value, float) and math.isnan(value): + return (float, "nan") + if is_dataclass(value) and not isinstance(value, type): + return ( + type(value), + tuple( + (field.name, snapshot(field_value(value, field.name))) + for field in fields(value) + if field.name != "resume" + ), + ) + if isinstance(value, Mapping): + return tuple((key, snapshot(item)) for key, item in OBJECT_MAPPING.validate_python(value).items()) + if isinstance(value, (tuple, list)): + return tuple(snapshot(item) for item in OBJECT_TUPLE.validate_python(value)) + return value + + +def compare(primary: object, reference: object) -> None: + if snapshot(primary) != snapshot(reference): + raise ShadowMismatch(f"auth shadow mismatch at {type(primary).__name__}") + + +@dataclass(frozen=True, slots=True) +class _RecordedEffect: + signature: object + answers: tuple[object, ...] + + +def _replay(reference: CoreStep, record: _RecordedEffect) -> CoreStep: + if snapshot(reference) != record.signature or isinstance(reference, (decision.Allow, decision.Deny)): + raise ShadowMismatch(f"auth shadow effect mismatch at {type(reference).__name__}") + return _resume(reference.resume, record.answers) + + +def _resume(continuation: Callable[..., CoreStep], answers: tuple[object, ...]) -> CoreStep: + return continuation(*answers) + + +def _shadow( + primary: CoreStep, reference: Callable[[], CoreStep], records: tuple[_RecordedEffect, ...] = () +) -> CoreStep: + if isinstance(primary, (decision.Allow, decision.Deny)): + compare(primary, reduce(_replay, records, reference())) + return primary + primary_resume: Final[Callable[..., CoreStep]] = primary.resume + signature: Final = snapshot(primary) + + def resume(*recorded_answers: object) -> CoreStep: + recorded: Final = _RecordedEffect(signature, recorded_answers) + next_primary: Final = _resume(primary_resume, recorded_answers) + return _shadow(next_primary, reference, (*records, recorded)) + + return replace(primary, resume=resume) + + +class ShadowCore: + def __init__(self, primary: decision.AuthCore | None = None, reference: decision.AuthCore | None = None) -> None: + self.primary: Final = primary if primary is not None else RustCore() + self.reference: Final = reference if reference is not None else PythonCore() + + def start(self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow(self.primary.start(facts, settings, now), lambda: self.reference.start(facts, settings, now)), + ) + + def start_resolved_key( + self, facts: decision.RequestFacts, settings: decision.Settings, now: datetime + ) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow( + self.primary.start_resolved_key(facts, settings, now), + lambda: self.reference.start_resolved_key(facts, settings, now), + ), + ) + + def check_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow(self.primary.check_model_budgets(), self.reference.check_model_budgets) + ) + + def check_user_model_budgets(self) -> decision.Step: + return AUTH_STEP.validate_python( + _shadow(self.primary.check_user_model_budgets(), self.reference.check_user_model_budgets) + ) + + def check_team_member_model_budgets(self) -> decision.TeamMemberStep: + return MEMBER_STEP.validate_python( + _shadow(self.primary.check_team_member_model_budgets(), self.reference.check_team_member_model_budgets), + ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4514882a45f..71a47df48ca 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6720,6 +6720,53 @@ def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: 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: + # This message is returned to the caller, and key_name has no enforced + # shape (a direct DB write bypasses abbreviate_api_key), so echo it only + # when it still looks masked and fall back to the alias otherwise. + 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, @@ -6734,47 +6781,13 @@ async def virtual_key_max_budget_check( if valid_token.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend - fallback_spend: Final = valid_token.spend or 0.0 - counter_key: Final = key_counter(str(valid_token.token)) - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) spend: Final = await get_current_spend( - counter_key=counter_key, - fallback_spend=fallback_spend, + counter_key=key_counter(str(valid_token.token)), + fallback_spend=valid_token.spend or 0.0, max_budget=valid_token.max_budget, ) - - #################################### - # collect information for alerting # - #################################### - - user_email = None - # Check if the token has any user id information - if user_obj is not None: - user_email = user_obj.user_email - - 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_email, - 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, - ) - ) - - #################################### - # collect information for alerting # - #################################### + 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 @@ -6782,22 +6795,7 @@ async def virtual_key_max_budget_check( if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: if _apply_budget_exceeded_throttle(valid_token): return - # This message is returned to the caller, and key_name has no enforced - # shape (a direct DB write bypasses abbreviate_api_key), so echo it only - # when it still looks masked and fall back to the alias otherwise. - 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 - ) - raise litellm.BudgetExceededError( - current_cost=spend, - max_budget=valid_token.max_budget, - message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", - entity_type=Litellm_EntityType.KEY.value, - entity_id=valid_token.token, - ) + raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget) _virtual_key_max_budget_check: Final = virtual_key_max_budget_check @@ -6858,32 +6856,39 @@ async def virtual_key_soft_budget_check( """ if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: - verbose_proxy_logger.debug( - "Crossed Soft Budget for token %s, spend %s, soft_budget %s", - valid_token.token, - valid_token.spend, - valid_token.soft_budget, - ) - 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, - ) + send_soft_budget_alert(valid_token, proxy_logging_obj, user_obj) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) + +def send_soft_budget_alert( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + verbose_proxy_logger.debug( + "Crossed Soft Budget for token %s, spend %s, soft_budget %s", + valid_token.token, + valid_token.spend, + valid_token.soft_budget, + ) + 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, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, ) + ) _virtual_key_soft_budget_check: Final = virtual_key_soft_budget_check diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py new file mode 100644 index 00000000000..5b872561e6f --- /dev/null +++ b/litellm/proxy/auth/auth_denial.py @@ -0,0 +1,146 @@ +from typing import Final, assert_never + +from fastapi import HTTPException, status + +import litellm +from litellm._auth.decision import ( + Denial, + ExpiredKey, + KeyBlocked, + KeyBudgetExceeded, + KeyWithoutToken, + MalformedKey, + MasterKeyOnlyRoute, + MissingKey, + ModelBudgetExceeded, + NoDatabase, + NotAVirtualKey, + OwnerDeactivated, + TeamMemberBudgetExceeded, + UserAuthDisabled, +) +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.types.proxy.auth.user_api_key_auth import is_jwt + +_JWT_AUTH_DISABLED_HINT: Final = ( + " This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a" + " virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate" + " with JWTs." +) + + +def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException: + masked: Final = f"{token[:4]}****{token[-4:]}" if len(token) > 8 else "****" + hint: Final = _JWT_AUTH_DISABLED_HINT if not jwt_auth_enabled and is_jwt(token=token) else "" + error: Final = HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={masked}, expected to start with 'sk-'.{hint}", + ) + # Log routing classifies this 401 by where it was raised, never by its message text. + setattr(error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) + return error + + +def denial_exception( + denial: Denial, + token: str | None, + route: str, + jwt_auth_enabled: bool, + key: UserAPIKeyAuth | None, +) -> ProxyException | HTTPException | litellm.BudgetExceededError: + match denial: + case MissingKey(): + return ProxyException( + message="Authentication Error, No api key passed in.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case MalformedKey(): + return ProxyException( + message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case UserAuthDisabled(): + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="'allow_user_auth' not set or set to False", + ) + case ExpiredKey(expires=expires, now=now, key_label=key_label): + return ProxyException( + message=f"Authentication Error - Expired Key. Key Expiry time {expires} and current time {now}", + type=ProxyErrorTypes.expired_key, + code=status.HTTP_401_UNAUTHORIZED, + param=abbreviate_api_key(api_key=key_label), + ) + case MasterKeyOnlyRoute(): + return ProxyException( + message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case NoDatabase(): + return ProxyException( + message="No connected db.", + type=ProxyErrorTypes.no_db_connection, + code=400, + param=None, + ) + case NotAVirtualKey(): + return _not_a_virtual_key(token or "", jwt_auth_enabled) + case KeyBlocked(): + return ProxyException( + message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case OwnerDeactivated(user_id=user_id): + return ProxyException( + message=( + f"Authentication Error, User={user_id} has been deactivated via SCIM. " + "Keys owned by this user cannot be used." + ), + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=spend, max_budget=max_budget): + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=( + f"Budget has been exceeded! TeamMember={user_id}:{team_id} " + f"Current cost: {spend}, Max budget: {max_budget}" + ), + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{user_id}:{team_id}", + ) + case KeyBudgetExceeded(spend=spend, max_budget=max_budget): + return key_budget_exceeded_error(key or UserAPIKeyAuth(), spend, max_budget) + case KeyWithoutToken(): + return HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key, no token associated" + ) + case ModelBudgetExceeded( + message=message, + current_cost=current_cost, + max_budget=max_budget, + entity_type=entity_type, + entity_id=entity_id, + ): + return litellm.BudgetExceededError( + current_cost=current_cost, + max_budget=max_budget, + message=message, + entity_type=entity_type, + entity_id=entity_id, + ) + case _: + assert_never(denial) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a52ad14e293..a97335f5061 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -9,13 +9,11 @@ Returns a UserAPIKeyAuth object if the API key is valid import asyncio import fnmatch -import re -import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime, timezone from enum import Enum from types import MappingProxyType -from typing import Final, NamedTuple, Protocol, Union, cast +from typing import Final, NamedTuple, Protocol, Union, assert_never, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -24,14 +22,45 @@ from pydantic import TypeAdapter from starlette.exceptions import WebSocketException import litellm +from litellm._auth.credential import ( + PresentedKeys, + WebSocketDenial, + WebSocketKeys, +) +from litellm._auth.decision import ( + AdminKey, + Allow, + AuthCore, + CachedAdminKey, + Deny, + FallbackTarget, + Grant, + KeyFacts, + MasterKey, + ModelBudgetExceeded, + ModelBudgetScope, + NoMasterKey, + NonAdminKey, + Outcome, + OwnerFacts, + PresentedKey, + RequestFacts, + Settings, + Step, + TeamMemberBudget, + TeamMemberModelBudgetScope, + UserAuthRoute, + VirtualKey, + WithinModelBudgets, +) +from litellm._auth.driver import drive, drive_model_budgets, drive_team_member_model_budgets +from litellm._auth.runtime import get_auth_core, scheme_credential, select_api_key, select_websocket_api_key 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 from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, GLOBAL_PROXY_SPEND_CACHE_KEY, - INVALID_VIRTUAL_KEY_ERROR_MARKER, - INVALID_VIRTUAL_KEY_ERROR_MESSAGE, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY, @@ -69,6 +98,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo check_end_user_budget, common_checks, delete_cache_key_object, + get_default_end_user_budget, get_end_user_object, get_jwt_key_mapping_object, get_key_end_user_budget_id, @@ -91,12 +121,15 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo 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, - virtual_key_soft_budget_check, + virtual_key_soft_budget_check, # pyright: ignore[reportUnusedImport] # backwards-compatible package export ) +from litellm.proxy.auth.auth_denial import denial_exception from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_method import AuthMethod from litellm.proxy.auth.auth_object_prefetch import ( @@ -121,6 +154,7 @@ from litellm.proxy.auth.auth_utils import ( request_dispatched_to_provider_pass_through, route_in_additonal_public_routes, ) +from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, should_throttle_budget_exceeded from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler @@ -279,6 +313,9 @@ _CLAUDE_MODEL_ROUTES: Final = frozenset( ) _CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized" _GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_PASS_THROUGH_HEADERS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( + Mapping[str, object] | None +) async def _normalize_claude_model( @@ -409,13 +446,9 @@ def _budget_fallback_target_model(target: object) -> str | None: return model if isinstance(model, str) else None -class _KeyModelBudgetLimiter(Protocol): +class _KeyAndUserModelBudgetLimiter(Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... - async def get_fallback_model_within_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> str | None: ... - - -class _UserModelBudgetLimiter(Protocol): 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: ... @@ -542,28 +575,6 @@ async def _read_user_model_max_budget( return user_obj.model_max_budget if user_obj is not None else None -async def _check_user_model_budget( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _UserModelBudgetLimiter, - models: list[str], -) -> None: - """Enforce the internal user's own `model_max_budget` across the request's models. - - Separate from the key check: a user's per-model budget caps every key they - own, so a caller cannot escape it by minting another key. - """ - user_model_max_budget: Final = valid_token.user_model_max_budget - if valid_token.user_id is None or not isinstance(user_model_max_budget, Mapping) or not user_model_max_budget: - return - for model_name in models: - 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_name, - team_id=valid_token.team_id, - ) - - async def _check_team_model_budget( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _TeamModelBudgetLimiter, @@ -780,6 +791,366 @@ async def _is_within_all_fallback_model_budgets( 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 + + +class _ModelBudgetEffects: + def __init__( + self, + *, + request: Request, + request_data: dict[str, object], # mutable-ok: budget fallbacks rewrite the request model in place + valid_token: UserAPIKeyAuth, + route: str, + limiter: _KeyAndUserModelBudgetLimiter, + llm_router: litellm.Router | None, + llm_model_list: Sequence[object] | None, + ) -> None: + self._request: Final = request + self._request_data: Final = request_data + self._valid_token: Final = valid_token + self._route: Final = route + self._limiter: Final = limiter + self._llm_router: Final = llm_router + self._llm_model_list: Final = llm_model_list + + @property + def valid_token(self) -> UserAPIKeyAuth: + return self._valid_token + + @property + def route(self) -> str: + return self._route + + async def read_model_budget_scope(self) -> ModelBudgetScope: + requested: Final = self._request_data.get("model") + user_model_max_budget: Final = self._valid_token.user_model_max_budget + budget_fallbacks: Final[Mapping[str, list[str]]] = self._valid_token.budget_fallbacks or {} + return ModelBudgetScope( + models=tuple( + _get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=self._request_data, + route=self._route, + request=self._request, + llm_router=self._llm_router, + team_id=self._valid_token.team_id, + ) + ) + ), + 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_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), + ) + + 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 authorize_budget_fallback(self, model: str) -> bool: + return await _is_fallback_model_authorized( + 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: + 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, + ) + ) + + +async def _enforce_model_budgets(step: Step, effects: _ModelBudgetEffects) -> None: + match await drive_model_budgets(step, effects): + case Deny(denial=denial): + raise denial_exception( + denial, + token=effects.valid_token.token, + route=effects.route, + jwt_auth_enabled=False, + key=effects.valid_token, + ) + case Allow(): + return + + +def _reroute_request( + request_data: dict[str, object], # mutable-ok: the fallback model replaces the requested one in place + request: Request, + model: str, +) -> None: + rewrite_request_model(request_data=request_data, request=request, model=model) + if "path_params" in request.scope: + request.scope["path_params"]["model"] = model # rebind-ok: route handlers use the fallback model + + +def _budget_fallback_candidates(targets: Sequence[object]) -> tuple[str, ...]: + return tuple(model for target in targets if (model := _budget_fallback_target_model(target)) is not None) + + +class _TeamMemberModelBudgetEffects: + def __init__( + self, + *, + valid_token: UserAPIKeyAuth, + limiter: _EveryModelBudgetLimiter, + models: Sequence[str], + request_data: dict[str, object], # mutable-ok: fallback rewrites the request body + request: Request, + llm_model_list: Sequence[object] | None, + llm_router: litellm.Router | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + team_object: LiteLLM_TeamTable | None, + end_user_object: LiteLLM_EndUserTable | None, + route: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + ) -> None: + self._valid_token: Final = valid_token + self._limiter: Final = limiter + self._models: Final = tuple(models) + self._request_data: Final = request_data + self._request: Final = request + self._llm_model_list: Final = llm_model_list + self._llm_router: Final = llm_router + self._project_object: Final = project_object + self._team_object: Final = team_object + self._end_user_object: Final = end_user_object + self._route: Final = route + self._prisma_client: Final = prisma_client + self._user_api_key_cache: Final = user_api_key_cache + self._proxy_logging_obj: Final = proxy_logging_obj + + def _resolved_dependencies(self) -> tuple[PrismaClient | None, UserApiKeyCache, ProxyLogging]: + from litellm.proxy import proxy_server + + return ( + self._prisma_client if self._prisma_client is not None else proxy_server.prisma_client, + self._user_api_key_cache if self._user_api_key_cache is not None else proxy_server.user_api_key_cache, + self._proxy_logging_obj if self._proxy_logging_obj is not None else proxy_server.proxy_logging_obj, + ) + + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: + budget: Final = self._valid_token.team_member_model_max_budget + requested: Final = self._request_data.get("model") + 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), + ) + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: + 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 + ) + ) + + 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) + ): + return () + from litellm.proxy.proxy_server import proxy_config + + prisma_client, _, proxy_logging_obj = self._resolved_dependencies() + try: + router_settings: Final[dict[str, object] | None] = ( + cast( # cast-ok: proxy settings return an unparameterized dict + dict[str, object] | None, + await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=self._valid_token, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. + return () + + configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None + router_fallbacks: Final[Sequence[object] | None] = ( + cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below + if isinstance(configured_fallbacks, list) + else _router_fallbacks(llm_router) + ) + if router_fallbacks is None: + return () + fallback_models, _ = request_fallback_chain(list(router_fallbacks), self._request_data, model) + model_group_aliases: Final[object] = ( + router_settings.get("model_group_alias") if router_settings is not None else None + ) + key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime + object, self._valid_token.aliases + ) + resolved: Final[tuple[tuple[str, str], ...]] = tuple( + ( + candidate, + target + if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None + else candidate, + ) + for candidate in _budget_fallback_candidates(fallback_models or ()) + ) + return tuple( + FallbackTarget( + requested=candidate, + model=fallback_model, + served_model=resolve_served_model( + fallback_model, litellm.model_alias_map, key_aliases, model_group_aliases + ), + ) + 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 _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( + 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, + 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, + ) + 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) + + async def _check_team_member_model_budget_with_fallback( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _EveryModelBudgetLimiter, @@ -796,352 +1167,29 @@ async def _check_team_member_model_budget_with_fallback( user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> None: - team_member_model_max_budget: Final = valid_token.team_member_model_max_budget - if ( - valid_token.user_id is None - or valid_token.team_id is None - or not isinstance(team_member_model_max_budget, Mapping) - or not team_member_model_max_budget - ): - return - - user_id: Final = valid_token.user_id - team_id: Final = valid_token.team_id - for model_name in models: - await _check_team_member_model_budget_for_model( - user_id=user_id, - team_id=team_id, - team_member_model_max_budget=team_member_model_max_budget, - model_name=model_name, - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - request_data=request_data, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - project_object=project_object, - team_object=team_object, - end_user_object=end_user_object, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - -async def _check_team_member_model_budget_for_model( - user_id: str, - team_id: str, - team_member_model_max_budget: Mapping[str, object], - model_name: str, - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router | None, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> None: - try: - await model_max_budget_limiter.is_team_member_within_model_budget( - user_id=user_id, - team_id=team_id, - team_member_model_max_budget=team_member_model_max_budget, - model=model_name, - ) - except litellm.BudgetExceededError: - if request_data.get("model") != model_name or llm_router is None or fallbacks_disabled(request_data): - raise - fallback_succeeded: Final = await _try_team_member_model_budget_fallback( - user_id=user_id, - team_id=team_id, - model_name=model_name, - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - request_data=request_data, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - project_object=project_object, - team_object=team_object, - end_user_object=end_user_object, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if not fallback_succeeded: - raise - - -async def _try_team_member_model_budget_fallback( - user_id: str, - team_id: str, - model_name: str, - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> bool: - if request.headers.get("content-type", "").lower().startswith( - "multipart/form-data" - ) or references_provider_scoped_resource(request_data): - return False - - from litellm.proxy import proxy_server - from litellm.proxy.proxy_server import proxy_config - - resolved_prisma_client: Final = prisma_client if prisma_client is not None else proxy_server.prisma_client - resolved_user_api_key_cache: Final = ( - user_api_key_cache if user_api_key_cache is not None else proxy_server.user_api_key_cache + effects: Final = _TeamMemberModelBudgetEffects( + valid_token=valid_token, + limiter=model_max_budget_limiter, + models=models, + request_data=request_data, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + project_object=project_object, + team_object=team_object, + end_user_object=end_user_object, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) - resolved_proxy_logging_obj: Final = ( - proxy_logging_obj if proxy_logging_obj is not None else proxy_server.proxy_logging_obj - ) - - try: - router_settings: Final[dict[str, object] | None] = ( - cast( # cast-ok: proxy settings return an unparameterized dict - dict[str, object] | None, - await proxy_config.get_hierarchical_router_settings( - user_api_key_dict=valid_token, - prisma_client=resolved_prisma_client, - proxy_logging_obj=resolved_proxy_logging_obj, - ), + match await drive_team_member_model_budgets(_AUTH_CORE.check_team_member_model_budgets(), effects): + case Deny(denial=denial): + raise denial_exception( + denial, token=valid_token.token, route=route, jwt_auth_enabled=False, key=valid_token ) - ) - except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. - return False - - configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None - router_fallbacks: Final[Sequence[object] | None] = ( - cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below - if isinstance(configured_fallbacks, list) - else _router_fallbacks(llm_router) - ) - if router_fallbacks is None: - return False - - fallback_models, _ = request_fallback_chain(list(router_fallbacks), request_data, model_name) - fallback_targets: Final[Sequence[object]] = fallback_models or () - fallback_candidates: Final[tuple[str, ...]] = tuple( - model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None - ) - if not fallback_candidates: - return False - - model_group_aliases: Final[object] = ( - router_settings.get("model_group_alias") if router_settings is not None else None - ) - key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime - object, valid_token.aliases - ) - resolved_fallback_candidates: Final[tuple[tuple[str, str], ...]] = tuple( - ( - candidate, - target if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None else candidate, - ) - for candidate in fallback_candidates - ) - fallback_candidates_with_served_models: Final = ( - ( - fallback_candidate, - fallback_model, - resolve_served_model( - fallback_model, - litellm.model_alias_map, - key_aliases, - model_group_aliases, - ), - ) - for fallback_candidate, fallback_model in resolved_fallback_candidates - ) - - try: - team_membership: Final = await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - proxy_logging_obj=resolved_proxy_logging_obj, - ) - except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback - return False - - for fallback_candidate, fallback_model, served_model in fallback_candidates_with_served_models: - if fallback_model == model_name: - continue - try: - if ( - litellm.block_requests_for_models_without_pricing - and RouteChecks.is_llm_api_route(route=route) - and _unpriced_models_in_request(model=served_model, llm_router=llm_router) - ): - continue - if end_user_object is not None and end_user_object.models: - can_customer_access_model( - model=fallback_model, - end_user_object=end_user_object, - llm_router=llm_router, - valid_token=valid_token, - ) - if not is_model_cost_zero( - model=served_model, llm_router=llm_router - ) and not await _is_within_all_fallback_model_budgets( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=tuple( - dict.fromkeys( - ( - fallback_candidate, - fallback_model, - served_model, - llm_router.registry().resolve(served_model), - ) - ) - ), - ): - continue - if not await _is_fallback_model_authorized( - model=fallback_model, - valid_token=valid_token, - llm_model_list=llm_model_list, - llm_router=llm_router, - prisma_client=prisma_client, - team_object=team_object, - team_membership=team_membership, - project_object=project_object, - refuse_agent_keys=True, - ): - continue - valid_token.matched_model_access_groups = [] # rebind-ok: clear prior groups before evaluating this candidate - matched_groups = await stamp_matched_model_access_groups( - model=fallback_model, - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - llm_router=llm_router, - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - proxy_logging_obj=resolved_proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=True, - ) - valid_token.matched_model_access_groups = list(matched_groups) # rebind-ok: spend uses the candidate groups - if matched_groups: - await _model_access_group_max_budget_check( - matched_model_access_groups=list(matched_groups), - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - ) - 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", - fallback_model.replace("\r", "").replace("\n", ""), - exc_info=True, - ) - continue - rewrite_request_model(request_data=request_data, request=request, model=fallback_model) - if "path_params" in request.scope: - request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model - return True - return False - - -async def _check_key_model_budget_with_fallback( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _KeyModelBudgetLimiter, - model_name: str, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None = None, - llm_router: litellm.Router | None = None, -) -> None: - """ - Enforce the key's per-model budget for `model_name`. If exceeded and the - key has a `budget_fallbacks` chain configured for `model_name`, reroute - the request to the first fallback model still within its own budget - instead of rejecting the request. - - The selected fallback is validated against the key's model-access - allowlist and the team's model restrictions so that budget_fallbacks - cannot bypass model authorization. The rewrite is persisted to the - parsed-body cache, Starlette's JSON cache (``request._json``), and - path parameters so that downstream handlers see the final model - regardless of whether they consume ``_read_request_body()``, - ``request.json()``, or the path ``model`` parameter. - - Fallback is only attempted when ``model_name`` matches the top-level - ``request_data["model"]``; models extracted from nested fields - (``session.model``, ``completion.model``, etc.) are not rewritable - and raise immediately. - - Raises: - BudgetExceededError: if `model_name` is over budget and no configured - fallback is within budget either (or the fallback is not authorized). - """ - try: - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=model_name, - ) - except litellm.BudgetExceededError as e: - if request_data.get("model") != model_name: - raise e - fallback_model: Final = await model_max_budget_limiter.get_fallback_model_within_budget( - user_api_key_dict=valid_token, - model=model_name, - ) - if fallback_model is None: - raise e - if not await _is_fallback_model_authorized( - model=fallback_model, - valid_token=valid_token, - llm_model_list=llm_model_list, - llm_router=llm_router, - ): - raise - rewrite_request_model(request_data=request_data, request=request, model=fallback_model) - if "path_params" in request.scope: - request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model - - -def _get_bearer_token_or_received_api_key(api_key: str) -> str: - if api_key.startswith("Bearer "): # ensure Bearer token passed in - api_key = api_key.replace("Bearer ", "") # extract the token - elif api_key.startswith("Basic "): - api_key = api_key.replace("Basic ", "") # handle langfuse input - elif api_key.startswith("bearer "): - api_key = api_key.replace("bearer ", "") - elif api_key.startswith("AWS4-HMAC-SHA256"): - # Handle AWS Signature V4 format from LangChain - # Format: AWS4-HMAC-SHA256 Credential=Bearer $LITELLM_MASTER_KEY/date/region/service/aws4_request, SignedHeaders=..., Signature=... - # Extract the Bearer token from the Credential field - match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r"Credential=([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - - return api_key + case Allow(): + return def _routing_selector_matches_claim( @@ -1236,29 +1284,7 @@ def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> def get_bearer_token( api_key: str, ) -> str: - if api_key.startswith("Bearer "): # ensure Bearer token passed in - api_key = api_key.replace("Bearer ", "") # extract the token - elif api_key.startswith("Basic "): - api_key = api_key.replace("Basic ", "") # handle langfuse input - elif api_key.startswith("bearer "): - api_key = api_key.replace("bearer ", "") - elif api_key.startswith("AWS4-HMAC-SHA256"): - # Handle AWS Signature V4 format from LangChain - # Format: AWS4-HMAC-SHA256 Credential=Bearer $LITELLM_MASTER_KEY/date/region/service/aws4_request, SignedHeaders=..., Signature=... - # Extract the Bearer token from the Credential field - match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r"Credential=([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - api_key = "" - else: - api_key = "" - return api_key + return scheme_credential(api_key) or "" _get_bearer_token: Final = get_bearer_token @@ -1295,6 +1321,14 @@ def _apply_budget_limits_to_end_user_params( verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) +_WEBSOCKET_DENIAL_DETAIL: Final = MappingProxyType( + { + WebSocketDenial.NO_API_KEY: "No API key provided", + WebSocketDenial.INVALID_AUTHORIZATION: "Invalid Authorization header format", + } +) + + async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model", "")) @@ -1355,27 +1389,17 @@ async def user_api_key_auth_websocket_for_model( request.body = return_body - authorization: Final = websocket.headers.get("authorization") - # If no Authorization header, try the api-key header - if not authorization: - api_key = websocket.headers.get("api-key") - if not api_key: - # Try extracting from WebSocket subprotocol (browser clients) - for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): - protocol = protocol.strip() - if protocol.startswith("openai-insecure-api-key."): - api_key = protocol[len("openai-insecure-api-key.") :] - break - if not api_key: - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="No API key provided") - else: - # Extract the API key from the Bearer token - if not authorization.startswith("Bearer "): - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="Invalid Authorization header format") - - api_key = authorization[len("Bearer ") :].strip() + selected: Final = select_websocket_api_key( + WebSocketKeys( + authorization=websocket.headers.get("authorization"), + api_key=websocket.headers.get("api-key"), + subprotocols=websocket.headers.get("sec-websocket-protocol"), + ) + ) + if isinstance(selected, WebSocketDenial): + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + raise HTTPException(status_code=403, detail=_WEBSOCKET_DENIAL_DETAIL[selected]) + api_key: Final = selected # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed @@ -1503,6 +1527,43 @@ def get_rbac_role(jwt_handler: JWTHandler, scopes: list[str]) -> str: return LitellmUserRoles.TEAM +def _query_api_key(route: str, request: Request) -> str | None: + from litellm.proxy.auth.route_checks import RouteChecks + from litellm.proxy.common_utils.http_parsing_utils import ( + safe_get_request_query_params, + ) + + if not RouteChecks.is_generate_content_route(route=route): + return None + return safe_get_request_query_params(request).get("key") + + +def _pass_through_header_api_key(endpoint: Mapping[str, object], route: str, request: Request) -> str | None: + if endpoint.get("path", "") != route: + return None + headers: Final = _PASS_THROUGH_HEADERS_ADAPTER.validate_python(endpoint.get("headers")) + header_name: Final = headers.get("litellm_user_api_key", "") if headers is not None else None + return request.headers.get(header_name) if isinstance(header_name, str) else None + + +def _pass_through_api_key( + pass_through_endpoints: Sequence[Mapping[str, object]] | None, route: str, request: Request +) -> str | None: + header_values: Final = tuple( + value + for value in ( + _pass_through_header_api_key(endpoint, route, request) for endpoint in pass_through_endpoints or () + ) + if value is not None + ) + return header_values[-1] if header_values else None + + +def _header_or_none(value: object) -> str | None: + """Direct callers of ``user_api_key_auth`` leave unpassed headers as their ``fastapi.Security`` defaults.""" + return value if isinstance(value, str) else None + + def get_api_key( custom_litellm_key_header: str | None, api_key: str, @@ -1510,57 +1571,29 @@ def get_api_key( anthropic_api_key_header: str | None, google_ai_studio_api_key_header: str | None, azure_apim_header: str | None, - pass_through_endpoints: list[dict] | None, + pass_through_endpoints: Sequence[Mapping[str, object]] | None, route: str, request: Request, + *, + custom_header: str | None = None, ) -> tuple[str, str | None]: - """ - Returns: - Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key - """ - from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import ( - safe_get_request_query_params, + selected: Final = select_api_key( + PresentedKeys( + custom_header=custom_header, + litellm_header=_header_or_none(custom_litellm_key_header), + authorization=_header_or_none(api_key), + api_key=_header_or_none(azure_api_key_header), + x_api_key=_header_or_none(anthropic_api_key_header), + google=_header_or_none(google_ai_studio_api_key_header), + apim=_header_or_none(azure_apim_header), + query=_query_api_key(route, request), + pass_through=_pass_through_api_key(pass_through_endpoints, route, request), + ) ) - - api_key = api_key - passed_in_key: str | None = None - if isinstance(custom_litellm_key_header, str): - passed_in_key = custom_litellm_key_header - api_key = _get_bearer_token_or_received_api_key(custom_litellm_key_header) - elif isinstance(api_key, str) and len(api_key) > 0: - passed_in_key = api_key - api_key = get_bearer_token(api_key=api_key) - elif isinstance(azure_api_key_header, str): - passed_in_key = azure_api_key_header - api_key = azure_api_key_header - elif isinstance(anthropic_api_key_header, str): - passed_in_key = anthropic_api_key_header - api_key = anthropic_api_key_header - elif isinstance(google_ai_studio_api_key_header, str): - passed_in_key = google_ai_studio_api_key_header - api_key = google_ai_studio_api_key_header - elif isinstance(azure_apim_header, str): - passed_in_key = azure_apim_header - api_key = azure_apim_header - elif ( - RouteChecks.is_generate_content_route(route=route) - and request is not None - and safe_get_request_query_params(request).get("key") - ): - google_auth_key: Final[str] = safe_get_request_query_params(request).get("key") or "" - passed_in_key = google_auth_key - api_key = google_auth_key - elif pass_through_endpoints is not None: - for endpoint in pass_through_endpoints: - if endpoint.get("path", "") == route: - headers: dict | None = endpoint.get("headers", None) - if headers is not None: - header_key: str = headers.get("litellm_user_api_key", "") - if request.headers.get(header_key) is not None: - api_key = request.headers.get(header_key) or "" - passed_in_key = api_key - return api_key, passed_in_key + if selected is None: + return api_key, None + verbose_proxy_logger.debug("Selected API key from %s", selected.source.value) + return selected.token, selected.presented async def check_api_key_for_custom_headers_or_pass_through_endpoints( @@ -1618,12 +1651,6 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( # the lookup and return None (caller proceeds to auth_builder). _JWT_PROXY_ADMIN_SENTINEL: Final = "__JWT_PROXY_ADMIN__" -_JWT_AUTH_DISABLED_HINT = ( - " This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a" - " virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate" - " with JWTs." -) - class _JWTKeyMappingFailure(Enum): NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied." @@ -2236,10 +2263,518 @@ def _custom_auth_token_for_this_request(response: object) -> UserAPIKeyAuth: return UserAPIKeyAuth.model_validate(response).model_copy() +_AUTH_CORE: Final[AuthCore] = get_auth_core() + + +class _ResolvedEndUser(NamedTuple): + end_user_id: str | None + end_user_object: LiteLLM_EndUserTable | None + budget: LiteLLM_BudgetTable | None + + @property + def params(self) -> dict[str, object]: + if not self.end_user_id: + return {} + region: Final = ( + {} if self.end_user_object is None else {"allowed_model_region": self.end_user_object.allowed_model_region} + ) + params: Final[dict[str, object]] = {"end_user_id": self.end_user_id, **region} + if self.budget is not None: + _apply_budget_limits_to_end_user_params( + end_user_params=params, budget_info=self.budget, end_user_id=self.end_user_id + ) + return params + + +async def _resolve_end_user( + request: Request, + request_data: Mapping[str, object], + token_end_user_id: str | None, + api_key: str, + key_is_resolved: bool, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + route: str, +) -> _ResolvedEndUser: + raw_end_user_id: Final = token_end_user_id or get_end_user_id_from_request_body( + request_data, safe_get_request_headers(request) + ) + end_user_id: Final = await resolve_and_validate_end_user_id( + raw_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, + ) + if prisma_client is not None: + await prefetch_identity_keys( + _identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=key_is_resolved), + user_api_key_cache=user_api_key_cache, + before_read=arm_routing_cooldowns, + ) + if not end_user_id: + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) + try: + with tracer.trace("litellm.proxy.auth.get_end_user_object"): + end_user_object: 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_object is not None: + return _ResolvedEndUser( + end_user_id=end_user_id, end_user_object=end_user_object, budget=end_user_object.litellm_budget_table + ) + if litellm.max_end_user_budget_id is None: + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) + default_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, + ) + return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=default_budget) + except litellm.BudgetExceededError: + 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) + + +def _proxy_model_list() -> list[object] | None: + from litellm.proxy.proxy_server import llm_model_list + + return llm_model_list + + +_MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value) + + +def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: PrismaClient | None) -> Settings: + return Settings( + master_key=master_key, + allow_user_auth=allow_user_auth, + has_db=prisma_client is not None, + master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, + proxy_max_budget=litellm.max_budget, + ) + + +def _aware_expiry(expires: datetime | str) -> datetime: + expiry: Final = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + return ( + expiry if expiry.tzinfo is not None and expiry.utcoffset() is not None else expiry.replace(tzinfo=timezone.utc) + ) + + +def _presented_key(key: UserAPIKeyAuth) -> PresentedKey: + if key.user_role != LitellmUserRoles.PROXY_ADMIN: + return NonAdminKey() + return AdminKey(expires=None if key.expires is None else _aware_expiry(key.expires)) + + +def _with_end_user(key: UserAPIKeyAuth, end_user: _ResolvedEndUser, parent_otel_span: Span | None) -> UserAPIKeyAuth: + with_end_user: Final = update_valid_token_with_end_user_params(valid_token=key, end_user_params=end_user.params) + with_end_user.parent_otel_span = parent_otel_span + if end_user.end_user_object is not None: + with_end_user.end_user_object_permission = end_user.end_user_object.object_permission + return with_end_user + + +def _reset_end_user_fields(key: UserAPIKeyAuth, end_user: _ResolvedEndUser) -> UserAPIKeyAuth: + key.end_user_id = end_user.end_user_id or None + key.end_user_tpm_limit = end_user.budget.tpm_limit if end_user.budget is not None else None + key.end_user_rpm_limit = end_user.budget.rpm_limit if end_user.budget is not None else None + key.end_user_tpd_limit = end_user.budget.tpd_limit if end_user.budget is not None else None + key.allowed_model_region = ( + end_user.end_user_object.allowed_model_region if end_user.end_user_object is not None else None + ) + return key + + +async def _refresh_team_fields_from_cache( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID: + return + ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token + try: + team_obj: Final[LiteLLM_TeamTableCachedObj] = 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=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ) + + if ( + team_obj.last_refreshed_at is not None + and valid_token.last_refreshed_at is not None + and team_obj.last_refreshed_at > valid_token.last_refreshed_at + ): + team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) + + for k, v in team_obj_dict.items(): + field_name = f"team_{k}" + if field_name in valid_token.__fields__: + setattr(valid_token, field_name, v) + except Exception as e: + verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. + + +class _BuilderEffects: + def __init__( + self, + *, + request: Request, + request_data: dict[str, object], # mutable-ok: model access checks rewrite the request model in place + token: str, + preresolved_key: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + route: str, + ) -> None: + self._request: Final = request + self._request_data: Final = request_data + self._token: Final = token + self._preresolved_key: Final = preresolved_key + self._prisma_client: Final = prisma_client + self._user_api_key_cache: Final = user_api_key_cache + self._parent_otel_span: Final = parent_otel_span + self._proxy_logging_obj: Final = proxy_logging_obj + self._route: Final = route + self.end_user: _ResolvedEndUser = _ResolvedEndUser(end_user_id=None, end_user_object=None, budget=None) + self.resolved_key: UserAPIKeyAuth | None = preresolved_key + self.owner: LiteLLM_UserTable | None = None + + @property + def token(self) -> str: + return self._token + + @property + def parent_otel_span(self) -> Span | None: + return self._parent_otel_span + + @property + def user_api_key_cache(self) -> UserApiKeyCache: + return self._user_api_key_cache + + @property + def proxy_logging_obj(self) -> ProxyLogging: + return self._proxy_logging_obj + + @property + def key(self) -> UserAPIKeyAuth: + if self.resolved_key is None: + raise RuntimeError("the auth core granted a key it never resolved") + return self.resolved_key + + def _identity_store(self, *, check_cache_only: bool) -> IdentityStore: + return IdentityStore( + self._prisma_client, + self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + check_cache_only=check_cache_only, + ) + + async def resolve_end_user(self, prefetch_key: bool) -> None: + self.end_user = await _resolve_end_user( + request=self._request, + request_data=self._request_data, + token_end_user_id=self._preresolved_key.end_user_id if self._preresolved_key is not None else None, + api_key=self._token, + key_is_resolved=not prefetch_key or self._preresolved_key is not None, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + route=self._route, + ) + + async def _cached_key(self, hashed_token: str) -> UserAPIKeyAuth | None: + try: + with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): + cached: Final = IdentityStore.key_from_principal( + await self._identity_store(check_cache_only=True).resolve(hashed_token=hashed_token) + ) + except Exception: + verbose_logger.debug("api key not found in cache.") + return None + # Key-cache entries are written only after the proxy validated a virtual key or the master key, but + # via_virtual_key is exclude=True so serialization drops it. Restore it at this trusted boundary. + cached.via_virtual_key = True + return cached + + def _ui_login_key(self) -> UserAPIKeyAuth | None: + # Decryption fails closed for anything that is not a genuine UI login blob. + if self._token.startswith("sk-") or get_secret_bool("EXPERIMENTAL_UI_LOGIN") is False: + return None + return ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(self._token) + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: + found: Final = self._preresolved_key or await self._cached_key(hashed_token) or self._ui_login_key() + if found is None: + return None + self.resolved_key = ( + await _refresh_session_token_grants( + valid_token=found, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + if found.is_session_token and self._prisma_client is not None + else found + ) + return _presented_key(self.resolved_key) + + async def evict_presented_key(self, hashed_token: str) -> None: + await delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=self._user_api_key_cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def load_key(self, hashed_token: str) -> None: + with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): + self.resolved_key = IdentityStore.key_from_principal( + await self._identity_store(check_cache_only=False).resolve(hashed_token=hashed_token) + ) + + async def prepare_key(self, loaded_from_db: bool) -> None: + if loaded_from_db: + _reset_end_user_fields(self.key, self.end_user) + return + await _refresh_team_fields_from_cache( + valid_token=self.key, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def inspect_key(self) -> KeyFacts: + key: Final = _update_key_budget_with_temp_budget_increase(self.key) + self.resolved_key = key + return KeyFacts( + has_token=key.token is not None, + blocked=key.blocked is True, + expires=None if key.expires is None else _aware_expiry(key.expires), + spend=key.spend, + soft_budget=key.soft_budget, + max_budget=key.max_budget, + team_member_spend=key.team_member_spend, + user_id=key.user_id, + team_id=key.team_id, + budget_throttle_pct=budget_throttle_percentage() if should_throttle_budget_exceeded(key) else None, + ) + + async def enforce_model_access(self) -> None: + from litellm.proxy.proxy_server import llm_router + + await enforce_key_and_fallback_model_access( + valid_token=self.key, + request_data=self._request_data, + route=self._route, + request=self._request, + llm_model_list=_proxy_model_list(), + llm_router=llm_router, + ) + await _prefetch_referenced_auth_objects( + self.key, + end_user_id=self.end_user.end_user_id, + user_api_key_cache=self._user_api_key_cache, + prisma_client=self._prisma_client, + ) + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: + try: + with tracer.trace("litellm.proxy.auth.get_user_object"): + owner: Final = await get_user_object( + user_id=user_id, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + user_id_upsert=False, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except Exception as e: + verbose_logger.debug("Unable to get user %s from db/cache, continuing without it: %s", user_id, e) + return None + self.owner = owner + if owner is None: + return None + key: Final = self.key + # The joint verification-token view carries the key's columns only, so the user's own per-model budget + # reaches enforcement and the post-call increment through the row fetched here. + 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: + from litellm.proxy.proxy_server import llm_router + + model: Final = _get_model_from_request_context( + request_data=self._request_data, + route=self._route, + request=self._request, + llm_router=llm_router, + team_id=self.key.team_id, + ) + return request_skips_budget_checks( + 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.proxy.proxy_server import get_current_spend + + if self._prisma_client is None: + return None + cache_key: Final = team_membership_auth_cache_key(team_id=team_id, user_id=user_id) + cached: Final = await self._user_api_key_cache.async_get_cache(key=cache_key, model_type=LiteLLM_TeamMembership) + 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: + return 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, + max_budget=max_budget, + ) + return TeamMemberBudget(spend=spend, max_budget=max_budget) + + async def _load_team_membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None: + if self._prisma_client is None: + return None + row: Final = await TeamMembershipRepository(self._prisma_client).table.find_first( + where={"user_id": user_id, "team_id": team_id}, + include={"litellm_budget_table": True}, + ) + if row is None: + return None + membership: Final = LiteLLM_TeamMembership.model_validate(row.model_dump()) + await self._user_api_key_cache.async_set_cache( + key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + value=membership, + model_type=LiteLLM_TeamMembership, + ttl=5, + ) + return membership + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: + # common_checks sends this alert on requests that get past here, so only the rejected request sends it here. + team_member_max_budget_alert_check( + team_id=team_id, + team_alias=self.key.team_alias, + team_metadata=self.key.team_metadata, + organization_id=self.key.org_id, + user_id=user_id, + user_email=self.owner.user_email if self.owner is not None else None, + proxy_logging_obj=self._proxy_logging_obj, + spend=budget.spend, + max_budget=budget.max_budget, + ) + + async def send_key_budget_alerts(self) -> None: + await virtual_key_max_budget_alert_check( + valid_token=self.key, proxy_logging_obj=self._proxy_logging_obj, user_obj=self.owner + ) + + async def read_key_spend(self, max_budget: float) -> float: + from litellm.proxy.proxy_server import get_current_spend + + # Redis-first cross-pod counter, falling back to the key object's recorded spend. + return await get_current_spend( + counter_key=f"spend:key:{self.key.token}", + fallback_spend=self.key.spend or 0.0, + max_budget=max_budget, + ) + + async def alert_key_budget(self, spend: float) -> None: + send_token_budget_alert(self.key, spend, self._proxy_logging_obj, self.owner) + + async def alert_key_soft_budget(self) -> None: + send_soft_budget_alert(self.key, self._proxy_logging_obj, self.owner) + + def _model_budgets(self) -> _ModelBudgetEffects: + from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter + + return _ModelBudgetEffects( + request=self._request, + request_data=self._request_data, + valid_token=self.key, + route=self._route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), + ) + + 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: + return await self._model_budgets().check_key_model_budget(model) + + async def authorize_budget_fallback(self, model: str) -> bool: + 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: + return await self._model_budgets().check_user_model_budget(model) + + async def load_key_context(self) -> None: + await _load_team_and_project( + valid_token=self.key, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def read_proxy_spend(self) -> float | None: + if self._prisma_client is None: + return None + with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): + return await fetch_global_spend_with_event_coordination( + cache_key=GLOBAL_PROXY_SPEND_CACHE_KEY, + user_api_key_cache=self._user_api_key_cache, + prisma_client=self._prisma_client, + ) + + async def alert_proxy_budget(self, spend: float) -> None: + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + call_info: Final = CallInfo( + token=self.key.token, + spend=spend, + max_budget=litellm.max_budget, + user_id=litellm_proxy_admin_name, + team_id=self.key.team_id, + event_group=Litellm_EntityType.PROXY, + ) + asyncio.create_task(self._proxy_logging_obj.budget_alerts(type="proxy_budget", user_info=call_info)) + + async def user_api_key_auth_builder( request: Request, api_key: str, - azure_api_key_header: str, + azure_api_key_header: str | None, anthropic_api_key_header: str | None, google_ai_studio_api_key_header: str | None, azure_apim_header: str | None, @@ -2281,8 +2816,13 @@ async def user_api_key_auth_builder( route=route, ) pass_through_endpoints: Final[list[dict] | None] = general_settings.get("pass_through_endpoints", None) - ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header - api_key, passed_in_key = get_api_key( + custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name") + api_key, _ = get_api_key( + custom_header=( + request.headers.get(custom_litellm_key_header_name) + if custom_litellm_key_header_name is not None + else None + ), custom_litellm_key_header=custom_litellm_key_header, api_key=api_key, azure_api_key_header=azure_api_key_header, @@ -2293,12 +2833,6 @@ async def user_api_key_auth_builder( route=route, request=request, ) - custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name") - if custom_litellm_key_header_name is not None and custom_litellm_key_header_name in request.headers: - api_key = get_api_key_from_custom_header( - request=request, - custom_litellm_key_header_name=custom_litellm_key_header_name, - ) if open_telemetry_logger is not None: # Reuse the span created by user_api_key_auth (before body parse) @@ -2613,17 +3147,16 @@ async def user_api_key_auth_builder( # zero-cost model would be refused here and allowed there, # while the log above claims all budget checks were skipped. if not skip_budget_checks: - await _check_user_model_budget( - valid_token=cast(UserAPIKeyAuth, valid_token), - model_max_budget_limiter=model_max_budget_limiter, - models=_get_model_names_for_budget_checks( - model=_get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) + await _enforce_model_budgets( + _AUTH_CORE.check_user_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=cast(UserAPIKeyAuth, valid_token), + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), ), ) @@ -2642,343 +3175,46 @@ async def user_api_key_auth_builder( api_key = response elif isinstance(response, UserAPIKeyAuth): return response - if master_key is None: - if isinstance(api_key, str): - return UserAPIKeyAuth( - api_key=api_key, - user_role=LitellmUserRoles.INTERNAL_USER, - parent_otel_span=parent_otel_span, - ) - else: - return UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - parent_otel_span=parent_otel_span, - ) - elif api_key is None: # only require api key if master key is set - raise ProxyException( - message="Authentication Error, No api key passed in.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, + if master_key is not None and not isinstance(master_key, str): + raise HTTPException( + status_code=500, + detail={f"Master key must be a valid string. Current type={type(master_key)}"}, ) - elif api_key == "": - # missing 'Bearer ' prefix - raise ProxyException( - message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - if route == "/user/auth": - if general_settings.get("allow_user_auth", False) is True: - return UserAPIKeyAuth() - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="'allow_user_auth' not set or set to False", - ) - - ## Check END-USER OBJECT - _end_user_object = None - end_user_params: Final = {} - - raw_end_user_id: Final = ( - valid_token.end_user_id if valid_token is not None else None - ) or get_end_user_id_from_request_body(request_data, safe_get_request_headers(request)) - end_user_id = await resolve_and_validate_end_user_id( - raw_end_user_id=raw_end_user_id, + effects: Final = _BuilderEffects( + request=request, + request_data=request_data, + token=api_key, + preresolved_key=valid_token, 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 prisma_client is not None: - await prefetch_identity_keys( - _identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=valid_token is not None), - user_api_key_cache=user_api_key_cache, - before_read=arm_routing_cooldowns, - ) - if end_user_id: - try: - end_user_params["end_user_id"] = end_user_id - - with tracer.trace("litellm.proxy.auth.get_end_user_object"): - _end_user_object = 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_object is not None: - end_user_params["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=end_user_id, - ) - elif litellm.max_end_user_budget_id is not None: - # End user doesn't exist yet, but apply default budget limits if configured - from litellm.proxy.auth.auth_checks import ( - get_default_end_user_budget, - ) - - default_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, - ) - if default_budget is not None: - _apply_budget_limits_to_end_user_params( - end_user_params=end_user_params, - budget_info=default_budget, - end_user_id=end_user_id, - ) - except Exception as e: - if isinstance(e, litellm.BudgetExceededError): - raise e - verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) - - ### CHECK IF ADMIN ### - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead - ### CHECK IF ADMIN ### - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead - if valid_token is None: - ## Check CACHE - try: - with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ).resolve(hashed_token=hash_token(api_key)) - ) - # Key-cache entries are written only after the proxy validated a - # virtual key or the master key, but via_virtual_key is exclude=True - # so serialization drops it; restore it at this trusted boundary. - # The UI-login JWT fallback below constructs its token from a - # decrypted blob, not this cache, and stays unmarked. - if isinstance(valid_token, UserAPIKeyAuth): - valid_token.via_virtual_key = True - except Exception: - verbose_logger.debug("api key not found in cache.") - valid_token = None - - ## Check UI/CLI Hash Key - # Attempt decryption for non-sk- tokens unless the operator has - # explicitly set EXPERIMENTAL_UI_LOGIN=false to disable it. - # Unset (None) keeps the new default of always attempting decryption; - # decryption fails closed for anything that is not a genuine blob. - if ( - valid_token is None - and not api_key.startswith("sk-") - and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False - ): - valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) - - if valid_token is not None and valid_token.is_session_token and prisma_client is not None: - valid_token = await _refresh_session_token_grants( # rebind-ok: later checks read this name - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN - ): - if valid_token.expires is not None: - current_time = 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: - await delete_cache_key_object( - hashed_token=hash_token(api_key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - raise ProxyException( - message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", - type=ProxyErrorTypes.expired_key, - code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=api_key), - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params - ) - valid_token.parent_otel_span = parent_otel_span - if _end_user_object is not None: - valid_token.end_user_object_permission = _end_user_object.object_permission - - return valid_token - - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.team_id is not None - and valid_token.team_id != UI_TEAM_ID - ): - ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token - try: - team_obj: Final[LiteLLM_TeamTableCachedObj] = 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=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ) - - if ( - team_obj.last_refreshed_at is not None - and valid_token.last_refreshed_at is not None - and team_obj.last_refreshed_at > valid_token.last_refreshed_at - ): - team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) - - for k, v in team_obj_dict.items(): - field_name = f"team_{k}" - if field_name in valid_token.__fields__: - setattr(valid_token, field_name, v) - except Exception as e: - verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. - try: - is_master_key_valid = secrets.compare_digest(api_key, master_key) - except Exception: - is_master_key_valid = False - - ## VALIDATE MASTER KEY ## - if not isinstance(master_key, str): - raise HTTPException( - status_code=500, - detail={f"Master key must be a valid string. Current type={type(master_key)}"}, + outcome: Final = await drive( + _AUTH_CORE.start( + RequestFacts( + route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route) + ), + _core_settings( + master_key=master_key, + allow_user_auth=general_settings.get("allow_user_auth", False) is True, + prisma_client=prisma_client, + ), + datetime.now(timezone.utc), + ), + effects, ) - - if is_master_key_valid: - # Substitute a stable alias for the raw master key so neither the - # master key nor its hash propagates into spend logs, Prometheus - # /metrics labels, audit trails, rate-limit buckets, or any other - # downstream consumer of UserAPIKeyAuth.api_key. - _user_api_key_obj = await _return_user_api_key_auth_obj( - user_obj=None, - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, - parent_otel_span=parent_otel_span, - valid_token_dict={ - **end_user_params, - "user_id": litellm_proxy_admin_name, - }, - route=route, - start_time=start_time, - ) - asyncio.create_task( - cache_key_object( - hashed_token=hash_token(master_key), - user_api_key_obj=_user_api_key_obj, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - - _user_api_key_obj = update_valid_token_with_end_user_params( - valid_token=_user_api_key_obj, end_user_params=end_user_params - ) - _user_api_key_obj.via_virtual_key = True - - return _user_api_key_obj - - ## IF it's not a master key - ## Route should not be in master_key_only_routes - if route in LiteLLMRoutes.master_key_only_routes.value: - raise ProxyException( - message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - ## Check DB - - if ( - prisma_client is None - ): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error - raise ProxyException( - message="No connected db.", - type=ProxyErrorTypes.no_db_connection, - code=400, - param=None, - ) - - if valid_token is None: - if isinstance(api_key, str): # if generated token, make sure it starts with sk-. - _masked_key: Final = f"{api_key[:4]}****{api_key[-4:]}" if len(api_key) > 8 else "****" - if not api_key.startswith("sk-"): - _hint = _JWT_AUTH_DISABLED_HINT if not enable_jwt_auth and JWTHandler.is_jwt(token=api_key) else "" - _malformed_key_error = HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=( - f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={_masked_key}, " - f"expected to start with 'sk-'.{_hint}" - ), - ) # prevent token hashes from being used - # Stamp provenance here so log routing classifies this 401 by - # where it was raised, never by its message text. - setattr(_malformed_key_error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) - raise _malformed_key_error - else: - verbose_logger.warning( - "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( - type(api_key) if api_key is not None else "None" - ) - ) - if api_key.startswith("sk-"): - api_key = hash_token(token=api_key) - - with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ).resolve(hashed_token=api_key) - ) - # update end-user params on valid token - # These can change per request - it's important to update them here - valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") - valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit") - valid_token.allowed_model_region = end_user_params.get("allowed_model_region") - - return await validate_resolved_virtual_key( - request=request, - request_data=request_data, - valid_token=valid_token, - api_key=api_key, + finally: + valid_token = effects.resolved_key # rebind-ok: the exception handler below reports the resolved identity + return await _authenticated( + outcome, + effects, + admin_user_id=litellm_proxy_admin_name, route=route, start_time=start_time, - parent_otel_span=parent_otel_span, - end_user_id=end_user_id, - end_user_params=cast( # cast-ok: builder assembles this dict from validated end-user fields - dict[str, object], end_user_params - ), - _end_user_object=_end_user_object, + jwt_auth_enabled=enable_jwt_auth, ) except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( @@ -2995,417 +3231,194 @@ async def user_api_key_auth_builder( _user_api_key_auth_builder: Final = user_api_key_auth_builder -async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of existing shared authorization checks - request: Request, - request_data: dict[str, object], - valid_token: UserAPIKeyAuth | None, - api_key: str, - route: str, - start_time: datetime, +async def _load_team_and_project( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, - end_user_id: str | None, - end_user_params: dict[str, object], - _end_user_object: LiteLLM_EndUserTable | None, -) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - llm_model_list, - llm_router, - model_max_budget_limiter, - prisma_client, - proxy_logging_obj, - user_api_key_cache, + proxy_logging_obj: ProxyLogging, +) -> None: + team_obj: Final = ( + None + if valid_token.team_id is None + else await _key_team_object( + valid_token, valid_token.team_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj + ) + ) + valid_token.team_object_permission = team_obj.object_permission if team_obj is not None else None + if team_obj is not None: + # Guardrails or other metadata added after the key was cached apply without a cache eviction. + valid_token.team_metadata = team_obj.metadata + if valid_token.project_id is None: + return + project_obj: 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_obj is not None: + valid_token.project_metadata = project_obj.metadata + valid_token.project_alias = project_obj.project_alias + + +async def _key_team_object( + valid_token: UserAPIKeyAuth, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj: + fetched: Final = ( + None + if team_id == UI_TEAM_ID + else await _team_object_or_none(team_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj) + ) + if fetched is not None: + return fetched + return LiteLLM_TeamTableCachedObj( + team_id=team_id, + max_budget=valid_token.team_max_budget, + soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, + spend=valid_token.team_spend, + tpm_limit=valid_token.team_tpm_limit, + rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, + blocked=valid_token.team_blocked, + models=_token_team_models(valid_token), + metadata=valid_token.team_metadata, + object_permission_id=valid_token.team_object_permission_id, + object_permission=await _resolve_object_permission_for_unresolvable_team( + object_permission_id=valid_token.team_object_permission_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, + ), ) - if valid_token is not None: - valid_token = _update_key_budget_with_temp_budget_increase(valid_token) - user_obj: LiteLLM_UserTable | None = None - valid_token_dict: dict = {} - if valid_token is not None: - # Got Valid Token from Cache, DB - # Run checks for - # 1. If token can call model - ## 1a. If token can call fallback models (if client-side fallbacks given) - # 2. If user_id for this token is in budget - # 3. If the user spend within their own team is within budget - # 4. If 'user' passed to /chat/completions, /embeddings endpoint is in budget - # 5. If token is expired - # 6. If token spend is under Budget for the token - # 7. If token spend per model is under budget per model - # 8. If token spend is under team budget - # 9. If team spend is under team budget - - ## base case ## key is disabled - if valid_token.blocked is True: - raise ProxyException( - message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) - await _prefetch_referenced_auth_objects( - valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client - ) - - # Check 2. If user_id for this token is in budget - done in common_checks() - if valid_token.user_id is not None: - try: - with tracer.trace("litellm.proxy.auth.get_user_object"): - user_obj = 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=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - verbose_logger.debug( - "litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - %s", - e, - ) - user_obj = None - - if user_obj is not None: - # The joint verification-token view carries the key's columns only, so the - # user's own per-model budget reaches enforcement and the post-call - # increment through the row fetched here. - valid_token.user_model_max_budget = user_obj.model_max_budget - - if ( - user_obj is not None - and isinstance(user_obj.metadata, dict) - and user_obj.metadata.get("scim_active") is False - ): - raise ProxyException( - message=( - f"Authentication Error, User={valid_token.user_id} has been deactivated via SCIM. " - "Keys owned by this user cannot be used." - ), - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - # Check 2a. Check if model has zero cost - if so, skip all budget checks - model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - skip_budget_checks: Final = request_skips_budget_checks( - route=route, model=model, llm_router=llm_router, request_body=request_data - ) - - # Check 3. Check if user is in their team budget - if not skip_budget_checks and valid_token.team_member_spend is not None: - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - if prisma_client is not None and _user_id is not None and _team_id is not None: - _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) - - team_member_info = await user_api_key_cache.async_get_cache( - key=_cache_key, - model_type=LiteLLM_TeamMembership, - ) - if team_member_info is None: - # read from DB - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, - ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership.model_validate(_db_member.model_dump()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) - - if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( - now=datetime.now(timezone.utc), - ) - if team_member_budget is not None and team_member_budget > 0: - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = valid_token.team_member_spend - if valid_token.user_id is not None and valid_token.team_id is not None: - team_member_spend = await get_current_spend( - counter_key=team_member_counter(valid_token.user_id, valid_token.team_id), - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - if team_member_spend >= team_member_budget: - # common_checks sends this alert on requests that get past here, so only the - # request rejected here sends it from the builder. - team_member_max_budget_alert_check( - team_id=_team_id, - team_alias=valid_token.team_alias, - team_metadata=valid_token.team_metadata, - organization_id=valid_token.org_id, - user_id=_user_id, - user_email=user_obj.user_email if user_obj is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) - _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - message=( - f"Budget has been exceeded! TeamMember={_entity_id} " - f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" - ), - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=_entity_id, - ) - - # Check 3. If token is expired - if valid_token.expires is not None: - current_time = 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) - verbose_proxy_logger.debug( - "Checking if token expired, expiry time %s and current time %s", expiry_time, current_time - ) - if expiry_time < current_time: - # Token exists but is expired. - raise ProxyException( - message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", - type=ProxyErrorTypes.expired_key, - code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=api_key), - ) - - if not skip_budget_checks: - with tracer.trace("litellm.proxy.auth.budget_checks"): - # Check 4. Max Budget Alert Check (runs before budget enforcement - # so multi-threshold 100% alerts fire on the request that crosses - # max_budget, before BudgetExceededError is raised below) - await virtual_key_max_budget_alert_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Spend is under budget - if RouteChecks.is_llm_api_route(route=route): - await virtual_key_max_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 6. Soft Budget Check - await virtual_key_soft_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Model Spend is under Model budget - max_budget_per_model: Final = valid_token.model_max_budget - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - if ( - max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - and prisma_client is not None - and current_models - and valid_token.token is not None - ): - ## GET THE SPEND FOR THIS MODEL - for model_name in current_models: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, llm_model_list - ), - llm_router=llm_router, - ) - - # Recompute after a potential budget-fallback rewrite so - # the user check below validates the final model - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - # Check 5a. Internal user model_max_budget - if current_models: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=current_models, - ) - - # Check 6: Additional Common Checks across jwt + key auth - if valid_token.team_id is not None: - try: - if valid_token.team_id == UI_TEAM_ID: - raise TeamNotFoundError(team_id=UI_TEAM_ID) - with tracer.trace("litellm.proxy.auth.get_team_object"): - _team_obj = 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=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except HTTPException: - token_team_models: Final = _token_team_models(valid_token) - _team_obj = LiteLLM_TeamTableCachedObj( - team_id=valid_token.team_id, - max_budget=valid_token.team_max_budget, - soft_budget=valid_token.team_soft_budget, - model_max_budget=valid_token.team_model_max_budget, - spend=valid_token.team_spend, - tpm_limit=valid_token.team_tpm_limit, - rpm_limit=valid_token.team_rpm_limit, - tpd_limit=valid_token.team_tpd_limit, - blocked=valid_token.team_blocked, - models=token_team_models, - metadata=valid_token.team_metadata, - object_permission_id=valid_token.team_object_permission_id, - object_permission=await _resolve_object_permission_for_unresolvable_team( - object_permission_id=valid_token.team_object_permission_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, - ), - ) - else: - _team_obj = None - - if _team_obj is not None: - valid_token.team_object_permission = _team_obj.object_permission - # Keep team_metadata in sync with the freshly fetched team so that - # guardrails (or any other metadata) added after the key was cached - # are picked up on subsequent requests without a cache eviction. - valid_token.team_metadata = _team_obj.metadata - else: - valid_token.team_object_permission = None - - # Fetch project object if key belongs to a project - _project_obj = None - if valid_token.project_id is not None: - _project_obj = await get_project_object( - project_id=valid_token.project_id, +async def _team_object_or_none( + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj | None: + try: + with tracer.trace("litellm.proxy.auth.get_team_object"): + return await get_team_object( + team_id=team_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, ) - if _project_obj is not None: - valid_token.project_metadata = _project_obj.metadata - valid_token.project_alias = _project_obj.project_alias + except HTTPException: + return None - global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget - cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY - with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): - global_proxy_spend = await fetch_global_spend_with_event_coordination( # rebind-ok: pre-existing rebinding on a rename-only line - cache_key=cache_key, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) - if global_proxy_spend is not None: - call_info: Final = CallInfo( - token=valid_token.token, - spend=global_proxy_spend, - max_budget=litellm.max_budget, - user_id=litellm_proxy_admin_name, - team_id=valid_token.team_id, - event_group=Litellm_EntityType.PROXY, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="proxy_budget", - user_info=call_info, - ) - ) - # Token passed all checks - if valid_token is None: - raise HTTPException(401, detail="Invalid API key") - if valid_token.token is None: - raise HTTPException(401, detail="Invalid API key, no token associated") - api_key = valid_token.token - - valid_token_dict = valid_token.model_dump(exclude_none=True) - valid_token_dict.pop("token", None) - # budget_throttle_pct is excluded from model_dump (it must not leak - # into serialized responses), so carry the request-scoped decision - # forward by hand to the auth object the rate limiter receives. - if valid_token.budget_throttle_pct is not None: - valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct - - if _end_user_object is not None or litellm.max_end_user_budget_id is not None: - valid_token_dict.update(end_user_params) - if _end_user_object is not None: - valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission - - # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions - # sso/login, ui/login, /key functions and /user functions - # this will never be allowed to call /chat/completions - - if valid_token is None: - # No token was found when looking up in the DB - raise ProxyException( - message="Authentication Error, Invalid proxy server token passed", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, +async def _master_key_auth( + effects: _BuilderEffects, admin_user_id: str, route: str, start_time: datetime +) -> UserAPIKeyAuth: + # A stable alias replaces the raw master key so neither the key nor its hash reaches spend logs, + # metrics labels, audit trails or rate-limit buckets. + master_key_auth: Final = await _return_user_api_key_auth_obj( + user_obj=None, + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + parent_otel_span=effects.parent_otel_span, + valid_token_dict={**effects.end_user.params, "user_id": admin_user_id}, + route=route, + start_time=start_time, + ) + asyncio.create_task( + cache_key_object( + hashed_token=hash_token(effects.token), + user_api_key_obj=master_key_auth, + user_api_key_cache=effects.user_api_key_cache, + proxy_logging_obj=effects.proxy_logging_obj, ) - if valid_token_dict is not None: - virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( - user_obj=user_obj, - api_key=api_key, - parent_otel_span=parent_otel_span, - valid_token_dict=valid_token_dict, - route=route, - start_time=start_time, - ) - virtual_key_auth_obj.via_virtual_key = True - return virtual_key_auth_obj + ) + with_end_user: Final = _with_end_user(master_key_auth, effects.end_user, effects.parent_otel_span) + with_end_user.via_virtual_key = True + return with_end_user + + +async def _authenticated( + outcome: Outcome, + effects: _BuilderEffects, + admin_user_id: str, + route: str, + start_time: datetime, + jwt_auth_enabled: bool, +) -> UserAPIKeyAuth: + match outcome: + case Deny(denial=denial): + raise denial_exception( + denial, token=effects.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=effects.resolved_key + ) + case Allow(grant=grant): + return await _granted(grant, effects, admin_user_id, route, start_time) + + +async def _granted( + grant: Grant, effects: _BuilderEffects, admin_user_id: str, route: str, start_time: datetime +) -> UserAPIKeyAuth: + match grant: + case NoMasterKey(token=token): + return UserAPIKeyAuth( + api_key=token, user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=effects.parent_otel_span + ) + case UserAuthRoute(): + return UserAPIKeyAuth() + case MasterKey(): + return await _master_key_auth(effects, admin_user_id, route, start_time) + case CachedAdminKey(): + return _with_end_user(effects.key, effects.end_user, effects.parent_otel_span) + case VirtualKey(budget_throttle_pct=budget_throttle_pct): + return await _virtual_key_auth(effects, budget_throttle_pct, route, start_time) + case WithinModelBudgets(): + raise RuntimeError("the request auth core never ends on a model budget check") + case _: + assert_never(grant) + + +async def _virtual_key_auth( + effects: _BuilderEffects, budget_throttle_pct: float | None, route: str, start_time: datetime +) -> UserAPIKeyAuth: + key: Final = effects.key + token: Final = key.token + if token is None: + raise RuntimeError("the auth core granted a virtual key without a token") + end_user: Final = effects.end_user + # budget_throttle_pct is excluded from model_dump so it never leaks into serialized responses. + throttle: Final = {} if budget_throttle_pct is None else {"budget_throttle_pct": budget_throttle_pct} + end_user_params: Final = ( + end_user.params if end_user.end_user_object is not None or litellm.max_end_user_budget_id is not None else {} + ) + end_user_permission: Final = ( + {} + if end_user.end_user_object is None + else {"end_user_object_permission": end_user.end_user_object.object_permission} + ) + virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( + user_obj=effects.owner, + api_key=token, + parent_otel_span=effects.parent_otel_span, + valid_token_dict={ + **{k: v for k, v in key.model_dump(exclude_none=True).items() if k != "token"}, + **throttle, + **end_user_params, + **end_user_permission, + }, + route=route, + start_time=start_time, + ) + virtual_key_auth_obj.via_virtual_key = True + return virtual_key_auth_obj async def _safe_fetch(label: str, awaitable): @@ -4407,7 +4420,7 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque async def user_api_key_auth( request: Request, api_key: str = fastapi.Security(api_key_header), - azure_api_key_header: str = fastapi.Security(azure_api_key_header), + azure_api_key_header: str | None = fastapi.Security(azure_api_key_header), anthropic_api_key_header: str | None = fastapi.Security(anthropic_api_key_header), google_ai_studio_api_key_header: str | None = fastapi.Security(google_ai_studio_api_key_header), azure_apim_header: str | None = fastapi.Security(azure_apim_header), @@ -4557,41 +4570,6 @@ async def _return_user_api_key_auth_obj( return UserAPIKeyAuth.model_validate(user_api_key_kwargs) -def get_api_key_from_custom_header(request: Request, custom_litellm_key_header_name: str) -> str: - """ - Get API key from custom header - - Args: - request (Request): Request object - custom_litellm_key_header_name (str): Custom header name - - Returns: - Optional[str]: API key - """ - api_key: str = "" - # use this as the virtual key passed to litellm proxy - custom_litellm_key_header_name = custom_litellm_key_header_name.lower() - _headers: Final = {k.lower(): v for k, v in request.headers.items()} - verbose_proxy_logger.debug( - "searching for custom_litellm_key_header_name= %s, in headers=%s", - custom_litellm_key_header_name, - _headers, - ) - custom_api_key: Final = _headers.get(custom_litellm_key_header_name) - if custom_api_key: - api_key = get_bearer_token(api_key=custom_api_key) # rebind-ok: pre-existing rebinding on a rename-only line - verbose_proxy_logger.debug( - "Found custom API key using header: %s, setting api_key=%s", - custom_litellm_key_header_name, - abbreviate_api_key(api_key), - ) - else: - verbose_proxy_logger.exception( - "No LiteLLM Virtual Key pass. Please set header=%s: Bearer ", custom_litellm_key_header_name - ) - return api_key - - 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: @@ -4802,14 +4780,13 @@ async def _run_post_custom_auth_checks( llm_router=llm_router, ) - current_model = _get_model_from_request_context( + current_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, ) - current_models = _get_model_names_for_budget_checks(model=current_model) # A zero-cost model cannot move any counter, so refusing it means refusing on # spend some other model accrued. The JWT and virtual-key paths already skip @@ -4821,41 +4798,6 @@ async def _run_post_custom_auth_checks( else False ) - # 3. Check key-level model_max_budget - max_budget_per_model: Final = valid_token.model_max_budget - if ( - not skip_budget_checks - and max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - and current_models - and valid_token.token is not None - ): - for model_name in current_models: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, llm_model_list - ), - llm_router=llm_router, - ) - - # Recompute after a potential budget-fallback rewrite so - # the end-user check below validates the final model - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - # 3b. Attach and check the internal user's model_max_budget. # Custom auth builds its own token, so unlike the main path nothing has # loaded the user row yet. The attach is unconditional because the post-call # spend hook reads this field off the token: gating it on the same condition @@ -4869,14 +4811,33 @@ 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 and current_models: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=current_models, + if not skip_budget_checks: + await _enforce_model_budgets( + _AUTH_CORE.check_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=valid_token, + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=cast( # cast-ok: router model lists are deployment sequences + Sequence[object] | None, llm_model_list + ), + ), ) - # 4. Check end-user model_max_budget + # Read after the budget checks, which may have rerouted the request to a fallback model. + current_models: Final = _get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + ) + end_user_mmb: Final = valid_token.end_user_model_max_budget if ( not skip_budget_checks @@ -4912,6 +4873,51 @@ async def _run_post_custom_auth_checks( return valid_token +async def _authorize_resolved_key( + identity: UserAPIKeyAuth, + key_hash: str, + request: Request, + request_data: dict[str, object], # mutable-ok: model access checks rewrite the request model in place + route: str, +) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + master_key, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + start_time: Final = datetime.now(timezone.utc) + effects: Final = _BuilderEffects( + request=request, + request_data=request_data, + token=key_hash, + preresolved_key=identity, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + outcome: Final = await drive( + _AUTH_CORE.start_resolved_key( + RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)), + _core_settings(master_key=master_key, allow_user_auth=False, prisma_client=prisma_client), + start_time, + ), + effects, + ) + return await _authenticated( + outcome, + effects, + admin_user_id=litellm_proxy_admin_name, + route=route, + start_time=start_time, + jwt_auth_enabled=False, + ) + + async def authorize_internal_virtual_key( key_hash: str, request: Request, request_data: dict[str, object] ) -> UserAPIKeyAuth: @@ -4925,17 +4931,8 @@ async def authorize_internal_virtual_key( ) route: Final = get_request_route(request=request) await pre_db_read_auth_checks(request_data=request_data, request=request, route=route) - auth: Final = await validate_resolved_virtual_key( - request=request, - request_data=request_data, - valid_token=identity, - api_key=key_hash, - route=route, - start_time=datetime.now(timezone.utc), - parent_otel_span=None, - end_user_id=None, - end_user_params={}, - _end_user_object=None, + auth: Final = await _authorize_resolved_key( + identity=identity, key_hash=key_hash, request=request, request_data=request_data, route=route ) auth.budget_reservation = None recovered: Final = await _authorize_authenticated_request( diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 4daef2592f7..0b389e3dfca 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -923,21 +923,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) - @with_service_target("model_budgets") - async def get_fallback_model_within_budget( - self, - user_api_key_dict: UserAPIKeyAuth, - model: str, - ) -> str | None: - budget_fallbacks: Final[dict[str, list[str]]] = user_api_key_dict.budget_fallbacks or {} - for fallback_model in budget_fallbacks.get(model, []): - try: - await self.is_key_within_model_budget(user_api_key_dict=user_api_key_dict, model=fallback_model) - return fallback_model - except litellm.BudgetExceededError: - continue - return None - @with_service_target("model_budgets") async def is_user_within_model_budget( self, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d54eda45160..e368a282240 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6586,9 +6586,7 @@ async def regenerate_key_fn( # value in that field. regenerate_target_key: Final = data.key if data and data.key else key is_master_key_regeneration: Final = ( - data is not None - and data.new_master_key is not None - and is_master_key(api_key=regenerate_target_key, _master_key=master_key) + data is not None and data.new_master_key is not None and is_master_key(regenerate_target_key, master_key) ) if ( @@ -6616,7 +6614,7 @@ async def regenerate_key_fn( detail={"error": "DB not connected. prisma_client is None"}, ) - _is_master_key_valid: Final = is_master_key(api_key=key, _master_key=master_key) + _is_master_key_valid: Final = is_master_key(key, master_key) if master_key is not None and data and _is_master_key_valid: if data.new_master_key is None: diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index be24b1630b0..948d2613008 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -19,6 +19,7 @@ from starlette.types import Scope from typing_extensions import assert_never import litellm +from litellm._auth.runtime import litellm_header_token from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.constants import MCP_ALL_TOOLS_WILDCARD, MCP_PEEKED_BODY_SCOPE_KEY @@ -48,7 +49,6 @@ from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_age from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports - _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export oauth2_introspection_applies, run_centralized_common_checks, @@ -110,7 +110,7 @@ _CACHED_VALUE: Final = TypeAdapter(object) def _normalize_caller_admission_credential(value: str) -> str: - return _get_bearer_token_or_received_api_key(strip_auth_scheme(value, "Bearer")).strip() + return litellm_header_token(strip_auth_scheme(value, "Bearer")).strip() def _admission_request(scope: Scope) -> Request: @@ -797,7 +797,7 @@ class MCPRequestHandler: # for a delegated server, so validate it: identity / spend / rate # limits resolve and any stored upstream token can be forwarded. validated_user_api_key_auth = await user_api_key_auth( - api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}", + api_key=f"Bearer {litellm_header_token(litellm_api_key)}", request=request, ) elif MCPRequestHandler._target_servers_are_true_passthrough( @@ -1278,7 +1278,7 @@ class MCPRequestHandler: permissions, budgets, rate limits) and the sealed upstream token is injected under the server's per-server auth-header key, while the leak-defense chokepoint strips the envelope ``Authorization`` itself from egress.""" - presented_token: Final = _get_bearer_token_or_received_api_key(litellm_api_key) + presented_token: Final = litellm_header_token(litellm_api_key) explicit_auth: Final = await user_api_key_auth(api_key=f"Bearer {presented_token}", request=request) result: Final = await MCPRequestHandler._open_dcr_bridge_envelope( server=server, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ba8153ebb6d..0b12a940f73 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -31,6 +31,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider +from litellm._auth.runtime import scheme_credential from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( @@ -99,7 +100,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul _get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export api_key_header, custom_litellm_key_header, - get_bearer_token, + get_bearer_token, # pyright: ignore[reportUnusedImport] # backwards-compatible package export is_no_auth_dev_mode, user_api_key_auth, user_api_key_auth_accepted_websocket_for_model, @@ -2858,17 +2859,7 @@ _CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL: Final = ( def _normalize_credential_value(value: str) -> str: - """Reduce a header value to the bare token, matching how ``user_api_key_auth`` - reads a caller's key. - - Reuses the auth module's ``_get_bearer_token`` so the caller-key comparison - strips exactly the schemes authentication accepts (``Bearer`` / ``bearer`` / - ``Basic`` / ``AWS4-HMAC-SHA256`` credential), rather than re-deriving a - narrower normalization here. ``_get_bearer_token`` returns ``""`` for a value - with no recognized scheme prefix, so a bare token (or a real Google - credential that carries no scheme) falls back to its own value. - """ - return get_bearer_token(value) or value + return scheme_credential(value) or value _VERTEX_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-goog-api-key"}) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 15b2d18eeae..a50067e00f2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -2,7 +2,6 @@ import ipaddress import json import os import re -import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt @@ -14,6 +13,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm +from litellm._auth.key import is_master_key as is_master_key # noqa: PLC0414 # public re-export from litellm._logging import verbose_proxy_logger from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, @@ -135,16 +135,6 @@ def _get_max_string_length_prompt_in_db() -> int: return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB -def is_master_key(api_key: str | None, _master_key: str | None) -> bool: - """ - Raw-only constant-time master-key comparison. The hashed form is never - considered equivalent — only the raw master-key string matches. - """ - if _master_key is None or api_key is None: - return False - return secrets.compare_digest(api_key, _master_key) - - _is_master_key: Final = is_master_key diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b9189f143f..3c2ba2092d0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -141,6 +141,7 @@ from litellm import ( ModelResponseStream, Router, ) +from litellm._auth.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm._internal_context import service_target from litellm._logging import redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes @@ -5878,11 +5879,8 @@ class PrismaClient: # Default to success if metadata parsing fails return "success" - def hash_token(self, token: str): - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token + def hash_token(self, token: str) -> str: + return hash_token(token) def jsonify_object(self, data: Mapping[str, object]) -> dict[str, object]: db_data: Final[dict[str, object]] = copy.deepcopy(dict(data)) @@ -8457,15 +8455,6 @@ async def send_email( return True -def hash_token(token: str): - import hashlib - - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - def _is_pbkdf2_row(stored: str) -> bool: return re.fullmatch(r"pbkdf2:[A-Za-z0-9+/]{64}", stored) is not None @@ -8473,7 +8462,6 @@ def _is_pbkdf2_row(stored: str) -> bool: def hash_password(password: str) -> str: """Hash a password using scrypt with a random salt, or PBKDF2 in FIPS mode.""" import base64 - import hashlib import os if fips_mode_active(): @@ -8488,7 +8476,6 @@ def hash_password(password: str) -> str: def verify_password(password: str, stored: str) -> bool: """Verify a password against a stored hash. Supports scrypt and SHA256, or only PBKDF2 in FIPS mode.""" import base64 - import hashlib import secrets if fips_mode_active(): diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 77af87f5eb3..8a7c43372d8 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -22,6 +22,17 @@ def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... +@final +class NativeAuthContinuation: + def resume(self, answer: object) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... + +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_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: ... + @final class NativeClickHouseSpendConfig: def __new__(cls, database: str, url: str, retention_days: int) -> NativeClickHouseSpendConfig: ... diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 1183c36eda7..92067be067a 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -9,6 +9,7 @@ from typing import ( from pydantic import ConfigDict, Field, PrivateAttr, model_validator from typing_extensions import TypeIs # noqa: TID251 # narrows the before-validator input without copying it +from litellm._auth.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.team import Member from litellm.models.verification_token import LiteLLM_VerificationToken @@ -37,15 +38,6 @@ def is_jwt(token: str | None) -> bool: return len(parts) == 3 -def hash_token(token: str) -> str: - import hashlib - - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - _UntypedDict: TypeAlias = dict[Any, Any] diff --git a/tests/unit/_auth/__init__.py b/tests/unit/_auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/_auth/conftest.py b/tests/unit/_auth/conftest.py new file mode 100644 index 00000000000..1e171ca0db5 --- /dev/null +++ b/tests/unit/_auth/conftest.py @@ -0,0 +1,29 @@ +from collections.abc import Callable +from typing import Final, Protocol, cast + +import pytest + +from litellm._auth import runtime, rust_core +from litellm._auth.decision import AuthCore +from litellm._auth.python_core import PythonCore +from litellm._auth.shadow import ShadowCore +from litellm._auth.credential import PresentedKeys, SelectedKey, WebSocketDenial, WebSocketKeys + + +class CredentialSelector(Protocol): + def select_api_key(self, keys: PresentedKeys) -> SelectedKey | None: ... + def select_websocket_api_key(self, keys: WebSocketKeys) -> str | WebSocketDenial: ... + def scheme_credential(self, value: str) -> str | None: ... + def litellm_header_token(self, value: str) -> str: ... + + +@pytest.fixture(params=(PythonCore, rust_core.RustCore, ShadowCore), ids=("python", "rust", "shadow")) +def core(request: pytest.FixtureRequest) -> AuthCore: + factory: Final = cast(Callable[[], AuthCore], request.param) + 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 diff --git a/tests/unit/_auth/test_credential.py b/tests/unit/_auth/test_credential.py new file mode 100644 index 00000000000..5c9c21a1d10 --- /dev/null +++ b/tests/unit/_auth/test_credential.py @@ -0,0 +1,161 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from tests.unit._auth.conftest import CredentialSelector + +from litellm._auth.credential import ( + KeySource, + PresentedKeys, + SelectedKey, + WebSocketDenial, + WebSocketKeys, +) + +AWS4_BEARER: Final = ( + "AWS4-HMAC-SHA256 Credential=Bearer token/date/region/service/aws4_request, SignedHeaders=host, Signature=sig" +) +AWS4_RAW: Final = ( + "AWS4-HMAC-SHA256 Credential=token/date/region/service/aws4_request, SignedHeaders=host, Signature=sig" +) + +PRECEDENCE: Final = ( + ("custom_header", KeySource.CUSTOM_HEADER, "Bearer {}"), + ("litellm_header", KeySource.LITELLM_HEADER, "Bearer {}"), + ("authorization", KeySource.AUTHORIZATION, "Bearer {}"), + ("api_key", KeySource.API_KEY, "{}"), + ("x_api_key", KeySource.X_API_KEY, "{}"), + ("google", KeySource.GOOGLE, "{}"), + ("apim", KeySource.APIM, "{}"), + ("query", KeySource.QUERY, "{}"), + ("pass_through", KeySource.PASS_THROUGH, "{}"), +) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("Bearer token", "token"), + ("bearer token", "token"), + ("Basic dXNlcjpwYXNz", "dXNlcjpwYXNz"), + (AWS4_BEARER, "token"), + (AWS4_RAW, "token"), + ("Bearer ", ""), + ("Bearer token", " token"), + ("BEARER token", "token"), + ("BASIC dXNlcjpwYXNz", "dXNlcjpwYXNz"), + ("AWS4-HMAC-SHA256 SignedHeaders=host", None), + ("token", None), + ("Bearer", None), + ("Token token", None), + ("Bearertoken", None), + ], +) +def test_scheme_credential_strips_a_known_scheme_in_any_case( + credentials: CredentialSelector, value: str, expected: str | None +): + assert credentials.scheme_credential(value) == expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("token", "token"), + ("Bearer token", "token"), + (AWS4_BEARER, "token"), + ("AWS4-HMAC-SHA256 malformed", "AWS4-HMAC-SHA256 malformed"), + ("Bearer ", ""), + ], +) +def test_litellm_header_strips_a_known_scheme_or_keeps_the_raw_value( + credentials: CredentialSelector, value: str, expected: str +): + assert credentials.litellm_header_token(value) == expected + + +@pytest.mark.parametrize("rank", range(len(PRECEDENCE))) +def test_higher_priority_source_wins_over_every_lower_one(credentials: CredentialSelector, rank: int): + higher, source, template = PRECEDENCE[rank] + winner: Final = template.format(higher) + losers: Final = {field: lower_template.format("loser") for field, _, lower_template in PRECEDENCE[rank + 1 :]} + assert credentials.select_api_key(PresentedKeys(**losers, **{higher: winner})) == SelectedKey( + source=source, presented=winner, token=higher + ) + + +@pytest.mark.parametrize( + ("keys", "expected"), + [ + (PresentedKeys(custom_header="custom", x_api_key="lower"), ""), + (PresentedKeys(custom_header="", x_api_key="lower"), ""), + (PresentedKeys(authorization="token", x_api_key="lower"), ""), + (PresentedKeys(authorization="Token x", x_api_key="lower"), ""), + ], +) +def test_selected_source_without_a_known_scheme_yields_an_empty_token( + credentials: CredentialSelector, keys: PresentedKeys, expected: str +): + selected: Final = credentials.select_api_key(keys) + assert selected is not None + assert selected.token == expected + + +@pytest.mark.parametrize( + ("keys", "expected_source"), + [ + (PresentedKeys(authorization="", x_api_key="lower"), KeySource.X_API_KEY), + (PresentedKeys(query="", apim="lower"), KeySource.APIM), + (PresentedKeys(litellm_header=""), KeySource.LITELLM_HEADER), + (PresentedKeys(api_key=""), KeySource.API_KEY), + ], +) +def test_only_empty_authorization_and_query_fall_through( + credentials: CredentialSelector, keys: PresentedKeys, expected_source: KeySource +): + selected: Final = credentials.select_api_key(keys) + assert selected is not None + assert selected.source is expected_source + + +@pytest.mark.parametrize("keys", [PresentedKeys(), PresentedKeys(authorization=""), PresentedKeys(query="")]) +def test_nothing_presented_selects_nothing(credentials: CredentialSelector, keys: PresentedKeys): + assert credentials.select_api_key(keys) is None + + +def test_presented_keeps_the_raw_value_of_the_selected_source( + credentials: CredentialSelector, +): + keys: Final = PresentedKeys(authorization=AWS4_BEARER) + assert credentials.select_api_key(keys) == SelectedKey(KeySource.AUTHORIZATION, AWS4_BEARER, "token") + assert credentials.select_api_key(replace(keys, custom_header="Basic abc")) == SelectedKey( + KeySource.CUSTOM_HEADER, "Basic abc", "abc" + ) + + +@pytest.mark.parametrize( + ("keys", "expected"), + [ + (WebSocketKeys(authorization="Bearer token"), "token"), + (WebSocketKeys(authorization="bearer token "), "token"), + ( + WebSocketKeys(authorization="Bearer token", api_key="other", subprotocols="openai-insecure-api-key.x"), + "token", + ), + (WebSocketKeys(authorization="Basic token"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(authorization="token"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(authorization="Basic token", api_key="other"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(api_key="token"), "token"), + (WebSocketKeys(authorization="", api_key="token"), "token"), + (WebSocketKeys(api_key="token", subprotocols="openai-insecure-api-key.other"), "token"), + (WebSocketKeys(subprotocols="realtime, openai-insecure-api-key.token, openai-beta.realtime-v1"), "token"), + (WebSocketKeys(api_key="", subprotocols="openai-insecure-api-key.token"), "token"), + (WebSocketKeys(subprotocols="realtime"), WebSocketDenial.NO_API_KEY), + (WebSocketKeys(subprotocols="openai-insecure-api-key."), WebSocketDenial.NO_API_KEY), + (WebSocketKeys(), WebSocketDenial.NO_API_KEY), + ], +) +def test_websocket_takes_a_bearer_authorization_then_api_key_then_subprotocol( + credentials: CredentialSelector, keys: WebSocketKeys, expected: str | WebSocketDenial +): + assert credentials.select_websocket_api_key(keys) == expected diff --git a/tests/unit/_auth/test_key.py b/tests/unit/_auth/test_key.py new file mode 100644 index 00000000000..1d41920e4dd --- /dev/null +++ b/tests/unit/_auth/test_key.py @@ -0,0 +1,35 @@ +import hashlib + +import pytest + +from litellm._auth.key import hash_token, is_master_key + + +@pytest.mark.parametrize("token", ["sk-1234", "", "sk-ünïcode"]) +def test_hash_token_is_sha256_hex_of_utf8(token: str): + assert hash_token(token) == hashlib.sha256(token.encode()).hexdigest() + + +def test_hash_token_differs_per_token(): + assert hash_token("sk-a") != hash_token("sk-b") + + +@pytest.mark.parametrize( + ("token", "master_key", "expected"), + [ + ("sk-master", "sk-master", True), + ("sk-master", "sk-other", False), + ("sk-mast", "sk-master", False), + (None, "sk-master", False), + ("sk-master", None, False), + (None, None, False), + ("sk-mästér", "sk-mästér", True), + ("sk-mästér", "sk-master", False), + ], +) +def test_is_master_key_matches_only_the_exact_raw_key(token: str | None, master_key: str | None, expected: bool): + assert is_master_key(token, master_key) is expected + + +def test_hashed_master_key_is_not_the_master_key(): + assert is_master_key(hash_token("sk-master"), "sk-master") is False diff --git a/tests/unit/_auth/test_model_budgets.py b/tests/unit/_auth/test_model_budgets.py new file mode 100644 index 00000000000..5c704f8b465 --- /dev/null +++ b/tests/unit/_auth/test_model_budgets.py @@ -0,0 +1,327 @@ +from dataclasses import dataclass, replace +from typing import Final + +import pytest + +from litellm._auth.decision import ( + Allow, + Deny, + FallbackTarget, + ModelBudgetExceeded, + ModelBudgetScope, + Outcome, + Step, + TeamMemberModelBudgetScope, + WithinModelBudgets, +) +from litellm._auth.driver import drive_model_budgets, drive_team_member_model_budgets +from litellm._auth.decision import AuthCore + +Call = tuple[str, str | None] + +SCOPE: Final = ModelBudgetScope( + models=("gpt-4o",), + rewritable_model="gpt-4o", + key_has_model_budgets=True, + key_budget_fallbacks={"gpt-4o": ("gpt-4o-mini", "haiku")}, + user_has_model_budgets=True, +) +ALLOWED: Final = Allow(WithinModelBudgets()) + + +def _exceeded(entity_type: str, model: str) -> ModelBudgetExceeded: + return ModelBudgetExceeded( + message=f"{entity_type} over budget for {model}", + current_cost=2.0, + max_budget=1.0, + entity_type=entity_type, + entity_id="id", + ) + + +@dataclass +class FakeRequest: + scope: ModelBudgetScope + key_over: frozenset[str] = frozenset() + user_over: frozenset[str] = frozenset() + unauthorized: frozenset[str] = frozenset() + calls: tuple[Call, ...] = () + + def _record(self, name: str, model: str | None = None) -> None: + self.calls = (*self.calls, (name, model)) + + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read") + return self.scope + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("key", model) + return _exceeded("key", model) if model in self.key_over else None + + async def authorize_budget_fallback(self, model: str) -> bool: + self._record("authorize", model) + return model not in self.unauthorized + + async def rewrite_request_model(self, model: str) -> None: + self._record("rewrite", model) + self.scope = replace( + self.scope, + models=tuple(model if m == self.scope.rewritable_model else m for m in self.scope.models), + rewritable_model=model, + ) + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("user", model) + return _exceeded("user", model) if model in self.user_over else None + + +async def _run(step: Step, request: FakeRequest) -> tuple[Outcome, tuple[Call, ...]]: + outcome: Final = await drive_model_budgets(step, request) + return outcome, request.calls + + +@pytest.mark.asyncio +async def test_within_budget_checks_the_key_then_the_user_without_rerouting( + core: AuthCore, +): + assert await _run(core.check_model_budgets(), FakeRequest(SCOPE)) == ( + ALLOWED, + (("read", None), ("key", "gpt-4o"), ("read", None), ("user", "gpt-4o")), + ) + + +@pytest.mark.asyncio +async def test_over_budget_reroutes_to_the_first_fallback_within_budget_and_checks_the_user_on_it( + core: AuthCore, +): + request: Final = FakeRequest(SCOPE, key_over=frozenset({"gpt-4o", "gpt-4o-mini"})) + assert await _run(core.check_model_budgets(), request) == ( + ALLOWED, + ( + ("read", None), + ("key", "gpt-4o"), + ("key", "gpt-4o-mini"), + ("key", "haiku"), + ("authorize", "haiku"), + ("rewrite", "haiku"), + ("read", None), + ("user", "haiku"), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "key_over", "unauthorized"), + [ + (replace(SCOPE, rewritable_model=None), {"gpt-4o"}, set[str]()), + (replace(SCOPE, rewritable_model="other"), {"gpt-4o"}, set[str]()), + (replace(SCOPE, key_budget_fallbacks={}), {"gpt-4o"}, set[str]()), + (SCOPE, {"gpt-4o", "gpt-4o-mini", "haiku"}, set[str]()), + (SCOPE, {"gpt-4o"}, {"gpt-4o-mini"}), + ], + ids=["nested_model", "different_top_level_model", "no_fallbacks", "every_fallback_over", "fallback_unauthorized"], +) +async def test_an_unreroutable_over_budget_model_is_denied_with_its_own_budget_error( + core: AuthCore, scope: ModelBudgetScope, key_over: set[str], unauthorized: set[str] +): + request: Final = FakeRequest(scope, key_over=frozenset(key_over), unauthorized=frozenset(unauthorized)) + outcome, calls = await _run(core.check_model_budgets(), request) + assert outcome == Deny(_exceeded("key", "gpt-4o")) + assert not any(name in ("rewrite", "user") for name, _ in calls) + + +@pytest.mark.asyncio +async def test_a_later_model_is_checked_after_an_earlier_one_was_rerouted( + core: AuthCore, +): + scope: Final = replace(SCOPE, models=("gpt-4o", "embed")) + request: Final = FakeRequest(scope, key_over=frozenset({"gpt-4o", "embed"})) + outcome, calls = await _run(core.check_model_budgets(), request) + assert outcome == Deny(_exceeded("key", "embed")) + assert calls[-2:] == (("rewrite", "gpt-4o-mini"), ("key", "embed")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "expected_calls"), + [ + ( + replace(SCOPE, key_has_model_budgets=False), + (("read", None), ("user", "gpt-4o")), + ), + ( + replace(SCOPE, user_has_model_budgets=False), + (("read", None), ("key", "gpt-4o"), ("read", None)), + ), + (replace(SCOPE, models=()), (("read", None),)), + ], + ids=["no_key_budgets", "no_user_budgets", "no_models"], +) +async def test_only_budgets_that_exist_are_checked( + core: AuthCore, scope: ModelBudgetScope, expected_calls: tuple[Call, ...] +): + assert await _run(core.check_model_budgets(), FakeRequest(scope)) == (ALLOWED, expected_calls) + + +@pytest.mark.asyncio +async def test_a_user_over_budget_is_denied_on_every_path( + core: AuthCore, +): + over: Final = frozenset({"gpt-4o"}) + for step in (core.check_model_budgets(), core.check_user_model_budgets()): + outcome, _ = await _run(step, FakeRequest(SCOPE, user_over=over)) + assert outcome == Deny(_exceeded("user", "gpt-4o")) + + +@pytest.mark.asyncio +async def test_the_user_only_check_never_reads_key_budgets( + core: AuthCore, +): + request: Final = FakeRequest(SCOPE, key_over=frozenset({"gpt-4o"})) + assert await _run(core.check_user_model_budgets(), request) == ( + ALLOWED, + (("read", None), ("user", "gpt-4o")), + ) + + +MEMBER_SCOPE: Final = TeamMemberModelBudgetScope( + models=("gpt-4o",), rewritable_model="gpt-4o", has_model_budgets=True, can_fall_back=True +) + + +def _target(model: str) -> FallbackTarget: + return FallbackTarget(requested=model, model=model, served_model=model) + + +@dataclass +class FakeTeamMemberRequest: + scope: TeamMemberModelBudgetScope + over: frozenset[str] = frozenset() + fallbacks: tuple[str, ...] = ("gpt-4o", "unusable", "gpt-4o-mini", "haiku") + unusable: frozenset[str] = frozenset({"unusable"}) + calls: tuple[Call, ...] = () + targets: tuple[FallbackTarget, ...] | None = None + vetted: tuple[FallbackTarget, ...] = () + + def _record(self, name: str, model: str | None = None) -> None: + self.calls = (*self.calls, (name, model)) + + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: + self._record("read") + return self.scope + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("member", model) + return _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: + self._record("vet", target.model) + self.vetted = (*self.vetted, target) + return target.model not in self.unusable + + async def reroute_team_member_request(self, model: str) -> None: + self._record("reroute", model) + + +async def _run_member(core: AuthCore, request: FakeTeamMemberRequest) -> tuple[Outcome, tuple[Call, ...]]: + outcome: Final = await drive_team_member_model_budgets(core.check_team_member_model_budgets(), request) + return outcome, request.calls + + +@pytest.mark.asyncio +async def test_large_user_budget_scopes_do_not_exhaust_the_driver_stack(core: AuthCore): + count: Final = 600 + scope: Final = replace(SCOPE, models=("model",) * count, key_has_model_budgets=False) + request: Final = FakeRequest(scope) + assert await drive_model_budgets(core.check_user_model_budgets(), request) == ALLOWED + assert request.calls == (("read", None),) + (("user", "model"),) * count + + +@pytest.mark.asyncio +async def test_member_fallback_preserves_requested_resolved_and_served_aliases(core: AuthCore): + target: Final = FallbackTarget(requested="public-alias", model="resolved-model", served_model="deployment") + request: Final = FakeTeamMemberRequest(MEMBER_SCOPE, over=frozenset({"gpt-4o"}), targets=(target,)) + assert await _run_member(core, request) == ( + ALLOWED, + ( + ("read", None), + ("member", "gpt-4o"), + ("fallbacks", "gpt-4o"), + ("vet", "resolved-model"), + ("reroute", "resolved-model"), + ), + ) + assert request.vetted == (target,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "expected_calls"), + [ + (MEMBER_SCOPE, (("read", None), ("member", "gpt-4o"))), + (replace(MEMBER_SCOPE, has_model_budgets=False), (("read", None),)), + (replace(MEMBER_SCOPE, models=()), (("read", None),)), + ], + ids=["within_budget", "no_member_budgets", "no_models"], +) +async def test_a_team_member_within_budget_is_allowed( + core: AuthCore, scope: TeamMemberModelBudgetScope, expected_calls: tuple[Call, ...] +): + assert await _run_member(core, FakeTeamMemberRequest(scope)) == (ALLOWED, expected_calls) + + +@pytest.mark.asyncio +async def test_an_over_budget_team_member_reroutes_to_the_first_usable_fallback_other_than_itself( + core: AuthCore, +): + request: Final = FakeTeamMemberRequest( + replace(MEMBER_SCOPE, models=("gpt-4o", "embed")), over=frozenset({"gpt-4o"}) + ) + assert await _run_member(core, request) == ( + ALLOWED, + ( + ("read", None), + ("member", "gpt-4o"), + ("fallbacks", "gpt-4o"), + ("vet", "unusable"), + ("vet", "gpt-4o-mini"), + ("reroute", "gpt-4o-mini"), + ("member", "embed"), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "unusable"), + [ + (replace(MEMBER_SCOPE, rewritable_model=None), frozenset[str]()), + (replace(MEMBER_SCOPE, can_fall_back=False), frozenset[str]()), + (MEMBER_SCOPE, frozenset({"unusable", "gpt-4o-mini", "haiku"})), + ], + ids=["nested_model", "fallbacks_unavailable", "no_usable_fallback"], +) +async def test_an_over_budget_team_member_without_a_usable_fallback_is_denied_with_its_own_error( + core: AuthCore, scope: TeamMemberModelBudgetScope, unusable: frozenset[str] +): + outcome, calls = await _run_member( + core, FakeTeamMemberRequest(scope, over=frozenset({"gpt-4o"}), unusable=unusable) + ) + assert outcome == Deny(_exceeded("member", "gpt-4o")) + assert not any(name == "reroute" for name, _ in calls) + + +@pytest.mark.asyncio +async def test_a_later_model_is_not_rerouted_once_the_request_model_changed( + core: AuthCore, +): + scope: Final = replace(MEMBER_SCOPE, models=("gpt-4o", "gpt-4o")) + request: Final = FakeTeamMemberRequest(scope, over=frozenset({"gpt-4o"})) + outcome, calls = await _run_member(core, request) + assert outcome == Deny(_exceeded("member", "gpt-4o")) + assert [name for name, _ in calls].count("reroute") == 1 diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py new file mode 100644 index 00000000000..4f1fe014d0f --- /dev/null +++ b/tests/unit/_auth/test_python_core.py @@ -0,0 +1,628 @@ +import math +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from typing import Final +from zoneinfo import ZoneInfo + +import pytest + +from litellm._auth.decision import ( + AdminKey, + Allow, + CachedAdminKey, + Deny, + ExpiredKey, + KeyBlocked, + KeyBudgetExceeded, + KeyFacts, + ModelBudgetExceeded, + ModelBudgetScope, + KeyWithoutToken, + MalformedKey, + MasterKey, + MasterKeyOnlyRoute, + MissingKey, + NoDatabase, + NoMasterKey, + NonAdminKey, + NotAVirtualKey, + Outcome, + OwnerDeactivated, + OwnerFacts, + PresentedKey, + RequestFacts, + Settings, + Step, + TeamMemberBudget, + TeamMemberBudgetExceeded, + UserAuthDisabled, + UserAuthRoute, + VirtualKey, +) +from litellm._auth.driver import drive +from litellm._auth.key import hash_token +from litellm._auth.decision import AuthCore + +MASTER: Final = "sk-master" +NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc) +SETTINGS: Final = Settings( + master_key=MASTER, + allow_user_auth=False, + has_db=True, + master_key_only_routes=frozenset({"/global/spend/reset"}), + proxy_max_budget=0.0, +) +PLAIN_KEY: Final = KeyFacts( + has_token=True, + blocked=False, + expires=None, + spend=0.0, + soft_budget=None, + max_budget=None, + team_member_spend=None, + user_id=None, + team_id=None, + budget_throttle_pct=None, +) +CHECKED: Final = ( + "prepare_key", + "inspect_key", + "enforce_model_access", + "check_budget_exemption", + "send_key_budget_alerts", + "read_model_budget_scope", + "load_key_context", +) + +NO_MODEL_BUDGETS: Final = ModelBudgetScope( + models=(), + rewritable_model=None, + key_has_model_budgets=False, + key_budget_fallbacks={}, + user_has_model_budgets=False, +) + +Call = tuple[str, object] + + +@dataclass +class RecordingHandler: + presented: PresentedKey | None = None + key: KeyFacts = PLAIN_KEY + owner: OwnerFacts | None = None + exempt: bool = False + team_member_budget: TeamMemberBudget | None = None + key_spend: float = 0.0 + proxy_spend: float | None = None + model_budget_scope: ModelBudgetScope = NO_MODEL_BUDGETS + user_model_budget_exceeded: ModelBudgetExceeded | None = None + calls: tuple[Call, ...] = () + + def _record(self, name: str, payload: object = None) -> None: + self.calls = (*self.calls, (name, payload)) + + async def resolve_end_user(self, prefetch_key: bool) -> None: + self._record("resolve_end_user", prefetch_key) + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: + self._record("resolve_presented_key", hashed_token) + return self.presented + + async def evict_presented_key(self, hashed_token: str) -> None: + self._record("evict_presented_key", hashed_token) + + async def load_key(self, hashed_token: str) -> None: + self._record("load_key", hashed_token) + + async def prepare_key(self, loaded_from_db: bool) -> None: + self._record("prepare_key", loaded_from_db) + + async def inspect_key(self) -> KeyFacts: + self._record("inspect_key") + return self.key + + async def enforce_model_access(self) -> None: + self._record("enforce_model_access") + + 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: + self._record("check_budget_exemption") + return self.exempt + + 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)) + return self.team_member_budget + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: + self._record("alert_team_member_budget", budget) + + async def send_key_budget_alerts(self) -> None: + self._record("send_key_budget_alerts") + + async def read_key_spend(self, max_budget: float) -> float: + self._record("read_key_spend", max_budget) + return self.key_spend + + async def alert_key_budget(self, spend: float) -> None: + self._record("alert_key_budget", spend) + + async def alert_key_soft_budget(self) -> None: + self._record("alert_key_soft_budget") + + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read_model_budget_scope") + return self.model_budget_scope + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("check_key_model_budget", model) + return None + + async def authorize_budget_fallback(self, model: str) -> bool: + self._record("authorize_budget_fallback", model) + return True + + 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: + self._record("check_user_model_budget", model) + return self.user_model_budget_exceeded + + async def load_key_context(self) -> None: + self._record("load_key_context") + + async def read_proxy_spend(self) -> float | None: + self._record("read_proxy_spend") + return self.proxy_spend + + async def alert_proxy_budget(self, spend: float) -> None: + self._record("alert_proxy_budget", spend) + + +def _facts(token: str | None, route: str = "/chat/completions", is_llm_api_route: bool = True) -> RequestFacts: + return RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route) + + +async def _run(step: Step, handler: RecordingHandler) -> tuple[Outcome, tuple[str, ...]]: + outcome: Final = await drive(step, handler) + return outcome, tuple(name for name, _ in handler.calls) + + +async def _decide( + core: AuthCore, + token: str | None, + *, + route: str = "/chat/completions", + settings: Settings = SETTINGS, + handler: RecordingHandler | None = None, +) -> tuple[Outcome, tuple[str, ...]]: + return await _run(core.start(_facts(token, route), settings, NOW), handler or RecordingHandler()) + + +async def _check( + core: AuthCore, + key: KeyFacts, + *, + is_llm_api_route: bool = True, + settings: Settings = SETTINGS, + owner: OwnerFacts | None = None, + exempt: bool = False, + team_member_budget: TeamMemberBudget | None = None, + key_spend: float = 0.0, +) -> tuple[Outcome, tuple[str, ...]]: + handler: Final = RecordingHandler( + key=key, owner=owner, exempt=exempt, team_member_budget=team_member_budget, key_spend=key_spend + ) + facts: Final = _facts("hash", is_llm_api_route=is_llm_api_route) + return await _run(core.start_resolved_key(facts, settings, NOW), handler) + + +def _allowed(budget_throttle_pct: float | None = None) -> Outcome: + return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "token", "settings", "expected"), + [ + ("/chat/completions", None, Settings(None, False, True, frozenset(), 0.0), Allow(NoMasterKey(token=None))), + ("/chat/completions", "sk-x", Settings(None, False, False, frozenset(), 0.0), Allow(NoMasterKey(token="sk-x"))), + ("/chat/completions", None, SETTINGS, Deny(MissingKey())), + ("/chat/completions", "", SETTINGS, Deny(MalformedKey())), + ("/user/auth", "", Settings(MASTER, True, True, frozenset(), 0.0), Deny(MalformedKey())), + ("/user/auth", "sk-x", Settings(MASTER, True, True, frozenset(), 0.0), Allow(UserAuthRoute())), + ("/user/auth", "sk-x", SETTINGS, Deny(UserAuthDisabled())), + ("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())), + ], +) +async def test_credential_decisions_do_no_io( + core: AuthCore, route: str, token: str | None, settings: Settings, expected: Outcome +): + assert await _decide(core, token, route=route, settings=settings) == (expected, ()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"]) +async def test_master_key_resolves_end_user_without_key_prefetch(core: AuthCore, route: str): + handler: Final = RecordingHandler() + assert await _decide(core, MASTER, route=route, handler=handler) == (Allow(MasterKey()), ("resolve_end_user",)) + assert handler.calls == (("resolve_end_user", False),) + + +@pytest.mark.asyncio +async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked( + core: AuthCore, +): + handler: Final = RecordingHandler() + outcome, calls = await _decide(core, "sk-x", handler=handler) + assert outcome == _allowed() + assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", *CHECKED) + assert handler.calls[2] == ("load_key", hash_token("sk-x")) + assert handler.calls[3] == ("prepare_key", True) + + +@pytest.mark.asyncio +async def test_cached_virtual_key_is_checked_without_a_db_load( + core: AuthCore, +): + handler: Final = RecordingHandler(presented=NonAdminKey()) + outcome, calls = await _decide(core, "sk-x", handler=handler) + assert outcome == _allowed() + assert calls == ("resolve_end_user", "resolve_presented_key", *CHECKED) + assert handler.calls[2] == ("prepare_key", False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW]) +async def test_unexpired_cached_admin_key_skips_route_db_and_key_checks(core: AuthCore, expires: datetime | None): + no_db: Final = Settings(MASTER, False, False, frozenset({"/global/spend/reset"}), 0.0) + outcome, calls = await _decide( + core, + "sk-admin", + route="/global/spend/reset", + settings=no_db, + handler=RecordingHandler(presented=AdminKey(expires=expires)), + ) + assert (outcome, calls) == (Allow(CachedAdminKey()), ("resolve_end_user", "resolve_presented_key")) + + +@pytest.mark.asyncio +async def test_expired_cached_admin_key_is_evicted_then_denied_with_the_presented_key( + core: AuthCore, +): + expires: Final = NOW - timedelta(seconds=1) + handler: Final = RecordingHandler(presented=AdminKey(expires=expires)) + outcome, calls = await _decide(core, "sk-admin", handler=handler) + assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label="sk-admin")) + assert calls == ("resolve_end_user", "resolve_presented_key", "evict_presented_key") + assert handler.calls[-1] == ("evict_presented_key", hash_token("sk-admin")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("presented", [None, NonAdminKey()]) +async def test_non_admin_key_cannot_use_master_key_only_routes(core: AuthCore, presented: PresentedKey | None): + outcome, calls = await _decide( + core, "sk-x", route="/global/spend/reset", handler=RecordingHandler(presented=presented) + ) + assert (outcome, calls) == (Deny(MasterKeyOnlyRoute()), ("resolve_end_user", "resolve_presented_key")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("presented", [None, NonAdminKey()]) +async def test_virtual_key_needs_a_database(core: AuthCore, presented: PresentedKey | None): + no_db: Final = Settings(MASTER, False, False, frozenset(), 0.0) + outcome, _ = await _decide(core, "sk-x", settings=no_db, handler=RecordingHandler(presented=presented)) + assert outcome == Deny(NoDatabase()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token", ["eyJhbGciOiJIUzI1NiJ9.e30.sig", "a" * 64, "sk"]) +async def test_uncached_token_without_virtual_key_prefix_is_rejected_before_the_db(core: AuthCore, token: str): + assert await _decide(core, token) == (Deny(NotAVirtualKey()), ("resolve_end_user", "resolve_presented_key")) + + +@pytest.mark.asyncio +async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with( + core: AuthCore, +): + expires: Final = NOW - timedelta(seconds=1) + outcome, _ = await _decide(core, "sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires))) + assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x"))) + + +@pytest.mark.asyncio +async def test_blocked_key_is_denied_before_model_access( + core: AuthCore, +): + assert await _check(core, replace(PLAIN_KEY, blocked=True)) == (Deny(KeyBlocked()), ("inspect_key",)) + + +@pytest.mark.asyncio +async def test_deactivated_owner_is_denied_after_model_access( + core: AuthCore, +): + outcome, calls = await _check(core, replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True)) + assert (outcome, calls) == ( + Deny(OwnerDeactivated(user_id="u1")), + ("inspect_key", "enforce_model_access", "load_key_owner"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", [None, OwnerFacts(deactivated=False)]) +async def test_active_or_missing_owner_continues(core: AuthCore, owner: OwnerFacts | None): + outcome, _ = await _check(core, replace(PLAIN_KEY, user_id="u1"), owner=owner) + assert outcome == _allowed() + + +@pytest.mark.asyncio +async def test_exempt_request_skips_every_budget_check_but_not_expiry( + core: AuthCore, +): + key: Final = replace( + PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=99.0, max_budget=1.0, spend=5.0, soft_budget=1.0 + ) + assert await _check(core, key, exempt=True) == ( + _allowed(), + ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption", "load_key_context"), + ) + expired: Final = replace(key, expires=NOW - timedelta(seconds=1)) + outcome, _ = await _check(core, expired, exempt=True) + assert outcome == Deny(ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("budget", "denied"), + [ + (TeamMemberBudget(spend=10.0, max_budget=10.0), True), + (TeamMemberBudget(spend=10.1, max_budget=10.0), True), + (TeamMemberBudget(spend=9.9, max_budget=10.0), False), + (None, False), + ], +) +async def test_team_member_budget_alerts_then_denies_at_the_limit( + core: AuthCore, budget: TeamMemberBudget | None, denied: bool +): + key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0) + outcome, calls = await _check(core, key, team_member_budget=budget) + if denied: + assert budget is not None + assert outcome == Deny( + TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=budget.spend, max_budget=budget.max_budget) + ) + assert calls[-2:] == ("read_team_member_budget", "alert_team_member_budget") + else: + assert outcome == _allowed() + assert "alert_team_member_budget" not in calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key", + [ + replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=None), + replace(PLAIN_KEY, user_id=None, team_id="t1", team_member_spend=0.0), + replace(PLAIN_KEY, user_id="u1", team_id=None, team_member_spend=0.0), + ], +) +async def test_team_member_budget_is_read_only_for_a_member_with_tracked_spend(core: AuthCore, key: KeyFacts): + _, calls = await _check(core, key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert "read_team_member_budget" not in calls + + +@pytest.mark.asyncio +async def test_team_member_budget_needs_a_database( + core: AuthCore, +): + key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0) + no_db: Final = Settings(MASTER, False, False, frozenset(), 0.0) + _, calls = await _check(core, key, settings=no_db, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert "read_team_member_budget" not in calls + + +@pytest.mark.asyncio +async def test_team_member_budget_is_checked_before_expiry( + core: AuthCore, +): + key: Final = replace( + PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0, expires=NOW - timedelta(seconds=1) + ) + outcome, _ = await _check(core, key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert outcome == Deny(TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=1.0, max_budget=0.5)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend", "max_budget", "throttle", "expected"), + [ + (5.0, 5.0, None, Deny(KeyBudgetExceeded(spend=5.0, max_budget=5.0))), + (6.0, 5.0, None, Deny(KeyBudgetExceeded(spend=6.0, max_budget=5.0))), + (4.9, 5.0, None, _allowed()), + (6.0, 5.0, 0.5, _allowed(budget_throttle_pct=0.5)), + (4.9, 5.0, 0.5, _allowed()), + ( + 1e9, + math.inf, + None, + _allowed(), + ), + ( + 1e9, + math.nan, + None, + _allowed(), + ), + ], +) +async def test_key_budget_denies_or_throttles_at_the_limit( + core: AuthCore, spend: float, max_budget: float, throttle: float | None, expected: Outcome +): + key: Final = replace(PLAIN_KEY, max_budget=max_budget, budget_throttle_pct=throttle) + handler: Final = RecordingHandler(key=key, key_spend=spend) + outcome, calls = await _run(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert outcome == expected + assert calls[3:6] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") + assert handler.calls[5] == ("alert_key_budget", spend) + assert ("read_model_budget_scope" in calls) == isinstance(expected, Allow) + + +@pytest.mark.asyncio +async def test_key_budget_is_enforced_only_on_llm_routes( + core: AuthCore, +): + key: Final = replace(PLAIN_KEY, max_budget=1.0) + outcome, calls = await _check(core, key, is_llm_api_route=False, key_spend=2.0) + assert outcome == _allowed() + assert "read_key_spend" not in calls + assert "read_model_budget_scope" in calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend", "soft_budget", "alerted"), + [ + (5.0, 5.0, True), + (6.0, 5.0, True), + (4.9, 5.0, False), + (10.0, None, False), + (10.0, 0.0, False), + (10.0, math.nan, False), + ], +) +async def test_soft_budget_alerts_at_the_threshold_before_model_budgets( + core: AuthCore, spend: float, soft_budget: float | None, alerted: bool +): + outcome, calls = await _check(core, replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget)) + assert outcome == _allowed() + soft_alert: Final = ("alert_key_soft_budget",) if alerted else () + assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context") + + +@pytest.mark.asyncio +async def test_throttled_key_still_gets_soft_and_model_budget_checks( + core: AuthCore, +): + key: Final = replace(PLAIN_KEY, max_budget=5.0, budget_throttle_pct=0.5, spend=6.0, soft_budget=1.0) + outcome, calls = await _check(core, key, key_spend=6.0) + assert outcome == _allowed(budget_throttle_pct=0.5) + assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "read_model_budget_scope", "load_key_context") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exempt", [False, True]) +@pytest.mark.parametrize( + ("proxy_max_budget", "has_db", "proxy_spend", "tail"), + [ + (0.0, True, 50.0, (("load_key_context", None),)), + (-1.0, True, 50.0, (("load_key_context", None),)), + (math.nan, True, 50.0, (("load_key_context", None),)), + (10.0, False, 50.0, (("load_key_context", None),)), + (10.0, True, None, (("load_key_context", None), ("read_proxy_spend", None))), + ( + 10.0, + True, + 50.0, + (("load_key_context", None), ("read_proxy_spend", None), ("alert_proxy_budget", 50.0)), + ), + ], +) +async def test_proxy_budget_only_alerts_and_runs_for_exempt_requests_too( + core: AuthCore, + exempt: bool, + proxy_max_budget: float, + has_db: bool, + proxy_spend: float | None, + tail: tuple[Call, ...], +): + settings: Final = replace(SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db) + handler: Final = RecordingHandler(exempt=exempt, proxy_spend=proxy_spend) + outcome, _ = await _run(core.start_resolved_key(_facts("hash"), settings, NOW), handler) + assert outcome == _allowed() + assert handler.calls[-len(tail) :] == tail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exempt", [False, True]) +async def test_key_without_token_is_denied_after_its_context_loads(core: AuthCore, exempt: bool): + settings: Final = replace(SETTINGS, proxy_max_budget=10.0) + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, has_token=False), exempt=exempt, proxy_spend=1.0) + outcome, calls = await _run(core.start_resolved_key(_facts("hash"), settings, NOW), handler) + assert outcome == Deny(KeyWithoutToken()) + assert calls[-3:] == ("load_key_context", "read_proxy_spend", "alert_proxy_budget") + + +@pytest.mark.asyncio +async def test_a_model_budget_denial_stops_the_key_chain_before_its_context_loads( + core: AuthCore, +): + exceeded: Final = ModelBudgetExceeded( + message="over", current_cost=2.0, max_budget=1.0, entity_type="user", entity_id="u-1" + ) + scope: Final = replace(NO_MODEL_BUDGETS, models=("gpt-4o",), user_has_model_budgets=True) + handler: Final = RecordingHandler(model_budget_scope=scope, user_model_budget_exceeded=exceeded) + outcome, calls = await _run(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert outcome == Deny(exceeded) + assert calls[-2:] == ("read_model_budget_scope", "check_user_model_budget") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("now", "expires", "expired"), + [ + (NOW, NOW, False), + (NOW, NOW - timedelta(microseconds=1), True), + (NOW, NOW.astimezone(timezone(timedelta(hours=9))), False), + (NOW.replace(tzinfo=None), NOW.replace(tzinfo=None) - timedelta(microseconds=1), True), + ( + datetime(2026, 11, 1, 1, 20, fold=1, tzinfo=ZoneInfo("America/New_York")), + datetime(2026, 11, 1, 1, 30, fold=0, tzinfo=ZoneInfo("America/New_York")), + False, + ), + ( + datetime(2026, 11, 1, 1, 20, fold=1, tzinfo=ZoneInfo("America/New_York")), + datetime(2026, 11, 1, 1, 30, fold=0, tzinfo=ZoneInfo.no_cache("America/New_York")), + True, + ), + ], +) +async def test_datetime_comparisons_preserve_timezone_identity_and_public_values( + core: AuthCore, now: datetime, expires: datetime, expired: bool +): + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, expires=expires)) + outcome: Final = await drive(core.start_resolved_key(_facts("hash"), SETTINGS, now), handler) + if expired: + assert outcome == Deny(ExpiredKey(expires, now, "hash")) + assert isinstance(outcome, Deny) and isinstance(outcome.denial, ExpiredKey) + assert outcome.denial.expires is expires + assert outcome.denial.now is now + else: + assert outcome == _allowed() + + +@pytest.mark.asyncio +async def test_mixed_naive_and_aware_expiry_is_terminal_at_the_original_check(core: AuthCore): + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, expires=NOW.replace(tzinfo=None))) + with pytest.raises(TypeError, match="offset-naive and offset-aware"): + await drive(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert handler.calls == (("inspect_key", None), ("enforce_model_access", None), ("check_budget_exemption", None)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget", [math.nan, math.inf, -math.inf, 0.0]) +async def test_team_member_nonfinite_budgets_keep_the_reference_comparison(core: AuthCore, budget: float): + handler: Final = RecordingHandler( + key=replace(PLAIN_KEY, user_id="u", team_id="t", team_member_spend=0), + team_member_budget=TeamMemberBudget(1.0, budget), + ) + outcome: Final = await drive(core.start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + if budget == math.inf: + assert outcome == _allowed() + else: + assert isinstance(outcome, Deny) and isinstance(outcome.denial, TeamMemberBudgetExceeded) + assert outcome.denial.spend == 1.0 + assert outcome.denial.max_budget == budget or (math.isnan(outcome.denial.max_budget) and math.isnan(budget)) + assert handler.calls[-1][0] == "alert_team_member_budget" diff --git a/tests/unit/_auth/test_runtime.py b/tests/unit/_auth/test_runtime.py new file mode 100644 index 00000000000..830b743ee6e --- /dev/null +++ b/tests/unit/_auth/test_runtime.py @@ -0,0 +1,36 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.decision import Deny, InspectKey, KeyBlocked, RequestFacts +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")) +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) + else: + monkeypatch.setenv("LITELLM_AUTH_CORE", backend) + 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) == 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 new file mode 100644 index 00000000000..96e39e0ffa2 --- /dev/null +++ b/tests/unit/_auth/test_rust_core.py @@ -0,0 +1,126 @@ +from dataclasses import replace +from typing import Final, cast + +import pytest + +from litellm._auth.decision import ( + Allow, + AuthCore, + Deny, + KeyBlocked, + NoMasterKey, + RequestFacts, + ResolveEndUser, + Step, + WithinModelBudgets, +) +from litellm._auth.driver import drive +from litellm._auth.rust_core import RustCore +from litellm._auth.shadow import ShadowCore, ShadowMismatch +from litellm.rust_bridge import _native +from tests.unit._auth.test_python_core import NOW, PLAIN_KEY, SETTINGS, RecordingHandler + + +class ScriptedCore: + def __init__(self, step: Step) -> None: + self.step: Final = step + self.calls = 0 + + def check_model_budgets(self) -> Step: + self.calls += 1 + return self.step + + +FACTS: Final = RequestFacts("/chat/completions", "hash", True) + + +def test_native_continuations_are_one_shot_and_own_their_pending_decision(): + payload, continuation = _native.request_auth_budgets("Key") + assert payload == {"kind": "ReadModelBudgetScope"} + assert isinstance(continuation, _native.NativeAuthContinuation) + outcome, done = continuation.resume( + { + "kind": "ModelScope", + "value": { + "models": [], + "rewritable_model": None, + "key_has_model_budgets": False, + "key_budget_fallbacks": {}, + "user_has_model_budgets": False, + }, + } + ) + assert outcome == {"kind": "Allow", "grant": {"kind": "WithinModelBudgets"}} + assert done is None + with pytest.raises(RuntimeError, match="already consumed"): + continuation.resume({"kind": "Unit"}) + + +def test_invalid_native_answer_consumes_the_continuation_without_fallback(): + _, continuation = _native.request_auth_budgets("Key") + assert continuation is not None + with pytest.raises(RuntimeError, match="invalid auth effect answer"): + continuation.resume({"kind": "Bool", "value": False}) + with pytest.raises(RuntimeError, match="already consumed"): + continuation.resume({"kind": "Unit"}) + + +@pytest.mark.asyncio +async def test_shadow_replays_only_after_the_primary_finishes_and_uses_one_io_trace(): + primary: Final = ScriptedCore(ResolveEndUser(False, lambda: Allow(WithinModelBudgets()))) + reference: Final = ScriptedCore(ResolveEndUser(False, lambda: Allow(WithinModelBudgets()))) + step: Final = ShadowCore(cast(AuthCore, primary), cast(AuthCore, reference)).check_model_budgets() + assert reference.calls == 0 + handler: Final = RecordingHandler() + assert await drive(step, handler) == Allow(WithinModelBudgets()) + assert handler.calls == (("resolve_end_user", False),) + assert (primary.calls, reference.calls) == (1, 1) + + +@pytest.mark.asyncio +async def test_shadow_catches_effect_payload_mismatch_without_a_second_io(): + primary: Final = ScriptedCore(ResolveEndUser(True, lambda: Allow(WithinModelBudgets()))) + reference: Final = ScriptedCore(ResolveEndUser(False, lambda: Allow(WithinModelBudgets()))) + handler: Final = RecordingHandler() + with pytest.raises(ShadowMismatch, match="effect mismatch"): + await drive(ShadowCore(cast(AuthCore, primary), cast(AuthCore, reference)).check_model_budgets(), handler) + assert handler.calls == (("resolve_end_user", True),) + + +@pytest.mark.asyncio +async def test_shadow_catches_outcome_mismatch_without_exposing_credentials(): + with pytest.raises(ShadowMismatch) as error: + ShadowCore( + cast(AuthCore, ScriptedCore(Allow(NoMasterKey("private-credential")))), + cast(AuthCore, ScriptedCore(Deny(KeyBlocked()))), + ).check_model_budgets() + assert str(error.value) == "auth shadow mismatch at Allow" + + +@pytest.mark.asyncio +async def test_native_requests_keep_independent_continuations(): + core: Final = RustCore() + first: Final = core.start_resolved_key(replace(FACTS, token="one"), SETTINGS, NOW) + second: Final = core.start_resolved_key(replace(FACTS, token="two"), SETTINGS, NOW) + blocked: Final = RecordingHandler(key=replace(PLAIN_KEY, blocked=True)) + plain: Final = RecordingHandler() + assert await drive(second, blocked) == Deny(KeyBlocked()) + assert isinstance(await drive(first, plain), Allow) + assert blocked.calls == (("inspect_key", None),) + assert plain.calls[-1] == ("load_key_context", None) + + +@pytest.mark.asyncio +async def test_effect_failure_after_native_entry_is_never_replayed(core: AuthCore): + failure: Final = RuntimeError("effect failed") + + class FailingHandler(RecordingHandler): + async def enforce_model_access(self) -> None: + self._record("enforce_model_access") + raise failure + + handler: Final = FailingHandler() + with pytest.raises(RuntimeError) as error: + await drive(core.start_resolved_key(FACTS, SETTINGS, NOW), handler) + assert error.value is failure + assert handler.calls == (("inspect_key", None), ("enforce_model_access", None)) 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 19f9d0dcc73..e8ac0e96760 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -15,10 +15,12 @@ from litellm._logging import verbose_proxy_logger import json import logging import litellm +from litellm._auth.python_core import PythonCore from litellm.proxy.auth.user_api_key_auth import ( + _enforce_model_budgets, + _ModelBudgetEffects, user_api_key_auth, UserAPIKeyAuth, - get_api_key_from_custom_header, ) from fastapi import WebSocket, HTTPException, status @@ -26,6 +28,23 @@ from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from tests._master_key import MASTER_KEY +async def _check_user_model_budgets(valid_token: UserAPIKeyAuth, limiter: object, model: str) -> None: + request: Final = MagicMock() + request.scope = {} + await _enforce_model_budgets( + PythonCore().check_user_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data={"model": model}, + valid_token=valid_token, + route="/chat/completions", + limiter=limiter, + llm_router=None, + llm_model_list=None, + ), + ) + + class Request: def __init__(self, client_ip: Optional[str] = None, headers: Optional[dict] = None): self.client = MagicMock() @@ -506,50 +525,6 @@ async def test_auth_not_connected_to_db(): assert valid_token.token == "failed-to-connect-to-db" -def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api_key): - verbose_proxy_logger.setLevel(logging.DEBUG) - request = MagicMock(spec=Request) - request.headers = headers - api_key = get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=custom_header_name) - assert api_key == expected_api_key - - -def test_get_api_key_from_custom_header_bearer_token(): - token = "sk-" + "1" * 8 - _assert_api_key_from_custom_header( - headers={"x-custom-api-key": f"Bearer {token}"}, - custom_header_name="x-custom-api-key", - expected_api_key=token, - ) - - -def test_get_api_key_from_custom_header_empty_value(): - _assert_api_key_from_custom_header( - headers={"x-custom-api-key": ""}, - custom_header_name="x-custom-api-key", - expected_api_key="", - ) - - -def test_get_api_key_from_custom_header_missing_header(): - _assert_api_key_from_custom_header( - headers={}, - custom_header_name="X-Custom-API-Key", - expected_api_key="", - ) - - -def test_get_api_key_from_custom_header_different_casing(): - token = "sk-" + "1" * 8 - _assert_api_key_from_custom_header( - headers={"X-CUSTOM-API-KEY": f"Bearer {token}"}, - custom_header_name="X-Custom-API-Key", - expected_api_key=token, - ) - - - - @pytest.mark.parametrize( "user_role, auth_user_id, requested_user_id, expected_result", [ @@ -1426,11 +1401,12 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ stored on LiteLLM_UserTable, accepted by /user/new and /user/update, and read by nothing, so a user-level per-model budget never blocked anything. """ - from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget - calls = [] class _Limiter: + async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: + return True + 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: @@ -1442,11 +1418,7 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ user_id=user_id, user_model_max_budget=user_model_max_budget, ) - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=_Limiter(), - models=["gpt-4"], - ) + await _check_user_model_budgets(valid_token, _Limiter(), model="gpt-4") assert len(calls) == expected_calls if expected_calls: assert calls[0] == ("u-1", user_model_max_budget, "gpt-4") @@ -1588,7 +1560,6 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove nothing ever reads it, which looks enforced and is not. """ from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth - from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key from litellm.proxy.proxy_server import model_max_budget_limiter @@ -1616,18 +1587,11 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove if expect_refusal: with pytest.raises(litellm.BudgetExceededError) as exc: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=[model], - ) + await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) assert exc.value.entity_type == Litellm_EntityType.USER.value + assert exc.value.entity_id == user_id else: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=[model], - ) + await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) def test_jwt_path_enforces_the_user_model_budget_before_returning(): @@ -1645,7 +1609,7 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) def calls_before_each_return(node): @@ -1654,7 +1618,7 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning(): if isinstance(child, ast.Call): fn = child.func name = getattr(fn, "id", None) or getattr(fn, "attr", None) - if name == "_check_user_model_budget": + if name == "_enforce_model_budgets": seen_check.append(child.lineno) return seen_check @@ -1690,7 +1654,7 @@ def test_every_jwt_branch_carries_the_user_model_budget(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) assignments = [ @@ -1708,7 +1672,7 @@ def test_every_jwt_branch_carries_the_user_model_budget(): assert "auto_registered" in targets, ( f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}" ) - assert "valid_token" in targets, "the virtual-key path must carry the user's model budget" + assert "key" in targets, "the virtual-key path must carry the user's model budget" @pytest.mark.asyncio @@ -1815,8 +1779,10 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): whether the caller used a JWT or a virtual key, and makes that log a lie. Structural rather than behavioural on purpose: the defect is a call site - sitting outside a guard, and driving both auth paths to a zero-cost model - would prove it for the two requests exercised rather than for every site. + sitting outside a guard, and driving the JWT path to a zero-cost model + would prove it for the request exercised rather than for every site. The + virtual-key path's exemption is decided by the auth core and covered in + tests/unit/_auth/test_python_core.py. """ import ast import inspect @@ -1824,11 +1790,7 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): from litellm.proxy.auth import user_api_key_auth as auth_module - tree = ast.parse( - textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) - + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) - ) + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool: for parent in ast.walk(node): @@ -1848,13 +1810,13 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): calls = [ node for node in ast.walk(tree) - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget" + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_enforce_model_budgets" ] - assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}" + assert len(calls) == 1, f"expected the JWT call site, found {len(calls)}" unguarded = [c for c in calls if not guarded_by_skip(tree, c)] assert not unguarded, ( - f"{len(unguarded)} _check_user_model_budget call(s) run even when " + f"{len(unguarded)} _enforce_model_budgets call(s) run even when " "skip_budget_checks is set, so a zero-cost model is enforced on one auth path and not the other" ) @@ -1882,8 +1844,7 @@ def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag" budget_calls = ( - "_check_key_model_budget_with_fallback", - "_check_user_model_budget", + "_enforce_model_budgets", "is_end_user_within_model_budget", ) @@ -1972,7 +1933,7 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) # Half one: the shared block copies the user row's budget onto the token. @@ -2154,7 +2115,7 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods( with ( patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.auth.handle_jwt.JWTAuthManager.auth_builder", return_value=jwt_response), - patch("litellm.proxy.auth.user_api_key_auth._check_user_model_budget", user_model_budget_check), + patch("litellm.proxy.auth.user_api_key_auth._enforce_model_budgets", user_model_budget_check), ): result: Final = await user_api_key_auth_builder( request=mock_request, 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 8440d3d2ab2..8b9d682581a 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,10 +74,12 @@ from litellm.proxy.auth.auth_checks import ( cache_key_object, jwt_key_mapping_cache_key, ) +from litellm._auth.python_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 ( - _check_key_model_budget_with_fallback, + _enforce_model_budgets, + _ModelBudgetEffects, _check_team_member_model_budget_with_fallback, _ensure_litellm_received_at_on_request_state, _ensure_parent_otel_span_on_request_state, @@ -753,6 +755,61 @@ async def _authenticate_through_custom_auth() -> UserAPIKeyAuth: ) +async def _authenticate_master_key(request_data: dict[str, object]) -> UserAPIKeyAuth: + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder + + request: Final = Request(scope={"type": "http", "headers": []}) + request._url = URL(url="/chat/completions") + return await user_api_key_auth_builder( + request=request, + api_key="Bearer sk-master-key", + azure_api_key_header=None, + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + + +@pytest.mark.asyncio +async def test_master_key_request_does_not_inherit_previous_requests_end_user_limits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as _proxy_server_mod + + proxy_attrs: Final = _proxy_server_attrs_for_custom_auth(user_custom_auth=None) + for attr, val in proxy_attrs.items(): + monkeypatch.setattr(_proxy_server_mod, attr, val, raising=False) + monkeypatch.setattr(_proxy_server_mod, "user_api_key_cache", UserApiKeyCache()) + capped: Final = LiteLLM_EndUserTable( + user_id="eu-capped", + blocked=False, + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=5, rpm_limit=7), + ) + + async def end_user_object(end_user_id: str, **_: object) -> LiteLLM_EndUserTable | None: + return capped if end_user_id == "eu-capped" else None + + with patch("litellm.proxy.auth.user_api_key_auth.get_end_user_object", side_effect=end_user_object): + capped_request: Final = await _authenticate_master_key({"user": "eu-capped"}) + await asyncio.sleep(0.05) + next_request: Final = await _authenticate_master_key({}) + + assert (capped_request.end_user_id, capped_request.end_user_tpm_limit, capped_request.end_user_rpm_limit) == ( + "eu-capped", + 5, + 7, + ) + assert (next_request.end_user_id, next_request.end_user_tpm_limit, next_request.end_user_rpm_limit) == ( + None, + None, + None, + ) + assert next_request.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + + @pytest.mark.asyncio async def test_user_custom_auth_returning_one_shared_token_keeps_each_requests_end_user_apart( monkeypatch: pytest.MonkeyPatch, @@ -5487,8 +5544,6 @@ async def _run_centralized_checks_for_end_user( setattr(_proxy_server_mod, attribute, value) - - @pytest.mark.parametrize( ("make_token", "make_end_user", "gpt_4o_spend", "model", "route", "must_raise"), [ @@ -9540,10 +9595,43 @@ class TestJWTAuthUserEmail: assert result.user_email == "admin@example.com" -class TestCheckKeyModelBudgetWithFallback: - """`_check_key_model_budget_with_fallback` must reroute a request to the - first configured `budget_fallbacks` entry still within its own budget, - and only raise `BudgetExceededError` when no fallback is available.""" +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 + + limiter: Final = AsyncMock() + limiter.is_key_within_model_budget.side_effect = is_key_within_model_budget + return limiter + + +async def _check_key_model_budgets( + valid_token: UserAPIKeyAuth, + limiter: object, + request_data: dict[str, object], + request: object, + llm_router: litellm.Router | None = None, +) -> None: + await _enforce_model_budgets( + PythonCore().check_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=valid_token, + route="/chat/completions", + limiter=limiter, + llm_router=llm_router, + llm_model_list=None, + ), + ) + + +_KEY_MODEL_BUDGET: Final = {"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}} + + +class TestKeyModelBudgetFallback: + """An over-budget key reroutes to the first `budget_fallbacks` entry it may call and can still afford.""" def _make_request(self): request = MagicMock() @@ -9552,162 +9640,64 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_within_budget_does_not_reroute(self): - valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) - limiter = AsyncMock() - limiter.is_key_within_model_budget.return_value = True + valid_token = UserAPIKeyAuth( + token="test-key", model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} + ) request_data = {"model": "gpt-4o"} request = self._make_request() - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for(), request_data, request) assert request_data["model"] == "gpt-4o" - limiter.get_fallback_model_within_budget.assert_not_awaited() assert "parsed_body" not in request.scope @pytest.mark.asyncio - async def test_exceeded_budget_reroutes_to_fallback(self): + async def test_exceeded_budget_reroutes_to_fallback_visible_to_a_later_body_read(self): valid_token = UserAPIKeyAuth( token="test-key", + model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini", "claude-haiku"]}, ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for("gpt-4o"), request_data, request) assert request_data["model"] == "gpt-4o-mini" - limiter.get_fallback_model_within_budget.assert_awaited_once_with(user_api_key_dict=valid_token, model="gpt-4o") - # the rerouted model must be visible to a later, separate - # `_read_request_body` call on the same `request` (route handlers - # re-parse the body from this cache instead of reusing the dict). - cached_keys, cached_body = request.scope["parsed_body"] + _, cached_body = request.scope["parsed_body"] assert cached_body["model"] == "gpt-4o-mini" @pytest.mark.asyncio - async def test_raises_when_every_fallback_also_exceeded(self): - valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = None + @pytest.mark.parametrize( + ("token_fields", "over_budget"), + [ + ({"budget_fallbacks": {"gpt-4o": ["gpt-4o-mini"]}}, ("gpt-4o", "gpt-4o-mini")), + ({"models": ["gpt-4o"], "budget_fallbacks": {"gpt-4o": ["restricted-model"]}}, ("gpt-4o",)), + ( + { + "models": ["gpt-4o", "restricted-model"], + "team_id": "team-1", + "team_models": ["gpt-4o"], + "budget_fallbacks": {"gpt-4o": ["restricted-model"]}, + }, + ("gpt-4o",), + ), + ], + ids=["every_fallback_over_budget", "fallback_outside_key_models", "fallback_outside_team_models"], + ) + async def test_raises_the_original_budget_error_when_no_fallback_is_usable( + self, token_fields: dict[str, object], over_budget: tuple[str, ...] + ): + valid_token = UserAPIKeyAuth(token="test-key", model_max_budget=_KEY_MODEL_BUDGET, **token_fields) request_data = {"model": "gpt-4o"} - request = self._make_request() with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, + await _check_key_model_budgets( + valid_token, _key_over_budget_for(*over_budget), request_data, self._make_request() ) - assert exc_info.value is original_error - assert request_data["model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_raises_when_fallback_not_authorized(self): - """If the fallback model is within budget but the key is not allowed - to call it, the original BudgetExceededError must be raised instead - of rerouting to an unauthorized model.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o"], - budget_fallbacks={"gpt-4o": ["restricted-model"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = "restricted-model" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert exc_info.value is original_error - assert request_data["model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_reroute_succeeds_when_fallback_is_authorized(self): - """If the fallback model is within budget AND authorized, the request - must be rerouted to it.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "gpt-4o-mini"], - budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, - ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert request_data["model"] == "gpt-4o-mini" - - @pytest.mark.asyncio - async def test_raises_when_fallback_blocked_by_team_models(self): - """If the key allows the fallback but the team does not, the original - BudgetExceededError must be raised.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "restricted-model"], - team_id="team-1", - team_models=["gpt-4o"], - budget_fallbacks={"gpt-4o": ["restricted-model"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = "restricted-model" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert exc_info.value is original_error + assert exc_info.value.message == "over budget for gpt-4o" + assert (exc_info.value.current_cost, exc_info.value.max_budget) == (10, 5) assert request_data["model"] == "gpt-4o" @pytest.mark.asyncio @@ -9718,56 +9708,18 @@ class TestCheckKeyModelBudgetWithFallback: valid_token = UserAPIKeyAuth( token="test-key", models=["gpt-4o", "gpt-4o-mini"], + model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() request.scope["path_params"] = {"model": "gpt-4o"} - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for("gpt-4o"), request_data, request) assert request_data["model"] == "gpt-4o-mini" assert request.scope["path_params"]["model"] == "gpt-4o-mini" - @pytest.mark.asyncio - async def test_raises_when_model_from_nested_field(self): - """Budget fallback must not attempt rewrite when the checked model - came from a nested field (e.g. session.model) rather than the - top-level request_data['model'].""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "gpt-4o-mini"], - budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - request_data = {"session": {"model": "gpt-4o"}} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) - - assert exc_info.value is original_error - assert "model" not in request_data - class TestCheckTeamMemberModelBudgetWithFallback: @pytest.fixture(autouse=True) @@ -10786,13 +10738,7 @@ class TestCheckTeamMemberModelBudgetWithFallback: request_data: Final = {"model": "gpt-4o"} request: Final = self._request("gpt-4o") - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, limiter, request_data, request) assert request_data["model"] == "gpt-4o-mini" 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 21cbe8667ff..55beaf5254d 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 @@ -510,70 +510,6 @@ async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queue assert limiter.redis_increment_operation_queue == [] -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( - budget_limiter, -): - user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={}) - assert await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") is None - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_first_within_budget( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}}, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=1.0): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result == "gpt-4o-mini" - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_skips_exhausted_fallback( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={ - "gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}, - "claude-haiku": {"budget_limit": 100.0, "time_period": "1d"}, - }, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - - async def _spend_for_model(entity_type, entity_id, model, resolved): - return 150.0 if resolved.budget_model == "gpt-4o-mini" else 1.0 - - with patch.object( - budget_limiter, - "_get_spend_for_model_budget", - side_effect=_spend_for_model, - ): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result == "claude-haiku" - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_none_when_chain_exhausted( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={ - "gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}, - "claude-haiku": {"budget_limit": 100.0, "time_period": "1d"}, - }, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result is None - - @pytest.mark.asyncio async def test_async_log_success_event_skips_redis_push_without_redis(budget_limiter): """When dual_cache has no Redis backend, do not await _push_in_memory_increments_to_redis.""" diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index af8bc254e2c..a64dab0bf9b 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -51,7 +51,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_session_id_for_spend_log, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, - is_master_key, redact_logged_api_key, _redact_prompt_leaks_in_error_string, sanitize_error_information_for_spend_logs, @@ -3484,37 +3483,6 @@ def test_get_logging_payload_includes_request_duration_ms(): assert payload["request_duration_ms"] == 3000 -class TestIsMasterKey: - """Tests for _is_master_key handling None inputs without raising TypeError.""" - - def test_none_api_key_returns_false(self): - """Regression: _is_master_key(None, 'sk-master') should return False, not raise TypeError.""" - assert is_master_key(api_key=None, _master_key="sk-master-key") is False - - def test_none_master_key_returns_false(self): - assert is_master_key(api_key="sk-some-key", _master_key=None) is False - - def test_both_none_returns_false(self): - assert is_master_key(api_key=None, _master_key=None) is False - - def test_matching_key_returns_true(self): - assert is_master_key(api_key="sk-master", _master_key="sk-master") is True - - def test_non_matching_key_returns_false(self): - assert is_master_key(api_key="sk-other", _master_key="sk-master") is False - - def test_master_key_hash_is_rejected(self): - """ - ``_is_master_key`` must not accept ``hash_token(master_key)`` as - equivalent to the raw master key — only the raw value matches. - """ - from litellm.proxy.utils import hash_token - - master = "sk-master-key-123" - hashed = hash_token(master) - assert is_master_key(api_key=hashed, _master_key=master) is False - - def test_sanitize_request_body_strips_secret_fields(): """ secret_fields contains raw HTTP headers (including Authorization Bearer