refactor(auth): drive model limits and exemption reads through native facts

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

@@ -39,32 +39,32 @@ mod _native {
     use crate::model_registry::bundled_model_cost_map;
     #[pymodule_export]
     use crate::request_auth::{
-    NativeAuthContinuation,
-    request_auth_access,
-    request_auth_admission,
-    request_auth_budget_alert,
-    request_auth_budget_price,
-    request_auth_budget_route,
-    request_auth_budget_window_start,
-    request_auth_budgets,
-    request_auth_common_budget,
-    request_auth_fallback_access,
-    request_auth_litellm_header_token,
-    request_auth_managed,
-    request_auth_member_limit,
-    request_auth_model_access,
-    request_auth_model_budget_check,
-    request_auth_model_budget_query,
-    request_auth_model_budget_resolve,
-    request_auth_model_groups,
-    request_auth_org_admin,
-    request_auth_org_route,
-    request_auth_route_access,
-    request_auth_scheme_credential,
-    request_auth_select_api_key,
-    request_auth_select_websocket_api_key,
-    request_auth_start,
-    ||||||| parent of da364c75e2 (refactor(auth): port model and route policy to native decisions) NativeAuthContinuation,
+        NativeAuthContinuation,
+        request_auth_access,
+        request_auth_admission,
+        request_auth_budget_alert,
+        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_fallback_access,
+        request_auth_litellm_header_token,
+        request_auth_managed,
+        request_auth_member_limit,
+        request_auth_model_access,
+        request_auth_model_budget_check,
+        request_auth_model_budget_query,
+        request_auth_model_budget_resolve,
+        request_auth_model_groups,
+        request_auth_org_admin,
+        request_auth_org_route,
+        request_auth_route_access,
+        request_auth_scheme_credential,
+        request_auth_select_api_key,
+        request_auth_select_websocket_api_key,
+        request_auth_start,
     };
     #[pymodule_export]
     use crate::router::{RouterCall, RouterCoreHandle};

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

@@ -277,3 +277,15 @@ pub(crate) fn request_auth_fallback_access(py: Python<'_>, facts: &Bound<'_, PyA
     let decision = py.detach(move || core::fallback_access::decide(facts)).map_err(map_error)?;
     to_py(py, &decision)
 }
+
+#[pyfunction]
+pub(crate) fn request_auth_budget_exemption(
+    py: Python<'_>,
+    facts: &Bound<'_, PyAny>,
+) -> PyResult<Py<PyAny>> {
+    let facts: core::budget_pricing::ExemptionFacts = from_py(facts)?;
+    to_py(
+        py,
+        &py.detach(move || core::budget_pricing::exemption(&facts)),
+    )
+}

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

@@ -117,6 +117,63 @@ pub fn route(facts: &RouteBudgetFacts) -> RouteBudgetDecision {
     if mcp_discovery { Skip } else { ReadPricing }
 }
 
+#[derive(Clone, Debug, Deserialize)]
+pub struct ExemptionFacts {
+    pub route: RouteBudgetFacts,
+    pub model_input_present: bool,
+    pub router_present: bool,
+    pub model_count: usize,
+    #[serde(default)]
+    pub cache_read: bool,
+    pub prices: Vec<ModelPriceFacts>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+pub struct ExemptionDecision {
+    pub read_pass_through: bool,
+    pub read_price_cache: bool,
+    pub read_model: Option<usize>,
+    pub exempt: bool,
+}
+
+pub fn exemption(facts: &ExemptionFacts) -> ExemptionDecision {
+    let done = |exempt| ExemptionDecision {
+        read_pass_through: false,
+        read_price_cache: false,
+        read_model: None,
+        exempt,
+    };
+    match route(&facts.route) {
+        RouteBudgetDecision::ReadPassThrough => {
+            return ExemptionDecision {
+                read_pass_through: true,
+                ..done(false)
+            };
+        }
+        RouteBudgetDecision::Skip => return done(true),
+        RouteBudgetDecision::ReadPricing => (),
+    }
+    if !facts.model_input_present || !facts.router_present {
+        return done(false);
+    }
+    if !facts.cache_read {
+        return ExemptionDecision {
+            read_price_cache: true,
+            ..done(false)
+        };
+    }
+    if facts.prices.iter().any(|facts| !price(facts).zero_cost) {
+        return done(false);
+    }
+    if facts.prices.len() < facts.model_count {
+        return ExemptionDecision {
+            read_model: Some(facts.prices.len()),
+            ..done(false)
+        };
+    }
+    done(true)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -154,4 +211,51 @@ mod tests {
             }
         );
     }
+    #[rstest]
+    #[case::read_first(vec![], Some(0), false)]
+    #[case::read_second(vec![true], Some(1), false)]
+    #[case::priced_first(vec![false], None, false)]
+    #[case::all_free(vec![true, true], None, true)]
+    #[case::priced_second(vec![true, false], None, false)]
+    fn exemption_reads_models_in_order(
+        #[case] cached: Vec<bool>,
+        #[case] read: Option<usize>,
+        #[case] exempt: bool,
+    ) {
+        let facts = ExemptionFacts {
+            route: RouteBudgetFacts {
+                side_effect_route: false,
+                discovery_route: false,
+                llm_route: true,
+                route: "/chat/completions".to_owned(),
+                method: None,
+                body_present: false,
+                registered_pass_through: None,
+            },
+            model_input_present: true,
+            router_present: true,
+            model_count: 2,
+            cache_read: true,
+            prices: cached
+                .into_iter()
+                .map(|value| ModelPriceFacts {
+                    cached: Some(value),
+                    lookup_failed: false,
+                    group_found: false,
+                    input_price: None,
+                    output_price: None,
+                    deployments: None,
+                })
+                .collect(),
+        };
+        assert_eq!(
+            exemption(&facts),
+            ExemptionDecision {
+                read_pass_through: false,
+                read_price_cache: false,
+                read_model: read,
+                exempt
+            }
+        );
+    }
 }

