refactor(auth): stage model budget aliases in native core

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

@@ -81,10 +81,12 @@ mod _native {
         request_auth_litellm_header_token, request_auth_llm_route, request_auth_managed,
         request_auth_managed_model_grants, request_auth_managed_models, request_auth_member_limit,
         request_auth_member_models, request_auth_model_access, request_auth_model_attribution,
-        request_auth_model_budget_check, request_auth_model_budget_query,
-        request_auth_model_budget_resolve, request_auth_model_groups, request_auth_model_helpers,
-        request_auth_model_projection, request_auth_org_admin, request_auth_org_admin_raw,
-        request_auth_org_route, request_auth_owner_teams, request_auth_personal_budget,
+        request_auth_model_budget_aliases, request_auth_model_budget_check,
+        request_auth_model_budget_query, request_auth_model_budget_resolve,
+        request_auth_model_budget_scope, request_auth_model_budget_team_exclusions,
+        request_auth_model_groups, request_auth_model_helpers, request_auth_model_projection,
+        request_auth_org_admin, request_auth_org_admin_raw, request_auth_org_route,
+        request_auth_owner_teams, request_auth_personal_budget,
         request_auth_personal_model_listing, request_auth_personal_model_request,
         request_auth_personal_models, request_auth_personal_owner, request_auth_reservation,
         request_auth_resolved_models, request_auth_route_access, request_auth_route_enable,
@@ -205,6 +207,9 @@ mod tests {
                 "request_auth_model_access",
                 "request_auth_model_attribution",
                 "request_auth_model_budget_check",
+                "request_auth_model_budget_aliases",
+                "request_auth_model_budget_scope",
+                "request_auth_model_budget_team_exclusions",
                 "request_auth_model_budget_query",
                 "request_auth_model_budget_resolve",
                 "request_auth_model_groups",

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

@@ -200,6 +200,42 @@ pub(crate) fn request_auth_model_budget_resolve(
     )
 }
 
+#[pyfunction]
+pub(crate) fn request_auth_model_budget_aliases(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::model_budget_aliases::Facts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::model_budget_aliases::decide(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_model_budget_scope(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::model_budget_aliases::ScopeFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::model_budget_aliases::scope(&facts)),
+    )
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_model_budget_team_exclusions(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: Vec<core::model_budget_data::ResolveFacts> = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::model_budget_data::team_exclusions(&facts)),
+    )
+}
+
 #[pyfunction]
 pub(crate) fn request_auth_model_budget_query(
     py: Python<'_>,

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

@@ -24,6 +24,7 @@ pub mod managed_models;
 pub mod member_models;
 pub mod model_access;
 pub mod model_attribution;
+pub mod model_budget_aliases;
 pub mod model_budget_data;
 mod model_budgets;
 pub mod model_groups;

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

@@ -0,0 +1,250 @@
+use serde::{Deserialize, Serialize};
+
+use crate::model_budget_data::{Snapshot, check};
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Row {
+    pub key: String,
+    pub names: Vec<String>,
+    pub prefix_known: Option<bool>,
+    pub provider_read: bool,
+    pub provider: Option<String>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Facts {
+    pub model: String,
+    pub rows: Vec<Row>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)]
+pub enum Read {
+    Entry,
+    Prefix,
+    Provider,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct Decision {
+    pub read: Option<Read>,
+    pub key: String,
+    pub models: Vec<String>,
+}
+
+fn reading(read: Read, key: &str) -> Decision {
+    Decision {
+        read: Some(read),
+        key: key.to_owned(),
+        models: vec![],
+    }
+}
+
+fn row<'a>(facts: &'a Facts, key: &str) -> Result<&'a Row, Decision> {
+    facts
+        .rows
+        .iter()
+        .find(|row| row.key == key)
+        .ok_or_else(|| reading(Read::Entry, key))
+}
+
+fn first(row: &Row) -> Option<&str> {
+    row.names
+        .first()
+        .filter(|name| **name != row.key)
+        .map(String::as_str)
+}
+
+fn provider(row: &Row) -> Result<Option<&str>, Decision> {
+    if let Some((prefix, _)) = row.key.split_once('/') {
+        match row.prefix_known {
+            Some(true) => return Ok(Some(prefix)),
+            Some(false) => (),
+            None => return Err(reading(Read::Prefix, &row.key)),
+        }
+    }
+    if !row.provider_read {
+        return Err(reading(Read::Provider, &row.key));
+    }
+    Ok(row.provider.as_deref())
+}
+
+fn target(facts: &Facts, key: &str, requested_provider: &str) -> Result<Option<String>, Decision> {
+    let row = row(facts, key)?;
+    let Some(name) = first(row) else {
+        return Ok(None);
+    };
+    if provider(row)? != Some(requested_provider) {
+        return Ok(None);
+    }
+    Ok(Some(
+        name.strip_prefix(&format!("{requested_provider}/"))
+            .unwrap_or(name)
+            .to_owned(),
+    ))
+}
+
+fn models(facts: &Facts) -> Result<Vec<String>, Decision> {
+    let Some((provider, name)) = facts.model.split_once('/') else {
+        return Ok(match first(row(facts, &facts.model)?) {
+            Some(target) => vec![facts.model.clone(), target.to_owned()],
+            None => vec![facts.model.clone()],
+        });
+    };
+    let direct = target(facts, &facts.model, provider)?.filter(|name| !name.is_empty());
+    let resolved = match direct {
+        Some(target) => Some(target),
+        None => target(facts, name, provider)?.filter(|name| !name.is_empty()),
+    };
+    Ok(match resolved {
+        Some(target) if target != name => vec![facts.model.clone(), format!("{provider}/{target}")],
+        _ => vec![facts.model.clone()],
+    })
+}
+
+pub fn decide(facts: &Facts) -> Decision {
+    match models(facts) {
+        Ok(models) => Decision {
+            read: None,
+            key: facts.model.clone(),
+            models,
+        },
+        Err(read) => read,
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct ScopeFacts {
+    pub models: Vec<String>,
+    pub snapshots: Vec<Snapshot>,
+    pub budget_present: bool,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
+pub struct ScopeDecision {
+    pub read_index: Option<usize>,
+    pub result_index: Option<usize>,
+}
+
+pub fn scope(facts: &ScopeFacts) -> ScopeDecision {
+    if let Some(index) = facts.snapshots.iter().position(|value| check(value).denied) {
+        return ScopeDecision {
+            read_index: None,
+            result_index: Some(index),
+        };
+    }
+    if !facts.budget_present {
+        return ScopeDecision {
+            read_index: None,
+            result_index: None,
+        };
+    }
+    if facts.snapshots.len() < facts.models.len() {
+        return ScopeDecision {
+            read_index: Some(facts.snapshots.len()),
+            result_index: None,
+        };
+    }
+    ScopeDecision {
+        read_index: None,
+        result_index: facts.snapshots.len().checked_sub(1),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::{fixture, rstest};
+
+    #[fixture]
+    fn alias() -> Row {
+        Row {
+            key: "provider/latest".into(),
+            names: vec!["provider/target".into()],
+            prefix_known: Some(true),
+            provider_read: false,
+            provider: None,
+        }
+    }
+
+    #[rstest]
+    #[case::catalog_first(vec![], Some(Read::Entry), "provider/latest", vec![])]
+    #[case::prefix_before_metadata(vec![Row { prefix_known: None, ..alias() }], Some(Read::Prefix), "provider/latest", vec![])]
+    #[case::metadata_for_unknown_provider(vec![Row { prefix_known: Some(false), ..alias() }], Some(Read::Provider), "provider/latest", vec![])]
+    #[case::known_prefix_ignores_metadata(vec![alias()], None, "provider/latest", vec!["provider/latest", "provider/target"])]
+    #[case::absent_qualified_reads_bare(vec![Row { names: vec![], ..alias() }], Some(Read::Entry), "latest", vec![])]
+    #[case::matching_metadata(vec![Row { prefix_known: Some(false), provider_read: true, provider: Some("provider".into()), ..alias() }], None, "provider/latest", vec!["provider/latest", "provider/target"])]
+    #[case::wrong_provider_reads_bare(vec![Row { prefix_known: Some(false), provider_read: true, provider: Some("other".into()), ..alias() }], Some(Read::Entry), "latest", vec![])]
+    fn lazy_alias_reads(
+        #[case] rows: Vec<Row>,
+        #[case] read: Option<Read>,
+        #[case] key: &str,
+        #[case] models: Vec<&str>,
+    ) {
+        let result = decide(&Facts {
+            model: "provider/latest".into(),
+            rows,
+        });
+        assert_eq!(
+            (result.read, result.key.as_str(), result.models),
+            (read, key, models.into_iter().map(str::to_owned).collect())
+        );
+    }
+
+    #[rstest]
+    #[case::bare_alias("latest", vec!["target"], vec!["latest", "target"])]
+    #[case::self("latest", vec!["latest", "other"], vec!["latest"])]
+    #[case::empty_row("latest", vec![], vec!["latest"])]
+    #[case::empty_bare_target("latest", vec![""], vec!["latest", ""])]
+    fn bare_alias(#[case] model: &str, #[case] names: Vec<&str>, #[case] expected: Vec<&str>) {
+        let result = decide(&Facts {
+            model: model.into(),
+            rows: vec![Row {
+                key: model.into(),
+                names: names.into_iter().map(str::to_owned).collect(),
+                ..alias()
+            }],
+        });
+        assert_eq!(result.read, None);
+        assert_eq!(result.models, expected);
+    }
+
+    fn snapshot(limit: f64, spend: f64) -> Snapshot {
+        Snapshot {
+            max_budget: Some(limit),
+            local_spend: Some(spend),
+            legacy_spend: None,
+            shared_spend: None,
+            shared_started_at: None,
+            shared_window_seconds: None,
+            now: 0.0,
+        }
+    }
+
+    #[rstest]
+    #[case::first(vec![], true, Some(0), None)]
+    #[case::target_after_allowed(vec![snapshot(1.0, 0.0)], true, Some(1), None)]
+    #[case::alias_denial(vec![snapshot(1.0, 1.0)], true, None, Some(0))]
+    #[case::target_denial(vec![snapshot(1.0, 0.0), snapshot(2.0, 3.0)], true, None, Some(1))]
+    #[case::all_allowed(vec![snapshot(1.0, 0.0), snapshot(2.0, 0.0)], true, None, Some(1))]
+    #[case::no_budgets(vec![], false, None, None)]
+    #[case::retain_denial(vec![snapshot(1.0, 1.0)], false, None, Some(0))]
+    #[case::nan_continues(vec![snapshot(f64::NAN, 1.0)], true, Some(1), None)]
+    fn scope_order(
+        #[case] snapshots: Vec<Snapshot>,
+        #[case] budget_present: bool,
+        #[case] read_index: Option<usize>,
+        #[case] result_index: Option<usize>,
+    ) {
+        assert_eq!(
+            scope(&ScopeFacts {
+                models: vec!["alias".into(), "target".into()],
+                snapshots,
+                budget_present
+            }),
+            ScopeDecision {
+                read_index,
+                result_index
+            }
+        );
+    }
+}

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

@@ -25,6 +25,8 @@ pub struct ResolveFacts {
     pub key_entries: Vec<Entry>,
     pub team: bool,
     pub access_groups: Option<Vec<String>>,
+    #[serde(default)]
+    pub excluded_entries: Vec<String>,
 }
 
 #[derive(Clone, Debug, Serialize)]
@@ -60,6 +62,18 @@ fn exact<'a>(entries: &'a [Entry], names: &[&str]) -> Option<&'a Entry> {
 }
 
 pub fn resolve(facts: &ResolveFacts) -> Resolution {
+    if !facts.excluded_entries.is_empty() {
+        return resolve(&ResolveFacts {
+            entries: facts
+                .entries
+                .iter()
+                .filter(|entry| !facts.excluded_entries.contains(&entry.name))
+                .cloned()
+                .collect(),
+            excluded_entries: vec![],
+            ..facts.clone()
+        });
+    }
     let names = candidates(facts);
     let key_replaces = facts.team
         && exact(&facts.key_entries, &names)
@@ -125,6 +139,24 @@ pub fn resolve(facts: &ResolveFacts) -> Resolution {
     }
 }
 
+pub fn team_exclusions(facts: &[ResolveFacts]) -> Vec<String> {
+    facts
+        .iter()
+        .flat_map(|facts| {
+            let names = candidates(facts);
+            if exact(&facts.key_entries, &names)
+                .is_some_and(|entry| entry.max_budget.is_some_and(|limit| limit >= 0.0))
+            {
+                names.into_iter().map(str::to_owned).collect::<Vec<_>>()
+            } else {
+                vec![]
+            }
+        })
+        .collect::<std::collections::BTreeSet<_>>()
+        .into_iter()
+        .collect()
+}
+
 fn glob(pattern: &str, name: &str) -> bool {
     let parts: Vec<&str> = pattern.split('*').collect();
     if parts.len() == 1 {
@@ -260,4 +292,40 @@ mod tests {
     fn wildcard(#[case] pattern: &str, #[case] name: &str, #[case] matches: bool) {
         assert_eq!(glob(pattern, name), matches);
     }
+
+    #[rstest]
+    #[case::alias_zero("alias", Some(0.0), vec!["alias"])]
+    #[case::target_positive("target", Some(1.0), vec!["target"])]
+    #[case::target_infinite("target", Some(f64::INFINITY), vec!["target"])]
+    #[case::negative("target", Some(-1.0), vec![])]
+    #[case::nan("target", Some(f64::NAN), vec![])]
+    #[case::missing("alias", None, vec![])]
+    fn team_alias_overrides(
+        #[case] key: &str,
+        #[case] limit: Option<f64>,
+        #[case] expected: Vec<&str>,
+    ) {
+        let entry = Entry {
+            name: key.into(),
+            max_budget: limit,
+            duration: Some("1h".into()),
+            duration_seconds: Some(3600.0),
+        };
+        let alias = ResolveFacts {
+            model: "alias".into(),
+            bedrock_base: "alias".into(),
+            bedrock_provider: String::new(),
+            entries: vec![],
+            key_entries: vec![entry],
+            team: true,
+            access_groups: None,
+            excluded_entries: vec![],
+        };
+        let target = ResolveFacts {
+            model: "target".into(),
+            bedrock_base: "target".into(),
+            ..alias.clone()
+        };
+        assert_eq!(team_exclusions(&[alias, target]), expected);
+    }
 }

litellm/_auth/model_budget_aliases.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+from typing import Final, Literal
+
+from pydantic import TypeAdapter
+
+from litellm._auth.failure import native_boundary
+from litellm._auth.model_budget_data import ModelBudgetSnapshot
+from litellm.rust_bridge import _native
+
+
+@dataclass(frozen=True, slots=True)
+class Row:
+    key: str
+    names: tuple[str, ...]
+    prefix_known: bool | None = None
+    provider_read: bool = False
+    provider: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class Facts:
+    model: str
+    rows: tuple[Row, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class Decision:
+    read: Literal["Entry", "Prefix", "Provider"] | None
+    key: str
+    models: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class ScopeFacts:
+    models: tuple[str, ...]
+    snapshots: tuple[ModelBudgetSnapshot, ...] = ()
+    budget_present: bool = True
+
+
+@dataclass(frozen=True, slots=True)
+class ScopeDecision:
+    read_index: int | None
+    result_index: int | None
+
+
+_FACTS: Final = TypeAdapter(Facts)
+_DECISION: Final = TypeAdapter(Decision)
+_SCOPE: Final = TypeAdapter(ScopeFacts)
+_SCOPE_DECISION: Final = TypeAdapter(ScopeDecision)
+
+
+@native_boundary
+def decide(facts: Facts) -> Decision:
+    return _DECISION.validate_python(_native.request_auth_model_budget_aliases(_FACTS.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes typed native facts
+
+
+@native_boundary
+def scope(facts: ScopeFacts) -> ScopeDecision:
+    return _SCOPE_DECISION.validate_python(_native.request_auth_model_budget_scope(_SCOPE.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes typed native facts

litellm/_auth/model_budget_data.py

@@ -26,6 +26,7 @@ class ModelBudgetResolveFacts:
     key_entries: tuple[ModelBudgetEntry, ...] = ()
     team: bool = False
     access_groups: tuple[str, ...] | None = None
+    excluded_entries: tuple[str, ...] = ()
 
 
 @dataclass(frozen=True, slots=True)
@@ -69,6 +70,13 @@ _RESOLUTION: Final = TypeAdapter(ModelBudgetResolution)
 _QUERY: Final = TypeAdapter(ModelBudgetQueryFacts)
 _SNAPSHOT: Final = TypeAdapter(ModelBudgetSnapshot)
 _DECISION: Final = TypeAdapter(ModelBudgetDecision)
+_TEAM_FACTS: Final = TypeAdapter(tuple[ModelBudgetResolveFacts, ...])
+_NAMES: Final = TypeAdapter(tuple[str, ...])
+
+
+@native_boundary
+def team_exclusions(facts: tuple[ModelBudgetResolveFacts, ...]) -> tuple[str, ...]:
+    return _NAMES.validate_python(_native.request_auth_model_budget_team_exclusions(_TEAM_FACTS.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes typed native facts
 
 
 @native_boundary

litellm/proxy/auth/common_budget_adapter.py

@@ -22,6 +22,7 @@ from litellm.proxy._types import (
     UserAPIKeyAuth,
 )
 from litellm.proxy.utils import ProxyLogging
+from litellm.redis.stores.spend_counters import end_user_counter, key_counter, team_counter, user_counter
 
 _METADATA: Final = TypeAdapter(Mapping[str, object])
 AlertKind: TypeAlias = Literal[
@@ -136,7 +137,7 @@ async def team_max_budget_check(
         return
     facts, result = await read_budget(
         BudgetFacts("Team", team_object.max_budget, token_present=valid_token is not None),
-        CounterRead(f"spend:team:{team_object.team_id}", team_object.spend or 0.0, team_object.max_budget),
+        CounterRead(team_counter(team_object.team_id), team_object.spend or 0.0, team_object.max_budget),
     )
     if result.alert is not None and valid_token is not None:
         emit_alert(
@@ -175,7 +176,7 @@ async def virtual_key_max_budget_check(
             tpm_limit=valid_token.tpm_limit,
             rpm_limit=valid_token.rpm_limit,
         ),
-        CounterRead(f"spend:key:{valid_token.token}", valid_token.spend or 0.0, valid_token.max_budget),
+        CounterRead(key_counter(str(valid_token.token)), valid_token.spend or 0.0, valid_token.max_budget),
     )
     if result.alert is not None and facts.spend is not None:
         send_token_budget_alert(valid_token, facts.spend, proxy_logging_obj, user_obj)
@@ -231,7 +232,7 @@ async def user_max_budget_check(
             team_key=team_object is not None and team_object.team_id is not None,  # pyright: ignore[reportUnnecessaryComparison]  # retain the legacy None guard for constructed cache rows
             apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
         ),
-        CounterRead(f"spend:user:{user_object.user_id}", user_object.spend or 0.0, user_object.max_budget),
+        CounterRead(user_counter(user_object.user_id), user_object.spend or 0.0, user_object.max_budget),
     )
     if result.alert is not None:
         emit_alert(
@@ -262,7 +263,7 @@ async def check_end_user_budget(end_user_obj: LiteLLM_EndUserTable, route: str)
     )
     facts, result = await read_budget(
         BudgetFacts("EndUser", limit, info_route=RouteChecks.is_info_route(route)),
-        CounterRead(f"spend:end_user:{end_user_obj.user_id}", end_user_obj.spend or 0.0, limit, authoritative=True),
+        CounterRead(end_user_counter(end_user_obj.user_id), end_user_obj.spend or 0.0, limit, authoritative=True),
     )
     if result.denied:
         raise_budget(

litellm/proxy/auth/common_budget_lookups.py

@@ -22,9 +22,14 @@ from litellm.proxy.auth.common_budget_adapter import (
     resolve_effective_team_member_budget,
 )
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
-from litellm.redis.stores.spend_counters import model_access_group_counter
 from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
 from litellm.proxy.utils import PrismaClient, ProxyLogging
+from litellm.redis.stores.spend_counters import (
+    model_access_group_counter,
+    org_counter,
+    tag_counter,
+    team_member_counter,
+)
 from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
 
 
@@ -47,7 +52,7 @@ async def organization_max_budget_check(
     carry_organization_budget_state(valid_token=valid_token, org_table=org)
     limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None
     facts, result = await read_budget(
-        BudgetFacts("Organization", limit), CounterRead(f"spend:org:{org_id}", org.spend or 0.0, limit)
+        BudgetFacts("Organization", limit), CounterRead(org_counter(org_id), org.spend or 0.0, limit)
     )
     if result.alert is not None:
         info: Final = key_alert_info(valid_token, facts.spend, limit, Litellm_EntityType.ORGANIZATION)
@@ -89,7 +94,7 @@ async def check_tag_row(name: str, row: LiteLLM_TagTable | None) -> None:
     )
     facts, result = await read_budget(
         BudgetFacts("Tag", limit),
-        CounterRead(f"spend:tag:{name}", (row.spend or 0.0) if row is not None else 0.0, limit, authoritative=True),
+        CounterRead(tag_counter(name), (row.spend or 0.0) if row is not None else 0.0, limit, authoritative=True),
     )
     if result.denied:
         raise_budget(
@@ -164,7 +169,7 @@ async def check_team_member_budget(
     facts, result = await read_budget(
         BudgetFacts("TeamMember", limit),
         CounterRead(
-            f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
+            team_member_counter(valid_token.user_id, team_object.team_id),
             (membership.spend or 0.0) if membership is not None else 0.0,
             limit,
         ),
@@ -200,9 +205,7 @@ async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget |
     limit: Final = row.max_budget if row is not None else None
     facts, _decision = await read_budget(
         BudgetFacts("AccessGroup", limit),
-        CounterRead(
-            model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True
-        ),
+        CounterRead(model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True),
     )
     return AccessGroupBudgetReading(name, facts)
 

litellm/proxy/auth/model_budget_adapter.py

@@ -6,7 +6,7 @@ from typing import Final
 from pydantic import TypeAdapter
 
 import litellm
-from litellm._auth import model_budget_data
+from litellm._auth import model_budget_aliases, model_budget_data
 from litellm._auth.failure import NativeAuthFailure
 from litellm._auth.model_budget_data import (
     ModelBudgetEntity,
@@ -19,6 +19,7 @@ from litellm._auth.model_budget_data import (
 from litellm.caching.caching import DualCache
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
 from litellm.llms.bedrock.common_utils import get_bedrock_base_model
 from litellm.proxy._types import Litellm_EntityType
 from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows
@@ -57,6 +58,10 @@ def _spend(raw: object) -> float:
         return 0.0
 
 
+def _price_entry(model: str) -> object:
+    return _OBJECT.validate_python(litellm.model_cost.get(model))  # pyright: ignore[reportUnknownMemberType]  # raw catalog IO
+
+
 @dataclass(frozen=True, slots=True)
 class ModelBudgetReading:
     snapshot: ModelBudgetSnapshot
@@ -72,6 +77,7 @@ class ModelBudgetIO:
     process_started_at: float
     clock: Callable[[], float] = time.time
     monotonic: Callable[[], float] = time.monotonic
+    price_entry: Callable[[str], object] = _price_entry
 
     async def counter(self, key: str) -> float:
         redis: Final = self.cache.redis_cache
@@ -81,6 +87,70 @@ class ModelBudgetIO:
         return _spend(raw)
 
 
+@dataclass(frozen=True, slots=True)
+class AliasReading:
+    row: model_budget_aliases.Row
+    source: AliasedCostEntry | None
+
+
+def read_alias_models(model: str, io: ModelBudgetIO, readings: tuple[AliasReading, ...] = ()) -> tuple[str, ...]:
+    decision: Final = model_budget_aliases.decide(
+        model_budget_aliases.Facts(model, tuple(reading.row for reading in readings))
+    )
+    match decision.read:
+        case None:
+            return decision.models
+        case "Entry":
+            raw: Final = io.price_entry(decision.key)
+            source: Final = raw if isinstance(raw, AliasedCostEntry) else None
+            entry: Final = AliasReading(
+                model_budget_aliases.Row(decision.key, source.names if source is not None else ()), source
+            )
+            return read_alias_models(model, io, (*readings, entry))
+        case "Prefix":
+            prefix_known: Final = decision.key.partition("/")[0] in litellm.provider_list
+            updated: Final = tuple(
+                replace(reading, row=replace(reading.row, prefix_known=prefix_known))
+                if reading.row.key == decision.key
+                else reading
+                for reading in readings
+            )
+            return read_alias_models(model, io, updated)
+        case "Provider":
+            selected: Final = next(reading for reading in readings if reading.row.key == decision.key)
+            assert selected.source is not None
+            provider: Final = selected.source.get("litellm_provider")
+            provided: Final = replace(
+                selected,
+                row=replace(selected.row, provider_read=True, provider=provider if isinstance(provider, str) else None),
+            )
+            return read_alias_models(
+                model, io, tuple(provided if reading.row.key == decision.key else reading for reading in readings)
+            )
+
+
+def _resolve_facts(
+    model: str,
+    budgets: Mapping[str, object],
+    io: ModelBudgetIO,
+    key_budgets: Mapping[str, object] | None = None,
+    *,
+    team: bool = False,
+    excluded_entries: tuple[str, ...] = (),
+) -> ModelBudgetResolveFacts:
+    base: Final = get_bedrock_base_model(model)
+    raw_price: Final = as_str_mapping(io.price_entry(base)) or {}
+    return ModelBudgetResolveFacts(
+        model=model,
+        bedrock_base=base,
+        bedrock_provider=str(raw_price.get("litellm_provider") or ""),
+        entries=tuple(_entry(name, config) for name, config in budgets.items()),
+        key_entries=tuple(_entry(name, config) for name, config in (key_budgets or {}).items()),
+        team=team,
+        excluded_entries=excluded_entries,
+    )
+
+
 def resolve_model_budget(
     model: str,
     budgets: Mapping[str, object],
@@ -89,17 +159,9 @@ def resolve_model_budget(
     key_budgets: Mapping[str, object] | None = None,
     *,
     team: bool = False,
+    excluded_entries: tuple[str, ...] = (),
 ) -> ModelBudgetResolution:
-    base: Final = get_bedrock_base_model(model)
-    raw_price: Final = as_str_mapping(_OBJECT.validate_python(litellm.model_cost.get(base))) or {}  # pyright: ignore[reportUnknownMemberType]  # project the legacy pricing map value
-    facts: Final = ModelBudgetResolveFacts(
-        model=model,
-        bedrock_base=base,
-        bedrock_provider=str(raw_price.get("litellm_provider") or ""),
-        entries=tuple(_entry(name, config) for name, config in budgets.items()),
-        key_entries=tuple(_entry(name, config) for name, config in (key_budgets or {}).items()),
-        team=team,
-    )
+    facts: Final = _resolve_facts(model, budgets, io, key_budgets, team=team, excluded_entries=excluded_entries)
     initial: Final = model_budget_data.resolve(facts)
     if not initial.read_groups:
         return initial
@@ -116,7 +178,55 @@ async def read_model_budget(
     budgets: Mapping[str, object] | None,
     key_budgets: Mapping[str, object] | None = None,
 ) -> ModelBudgetReading:
-    resolution: Final = resolve_model_budget(model, budgets or {}, io, team_id, key_budgets, team=entity == "Team")
+    models: Final = read_alias_models(model, io)
+    exclusions: Final = (
+        model_budget_data.team_exclusions(
+            tuple(_resolve_facts(candidate, {}, io, key_budgets, team=True) for candidate in models)
+        )
+        if entity == "Team"
+        else ()
+    )
+    scope_models: Final = read_alias_models(model, io) if entity == "Team" else models
+    return await _read_scope(io, entity, entity_id, team_id, scope_models, budgets, exclusions)
+
+
+async def _read_scope(
+    io: ModelBudgetIO,
+    entity: ModelBudgetEntity,
+    entity_id: str | None,
+    team_id: str | None,
+    models: tuple[str, ...],
+    budgets: Mapping[str, object] | None,
+    excluded_entries: tuple[str, ...],
+    readings: tuple[ModelBudgetReading, ...] = (),
+) -> ModelBudgetReading:
+    decision: Final = model_budget_aliases.scope(
+        model_budget_aliases.ScopeFacts(models, tuple(reading.snapshot for reading in readings), bool(budgets))
+    )
+    if decision.read_index is not None:
+        reading: Final = await _read_single(
+            io, entity, entity_id, team_id, models[decision.read_index], budgets, excluded_entries
+        )
+        return await _read_scope(
+            io, entity, entity_id, team_id, models, budgets, excluded_entries, (*readings, reading)
+        )
+    return (
+        readings[decision.result_index]
+        if decision.result_index is not None
+        else ModelBudgetReading(ModelBudgetSnapshot(), None, False)
+    )
+
+
+async def _read_single(
+    io: ModelBudgetIO,
+    entity: ModelBudgetEntity,
+    entity_id: str | None,
+    team_id: str | None,
+    model: str,
+    budgets: Mapping[str, object] | None,
+    excluded_entries: tuple[str, ...],
+) -> ModelBudgetReading:
+    resolution: Final = resolve_model_budget(model, budgets or {}, io, team_id, excluded_entries=excluded_entries)
     entry: Final = resolution.entry
     if entry is None:
         return ModelBudgetReading(ModelBudgetSnapshot(), None, False)

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -1046,28 +1046,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         model: str,
         exceeded_message: str,
         key_model_max_budget: Mapping[str, object] | None = None,
-    ) -> bool:
-        for candidate in expand_model_latest_alias(model):
-            await self._is_single_entity_within_model_budget(
-                entity_type=entity_type,
-                entity_id=entity_id,
-                team_id=team_id,
-                model_max_budget=model_max_budget,
-                model=candidate,
-                exceeded_message=exceeded_message,
-                key_model_max_budget=key_model_max_budget,
-            )
-        return True
-
-    async def _is_single_entity_within_model_budget(
-        self,
-        entity_type: Litellm_EntityType,
-        entity_id: str | None,
-        team_id: str | None,
-        model_max_budget: Mapping[str, object] | None,
-        model: str,
-        exceeded_message: str,
-        key_model_max_budget: Mapping[str, object] | None = None,
     ) -> bool:
         entities: Final[Mapping[Litellm_EntityType, ModelBudgetEntity]] = {
             Litellm_EntityType.KEY: "Key",

litellm/rust_bridge/_native.pyi

@@ -26,7 +26,9 @@ def native_redis_breaker_states() -> list[str]: ...
 class NativeAuthContinuation:
     def resume(self, answer: object) -> tuple[dict[str, object], NativeAuthContinuation | None]: ...
 
-def request_auth_start(facts: object, settings: object, now: object, resolved: bool) -> tuple[dict[str, object], NativeAuthContinuation | None]: ...
+def request_auth_start(
+    facts: object, settings: object, now: object, resolved: bool
+) -> tuple[dict[str, object], NativeAuthContinuation | None]: ...
 def request_auth_budgets(entry: str) -> tuple[dict[str, object], NativeAuthContinuation | None]: ...
 def request_auth_access(facts: object) -> dict[str, object]: ...
 def request_auth_agent_access(facts: object) -> str: ...
@@ -51,6 +53,9 @@ def request_auth_budget_route(facts: object) -> str: ...
 def request_auth_model_budget_resolve(facts: object) -> dict[str, object]: ...
 def request_auth_model_budget_query(facts: object) -> dict[str, object]: ...
 def request_auth_model_budget_check(facts: object) -> dict[str, object]: ...
+def request_auth_model_budget_aliases(facts: object) -> dict[str, object]: ...
+def request_auth_model_budget_scope(facts: object) -> dict[str, object]: ...
+def request_auth_model_budget_team_exclusions(facts: object) -> object: ...
 def request_auth_select_api_key(keys: object) -> dict[str, str] | None: ...
 def request_auth_select_websocket_api_key(keys: object) -> dict[str, str]: ...
 def request_auth_scheme_credential(value: str) -> str | None: ...
@@ -196,9 +201,7 @@ class NativeRouterUsage:
     def attach(self, client: NativeRedisClient) -> bool: ...
     @property
     def has_redis(self) -> bool: ...
-    def increment(
-        self, counter: _UsageCounter, value: float, ttl_seconds: int
-    ) -> Future[float]: ...
+    def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ...
     def increment_deployment(
         self,
         deployment_id: str,
@@ -207,12 +210,8 @@ class NativeRouterUsage:
         tokens: float,
         requests: int,
     ) -> Future[None]: ...
-    def get_many(
-        self, counters: Sequence[_UsageCounter]
-    ) -> Future[list[float | None]]: ...
-    def increment_sync(
-        self, counter: _UsageCounter, value: float, ttl_seconds: int
-    ) -> float: ...
+    def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ...
+    def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ...
     def increment_deployment_sync(
         self,
         deployment_id: str,
@@ -221,9 +220,7 @@ class NativeRouterUsage:
         tokens: float,
         requests: int,
     ) -> None: ...
-    def get_many_sync(
-        self, counters: Sequence[_UsageCounter]
-    ) -> list[float | None]: ...
+    def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ...
     def flush_local_counters(self) -> None: ...
     def snapshot(self) -> dict[str, float]: ...
 
@@ -399,7 +396,9 @@ class NativeAuthCache:
     def get_sync(self, key: str) -> object: ...
     def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ...
     def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
-    def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ...
+    def declare_set(
+        self, key: str, value: object, ttl_seconds: float | None, send: NativeSend
+    ) -> NativePending[None]: ...
     def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ...
     def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ...
     def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ...
@@ -1137,7 +1136,9 @@ class RouterCoreHandle:
         callbacks: Mapping[str, bool] | None = None,
     ) -> RouterCall[_SuccessReportSteps]: ...
     # Count a success against any id, driven like `report_failure`; `done` carries `key`.
-    def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ...
+    def count_success(
+        self, id: str, callbacks: Mapping[str, bool] | None = None
+    ) -> RouterCall[_SuccessReportSteps]: ...
     # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment.
     def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ...
     # Every deployment as a `model_list` entry, in config order.
@@ -1236,52 +1237,29 @@ class RouterCall(Generic[_S]):
     def next(self, outcome: object) -> _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: ...
-
 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: ...
-
 def request_auth_budget_lookup_organization(facts: object) -> object: ...
 def request_auth_budget_lookup_member(facts: object) -> object: ...
 def request_auth_budget_lookup_emails(facts: object) -> object: ...
 def request_auth_budget_lookup_soft(facts: object) -> bool: ...
-
 def request_auth_route_role(facts: object) -> object: ...
-
 def request_auth_route_role_api(facts: object) -> object: ...
-
 def request_auth_route_role_jwt(facts: object) -> object: ...
-
 def request_auth_route_role_inside(facts: object) -> object: ...
-
 def request_auth_route_role_expand(facts: object) -> object: ...
-
 def request_auth_route_role_listed(facts: object) -> object: ...
-
 def request_auth_org_admin_raw(facts: object) -> object: ...

tests/unit/_auth/model_alias_reference.py

@@ -0,0 +1,38 @@
+from collections.abc import Mapping, Sequence
+from typing import Final
+
+from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
+
+
+def _entry(key: str, catalog: Mapping[str, object]) -> AliasedCostEntry | None:
+    value: Final = catalog.get(key)
+    return value if isinstance(value, AliasedCostEntry) and value.names[:1] not in ((), (key,)) else None
+
+
+def _provider(key: str, value: AliasedCostEntry, providers: Sequence[str]) -> str | None:
+    prefix, separator, _ = key.partition("/")
+    if separator and prefix in providers:
+        return prefix
+    provider: Final = value.get("litellm_provider")
+    return provider if isinstance(provider, str) else None
+
+
+def _target(key: str, provider: str, catalog: Mapping[str, object], providers: Sequence[str]) -> str | None:
+    value: Final = _entry(key, catalog)
+    if value is None or _provider(key, value, providers) != provider:
+        return None
+    return value.names[0].removeprefix(f"{provider}/")
+
+
+def expand(model: str, catalog: Mapping[str, object], providers: Sequence[str]) -> tuple[str, ...]:
+    """Alias expansion frozen from main 2f430873aead937b9112e683703956eb1ab179af."""
+    provider, separator, name = model.partition("/")
+    if not separator:
+        value: Final = _entry(model, catalog)
+        return (model,) if value is None else (model, value.names[0])
+    target: Final = (
+        _target(f"{provider}/{name}", provider, catalog, providers)
+        or _target(name, provider, catalog, providers)
+        or name
+    )
+    return (model, f"{provider}/{target}") if target != name else (model,)

tests/unit/_auth/test_model_budget_aliases.py

@@ -0,0 +1,66 @@
+from collections.abc import Mapping
+from typing import Final
+from unittest.mock import Mock, call
+
+import pytest
+
+import litellm
+from litellm._auth import model_budget_aliases as native
+from litellm._auth.model_budget_data import ModelBudgetSnapshot
+from litellm.caching.caching import DualCache
+from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
+from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_alias_models
+from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows
+
+from .model_alias_reference import expand
+
+
+@pytest.mark.parametrize("model", ("latest", "openai/latest", "unknown/latest", "bedrock/latest"))
+@pytest.mark.parametrize("names", ((), ("latest",), ("openai/latest",), ("target",), ("",), ("openai/target",)))
+@pytest.mark.parametrize("provider", (None, "openai", "unknown", "bedrock_converse", 42))
+def test_alias_lookup_matches_frozen_main(model: str, names: tuple[str, ...], provider: object) -> None:
+    row: Final = AliasedCostEntry({"litellm_provider": provider}, names)
+    catalog: Final[Mapping[str, object]] = {model: row, "latest": row}
+    expected: Final = expand(model, catalog, litellm.provider_list)
+    lookup: Final = Mock(side_effect=catalog.get)
+    io: Final = ModelBudgetIO(
+        DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lookup
+    )
+    assert read_alias_models(model, io) == expected
+    assert lookup.call_args_list[0].args == (model,)
+
+
+@pytest.mark.parametrize("raw", (None, {}, {"alias_of": "target"}, {"names": ["target"]}))
+def test_plain_catalog_rows_do_not_become_aliases(raw: object) -> None:
+    io: Final = ModelBudgetIO(
+        DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lambda _: raw
+    )
+    assert read_alias_models("latest", io) == ("latest",)
+
+
+@pytest.mark.parametrize("direct", (True, False))
+def test_catalog_reads_bare_name_only_when_qualified_alias_misses(direct: bool) -> None:
+    row: Final = AliasedCostEntry({"litellm_provider": "openai"}, ("target", "latest"))
+    catalog: Final = {"openai/latest": row if direct else None, "latest": row}
+    lookup: Final = Mock(side_effect=catalog.get)
+    io: Final = ModelBudgetIO(
+        DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lookup
+    )
+    assert read_alias_models("openai/latest", io) == ("openai/latest", "openai/target")
+    assert lookup.call_args_list == ([call("openai/latest")] if direct else [call("openai/latest"), call("latest")])
+
+
+@pytest.mark.parametrize(
+    ("snapshots", "expected"),
+    (
+        ((), native.ScopeDecision(0, None)),
+        ((ModelBudgetSnapshot(1, 0),), native.ScopeDecision(1, None)),
+        ((ModelBudgetSnapshot(1, 1),), native.ScopeDecision(None, 0)),
+        ((ModelBudgetSnapshot(1, 0), ModelBudgetSnapshot(2, 2)), native.ScopeDecision(None, 1)),
+        ((ModelBudgetSnapshot(1, 0), ModelBudgetSnapshot(2, 0)), native.ScopeDecision(None, 1)),
+    ),
+)
+def test_scope_reads_only_until_first_denial(
+    snapshots: tuple[ModelBudgetSnapshot, ...], expected: native.ScopeDecision
+) -> None:
+    assert native.scope(native.ScopeFacts(("alias", "target"), snapshots)) == expected

tests/unit/_auth/test_model_budget_data.py

@@ -55,6 +55,27 @@ def entry(name: str, limit: float | None = 1, duration: str | None = "1h") -> Mo
     return ModelBudgetEntry(name, limit, duration, 3600 if duration is not None else None)
 
 
+@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan))
+@pytest.mark.parametrize("overrides", (("alias",), ("target",), ("alias", "target")))
+def test_team_exact_entries_are_removed_across_all_alias_names(limit: float | None, overrides: tuple[str, ...]) -> None:
+    key: Final = tuple(entry(name, limit) for name in overrides)
+    facts: Final = tuple(ModelBudgetResolveFacts(model, model, "", (), key, True) for model in ("alias", "target"))
+    excluded: Final = native.team_exclusions(facts)
+    assert excluded == (tuple(sorted(overrides)) if limit is not None and limit >= 0 else ())
+    resolution: Final = native.resolve(
+        ModelBudgetResolveFacts(
+            "alias",
+            "alias",
+            "",
+            (entry("alias"), entry("target"), entry("*")),
+            access_groups=("alias", "target"),
+            excluded_entries=excluded,
+        )
+    )
+    expected: Final = "alias" if "alias" not in excluded else "target" if "target" not in excluded else "*"
+    assert resolution.entry is not None and resolution.entry.name == expected
+
+
 @pytest.mark.parametrize(
     ("entries", "groups", "expected"),
     (

tests/unit/proxy/auth/test_model_budget_adapter.py

@@ -1,14 +1,22 @@
 from collections.abc import Awaitable, Mapping
 from dataclasses import dataclass, field
 from typing import Final
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, call
 
 import pytest
 
 import litellm
 from litellm._auth import model_budget_data
+from litellm._auth.budget_pricing import ModelPriceFacts
+from litellm._auth.fallback_candidate import (
+    CandidateBudgetRead,
+    CandidateBudgetScope,
+    FallbackCandidateFacts,
+    fallback_candidate,
+)
 from litellm._auth.model_budget_data import ModelBudgetEntity
 from litellm.caching.caching import DualCache
+from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
 from litellm.proxy._types import Litellm_EntityType
 from litellm.proxy.auth.model_budget_adapter import (
     ModelBudgetIO,
@@ -136,3 +144,164 @@ async def test_shared_read_failure_is_terminal_and_not_replayed() -> None:
         await read_model_budget(io, "User", "id", None, "model", {"model": {"max_budget": 1, "budget_duration": "1h"}})
     assert raised.value is failure
     assert windows.read.await_count == 1
+
+
+@pytest.fixture
+def alias_catalog(monkeypatch: pytest.MonkeyPatch) -> Mapping[str, object]:
+    catalog: Final = {"budget-latest": AliasedCostEntry({}, ("budget-target", "budget-latest"))}
+    monkeypatch.setattr(litellm, "model_cost", catalog)
+    return catalog
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("entity", ("Key", "User", "Team", "TeamMember", "EndUser"))
+@pytest.mark.parametrize(
+    ("alias_spend", "target_spend", "expected"),
+    ((1.0, 2.0, "budget-latest"), (0.0, 2.0, "budget-target"), (0.0, 0.0, None)),
+)
+async def test_raw_alias_scope_preserves_first_denial_and_counter_order(
+    alias_catalog: Mapping[str, object],
+    entity: ModelBudgetEntity,
+    alias_spend: float,
+    target_spend: float,
+    expected: str | None,
+) -> None:
+    cache: Final = DualCache()
+    scope: Final = {
+        "Key": Litellm_EntityType.KEY,
+        "User": Litellm_EntityType.USER,
+        "Team": Litellm_EntityType.TEAM,
+        "TeamMember": Litellm_EntityType.TEAM_MEMBER,
+        "EndUser": Litellm_EntityType.END_USER,
+    }[entity]
+    alias_key: Final = reference.model_budget_spend_cache_key(scope, "id", "budget-latest", "1h")
+    target_key: Final = reference.model_budget_spend_cache_key(scope, "id", "budget-target", "1h")
+    cache.in_memory_cache.set_cache(alias_key, alias_spend)  # pyright: ignore[reportUnknownMemberType]  # legacy cache setter has no annotations
+    cache.in_memory_cache.set_cache(target_key, target_spend)  # pyright: ignore[reportUnknownMemberType]  # legacy cache setter has no annotations
+    shared_read: Final = AsyncMock(return_value=None)
+    windows: Final = Windows(read=shared_read)
+    io: Final = ModelBudgetIO(
+        cache, groups, windows, 0, clock=lambda: 100, monotonic=lambda: 0, price_entry=alias_catalog.get
+    )
+    reading: Final = await read_model_budget(
+        io,
+        entity,
+        "id",
+        None,
+        "budget-latest",
+        {
+            "budget-latest": {"max_budget": 1, "budget_duration": "1h"},
+            "budget-target": {"max_budget": 2, "budget_duration": "1h"},
+        },
+    )
+    decision: Final = model_budget_data.check(reading.snapshot)
+    assert decision.denied is (expected is not None)
+    assert reading.budget_model == (expected or "budget-target")
+    assert decision.current_spend == (alias_spend if expected == "budget-latest" else target_spend)
+    candidate: Final = fallback_candidate(
+        FallbackCandidateFacts(
+            True,
+            True,
+            True,
+            membership="Loaded",
+            price=ModelPriceFacts(cached=False),
+            budget_models=("budget-latest",),
+            budget_scopes=(CandidateBudgetScope(entity, True, True),),
+            budgets=(CandidateBudgetRead(0, entity, reading.snapshot),),
+        )
+    )
+    assert candidate.kind == ("Complete" if expected is not None else "ReadAuthorization")
+    assert candidate.usable is False
+    expected_shared: Final = (
+        [] if entity == "EndUser" or alias_spend == 1 else [alias_key] if target_spend == 2 else [alias_key, target_key]
+    )
+    assert shared_read.call_args_list == [call(spend_key=key, budget_duration="1h") for key in expected_shared]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("overrides", (("budget-latest",), ("budget-target",), ("budget-latest", "budget-target")))
+async def test_team_overrides_apply_across_alias_and_target(
+    alias_catalog: Mapping[str, object], overrides: tuple[str, ...]
+) -> None:
+    io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0, price_entry=alias_catalog.get)
+    budgets: Final = {name: {"max_budget": 0, "budget_duration": "1h"} for name in ("budget-latest", "budget-target")}
+    key: Final = {name: {"max_budget": 1, "budget_duration": "1h"} for name in overrides}
+    reading: Final = await read_model_budget(io, "Team", "id", "id", "budget-latest", budgets, key)
+    expected: Final = next((name for name in budgets if name not in overrides), None)
+    assert model_budget_data.check(reading.snapshot).denied is (expected is not None)
+    assert reading.budget_model == expected
+
+
+@pytest.mark.asyncio
+async def test_shared_native_model_flow_enforces_target_before_user_scope(alias_catalog: Mapping[str, object]) -> None:
+    from starlette.requests import Request
+
+    from litellm._auth.decision import Deny, ModelBudgetExceeded
+    from litellm._auth.driver import drive_model_budgets
+    from litellm._auth.rust_core import RustCore
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.auth.user_api_key_auth import (
+        _ModelBudgetEffects,  # pyright: ignore[reportPrivateUsage]  # exercise the production raw effect boundary
+    )
+
+    token: Final = UserAPIKeyAuth(
+        token="id",
+        user_id="owner",
+        key_alias="alias-key",
+        model_max_budget={"budget-target": {"max_budget": 0, "budget_duration": "1h"}},
+        user_model_max_budget={"budget-latest": {"max_budget": 0, "budget_duration": "1h"}},
+    )
+    effects: Final = _ModelBudgetEffects(
+        request=Request({"type": "http", "headers": [], "path": "/chat/completions"}),
+        request_data={"model": "budget-latest"},
+        valid_token=token,
+        route="/chat/completions",
+        limiter=reference._PROXY_VirtualKeyModelMaxBudgetLimiter(DualCache(), clock=lambda: 100),  # pyright: ignore[reportPrivateUsage]  # concrete limiter integration
+        llm_router=None,
+        llm_model_list=None,
+    )
+    outcome: Final = await drive_model_budgets(RustCore().check_model_budgets(), effects)
+    assert isinstance(outcome, Deny) and isinstance(outcome.denial, ModelBudgetExceeded)
+    assert (outcome.denial.entity_type, outcome.denial.entity_id, outcome.denial.max_budget) == ("key", "id", 0.0)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("authorization_denied", (False, True))
+async def test_fallback_probe_retains_original_denial_without_authorization_retries(
+    alias_catalog: Mapping[str, object], authorization_denied: bool
+) -> None:
+    from starlette.requests import Request
+
+    from litellm._auth.decision import Deny, ModelBudgetExceeded
+    from litellm._auth.driver import drive_model_budgets
+    from litellm._auth.rust_core import RustCore
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.auth.user_api_key_auth import (
+        _ModelBudgetEffects,  # pyright: ignore[reportPrivateUsage]  # exercise the production raw effect boundary
+    )
+
+    data: Final[dict[str, object]] = {"model": "initial"}
+    token: Final = UserAPIKeyAuth(
+        token="id",
+        models=["initial", "later"],
+        model_max_budget={
+            "initial": {"max_budget": 0, "budget_duration": "1h"},
+            "budget-target": {"max_budget": 1 if authorization_denied else 0, "budget_duration": "1h"},
+        },
+        budget_fallbacks={"initial": ["budget-latest", "later"] if authorization_denied else ["budget-latest"]},
+    )
+    effects: Final = _ModelBudgetEffects(
+        request=Request({"type": "http", "headers": [], "path": "/chat/completions"}),
+        request_data=data,
+        valid_token=token,
+        route="/chat/completions",
+        limiter=reference._PROXY_VirtualKeyModelMaxBudgetLimiter(  # pyright: ignore[reportPrivateUsage]  # concrete limiter integration
+            DualCache(), clock=lambda: 100, shared_windows=Windows(read=AsyncMock(return_value=None))
+        ),
+        llm_router=None,
+        llm_model_list=None,
+    )
+    outcome: Final = await drive_model_budgets(RustCore().check_model_budgets(), effects)
+    assert isinstance(outcome, Deny) and isinstance(outcome.denial, ModelBudgetExceeded)
+    assert "model=initial" in outcome.denial.message
+    assert data == {"model": "initial"}