refactor(auth): move end-user defaults and identity stages into Rust

litellm-rust/crates/python-bridge/src/lib.rs

@@ -44,7 +44,10 @@ mod _native {
         request_auth_budget_window_start, request_auth_budgets, request_auth_common_budget,
         request_auth_context_end_user, request_auth_context_member, request_auth_context_models,
         request_auth_context_percentage, request_auth_context_temporary,
-        request_auth_context_throttle, request_auth_context_token,
+        request_auth_context_throttle, request_auth_context_token, request_auth_end_user_apply,
+        request_auth_end_user_assignments, request_auth_end_user_default_id,
+        request_auth_end_user_lookup, request_auth_end_user_normalize_id,
+        request_auth_end_user_registry, request_auth_end_user_validation,
         request_auth_litellm_header_token, request_auth_member_limit,
         request_auth_model_budget_check, request_auth_model_budget_query,
         request_auth_model_budget_resolve, request_auth_scheme_credential,

litellm-rust/crates/python-bridge/src/request_auth.rs

@@ -315,3 +315,87 @@ pub(crate) fn request_auth_context_percentage(
     let facts: core::budget_context::NumericFacts = from_py(facts)?;
     Ok(py.detach(move || core::budget_context::percentage(&facts)))
 }
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_normalize_id(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::IdFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::normalize_id(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_default_id(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::DefaultFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::default_id(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_apply(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::ApplyFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::apply_default(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_registry(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::RegistryFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::registry(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_assignments(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::AssignmentFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::assignments(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_lookup(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::LookupFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::lookup(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_end_user_validation(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::end_user_context::ValidationFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::end_user_context::validation(&facts)),
+    )
+}

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

@@ -0,0 +1,428 @@
+use serde::{Deserialize, Serialize};
+use strum::VariantArray;
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct IdFacts {
+    pub value: Option<String>,
+}
+
+pub fn normalize_id(facts: &IdFacts) -> Option<String> {
+    facts.value.clone().filter(|value| !value.is_empty())
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct DefaultFacts {
+    pub prisma_present: bool,
+    pub explicit_id: Option<String>,
+    pub global_read: bool,
+    pub global_id: Option<String>,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct DefaultDecision {
+    pub read_global: bool,
+    pub budget_id: Option<String>,
+}
+
+pub fn default_id(facts: &DefaultFacts) -> DefaultDecision {
+    if facts.explicit_id.is_none() && !facts.global_read {
+        return DefaultDecision {
+            read_global: true,
+            budget_id: None,
+        };
+    }
+    if !facts.prisma_present {
+        return DefaultDecision {
+            read_global: false,
+            budget_id: None,
+        };
+    }
+    DefaultDecision {
+        read_global: false,
+        budget_id: facts
+            .explicit_id
+            .clone()
+            .or_else(|| facts.global_id.clone()),
+    }
+}
+
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct ApplyFacts {
+    pub own_budget_id: Option<String>,
+    pub own_budget_present: bool,
+    pub key_id: Option<String>,
+    pub global_read: bool,
+    pub global_id: Option<String>,
+    pub default_read: bool,
+    pub default_present: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum ApplyDecision {
+    Keep,
+    ReadGlobal,
+    ReadDefault,
+    Copy,
+}
+
+pub fn apply_default(facts: &ApplyFacts) -> ApplyDecision {
+    if facts.own_budget_id.is_some() && facts.own_budget_present {
+        return ApplyDecision::Keep;
+    }
+    if facts.key_id.is_none() {
+        if !facts.global_read {
+            return ApplyDecision::ReadGlobal;
+        }
+        if facts.global_id.is_none() {
+            return ApplyDecision::Keep;
+        }
+    }
+    if !facts.default_read {
+        return ApplyDecision::ReadDefault;
+    }
+    if facts.default_present {
+        ApplyDecision::Copy
+    } else {
+        ApplyDecision::Keep
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct RegistryFacts {
+    pub end_user_id: String,
+    pub global_id: Option<String>,
+    pub key_id: Option<String>,
+    pub validate_id: bool,
+    pub token_max_present: bool,
+    pub registry_read: bool,
+    pub registry: Option<Vec<String>>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum RegistryDecision {
+    Read,
+    Restricted,
+    Unrestricted,
+}
+
+pub fn registry(facts: &RegistryFacts) -> RegistryDecision {
+    if facts.global_id.is_some()
+        || facts.key_id.is_some()
+        || facts.validate_id
+        || facts.token_max_present
+    {
+        return RegistryDecision::Restricted;
+    }
+    if !facts.registry_read {
+        return RegistryDecision::Read;
+    }
+    if facts
+        .registry
+        .as_ref()
+        .is_some_and(|ids| !ids.contains(&facts.end_user_id))
+    {
+        RegistryDecision::Unrestricted
+    } else {
+        RegistryDecision::Restricted
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum Field {
+    EndUserId,
+    Tpm,
+    Rpm,
+    Tpd,
+    MaxBudget,
+    Models,
+    Region,
+}
+
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct Fields {
+    pub tpm: bool,
+    pub rpm: bool,
+    pub tpd: bool,
+    pub max_budget: bool,
+    pub models: bool,
+    pub region: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum AssignmentMode {
+    Params,
+    Token,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct AssignmentFacts {
+    pub mode: AssignmentMode,
+    pub fields: Fields,
+}
+
+pub fn assignments(facts: &AssignmentFacts) -> Vec<Field> {
+    let fields = &facts.fields;
+    match facts.mode {
+        AssignmentMode::Params => [
+            (Field::Tpm, fields.tpm),
+            (Field::Rpm, fields.rpm),
+            (Field::Tpd, fields.tpd),
+            (Field::MaxBudget, fields.max_budget),
+            (Field::Models, fields.models),
+        ]
+        .into_iter()
+        .filter_map(|(field, present)| present.then_some(field))
+        .collect(),
+        AssignmentMode::Token => [
+            (Field::EndUserId, true),
+            (Field::Tpm, fields.tpm),
+            (Field::Rpm, fields.rpm),
+            (Field::Tpd, fields.tpd),
+            (Field::Region, fields.region),
+            (Field::Models, fields.models),
+        ]
+        .into_iter()
+        .filter_map(|(field, present)| present.then_some(field))
+        .collect(),
+    }
+}
+
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct LookupFacts {
+    pub own_budget_present: bool,
+    pub key_id: Option<String>,
+    pub end_user_read: bool,
+    pub end_user_present: bool,
+    pub global_read: bool,
+    pub global_id: Option<String>,
+    pub default_read: bool,
+    pub default_present: bool,
+    pub token_max_present: bool,
+    pub applied: bool,
+    pub read_failed: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum LookupAction {
+    ReadEndUser,
+    ReadGlobal,
+    ReadDefault,
+    ApplyOwn,
+    ApplyDefault,
+    Done,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct LookupDecision {
+    pub action: LookupAction,
+    pub set_max_budget: bool,
+    pub project_budget: bool,
+}
+
+pub fn lookup(facts: &LookupFacts) -> LookupDecision {
+    let action = if facts.read_failed || facts.applied {
+        LookupAction::Done
+    } else if !facts.end_user_read {
+        LookupAction::ReadEndUser
+    } else if facts.end_user_present {
+        LookupAction::ApplyOwn
+    } else if facts.key_id.is_none() && !facts.global_read {
+        LookupAction::ReadGlobal
+    } else if facts.key_id.is_none() && facts.global_id.is_none() {
+        LookupAction::Done
+    } else if !facts.default_read {
+        LookupAction::ReadDefault
+    } else if facts.default_present {
+        LookupAction::ApplyDefault
+    } else {
+        LookupAction::Done
+    };
+    LookupDecision {
+        project_budget: action == LookupAction::ApplyDefault
+            || (action == LookupAction::ApplyOwn && facts.own_budget_present),
+        set_max_budget: facts.applied
+            && !facts.read_failed
+            && !facts.end_user_present
+            && facts.default_present
+            && !facts.token_max_present,
+        action,
+    }
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum ValidationMode {
+    #[default]
+    Resolve,
+    Exists,
+}
+
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct ValidationFacts {
+    pub mode: ValidationMode,
+    pub end_user_id: Option<String>,
+    pub enabled: bool,
+    pub prisma_present: bool,
+    pub global_id: Option<String>,
+    pub key_id: Option<String>,
+    pub cache_read: bool,
+    pub cached: Option<String>,
+    pub end_user_read: bool,
+    pub end_user_present: bool,
+    pub user_read: bool,
+    pub user_present: bool,
+    pub cache_written: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)]
+pub enum ValidationRead {
+    Cache,
+    EndUser,
+    User,
+    WriteCache,
+}
+
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct ValidationDecision {
+    pub read: Option<ValidationRead>,
+    pub end_user_id: Option<String>,
+    pub valid: bool,
+    pub cache_value: Option<String>,
+    pub cache_ttl: u64,
+    pub user_email: Option<String>,
+}
+
+fn validation_result(facts: &ValidationFacts, valid: bool) -> ValidationDecision {
+    let retain = valid
+        || facts.key_id.is_some()
+        || facts.global_id.as_ref().is_some_and(|id| !id.is_empty());
+    ValidationDecision {
+        end_user_id: facts.end_user_id.clone().filter(|_| retain),
+        valid,
+        ..ValidationDecision::default()
+    }
+}
+
+pub fn validation(facts: &ValidationFacts) -> ValidationDecision {
+    let Some(id) = &facts.end_user_id else {
+        return ValidationDecision::default();
+    };
+    if facts.mode == ValidationMode::Resolve {
+        if !facts.enabled || !facts.prisma_present {
+            return ValidationDecision {
+                end_user_id: Some(id.clone()),
+                ..ValidationDecision::default()
+            };
+        }
+        if !facts.cache_read {
+            return ValidationDecision {
+                read: Some(ValidationRead::Cache),
+                ..ValidationDecision::default()
+            };
+        }
+        match facts.cached.as_deref() {
+            Some("valid") => return validation_result(facts, true),
+            Some("invalid") => return validation_result(facts, false),
+            _ => {}
+        }
+    }
+    if !facts.end_user_read {
+        return ValidationDecision {
+            read: Some(ValidationRead::EndUser),
+            ..ValidationDecision::default()
+        };
+    }
+    if !facts.end_user_present && !facts.user_read {
+        return ValidationDecision {
+            read: Some(ValidationRead::User),
+            user_email: id.contains('@').then(|| id.clone()),
+            ..ValidationDecision::default()
+        };
+    }
+    let valid = facts.end_user_present || facts.user_present;
+    if facts.mode == ValidationMode::Resolve && !facts.cache_written {
+        return ValidationDecision {
+            read: Some(ValidationRead::WriteCache),
+            valid,
+            cache_value: Some(if valid { "valid" } else { "invalid" }.to_owned()),
+            cache_ttl: if valid { 300 } else { 60 },
+            ..ValidationDecision::default()
+        };
+    }
+    validation_result(facts, valid)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::absent(None, None)]
+    #[case::empty(Some(""), None)]
+    #[case::space(Some(" "), Some(" "))]
+    #[case::id(Some("budget"), Some("budget"))]
+    fn normalizes_metadata_id(#[case] value: Option<&str>, #[case] expected: Option<&str>) {
+        assert_eq!(
+            normalize_id(&IdFacts {
+                value: value.map(str::to_owned)
+            })
+            .as_deref(),
+            expected
+        );
+    }
+
+    #[rstest]
+    #[case::explicit(Some("key"), Some("global"), Some("key"))]
+    #[case::empty_explicit(Some(""), Some("global"), Some(""))]
+    #[case::global(None, Some("global"), Some("global"))]
+    #[case::absent(None, None, None)]
+    fn selects_default(
+        #[case] explicit: Option<&str>,
+        #[case] global: Option<&str>,
+        #[case] expected: Option<&str>,
+    ) {
+        let facts = DefaultFacts {
+            prisma_present: true,
+            explicit_id: explicit.map(str::to_owned),
+            global_read: true,
+            global_id: global.map(str::to_owned),
+        };
+        assert_eq!(default_id(&facts).budget_id.as_deref(), expected);
+    }
+
+    #[rstest]
+    #[case::none(None, None, None)]
+    #[case::empty_global(Some(""), None, None)]
+    #[case::global(Some("budget"), None, Some("missing"))]
+    #[case::key(None, Some(""), Some("missing"))]
+    fn retains_unknown_ids(
+        #[case] global: Option<&str>,
+        #[case] key: Option<&str>,
+        #[case] expected: Option<&str>,
+    ) {
+        let facts = ValidationFacts {
+            end_user_id: Some("missing".to_owned()),
+            enabled: true,
+            prisma_present: true,
+            global_id: global.map(str::to_owned),
+            key_id: key.map(str::to_owned),
+            cache_read: true,
+            cached: Some("invalid".to_owned()),
+            ..ValidationFacts::default()
+        };
+        assert_eq!(validation(&facts).end_user_id.as_deref(), expected);
+    }
+
+    #[rstest]
+    #[case::missing(AssignmentMode::Params, vec![])]
+    #[case::token_identity(AssignmentMode::Token, vec![Field::EndUserId])]
+    fn absent_fields(#[case] mode: AssignmentMode, #[case] expected: Vec<Field>) {
+        assert_eq!(
+            assignments(&AssignmentFacts {
+                mode,
+                fields: Fields::default()
+            }),
+            expected
+        );
+    }
+}

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

@@ -5,6 +5,7 @@ pub mod budget_context;
 pub mod budget_pricing;
 pub mod common_budgets;
 pub mod credential;
+pub mod end_user_context;
 mod error;
 mod key;
 pub mod model_budget_data;

litellm/_auth/end_user_context.py

@@ -0,0 +1,171 @@
+from dataclasses import dataclass
+from typing import Final, Literal, TypeAlias
+
+from pydantic import TypeAdapter
+
+from litellm._auth.failure import native_boundary
+from litellm.rust_bridge import _native
+
+
+@dataclass(frozen=True, slots=True)
+class IdFacts:
+    value: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class DefaultFacts:
+    prisma_present: bool
+    explicit_id: str | None
+    global_id: str | None
+    global_read: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class DefaultDecision:
+    read_global: bool
+    budget_id: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class ApplyFacts:
+    own_budget_id: str | None
+    own_budget_present: bool
+    key_id: str | None
+    global_read: bool = False
+    global_id: str | None = None
+    default_read: bool = False
+    default_present: bool = False
+
+
+ApplyDecision: TypeAlias = Literal["Keep", "ReadGlobal", "ReadDefault", "Copy"]
+
+
+@dataclass(frozen=True, slots=True)
+class RegistryFacts:
+    end_user_id: str
+    global_id: str | None
+    key_id: str | None
+    validate_id: bool
+    token_max_present: bool
+    registry_read: bool = False
+    registry: tuple[str, ...] | None = None
+
+
+RegistryDecision: TypeAlias = Literal["Read", "Restricted", "Unrestricted"]
+Field: TypeAlias = Literal["EndUserId", "Tpm", "Rpm", "Tpd", "MaxBudget", "Models", "Region"]
+
+
+@dataclass(frozen=True, slots=True)
+class Fields:
+    tpm: bool = False
+    rpm: bool = False
+    tpd: bool = False
+    max_budget: bool = False
+    models: bool = False
+    region: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class AssignmentFacts:
+    mode: Literal["Params", "Token"]
+    fields: Fields
+
+
+@dataclass(frozen=True, slots=True)
+class LookupFacts:
+    key_id: str | None
+    own_budget_present: bool = False
+    end_user_read: bool = False
+    end_user_present: bool = False
+    global_read: bool = False
+    global_id: str | None = None
+    default_read: bool = False
+    default_present: bool = False
+    token_max_present: bool = False
+    applied: bool = False
+    read_failed: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class LookupDecision:
+    action: Literal["ReadEndUser", "ReadGlobal", "ReadDefault", "ApplyOwn", "ApplyDefault", "Done"]
+    set_max_budget: bool
+    project_budget: bool
+
+
+@dataclass(frozen=True, slots=True)
+class ValidationFacts:
+    mode: Literal["Resolve", "Exists"]
+    end_user_id: str | None
+    enabled: bool = False
+    prisma_present: bool = False
+    global_id: str | None = None
+    key_id: str | None = None
+    cache_read: bool = False
+    cached: str | None = None
+    end_user_read: bool = False
+    end_user_present: bool = False
+    user_read: bool = False
+    user_present: bool = False
+    cache_written: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class ValidationDecision:
+    read: Literal["Cache", "EndUser", "User", "WriteCache"] | None
+    end_user_id: str | None
+    valid: bool
+    cache_value: str | None
+    cache_ttl: int
+    user_email: str | None
+
+
+_ID: Final = TypeAdapter(IdFacts)
+_STRING: Final = TypeAdapter[str | None](str | None)
+_DEFAULT: Final = TypeAdapter(DefaultFacts)
+_DEFAULT_RESULT: Final = TypeAdapter(DefaultDecision)
+_APPLY: Final = TypeAdapter(ApplyFacts)
+_APPLY_RESULT: Final = TypeAdapter[ApplyDecision](ApplyDecision)
+_REGISTRY: Final = TypeAdapter(RegistryFacts)
+_REGISTRY_RESULT: Final = TypeAdapter[RegistryDecision](RegistryDecision)
+_ASSIGNMENTS: Final = TypeAdapter(AssignmentFacts)
+_FIELDS: Final = TypeAdapter(tuple[Field, ...])
+_LOOKUP: Final = TypeAdapter(LookupFacts)
+_LOOKUP_RESULT: Final = TypeAdapter(LookupDecision)
+_VALIDATION: Final = TypeAdapter(ValidationFacts)
+_VALIDATION_RESULT: Final = TypeAdapter(ValidationDecision)
+
+
+@native_boundary
+def normalize_id(facts: IdFacts) -> str | None:
+    return _STRING.validate_python(_native.request_auth_end_user_normalize_id(_ID.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts
+
+
+@native_boundary
+def default_id(facts: DefaultFacts) -> DefaultDecision:
+    return _DEFAULT_RESULT.validate_python(_native.request_auth_end_user_default_id(_DEFAULT.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts
+
+
+@native_boundary
+def apply_default(facts: ApplyFacts) -> ApplyDecision:
+    return _APPLY_RESULT.validate_python(_native.request_auth_end_user_apply(_APPLY.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts
+
+
+@native_boundary
+def registry(facts: RegistryFacts) -> RegistryDecision:
+    return _REGISTRY_RESULT.validate_python(_native.request_auth_end_user_registry(_REGISTRY.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts
+
+
+@native_boundary
+def assignments(facts: AssignmentFacts) -> tuple[Field, ...]:
+    return _FIELDS.validate_python(_native.request_auth_end_user_assignments(_ASSIGNMENTS.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts
+
+
+@native_boundary
+def lookup(facts: LookupFacts) -> LookupDecision:
+    return _LOOKUP_RESULT.validate_python(_native.request_auth_end_user_lookup(_LOOKUP.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts
+
+
+@native_boundary
+def validation(facts: ValidationFacts) -> ValidationDecision:
+    return _VALIDATION_RESULT.validate_python(_native.request_auth_end_user_validation(_VALIDATION.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen typed facts

litellm/proxy/auth/budget_context_adapter.py

@@ -161,13 +161,13 @@ async def _end_user_row(
 
 
 async def resolve_default_end_user_budget(
-    prisma_client: PrismaClient,
+    prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     key_end_user_budget_id: str | None,
     parent_otel_span: Span | None = None,
 ) -> LiteLLM_BudgetTable | None:
     return await _end_user_row(
-        core.EndUserFacts("Default", False, None, prisma_client is not None, key_end_user_budget_id),  # pyright: ignore[reportUnnecessaryComparison]  # legacy callers can pass absent database clients at runtime
+        core.EndUserFacts("Default", False, None, prisma_client is not None, key_end_user_budget_id),
         EndUserIO(prisma_client, user_api_key_cache, parent_otel_span),
         None,
     )

litellm/proxy/auth/end_user_context_adapter.py

@@ -0,0 +1,406 @@
+from collections.abc import Mapping, MutableMapping
+from dataclasses import replace
+from types import MappingProxyType
+from typing import Final, assert_never
+
+from opentelemetry.trace import Span
+
+import litellm
+from litellm._auth import end_user_context as core
+from litellm._auth.failure import NativeAuthFailure
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_EndUserTable, UserAPIKeyAuth
+from litellm.proxy.auth.budget_context_adapter import mapping_value, resolve_default_end_user_budget
+from litellm.proxy.auth.end_user_context_io import EndUserIO, read_default_budget
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+
+_PARAM_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType(
+    {
+        "EndUserId": "end_user_id",
+        "Tpm": "end_user_tpm_limit",
+        "Rpm": "end_user_rpm_limit",
+        "Tpd": "end_user_tpd_limit",
+        "MaxBudget": "end_user_max_budget",
+        "Models": "end_user_model_max_budget",
+        "Region": "allowed_model_region",
+    }
+)
+_BUDGET_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType(
+    {
+        "Tpm": "tpm_limit",
+        "Rpm": "rpm_limit",
+        "Tpd": "tpd_limit",
+        "MaxBudget": "max_budget",
+        "Models": "model_max_budget",
+    }
+)
+
+
+def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None:
+    raw: Final = key_metadata.get("end_user_budget_id") if key_metadata is not None else None
+    return core.normalize_id(core.IdFacts(raw if isinstance(raw, str) else None))
+
+
+def select_default_id(facts: core.DefaultFacts) -> str | None:
+    decision: Final = core.default_id(facts)
+    if decision.read_global:
+        return select_default_id(replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id))
+    return decision.budget_id
+
+
+async def get_default_end_user_budget(
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None = None,
+    budget_id: str | None = None,
+) -> LiteLLM_BudgetTable | None:
+    selected: Final = select_default_id(core.DefaultFacts(prisma_client is not None, budget_id, None))
+    if selected is None:
+        return None
+    return await read_default_budget(selected, EndUserIO(prisma_client, user_api_key_cache, parent_otel_span))
+
+
+async def _apply_default(
+    facts: core.ApplyFacts, row: LiteLLM_EndUserTable, io: EndUserIO, default: LiteLLM_BudgetTable | None = None
+) -> LiteLLM_EndUserTable:
+    decision: Final = core.apply_default(facts)
+    match decision:
+        case "Keep":
+            return row
+        case "ReadGlobal":
+            return await _apply_default(
+                replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), row, io
+            )
+        case "ReadDefault":
+            budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span)
+            return await _apply_default(
+                replace(facts, default_read=True, default_present=budget is not None), row, io, budget
+            )
+        case "Copy":
+            assert default is not None
+            verbose_proxy_logger.debug("Applied default budget %s to end user %s", default.budget_id, row.user_id)
+            return row.model_copy(update=MappingProxyType({"litellm_budget_table": default}))
+        case _:
+            assert_never(decision)
+
+
+async def apply_default_budget_to_end_user(
+    end_user_obj: LiteLLM_EndUserTable,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None = None,
+    key_end_user_budget_id: str | None = None,
+) -> LiteLLM_EndUserTable:
+    return await _apply_default(
+        core.ApplyFacts(end_user_obj.budget_id, end_user_obj.litellm_budget_table is not None, key_end_user_budget_id),
+        end_user_obj,
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span),
+    )
+
+
+async def _registry(facts: core.RegistryFacts, io: EndUserIO) -> bool:
+    from litellm.proxy.auth import auth_checks
+
+    decision: Final = core.registry(facts)
+    match decision:
+        case "Restricted":
+            return False
+        case "Unrestricted":
+            return True
+        case "Read":
+            assert io.prisma is not None
+            ids: Final = await auth_checks._load_end_user_restricted_registry(  # pyright: ignore[reportPrivateUsage]  # raw cached registry IO
+                prisma_client=io.prisma, user_api_key_cache=io.cache
+            )
+            return await _registry(
+                replace(facts, registry_read=True, registry=tuple(ids) if ids is not None else None), io
+            )
+        case _:
+            assert_never(decision)
+
+
+async def end_user_is_known_unrestricted(
+    end_user_id: str,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    token_end_user_max_budget: float | None,
+    key_end_user_budget_id: str | None = None,
+) -> bool:
+    return await _registry(
+        core.RegistryFacts(
+            end_user_id,
+            litellm.max_end_user_budget_id,
+            key_end_user_budget_id,
+            litellm.validate_end_user_id_in_db,
+            token_end_user_max_budget is not None,
+        ),
+        EndUserIO(prisma_client, user_api_key_cache),
+    )
+
+
+def _budget_fields(row: LiteLLM_BudgetTable | None) -> core.Fields:
+    return core.Fields(
+        tpm=getattr(row, "tpm_limit", None) is not None,
+        rpm=getattr(row, "rpm_limit", None) is not None,
+        tpd=getattr(row, "tpd_limit", None) is not None,
+        max_budget=getattr(row, "max_budget", None) is not None,
+        models=getattr(row, "model_max_budget", None) is not None,
+    )
+
+
+def _budget_params(row: LiteLLM_BudgetTable | None) -> Mapping[str, object]:
+    selected: Final = core.assignments(core.AssignmentFacts("Params", _budget_fields(row)))
+    return MappingProxyType({_PARAM_NAMES[field]: getattr(row, _BUDGET_NAMES[field]) for field in selected})
+
+
+def apply_budget_limits_to_end_user_params(
+    end_user_params: MutableMapping[str, object],  # mutable-ok: legacy request-local output parameter
+    budget_info: LiteLLM_BudgetTable,
+    end_user_id: str | None,
+) -> None:
+    for name, value in _budget_params(budget_info).items():
+        end_user_params[name] = value  # rebind-ok: native-selected legacy request parameter projection
+    verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id)
+
+
+def _token_fields(end_user_params: Mapping[str, object]) -> tuple[core.Field, ...]:
+    fields: Final = core.Fields(
+        tpm=end_user_params.get("end_user_tpm_limit") is not None,
+        rpm=end_user_params.get("end_user_rpm_limit") is not None,
+        tpd=end_user_params.get("end_user_tpd_limit") is not None,
+        models=end_user_params.get("end_user_model_max_budget") is not None,
+        region=end_user_params.get("allowed_model_region") is not None,
+    )
+    return core.assignments(core.AssignmentFacts("Token", fields))
+
+
+def _project_token(
+    valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object], fields: tuple[core.Field, ...]
+) -> UserAPIKeyAuth:
+    for field in fields:
+        setattr(valid_token, _PARAM_NAMES[field], end_user_params.get(_PARAM_NAMES[field]))
+    return valid_token
+
+
+def update_valid_token_with_end_user_params(
+    valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object]
+) -> UserAPIKeyAuth:
+    return _project_token(valid_token, end_user_params, _token_fields(end_user_params))
+
+
+async def _validation(facts: core.ValidationFacts, io: EndUserIO) -> core.ValidationDecision:
+    from litellm.proxy.auth import auth_checks
+
+    decision: Final = core.validation(facts)
+    match decision.read:
+        case None:
+            return decision
+        case "Cache":
+            cached: Final[object] = await auth_checks._raw_cache(io.cache).async_get_cache(  # pyright: ignore[reportPrivateUsage]  # shared raw cache IO
+                key=f"end_user_validation:{facts.end_user_id}"
+            )
+            return await _validation(
+                replace(facts, cache_read=True, cached=cached if isinstance(cached, str) else None), io
+            )
+        case "EndUser":
+            try:
+                row: Final = await auth_checks.get_end_user_object(
+                    end_user_id=facts.end_user_id,
+                    prisma_client=io.prisma,
+                    user_api_key_cache=io.cache,
+                    parent_otel_span=io.span,
+                    proxy_logging_obj=io.logging,
+                    route=io.route,
+                )
+            except NativeAuthFailure:
+                raise
+            except Exception as error:  # noqa: BLE001  # preserve end-user validation lookup failure handling
+                verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", error)
+                return await _validation(replace(facts, end_user_read=True), io)
+            return await _validation(replace(facts, end_user_read=True, end_user_present=row is not None), io)
+        case "User":
+            assert facts.end_user_id is not None and io.prisma is not None
+            try:
+                user: Final = await auth_checks.get_user_object(
+                    user_id=facts.end_user_id,
+                    prisma_client=io.prisma,
+                    user_api_key_cache=io.cache,
+                    user_id_upsert=False,
+                    parent_otel_span=io.span,
+                    proxy_logging_obj=io.logging,
+                    check_db_only=False,
+                    user_email=decision.user_email,
+                )
+            except NativeAuthFailure:
+                raise
+            except Exception as error:  # noqa: BLE001  # preserve secondary identity lookup failure handling
+                verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", error)
+                return await _validation(replace(facts, user_read=True), io)
+            return await _validation(replace(facts, user_read=True, user_present=user is not None), io)
+        case "WriteCache":
+            await io.cache.async_set_cache(
+                key=f"end_user_validation:{facts.end_user_id}", value=decision.cache_value, ttl=decision.cache_ttl
+            )
+            return await _validation(replace(facts, cache_written=True), io)
+        case _:
+            assert_never(decision.read)
+
+
+async def resolve_and_validate_end_user_id(
+    raw_end_user_id: str | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None = None,
+    proxy_logging_obj: ProxyLogging | None = None,
+    route: str = "",
+    key_end_user_budget_id: str | None = None,
+) -> str | None:
+    decision: Final = await _validation(
+        core.ValidationFacts(
+            "Resolve",
+            raw_end_user_id,
+            litellm.validate_end_user_id_in_db,
+            prisma_client is not None,
+            litellm.max_end_user_budget_id,
+            key_end_user_budget_id,
+        ),
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route),
+    )
+    return decision.end_user_id
+
+
+async def end_user_id_exists_in_db(
+    end_user_id: str,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None = None,
+    proxy_logging_obj: ProxyLogging | None = None,
+    route: str = "",
+) -> bool:
+    decision: Final = await _validation(
+        core.ValidationFacts("Exists", end_user_id),
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route),
+    )
+    return decision.valid
+
+
+def _complete_lookup(
+    decision: core.LookupDecision,
+    token: UserAPIKeyAuth,
+    row: LiteLLM_EndUserTable | None,
+    default: LiteLLM_BudgetTable | None,
+) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]:
+    if decision.set_max_budget:
+        assert default is not None
+        try:
+            token.end_user_max_budget = default.max_budget  # rebind-ok: request-local native projection
+        except (NativeAuthFailure, litellm.BudgetExceededError):
+            raise
+        except Exception as error:  # noqa: BLE001  # preserve request-local setter failure handling
+            verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error)
+    return token, row
+
+
+async def _lookup(
+    facts: core.LookupFacts,
+    token: UserAPIKeyAuth,
+    io: EndUserIO,
+    row: LiteLLM_EndUserTable | None = None,
+    default: LiteLLM_BudgetTable | None = None,
+) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]:
+    from litellm.proxy.auth.auth_checks import get_end_user_object
+
+    decision: Final = core.lookup(facts)
+    match decision.action:
+        case "Done":
+            return _complete_lookup(decision, token, row, default)
+        case "ReadGlobal":
+            return await _lookup(
+                replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), token, io, row
+            )
+        case "ReadEndUser":
+            try:
+                end_user: Final = await get_end_user_object(
+                    end_user_id=token.end_user_id,
+                    prisma_client=io.prisma,
+                    user_api_key_cache=io.cache,
+                    parent_otel_span=io.span,
+                    proxy_logging_obj=io.logging,
+                    route=io.route,
+                    token_end_user_max_budget=token.end_user_max_budget,
+                    key_end_user_budget_id=facts.key_id,
+                )
+            except (NativeAuthFailure, litellm.BudgetExceededError):
+                raise
+            except Exception as error:  # noqa: BLE001  # preserve the legacy lookup IO failure path
+                verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error)
+                return await _lookup(replace(facts, read_failed=True), token, io, row)
+            return await _lookup(
+                replace(
+                    facts,
+                    end_user_read=True,
+                    end_user_present=end_user is not None,
+                    own_budget_present=end_user is not None and end_user.litellm_budget_table is not None,
+                ),
+                token,
+                io,
+                end_user,
+            )
+        case "ReadDefault":
+            try:
+                budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span)
+            except (NativeAuthFailure, litellm.BudgetExceededError):
+                raise
+            except Exception as error:  # noqa: BLE001  # preserve the legacy default lookup IO failure path
+                verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error)
+                return await _lookup(replace(facts, read_failed=True), token, io, row)
+            return await _lookup(
+                replace(facts, default_read=True, default_present=budget is not None), token, io, row, budget
+            )
+        case "ApplyOwn" | "ApplyDefault":
+            budget_row: Final = (
+                row.litellm_budget_table if decision.action == "ApplyOwn" and row is not None else default
+            )
+            params: Final = MappingProxyType(
+                {
+                    "end_user_id": token.end_user_id,
+                    "allowed_model_region": getattr(row, "allowed_model_region", None),
+                    **_budget_params(budget_row),
+                }
+            )
+            fields: Final = _token_fields(params)
+            try:
+                if decision.project_budget:
+                    verbose_proxy_logger.debug("Applied budget limits to end user %s", token.end_user_id or "")
+                _project_token(token, params, fields)
+            except (NativeAuthFailure, litellm.BudgetExceededError):
+                raise
+            except Exception as error:  # noqa: BLE001  # preserve request-local setter failure handling
+                verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error)
+                return await _lookup(replace(facts, read_failed=True), token, io, row)
+            return await _lookup(
+                replace(facts, applied=True, token_max_present=token.end_user_max_budget is not None),
+                token,
+                io,
+                row,
+                default,
+            )
+        case _:
+            assert_never(decision.action)
+
+
+async def lookup_end_user_and_apply_budget(
+    valid_token: UserAPIKeyAuth,
+    route: str,
+    parent_otel_span: Span | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    proxy_logging_obj: ProxyLogging | None,
+) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]:
+    key_id: Final = get_key_end_user_budget_id(mapping_value(getattr(valid_token, "metadata", None)))
+    return await _lookup(
+        core.LookupFacts(key_id),
+        valid_token,
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route),
+    )

litellm/proxy/auth/end_user_context_io.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from typing import Final
+
+from opentelemetry.trace import Span
+
+from litellm._auth.failure import NativeAuthFailure
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import LiteLLM_BudgetTable
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+from litellm.repositories.budget_repository import BudgetRepository
+
+
+@dataclass(frozen=True, slots=True)
+class EndUserIO:
+    prisma: PrismaClient | None
+    cache: UserApiKeyCache
+    span: Span | None = None
+    logging: ProxyLogging | None = None
+    route: str = ""
+
+
+async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTable | None:
+    from litellm.proxy.auth import auth_checks
+
+    assert io.prisma is not None
+    key: Final = f"default_end_user_budget:{budget_id}"
+    cached: Final = await io.cache.async_get_cache(key=key, model_type=LiteLLM_BudgetTable)
+    if cached is not None:
+        return cached
+    try:
+        table: Final = auth_checks._dictable_table(BudgetRepository(io.prisma), "budget")  # pyright: ignore[reportPrivateUsage]  # shared raw repository projection
+        record: Final = await table.find_unique(where={"budget_id": budget_id})
+        if record is None:
+            verbose_proxy_logger.warning(
+                "Default end user budget not found in database: %s", budget_id.replace("\r", "").replace("\n", "")
+            )
+            return None
+        row: Final = LiteLLM_BudgetTable.model_validate(record.dict())
+        await io.cache.async_set_cache(
+            key=key, value=row, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(io.cache)
+        )
+    except NativeAuthFailure:
+        raise
+    except Exception as error:  # noqa: BLE001  # preserve the default budget DB/cache IO failure contract
+        verbose_proxy_logger.error("Error fetching default end user budget: %s", error)
+        return None
+
+    else:
+        return row

litellm/rust_bridge/_native.pyi

@@ -826,3 +826,17 @@ def request_auth_context_temporary(facts: object) -> object: ...
 def request_auth_context_throttle(facts: object) -> object: ...
 
 def request_auth_context_percentage(facts: object) -> float | None: ...
+
+def request_auth_end_user_normalize_id(facts: object) -> object: ...
+
+def request_auth_end_user_default_id(facts: object) -> object: ...
+
+def request_auth_end_user_apply(facts: object) -> object: ...
+
+def request_auth_end_user_registry(facts: object) -> object: ...
+
+def request_auth_end_user_assignments(facts: object) -> object: ...
+
+def request_auth_end_user_lookup(facts: object) -> object: ...
+
+def request_auth_end_user_validation(facts: object) -> object: ...

tests/unit/proxy/auth/test_end_user_context_adapter.py

@@ -0,0 +1,422 @@
+from collections.abc import Mapping
+from typing import Final
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+import litellm
+from litellm._auth.failure import NativeAuthRuntimeError
+from litellm.models.budget import LiteLLM_BudgetTable
+from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_UserTable, UserAPIKeyAuth
+from litellm.proxy.auth import auth_checks, user_api_key_auth as legacy
+from litellm.proxy.auth import end_user_context_adapter as adapter
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+
+@pytest.fixture(params=(False, True), ids=("reference", "native"))
+def native(request: pytest.FixtureRequest) -> bool:
+    return request.param is True  # pyright: ignore[reportAny]  # boolean fixture parameter
+
+
+@pytest.mark.parametrize(
+    "value,expected", ((None, None), ("", None), (" ", " "), ("budget", "budget"), (0, None), (False, None), ({}, None))
+)
+def test_metadata_id_keeps_only_nonempty_strings(native: bool, value: object, expected: str | None) -> None:
+    function: Final = adapter.get_key_end_user_budget_id if native else auth_checks.get_key_end_user_budget_id
+    assert function({"end_user_budget_id": value}) == expected
+    assert function(None) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("explicit", (None, "", "key"))
+@pytest.mark.parametrize("global_id", (None, "", "global"))
+@pytest.mark.parametrize("db_present", (False, True))
+async def test_default_budget_id_and_cache_selection(
+    native: bool, explicit: str | None, global_id: str | None, db_present: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id)
+    cache: Final = UserApiKeyCache()
+    selected: Final = explicit if explicit is not None else global_id
+    budget: Final = LiteLLM_BudgetTable(budget_id=selected, max_budget=0)
+    await cache.async_set_cache(key=f"default_end_user_budget:{selected}", value=budget, model_type=LiteLLM_BudgetTable)
+    function: Final = adapter.get_default_end_user_budget if native else auth_checks.get_default_end_user_budget
+    result: Final = await function(Mock() if db_present else None, cache, budget_id=explicit)
+    assert result == (budget if db_present and selected is not None else None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("own_id", (None, "", "own"))
+@pytest.mark.parametrize("own_present", (False, True))
+@pytest.mark.parametrize("key_id", (None, "", "key"))
+@pytest.mark.parametrize("global_id", (None, "", "global"))
+@pytest.mark.parametrize("found", (False, True))
+async def test_default_application_preserves_own_budget_and_copies_only_resolved_defaults(
+    native: bool,
+    own_id: str | None,
+    own_present: bool,
+    key_id: str | None,
+    global_id: str | None,
+    found: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id)
+    own: Final = LiteLLM_BudgetTable(budget_id="own", max_budget=2)
+    default: Final = LiteLLM_BudgetTable(budget_id="default", max_budget=3) if found else None
+    row: Final = LiteLLM_EndUserTable(
+        user_id="eu", blocked=False, budget_id=own_id, litellm_budget_table=own if own_present else None
+    )
+    load: Final = AsyncMock(return_value=default)
+    monkeypatch.setattr(auth_checks, "get_default_end_user_budget", load)
+    function: Final = (
+        adapter.apply_default_budget_to_end_user if native else auth_checks._apply_default_budget_to_end_user  # pyright: ignore[reportPrivateUsage]  # Phase 1 parity reference
+    )  # pyright: ignore[reportPrivateUsage]  # Phase 1 parity reference
+    result: Final = await function(row, Mock(), UserApiKeyCache(), key_end_user_budget_id=key_id)
+    unchanged: Final = (own_id is not None and own_present) or (key_id is None and global_id is None) or not found
+    assert (result is row) is unchanged
+    assert result.litellm_budget_table is (row.litellm_budget_table if unchanged else default)
+    assert row.litellm_budget_table is (own if own_present else None)
+    expected_ids: Final = (
+        ()
+        if own_id is not None and own_present
+        else (
+            ((key_id,) if key_id is not None else ())
+            + ((global_id,) if global_id is not None and (key_id is None or not found) else ())
+        )
+    )
+    assert tuple(call.kwargs.get("budget_id", global_id) for call in load.await_args_list) == expected_ids
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("global_id", (None, "", "global"))
+@pytest.mark.parametrize("key_id", (None, "", "key"))
+@pytest.mark.parametrize("validate", (False, True))
+@pytest.mark.parametrize("max_budget", (None, 0))
+@pytest.mark.parametrize("ids", (None, frozenset[str](), frozenset(("eu",))))
+async def test_unrestricted_registry_requires_absence_of_all_budget_and_validation_constraints(
+    native: bool,
+    global_id: str | None,
+    key_id: str | None,
+    validate: bool,
+    max_budget: int | None,
+    ids: frozenset[str] | None,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id)
+    monkeypatch.setattr(litellm, "validate_end_user_id_in_db", validate)
+    load: Final = AsyncMock(return_value=ids)
+    monkeypatch.setattr(auth_checks, "_load_end_user_restricted_registry", load)
+    function: Final = adapter.end_user_is_known_unrestricted if native else auth_checks._end_user_is_known_unrestricted  # pyright: ignore[reportPrivateUsage]  # Phase 1 parity reference
+    result: Final = await function("eu", Mock(), UserApiKeyCache(), max_budget, key_id)
+    blocked: Final = global_id is not None or key_id is not None or validate or max_budget is not None
+    assert result is (not blocked and ids is not None and "eu" not in ids)
+    assert load.await_count == (0 if blocked else 1)
+
+
+@pytest.mark.parametrize("value", (None, 0, 1))
+@pytest.mark.parametrize("models", (None, {}, {"model": {"budget_limit": 0}}))
+def test_budget_parameter_projection_keeps_zero_and_empty_model_limits(
+    native: bool,
+    value: int | None,
+    models: Mapping[str, object] | None,
+) -> None:
+    budget: Final = LiteLLM_BudgetTable(
+        tpm_limit=value,
+        rpm_limit=value,
+        tpd_limit=value,
+        max_budget=value,
+        model_max_budget=dict(models) if models is not None else None,
+    )
+    params: Final[dict[str, object]] = {"existing": True}  # mutable-ok: legacy output argument
+    if native:
+        adapter.apply_budget_limits_to_end_user_params(params, budget, "eu")
+    else:
+        legacy._apply_budget_limits_to_end_user_params(params, budget, "eu")  # pyright: ignore[reportUnknownMemberType, reportPrivateUsage]  # Phase 1 unparameterized reference dictionary
+    expected: Final = {
+        "existing": True,
+        **(
+            {
+                name: value
+                for name in ("end_user_tpm_limit", "end_user_rpm_limit", "end_user_tpd_limit", "end_user_max_budget")
+            }
+            if value is not None
+            else {}
+        ),
+        **({"end_user_model_max_budget": models} if models is not None else {}),
+    }
+    assert params == expected
+    assert tuple(params) == tuple(expected)
+    if models is not None:
+        assert params["end_user_model_max_budget"] is getattr(budget, "model_max_budget")
+
+
+@pytest.mark.parametrize("value", (None, 0, "opaque"))
+def test_token_projection_preserves_unset_limits_and_never_copies_max_budget(native: bool, value: object) -> None:
+    token: Final = UserAPIKeyAuth(
+        end_user_id="old",
+        end_user_tpm_limit=9,
+        end_user_rpm_limit=9,
+        end_user_tpd_limit=9,
+        end_user_max_budget=9,
+        end_user_model_max_budget={"old": {}},
+        allowed_model_region="eu",
+    )
+    params: Final = {
+        name: value
+        for name in (
+            "end_user_tpm_limit",
+            "end_user_rpm_limit",
+            "end_user_tpd_limit",
+            "end_user_max_budget",
+            "end_user_model_max_budget",
+            "allowed_model_region",
+        )
+    }
+    result: Final = (
+        adapter.update_valid_token_with_end_user_params(token, params)
+        if native
+        else legacy.update_valid_token_with_end_user_params(token, params)  # pyright: ignore[reportUnknownMemberType]  # Phase 1 unparameterized reference dictionary
+    )
+    assert result is token
+    assert token.end_user_id is None
+    assert token.end_user_tpm_limit == (9 if value is None else value)
+    assert token.end_user_rpm_limit == (9 if value is None else value)
+    assert token.end_user_tpd_limit == (9 if value is None else value)
+    assert token.allowed_model_region == ("eu" if value is None else value)
+    assert token.end_user_model_max_budget == ({"old": {}} if value is None else value)
+    assert token.end_user_max_budget == 9
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cached", ("valid", "invalid", None, 0))
+@pytest.mark.parametrize("global_id,key_id", ((None, None), ("", None), ("global", None), (None, "")))
+@pytest.mark.parametrize("end_found,user_found", ((False, False), (True, False), (False, True)))
+@pytest.mark.parametrize("identity", ("eu", "email@example.test"))
+async def test_validation_uses_raw_row_presence_and_preserves_unknown_ids_with_defaults(
+    native: bool,
+    cached: object,
+    global_id: str | None,
+    key_id: str | None,
+    end_found: bool,
+    user_found: bool,
+    identity: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True)
+    monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id)
+    cache: Final = UserApiKeyCache()
+    await cache.async_set_cache(key=f"end_user_validation:{identity}", value=cached)
+    end: Final = AsyncMock(return_value=LiteLLM_EndUserTable(user_id=identity, blocked=False) if end_found else None)
+    user: Final = AsyncMock(return_value=LiteLLM_UserTable(user_id=identity) if user_found else None)
+    monkeypatch.setattr(auth_checks, "get_end_user_object", end)
+    monkeypatch.setattr(auth_checks, "get_user_object", user)
+    function: Final = (
+        adapter.resolve_and_validate_end_user_id if native else auth_checks.resolve_and_validate_end_user_id
+    )
+    result: Final = await function(identity, Mock(), cache, key_end_user_budget_id=key_id)
+    valid: Final = cached == "valid" or (cached != "invalid" and (end_found or user_found))
+    assert result == (identity if valid or bool(global_id) or key_id is not None else None)
+    assert end.await_count == (0 if cached in ("valid", "invalid") else 1)
+    assert user.await_count == (0 if cached in ("valid", "invalid") or end_found else 1)
+    if user.await_count:
+        assert user.await_args is not None
+        assert user.await_args.kwargs["user_email"] == (identity if "@" in identity else None)
+    assert await cache.async_get_cache(key=f"end_user_validation:{identity}") == ("valid" if valid else "invalid")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("identity", (None, "", "eu"))
+@pytest.mark.parametrize("enabled,db", ((False, False), (False, True), (True, False)))
+async def test_validation_disabled_or_without_db_skips_all_lookups(
+    native: bool,
+    identity: str | None,
+    enabled: bool,
+    db: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "validate_end_user_id_in_db", enabled)
+    cache: Final = Mock(spec=UserApiKeyCache)
+    read: Final = AsyncMock(side_effect=AssertionError("cache must not be read"))
+    cache.async_get_cache = read
+    function: Final = (
+        adapter.resolve_and_validate_end_user_id if native else auth_checks.resolve_and_validate_end_user_id
+    )
+    assert await function(identity, Mock() if db else None, cache) == identity
+    read.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("row_mode", ("missing", "empty", "budget"))
+@pytest.mark.parametrize("initial_max", (None, 0, 7))
+@pytest.mark.parametrize(
+    "key_id,global_id,found",
+    ((None, None, False), ("key", None, True), (None, "global", True), ("key", "global", False)),
+)
+async def test_lookup_uses_row_before_default_and_only_fills_missing_max_budget(
+    native: bool,
+    row_mode: str,
+    initial_max: int | None,
+    key_id: str | None,
+    global_id: str | None,
+    found: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "max_end_user_budget_id", global_id)
+    budget: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=5, tpm_limit=0, model_max_budget={})
+    row: Final = (
+        LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if row_mode == "budget" else None)
+        if row_mode != "missing"
+        else None
+    )
+    end: Final = AsyncMock(return_value=row)
+    default: Final = AsyncMock(return_value=budget if found else None)
+    monkeypatch.setattr(auth_checks, "get_end_user_object", end)
+    monkeypatch.setattr(legacy, "get_end_user_object", end)
+    monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default)
+    token: Final = UserAPIKeyAuth(
+        end_user_id="eu", end_user_max_budget=initial_max, end_user_tpm_limit=9, metadata={"end_user_budget_id": key_id}
+    )
+    result: Final = (
+        await adapter.lookup_end_user_and_apply_budget(
+            token, "/chat/completions", None, Mock(), UserApiKeyCache(), None
+        )
+        if native
+        else await legacy._lookup_end_user_and_apply_budget(  # pyright: ignore[reportUnknownMemberType, reportPrivateUsage]  # Phase 1 untyped reference boundary
+            token, "/chat/completions", None, Mock(), UserApiKeyCache(), None
+        )  # pyright: ignore[reportUnknownMemberType, reportPrivateUsage]  # Phase 1 untyped reference boundary
+    )
+    actual_token, actual_row = result
+    assert actual_token is token and actual_row is row
+    default_used: Final = row is None and (key_id is not None or global_id is not None) and found
+    assert token.end_user_max_budget == (5 if default_used and initial_max is None else initial_max)
+    assert token.end_user_tpm_limit == (0 if row_mode == "budget" or default_used else 9)
+    assert token.end_user_model_max_budget == ({} if row_mode == "budget" or default_used else None)
+    assert end.await_count == 1
+    assert default.await_count == (
+        0
+        if row is not None
+        else (int(key_id is not None) + int(global_id is not None and (key_id is None or not found)))
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native_failure", (False, True))
+async def test_lookup_io_failure_does_not_select_defaults_and_native_failure_is_terminal(
+    native_failure: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    failure: Final = NativeAuthRuntimeError("native") if native_failure else RuntimeError("io")
+    load: Final = AsyncMock(side_effect=failure)
+    default: Final = AsyncMock(side_effect=AssertionError("no defaults after lookup error"))
+    monkeypatch.setattr(auth_checks, "get_end_user_object", load)
+    monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default)
+    token: Final = UserAPIKeyAuth(end_user_id="eu", metadata={"end_user_budget_id": "key"})
+    if native_failure:
+        with pytest.raises(NativeAuthRuntimeError) as raised:
+            await adapter.lookup_end_user_and_apply_budget(token, "", None, Mock(), UserApiKeyCache(), None)
+        assert raised.value is failure
+    else:
+        assert await adapter.lookup_end_user_and_apply_budget(token, "", None, Mock(), UserApiKeyCache(), None) == (
+            token,
+            None,
+        )
+    default.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_native_validation_failure_stops_before_user_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True)
+    failure: Final = NativeAuthRuntimeError("native")
+    monkeypatch.setattr(auth_checks, "get_end_user_object", AsyncMock(side_effect=failure))
+    user: Final = AsyncMock()
+    monkeypatch.setattr(auth_checks, "get_user_object", user)
+    with pytest.raises(NativeAuthRuntimeError) as raised:
+        await adapter.resolve_and_validate_end_user_id("eu", Mock(), UserApiKeyCache())
+    assert raised.value is failure
+    user.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("found", (False, True))
+async def test_default_budget_cache_miss_reads_database_and_caches_only_found_rows(
+    native: bool,
+    found: bool,
+) -> None:
+    cache: Final = UserApiKeyCache()
+    row: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=0) if found else None
+    read: Final = AsyncMock(return_value=row)
+    db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read)))
+    function: Final = adapter.get_default_end_user_budget if native else auth_checks.get_default_end_user_budget
+    assert await function(db, cache, budget_id="budget") == row
+    read.assert_awaited_once_with(where={"budget_id": "budget"})
+    assert await cache.async_get_cache(key="default_end_user_budget:budget", model_type=LiteLLM_BudgetTable) == row
+    if found:
+        assert await function(db, cache, budget_id="budget") == row
+        read.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native_failure", (False, True))
+async def test_default_budget_database_failure_keeps_native_fault_terminal(native_failure: bool) -> None:
+    failure: Final = NativeAuthRuntimeError("native") if native_failure else RuntimeError("io")
+    read: Final = AsyncMock(side_effect=failure)
+    db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read)))
+    if native_failure:
+        with pytest.raises(NativeAuthRuntimeError) as raised:
+            await adapter.get_default_end_user_budget(db, UserApiKeyCache(), budget_id="budget")
+        assert raised.value is failure
+    else:
+        assert await adapter.get_default_end_user_budget(db, UserApiKeyCache(), budget_id="budget") is None
+    read.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_default_budget_cache_failure_is_not_swallowed(native: bool) -> None:
+    failure: Final = RuntimeError("cache unavailable")
+    cache: Final = Mock(spec=UserApiKeyCache)
+    read: Final = AsyncMock(side_effect=failure)
+    cache.async_get_cache = read
+    function: Final = adapter.get_default_end_user_budget if native else auth_checks.get_default_end_user_budget
+    with pytest.raises(RuntimeError) as raised:
+        await function(Mock(), cache, budget_id="budget")
+    assert raised.value is failure
+    read.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("own_failure", (False, True))
+@pytest.mark.parametrize("user_found", (False, True))
+async def test_validation_lookup_errors_select_secondary_identity_read(
+    native: bool,
+    own_failure: bool,
+    user_found: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    own: Final = AsyncMock(side_effect=RuntimeError("end-user read")) if own_failure else AsyncMock(return_value=None)
+    user: Final = (
+        AsyncMock(return_value=LiteLLM_UserTable(user_id="eu"))
+        if user_found
+        else AsyncMock(side_effect=RuntimeError("user read"))
+    )
+    monkeypatch.setattr(auth_checks, "get_end_user_object", own)
+    monkeypatch.setattr(auth_checks, "get_user_object", user)
+    function: Final = adapter.end_user_id_exists_in_db if native else auth_checks._end_user_id_exists_in_db  # pyright: ignore[reportPrivateUsage]  # Phase 1 parity reference
+    assert await function("eu", Mock(), UserApiKeyCache()) is user_found
+    own.assert_awaited_once()
+    user.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_lookup_budget_denial_preserves_exception_identity(native: bool, monkeypatch: pytest.MonkeyPatch) -> None:
+    failure: Final = litellm.BudgetExceededError(current_cost=2, max_budget=1)
+    read: Final = AsyncMock(side_effect=failure)
+    monkeypatch.setattr(auth_checks, "get_end_user_object", read)
+    monkeypatch.setattr(legacy, "get_end_user_object", read)
+    invocation: Final = (
+        adapter.lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None)
+        if native
+        else legacy._lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None)  # pyright: ignore[reportPrivateUsage, reportUnknownMemberType]  # Phase 1 untyped reference boundary
+    )
+    with pytest.raises(litellm.BudgetExceededError) as raised:
+        await invocation
+    assert raised.value is failure
+    read.assert_awaited_once()