litellm/_auth/budget_pricing.py

@@ -58,3 +58,29 @@ def price(facts: ModelPriceFacts) -> PriceDecision:
 
 def route(facts: RouteBudgetFacts) -> RouteBudgetDecision:
     return _ROUTE_DECISION.validate_python(_native.request_auth_budget_route(_ROUTE.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes the validated frozen facts for the native boundary
+
+
+@dataclass(frozen=True, slots=True)
+class ExemptionFacts:
+    route: RouteBudgetFacts
+    model_input_present: bool
+    router_present: bool
+    model_count: int
+    prices: tuple[ModelPriceFacts, ...] = ()
+    cache_read: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class ExemptionDecision:
+    read_pass_through: bool
+    read_price_cache: bool
+    read_model: int | None
+    exempt: bool
+
+
+_EXEMPTION: Final = TypeAdapter(ExemptionFacts)
+_EXEMPTION_DECISION: Final = TypeAdapter(ExemptionDecision)
+
+
+def exemption(facts: ExemptionFacts) -> ExemptionDecision:
+    return _EXEMPTION_DECISION.validate_python(_native.request_auth_budget_exemption(_EXEMPTION.dump_python(facts)))  # pyright: ignore[reportAny]  # Pydantic serializes frozen facts

litellm/proxy/auth/budget_pricing_adapter.py

@@ -56,7 +56,7 @@ def read_model_price_facts(model: str, router: Router, cached: bool | None = Non
             input_price=group.input_cost_per_token if group is not None else None,
             output_price=group.output_cost_per_token if group is not None else None,
         )
-    except Exception:
+    except Exception:  # noqa: BLE001  # preserve lookup failure as non-exempt without caching
         return ModelPriceFacts(lookup_failed=True)
     if not budget_pricing.price(facts).read_deployments:
         return facts
@@ -64,18 +64,22 @@ def read_model_price_facts(model: str, router: Router, cached: bool | None = Non
         deployments: Final = tuple(
             _deployment_price(deployment) for deployment in _deployments_served_for(model, router)
         )
-    except Exception:
+    except Exception:  # noqa: BLE001  # preserve lookup failure as non-exempt without caching
         return ModelPriceFacts(lookup_failed=True)
     return replace(facts, deployments=deployments)
 
 
-def _model_zero(model: str, router: Router, cache: Mapping[str, bool] | None) -> bool:
-    cached: Final = cache.get(model) if cache is not None else None
+def _read_cached_model_price(model: str, router: Router, cache: Mapping[str, bool] | None) -> ModelPriceFacts:
+    cached: Final = cache.get(model) if isinstance(cache, dict) else None
     facts: Final = read_model_price_facts(model, router, cached)
     decision: Final = budget_pricing.price(facts)
     if isinstance(cache, dict) and decision.cache_value is not None:
-        cache[model] = decision.cache_value  # rebind-ok: preserve registry cache writes selected by native policy
-    return decision.zero_cost
+        cache[model] = decision.cache_value  # rebind-ok: native-selected registry cache write
+    return facts
+
+
+def read_model_cost_facts(model: str, router: Router) -> ModelPriceFacts:
+    return _read_cached_model_price(model, router, router.registry().zero_cost())
 
 
 def is_model_cost_zero(model: str | Sequence[str] | None, llm_router: Router | None) -> bool:
@@ -83,7 +87,7 @@ def is_model_cost_zero(model: str | Sequence[str] | None, llm_router: Router | N
         return budget_pricing.price(ModelPriceFacts()).zero_cost
     names: Final = (model,) if isinstance(model, str) else tuple(model)
     cache: Final = llm_router.registry().zero_cost()
-    return all(_model_zero(name, llm_router, cache) for name in names)
+    return all(budget_pricing.price(_read_cached_model_price(name, llm_router, cache)).zero_cost for name in names)
 
 
 def read_route_budget_facts(route: str, body: Mapping[str, object] | None = None) -> RouteBudgetFacts:
@@ -115,6 +119,51 @@ def route_skips_budget_checks(route: str) -> bool:
     return _route_decision(route) == "Skip"
 
 
+def read_budget_exemption_facts(
+    route: str,
+    model: str | Sequence[str] | None,
+    llm_router: Router | None,
+    *,
+    request_body: Mapping[str, object] | None = None,
+) -> budget_pricing.ExemptionFacts:
+    names: Final = (model,) if isinstance(model, str) else tuple(model or ())
+    initial: Final = budget_pricing.ExemptionFacts(
+        read_route_budget_facts(route, request_body),
+        model is not None,
+        llm_router is not None,
+        len(names),
+    )
+    return _read_exemption(initial, names, llm_router, None)
+
+
+def _read_exemption(
+    facts: budget_pricing.ExemptionFacts,
+    models: tuple[str, ...],
+    router: Router | None,
+    cache: Mapping[str, bool] | None,
+) -> budget_pricing.ExemptionFacts:
+    current = facts  # rebind-ok: native-selected IO driver state
+    selected_cache = cache  # rebind-ok: cache handle acquired at the first model read
+    while True:
+        decision = budget_pricing.exemption(current)
+        if decision.read_pass_through:
+            from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers
+
+            registered = InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=current.route.route)
+            current = replace(current, route=replace(current.route, registered_pass_through=registered))
+            continue
+        if decision.read_price_cache:
+            assert router is not None
+            selected_cache = router.registry().zero_cost()
+            current = replace(current, cache_read=True)
+            continue
+        if decision.read_model is None:
+            return current
+        assert router is not None
+        price = _read_cached_model_price(models[decision.read_model], router, selected_cache)
+        current = replace(current, prices=(*current.prices, price))
+
+
 def request_skips_budget_checks(
     route: str,
     model: str | Sequence[str] | None,
@@ -122,6 +171,5 @@ def request_skips_budget_checks(
     *,
     request_body: Mapping[str, object] | None = None,
 ) -> bool:
-    if _route_decision(route, request_body) == "Skip":
-        return True
-    return is_model_cost_zero(model, llm_router)
+    facts: Final = read_budget_exemption_facts(route, model, llm_router, request_body=request_body)
+    return budget_pricing.exemption(facts).exempt

litellm/proxy/auth/budget_windows_adapter.py

@@ -28,7 +28,7 @@ def _window(raw: object) -> BudgetWindow:
 def _seconds(duration: str) -> int | None:
     try:
         return duration_in_seconds(duration)
-    except Exception:
+    except Exception:  # noqa: BLE001  # preserve invalid-duration skip behavior
         return None
 
 

litellm/proxy/auth/common_budget_lookups.py

@@ -1,6 +1,8 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
 from typing import Final
 
+from litellm._auth import common_budgets
 from litellm._auth.common_budgets import BudgetFacts
 from litellm.proxy._types import (
     Litellm_EntityType,
@@ -51,7 +53,7 @@ async def organization_max_budget_check(
             proxy_logging_obj=proxy_logging_obj,
             include_budget_table=True,
         )
-    except Exception:
+    except Exception:  # noqa: BLE001  # preserve organization lookup failure handling
         return
     if org is None:
         return
@@ -218,3 +220,50 @@ async def check_team_member_budget(
             f"{valid_token.user_id}:{team_object.team_id}",
             f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {facts.spend}, Max budget: {limit}",
         )
+
+
+@dataclass(frozen=True, slots=True)
+class AccessGroupBudgetReading:
+    name: str
+    facts: BudgetFacts
+
+
+async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget | None) -> AccessGroupBudgetReading:
+    limit: Final = row.max_budget if row is not None else None
+    facts, _decision = await read_budget(
+        BudgetFacts("AccessGroup", limit),
+        CounterRead(
+            model_access_group_spend_counter_key(name), row.spend if row is not None else 0.0, limit, authoritative=True
+        ),
+    )
+    return AccessGroupBudgetReading(name, facts)
+
+
+async def read_access_group_budget_facts(
+    groups: Sequence[str],
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+) -> tuple[AccessGroupBudgetReading, ...]:
+    from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch
+
+    if prisma_client is None or not groups:
+        return ()
+    rows: Final = await get_model_access_group_budgets_batch(
+        access_group_names=groups,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+    )
+    return await _read_access_group_rows(tuple(groups), rows)
+
+
+async def _read_access_group_rows(
+    groups: tuple[str, ...],
+    rows: Mapping[str, ModelAccessGroupBudget],
+) -> tuple[AccessGroupBudgetReading, ...]:
+    readings: tuple[AccessGroupBudgetReading, ...] = ()  # rebind-ok: native-selected IO driver accumulator
+    for name in groups:
+        reading = await read_access_group_budget_row(name, rows.get(name))
+        readings = (*readings, reading)
+        if common_budgets.check(reading.facts).denied:
+            return readings
+    return readings

litellm/proxy/auth/model_budget_adapter.py

@@ -40,7 +40,7 @@ def _entry(name: str, raw: object) -> ModelBudgetEntry:
     try:
         config: Final = BudgetConfig.model_validate(raw)
         seconds: Final = duration_in_seconds(config.budget_duration) if config.budget_duration is not None else None
-    except Exception:
+    except Exception:  # noqa: BLE001  # preserve malformed entry fallthrough
         return ModelBudgetEntry(name, None, None, None)
     return ModelBudgetEntry(name, config.max_budget, config.budget_duration, seconds)
 

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -15,6 +15,8 @@ from typing import (
 )
 
 import litellm
+from litellm._auth import model_budget_data
+from litellm._auth.model_budget_data import ModelBudgetEntity
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -40,6 +42,7 @@ from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload
 
 if TYPE_CHECKING:
+    from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading
     from litellm.router import Router
 
 VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
@@ -947,9 +950,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             entity_type=Litellm_EntityType.TEAM,
             entity_id=team_id,
             team_id=team_id,
-            model_max_budget=team_model_budget_entries(
-                model=model, team_model_max_budget=team_model_max_budget, key_model_max_budget=key_model_max_budget
-            ),
+            model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
             model=model,
             exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}",
         )
@@ -973,6 +975,27 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             ),
         )
 
