From 044ebfe86b70d8b0acc9e296c0fff2a47d4f595b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 8 Nov 2026 23:44:01 +0000 Subject: [PATCH] refactor(auth): move end-user defaults and identity stages into Rust --- litellm-rust/crates/python-bridge/src/lib.rs | 5 +- .../crates/python-bridge/src/request_auth.rs | 84 ++++ .../request-auth-core/src/end_user_context.rs | 428 ++++++++++++++++++ .../crates/request-auth-core/src/lib.rs | 1 + litellm/_auth/end_user_context.py | 171 +++++++ litellm/proxy/auth/budget_context_adapter.py | 4 +- .../proxy/auth/end_user_context_adapter.py | 406 +++++++++++++++++ litellm/proxy/auth/end_user_context_io.py | 50 ++ litellm/rust_bridge/_native.pyi | 14 + .../auth/test_end_user_context_adapter.py | 422 +++++++++++++++++ 10 files changed, 1582 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/request-auth-core/src/end_user_context.rs create mode 100644 litellm/_auth/end_user_context.py create mode 100644 litellm/proxy/auth/end_user_context_adapter.py create mode 100644 litellm/proxy/auth/end_user_context_io.py create mode 100644 tests/unit/proxy/auth/test_end_user_context_adapter.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 51542654a02..9376693ce18 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/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, diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index e8a5bf6f5cf..f4c1deb388a 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/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> { + let facts: core::end_user_context::IdFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::normalize_id(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_default_id( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::DefaultFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::default_id(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_apply( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::ApplyFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::apply_default(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_registry( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::RegistryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::registry(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_assignments( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::AssignmentFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::assignments(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_lookup( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::LookupFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::lookup(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_end_user_validation( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::end_user_context::ValidationFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::end_user_context::validation(&facts)), + ) +} diff --git a/litellm-rust/crates/request-auth-core/src/end_user_context.rs b/litellm-rust/crates/request-auth-core/src/end_user_context.rs new file mode 100644 index 00000000000..ef7a5084bc2 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/end_user_context.rs @@ -0,0 +1,428 @@ +use serde::{Deserialize, Serialize}; +use strum::VariantArray; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct IdFacts { + pub value: Option, +} + +pub fn normalize_id(facts: &IdFacts) -> Option { + facts.value.clone().filter(|value| !value.is_empty()) +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DefaultFacts { + pub prisma_present: bool, + pub explicit_id: Option, + pub global_read: bool, + pub global_id: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct DefaultDecision { + pub read_global: bool, + pub budget_id: Option, +} + +pub fn default_id(facts: &DefaultFacts) -> DefaultDecision { + if facts.explicit_id.is_none() && !facts.global_read { + return DefaultDecision { + read_global: true, + budget_id: None, + }; + } + if !facts.prisma_present { + return DefaultDecision { + read_global: false, + budget_id: None, + }; + } + DefaultDecision { + read_global: false, + budget_id: facts + .explicit_id + .clone() + .or_else(|| facts.global_id.clone()), + } +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ApplyFacts { + pub own_budget_id: Option, + pub own_budget_present: bool, + pub key_id: Option, + pub global_read: bool, + pub global_id: Option, + pub default_read: bool, + pub default_present: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ApplyDecision { + Keep, + ReadGlobal, + ReadDefault, + Copy, +} + +pub fn apply_default(facts: &ApplyFacts) -> ApplyDecision { + if facts.own_budget_id.is_some() && facts.own_budget_present { + return ApplyDecision::Keep; + } + if facts.key_id.is_none() { + if !facts.global_read { + return ApplyDecision::ReadGlobal; + } + if facts.global_id.is_none() { + return ApplyDecision::Keep; + } + } + if !facts.default_read { + return ApplyDecision::ReadDefault; + } + if facts.default_present { + ApplyDecision::Copy + } else { + ApplyDecision::Keep + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct RegistryFacts { + pub end_user_id: String, + pub global_id: Option, + pub key_id: Option, + pub validate_id: bool, + pub token_max_present: bool, + pub registry_read: bool, + pub registry: Option>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum RegistryDecision { + Read, + Restricted, + Unrestricted, +} + +pub fn registry(facts: &RegistryFacts) -> RegistryDecision { + if facts.global_id.is_some() + || facts.key_id.is_some() + || facts.validate_id + || facts.token_max_present + { + return RegistryDecision::Restricted; + } + if !facts.registry_read { + return RegistryDecision::Read; + } + if facts + .registry + .as_ref() + .is_some_and(|ids| !ids.contains(&facts.end_user_id)) + { + RegistryDecision::Unrestricted + } else { + RegistryDecision::Restricted + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum Field { + EndUserId, + Tpm, + Rpm, + Tpd, + MaxBudget, + Models, + Region, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct Fields { + pub tpm: bool, + pub rpm: bool, + pub tpd: bool, + pub max_budget: bool, + pub models: bool, + pub region: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum AssignmentMode { + Params, + Token, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct AssignmentFacts { + pub mode: AssignmentMode, + pub fields: Fields, +} + +pub fn assignments(facts: &AssignmentFacts) -> Vec { + let fields = &facts.fields; + match facts.mode { + AssignmentMode::Params => [ + (Field::Tpm, fields.tpm), + (Field::Rpm, fields.rpm), + (Field::Tpd, fields.tpd), + (Field::MaxBudget, fields.max_budget), + (Field::Models, fields.models), + ] + .into_iter() + .filter_map(|(field, present)| present.then_some(field)) + .collect(), + AssignmentMode::Token => [ + (Field::EndUserId, true), + (Field::Tpm, fields.tpm), + (Field::Rpm, fields.rpm), + (Field::Tpd, fields.tpd), + (Field::Region, fields.region), + (Field::Models, fields.models), + ] + .into_iter() + .filter_map(|(field, present)| present.then_some(field)) + .collect(), + } +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct LookupFacts { + pub own_budget_present: bool, + pub key_id: Option, + pub end_user_read: bool, + pub end_user_present: bool, + pub global_read: bool, + pub global_id: Option, + pub default_read: bool, + pub default_present: bool, + pub token_max_present: bool, + pub applied: bool, + pub read_failed: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum LookupAction { + ReadEndUser, + ReadGlobal, + ReadDefault, + ApplyOwn, + ApplyDefault, + Done, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct LookupDecision { + pub action: LookupAction, + pub set_max_budget: bool, + pub project_budget: bool, +} + +pub fn lookup(facts: &LookupFacts) -> LookupDecision { + let action = if facts.read_failed || facts.applied { + LookupAction::Done + } else if !facts.end_user_read { + LookupAction::ReadEndUser + } else if facts.end_user_present { + LookupAction::ApplyOwn + } else if facts.key_id.is_none() && !facts.global_read { + LookupAction::ReadGlobal + } else if facts.key_id.is_none() && facts.global_id.is_none() { + LookupAction::Done + } else if !facts.default_read { + LookupAction::ReadDefault + } else if facts.default_present { + LookupAction::ApplyDefault + } else { + LookupAction::Done + }; + LookupDecision { + project_budget: action == LookupAction::ApplyDefault + || (action == LookupAction::ApplyOwn && facts.own_budget_present), + set_max_budget: facts.applied + && !facts.read_failed + && !facts.end_user_present + && facts.default_present + && !facts.token_max_present, + action, + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ValidationMode { + #[default] + Resolve, + Exists, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ValidationFacts { + pub mode: ValidationMode, + pub end_user_id: Option, + pub enabled: bool, + pub prisma_present: bool, + pub global_id: Option, + pub key_id: Option, + pub cache_read: bool, + pub cached: Option, + pub end_user_read: bool, + pub end_user_present: bool, + pub user_read: bool, + pub user_present: bool, + pub cache_written: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, VariantArray)] +pub enum ValidationRead { + Cache, + EndUser, + User, + WriteCache, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct ValidationDecision { + pub read: Option, + pub end_user_id: Option, + pub valid: bool, + pub cache_value: Option, + pub cache_ttl: u64, + pub user_email: Option, +} + +fn validation_result(facts: &ValidationFacts, valid: bool) -> ValidationDecision { + let retain = valid + || facts.key_id.is_some() + || facts.global_id.as_ref().is_some_and(|id| !id.is_empty()); + ValidationDecision { + end_user_id: facts.end_user_id.clone().filter(|_| retain), + valid, + ..ValidationDecision::default() + } +} + +pub fn validation(facts: &ValidationFacts) -> ValidationDecision { + let Some(id) = &facts.end_user_id else { + return ValidationDecision::default(); + }; + if facts.mode == ValidationMode::Resolve { + if !facts.enabled || !facts.prisma_present { + return ValidationDecision { + end_user_id: Some(id.clone()), + ..ValidationDecision::default() + }; + } + if !facts.cache_read { + return ValidationDecision { + read: Some(ValidationRead::Cache), + ..ValidationDecision::default() + }; + } + match facts.cached.as_deref() { + Some("valid") => return validation_result(facts, true), + Some("invalid") => return validation_result(facts, false), + _ => {} + } + } + if !facts.end_user_read { + return ValidationDecision { + read: Some(ValidationRead::EndUser), + ..ValidationDecision::default() + }; + } + if !facts.end_user_present && !facts.user_read { + return ValidationDecision { + read: Some(ValidationRead::User), + user_email: id.contains('@').then(|| id.clone()), + ..ValidationDecision::default() + }; + } + let valid = facts.end_user_present || facts.user_present; + if facts.mode == ValidationMode::Resolve && !facts.cache_written { + return ValidationDecision { + read: Some(ValidationRead::WriteCache), + valid, + cache_value: Some(if valid { "valid" } else { "invalid" }.to_owned()), + cache_ttl: if valid { 300 } else { 60 }, + ..ValidationDecision::default() + }; + } + validation_result(facts, valid) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::absent(None, None)] + #[case::empty(Some(""), None)] + #[case::space(Some(" "), Some(" "))] + #[case::id(Some("budget"), Some("budget"))] + fn normalizes_metadata_id(#[case] value: Option<&str>, #[case] expected: Option<&str>) { + assert_eq!( + normalize_id(&IdFacts { + value: value.map(str::to_owned) + }) + .as_deref(), + expected + ); + } + + #[rstest] + #[case::explicit(Some("key"), Some("global"), Some("key"))] + #[case::empty_explicit(Some(""), Some("global"), Some(""))] + #[case::global(None, Some("global"), Some("global"))] + #[case::absent(None, None, None)] + fn selects_default( + #[case] explicit: Option<&str>, + #[case] global: Option<&str>, + #[case] expected: Option<&str>, + ) { + let facts = DefaultFacts { + prisma_present: true, + explicit_id: explicit.map(str::to_owned), + global_read: true, + global_id: global.map(str::to_owned), + }; + assert_eq!(default_id(&facts).budget_id.as_deref(), expected); + } + + #[rstest] + #[case::none(None, None, None)] + #[case::empty_global(Some(""), None, None)] + #[case::global(Some("budget"), None, Some("missing"))] + #[case::key(None, Some(""), Some("missing"))] + fn retains_unknown_ids( + #[case] global: Option<&str>, + #[case] key: Option<&str>, + #[case] expected: Option<&str>, + ) { + let facts = ValidationFacts { + end_user_id: Some("missing".to_owned()), + enabled: true, + prisma_present: true, + global_id: global.map(str::to_owned), + key_id: key.map(str::to_owned), + cache_read: true, + cached: Some("invalid".to_owned()), + ..ValidationFacts::default() + }; + assert_eq!(validation(&facts).end_user_id.as_deref(), expected); + } + + #[rstest] + #[case::missing(AssignmentMode::Params, vec![])] + #[case::token_identity(AssignmentMode::Token, vec![Field::EndUserId])] + fn absent_fields(#[case] mode: AssignmentMode, #[case] expected: Vec) { + assert_eq!( + assignments(&AssignmentFacts { + mode, + fields: Fields::default() + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index eaf650b66c2..b2815d1306e 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/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; diff --git a/litellm/_auth/end_user_context.py b/litellm/_auth/end_user_context.py new file mode 100644 index 00000000000..457cf5e6af6 --- /dev/null +++ b/litellm/_auth/end_user_context.py @@ -0,0 +1,171 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class IdFacts: + value: str | None + + +@dataclass(frozen=True, slots=True) +class DefaultFacts: + prisma_present: bool + explicit_id: str | None + global_id: str | None + global_read: bool = False + + +@dataclass(frozen=True, slots=True) +class DefaultDecision: + read_global: bool + budget_id: str | None + + +@dataclass(frozen=True, slots=True) +class ApplyFacts: + own_budget_id: str | None + own_budget_present: bool + key_id: str | None + global_read: bool = False + global_id: str | None = None + default_read: bool = False + default_present: bool = False + + +ApplyDecision: TypeAlias = Literal["Keep", "ReadGlobal", "ReadDefault", "Copy"] + + +@dataclass(frozen=True, slots=True) +class RegistryFacts: + end_user_id: str + global_id: str | None + key_id: str | None + validate_id: bool + token_max_present: bool + registry_read: bool = False + registry: tuple[str, ...] | None = None + + +RegistryDecision: TypeAlias = Literal["Read", "Restricted", "Unrestricted"] +Field: TypeAlias = Literal["EndUserId", "Tpm", "Rpm", "Tpd", "MaxBudget", "Models", "Region"] + + +@dataclass(frozen=True, slots=True) +class Fields: + tpm: bool = False + rpm: bool = False + tpd: bool = False + max_budget: bool = False + models: bool = False + region: bool = False + + +@dataclass(frozen=True, slots=True) +class AssignmentFacts: + mode: Literal["Params", "Token"] + fields: Fields + + +@dataclass(frozen=True, slots=True) +class LookupFacts: + key_id: str | None + own_budget_present: bool = False + end_user_read: bool = False + end_user_present: bool = False + global_read: bool = False + global_id: str | None = None + default_read: bool = False + default_present: bool = False + token_max_present: bool = False + applied: bool = False + read_failed: bool = False + + +@dataclass(frozen=True, slots=True) +class LookupDecision: + action: Literal["ReadEndUser", "ReadGlobal", "ReadDefault", "ApplyOwn", "ApplyDefault", "Done"] + set_max_budget: bool + project_budget: bool + + +@dataclass(frozen=True, slots=True) +class ValidationFacts: + mode: Literal["Resolve", "Exists"] + end_user_id: str | None + enabled: bool = False + prisma_present: bool = False + global_id: str | None = None + key_id: str | None = None + cache_read: bool = False + cached: str | None = None + end_user_read: bool = False + end_user_present: bool = False + user_read: bool = False + user_present: bool = False + cache_written: bool = False + + +@dataclass(frozen=True, slots=True) +class ValidationDecision: + read: Literal["Cache", "EndUser", "User", "WriteCache"] | None + end_user_id: str | None + valid: bool + cache_value: str | None + cache_ttl: int + user_email: str | None + + +_ID: Final = TypeAdapter(IdFacts) +_STRING: Final = TypeAdapter[str | None](str | None) +_DEFAULT: Final = TypeAdapter(DefaultFacts) +_DEFAULT_RESULT: Final = TypeAdapter(DefaultDecision) +_APPLY: Final = TypeAdapter(ApplyFacts) +_APPLY_RESULT: Final = TypeAdapter[ApplyDecision](ApplyDecision) +_REGISTRY: Final = TypeAdapter(RegistryFacts) +_REGISTRY_RESULT: Final = TypeAdapter[RegistryDecision](RegistryDecision) +_ASSIGNMENTS: Final = TypeAdapter(AssignmentFacts) +_FIELDS: Final = TypeAdapter(tuple[Field, ...]) +_LOOKUP: Final = TypeAdapter(LookupFacts) +_LOOKUP_RESULT: Final = TypeAdapter(LookupDecision) +_VALIDATION: Final = TypeAdapter(ValidationFacts) +_VALIDATION_RESULT: Final = TypeAdapter(ValidationDecision) + + +@native_boundary +def normalize_id(facts: IdFacts) -> str | None: + return _STRING.validate_python(_native.request_auth_end_user_normalize_id(_ID.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def default_id(facts: DefaultFacts) -> DefaultDecision: + return _DEFAULT_RESULT.validate_python(_native.request_auth_end_user_default_id(_DEFAULT.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def apply_default(facts: ApplyFacts) -> ApplyDecision: + return _APPLY_RESULT.validate_python(_native.request_auth_end_user_apply(_APPLY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def registry(facts: RegistryFacts) -> RegistryDecision: + return _REGISTRY_RESULT.validate_python(_native.request_auth_end_user_registry(_REGISTRY.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def assignments(facts: AssignmentFacts) -> tuple[Field, ...]: + return _FIELDS.validate_python(_native.request_auth_end_user_assignments(_ASSIGNMENTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def lookup(facts: LookupFacts) -> LookupDecision: + return _LOOKUP_RESULT.validate_python(_native.request_auth_end_user_lookup(_LOOKUP.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts + + +@native_boundary +def validation(facts: ValidationFacts) -> ValidationDecision: + return _VALIDATION_RESULT.validate_python(_native.request_auth_end_user_validation(_VALIDATION.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen typed facts diff --git a/litellm/proxy/auth/budget_context_adapter.py b/litellm/proxy/auth/budget_context_adapter.py index aab2286515b..80a35f8cac5 100644 --- a/litellm/proxy/auth/budget_context_adapter.py +++ b/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, ) diff --git a/litellm/proxy/auth/end_user_context_adapter.py b/litellm/proxy/auth/end_user_context_adapter.py new file mode 100644 index 00000000000..aae5d73f770 --- /dev/null +++ b/litellm/proxy/auth/end_user_context_adapter.py @@ -0,0 +1,406 @@ +from collections.abc import Mapping, MutableMapping +from dataclasses import replace +from types import MappingProxyType +from typing import Final, assert_never + +from opentelemetry.trace import Span + +import litellm +from litellm._auth import end_user_context as core +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_EndUserTable, UserAPIKeyAuth +from litellm.proxy.auth.budget_context_adapter import mapping_value, resolve_default_end_user_budget +from litellm.proxy.auth.end_user_context_io import EndUserIO, read_default_budget +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +_PARAM_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType( + { + "EndUserId": "end_user_id", + "Tpm": "end_user_tpm_limit", + "Rpm": "end_user_rpm_limit", + "Tpd": "end_user_tpd_limit", + "MaxBudget": "end_user_max_budget", + "Models": "end_user_model_max_budget", + "Region": "allowed_model_region", + } +) +_BUDGET_NAMES: Final[Mapping[core.Field, str]] = MappingProxyType( + { + "Tpm": "tpm_limit", + "Rpm": "rpm_limit", + "Tpd": "tpd_limit", + "MaxBudget": "max_budget", + "Models": "model_max_budget", + } +) + + +def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: + raw: Final = key_metadata.get("end_user_budget_id") if key_metadata is not None else None + return core.normalize_id(core.IdFacts(raw if isinstance(raw, str) else None)) + + +def select_default_id(facts: core.DefaultFacts) -> str | None: + decision: Final = core.default_id(facts) + if decision.read_global: + return select_default_id(replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id)) + return decision.budget_id + + +async def get_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + budget_id: str | None = None, +) -> LiteLLM_BudgetTable | None: + selected: Final = select_default_id(core.DefaultFacts(prisma_client is not None, budget_id, None)) + if selected is None: + return None + return await read_default_budget(selected, EndUserIO(prisma_client, user_api_key_cache, parent_otel_span)) + + +async def _apply_default( + facts: core.ApplyFacts, row: LiteLLM_EndUserTable, io: EndUserIO, default: LiteLLM_BudgetTable | None = None +) -> LiteLLM_EndUserTable: + decision: Final = core.apply_default(facts) + match decision: + case "Keep": + return row + case "ReadGlobal": + return await _apply_default( + replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), row, io + ) + case "ReadDefault": + budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span) + return await _apply_default( + replace(facts, default_read=True, default_present=budget is not None), row, io, budget + ) + case "Copy": + assert default is not None + verbose_proxy_logger.debug("Applied default budget %s to end user %s", default.budget_id, row.user_id) + return row.model_copy(update=MappingProxyType({"litellm_budget_table": default})) + case _: + assert_never(decision) + + +async def apply_default_budget_to_end_user( + end_user_obj: LiteLLM_EndUserTable, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable: + return await _apply_default( + core.ApplyFacts(end_user_obj.budget_id, end_user_obj.litellm_budget_table is not None, key_end_user_budget_id), + end_user_obj, + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + ) + + +async def _registry(facts: core.RegistryFacts, io: EndUserIO) -> bool: + from litellm.proxy.auth import auth_checks + + decision: Final = core.registry(facts) + match decision: + case "Restricted": + return False + case "Unrestricted": + return True + case "Read": + assert io.prisma is not None + ids: Final = await auth_checks._load_end_user_restricted_registry( # pyright: ignore[reportPrivateUsage] # raw cached registry IO + prisma_client=io.prisma, user_api_key_cache=io.cache + ) + return await _registry( + replace(facts, registry_read=True, registry=tuple(ids) if ids is not None else None), io + ) + case _: + assert_never(decision) + + +async def end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, + key_end_user_budget_id: str | None = None, +) -> bool: + return await _registry( + core.RegistryFacts( + end_user_id, + litellm.max_end_user_budget_id, + key_end_user_budget_id, + litellm.validate_end_user_id_in_db, + token_end_user_max_budget is not None, + ), + EndUserIO(prisma_client, user_api_key_cache), + ) + + +def _budget_fields(row: LiteLLM_BudgetTable | None) -> core.Fields: + return core.Fields( + tpm=getattr(row, "tpm_limit", None) is not None, + rpm=getattr(row, "rpm_limit", None) is not None, + tpd=getattr(row, "tpd_limit", None) is not None, + max_budget=getattr(row, "max_budget", None) is not None, + models=getattr(row, "model_max_budget", None) is not None, + ) + + +def _budget_params(row: LiteLLM_BudgetTable | None) -> Mapping[str, object]: + selected: Final = core.assignments(core.AssignmentFacts("Params", _budget_fields(row))) + return MappingProxyType({_PARAM_NAMES[field]: getattr(row, _BUDGET_NAMES[field]) for field in selected}) + + +def apply_budget_limits_to_end_user_params( + end_user_params: MutableMapping[str, object], # mutable-ok: legacy request-local output parameter + budget_info: LiteLLM_BudgetTable, + end_user_id: str | None, +) -> None: + for name, value in _budget_params(budget_info).items(): + end_user_params[name] = value # rebind-ok: native-selected legacy request parameter projection + verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + + +def _token_fields(end_user_params: Mapping[str, object]) -> tuple[core.Field, ...]: + fields: Final = core.Fields( + tpm=end_user_params.get("end_user_tpm_limit") is not None, + rpm=end_user_params.get("end_user_rpm_limit") is not None, + tpd=end_user_params.get("end_user_tpd_limit") is not None, + models=end_user_params.get("end_user_model_max_budget") is not None, + region=end_user_params.get("allowed_model_region") is not None, + ) + return core.assignments(core.AssignmentFacts("Token", fields)) + + +def _project_token( + valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object], fields: tuple[core.Field, ...] +) -> UserAPIKeyAuth: + for field in fields: + setattr(valid_token, _PARAM_NAMES[field], end_user_params.get(_PARAM_NAMES[field])) + return valid_token + + +def update_valid_token_with_end_user_params( + valid_token: UserAPIKeyAuth, end_user_params: Mapping[str, object] +) -> UserAPIKeyAuth: + return _project_token(valid_token, end_user_params, _token_fields(end_user_params)) + + +async def _validation(facts: core.ValidationFacts, io: EndUserIO) -> core.ValidationDecision: + from litellm.proxy.auth import auth_checks + + decision: Final = core.validation(facts) + match decision.read: + case None: + return decision + case "Cache": + cached: Final[object] = await auth_checks._raw_cache(io.cache).async_get_cache( # pyright: ignore[reportPrivateUsage] # shared raw cache IO + key=f"end_user_validation:{facts.end_user_id}" + ) + return await _validation( + replace(facts, cache_read=True, cached=cached if isinstance(cached, str) else None), io + ) + case "EndUser": + try: + row: Final = await auth_checks.get_end_user_object( + end_user_id=facts.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + route=io.route, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve end-user validation lookup failure handling + verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", error) + return await _validation(replace(facts, end_user_read=True), io) + return await _validation(replace(facts, end_user_read=True, end_user_present=row is not None), io) + case "User": + assert facts.end_user_id is not None and io.prisma is not None + try: + user: Final = await auth_checks.get_user_object( + user_id=facts.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + user_id_upsert=False, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + check_db_only=False, + user_email=decision.user_email, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve secondary identity lookup failure handling + verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", error) + return await _validation(replace(facts, user_read=True), io) + return await _validation(replace(facts, user_read=True, user_present=user is not None), io) + case "WriteCache": + await io.cache.async_set_cache( + key=f"end_user_validation:{facts.end_user_id}", value=decision.cache_value, ttl=decision.cache_ttl + ) + return await _validation(replace(facts, cache_written=True), io) + case _: + assert_never(decision.read) + + +async def resolve_and_validate_end_user_id( + raw_end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", + key_end_user_budget_id: str | None = None, +) -> str | None: + decision: Final = await _validation( + core.ValidationFacts( + "Resolve", + raw_end_user_id, + litellm.validate_end_user_id_in_db, + prisma_client is not None, + litellm.max_end_user_budget_id, + key_end_user_budget_id, + ), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) + return decision.end_user_id + + +async def end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", +) -> bool: + decision: Final = await _validation( + core.ValidationFacts("Exists", end_user_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) + return decision.valid + + +def _complete_lookup( + decision: core.LookupDecision, + token: UserAPIKeyAuth, + row: LiteLLM_EndUserTable | None, + default: LiteLLM_BudgetTable | None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + if decision.set_max_budget: + assert default is not None + try: + token.end_user_max_budget = default.max_budget # rebind-ok: request-local native projection + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve request-local setter failure handling + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return token, row + + +async def _lookup( + facts: core.LookupFacts, + token: UserAPIKeyAuth, + io: EndUserIO, + row: LiteLLM_EndUserTable | None = None, + default: LiteLLM_BudgetTable | None = None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + from litellm.proxy.auth.auth_checks import get_end_user_object + + decision: Final = core.lookup(facts) + match decision.action: + case "Done": + return _complete_lookup(decision, token, row, default) + case "ReadGlobal": + return await _lookup( + replace(facts, global_read=True, global_id=litellm.max_end_user_budget_id), token, io, row + ) + case "ReadEndUser": + try: + end_user: Final = await get_end_user_object( + end_user_id=token.end_user_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + proxy_logging_obj=io.logging, + route=io.route, + token_end_user_max_budget=token.end_user_max_budget, + key_end_user_budget_id=facts.key_id, + ) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve the legacy lookup IO failure path + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace( + facts, + end_user_read=True, + end_user_present=end_user is not None, + own_budget_present=end_user is not None and end_user.litellm_budget_table is not None, + ), + token, + io, + end_user, + ) + case "ReadDefault": + try: + budget: Final = await resolve_default_end_user_budget(io.prisma, io.cache, facts.key_id, io.span) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve the legacy default lookup IO failure path + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace(facts, default_read=True, default_present=budget is not None), token, io, row, budget + ) + case "ApplyOwn" | "ApplyDefault": + budget_row: Final = ( + row.litellm_budget_table if decision.action == "ApplyOwn" and row is not None else default + ) + params: Final = MappingProxyType( + { + "end_user_id": token.end_user_id, + "allowed_model_region": getattr(row, "allowed_model_region", None), + **_budget_params(budget_row), + } + ) + fields: Final = _token_fields(params) + try: + if decision.project_budget: + verbose_proxy_logger.debug("Applied budget limits to end user %s", token.end_user_id or "") + _project_token(token, params, fields) + except (NativeAuthFailure, litellm.BudgetExceededError): + raise + except Exception as error: # noqa: BLE001 # preserve request-local setter failure handling + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", error) + return await _lookup(replace(facts, read_failed=True), token, io, row) + return await _lookup( + replace(facts, applied=True, token_max_present=token.end_user_max_budget is not None), + token, + io, + row, + default, + ) + case _: + assert_never(decision.action) + + +async def lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Span | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[UserAPIKeyAuth, LiteLLM_EndUserTable | None]: + key_id: Final = get_key_end_user_budget_id(mapping_value(getattr(valid_token, "metadata", None))) + return await _lookup( + core.LookupFacts(key_id), + valid_token, + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj, route), + ) diff --git a/litellm/proxy/auth/end_user_context_io.py b/litellm/proxy/auth/end_user_context_io.py new file mode 100644 index 00000000000..83cf9784886 --- /dev/null +++ b/litellm/proxy/auth/end_user_context_io.py @@ -0,0 +1,50 @@ +from dataclasses import dataclass +from typing import Final + +from opentelemetry.trace import Span + +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_BudgetTable +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.budget_repository import BudgetRepository + + +@dataclass(frozen=True, slots=True) +class EndUserIO: + prisma: PrismaClient | None + cache: UserApiKeyCache + span: Span | None = None + logging: ProxyLogging | None = None + route: str = "" + + +async def read_default_budget(budget_id: str, io: EndUserIO) -> LiteLLM_BudgetTable | None: + from litellm.proxy.auth import auth_checks + + assert io.prisma is not None + key: Final = f"default_end_user_budget:{budget_id}" + cached: Final = await io.cache.async_get_cache(key=key, model_type=LiteLLM_BudgetTable) + if cached is not None: + return cached + try: + table: Final = auth_checks._dictable_table(BudgetRepository(io.prisma), "budget") # pyright: ignore[reportPrivateUsage] # shared raw repository projection + record: Final = await table.find_unique(where={"budget_id": budget_id}) + if record is None: + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", budget_id.replace("\r", "").replace("\n", "") + ) + return None + row: Final = LiteLLM_BudgetTable.model_validate(record.dict()) + await io.cache.async_set_cache( + key=key, value=row, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(io.cache) + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve the default budget DB/cache IO failure contract + verbose_proxy_logger.error("Error fetching default end user budget: %s", error) + return None + + else: + return row diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c8e429ab922..7c6d213e796 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/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: ... diff --git a/tests/unit/proxy/auth/test_end_user_context_adapter.py b/tests/unit/proxy/auth/test_end_user_context_adapter.py new file mode 100644 index 00000000000..66d66804629 --- /dev/null +++ b/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()