feat(auth): add the Rust AuthCore decision boundary

.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

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"

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" }

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

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};

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<PyAny>, Option<NativeAuthContinuation>);
+
+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<Step, core::Error>) -> PyResult<ProjectedStep> {
+    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<Option<Continuation>>,
+}
+
+#[pymethods]
+impl NativeAuthContinuation {
+    fn resume(&self, py: Python<'_>, answer: &Bound<'_, PyAny>) -> PyResult<ProjectedStep> {
+        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<ProjectedStep> {
+    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<ProjectedStep> {
+    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<Py<PyAny>> {
+    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<Py<PyAny>> {
+    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<String> {
+    credential::scheme_credential(value)
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_litellm_header_token(value: &str) -> String {
+    credential::litellm_header_token(value)
+}

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

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<SecretValue>,
+    pub litellm_header: Option<SecretValue>,
+    pub authorization: Option<SecretValue>,
+    pub api_key: Option<SecretValue>,
+    pub x_api_key: Option<SecretValue>,
+    pub google: Option<SecretValue>,
+    pub apim: Option<SecretValue>,
+    pub query: Option<SecretValue>,
+    pub pass_through: Option<SecretValue>,
+}
+
+#[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<SecretValue>,
+    pub api_key: Option<SecretValue>,
+    pub subprotocols: Option<SecretValue>,
+}
+
+#[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<String> {
+    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<SelectedKey> {
+    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"));
+    }
+}

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,
+}

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<Step, Error> {
+    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<PresentedKey>| {
+                    after_presented_key(presented, token, hashed_token, facts, settings, now)
+                },
+            )
+        },
+    )
+}
+
+pub fn start_resolved_key(
+    facts: RequestFacts,
+    settings: Settings,
+    now: DateTime,
+) -> Result<Step, Error> {
+    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<PresentedKey>,
+    token: SecretValue,
+    hashed_token: String,
+    facts: RequestFacts,
+    settings: Settings,
+    now: DateTime,
+) -> Result<Step, Error> {
+    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, Error> {
+    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<Step, Error> {
+    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<OwnerFacts>| {
+            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, Error> {
+    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<TeamMemberBudget>| {
+                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<Step, Error> {
+    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<Step, Error> {
+    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<f64>) -> Result<Step, Error> {
+    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<f64>) -> Result<Step, Error> {
+    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<f64>| match spend {
+                None => grant(key, throttle),
+                Some(spend) => Step::effect(Effect::AlertProxyBudget { spend }, move |(): ()| {
+                    grant(key, throttle)
+                }),
+            },
+        )
+    })
+}
+
+fn grant(key: KeyFacts, throttle: Option<f64>) -> Result<Step, Error> {
+    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"));
+    }
+}

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,
+};

litellm-rust/crates/request-auth-core/src/model_budgets.rs

@@ -0,0 +1,266 @@
+use crate::{Error, protocol::*};
+
+fn allowed() -> Result<Step, Error> {
+    Step::allow(Grant::WithinModelBudgets)
+}
+
+pub fn check_model_budgets() -> Result<Step, Error> {
+    key_budgets(Box::new(allowed))
+}
+
+pub fn check_user_model_budgets() -> Result<Step, Error> {
+    user_budgets(Box::new(allowed))
+}
+
+pub(crate) fn key_budgets(then: Then) -> Result<Step, Error> {
+    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, Error> {
+    Step::effect(
+        Effect::ReadModelBudgetScope,
+        move |scope: ModelBudgetScope| {
+            user_models(scope.user_has_model_budgets, scope.models, then)
+        },
+    )
+}
+
+fn key_models(
+    fallbacks: std::collections::BTreeMap<String, Vec<String>>,
+    models: Vec<String>,
+    rewritable: Option<String>,
+    then: Then,
+) -> Result<Step, Error> {
+    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<ModelBudgetExceeded>| {
+            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<dyn FnOnce(String) -> Result<Step, Error> + Send + Sync>;
+
+fn first_affordable(
+    candidates: Vec<String>,
+    exceeded: ModelBudgetExceeded,
+    found: Found,
+) -> Result<Step, Error> {
+    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<ModelBudgetExceeded>| {
+            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<String>, then: Then) -> Result<Step, Error> {
+    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<ModelBudgetExceeded>| match exceeded {
+            Some(exceeded) => Step::deny(exceeded),
+            None => user_models(enabled, rest, then),
+        },
+    )
+}
+
+pub fn check_team_member_model_budgets() -> Result<Step, Error> {
+    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<String>,
+    rewritable: Option<String>,
+    can_fall_back: bool,
+) -> Result<Step, Error> {
+    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<ModelBudgetExceeded>| {
+            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<FallbackTarget>| {
+                    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<FallbackTarget>,
+    exceeded: ModelBudgetExceeded,
+    found: Found,
+) -> Result<Step, Error> {
+    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")
+        );
+    }
+}

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<i64>,
+    pub timezone_id: u64,
+    pub reference: usize,
+}
+
+impl DateTime {
+    pub fn before(&self, other: &Self) -> Result<bool, Error> {
+        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<SecretValue>,
+    pub is_llm_api_route: bool,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Settings {
+    pub master_key: Option<SecretValue>,
+    pub allow_user_auth: bool,
+    pub has_db: bool,
+    pub master_key_only_routes: BTreeSet<String>,
+    pub proxy_max_budget: f64,
+}
+
+#[derive(Clone, veil::Redact, Serialize)]
+#[serde(tag = "kind")]
+pub enum Grant {
+    NoMasterKey {
+        #[redact(with = "[REDACTED]")]
+        token: Option<String>,
+    },
+    UserAuthRoute,
+    MasterKey,
+    CachedAdminKey,
+    VirtualKey {
+        budget_throttle_pct: Option<f64>,
+    },
+    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<String>,
+        entity_id: Option<String>,
+    },
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct ModelBudgetExceeded {
+    pub message: String,
+    pub current_cost: f64,
+    pub max_budget: f64,
+    pub entity_type: Option<String>,
+    pub entity_id: Option<String>,
+}
+
+impl From<ModelBudgetExceeded> 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<DateTime> },
+    NonAdminKey,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct KeyFacts {
+    pub has_token: bool,
+    pub blocked: bool,
+    pub expires: Option<DateTime>,
+    pub spend: f64,
+    pub soft_budget: Option<f64>,
+    pub max_budget: Option<f64>,
+    pub team_member_spend: Option<f64>,
+    pub user_id: Option<String>,
+    pub team_id: Option<String>,
+    pub budget_throttle_pct: Option<f64>,
+}
+
+#[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<String>,
+    pub rewritable_model: Option<String>,
+    pub key_has_model_budgets: bool,
+    pub key_budget_fallbacks: BTreeMap<String, Vec<String>>,
+    pub user_has_model_budgets: bool,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct TeamMemberModelBudgetScope {
+    pub models: Vec<String>,
+    pub rewritable_model: Option<String>,
+    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<PresentedKey>),
+    Key(KeyFacts),
+    Owner(Option<OwnerFacts>),
+    TeamBudget(Option<TeamMemberBudget>),
+    Spend(f64),
+    ProxySpend(Option<f64>),
+    ModelScope(ModelBudgetScope),
+    Exceeded(Option<ModelBudgetExceeded>),
+    TeamScope(TeamMemberModelBudgetScope),
+    Targets(Vec<FallbackTarget>),
+}
+
+pub trait FromAnswer: Sized {
+    fn from_answer(answer: Answer) -> Result<Self, Error>;
+}
+
+macro_rules! answer {
+    ($type:ty, $variant:ident) => {
+        impl FromAnswer for $type {
+            fn from_answer(answer: Answer) -> Result<Self, Error> {
+                match answer {
+                    Answer::$variant(value) => Ok(value),
+                    _ => Err(Error::InvalidAnswer),
+                }
+            }
+        }
+    };
+}
+answer!(bool, Bool);
+answer!(Option<PresentedKey>, Presented);
+answer!(KeyFacts, Key);
+answer!(Option<OwnerFacts>, Owner);
+answer!(Option<TeamMemberBudget>, TeamBudget);
+answer!(f64, Spend);
+answer!(Option<f64>, ProxySpend);
+answer!(ModelBudgetScope, ModelScope);
+answer!(Option<ModelBudgetExceeded>, Exceeded);
+answer!(TeamMemberModelBudgetScope, TeamScope);
+answer!(Vec<FallbackTarget>, Targets);
+
+impl FromAnswer for () {
+    fn from_answer(answer: Answer) -> Result<Self, Error> {
+        match answer {
+            Answer::Unit => Ok(()),
+            _ => Err(Error::InvalidAnswer),
+        }
+    }
+}
+
+pub type Continuation = Box<dyn FnOnce(Answer) -> Result<Step, Error> + Send + Sync>;
+pub type Then = Box<dyn FnOnce() -> Result<Step, Error> + Send + Sync>;
+
+pub struct Step {
+    pub node: Node,
+    pub resume: Option<Continuation>,
+}
+
+impl Step {
+    pub fn allow(grant: Grant) -> Result<Self, Error> {
+        Ok(Self {
+            node: Node::Outcome(Outcome::Allow { grant }),
+            resume: None,
+        })
+    }
+    pub fn deny(denial: impl Into<Denial>) -> Result<Self, Error> {
+        Ok(Self {
+            node: Node::Outcome(Outcome::Deny {
+                denial: denial.into(),
+            }),
+            resume: None,
+        })
+    }
+    pub fn effect<T: FromAnswer>(
+        effect: Effect,
+        resume: impl FnOnce(T) -> Result<Self, Error> + Send + Sync + 'static,
+    ) -> Result<Self, Error> {
+        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<i64>,
+        #[case] zone: u64,
+        #[case] other_wall: i64,
+        #[case] other_offset: Option<i64>,
+        #[case] other_zone: u64,
+        #[case] expected: Result<bool, ()>,
+    ) {
+        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)
+        ));
+    }
+}

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

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

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: ...

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)

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())

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),
+    )

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))

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

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)

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),
+        )

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

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)

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,

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:

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,

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"})

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
 
 

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():

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: ...

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]
 
 

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

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

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

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

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"

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()

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))

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,

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"
 

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."""

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