+    async def read_model_budget_raw(
+        self,
+        entity: ModelBudgetEntity,
+        entity_id: str | None,
+        team_id: str | None,
+        model: str,
+        budgets: Mapping[str, object] | None,
+        key_budgets: Mapping[str, object] | None = None,
+    ) -> "ModelBudgetReading":
+        from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_model_budget
+
+        io: Final = ModelBudgetIO(
+            self.dual_cache,
+            self._model_access_groups,
+            self._shared_windows,
+            _PROCESS_STARTED_AT,
+            clock=self._clock,
+            monotonic=time.monotonic,
+        )
+        return await read_model_budget(io, entity, entity_id, team_id, model, budgets, key_budgets)
+
     async def _is_entity_within_model_budget(
         self,
         entity_type: Litellm_EntityType,
@@ -981,45 +1004,32 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         model_max_budget: Mapping[str, object] | None,
         model: str,
         exceeded_message: str,
+        key_model_max_budget: Mapping[str, object] | None = None,
     ) -> bool:
-        if not model_max_budget:
-            return True
-        resolved: Final = resolve_model_budget(
-            model=model,
-            model_max_budget=model_max_budget,
-            model_access_groups=self._model_access_groups,
-            team_id=team_id,
+        entities: Final[Mapping[Litellm_EntityType, ModelBudgetEntity]] = {
+            Litellm_EntityType.KEY: "Key",
+            Litellm_EntityType.USER: "User",
+            Litellm_EntityType.TEAM: "Team",
+            Litellm_EntityType.TEAM_MEMBER: "TeamMember",
+            Litellm_EntityType.END_USER: "EndUser",
+        }
+        reading: Final = await self.read_model_budget_raw(
+            entities[entity_type],
+            entity_id,
+            team_id,
+            model,
+            model_max_budget,
+            key_model_max_budget,
         )
-        if resolved is None:
-            verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value)
-            return True
-
-        max_budget: Final = resolved.budget_config.max_budget
-        if max_budget is None or max_budget < 0:
-            return True
-
-        local_spend: Final = await self._get_spend_for_model_budget(
-            entity_type=entity_type,
-            entity_id=entity_id,
-            model=model,
-            resolved=resolved,
-        )
-        current_spend: Final = (
-            local_spend
-            if local_spend >= max_budget
-            else await self._with_shared_spend(
-                entity_type=entity_type, entity_id=entity_id, resolved=resolved, local_spend=local_spend
-            )
-        )
-        if current_spend >= max_budget:
+        decision: Final = model_budget_data.check(reading.snapshot)
+        if decision.denied:
+            assert reading.snapshot.max_budget is not None
             raise litellm.BudgetExceededError(
-                message=(
-                    exceeded_message
-                    if resolved.budget_model in _budget_model_candidates(model)
-                    else f"{exceeded_message}, budget_model={resolved.budget_model}"
-                ),
-                current_cost=current_spend,
-                max_budget=max_budget,
+                message=exceeded_message
+                if reading.named_model
+                else f"{exceeded_message}, budget_model={reading.budget_model}",
+                current_cost=decision.current_spend,
+                max_budget=reading.snapshot.max_budget,
                 entity_type=entity_type.value,
                 entity_id=entity_id,
             )

