refactor(auth): move budget context selection into native stages

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

@@ -42,6 +42,9 @@ mod _native {
         NativeAuthContinuation, request_auth_budget_alert, request_auth_budget_alert_plan,
         request_auth_budget_exemption, request_auth_budget_price, request_auth_budget_route,
         request_auth_budget_window_start, request_auth_budgets, request_auth_common_budget,
+        request_auth_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_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

@@ -234,3 +234,84 @@ pub(crate) fn request_auth_budget_alert_plan(
     let facts: core::budget_alerts::Facts = from_py(facts)?;
     to_py(py, &py.detach(move || core::budget_alerts::plan(&facts)))
 }
+
+#[pyfunction]
+pub(crate) fn request_auth_context_models(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_context::ModelCheckFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_context::model_checks(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_context_end_user(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_context::EndUserFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_context::end_user_source(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_context_token(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_context::TokenDefaultsFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_context::token_assignments(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_context_member(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_context::MemberFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_context::member_source(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_context_temporary(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_context::TemporaryFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_context::temporary(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_context_throttle(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_context::ThrottleFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_context::throttle(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_context_percentage(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Option<f64>> {
+    let facts: core::budget_context::NumericFacts = from_py(facts)?;
+    Ok(py.detach(move || core::budget_context::percentage(&facts)))
+}

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

@@ -0,0 +1,481 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)]
+pub enum ModelScope {
+    Team,
+    EndUser,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct ModelCheckFacts {
+    pub scope: ModelScope,
+    pub entity_id: Option<String>,
+    pub budget_count: usize,
+    pub model_count: usize,
+}
+
+pub fn model_checks(facts: &ModelCheckFacts) -> Vec<usize> {
+    if facts.budget_count == 0
+        || facts.entity_id.is_none()
+        || (facts.scope == ModelScope::EndUser && facts.entity_id.as_deref() == Some(""))
+    {
+        return vec![];
+    }
+    (0..facts.model_count).collect()
+}
+
+#[derive(Clone, Debug, Default, Deserialize)]
+pub struct BudgetFields {
+    pub max_budget: Option<f64>,
+    pub tpm_limit: Option<i64>,
+    pub rpm_limit: Option<i64>,
+    pub tpd_limit: Option<i64>,
+    pub model_budget_count: Option<usize>,
+}
+
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)]
+pub enum EndUserMode {
+    Models,
+    Token,
+    Default,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct EndUserFacts {
+    pub mode: EndUserMode,
+    pub object_present: bool,
+    pub end_user_id: Option<String>,
+    pub prisma_present: bool,
+    pub key_budget_id: Option<String>,
+    pub key_read: bool,
+    pub key_budget: Option<BudgetFields>,
+    pub global_config_read: bool,
+    pub global_budget_id: Option<String>,
+    pub global_read: bool,
+    pub global_budget: Option<BudgetFields>,
+}
+
+#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)]
+pub enum EndUserSource {
+    Own,
+    ReadKey,
+    Key,
+    ReadGlobalConfig,
+    ReadGlobal,
+    Global,
+    None,
+}
+
+pub fn end_user_source(facts: &EndUserFacts) -> EndUserSource {
+    use EndUserSource::*;
+    if facts.mode != EndUserMode::Default && facts.object_present {
+        return Own;
+    }
+    if facts.mode == EndUserMode::Models
+        && (!facts.prisma_present || facts.end_user_id.as_deref().is_none_or(str::is_empty))
+    {
+        return None;
+    }
+    if facts.key_budget_id.is_some() {
+        if !facts.key_read {
+            return ReadKey;
+        }
+        if facts.key_budget.is_some() {
+            return Key;
+        }
+    }
+    if !facts.global_config_read {
+        return ReadGlobalConfig;
+    }
+    if facts.global_budget_id.is_none() {
+        return None;
+    }
+    if !facts.global_read {
+        return ReadGlobal;
+    }
+    if facts.global_budget.is_some() {
+        Global
+    } else {
+        None
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct TokenDefaultsFacts {
+    pub current: BudgetFields,
+    pub keep_token_limits: bool,
+    pub default_present: bool,
+}
+
+#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)]
+pub enum TokenField {
+    MaxBudget,
+    Tpm,
+    Rpm,
+    Tpd,
+    Models,
+}
+
+pub fn token_assignments(facts: &TokenDefaultsFacts) -> Vec<TokenField> {
+    if !facts.default_present {
+        return vec![];
+    }
+    [
+        (TokenField::MaxBudget, facts.current.max_budget.is_none()),
+        (TokenField::Tpm, facts.current.tpm_limit.is_none()),
+        (TokenField::Rpm, facts.current.rpm_limit.is_none()),
+        (TokenField::Tpd, facts.current.tpd_limit.is_none()),
+        (
+            TokenField::Models,
+            facts.current.model_budget_count.is_none(),
+        ),
+    ]
+    .into_iter()
+    .filter_map(|(field, missing)| (!facts.keep_token_limits || missing).then_some(field))
+    .collect()
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct MembershipSnapshot {
+    pub budget_id: Option<String>,
+    pub model_budget_count: Option<usize>,
+}
+
+#[derive(Clone, Copy, Debug, Deserialize, strum::VariantArray)]
+pub enum LookupFailure {
+    Membership,
+    Default,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct MemberFacts {
+    pub user_id: Option<String>,
+    pub team_id: Option<String>,
+    pub prisma_present: bool,
+    pub default_id: Option<String>,
+    pub skip_budget_checks: bool,
+    pub membership_read: bool,
+    pub membership: Option<MembershipSnapshot>,
+    pub default_read: bool,
+    pub default_budget: Option<BudgetFields>,
+    pub lookup_failure: Option<LookupFailure>,
+}
+
+#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)]
+pub enum MemberSource {
+    Done,
+    ReadMembership,
+    ReadDefault,
+    Member,
+    Default,
+    Clear,
+    RaiseLookup,
+    SkipLookup,
+}
+
+pub fn member_source(facts: &MemberFacts) -> MemberSource {
+    use MemberSource::*;
+    if facts.user_id.is_none() || facts.team_id.is_none() || !facts.prisma_present {
+        return Done;
+    }
+    if facts.lookup_failure.is_some() {
+        return if facts.skip_budget_checks {
+            SkipLookup
+        } else {
+            RaiseLookup
+        };
+    }
+    if !facts.membership_read {
+        return ReadMembership;
+    }
+    if facts.membership.as_ref().is_some_and(|member| {
+        member.budget_id.is_some()
+            && member.budget_id != facts.default_id
+            && member.model_budget_count.is_some_and(|count| count > 0)
+    }) {
+        return Member;
+    }
+    if facts.default_id.is_none() {
+        return Done;
+    }
+    if !facts.default_read {
+        return ReadDefault;
+    }
+    if facts
+        .default_budget
+        .as_ref()
+        .is_some_and(|budget| budget.model_budget_count.is_some())
+    {
+        Default
+    } else {
+        Clear
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct IncreaseFacts {
+    pub number: Option<f64>,
+    pub truthy: bool,
+    pub type_name: String,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct TemporaryFacts {
+    pub max_budget: Option<f64>,
+    pub metadata_read: bool,
+    pub has_increase: bool,
+    pub has_expiry: bool,
+    pub expires_at: Option<i64>,
+    pub now: Option<i64>,
+    pub increase: Option<IncreaseFacts>,
+}
+
+#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)]
+pub enum TemporaryRead {
+    Metadata,
+    Expiry,
+    Now,
+    Increase,
+}
+
+#[derive(Clone, Debug, Default, Serialize, PartialEq)]
+pub struct TemporaryDecision {
+    pub read: Option<TemporaryRead>,
+    pub copy: bool,
+    pub max_budget: Option<f64>,
+    pub invalid_increase_type: Option<String>,
+}
+
+fn temp_read(read: TemporaryRead) -> TemporaryDecision {
+    TemporaryDecision {
+        read: Some(read),
+        ..TemporaryDecision::default()
+    }
+}
+
+pub fn temporary(facts: &TemporaryFacts) -> TemporaryDecision {
+    let Some(limit) = facts.max_budget else {
+        return TemporaryDecision::default();
+    };
+    if !facts.metadata_read {
+        return temp_read(TemporaryRead::Metadata);
+    }
+    if !facts.has_increase || !facts.has_expiry {
+        return TemporaryDecision::default();
+    }
+    let Some(expiry) = facts.expires_at else {
+        return temp_read(TemporaryRead::Expiry);
+    };
+    let Some(now) = facts.now else {
+        return temp_read(TemporaryRead::Now);
+    };
+    if expiry <= now {
+        return TemporaryDecision::default();
+    }
+    let Some(increase) = &facts.increase else {
+        return temp_read(TemporaryRead::Increase);
+    };
+    if !increase.truthy {
+        return TemporaryDecision::default();
+    }
+    match increase.number {
+        Some(amount) => TemporaryDecision {
+            copy: true,
+            max_budget: Some(limit + amount),
+            ..TemporaryDecision::default()
+        },
+        None => TemporaryDecision {
+            invalid_increase_type: Some(increase.type_name.clone()),
+            ..TemporaryDecision::default()
+        },
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct NumericFacts {
+    pub number: Option<f64>,
+    pub boolean: bool,
+}
+
+pub fn percentage(facts: &NumericFacts) -> Option<f64> {
+    if facts.boolean {
+        return None;
+    }
+    facts.number.filter(|pct| *pct > 0.0 && *pct <= 1.0)
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct ThrottleFacts {
+    pub opt_in: bool,
+    pub tpm_limit: Option<i64>,
+    pub rpm_limit: Option<i64>,
+    pub percentage: Option<NumericFacts>,
+}
+
+#[derive(Clone, Debug, Default, Serialize, PartialEq)]
+pub struct ThrottleDecision {
+    pub read_percentage: bool,
+    pub percentage: Option<f64>,
+}
+
+pub fn throttle(facts: &ThrottleFacts) -> ThrottleDecision {
+    if !facts.opt_in || (facts.tpm_limit.is_none() && facts.rpm_limit.is_none()) {
+        return ThrottleDecision::default();
+    }
+    match &facts.percentage {
+        None => ThrottleDecision {
+            read_percentage: true,
+            ..ThrottleDecision::default()
+        },
+        Some(raw) => ThrottleDecision {
+            percentage: percentage(raw),
+            ..ThrottleDecision::default()
+        },
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::team_empty(ModelScope::Team, Some(""), vec![0, 1])]
+    #[case::team_none(ModelScope::Team, None, vec![])]
+    #[case::enduser_empty(ModelScope::EndUser, Some(""), vec![])]
+    #[case::enduser(ModelScope::EndUser, Some("user"), vec![0, 1])]
+    fn model_identity_gates(
+        #[case] scope: ModelScope,
+        #[case] identity: Option<&str>,
+        #[case] expected: Vec<usize>,
+    ) {
+        assert_eq!(
+            model_checks(&ModelCheckFacts {
+                scope,
+                entity_id: identity.map(str::to_owned),
+                budget_count: 1,
+                model_count: 2
+            }),
+            expected
+        );
+    }
+
+    #[rstest]
+    #[case::zero(0.0, false, None)]
+    #[case::negative(-0.5, false, None)]
+    #[case::nan(f64::NAN, false, None)]
+    #[case::infinite(f64::INFINITY, false, None)]
+    #[case::one(1.0, false, Some(1.0))]
+    #[case::boolean(1.0, true, None)]
+    fn numeric_percentages(
+        #[case] number: f64,
+        #[case] boolean: bool,
+        #[case] expected: Option<f64>,
+    ) {
+        assert_eq!(
+            percentage(&NumericFacts {
+                number: Some(number),
+                boolean
+            }),
+            expected
+        );
+    }
+
+    #[rstest]
+    #[case::expired(10, 10, 3.0, false, None)]
+    #[case::before(9, 10, 3.0, false, None)]
+    #[case::active(11, 10, 3.0, true, Some(13.0))]
+    #[case::negative(11, 10, -3.0, true, Some(7.0))]
+    #[case::zero(11, 10, 0.0, false, None)]
+    fn temporary_copy_boundary(
+        #[case] expiry: i64,
+        #[case] now: i64,
+        #[case] amount: f64,
+        #[case] copy: bool,
+        #[case] expected: Option<f64>,
+    ) {
+        let facts = TemporaryFacts {
+            max_budget: Some(10.0),
+            metadata_read: true,
+            has_increase: true,
+            has_expiry: true,
+            expires_at: Some(expiry),
+            now: Some(now),
+            increase: Some(IncreaseFacts {
+                number: Some(amount),
+                truthy: amount != 0.0,
+                type_name: "float".to_owned(),
+            }),
+        };
+        let result = temporary(&facts);
+        assert_eq!(
+            (result.copy, result.max_budget, result.read),
+            (copy, expected, None)
+        );
+    }
+
+    #[rstest]
+    #[case::member_override(Some("member"), Some(1), Some("default"), false, MemberSource::Member)]
+    #[case::no_override(None, Some(1), None, false, MemberSource::Done)]
+    #[case::empty_override(
+        Some("member"),
+        Some(0),
+        Some("default"),
+        false,
+        MemberSource::ReadDefault
+    )]
+    #[case::current_default(
+        Some("default"),
+        Some(1),
+        Some("default"),
+        false,
+        MemberSource::ReadDefault
+    )]
+    #[case::missing_default(None, None, Some("default"), true, MemberSource::Clear)]
+    fn member_default_staging(
+        #[case] member_id: Option<&str>,
+        #[case] count: Option<usize>,
+        #[case] default_id: Option<&str>,
+        #[case] default_read: bool,
+        #[case] expected: MemberSource,
+    ) {
+        let facts = MemberFacts {
+            user_id: Some(String::new()),
+            team_id: Some(String::new()),
+            prisma_present: true,
+            default_id: default_id.map(str::to_owned),
+            skip_budget_checks: false,
+            membership_read: true,
+            membership: Some(MembershipSnapshot {
+                budget_id: member_id.map(str::to_owned),
+                model_budget_count: count,
+            }),
+            default_read,
+            default_budget: None,
+            lookup_failure: None,
+        };
+        assert_eq!(member_source(&facts), expected);
+    }
+
+    #[rstest]
+    #[case::replace(false, true, vec![TokenField::MaxBudget, TokenField::Tpm, TokenField::Rpm, TokenField::Tpd, TokenField::Models])]
+    #[case::keep_empty(true, true, vec![TokenField::Rpm, TokenField::Tpd])]
+    #[case::no_default(false, false, vec![])]
+    fn token_assignment_order(
+        #[case] keep: bool,
+        #[case] present: bool,
+        #[case] expected: Vec<TokenField>,
+    ) {
+        let facts = TokenDefaultsFacts {
+            current: BudgetFields {
+                max_budget: Some(0.0),
+                tpm_limit: Some(0),
+                model_budget_count: Some(0),
+                ..BudgetFields::default()
+            },
+            keep_token_limits: keep,
+            default_present: present,
+        };
+        assert_eq!(token_assignments(&facts), expected);
+    }
+}

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

@@ -1,6 +1,7 @@
 #![forbid(unsafe_code)]
 
 pub mod budget_alerts;
+pub mod budget_context;
 pub mod budget_pricing;
 pub mod common_budgets;
 pub mod credential;

litellm/_auth/budget_context.py

@@ -0,0 +1,174 @@
+from dataclasses import dataclass
+from typing import Final, Literal, TypeAlias
+
+from pydantic import TypeAdapter
+
+from litellm._auth.failure import native_boundary
+from litellm.rust_bridge import _native
+
+
+@dataclass(frozen=True, slots=True)
+class ModelCheckFacts:
+    scope: Literal["Team", "EndUser"]
+    entity_id: str | None
+    budget_count: int
+    model_count: int
+
+
+@dataclass(frozen=True, slots=True)
+class BudgetFields:
+    max_budget: float | None = None
+    tpm_limit: int | None = None
+    rpm_limit: int | None = None
+    tpd_limit: int | None = None
+    model_budget_count: int | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class EndUserFacts:
+    mode: Literal["Models", "Token", "Default"]
+    object_present: bool
+    end_user_id: str | None
+    prisma_present: bool
+    key_budget_id: str | None
+    key_read: bool = False
+    key_budget: BudgetFields | None = None
+    global_config_read: bool = False
+    global_budget_id: str | None = None
+    global_read: bool = False
+    global_budget: BudgetFields | None = None
+
+
+EndUserSource: TypeAlias = Literal["Own", "ReadKey", "Key", "ReadGlobalConfig", "ReadGlobal", "Global", "None"]
+
+
+@dataclass(frozen=True, slots=True)
+class TokenDefaultsFacts:
+    current: BudgetFields
+    keep_token_limits: bool
+    default_present: bool
+
+
+TokenField: TypeAlias = Literal["MaxBudget", "Tpm", "Rpm", "Tpd", "Models"]
+
+
+@dataclass(frozen=True, slots=True)
+class MembershipSnapshot:
+    budget_id: str | None
+    model_budget_count: int | None
+
+
+@dataclass(frozen=True, slots=True)
+class MemberFacts:
+    user_id: str | None
+    team_id: str | None
+    prisma_present: bool
+    default_id: str | None
+    skip_budget_checks: bool
+    membership_read: bool = False
+    membership: MembershipSnapshot | None = None
+    default_read: bool = False
+    default_budget: BudgetFields | None = None
+    lookup_failure: Literal["Membership", "Default"] | None = None
+
+
+MemberSource: TypeAlias = Literal[
+    "Done", "ReadMembership", "ReadDefault", "Member", "Default", "Clear", "RaiseLookup", "SkipLookup"
+]
+
+
+@dataclass(frozen=True, slots=True)
+class IncreaseFacts:
+    number: float | None
+    truthy: bool
+    type_name: str
+
+
+@dataclass(frozen=True, slots=True)
+class TemporaryFacts:
+    max_budget: float | None
+    metadata_read: bool = False
+    has_increase: bool = False
+    has_expiry: bool = False
+    expires_at: int | None = None
+    now: int | None = None
+    increase: IncreaseFacts | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class TemporaryDecision:
+    read: Literal["Metadata", "Expiry", "Now", "Increase"] | None = None
+    copy: bool = False
+    max_budget: float | None = None
+    invalid_increase_type: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class NumericFacts:
+    number: float | None
+    boolean: bool
+
+
+@dataclass(frozen=True, slots=True)
+class ThrottleFacts:
+    opt_in: bool
+    tpm_limit: int | None
+    rpm_limit: int | None
+    percentage: NumericFacts | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ThrottleDecision:
+    read_percentage: bool = False
+    percentage: float | None = None
+
+
+_MODELS: Final = TypeAdapter(ModelCheckFacts)
+_INDICES: Final = TypeAdapter(tuple[int, ...])
+_END: Final = TypeAdapter(EndUserFacts)
+_END_SOURCE: Final = TypeAdapter[EndUserSource](EndUserSource)
+_TOKEN: Final = TypeAdapter(TokenDefaultsFacts)
+_FIELDS: Final = TypeAdapter(tuple[TokenField, ...])
+_MEMBER: Final = TypeAdapter(MemberFacts)
+_MEMBER_SOURCE: Final = TypeAdapter[MemberSource](MemberSource)
+_TEMP: Final = TypeAdapter(TemporaryFacts)
+_TEMP_DECISION: Final = TypeAdapter(TemporaryDecision)
+_THROTTLE: Final = TypeAdapter(ThrottleFacts)
+_THROTTLE_DECISION: Final = TypeAdapter(ThrottleDecision)
+_NUMBER: Final = TypeAdapter(NumericFacts)
+_PERCENTAGE: Final = TypeAdapter[float | None](float | None)
+
+
+@native_boundary
+def model_checks(facts: ModelCheckFacts) -> tuple[int, ...]:
+    return _INDICES.validate_python(_native.request_auth_context_models(_MODELS.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts
+
+
+@native_boundary
+def end_user_source(facts: EndUserFacts) -> EndUserSource:
+    return _END_SOURCE.validate_python(_native.request_auth_context_end_user(_END.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts
+
+
+@native_boundary
+def token_assignments(facts: TokenDefaultsFacts) -> tuple[TokenField, ...]:
+    return _FIELDS.validate_python(_native.request_auth_context_token(_TOKEN.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts
+
+
+@native_boundary
+def member_source(facts: MemberFacts) -> MemberSource:
+    return _MEMBER_SOURCE.validate_python(_native.request_auth_context_member(_MEMBER.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts
+
+
+@native_boundary
+def temporary(facts: TemporaryFacts) -> TemporaryDecision:
+    return _TEMP_DECISION.validate_python(_native.request_auth_context_temporary(_TEMP.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts
+
+
+@native_boundary
+def percentage(facts: NumericFacts) -> float | None:
+    return _PERCENTAGE.validate_python(_native.request_auth_context_percentage(_NUMBER.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts
+
+
+@native_boundary
+def throttle(facts: ThrottleFacts) -> ThrottleDecision:
+    return _THROTTLE_DECISION.validate_python(_native.request_auth_context_throttle(_THROTTLE.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts

litellm/proxy/auth/budget_context_adapter.py

@@ -0,0 +1,377 @@
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, replace
+from typing import Final, Protocol, assert_never, cast  # noqa: TID251  # preserve identity of existing JSON mappings
+
+from opentelemetry.trace import Span
+
+import litellm
+from litellm._auth import budget_context as core
+from litellm._auth.failure import NativeAuthFailure
+from litellm._logging import verbose_proxy_logger
+from litellm.models.budget import LiteLLM_BudgetTable
+from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+
+
+class TeamModelLimiter(Protocol):
+    async def is_team_within_model_budget(
+        self,
+        team_id: str,
+        team_model_max_budget: Mapping[str, object],
+        key_model_max_budget: Mapping[str, object] | None,
+        model: str,
+    ) -> bool: ...
+
+
+class EndUserModelLimiter(Protocol):
+    async def is_end_user_within_model_budget(
+        self,
+        end_user_id: str,
+        end_user_model_max_budget: Mapping[str, object],
+        model: str,
+        team_id: str | None = None,
+    ) -> bool: ...
+
+
+def mapping_value(value: object) -> Mapping[str, object] | None:
+    if isinstance(value, Mapping):
+        return cast(Mapping[str, object], value)  # cast-ok: preserve the original JSON map's identity
+    return None
+
+
+def mapping_count(value: object) -> int | None:
+    mapping: Final = mapping_value(value)
+    return len(mapping) if mapping is not None else None
+
+
+def budget_fields(row: LiteLLM_BudgetTable | None) -> core.BudgetFields | None:
+    if row is None:
+        return None
+    return core.BudgetFields(
+        row.max_budget,
+        row.tpm_limit,
+        row.rpm_limit,
+        row.tpd_limit,
+        mapping_count(getattr(row, "model_max_budget", None)),
+    )  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # legacy budget table declares an unparameterized dict
+
+
+async def check_team_model_budget(
+    valid_token: UserAPIKeyAuth,
+    model_max_budget_limiter: TeamModelLimiter,
+    models: Sequence[str],
+) -> None:
+    budgets: Final = valid_token.team_model_max_budget
+    selected: Final = core.model_checks(
+        core.ModelCheckFacts("Team", valid_token.team_id, len(budgets) if budgets is not None else 0, len(models))
+    )
+    for index in selected:
+        assert valid_token.team_id is not None and budgets is not None
+        await model_max_budget_limiter.is_team_within_model_budget(
+            team_id=valid_token.team_id,
+            team_model_max_budget=budgets,
+            key_model_max_budget=mapping_value(getattr(valid_token, "model_max_budget", None)),
+            model=models[index],
+        )
+
+
+async def check_end_user_model_budget(
+    end_user_id: str | None,
+    end_user_model_max_budget: Mapping[str, object] | None,
+    model_max_budget_limiter: EndUserModelLimiter,
+    models: Sequence[str],
+    team_id: str | None,
+) -> None:
+    selected: Final = core.model_checks(
+        core.ModelCheckFacts(
+            "EndUser",
+            end_user_id,
+            len(end_user_model_max_budget) if end_user_model_max_budget is not None else 0,
+            len(models),
+        )
+    )
+    for index in selected:
+        assert end_user_id is not None and end_user_model_max_budget is not None
+        await model_max_budget_limiter.is_end_user_within_model_budget(
+            end_user_id=end_user_id,
+            end_user_model_max_budget=end_user_model_max_budget,
+            model=models[index],
+            team_id=team_id,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class EndUserIO:
+    prisma_client: PrismaClient | None
+    cache: UserApiKeyCache
+    span: Span | None
+
+
+async def _end_user_row(
+    facts: core.EndUserFacts,
+    io: EndUserIO,
+    own: LiteLLM_EndUserTable | None,
+    key: LiteLLM_BudgetTable | None = None,
+    proxy: LiteLLM_BudgetTable | None = None,
+) -> LiteLLM_BudgetTable | None:
+    from litellm.proxy.auth.auth_checks import get_default_end_user_budget
+
+    decision: Final = core.end_user_source(facts)
+    match decision:
+        case "Own":
+            assert own is not None
+            return own.litellm_budget_table
+        case "Key":
+            return key
+        case "Global":
+            return proxy
+        case "None":
+            return None
+        case "ReadKey":
+            row: Final = await get_default_end_user_budget(
+                prisma_client=io.prisma_client,
+                user_api_key_cache=io.cache,
+                parent_otel_span=io.span,
+                budget_id=facts.key_budget_id,
+            )
+            return await _end_user_row(
+                replace(facts, key_read=True, key_budget=budget_fields(row)), io, own, row, proxy
+            )
+        case "ReadGlobalConfig":
+            return await _end_user_row(
+                replace(facts, global_config_read=True, global_budget_id=litellm.max_end_user_budget_id),
+                io,
+                own,
+                key,
+                proxy,
+            )
+        case "ReadGlobal":
+            row_global: Final = await get_default_end_user_budget(
+                prisma_client=io.prisma_client,
+                user_api_key_cache=io.cache,
+                parent_otel_span=io.span,
+                budget_id=facts.global_budget_id,
+            )
+            return await _end_user_row(
+                replace(facts, global_read=True, global_budget=budget_fields(row_global)), io, own, key, row_global
+            )
+        case _:
+            assert_never(decision)
+
+
+async def resolve_default_end_user_budget(
+    prisma_client: PrismaClient,
+    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
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span),
+        None,
+    )
+
+
+async def end_user_model_max_budget(
+    end_user_id: str | None,
+    end_user_object: LiteLLM_EndUserTable | None,
+    key_end_user_budget_id: str | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None,
+) -> Mapping[str, object] | None:
+    row: Final = await _end_user_row(
+        core.EndUserFacts(
+            "Models", end_user_object is not None, end_user_id, prisma_client is not None, key_end_user_budget_id
+        ),
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span),
+        end_user_object,
+    )
+    return getattr(row, "model_max_budget", None) if row is not None else None  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # legacy row field
+
+
+async def apply_key_end_user_default_budget_to_token(
+    valid_token: UserAPIKeyAuth,
+    end_user_object: LiteLLM_EndUserTable | None,
+    key_end_user_budget_id: str,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: Span | None,
+    keep_token_limits: bool,
+) -> None:
+    row: Final = await _end_user_row(
+        core.EndUserFacts("Token", end_user_object is not None, valid_token.end_user_id, True, key_end_user_budget_id),
+        EndUserIO(prisma_client, user_api_key_cache, parent_otel_span),
+        end_user_object,
+    )
+    current: Final = core.BudgetFields(
+        valid_token.end_user_max_budget,
+        valid_token.end_user_tpm_limit,
+        valid_token.end_user_rpm_limit,
+        valid_token.end_user_tpd_limit,
+        mapping_count(valid_token.end_user_model_max_budget),
+    )
+    for field in core.token_assignments(core.TokenDefaultsFacts(current, keep_token_limits, row is not None)):
+        assert row is not None
+        match field:
+            case "MaxBudget":
+                valid_token.end_user_max_budget = (  # rebind-ok: request-local native projection
+                    row.max_budget
+                )
+            case "Tpm":
+                valid_token.end_user_tpm_limit = (  # rebind-ok: request-local native projection
+                    row.tpm_limit
+                )
+            case "Rpm":
+                valid_token.end_user_rpm_limit = (  # rebind-ok: request-local native projection
+                    row.rpm_limit
+                )
+            case "Tpd":
+                valid_token.end_user_tpd_limit = (  # rebind-ok: request-local native projection
+                    row.tpd_limit
+                )
+            case "Models":
+                valid_token.end_user_model_max_budget = row.model_max_budget  # pyright: ignore[reportUnknownMemberType]  # preserve the validated row dict and its identity  # rebind-ok: native-selected request-local budget assignment
+            case _:
+                assert_never(field)
+
+
+@dataclass(frozen=True, slots=True)
+class MemberIO:
+    prisma_client: PrismaClient | None
+    cache: UserApiKeyCache
+    logging: ProxyLogging | None
+
+
+@dataclass(frozen=True, slots=True)
+class MemberRead:
+    member: LiteLLM_TeamMembership | None = None
+    default: LiteLLM_BudgetTable | None = None
+    error: Exception | None = None
+
+
+async def _read_member(facts: core.MemberFacts, io: MemberIO) -> MemberRead:
+    from litellm.proxy.auth.auth_checks import get_team_membership
+
+    assert facts.user_id is not None and facts.team_id is not None
+    try:
+        row: Final = await get_team_membership(
+            user_id=facts.user_id,
+            team_id=facts.team_id,
+            prisma_client=io.prisma_client,
+            user_api_key_cache=io.cache,
+            proxy_logging_obj=io.logging,
+        )
+    except NativeAuthFailure:
+        raise
+    except Exception as error:  # noqa: BLE001  # only the database/cache read is caught; native failures rethrow
+        return MemberRead(error=error)
+    return MemberRead(member=row)
+
+
+async def _read_member_default(facts: core.MemberFacts, io: MemberIO) -> MemberRead:
+    from litellm.proxy.auth.auth_checks import get_team_member_default_budget
+
+    assert facts.default_id is not None
+    try:
+        row: Final = await get_team_member_default_budget(
+            budget_id=facts.default_id,
+            prisma_client=io.prisma_client,
+            user_api_key_cache=io.cache,
+            raise_on_lookup_error=True,
+        )
+    except NativeAuthFailure:
+        raise
+    except Exception as error:  # noqa: BLE001  # only the database/cache read is caught; native failures rethrow
+        return MemberRead(error=error)
+    return MemberRead(default=row)
+
+
+async def _member_model_budget(
+    token: UserAPIKeyAuth,
+    facts: core.MemberFacts,
+    io: MemberIO,
+    reading: MemberRead,
+) -> None:
+    decision: Final = core.member_source(facts)
+    match decision:
+        case "Done":
+            return
+        case "RaiseLookup":
+            assert reading.error is not None
+            raise reading.error
+        case "SkipLookup":
+            verbose_proxy_logger.warning(
+                "Skipping team-member model budget lookup after a lookup error", exc_info=reading.error
+            )
+            return
+        case "ReadMembership":
+            member_read: Final = await _read_member(facts, io)
+            member: Final = member_read.member
+            member_row: Final = member.litellm_budget_table if member is not None else None
+            snapshot: Final = (
+                core.MembershipSnapshot(member.budget_id, mapping_count(getattr(member_row, "model_max_budget", None)))
+                if member is not None
+                else None
+            )
+            await _member_model_budget(
+                token,
+                replace(
+                    facts,
+                    membership_read=True,
+                    membership=snapshot,
+                    lookup_failure="Membership" if member_read.error is not None else None,
+                ),
+                io,
+                member_read,
+            )
+        case "ReadDefault":
+            default_read: Final = await _read_member_default(facts, io)
+            await _member_model_budget(
+                token,
+                replace(
+                    facts,
+                    default_read=True,
+                    default_budget=budget_fields(default_read.default),
+                    lookup_failure="Default" if default_read.error is not None else None,
+                ),
+                io,
+                default_read,
+            )
+        case "Member":
+            assert reading.member is not None
+            token.team_member_model_max_budget = mapping_value(  # rebind-ok: request-local native projection
+                getattr(reading.member.litellm_budget_table, "model_max_budget", None)
+            )
+        case "Default":
+            assert reading.default is not None
+            token.team_member_model_max_budget = mapping_value(  # rebind-ok: request-local native projection
+                getattr(reading.default, "model_max_budget", None)
+            )
+        case "Clear":
+            token.team_member_model_max_budget = None  # rebind-ok: request-local native projection
+        case _:
+            assert_never(decision)
+
+
+async def load_team_member_default_model_budget(
+    user_api_key_auth_obj: UserAPIKeyAuth,
+    team_object: LiteLLM_TeamTable | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    skip_budget_checks: bool,
+    proxy_logging_obj: ProxyLogging | None,
+) -> None:
+    metadata: Final = mapping_value(getattr(team_object, "metadata", None))
+    default_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None
+    facts: Final = core.MemberFacts(
+        user_id=user_api_key_auth_obj.user_id if isinstance(user_api_key_auth_obj.user_id, str) else None,
+        team_id=user_api_key_auth_obj.team_id if isinstance(user_api_key_auth_obj.team_id, str) else None,
+        prisma_present=prisma_client is not None,
+        default_id=default_id if isinstance(default_id, str) else None,
+        skip_budget_checks=skip_budget_checks,
+    )
+    await _member_model_budget(
+        user_api_key_auth_obj, facts, MemberIO(prisma_client, user_api_key_cache, proxy_logging_obj), MemberRead()
+    )

litellm/proxy/auth/key_budget_context_adapter.py

@@ -0,0 +1,95 @@
+from collections.abc import Mapping
+from dataclasses import replace
+from datetime import datetime, timezone
+from typing import Final, assert_never
+
+import litellm
+from litellm._auth import budget_context as core
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.budget_context_adapter import mapping_value
+
+
+def numeric_facts(value: object) -> core.NumericFacts:
+    return core.NumericFacts(float(value) if isinstance(value, (int, float)) else None, isinstance(value, bool))
+
+
+def budget_throttle_percentage() -> float | None:
+    return core.percentage(numeric_facts(litellm.budget_exceeded_throttle_percentage))
+
+
+def read_key_throttle_facts(valid_token: UserAPIKeyAuth) -> core.ThrottleFacts:
+    metadata: Final = mapping_value(getattr(valid_token, "metadata", None))
+    facts: Final = core.ThrottleFacts(
+        (metadata or {}).get("throttle_on_budget_exceeded") is True,
+        valid_token.tpm_limit,
+        valid_token.rpm_limit,
+    )
+    decision: Final = core.throttle(facts)
+    return (
+        replace(facts, percentage=numeric_facts(litellm.budget_exceeded_throttle_percentage))
+        if decision.read_percentage
+        else facts
+    )
+
+
+def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool:
+    return core.throttle(read_key_throttle_facts(valid_token)).percentage is not None
+
+
+def inspect_key_budget_throttle(valid_token: UserAPIKeyAuth) -> float | None:
+    decision: Final = core.throttle(read_key_throttle_facts(valid_token))
+    return budget_throttle_percentage() if decision.percentage is not None else None
+
+
+def _microseconds(value: datetime) -> int:
+    delta: Final = value - datetime(1970, 1, 1, tzinfo=timezone.utc)
+    return (delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds
+
+
+def _temporary(
+    token: UserAPIKeyAuth,
+    facts: core.TemporaryFacts,
+    metadata: Mapping[str, object] | None,
+) -> UserAPIKeyAuth:
+    decision: Final = core.temporary(facts)
+    if decision.invalid_increase_type is not None:
+        raise TypeError(
+            f"unsupported operand type(s) for +: '{type(token.max_budget).__name__}' and '{decision.invalid_increase_type}'"
+        )
+    if decision.copy:
+        return token.model_copy(update={"max_budget": decision.max_budget})
+    match decision.read:
+        case None:
+            return token
+        case "Metadata":
+            raw: Final = mapping_value(getattr(token, "metadata", None))
+            return _temporary(
+                token,
+                replace(
+                    facts,
+                    metadata_read=True,
+                    has_increase=raw is not None and "temp_budget_increase" in raw,
+                    has_expiry=raw is not None and "temp_budget_expiry" in raw,
+                ),
+                raw,
+            )
+        case "Expiry":
+            assert metadata is not None
+            expiry: Final = datetime.fromisoformat(metadata["temp_budget_expiry"])  # pyright: ignore[reportArgumentType]  # preserve datetime's public TypeError for malformed metadata
+            aware: Final = expiry.replace(tzinfo=timezone.utc) if expiry.tzinfo is None else expiry
+            return _temporary(token, replace(facts, expires_at=_microseconds(aware)), metadata)
+        case "Now":
+            return _temporary(token, replace(facts, now=_microseconds(datetime.now(timezone.utc))), metadata)
+        case "Increase":
+            assert metadata is not None
+            value: Final = metadata["temp_budget_increase"]
+            raw_increase: Final = core.IncreaseFacts(
+                float(value) if isinstance(value, (int, float)) else None, bool(value), type(value).__name__
+            )
+            return _temporary(token, replace(facts, increase=raw_increase), metadata)
+        case _:
+            assert_never(decision.read)
+
+
+def update_key_budget_with_temp_budget_increase(valid_token: UserAPIKeyAuth) -> UserAPIKeyAuth:
+    return _temporary(valid_token, core.TemporaryFacts(valid_token.max_budget), None)

litellm/rust_bridge/_native.pyi

@@ -812,3 +812,17 @@ class RouterCall(Generic[_S]):
 def request_auth_budget_exemption(facts: object) -> dict[str, object]: ...
 
 def request_auth_budget_alert_plan(facts: object) -> dict[str, object]: ...
+
+def request_auth_context_models(facts: object) -> object: ...
+
+def request_auth_context_end_user(facts: object) -> object: ...
+
+def request_auth_context_token(facts: object) -> object: ...
+
+def request_auth_context_member(facts: object) -> object: ...
+
+def request_auth_context_temporary(facts: object) -> object: ...
+
+def request_auth_context_throttle(facts: object) -> object: ...
+
+def request_auth_context_percentage(facts: object) -> float | None: ...

tests/unit/proxy/auth/test_budget_context_adapter.py

@@ -0,0 +1,327 @@
+from collections.abc import Mapping
+from typing import Final, Literal
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+import litellm
+from litellm._auth import budget_context as core
+from litellm._auth.failure import NativeAuthRuntimeError
+from litellm.models.budget import LiteLLM_BudgetTable
+from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth
+from litellm.proxy.auth import auth_checks
+from litellm.proxy.auth import budget_context_adapter as adapter
+from litellm.proxy.auth import user_api_key_auth as legacy
+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]  # fixture parameters are the two literal booleans above
+
+
+@pytest.mark.parametrize("scope", ("Team", "EndUser"))
+@pytest.mark.parametrize("identity", (None, "", "entity"))
+@pytest.mark.parametrize("budget_count", (0, 1, 2))
+@pytest.mark.parametrize("model_count", (0, 1, 3))
+def test_native_model_read_selection(
+    scope: Literal["Team", "EndUser"], identity: str | None, budget_count: int, model_count: int
+) -> None:
+    facts: Final = core.ModelCheckFacts(scope, identity, budget_count, model_count)
+    expected: Final = (
+        tuple(range(model_count)) if budget_count and identity is not None and (scope == "Team" or identity) else ()
+    )
+    assert core.model_checks(facts) == expected
+
+
+@pytest.mark.parametrize("keep", (False, True))
+@pytest.mark.parametrize("present", (False, True))
+@pytest.mark.parametrize("limit", (None, 0, -1, 1))
+@pytest.mark.parametrize("models", (None, 0, 2))
+def test_native_token_assignment_table(keep: bool, present: bool, limit: int | None, models: int | None) -> None:
+    facts: Final = core.TokenDefaultsFacts(core.BudgetFields(limit, limit, limit, limit, models), keep, present)
+    expected: Final = tuple(
+        name
+        for name, value in (("MaxBudget", limit), ("Tpm", limit), ("Rpm", limit), ("Tpd", limit), ("Models", models))
+        if present and (not keep or value is None)
+    )
+    assert core.token_assignments(facts) == expected
+
+
+@pytest.mark.parametrize("mode", ("Models", "Token", "Default"))
+@pytest.mark.parametrize("object_present", (False, True))
+@pytest.mark.parametrize("identity", (None, "", "eu"))
+@pytest.mark.parametrize("db", (False, True))
+@pytest.mark.parametrize("key_id", (None, "", "key"))
+def test_end_user_initial_table(
+    mode: Literal["Models", "Token", "Default"],
+    object_present: bool,
+    identity: str | None,
+    db: bool,
+    key_id: str | None,
+) -> None:
+    facts: Final = core.EndUserFacts(mode, object_present, identity, db, key_id)
+    expected: Final = (
+        "Own"
+        if mode != "Default" and object_present
+        else "None"
+        if mode == "Models" and (not identity or not db)
+        else "ReadKey"
+        if key_id is not None
+        else "ReadGlobalConfig"
+    )
+    assert core.end_user_source(facts) == expected
+
+
+@pytest.mark.parametrize("key_row", (False, True))
+@pytest.mark.parametrize("global_id", (None, "", "proxy"))
+@pytest.mark.parametrize("global_read", (False, True))
+@pytest.mark.parametrize("global_row", (False, True))
+def test_end_user_fallback_stages(key_row: bool, global_id: str | None, global_read: bool, global_row: bool) -> None:
+    facts: Final = core.EndUserFacts(
+        "Default",
+        False,
+        None,
+        True,
+        "key",
+        key_read=True,
+        key_budget=core.BudgetFields() if key_row else None,
+        global_config_read=True,
+        global_budget_id=global_id,
+        global_read=global_read,
+        global_budget=core.BudgetFields() if global_row else None,
+    )
+    expected: Final = (
+        "Key"
+        if key_row
+        else "None"
+        if global_id is None
+        else "ReadGlobal"
+        if not global_read
+        else "Global"
+        if global_row
+        else "None"
+    )
+    assert core.end_user_source(facts) == expected
+
+
+@pytest.mark.parametrize("default_id", (None, "", "default"))
+@pytest.mark.parametrize("member_id", (None, "", "member", "default"))
+@pytest.mark.parametrize("member_models", (None, 0, 2))
+@pytest.mark.parametrize("default_models", (None, 0, 2))
+@pytest.mark.parametrize("skip", (False, True))
+def test_member_snapshot_table(
+    default_id: str | None, member_id: str | None, member_models: int | None, default_models: int | None, skip: bool
+) -> None:
+    facts: Final = core.MemberFacts(
+        "",
+        "",
+        True,
+        default_id,
+        skip,
+        membership_read=True,
+        membership=core.MembershipSnapshot(member_id, member_models),
+        default_read=True,
+        default_budget=core.BudgetFields(model_budget_count=default_models),
+    )
+    expected: Final = (
+        "Member"
+        if member_id is not None and member_id != default_id and member_models
+        else "Done"
+        if default_id is None
+        else "Default"
+        if default_models is not None
+        else "Clear"
+    )
+    assert core.member_source(facts) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("identity", (None, "", "entity"))
+@pytest.mark.parametrize("scope", ("Team", "EndUser"))
+@pytest.mark.parametrize("fail", (False, True))
+async def test_selected_model_reads_keep_order_and_stop_on_failure(
+    native: bool, identity: str | None, scope: Literal["Team", "EndUser"], fail: bool
+) -> None:
+    failure: Final = RuntimeError("counter failed")
+    read: Final = AsyncMock(side_effect=(True, failure, True) if fail else None, return_value=True)
+    limiter: Final = Mock()
+    limiter.is_team_within_model_budget = read
+    limiter.is_end_user_within_model_budget = read
+    caps: Final[Mapping[str, object]] = {"a": {"max_budget": 1}}
+    token: Final = UserAPIKeyAuth(team_id=identity, team_model_max_budget=dict(caps))
+    models: Final = ["a", "a", "b"]
+    coroutine: Final = (
+        (adapter.check_team_model_budget if native else legacy._check_team_model_budget)(token, limiter, models)  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+        if scope == "Team"
+        else (adapter.check_end_user_model_budget if native else legacy._check_end_user_model_budget)(  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+            identity, caps, limiter, models, "team"
+        )
+    )
+    active: Final = identity is not None and (scope == "Team" or bool(identity))
+    if fail and active:
+        with pytest.raises(RuntimeError, match="counter failed") as raised:
+            await coroutine
+        assert raised.value is failure
+    else:
+        await coroutine
+    assert read.await_count == (2 if fail else 3) if active else read.await_count == 0
+    assert tuple(item.kwargs["model"] for item in read.await_args_list) == tuple(models[: read.await_count])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("own", (False, True))
+@pytest.mark.parametrize("own_row", (False, True))
+@pytest.mark.parametrize("identity", (None, "", "eu"))
+@pytest.mark.parametrize("key_found", (False, True))
+async def test_end_user_source_read_order_and_identity(
+    native: bool, monkeypatch: pytest.MonkeyPatch, own: bool, own_row: bool, identity: str | None, key_found: bool
+) -> None:
+    budget: Final = LiteLLM_BudgetTable(model_max_budget={"m": {"max_budget": 1}})
+    end_user: Final = (
+        LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if own_row else None)
+        if own
+        else None
+    )
+    read: Final = AsyncMock(side_effect=(budget,) if key_found else (None, budget))
+    monkeypatch.setattr(auth_checks, "get_default_end_user_budget", read)
+    monkeypatch.setattr(litellm, "max_end_user_budget_id", "proxy")
+    fn: Final = adapter.end_user_model_max_budget if native else legacy._end_user_model_max_budget  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+    value: Final = await fn(identity, end_user, "key", Mock(), UserApiKeyCache(), None)
+    expected: Final = getattr(budget, "model_max_budget", None) if (own and own_row) or (not own and identity) else None
+    assert value is expected
+    assert read.await_count == (0 if own or not identity else 1 if key_found else 2)
+    if read.await_count:
+        assert read.await_args_list[0].kwargs["budget_id"] == "key"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("keep", (False, True))
+@pytest.mark.parametrize("present", (False, True))
+async def test_token_defaults_preserve_caps_and_projection_identity(native: bool, keep: bool, present: bool) -> None:
+    budget: Final = LiteLLM_BudgetTable(
+        max_budget=9, tpm_limit=8, rpm_limit=7, tpd_limit=6, model_max_budget={"m": {"max_budget": 5}}
+    )
+    token: Final = UserAPIKeyAuth(
+        end_user_max_budget=0 if present else None,
+        end_user_tpm_limit=0 if present else None,
+        end_user_model_max_budget={} if present else None,
+    )
+    original_models: Final = token.end_user_model_max_budget
+    own: Final = LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget)
+    fn: Final = (
+        adapter.apply_key_end_user_default_budget_to_token
+        if native
+        else legacy._apply_key_end_user_default_budget_to_token  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+    )
+    await fn(token, own, "key", Mock(), UserApiKeyCache(), None, keep)
+    assert token.end_user_max_budget == (0 if keep and present else 9)
+    assert token.end_user_tpm_limit == (0 if keep and present else 8)
+    assert (token.end_user_rpm_limit, token.end_user_tpd_limit) == (7, 6)
+    assert token.end_user_model_max_budget is (
+        original_models if keep and present else getattr(budget, "model_max_budget", None)
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("default_id", (None, "", "default"))
+@pytest.mark.parametrize("member_id", (None, "default", "override"))
+@pytest.mark.parametrize("member_caps", (None, {}, {"member": {"max_budget": 1}}))
+@pytest.mark.parametrize("default_caps", (None, {}, {"default": {"max_budget": 2}}))
+async def test_member_default_read_and_mutation_order(
+    native: bool,
+    monkeypatch: pytest.MonkeyPatch,
+    default_id: str | None,
+    member_id: str | None,
+    member_caps: Mapping[str, object] | None,
+    default_caps: Mapping[str, object] | None,
+) -> None:
+    member_budget: Final = LiteLLM_BudgetTable(model_max_budget=dict(member_caps) if member_caps is not None else None)
+    default_budget: Final = LiteLLM_BudgetTable(
+        model_max_budget=dict(default_caps) if default_caps is not None else None
+    )
+    membership: Final = LiteLLM_TeamMembership(
+        user_id="", team_id="", budget_id=member_id, litellm_budget_table=member_budget
+    )
+    member_read: Final = AsyncMock(return_value=membership)
+    default_read: Final = AsyncMock(return_value=default_budget)
+    monkeypatch.setattr(auth_checks, "get_team_membership", member_read)
+    monkeypatch.setattr(legacy, "get_team_membership", member_read)
+    monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read)
+    monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read)
+    token: Final = UserAPIKeyAuth(user_id="", team_id="", team_member_model_max_budget={"old": {}})
+    prior: Final = token.team_member_model_max_budget
+    team: Final = LiteLLM_TeamTable(team_id="", metadata={"team_member_budget_id": default_id})
+    fn: Final = (
+        adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+    )
+    await fn(token, team, Mock(), UserApiKeyCache(), False, None)
+    override: Final = member_id is not None and member_id != default_id and bool(member_caps)
+    expected: Final = (
+        getattr(member_budget, "model_max_budget", None)
+        if override
+        else getattr(default_budget, "model_max_budget", None)
+        if default_id is not None
+        else prior
+    )
+    assert token.team_member_model_max_budget is expected
+    member_read.assert_awaited_once()
+    assert default_read.await_count == (0 if override or default_id is None else 1)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("skip", (False, True))
+@pytest.mark.parametrize("during_default", (False, True))
+async def test_member_lookup_error_policy(
+    native: bool, monkeypatch: pytest.MonkeyPatch, skip: bool, during_default: bool
+) -> None:
+    error: Final = RuntimeError("budget read failed")
+    member_read: Final = AsyncMock(return_value=None, side_effect=None if during_default else error)
+    default_read: Final = AsyncMock(side_effect=error)
+    monkeypatch.setattr(auth_checks, "get_team_membership", member_read)
+    monkeypatch.setattr(legacy, "get_team_membership", member_read)
+    monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read)
+    monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read)
+    token: Final = UserAPIKeyAuth(user_id="u", team_id="t", team_member_model_max_budget={"old": {}})
+    prior: Final = token.team_member_model_max_budget
+    team: Final = LiteLLM_TeamTable(team_id="t", metadata={"team_member_budget_id": "default"})
+    fn: Final = (
+        adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+    )
+    pending: Final = fn(token, team, Mock(), UserApiKeyCache(), skip, None)
+    if skip:
+        await pending
+    else:
+        with pytest.raises(RuntimeError, match="budget read failed") as raised:
+            await pending
+        assert raised.value is error
+    assert token.team_member_model_max_budget is prior
+
+
+@pytest.mark.asyncio
+async def test_member_native_failure_never_becomes_skippable_lookup_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    error: Final = NativeAuthRuntimeError("native planner failed")
+    read: Final = AsyncMock(side_effect=error)
+    monkeypatch.setattr(auth_checks, "get_team_membership", read)
+    with pytest.raises(NativeAuthRuntimeError, match="native planner failed") as raised:
+        await adapter.load_team_member_default_model_budget(
+            UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None
+        )
+    assert raised.value is error
+    read.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_member_planner_failure_stops_before_io(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.rust_bridge import _native
+
+    native: Final = Mock(side_effect=RuntimeError("decode failed"))
+    read: Final = AsyncMock()
+    monkeypatch.setattr(_native, "request_auth_context_member", native)
+    monkeypatch.setattr(auth_checks, "get_team_membership", read)
+    with pytest.raises(NativeAuthRuntimeError, match="decode failed"):
+        await adapter.load_team_member_default_model_budget(
+            UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None
+        )
+    read.assert_not_awaited()
+    native.assert_called_once()

tests/unit/proxy/auth/test_key_budget_context_adapter.py

@@ -0,0 +1,108 @@
+import math
+from datetime import datetime, timezone
+from typing import Final
+from unittest.mock import Mock
+
+import pytest
+
+import litellm
+from litellm._auth import budget_context as core
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth import budget_throttle
+from litellm.proxy.auth import key_budget_context_adapter as adapter
+from litellm.proxy.auth import user_api_key_auth as legacy
+
+_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc)
+
+
+@pytest.mark.parametrize("native", (False, True), ids=("reference", "native"))
+@pytest.mark.parametrize(
+    "increase", (None, False, True, 0, -3, 3, float("nan"), float("inf"), "", "3", [], [3], {}, {"n": 3})
+)
+@pytest.mark.parametrize(
+    "expiry", ("2024-12-31", "2025-01-01T00:00:00+00:00", "2025-01-01T00:00:00.000001Z", "2026-01-01T09:00:00+09:00")
+)
+@pytest.mark.parametrize("limit", (None, 0, 10))
+def test_temporary_budget_reference_native(
+    native: bool, monkeypatch: pytest.MonkeyPatch, increase: object, expiry: str, limit: float | None
+) -> None:
+    now: Final = Mock(return_value=_NOW)
+    clock: Final = Mock(wraps=datetime, now=now)
+    monkeypatch.setattr(adapter, "datetime", clock)
+    monkeypatch.setattr(legacy, "datetime", clock)
+    token: Final = UserAPIKeyAuth(
+        max_budget=limit, metadata={"temp_budget_increase": increase, "temp_budget_expiry": expiry}
+    )
+    fn: Final = (
+        adapter.update_key_budget_with_temp_budget_increase
+        if native
+        else legacy._update_key_budget_with_temp_budget_increase  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+    )
+    active: Final = expiry.startswith("2026") or expiry == "2025-01-01T00:00:00.000001Z"
+    changes: Final = limit is not None and active and bool(increase)
+    if changes and not isinstance(increase, (int, float)):
+        with pytest.raises(TypeError, match="unsupported operand type"):
+            fn(token)
+        return
+    result: Final = fn(token)
+    assert (result is token) is not changes
+    if changes:
+        assert limit is not None and isinstance(increase, (int, float))
+        expected: Final = limit + increase
+        assert result.max_budget is not None
+        assert result.max_budget == expected or (math.isnan(result.max_budget) and math.isnan(expected))
+    assert token.max_budget == limit
+    assert now.call_count == (0 if limit is None else 1)
+
+
+@pytest.mark.parametrize("native", (False, True), ids=("reference", "native"))
+@pytest.mark.parametrize("expiry", (None, 5, "", "invalid"))
+@pytest.mark.parametrize("increase_present", (False, True))
+@pytest.mark.parametrize("limit", (None, 1))
+def test_temporary_metadata_error_timing(
+    native: bool, monkeypatch: pytest.MonkeyPatch, expiry: object, increase_present: bool, limit: float | None
+) -> None:
+    token: Final = UserAPIKeyAuth(
+        max_budget=limit,
+        metadata={"temp_budget_expiry": expiry, **({"temp_budget_increase": 0} if increase_present else {})},
+    )
+    now: Final = Mock(return_value=_NOW)
+    clock: Final = Mock(wraps=datetime, now=now)
+    monkeypatch.setattr(adapter, "datetime", clock)
+    monkeypatch.setattr(legacy, "datetime", clock)
+    fn: Final = (
+        adapter.update_key_budget_with_temp_budget_increase
+        if native
+        else legacy._update_key_budget_with_temp_budget_increase  # pyright: ignore[reportPrivateUsage]  # Phase 1 leaf is the unchanged reference implementation
+    )
+    if limit is not None and increase_present:
+        with pytest.raises(ValueError if isinstance(expiry, str) else TypeError):
+            fn(token)
+    else:
+        assert fn(token) is token
+    now.assert_not_called()
+
+
+@pytest.mark.parametrize("pct", (None, False, True, "", "0.5", 0, -0.5, 0.25, 1, 2, float("nan"), float("inf")))
+@pytest.mark.parametrize("opt_in", (None, False, True, 1, "true"))
+@pytest.mark.parametrize("limit", (None, 0, 1))
+def test_throttle_reference_native(
+    monkeypatch: pytest.MonkeyPatch, pct: object, opt_in: object, limit: int | None
+) -> None:
+    monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", pct)
+    token: Final = UserAPIKeyAuth(metadata={"throttle_on_budget_exceeded": opt_in}, tpm_limit=limit)
+    assert adapter.budget_throttle_percentage() == budget_throttle.budget_throttle_percentage()
+    assert adapter.should_throttle_budget_exceeded(token) == budget_throttle.should_throttle_budget_exceeded(token)
+    expected: Final = (
+        budget_throttle.budget_throttle_percentage() if budget_throttle.should_throttle_budget_exceeded(token) else None
+    )
+    assert adapter.inspect_key_budget_throttle(token) == expected
+
+
+@pytest.mark.parametrize("tpm", (None, 0, 1))
+@pytest.mark.parametrize("rpm", (None, 0, 1))
+@pytest.mark.parametrize("opt_in", (False, True))
+def test_throttle_native_selects_config_read(tpm: int | None, rpm: int | None, opt_in: bool) -> None:
+    assert core.throttle(core.ThrottleFacts(opt_in, tpm, rpm)).read_percentage is (
+        opt_in and (tpm is not None or rpm is not None)
+    )

tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py

@@ -131,26 +131,29 @@ def test_resolve_model_budget():
     assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None
 
 
-# Test is_key_within_model_budget
 @pytest.mark.asyncio
-async def test_is_key_within_model_budget(budget_limiter):
-    # Mock user API key dict
-    user_api_key = UserAPIKeyAuth(
+@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0))
+async def test_is_key_within_model_budget(
+    budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter,
+    spend: float,
+) -> None:
+    user_api_key: Final = UserAPIKeyAuth(
         token="test-key",
         key_alias="test-alias",
         model_max_budget={"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}},
     )
-
-    # Test when model is within budget
-    with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0):
-        assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True
-
-    # Test when model exceeds budget
-    with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0):
+    counter: Final = model_budget_spend_cache_key(
+        entity_type=Litellm_EntityType.KEY,
+        entity_id="test-key",
+        budget_model="gpt-4",
+        budget_duration="1d",
+    )
+    await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend)
+    if spend >= 100.0:
         with pytest.raises(litellm.BudgetExceededError):
             await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4")
-
-    # Test model not in budget config
+    else:
+        assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True
     assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True
 
 
@@ -224,38 +227,26 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim
         assert call_kwargs["response_cost"] == 0.05
 
 
-# Test is_end_user_within_model_budget
 @pytest.mark.asyncio
-async def test_is_end_user_within_model_budget(budget_limiter):
-    # Test when model is within budget
-    with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0):
-        assert (
-            await budget_limiter.is_end_user_within_model_budget(
-                "test-user",
-                {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}},
-                "gpt-4",
-            )
-            is True
-        )
-
-    # Test when model exceeds budget
-    with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0):
-        with pytest.raises(litellm.BudgetExceededError):
-            await budget_limiter.is_end_user_within_model_budget(
-                "test-user",
-                {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}},
-                "gpt-4",
-            )
-
-    # Test model not in budget config
-    assert (
-        await budget_limiter.is_end_user_within_model_budget(
-            "test-user",
-            {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}},
-            "non-existent",
-        )
-        is True
+@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0))
+async def test_is_end_user_within_model_budget(
+    budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter,
+    spend: float,
+) -> None:
+    counter: Final = model_budget_spend_cache_key(
+        entity_type=Litellm_EntityType.END_USER,
+        entity_id="test-user",
+        budget_model="gpt-4",
+        budget_duration="1d",
     )
+    await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend)
+    budgets: Final = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}
+    if spend >= 100.0:
+        with pytest.raises(litellm.BudgetExceededError):
+            await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4")
+    else:
+        assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") is True
+    assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "non-existent") is True
 
 
 # Test _get_spend_for_model_budget for the end-user scope