litellm/rust_bridge/_native.pyi

@@ -817,3 +817,5 @@ class RouterCall(Generic[_S]):
     # candidate ids, pick first; pre_route: see _PreRouteStep; cache: see _CacheStep; secrets, build and model_info:
     # see those steps; first call: None.
     def next(self, outcome: object) -> _S: ...
+
+def request_auth_budget_exemption(facts: object) -> dict[str, object]: ...

tests/unit/proxy/auth/budget_reference.py

@@ -0,0 +1,44 @@
+from collections.abc import Mapping
+from typing import Final
+
+import litellm
+from litellm.proxy._types import Litellm_EntityType
+from litellm.proxy.hooks.model_max_budget_limiter import (
+    _PROXY_VirtualKeyModelMaxBudgetLimiter,
+    entry_names_model,
+    resolve_model_budget,
+)
+
+
+async def model_allowed(
+    limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter,
+    entity_type: Litellm_EntityType,
+    entity_id: str | None,
+    team_id: str | None,
+    budgets: Mapping[str, object] | None,
+    model: str,
+    message: str,
+) -> bool:
+    if not budgets:
+        return True
+    resolved: Final = resolve_model_budget(model, budgets, limiter._model_access_groups, team_id)
+    if resolved is None:
+        return True
+    limit: Final = resolved.budget_config.max_budget
+    if limit is None or limit < 0:
+        return True
+    local: Final = await limiter._get_spend_for_model_budget(entity_type, entity_id, model, resolved)
+    current: Final = (
+        local if local >= limit else await limiter._with_shared_spend(entity_type, entity_id, resolved, local)
+    )
+    if current >= limit:
+        raise litellm.BudgetExceededError(
+            message=message
+            if entry_names_model(resolved.budget_model, model)
+            else f"{message}, budget_model={resolved.budget_model}",
+            current_cost=current,
+            max_budget=limit,
+            entity_type=entity_type.value,
+            entity_id=entity_id,
+        )
+    return True

tests/unit/proxy/auth/test_budget_pricing_adapter.py

@@ -0,0 +1,121 @@
+from types import SimpleNamespace
+from typing import Final
+from unittest.mock import Mock
+
+import pytest
+
+import litellm
+from litellm._auth.budget_pricing import ModelPriceFacts
+from litellm.proxy.auth import auth_checks as reference
+from litellm.proxy.auth import budget_pricing_adapter as native
+
+
+def router(input_price: float | None, output_price: float | None, deployments: tuple[object, ...] = ()) -> Mock:
+    view: Final = Mock()
+    view.group.return_value = SimpleNamespace(input_cost_per_token=input_price, output_cost_per_token=output_price)
+    view.zero_cost.return_value = {}
+    view.deployments.return_value = deployments
+    return Mock(registry=Mock(return_value=view))
+
+
+@pytest.mark.parametrize(("input_price", "output_price"), ((None, 0), (0, None), (1, 0), (0, 1)))
+def test_priced_or_undefined_group_does_not_read_deployments(
+    input_price: float | None, output_price: float | None
+) -> None:
+    fake: Final = router(input_price, output_price)
+    facts: Final = native.read_model_price_facts("model", fake)
+    assert facts == ModelPriceFacts(group_found=True, input_price=input_price, output_price=output_price)
+    fake.registry().deployments.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "route",
+    (
+        "/v1/models",
+        "/model/info",
+        "/health/services",
+        "/chat/completions",
+        "/mcp",
+        "/mcp-rest/tools/list",
+        "/mcp/tools/call",
+        "/key/info",
+    ),
+)
+@pytest.mark.parametrize("body", (None, {}, {"method": "tools/list"}, {"method": "tools/call"}, {"method": 1}))
+def test_route_and_mcp_decision_matches_production_reference(route: str, body: dict[str, object] | None) -> None:
+    assert native.request_skips_budget_checks(
+        route, None, None, request_body=body
+    ) == reference.request_skips_budget_checks(route, None, None, request_body=body)
+    assert native.route_skips_budget_checks(route) == reference.route_skips_budget_checks(route)
+
+
+def test_missing_models_and_empty_model_list_preserve_exemption_behavior() -> None:
+    fake: Final = router(0, 0)
+    assert native.is_model_cost_zero(None, fake) is False
+    assert native.is_model_cost_zero("model", None) is False
+    assert native.is_model_cost_zero([], fake) is True
+
+
+@pytest.mark.parametrize(
+    ("prices", "reads", "exempt"), (((0, 1, 0), 2, False), ((0, 0, 0), 3, True), ((1, 0, 0), 1, False))
+)
+def test_raw_exemption_prices_only_native_selected_models_and_preserves_cache_writes(
+    monkeypatch: pytest.MonkeyPatch, prices: tuple[float, ...], reads: int, exempt: bool
+) -> None:
+    from litellm._auth import budget_pricing
+
+    old: Final = router(0, 0)
+    new: Final = router(0, 0)
+    monkeypatch.setitem(
+        litellm.model_cost, "budget-zero-fixture", {"input_cost_per_token": 0, "output_cost_per_token": 0}
+    )
+    deployment: Final = {
+        "model_name": "model",
+        "litellm_params": {"model": "openai/fixture"},
+        "model_info": {"id": "budget-zero-fixture"},
+    }
+    old.registry().deployments.return_value = (deployment,)
+    new.registry().deployments.return_value = (deployment,)
+    old.registry().resolve.side_effect = lambda model: model
+    new.registry().resolve.side_effect = lambda model: model
+    old.registry().group.side_effect = tuple(
+        SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices
+    )
+    new.registry().group.side_effect = tuple(
+        SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices
+    )
+    old_cache: Final[dict[str, bool]] = {}
+    new_cache: Final[dict[str, bool]] = {}
+    old.registry().zero_cost.return_value = old_cache
+    new.registry().zero_cost.return_value = new_cache
+    models: Final = ["first", "second", "third"]
+    before: Final = reference.request_skips_budget_checks("/chat/completions", models, old)
+    facts: Final = native.read_budget_exemption_facts("/chat/completions", models, new)
+    assert budget_pricing.exemption(facts).exempt is before is exempt
+    assert len(facts.prices) == reads
+    assert old.registry().group.call_count == new.registry().group.call_count == reads
+    assert old_cache == new_cache
+
+
+def test_raw_exemption_route_skip_does_not_touch_model_registry() -> None:
+    fake: Final = router(0, 0)
+    facts: Final = native.read_budget_exemption_facts("/v1/models", ["first", "second"], fake)
+    assert facts.prices == ()
+    fake.registry.assert_not_called()
+
+
+def test_raw_cached_price_does_not_read_model_group_or_write_again() -> None:
+    fake: Final = router(0, 0)
+    fake.registry().zero_cost.return_value = {"model": False}
+    facts: Final = native.read_model_cost_facts("model", fake)
+    assert facts.cached is False
+    fake.registry().group.assert_not_called()
+
+
+def test_empty_model_list_still_reads_zero_cost_cache() -> None:
+    fake: Final = router(0, 0)
+    facts: Final = native.read_budget_exemption_facts("/chat/completions", [], fake)
+    assert facts.prices == ()
+    assert facts.cache_read is True
+    fake.registry().zero_cost.assert_called_once_with()
+    fake.registry().group.assert_not_called()

tests/unit/proxy/auth/test_common_budget_adapter.py

@@ -0,0 +1,132 @@
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from math import inf
+from typing import Final
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+import litellm
+from litellm._auth.common_budgets import BudgetFacts
+from litellm.models.budget import LiteLLM_BudgetTable
+from litellm.proxy import proxy_server
+from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamTable, UserAPIKeyAuth
+from litellm.proxy.auth import auth_checks as reference
+from litellm.proxy.auth import common_budget_adapter as native
+from litellm.proxy.utils import ProxyLogging
+
+
+@dataclass(frozen=True, slots=True)
+class Observed:
+    error: tuple[object, ...] | None
+    alerts: tuple[object, ...]
+    throttle: float | None
+
+
+async def observe(run: Callable[[], Awaitable[None]], logging: Mock, key: UserAPIKeyAuth) -> Observed:
+    try:
+        await run()
+    except litellm.BudgetExceededError as error:
+        failure: Final = (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id)
+        await asyncio.sleep(0)
+        return Observed(failure, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct)
+    await asyncio.sleep(0)
+    return Observed(None, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2)))
+async def test_team_error_envelope_and_alert_projection_match_reference(
+    monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float
+) -> None:
+    reader: Final = AsyncMock(return_value=spend)
+    monkeypatch.setattr(proxy_server, "get_current_spend", reader)
+    team: Final = LiteLLM_TeamTable(team_id="team", max_budget=limit, spend=0)
+    key: Final = UserAPIKeyAuth(token="hash", team_id="team")
+    old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock())
+    new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock())
+    before: Final = await observe(lambda: reference.team_max_budget_check(team, key, old_log), old_log, key)
+    after: Final = await observe(lambda: native.team_max_budget_check(team, key, new_log), new_log, key)
+    assert after == before
+    assert reader.await_count == (0 if limit is None else 2)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2)))
+@pytest.mark.parametrize("throttle", (None, 0.5, 2))
+async def test_key_error_envelope_alert_and_throttle_mutation_match_reference(
+    monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float, throttle: float | None
+) -> None:
+    monkeypatch.setattr(proxy_server, "get_current_spend", AsyncMock(return_value=spend))
+    monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", throttle)
+    old: Final = UserAPIKeyAuth(
+        token="hash", key_alias="alias", max_budget=limit, tpm_limit=10, metadata={"throttle_on_budget_exceeded": True}
+    )
+    new: Final = old.model_copy(deep=True)
+    old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock())
+    new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock())
+    before: Final = await observe(lambda: reference.virtual_key_max_budget_check(old, old_log), old_log, old)
+    after: Final = await observe(lambda: native.virtual_key_max_budget_check(new, new_log), new_log, new)
+    assert after == before
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ("/chat/completions", "/key/info"))
+@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 1), (1, 2), (0, 0), (-1, 0)))
+async def test_end_user_info_skip_and_strict_ceiling_match_reference(
+    monkeypatch: pytest.MonkeyPatch, route: str, limit: float | None, spend: float
+) -> None:
+    reader: Final = AsyncMock(return_value=spend)
+    monkeypatch.setattr(proxy_server, "get_current_spend", reader)
+    user: Final = LiteLLM_EndUserTable(
+        user_id="end", blocked=False, litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit)
+    )
+    key: Final = UserAPIKeyAuth()
+    old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock())
+    new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock())
+    before: Final = await observe(lambda: reference.check_end_user_budget(user, route), old_log, key)
+    after: Final = await observe(lambda: native.check_end_user_budget(user, route), new_log, key)
+    assert after == before
+
+
+@pytest.mark.asyncio
+async def test_spend_io_failure_is_terminal_and_does_not_replay() -> None:
+    failure: Final = RuntimeError("counter unavailable")
+    reader: Final = AsyncMock(side_effect=failure)
+    with pytest.raises(RuntimeError, match="counter unavailable") as raised:
+        await native.read_budget(BudgetFacts("Key", 1), native.CounterRead("spend:key:hash", 0, 1), reader)
+    assert raised.value is failure
+    assert reader.await_count == 1
+
+
+@pytest.mark.parametrize(
+    ("member", "default", "increase", "expiry"),
+    (
+        (1, 5, 2, datetime(2025, 1, 2)),
+        (None, 5, 2, datetime(2025, 1, 2)),
+        (1, 5, 2, datetime(2025, 1, 1, tzinfo=timezone.utc)),
+        (None, 0, 2, datetime(2025, 1, 2)),
+        (None, -1, 2, None),
+        (0, 5, None, None),
+    ),
+)
+def test_member_budget_selection_and_temporary_expiry_match_reference(
+    member: float | None, default: float | None, increase: float | None, expiry: datetime | None
+) -> None:
+    now: Final = datetime(2025, 1, 1, tzinfo=timezone.utc)
+    row: Final = LiteLLM_BudgetTable(max_budget=member, temp_budget_increase=increase, temp_budget_expiry=expiry)
+    team: Final = LiteLLM_BudgetTable(max_budget=default)
+    assert native.resolve_effective_team_member_budget(
+        row, team, now
+    ) == reference.resolve_effective_team_member_budget(row, team, now)
+
+
+@pytest.mark.parametrize(
+    "policy", (reference.resolve_effective_team_member_budget, native.resolve_effective_team_member_budget)
+)
+def test_naive_now_preserves_datetime_error(policy: Callable[..., float | None]) -> None:
+    row: Final = LiteLLM_BudgetTable(max_budget=1, temp_budget_increase=2, temp_budget_expiry=datetime(2025, 1, 2))
+    with pytest.raises(TypeError, match="offset-naive"):
+        policy(row, None, datetime(2025, 1, 1))

tests/unit/proxy/auth/test_model_budget_adapter.py

@@ -0,0 +1,138 @@
+from collections.abc import Awaitable, Mapping
+from dataclasses import dataclass, field
+from typing import Final
+from unittest.mock import AsyncMock
+
+import pytest
+
+import litellm
+from litellm._auth import model_budget_data
+from litellm._auth.model_budget_data import ModelBudgetEntity
+from litellm.caching.caching import DualCache
+from litellm.proxy._types import Litellm_EntityType
+from litellm.proxy.auth.model_budget_adapter import (
+    ModelBudgetIO,
+    enforce_model_budget,
+    read_model_budget,
+    resolve_model_budget,
+)
+from litellm.proxy.hooks import model_max_budget_limiter as reference
+from litellm.proxy.hooks.model_budget_shared_windows import SharedWindow
+
+from .budget_reference import model_allowed
+
+
+@dataclass(frozen=True, slots=True)
+class Windows:
+    value: SharedWindow | None = None
+    read: AsyncMock = field(default_factory=AsyncMock)
+
+    async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None:
+        raise AssertionError("request auth must not open a window")
+
+    async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None:
+        raise AssertionError("request auth must not write a counter")
+
+
+def groups(model: str, team_id: str | None) -> frozenset[str]:
+    return frozenset(("group",))
+
+
+@pytest.mark.parametrize(
+    ("model", "budgets"),
+    (
+        (
+            "provider/model",
+            {"model": {"max_budget": 1, "budget_duration": "1h"}, "*": {"max_budget": 2, "budget_duration": "1h"}},
+        ),
+        ("model", {"model": {"budget_duration": "bad"}, "*": {"max_budget": 2, "budget_duration": "1h"}}),
+        ("model", {"model": {}, "group": {"max_budget": 2, "budget_duration": "1h"}}),
+        (
+            "model",
+            {"m*": {"max_budget": 2, "budget_duration": "1h"}, "group": {"max_budget": 1, "budget_duration": "1h"}},
+        ),
+        (
+            "mødél",
+            {"m*é*": {"max_budget": 2, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}},
+        ),
+        (
+            "model",
+            {"model": {"max_budget": None, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}},
+        ),
+    ),
+)
+def test_resolution_matches_production_reference(model: str, budgets: Mapping[str, object]) -> None:
+    io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0)
+    old: Final = reference.resolve_model_budget(model, budgets, groups)
+    result: Final = resolve_model_budget(model, budgets, io, None)
+    assert result.read_groups is False
+    assert (result.entry.name if result.entry else None) == (old.budget_model if old else None)
+    assert (result.entry.max_budget if result.entry else None) == (old.budget_config.max_budget if old else None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("entity", "scope"),
+    (
+        ("Key", Litellm_EntityType.KEY),
+        ("User", Litellm_EntityType.USER),
+        ("Team", Litellm_EntityType.TEAM),
+        ("TeamMember", Litellm_EntityType.TEAM_MEMBER),
+        ("EndUser", Litellm_EntityType.END_USER),
+    ),
+)
+@pytest.mark.parametrize(
+    ("local", "shared", "now", "limit"),
+    ((0, 2, 100, 1), (1, 0, 100, 1), (0, 2, 3700, 1), (0, 0, 100, 0), (0, 2, 100, -1), (0, 2, 100, None)),
+)
+async def test_counter_and_shared_window_match_production_reference(
+    entity: ModelBudgetEntity, scope: Litellm_EntityType, local: float, shared: float, now: float, limit: float | None
+) -> None:
+    cache: Final = DualCache()
+    key: Final = reference.model_budget_spend_cache_key(scope, "id", "model", "1h")
+    cache.set_cache(key, local)
+    windows: Final = Windows(read=AsyncMock(return_value=SharedWindow(100, shared)))
+    budgets: Final = {"model": {"max_budget": limit, "budget_duration": "1h"}}
+    old: Final = reference._PROXY_VirtualKeyModelMaxBudgetLimiter(
+        cache, clock=lambda: now, model_access_groups=groups, shared_windows=windows
+    )
+    io: Final = ModelBudgetIO(cache, groups, windows, 0, clock=lambda: now, monotonic=lambda: 0)
+    before: Final = await observed_model_result(model_allowed(old, scope, "id", None, budgets, "model", "exceeded"))
+    after: Final = await observed_model_result(
+        enforce_model_budget(io, entity, "id", None, "model", budgets, "exceeded")
+    )
+    assert after == before
+
+
+async def observed_model_result(call: Awaitable[bool]) -> tuple[object, ...]:
+    try:
+        return (await call,)
+    except litellm.BudgetExceededError as error:
+        return (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id)
+
+
+@pytest.mark.asyncio
+async def test_legacy_counter_is_added_to_the_current_counter_before_shared_read() -> None:
+    cache: Final = DualCache()
+    cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "model", "1h"), 0.4)
+    cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "provider/model", "1h"), 0.6)
+    windows: Final = Windows(read=AsyncMock(side_effect=AssertionError("already over local limit")))
+    io: Final = ModelBudgetIO(cache, groups, windows, 0, monotonic=lambda: 1)
+    reading: Final = await read_model_budget(
+        io, "Key", "id", None, "provider/model", {"model": {"max_budget": 1, "budget_duration": "1h"}}
+    )
+    assert reading.snapshot.local_spend == 0.4
+    assert reading.snapshot.legacy_spend == 0.6
+    assert model_budget_data.check(reading.snapshot).denied is True
+    windows.read.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_shared_read_failure_is_terminal_and_not_replayed() -> None:
+    failure: Final = RuntimeError("database unavailable")
+    windows: Final = Windows(read=AsyncMock(side_effect=failure))
+    io: Final = ModelBudgetIO(DualCache(), groups, windows, 0, monotonic=lambda: 0)
+    with pytest.raises(RuntimeError, match="database unavailable") as raised:
+        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