refactor(auth): decide key budget checks in AuthCore

The key checks that ran first in validate_resolved_virtual_key now run in
the Python core, in the same order: blocked key, model access, deactivated
owner, budget exemption, team member budget, expiry, budget alerts and the
key's own max budget or throttle. Lookups, counters and alerts are effects
the proxy answers, and internal jobs run the same checks through
start_resolved_key.

The soft budget, per-model budgets and the team, project and global spend
lookups stay in _finish_virtual_key_auth for a later pass. The source-shape
tests now read the functions that replaced validate_resolved_virtual_key

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

litellm/_auth/AGENTS.md

@@ -12,4 +12,5 @@
 - A core does no IO and reads no globals. The same facts, settings, time and effect answers always give the same outcome
 - A `Step` is either an `Outcome` or an effect. Each effect carries its typed `resume` continuation, and `driver.drive` answers effects through an `EffectHandler` until an outcome comes back. The proxy's handler (`_BuilderEffects` in `litellm/proxy/auth/user_api_key_auth.py`) owns all cache and DB IO, so both cores share it
 - Projection stays in the proxy. `Allow` grants become `UserAPIKeyAuth` in the builder, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract
-- The checks in `validate_resolved_virtual_key` and the centralized common checks still run as projection after `Allow(VirtualKey(...))`. Move them into the core one at a time, as effects plus pure decisions, and delete the builder code each one replaces in the same change
+- The key checks run in the core, in the order the proxy has always applied them: blocked key, model access, SCIM-deactivated owner, budget exemption, team member budget, expiry, budget alerts, then the key's own max budget (or its throttle). `start_resolved_key` runs the same checks for keys the proxy resolved itself, such as internal jobs
+- The soft budget alert, per-model budgets, the user's model budget and the team, project and global spend lookups still run as projection in `_finish_virtual_key_auth`. Move them into the core one at a time, as effects plus pure decisions, and delete the builder code each one replaces in the same change

litellm/_auth/decision.py

@@ -8,6 +8,7 @@ from typing import Protocol, TypeAlias
 class RequestFacts:
     route: str
     token: str | None
+    is_llm_api_route: bool
 
 
 @dataclass(frozen=True, slots=True)
@@ -41,6 +42,8 @@ class CachedAdminKey:
 @dataclass(frozen=True, slots=True)
 class VirtualKey:
     loaded_from_db: bool
+    budget_checks_skipped: bool
+    budget_throttle_pct: float | None
 
 
 Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey
@@ -65,6 +68,7 @@ class UserAuthDisabled:
 class ExpiredKey:
     expires: datetime
     now: datetime
+    key_label: str
 
 
 @dataclass(frozen=True, slots=True)
@@ -82,8 +86,42 @@ class NotAVirtualKey:
     pass
 
 
+@dataclass(frozen=True, slots=True)
+class KeyBlocked:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class OwnerDeactivated:
+    user_id: str
+
+
+@dataclass(frozen=True, slots=True)
+class TeamMemberBudgetExceeded:
+    user_id: str
+    team_id: str
+    spend: float
+    max_budget: float
+
+
+@dataclass(frozen=True, slots=True)
+class KeyBudgetExceeded:
+    spend: float
+    max_budget: float
+
+
 Denial: TypeAlias = (
-    MissingKey | MalformedKey | UserAuthDisabled | ExpiredKey | MasterKeyOnlyRoute | NoDatabase | NotAVirtualKey
+    MissingKey
+    | MalformedKey
+    | UserAuthDisabled
+    | ExpiredKey
+    | MasterKeyOnlyRoute
+    | NoDatabase
+    | NotAVirtualKey
+    | KeyBlocked
+    | OwnerDeactivated
+    | TeamMemberBudgetExceeded
+    | KeyBudgetExceeded
 )
 
 
@@ -113,6 +151,28 @@ class NonAdminKey:
 PresentedKey: TypeAlias = AdminKey | NonAdminKey
 
 
+@dataclass(frozen=True, slots=True)
+class KeyFacts:
+    blocked: bool
+    expires: datetime | None
+    max_budget: float | None
+    team_member_spend: float | None
+    user_id: str | None
+    team_id: str | None
+    budget_throttle_pct: float | None
+
+
+@dataclass(frozen=True, slots=True)
+class OwnerFacts:
+    deactivated: bool
+
+
+@dataclass(frozen=True, slots=True)
+class TeamMemberBudget:
+    spend: float
+    max_budget: float
+
+
 @dataclass(frozen=True, slots=True)
 class ResolveEndUser:
     prefetch_key: bool
@@ -137,9 +197,85 @@ class LoadKey:
     resume: "Callable[[], Step]"
 
 
-Effect: TypeAlias = ResolveEndUser | ResolvePresentedKey | EvictPresentedKey | LoadKey
+@dataclass(frozen=True, slots=True)
+class PrepareKey:
+    loaded_from_db: bool
+    resume: "Callable[[], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class InspectKey:
+    resume: "Callable[[KeyFacts], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class EnforceModelAccess:
+    resume: "Callable[[], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class LoadKeyOwner:
+    user_id: str
+    resume: "Callable[[OwnerFacts | None], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class CheckBudgetExemption:
+    resume: "Callable[[bool], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class ReadTeamMemberBudget:
+    user_id: str
+    team_id: str
+    resume: "Callable[[TeamMemberBudget | None], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class AlertTeamMemberBudget:
+    user_id: str
+    team_id: str
+    budget: TeamMemberBudget
+    resume: "Callable[[], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class SendKeyBudgetAlerts:
+    resume: "Callable[[], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class ReadKeySpend:
+    max_budget: float
+    resume: "Callable[[float], Step]"
+
+
+@dataclass(frozen=True, slots=True)
+class AlertKeyBudget:
+    spend: float
+    resume: "Callable[[], Step]"
+
+
+Effect: TypeAlias = (
+    ResolveEndUser
+    | ResolvePresentedKey
+    | EvictPresentedKey
+    | LoadKey
+    | PrepareKey
+    | InspectKey
+    | EnforceModelAccess
+    | LoadKeyOwner
+    | CheckBudgetExemption
+    | ReadTeamMemberBudget
+    | AlertTeamMemberBudget
+    | SendKeyBudgetAlerts
+    | ReadKeySpend
+    | AlertKeyBudget
+)
 Step: TypeAlias = Outcome | Effect
 
 
 class AuthCore(Protocol):
     def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ...
+
+    def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ...

litellm/_auth/driver.py

@@ -1,15 +1,29 @@
 from typing import Protocol, assert_never
 
 from litellm._auth.decision import (
+    AlertKeyBudget,
+    AlertTeamMemberBudget,
     Allow,
+    CheckBudgetExemption,
     Deny,
+    Effect,
+    EnforceModelAccess,
     EvictPresentedKey,
+    InspectKey,
+    KeyFacts,
     LoadKey,
+    LoadKeyOwner,
     Outcome,
+    OwnerFacts,
+    PrepareKey,
     PresentedKey,
+    ReadKeySpend,
+    ReadTeamMemberBudget,
     ResolveEndUser,
     ResolvePresentedKey,
+    SendKeyBudgetAlerts,
     Step,
+    TeamMemberBudget,
 )
 
 
@@ -22,21 +36,72 @@ class EffectHandler(Protocol):
 
     async def load_key(self, hashed_token: str) -> None: ...
 
+    async def prepare_key(self, loaded_from_db: bool) -> None: ...
+
+    async def inspect_key(self) -> KeyFacts: ...
+
+    async def enforce_model_access(self) -> None: ...
+
+    async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ...
+
+    async def check_budget_exemption(self) -> bool: ...
+
+    async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: ...
+
+    async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: ...
+
+    async def send_key_budget_alerts(self) -> None: ...
+
+    async def read_key_spend(self, max_budget: float) -> float: ...
+
+    async def alert_key_budget(self, spend: float) -> None: ...
+
+
+async def _answer(step: Effect, handler: EffectHandler) -> Step:
+    match step:
+        case ResolveEndUser(prefetch_key=prefetch_key, resume=resume):
+            await handler.resolve_end_user(prefetch_key)
+            return resume()
+        case ResolvePresentedKey(hashed_token=hashed_token, resume=resume):
+            return resume(await handler.resolve_presented_key(hashed_token))
+        case EvictPresentedKey(hashed_token=hashed_token, resume=resume):
+            await handler.evict_presented_key(hashed_token)
+            return resume()
+        case LoadKey(hashed_token=hashed_token, resume=resume):
+            await handler.load_key(hashed_token)
+            return resume()
+        case PrepareKey(loaded_from_db=loaded_from_db, resume=resume):
+            await handler.prepare_key(loaded_from_db)
+            return resume()
+        case InspectKey(resume=resume):
+            return resume(await handler.inspect_key())
+        case EnforceModelAccess(resume=resume):
+            await handler.enforce_model_access()
+            return resume()
+        case LoadKeyOwner(user_id=user_id, resume=resume):
+            return resume(await handler.load_key_owner(user_id))
+        case CheckBudgetExemption(resume=resume):
+            return resume(await handler.check_budget_exemption())
+        case ReadTeamMemberBudget(user_id=user_id, team_id=team_id, resume=resume):
+            return resume(await handler.read_team_member_budget(user_id, team_id))
+        case AlertTeamMemberBudget(user_id=user_id, team_id=team_id, budget=budget, resume=resume):
+            await handler.alert_team_member_budget(user_id, team_id, budget)
+            return resume()
+        case SendKeyBudgetAlerts(resume=resume):
+            await handler.send_key_budget_alerts()
+            return resume()
+        case ReadKeySpend(max_budget=max_budget, resume=resume):
+            return resume(await handler.read_key_spend(max_budget))
+        case AlertKeyBudget(spend=spend, resume=resume):
+            await handler.alert_key_budget(spend)
+            return resume()
+        case _:
+            assert_never(step)
+
 
 async def drive(step: Step, handler: EffectHandler) -> Outcome:
     match step:
         case Allow() | Deny():
             return step
-        case ResolveEndUser(prefetch_key=prefetch_key, resume=resume):
-            await handler.resolve_end_user(prefetch_key)
-            return await drive(resume(), handler)
-        case ResolvePresentedKey(hashed_token=hashed_token, resume=resume):
-            return await drive(resume(await handler.resolve_presented_key(hashed_token)), handler)
-        case EvictPresentedKey(hashed_token=hashed_token, resume=resume):
-            await handler.evict_presented_key(hashed_token)
-            return await drive(resume(), handler)
-        case LoadKey(hashed_token=hashed_token, resume=resume):
-            await handler.load_key(hashed_token)
-            return await drive(resume(), handler)
         case _:
-            assert_never(step)
+            return await drive(await _answer(step, handler), handler)

litellm/_auth/python_core.py

@@ -1,14 +1,25 @@
+import math
+from dataclasses import dataclass
 from datetime import datetime
 from typing import Final
 
 from litellm._auth.decision import (
     AdminKey,
+    AlertKeyBudget,
+    AlertTeamMemberBudget,
     Allow,
     CachedAdminKey,
+    CheckBudgetExemption,
     Deny,
+    EnforceModelAccess,
     EvictPresentedKey,
     ExpiredKey,
+    InspectKey,
+    KeyBlocked,
+    KeyBudgetExceeded,
+    KeyFacts,
     LoadKey,
+    LoadKeyOwner,
     MalformedKey,
     MasterKey,
     MasterKeyOnlyRoute,
@@ -17,12 +28,20 @@ from litellm._auth.decision import (
     NoMasterKey,
     NonAdminKey,
     NotAVirtualKey,
+    OwnerDeactivated,
+    OwnerFacts,
+    PrepareKey,
     PresentedKey,
+    ReadKeySpend,
+    ReadTeamMemberBudget,
     RequestFacts,
     ResolveEndUser,
     ResolvePresentedKey,
+    SendKeyBudgetAlerts,
     Settings,
     Step,
+    TeamMemberBudget,
+    TeamMemberBudgetExceeded,
     UserAuthDisabled,
     UserAuthRoute,
     VirtualKey,
@@ -55,6 +74,11 @@ class PythonCore:
             ),
         )
 
+    def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step:
+        return _check_key(
+            _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=True, key_label=facts.token or "")
+        )
+
 
 def _after_presented_key(
     presented: PresentedKey | None,
@@ -67,7 +91,8 @@ def _after_presented_key(
     match presented:
         case AdminKey(expires=expires) if expires is not None and expires < now:
             return EvictPresentedKey(
-                hashed_token=hashed_token, resume=lambda: Deny(ExpiredKey(expires=expires, now=now))
+                hashed_token=hashed_token,
+                resume=lambda: Deny(ExpiredKey(expires=expires, now=now, key_label=token)),
             )
         case AdminKey():
             return Allow(CachedAdminKey())
@@ -78,7 +103,114 @@ def _after_presented_key(
     if not settings.has_db:
         return Deny(NoDatabase())
     if presented is not None:
-        return Allow(VirtualKey(loaded_from_db=False))
+        cached: Final = _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=False, key_label=token)
+        return PrepareKey(loaded_from_db=False, resume=lambda: _check_key(cached))
     if not token.startswith(_VIRTUAL_KEY_PREFIX):
         return Deny(NotAVirtualKey())
-    return LoadKey(hashed_token=hashed_token, resume=lambda: Allow(VirtualKey(loaded_from_db=True)))
+    loaded: Final = _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=True, key_label=hashed_token)
+    return LoadKey(
+        hashed_token=hashed_token,
+        resume=lambda: PrepareKey(loaded_from_db=True, resume=lambda: _check_key(loaded)),
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class _KeyCheck:
+    facts: RequestFacts
+    settings: Settings
+    now: datetime
+    loaded_from_db: bool
+    key_label: str
+
+
+def _allow(check: _KeyCheck, budget_checks_skipped: bool, budget_throttle_pct: float | None) -> Step:
+    return Allow(
+        VirtualKey(
+            loaded_from_db=check.loaded_from_db,
+            budget_checks_skipped=budget_checks_skipped,
+            budget_throttle_pct=budget_throttle_pct,
+        )
+    )
+
+
+def _check_key(check: _KeyCheck) -> Step:
+    return InspectKey(resume=lambda key: _after_inspect(check, key))
+
+
+def _after_inspect(check: _KeyCheck, key: KeyFacts) -> Step:
+    if key.blocked:
+        return Deny(KeyBlocked())
+    return EnforceModelAccess(resume=lambda: _check_owner(check, key))
+
+
+def _check_owner(check: _KeyCheck, key: KeyFacts) -> Step:
+    user_id: Final = key.user_id
+    if user_id is None:
+        return _check_budget_exemption(check, key)
+    return LoadKeyOwner(user_id=user_id, resume=lambda owner: _after_owner(check, key, user_id, owner))
+
+
+def _after_owner(check: _KeyCheck, key: KeyFacts, user_id: str, owner: OwnerFacts | None) -> Step:
+    if owner is not None and owner.deactivated:
+        return Deny(OwnerDeactivated(user_id=user_id))
+    return _check_budget_exemption(check, key)
+
+
+def _check_budget_exemption(check: _KeyCheck, key: KeyFacts) -> Step:
+    return CheckBudgetExemption(resume=lambda exempt: _check_team_member_budget(check, key, exempt))
+
+
+def _check_team_member_budget(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step:
+    user_id: Final = key.user_id
+    team_id: Final = key.team_id
+    if exempt or key.team_member_spend is None or not check.settings.has_db or user_id is None or team_id is None:
+        return _check_expiry(check, key, exempt)
+    return ReadTeamMemberBudget(
+        user_id=user_id,
+        team_id=team_id,
+        resume=lambda budget: _after_team_member_budget(check, key, user_id, team_id, budget),
+    )
+
+
+def _after_team_member_budget(
+    check: _KeyCheck, key: KeyFacts, user_id: str, team_id: str, budget: TeamMemberBudget | None
+) -> Step:
+    if budget is None or budget.spend < budget.max_budget:
+        return _check_expiry(check, key, exempt=False)
+    return AlertTeamMemberBudget(
+        user_id=user_id,
+        team_id=team_id,
+        budget=budget,
+        resume=lambda: Deny(
+            TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=budget.spend, max_budget=budget.max_budget)
+        ),
+    )
+
+
+def _check_expiry(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step:
+    if key.expires is not None and key.expires < check.now:
+        return Deny(ExpiredKey(expires=key.expires, now=check.now, key_label=check.key_label))
+    if exempt:
+        return _allow(check, budget_checks_skipped=True, budget_throttle_pct=None)
+    return SendKeyBudgetAlerts(resume=lambda: _check_key_budget(check, key))
+
+
+def _check_key_budget(check: _KeyCheck, key: KeyFacts) -> Step:
+    max_budget: Final = key.max_budget
+    if not check.facts.is_llm_api_route or max_budget is None:
+        return _allow(check, budget_checks_skipped=False, budget_throttle_pct=None)
+    return ReadKeySpend(
+        max_budget=max_budget,
+        resume=lambda spend: AlertKeyBudget(
+            spend=spend, resume=lambda: _after_key_spend(check, key, spend, max_budget)
+        ),
+    )
+
+
+def _after_key_spend(check: _KeyCheck, key: KeyFacts, spend: float, max_budget: float) -> Step:
+    # A NaN max_budget compares false against any spend, so a non-finite budget means no limit, never a bypass.
+    if not math.isfinite(max_budget) or spend < max_budget:
+        return _allow(check, budget_checks_skipped=False, budget_throttle_pct=None)
+    if key.budget_throttle_pct is not None:
+        return _allow(check, budget_checks_skipped=False, budget_throttle_pct=key.budget_throttle_pct)
+    return Deny(KeyBudgetExceeded(spend=spend, max_budget=max_budget))

litellm/proxy/auth/auth_checks.py

@@ -6707,6 +6707,53 @@ def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool:
     return True
 
 
+def send_token_budget_alert(
+    valid_token: UserAPIKeyAuth,
+    spend: float,
+    proxy_logging_obj: ProxyLogging,
+    user_obj: LiteLLM_UserTable | None,
+) -> None:
+    call_info: Final = CallInfo(
+        token=valid_token.token,
+        spend=spend,
+        max_budget=valid_token.max_budget,
+        soft_budget=valid_token.soft_budget,
+        user_id=valid_token.user_id,
+        team_id=valid_token.team_id,
+        organization_id=valid_token.org_id,
+        user_email=user_obj.user_email if user_obj is not None else None,
+        key_alias=valid_token.key_alias,
+        event_group=Litellm_EntityType.KEY,
+    )
+    asyncio.create_task(
+        proxy_logging_obj.budget_alerts(
+            type="token_budget",
+            user_info=call_info,
+        )
+    )
+
+
+def key_budget_exceeded_error(
+    valid_token: UserAPIKeyAuth, spend: float, max_budget: float
+) -> litellm.BudgetExceededError:
+    # This message is returned to the caller, and key_name has no enforced
+    # shape (a direct DB write bypasses abbreviate_api_key), so echo it only
+    # when it still looks masked and fall back to the alias otherwise.
+    key_label: Final = valid_token.key_alias or "key"
+    key_descriptor: Final = (
+        f"{key_label} ({valid_token.key_name})"
+        if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name)
+        else key_label
+    )
+    return litellm.BudgetExceededError(
+        current_cost=spend,
+        max_budget=max_budget,
+        message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {max_budget}",
+        entity_type=Litellm_EntityType.KEY.value,
+        entity_id=valid_token.token,
+    )
+
+
 async def virtual_key_max_budget_check(
     valid_token: UserAPIKeyAuth,
     proxy_logging_obj: ProxyLogging,
@@ -6721,47 +6768,13 @@ async def virtual_key_max_budget_check(
     if valid_token.max_budget is not None:
         from litellm.proxy.proxy_server import get_current_spend
 
-        fallback_spend: Final = valid_token.spend or 0.0
-        counter_key: Final = f"spend:key:{valid_token.token}"
-
         # Read spend from cross-pod counter (Redis-first) or cached object (fallback)
         spend: Final = await get_current_spend(
-            counter_key=counter_key,
-            fallback_spend=fallback_spend,
+            counter_key=f"spend:key:{valid_token.token}",
+            fallback_spend=valid_token.spend or 0.0,
             max_budget=valid_token.max_budget,
         )
-
-        ####################################
-        # collect information for alerting #
-        ####################################
-
-        user_email = None
-        # Check if the token has any user id information
-        if user_obj is not None:
-            user_email = user_obj.user_email
-
-        call_info: Final = CallInfo(
-            token=valid_token.token,
-            spend=spend,
-            max_budget=valid_token.max_budget,
-            soft_budget=valid_token.soft_budget,
-            user_id=valid_token.user_id,
-            team_id=valid_token.team_id,
-            organization_id=valid_token.org_id,
-            user_email=user_email,
-            key_alias=valid_token.key_alias,
-            event_group=Litellm_EntityType.KEY,
-        )
-        asyncio.create_task(
-            proxy_logging_obj.budget_alerts(
-                type="token_budget",
-                user_info=call_info,
-            )
-        )
-
-        ####################################
-        # collect information for alerting #
-        ####################################
+        send_token_budget_alert(valid_token, spend, proxy_logging_obj, user_obj)
 
         # Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False,
         # so a NaN max_budget would silently disable enforcement.  Treat a
@@ -6769,22 +6782,7 @@ async def virtual_key_max_budget_check(
         if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget:
             if _apply_budget_exceeded_throttle(valid_token):
                 return
-            # This message is returned to the caller, and key_name has no enforced
-            # shape (a direct DB write bypasses abbreviate_api_key), so echo it only
-            # when it still looks masked and fall back to the alias otherwise.
-            key_label: Final = valid_token.key_alias or "key"
-            key_descriptor: Final = (
-                f"{key_label} ({valid_token.key_name})"
-                if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name)
-                else key_label
-            )
-            raise litellm.BudgetExceededError(
-                current_cost=spend,
-                max_budget=valid_token.max_budget,
-                message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}",
-                entity_type=Litellm_EntityType.KEY.value,
-                entity_id=valid_token.token,
-            )
+            raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget)
 
 
 _virtual_key_max_budget_check: Final = virtual_key_max_budget_check

litellm/proxy/auth/auth_denial.py

@@ -2,18 +2,24 @@ from typing import Final, assert_never
 
 from fastapi import HTTPException, status
 
+import litellm
 from litellm._auth.decision import (
     Denial,
     ExpiredKey,
+    KeyBlocked,
+    KeyBudgetExceeded,
     MalformedKey,
     MasterKeyOnlyRoute,
     MissingKey,
     NoDatabase,
     NotAVirtualKey,
+    OwnerDeactivated,
+    TeamMemberBudgetExceeded,
     UserAuthDisabled,
 )
 from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_KEY_ERROR_MESSAGE
-from litellm.proxy._types import ProxyErrorTypes, ProxyException
+from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
+from litellm.proxy.auth.auth_checks import key_budget_exceeded_error
 from litellm.proxy.auth.auth_utils import abbreviate_api_key
 from litellm.types.proxy.auth.user_api_key_auth import is_jwt
 
@@ -37,8 +43,12 @@ def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException:
 
 
 def denial_exception(
-    denial: Denial, token: str | None, route: str, jwt_auth_enabled: bool
-) -> ProxyException | HTTPException:
+    denial: Denial,
+    token: str | None,
+    route: str,
+    jwt_auth_enabled: bool,
+    key: UserAPIKeyAuth | None,
+) -> ProxyException | HTTPException | litellm.BudgetExceededError:
     match denial:
         case MissingKey():
             return ProxyException(
@@ -59,12 +69,12 @@ def denial_exception(
                 status_code=status.HTTP_403_FORBIDDEN,
                 detail="'allow_user_auth' not set or set to False",
             )
-        case ExpiredKey(expires=expires, now=now):
+        case ExpiredKey(expires=expires, now=now, key_label=key_label):
             return ProxyException(
                 message=f"Authentication Error - Expired Key. Key Expiry time {expires} and current time {now}",
                 type=ProxyErrorTypes.expired_key,
                 code=status.HTTP_401_UNAUTHORIZED,
-                param=abbreviate_api_key(api_key=token or ""),
+                param=abbreviate_api_key(api_key=key_label),
             )
         case MasterKeyOnlyRoute():
             return ProxyException(
@@ -82,5 +92,35 @@ def denial_exception(
             )
         case NotAVirtualKey():
             return _not_a_virtual_key(token or "", jwt_auth_enabled)
+        case KeyBlocked():
+            return ProxyException(
+                message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.",
+                type=ProxyErrorTypes.auth_error,
+                param=None,
+                code=status.HTTP_401_UNAUTHORIZED,
+            )
+        case OwnerDeactivated(user_id=user_id):
+            return ProxyException(
+                message=(
+                    f"Authentication Error, User={user_id} has been deactivated via SCIM. "
+                    "Keys owned by this user cannot be used."
+                ),
+                type=ProxyErrorTypes.auth_error,
+                param=None,
+                code=status.HTTP_401_UNAUTHORIZED,
+            )
+        case TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=spend, max_budget=max_budget):
+            return litellm.BudgetExceededError(
+                current_cost=spend,
+                max_budget=max_budget,
+                message=(
+                    f"Budget has been exceeded! TeamMember={user_id}:{team_id} "
+                    f"Current cost: {spend}, Max budget: {max_budget}"
+                ),
+                entity_type=Litellm_EntityType.TEAM_MEMBER.value,
+                entity_id=f"{user_id}:{team_id}",
+            )
+        case KeyBudgetExceeded(spend=spend, max_budget=max_budget):
+            return key_budget_exceeded_error(key or UserAPIKeyAuth(), spend, max_budget)
         case _:
             assert_never(denial)

litellm/proxy/auth/user_api_key_auth.py

@@ -36,12 +36,15 @@ from litellm._auth.decision import (
     AuthCore,
     CachedAdminKey,
     Deny,
+    KeyFacts,
     MasterKey,
     NoMasterKey,
     NonAdminKey,
+    OwnerFacts,
     PresentedKey,
     RequestFacts,
     Settings,
+    TeamMemberBudget,
     UserAuthRoute,
     VirtualKey,
 )
@@ -114,6 +117,7 @@ from litellm.proxy.auth.auth_checks import (  # noqa: F401  # legacy module expo
     resolve_and_validate_end_user_id,
     resolve_default_end_user_budget,
     route_skips_budget_checks,
+    send_token_budget_alert,
     stamp_matched_model_access_groups,
     team_member_max_budget_alert_check,
     virtual_key_max_budget_alert_check,
@@ -145,6 +149,7 @@ from litellm.proxy.auth.auth_utils import (
     request_dispatched_to_provider_pass_through,
     route_in_additonal_public_routes,
 )
+from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, should_throttle_budget_exceeded
 from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
 from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context
 from litellm.proxy.auth.oauth2_check import Oauth2Handler
@@ -2288,9 +2293,24 @@ async def _master_key_auth(
     return with_end_user
 
 
+def _proxy_model_list() -> list[object] | None:
+    from litellm.proxy.proxy_server import llm_model_list
+
+    return llm_model_list
+
+
 _MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value)
 
 
+def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: PrismaClient | None) -> Settings:
+    return Settings(
+        master_key=master_key,
+        allow_user_auth=allow_user_auth,
+        has_db=prisma_client is not None,
+        master_key_only_routes=_MASTER_KEY_ONLY_ROUTES,
+    )
+
+
 def _aware_expiry(expires: datetime | str) -> datetime:
     expiry: Final = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires)
     return (
@@ -2363,7 +2383,7 @@ class _BuilderEffects:
         self,
         *,
         request: Request,
-        request_data: Mapping[str, object],
+        request_data: dict[str, object],  # mutable-ok: model access checks rewrite the request model in place
         token: str,
         preresolved_key: UserAPIKeyAuth | None,
         prisma_client: PrismaClient | None,
@@ -2383,6 +2403,7 @@ class _BuilderEffects:
         self._route: Final = route
         self.end_user: _ResolvedEndUser = _ResolvedEndUser(end_user_id=None, end_user_object=None, budget=None)
         self.resolved_key: UserAPIKeyAuth | None = preresolved_key
+        self.owner: LiteLLM_UserTable | None = None
 
     @property
     def key(self) -> UserAPIKeyAuth:
@@ -2463,6 +2484,156 @@ class _BuilderEffects:
                 await self._identity_store(check_cache_only=False).resolve(hashed_token=hashed_token)
             )
 
+    async def prepare_key(self, loaded_from_db: bool) -> None:
+        if loaded_from_db:
+            _reset_end_user_fields(self.key, self.end_user)
+            return
+        await _refresh_team_fields_from_cache(
+            valid_token=self.key,
+            prisma_client=self._prisma_client,
+            user_api_key_cache=self._user_api_key_cache,
+            parent_otel_span=self._parent_otel_span,
+            proxy_logging_obj=self._proxy_logging_obj,
+        )
+
+    async def inspect_key(self) -> KeyFacts:
+        key: Final = _update_key_budget_with_temp_budget_increase(self.key)
+        self.resolved_key = key
+        return KeyFacts(
+            blocked=key.blocked is True,
+            expires=None if key.expires is None else _aware_expiry(key.expires),
+            max_budget=key.max_budget,
+            team_member_spend=key.team_member_spend,
+            user_id=key.user_id,
+            team_id=key.team_id,
+            budget_throttle_pct=budget_throttle_percentage() if should_throttle_budget_exceeded(key) else None,
+        )
+
+    async def enforce_model_access(self) -> None:
+        from litellm.proxy.proxy_server import llm_router
+
+        await enforce_key_and_fallback_model_access(
+            valid_token=self.key,
+            request_data=self._request_data,
+            route=self._route,
+            request=self._request,
+            llm_model_list=_proxy_model_list(),
+            llm_router=llm_router,
+        )
+        await _prefetch_referenced_auth_objects(
+            self.key,
+            end_user_id=self.end_user.end_user_id,
+            user_api_key_cache=self._user_api_key_cache,
+            prisma_client=self._prisma_client,
+        )
+
+    async def load_key_owner(self, user_id: str) -> OwnerFacts | None:
+        try:
+            with tracer.trace("litellm.proxy.auth.get_user_object"):
+                owner: Final = await get_user_object(
+                    user_id=user_id,
+                    prisma_client=self._prisma_client,
+                    user_api_key_cache=self._user_api_key_cache,
+                    user_id_upsert=False,
+                    parent_otel_span=self._parent_otel_span,
+                    proxy_logging_obj=self._proxy_logging_obj,
+                )
+        except Exception as e:
+            verbose_logger.debug("Unable to get user %s from db/cache, continuing without it: %s", user_id, e)
+            return None
+        self.owner = owner
+        if owner is None:
+            return None
+        key: Final = self.key
+        # The joint verification-token view carries the key's columns only, so the user's own per-model budget
+        # reaches enforcement and the post-call increment through the row fetched here.
+        key.user_model_max_budget = owner.model_max_budget
+        return OwnerFacts(deactivated=isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False)
+
+    async def check_budget_exemption(self) -> bool:
+        from litellm.proxy.proxy_server import llm_router
+
+        model: Final = _get_model_from_request_context(
+            request_data=self._request_data,
+            route=self._route,
+            request=self._request,
+            llm_router=llm_router,
+            team_id=self.key.team_id,
+        )
+        return request_skips_budget_checks(
+            route=self._route, model=model, llm_router=llm_router, request_body=self._request_data
+        )
+
+    async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None:
+        from litellm.proxy.proxy_server import get_current_spend
+
+        if self._prisma_client is None:
+            return None
+        cache_key: Final = team_membership_auth_cache_key(team_id=team_id, user_id=user_id)
+        cached: Final = await self._user_api_key_cache.async_get_cache(key=cache_key, model_type=LiteLLM_TeamMembership)
+        membership: Final = cached if cached is not None else await self._load_team_membership(user_id, team_id)
+        if membership is None or membership.litellm_budget_table is None:
+            return None
+        max_budget: Final = membership.litellm_budget_table.effective_max_budget(now=datetime.now(timezone.utc))
+        if max_budget is None or max_budget <= 0:
+            return None
+        spend: Final = await get_current_spend(
+            counter_key=f"spend:team_member:{user_id}:{team_id}",
+            fallback_spend=self.key.team_member_spend or 0.0,
+            max_budget=max_budget,
+        )
+        return TeamMemberBudget(spend=spend, max_budget=max_budget)
+
+    async def _load_team_membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None:
+        if self._prisma_client is None:
+            return None
+        row: Final = await TeamMembershipRepository(self._prisma_client).table.find_first(
+            where={"user_id": user_id, "team_id": team_id},
+            include={"litellm_budget_table": True},
+        )
+        if row is None:
+            return None
+        membership: Final = LiteLLM_TeamMembership.model_validate(row.model_dump())
+        await self._user_api_key_cache.async_set_cache(
+            key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
+            value=membership,
+            model_type=LiteLLM_TeamMembership,
+            ttl=5,
+        )
+        return membership
+
+    async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None:
+        # common_checks sends this alert on requests that get past here, so only the rejected request sends it here.
+        team_member_max_budget_alert_check(
+            team_id=team_id,
+            team_alias=self.key.team_alias,
+            team_metadata=self.key.team_metadata,
+            organization_id=self.key.org_id,
+            user_id=user_id,
+            user_email=self.owner.user_email if self.owner is not None else None,
+            proxy_logging_obj=self._proxy_logging_obj,
+            spend=budget.spend,
+            max_budget=budget.max_budget,
+        )
+
+    async def send_key_budget_alerts(self) -> None:
+        await virtual_key_max_budget_alert_check(
+            valid_token=self.key, proxy_logging_obj=self._proxy_logging_obj, user_obj=self.owner
+        )
+
+    async def read_key_spend(self, max_budget: float) -> float:
+        from litellm.proxy.proxy_server import get_current_spend
+
+        # Redis-first cross-pod counter, falling back to the key object's recorded spend.
+        return await get_current_spend(
+            counter_key=f"spend:key:{self.key.token}",
+            fallback_spend=self.key.spend or 0.0,
+            max_budget=max_budget,
+        )
+
+    async def alert_key_budget(self, spend: float) -> None:
+        send_token_budget_alert(self.key, spend, self._proxy_logging_obj, self.owner)
+
 
 async def user_api_key_auth_builder(
     request: Request,
@@ -2885,23 +3056,28 @@ async def user_api_key_auth_builder(
             proxy_logging_obj=proxy_logging_obj,
             route=route,
         )
-        outcome: Final = await drive(
-            _AUTH_CORE.start(
-                RequestFacts(route=route, token=api_key),
-                Settings(
-                    master_key=master_key,
-                    allow_user_auth=general_settings.get("allow_user_auth", False) is True,
-                    has_db=prisma_client is not None,
-                    master_key_only_routes=_MASTER_KEY_ONLY_ROUTES,
+        try:
+            outcome: Final = await drive(
+                _AUTH_CORE.start(
+                    RequestFacts(
+                        route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)
+                    ),
+                    _core_settings(
+                        master_key=master_key,
+                        allow_user_auth=general_settings.get("allow_user_auth", False) is True,
+                        prisma_client=prisma_client,
+                    ),
+                    datetime.now(timezone.utc),
                 ),
-                datetime.now(timezone.utc),
-            ),
-            effects,
-        )
-        valid_token = effects.resolved_key  # rebind-ok: the exception handler below reports the resolved identity
+                effects,
+            )
+        finally:
+            valid_token = effects.resolved_key  # rebind-ok: the exception handler below reports the resolved identity
         match outcome:
             case Deny(denial=denial):
-                raise denial_exception(denial, token=api_key, route=route, jwt_auth_enabled=enable_jwt_auth)
+                raise denial_exception(
+                    denial, token=api_key, route=route, jwt_auth_enabled=enable_jwt_auth, key=effects.resolved_key
+                )
             case Allow(grant=grant):
                 match grant:
                     case NoMasterKey(token=token):
@@ -2925,28 +3101,19 @@ async def user_api_key_auth_builder(
                         )
                     case CachedAdminKey():
                         return _with_end_user(effects.key, effects.end_user, parent_otel_span)
-                    case VirtualKey(loaded_from_db=loaded_from_db):
-                        if not loaded_from_db:
-                            await _refresh_team_fields_from_cache(
-                                valid_token=effects.key,
-                                prisma_client=prisma_client,
-                                user_api_key_cache=user_api_key_cache,
-                                parent_otel_span=parent_otel_span,
-                                proxy_logging_obj=proxy_logging_obj,
-                            )
-                        return await validate_resolved_virtual_key(
+                    case VirtualKey(budget_checks_skipped=budget_checks_skipped, budget_throttle_pct=throttle_pct):
+                        if throttle_pct is not None:
+                            effects.key.budget_throttle_pct = throttle_pct
+                        return await _finish_virtual_key_auth(
                             request=request,
                             request_data=request_data,
-                            valid_token=(
-                                _reset_end_user_fields(effects.key, effects.end_user) if loaded_from_db else effects.key
-                            ),
-                            api_key=hash_token(api_key) if loaded_from_db else api_key,
+                            valid_token=effects.key,
                             route=route,
                             start_time=start_time,
                             parent_otel_span=parent_otel_span,
-                            end_user_id=effects.end_user.end_user_id,
-                            end_user_params=effects.end_user.params,
-                            _end_user_object=effects.end_user.end_user_object,
+                            end_user=effects.end_user,
+                            user_obj=effects.owner,
+                            skip_budget_checks=budget_checks_skipped,
                         )
     except Exception as e:
         return await UserAPIKeyAuthExceptionHandler.handle_authentication_error(
@@ -2963,417 +3130,203 @@ async def user_api_key_auth_builder(
 _user_api_key_auth_builder: Final = user_api_key_auth_builder
 
 
-async def validate_resolved_virtual_key(  # noqa: C901  # Preserve ordering of existing shared authorization checks
+async def _check_soft_and_model_budgets(
     request: Request,
-    request_data: dict[str, object],
-    valid_token: UserAPIKeyAuth | None,
-    api_key: str,
+    request_data: dict[str, object],  # mutable-ok: per-model budget fallbacks rewrite the request model in place
+    valid_token: UserAPIKeyAuth,
     route: str,
-    start_time: datetime,
-    parent_otel_span: Span | None,
-    end_user_id: str | None,
-    end_user_params: dict[str, object],
-    _end_user_object: LiteLLM_EndUserTable | None,
-) -> UserAPIKeyAuth:
-    from litellm.proxy.proxy_server import (
-        litellm_proxy_admin_name,
-        llm_model_list,
-        llm_router,
-        model_max_budget_limiter,
-        prisma_client,
-        proxy_logging_obj,
-        user_api_key_cache,
+    user_obj: LiteLLM_UserTable | None,
+) -> None:
+    from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter, prisma_client, proxy_logging_obj
+
+    # Check 6. Soft Budget Check
+    await virtual_key_soft_budget_check(
+        valid_token=valid_token,
+        proxy_logging_obj=proxy_logging_obj,
+        user_obj=user_obj,
     )
 
-    if valid_token is not None:
-        valid_token = _update_key_budget_with_temp_budget_increase(valid_token)
+    # Check 5. Token Model Spend is under Model budget
+    max_budget_per_model: Final = valid_token.model_max_budget
+    current_model = _get_model_from_request_context(
+        request_data=request_data,
+        route=route,
+        request=request,
+        llm_router=llm_router,
+        team_id=valid_token.team_id,
+    )
+    current_models = _get_model_names_for_budget_checks(model=current_model)
 
-    user_obj: LiteLLM_UserTable | None = None
-    valid_token_dict: dict = {}
-    if valid_token is not None:
-        # Got Valid Token from Cache, DB
-        # Run checks for
-        # 1. If token can call model
-        ## 1a. If token can call fallback models (if client-side fallbacks given)
-        # 2. If user_id for this token is in budget
-        # 3. If the user spend within their own team is within budget
-        # 4. If 'user' passed to /chat/completions, /embeddings endpoint is in budget
-        # 5. If token is expired
-        # 6. If token spend is under Budget for the token
-        # 7. If token spend per model is under budget per model
-        # 8. If token spend is under team budget
-        # 9. If team spend is under team budget
-
-        ## base case ## key is disabled
-        if valid_token.blocked is True:
-            raise ProxyException(
-                message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.",
-                type=ProxyErrorTypes.auth_error,
-                param=None,
-                code=status.HTTP_401_UNAUTHORIZED,
+    if max_budget_per_model and prisma_client is not None and current_models and valid_token.token is not None:
+        ## GET THE SPEND FOR THIS MODEL
+        for model_name in current_models:
+            await _check_key_model_budget_with_fallback(
+                valid_token=valid_token,
+                model_max_budget_limiter=model_max_budget_limiter,
+                model_name=model_name,
+                request_data=request_data,
+                request=request,
+                llm_model_list=cast(  # cast-ok: router model lists are deployment sequences
+                    Sequence[object] | None, _proxy_model_list()
+                ),
+                llm_router=llm_router,
             )
-        await enforce_key_and_fallback_model_access(
-            valid_token=valid_token,
-            request_data=request_data,
-            route=route,
-            request=request,
-            llm_model_list=llm_model_list,
-            llm_router=llm_router,
-        )
-        await _prefetch_referenced_auth_objects(
-            valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client
-        )
 
-        # Check 2. If user_id for this token is in budget - done in common_checks()
-        if valid_token.user_id is not None:
-            try:
-                with tracer.trace("litellm.proxy.auth.get_user_object"):
-                    user_obj = await get_user_object(
-                        user_id=valid_token.user_id,
-                        prisma_client=prisma_client,
-                        user_api_key_cache=user_api_key_cache,
-                        user_id_upsert=False,
-                        parent_otel_span=parent_otel_span,
-                        proxy_logging_obj=proxy_logging_obj,
-                    )
-            except Exception as e:
-                verbose_logger.debug(
-                    "litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - %s",
-                    e,
-                )
-                user_obj = None
-
-            if user_obj is not None:
-                # The joint verification-token view carries the key's columns only, so the
-                # user's own per-model budget reaches enforcement and the post-call
-                # increment through the row fetched here.
-                valid_token.user_model_max_budget = user_obj.model_max_budget
-
-            if (
-                user_obj is not None
-                and isinstance(user_obj.metadata, dict)
-                and user_obj.metadata.get("scim_active") is False
-            ):
-                raise ProxyException(
-                    message=(
-                        f"Authentication Error, User={valid_token.user_id} has been deactivated via SCIM. "
-                        "Keys owned by this user cannot be used."
-                    ),
-                    type=ProxyErrorTypes.auth_error,
-                    param=None,
-                    code=status.HTTP_401_UNAUTHORIZED,
-                )
-
-        # Check 2a. Check if model has zero cost - if so, skip all budget checks
-        model = _get_model_from_request_context(
+        # Recompute after a potential budget-fallback rewrite so
+        # the user check below validates the final model
+        current_model = _get_model_from_request_context(
             request_data=request_data,
             route=route,
             request=request,
             llm_router=llm_router,
             team_id=valid_token.team_id,
         )
-        skip_budget_checks: Final = request_skips_budget_checks(
-            route=route, model=model, llm_router=llm_router, request_body=request_data
+        current_models = _get_model_names_for_budget_checks(model=current_model)
+
+    # Check 5a. Internal user model_max_budget
+    if current_models:
+        await _check_user_model_budget(
+            valid_token=valid_token,
+            model_max_budget_limiter=model_max_budget_limiter,
+            models=current_models,
         )
 
-        # Check 3. Check if user is in their team budget
-        if not skip_budget_checks and valid_token.team_member_spend is not None:
-            _user_id: Final = valid_token.user_id
-            _team_id: Final = valid_token.team_id
-            if prisma_client is not None and _user_id is not None and _team_id is not None:
-                _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id)
 
-                team_member_info = await user_api_key_cache.async_get_cache(
-                    key=_cache_key,
-                    model_type=LiteLLM_TeamMembership,
-                )
-                if team_member_info is None:
-                    # read from DB
-                    _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
-                        where={
-                            "user_id": _user_id,
-                            "team_id": _team_id,
-                        },
-                        include={"litellm_budget_table": True},
-                    )
-                    if _db_member is not None:
-                        team_member_info = LiteLLM_TeamMembership.model_validate(_db_member.model_dump())
-                        await user_api_key_cache.async_set_cache(
-                            key=_cache_key,
-                            value=team_member_info,
-                            model_type=LiteLLM_TeamMembership,
-                            ttl=5,
-                        )
+async def _finish_virtual_key_auth(
+    request: Request,
+    request_data: dict[str, object],  # mutable-ok: per-model budget fallbacks rewrite the request model in place
+    valid_token: UserAPIKeyAuth,
+    route: str,
+    start_time: datetime,
+    parent_otel_span: Span | None,
+    end_user: _ResolvedEndUser,
+    user_obj: LiteLLM_UserTable | None,
+    skip_budget_checks: bool,
+) -> UserAPIKeyAuth:
+    from litellm.proxy.proxy_server import (
+        litellm_proxy_admin_name,
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
 
-                if team_member_info is not None and team_member_info.litellm_budget_table is not None:
-                    team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget(
-                        now=datetime.now(timezone.utc),
-                    )
-                    if team_member_budget is not None and team_member_budget > 0:
-                        # Read from cross-pod counter (Redis-first) if available
-                        from litellm.proxy.proxy_server import get_current_spend
-
-                        team_member_spend = valid_token.team_member_spend
-                        if valid_token.user_id is not None and valid_token.team_id is not None:
-                            team_member_spend = await get_current_spend(
-                                counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}",
-                                fallback_spend=team_member_spend,
-                                max_budget=team_member_budget,
-                            )
-                        if team_member_spend >= team_member_budget:
-                            # common_checks sends this alert on requests that get past here, so only the
-                            # request rejected here sends it from the builder.
-                            team_member_max_budget_alert_check(
-                                team_id=_team_id,
-                                team_alias=valid_token.team_alias,
-                                team_metadata=valid_token.team_metadata,
-                                organization_id=valid_token.org_id,
-                                user_id=_user_id,
-                                user_email=user_obj.user_email if user_obj is not None else None,
-                                proxy_logging_obj=proxy_logging_obj,
-                                spend=team_member_spend,
-                                max_budget=team_member_budget,
-                            )
-                            _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
-                            raise litellm.BudgetExceededError(
-                                current_cost=team_member_spend,
-                                max_budget=team_member_budget,
-                                message=(
-                                    f"Budget has been exceeded! TeamMember={_entity_id} "
-                                    f"Current cost: {team_member_spend}, Max budget: {team_member_budget}"
-                                ),
-                                entity_type=Litellm_EntityType.TEAM_MEMBER.value,
-                                entity_id=_entity_id,
-                            )
-
-        # Check 3. If token is expired
-        if valid_token.expires is not None:
-            current_time = datetime.now(timezone.utc)
-            if isinstance(valid_token.expires, datetime):
-                expiry_time = valid_token.expires
-            else:
-                expiry_time = datetime.fromisoformat(valid_token.expires)
-            if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None:
-                expiry_time = expiry_time.replace(tzinfo=timezone.utc)
-            verbose_proxy_logger.debug(
-                "Checking if token expired, expiry time %s and current time %s", expiry_time, current_time
+    if not skip_budget_checks:
+        with tracer.trace("litellm.proxy.auth.budget_checks"):
+            await _check_soft_and_model_budgets(
+                request=request, request_data=request_data, valid_token=valid_token, route=route, user_obj=user_obj
             )
-            if expiry_time < current_time:
-                # Token exists but is expired.
-                raise ProxyException(
-                    message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
-                    type=ProxyErrorTypes.expired_key,
-                    code=status.HTTP_401_UNAUTHORIZED,
-                    param=abbreviate_api_key(api_key=api_key),
-                )
 
-        if not skip_budget_checks:
-            with tracer.trace("litellm.proxy.auth.budget_checks"):
-                # Check 4. Max Budget Alert Check (runs before budget enforcement
-                # so multi-threshold 100% alerts fire on the request that crosses
-                # max_budget, before BudgetExceededError is raised below)
-                await virtual_key_max_budget_alert_check(
-                    valid_token=valid_token,
-                    proxy_logging_obj=proxy_logging_obj,
-                    user_obj=user_obj,
-                )
-
-                # Check 5. Token Spend is under budget
-                if RouteChecks.is_llm_api_route(route=route):
-                    await virtual_key_max_budget_check(
-                        valid_token=valid_token,
-                        proxy_logging_obj=proxy_logging_obj,
-                        user_obj=user_obj,
-                    )
-
-                # Check 6. Soft Budget Check
-                await virtual_key_soft_budget_check(
-                    valid_token=valid_token,
-                    proxy_logging_obj=proxy_logging_obj,
-                    user_obj=user_obj,
-                )
-
-                # Check 5. Token Model Spend is under Model budget
-                max_budget_per_model: Final = valid_token.model_max_budget
-                current_model = _get_model_from_request_context(
-                    request_data=request_data,
-                    route=route,
-                    request=request,
-                    llm_router=llm_router,
+    # Check 6: Additional Common Checks across jwt + key auth
+    if valid_token.team_id is not None:
+        try:
+            if valid_token.team_id == UI_TEAM_ID:
+                raise TeamNotFoundError(team_id=UI_TEAM_ID)
+            with tracer.trace("litellm.proxy.auth.get_team_object"):
+                _team_obj = await get_team_object(
                     team_id=valid_token.team_id,
-                )
-                current_models = _get_model_names_for_budget_checks(model=current_model)
-
-                if (
-                    max_budget_per_model is not None
-                    and isinstance(max_budget_per_model, dict)
-                    and len(max_budget_per_model) > 0
-                    and prisma_client is not None
-                    and current_models
-                    and valid_token.token is not None
-                ):
-                    ## GET THE SPEND FOR THIS MODEL
-                    for model_name in current_models:
-                        await _check_key_model_budget_with_fallback(
-                            valid_token=valid_token,
-                            model_max_budget_limiter=model_max_budget_limiter,
-                            model_name=model_name,
-                            request_data=request_data,
-                            request=request,
-                            llm_model_list=cast(  # cast-ok: router model lists are deployment sequences
-                                Sequence[object] | None, llm_model_list
-                            ),
-                            llm_router=llm_router,
-                        )
-
-                    # Recompute after a potential budget-fallback rewrite so
-                    # the user check below validates the final model
-                    current_model = _get_model_from_request_context(
-                        request_data=request_data,
-                        route=route,
-                        request=request,
-                        llm_router=llm_router,
-                        team_id=valid_token.team_id,
-                    )
-                    current_models = _get_model_names_for_budget_checks(model=current_model)
-
-                # Check 5a. Internal user model_max_budget
-                if current_models:
-                    await _check_user_model_budget(
-                        valid_token=valid_token,
-                        model_max_budget_limiter=model_max_budget_limiter,
-                        models=current_models,
-                    )
-
-        # Check 6: Additional Common Checks across jwt + key auth
-        if valid_token.team_id is not None:
-            try:
-                if valid_token.team_id == UI_TEAM_ID:
-                    raise TeamNotFoundError(team_id=UI_TEAM_ID)
-                with tracer.trace("litellm.proxy.auth.get_team_object"):
-                    _team_obj = await get_team_object(
-                        team_id=valid_token.team_id,
-                        prisma_client=prisma_client,
-                        user_api_key_cache=user_api_key_cache,
-                        parent_otel_span=parent_otel_span,
-                        proxy_logging_obj=proxy_logging_obj,
-                    )
-            except HTTPException:
-                token_team_models: Final = _token_team_models(valid_token)
-                _team_obj = LiteLLM_TeamTableCachedObj(
-                    team_id=valid_token.team_id,
-                    max_budget=valid_token.team_max_budget,
-                    soft_budget=valid_token.team_soft_budget,
-                    model_max_budget=valid_token.team_model_max_budget,
-                    spend=valid_token.team_spend,
-                    tpm_limit=valid_token.team_tpm_limit,
-                    rpm_limit=valid_token.team_rpm_limit,
-                    tpd_limit=valid_token.team_tpd_limit,
-                    blocked=valid_token.team_blocked,
-                    models=token_team_models,
-                    metadata=valid_token.team_metadata,
-                    object_permission_id=valid_token.team_object_permission_id,
-                    object_permission=await _resolve_object_permission_for_unresolvable_team(
-                        object_permission_id=valid_token.team_object_permission_id,
-                        prisma_client=prisma_client,
-                        user_api_key_cache=user_api_key_cache,
-                        parent_otel_span=parent_otel_span,
-                        proxy_logging_obj=proxy_logging_obj,
-                    ),
-                )
-        else:
-            _team_obj = None
-
-        if _team_obj is not None:
-            valid_token.team_object_permission = _team_obj.object_permission
-            # Keep team_metadata in sync with the freshly fetched team so that
-            # guardrails (or any other metadata) added after the key was cached
-            # are picked up on subsequent requests without a cache eviction.
-            valid_token.team_metadata = _team_obj.metadata
-        else:
-            valid_token.team_object_permission = None
-
-        # Fetch project object if key belongs to a project
-        _project_obj = None
-        if valid_token.project_id is not None:
-            _project_obj = await get_project_object(
-                project_id=valid_token.project_id,
-                prisma_client=prisma_client,
-                user_api_key_cache=user_api_key_cache,
-                proxy_logging_obj=proxy_logging_obj,
-            )
-            if _project_obj is not None:
-                valid_token.project_metadata = _project_obj.metadata
-                valid_token.project_alias = _project_obj.project_alias
-
-        global_proxy_spend = None
-        if litellm.max_budget > 0 and prisma_client is not None:  # user set proxy max budget
-            cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
-            with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"):
-                global_proxy_spend = await fetch_global_spend_with_event_coordination(  # rebind-ok: pre-existing rebinding on a rename-only line
-                    cache_key=cache_key,
-                    user_api_key_cache=user_api_key_cache,
                     prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    parent_otel_span=parent_otel_span,
+                    proxy_logging_obj=proxy_logging_obj,
                 )
+        except HTTPException:
+            token_team_models: Final = _token_team_models(valid_token)
+            _team_obj = LiteLLM_TeamTableCachedObj(
+                team_id=valid_token.team_id,
+                max_budget=valid_token.team_max_budget,
+                soft_budget=valid_token.team_soft_budget,
+                model_max_budget=valid_token.team_model_max_budget,
+                spend=valid_token.team_spend,
+                tpm_limit=valid_token.team_tpm_limit,
+                rpm_limit=valid_token.team_rpm_limit,
+                tpd_limit=valid_token.team_tpd_limit,
+                blocked=valid_token.team_blocked,
+                models=token_team_models,
+                metadata=valid_token.team_metadata,
+                object_permission_id=valid_token.team_object_permission_id,
+                object_permission=await _resolve_object_permission_for_unresolvable_team(
+                    object_permission_id=valid_token.team_object_permission_id,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    parent_otel_span=parent_otel_span,
+                    proxy_logging_obj=proxy_logging_obj,
+                ),
+            )
+    else:
+        _team_obj = None
 
-            if global_proxy_spend is not None:
-                call_info: Final = CallInfo(
-                    token=valid_token.token,
-                    spend=global_proxy_spend,
-                    max_budget=litellm.max_budget,
-                    user_id=litellm_proxy_admin_name,
-                    team_id=valid_token.team_id,
-                    event_group=Litellm_EntityType.PROXY,
-                )
-                asyncio.create_task(
-                    proxy_logging_obj.budget_alerts(
-                        type="proxy_budget",
-                        user_info=call_info,
-                    )
-                )
-        # Token passed all checks
-        if valid_token is None:
-            raise HTTPException(401, detail="Invalid API key")
-        if valid_token.token is None:
-            raise HTTPException(401, detail="Invalid API key, no token associated")
-        api_key = valid_token.token
+    if _team_obj is not None:
+        valid_token.team_object_permission = _team_obj.object_permission
+        # Keep team_metadata in sync with the freshly fetched team so that
+        # guardrails (or any other metadata) added after the key was cached
+        # are picked up on subsequent requests without a cache eviction.
+        valid_token.team_metadata = _team_obj.metadata
+    else:
+        valid_token.team_object_permission = None
 
-        valid_token_dict = valid_token.model_dump(exclude_none=True)
-        valid_token_dict.pop("token", None)
-        # budget_throttle_pct is excluded from model_dump (it must not leak
-        # into serialized responses), so carry the request-scoped decision
-        # forward by hand to the auth object the rate limiter receives.
-        if valid_token.budget_throttle_pct is not None:
-            valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct
-
-        if _end_user_object is not None or litellm.max_end_user_budget_id is not None:
-            valid_token_dict.update(end_user_params)
-        if _end_user_object is not None:
-            valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission
-
-    # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions
-    # sso/login, ui/login, /key functions and /user functions
-    # this will never be allowed to call /chat/completions
-
-    if valid_token is None:
-        # No token was found when looking up in the DB
-        raise ProxyException(
-            message="Authentication Error, Invalid proxy server token passed",
-            type=ProxyErrorTypes.auth_error,
-            param=None,
-            code=status.HTTP_401_UNAUTHORIZED,
+    # Fetch project object if key belongs to a project
+    _project_obj = None
+    if valid_token.project_id is not None:
+        _project_obj = await get_project_object(
+            project_id=valid_token.project_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
         )
-    if valid_token_dict is not None:
-        virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj(
-            user_obj=user_obj,
-            api_key=api_key,
-            parent_otel_span=parent_otel_span,
-            valid_token_dict=valid_token_dict,
-            route=route,
-            start_time=start_time,
-        )
-        virtual_key_auth_obj.via_virtual_key = True
-        return virtual_key_auth_obj
+        if _project_obj is not None:
+            valid_token.project_metadata = _project_obj.metadata
+            valid_token.project_alias = _project_obj.project_alias
+
+    global_proxy_spend = None
+    if litellm.max_budget > 0 and prisma_client is not None:  # user set proxy max budget
+        cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
+        with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"):
+            global_proxy_spend = await fetch_global_spend_with_event_coordination(  # rebind-ok: pre-existing rebinding on a rename-only line
+                cache_key=cache_key,
+                user_api_key_cache=user_api_key_cache,
+                prisma_client=prisma_client,
+            )
+
+        if global_proxy_spend is not None:
+            call_info: Final = CallInfo(
+                token=valid_token.token,
+                spend=global_proxy_spend,
+                max_budget=litellm.max_budget,
+                user_id=litellm_proxy_admin_name,
+                team_id=valid_token.team_id,
+                event_group=Litellm_EntityType.PROXY,
+            )
+            asyncio.create_task(
+                proxy_logging_obj.budget_alerts(
+                    type="proxy_budget",
+                    user_info=call_info,
+                )
+            )
+    if valid_token.token is None:
+        raise HTTPException(401, detail="Invalid API key, no token associated")
+
+    valid_token_dict: Final = valid_token.model_dump(exclude_none=True)
+    valid_token_dict.pop("token", None)
+    # budget_throttle_pct is excluded from model_dump (it must not leak
+    # into serialized responses), so carry the request-scoped decision
+    # forward by hand to the auth object the rate limiter receives.
+    if valid_token.budget_throttle_pct is not None:
+        valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct
+    if end_user.end_user_object is not None or litellm.max_end_user_budget_id is not None:
+        valid_token_dict.update(end_user.params)
+    if end_user.end_user_object is not None:
+        valid_token_dict["end_user_object_permission"] = end_user.end_user_object.object_permission
+
+    virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj(
+        user_obj=user_obj,
+        api_key=valid_token.token,
+        parent_otel_span=parent_otel_span,
+        valid_token_dict=valid_token_dict,
+        route=route,
+        start_time=start_time,
+    )
+    virtual_key_auth_obj.via_virtual_key = True
+    return virtual_key_auth_obj
 
 
 async def _safe_fetch(label: str, awaitable):
@@ -4832,6 +4785,61 @@ async def _run_post_custom_auth_checks(
     return valid_token
 
 
+async def _authorize_resolved_key(
+    identity: UserAPIKeyAuth,
+    key_hash: str,
+    request: Request,
+    request_data: dict[str, object],  # mutable-ok: model access checks rewrite the request model in place
+    route: str,
+) -> UserAPIKeyAuth:
+    from litellm.proxy.proxy_server import (
+        master_key,
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
+
+    start_time: Final = datetime.now(timezone.utc)
+    effects: Final = _BuilderEffects(
+        request=request,
+        request_data=request_data,
+        token=key_hash,
+        preresolved_key=identity,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        parent_otel_span=None,
+        proxy_logging_obj=proxy_logging_obj,
+        route=route,
+    )
+    outcome: Final = await drive(
+        _AUTH_CORE.start_resolved_key(
+            RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)),
+            _core_settings(master_key=master_key, allow_user_auth=False, prisma_client=prisma_client),
+            start_time,
+        ),
+        effects,
+    )
+    match outcome:
+        case Allow(grant=VirtualKey(budget_checks_skipped=budget_checks_skipped, budget_throttle_pct=throttle_pct)):
+            if throttle_pct is not None:
+                effects.key.budget_throttle_pct = throttle_pct
+            return await _finish_virtual_key_auth(
+                request=request,
+                request_data=request_data,
+                valid_token=effects.key,
+                route=route,
+                start_time=start_time,
+                parent_otel_span=None,
+                end_user=effects.end_user,
+                user_obj=effects.owner,
+                skip_budget_checks=budget_checks_skipped,
+            )
+        case Allow(grant=grant):
+            raise RuntimeError(f"a resolved key check granted {grant}")
+        case Deny(denial=denial):
+            raise denial_exception(denial, token=key_hash, route=route, jwt_auth_enabled=False, key=effects.key)
+
+
 async def authorize_internal_virtual_key(
     key_hash: str, request: Request, request_data: dict[str, object]
 ) -> UserAPIKeyAuth:
@@ -4845,17 +4853,8 @@ async def authorize_internal_virtual_key(
     )
     route: Final = get_request_route(request=request)
     await pre_db_read_auth_checks(request_data=request_data, request=request, route=route)
-    auth: Final = await validate_resolved_virtual_key(
-        request=request,
-        request_data=request_data,
-        valid_token=identity,
-        api_key=key_hash,
-        route=route,
-        start_time=datetime.now(timezone.utc),
-        parent_otel_span=None,
-        end_user_id=None,
-        end_user_params={},
-        _end_user_object=None,
+    auth: Final = await _authorize_resolved_key(
+        identity=identity, key_hash=key_hash, request=request, request_data=request_data, route=route
     )
     auth.budget_reservation = None
     recovered: Final = await _authorize_authenticated_request(

tests/unit/_auth/test_python_core.py

@@ -1,4 +1,5 @@
-from dataclasses import dataclass
+import math
+from dataclasses import dataclass, replace
 from datetime import datetime, timedelta, timezone
 from typing import Final
 
@@ -10,6 +11,9 @@ from litellm._auth.decision import (
     CachedAdminKey,
     Deny,
     ExpiredKey,
+    KeyBlocked,
+    KeyBudgetExceeded,
+    KeyFacts,
     MalformedKey,
     MasterKey,
     MasterKeyOnlyRoute,
@@ -19,9 +23,14 @@ from litellm._auth.decision import (
     NonAdminKey,
     NotAVirtualKey,
     Outcome,
+    OwnerDeactivated,
+    OwnerFacts,
     PresentedKey,
     RequestFacts,
     Settings,
+    Step,
+    TeamMemberBudget,
+    TeamMemberBudgetExceeded,
     UserAuthDisabled,
     UserAuthRoute,
     VirtualKey,
@@ -38,25 +47,95 @@ SETTINGS: Final = Settings(
     has_db=True,
     master_key_only_routes=frozenset({"/global/spend/reset"}),
 )
+PLAIN_KEY: Final = KeyFacts(
+    blocked=False,
+    expires=None,
+    max_budget=None,
+    team_member_spend=None,
+    user_id=None,
+    team_id=None,
+    budget_throttle_pct=None,
+)
+CHECKED: Final = (
+    "prepare_key",
+    "inspect_key",
+    "enforce_model_access",
+    "check_budget_exemption",
+    "send_key_budget_alerts",
+)
+
+Call = tuple[str, object]
 
 
 @dataclass
 class RecordingHandler:
     presented: PresentedKey | None = None
-    calls: tuple[tuple[str, object], ...] = ()
+    key: KeyFacts = PLAIN_KEY
+    owner: OwnerFacts | None = None
+    exempt: bool = False
+    team_member_budget: TeamMemberBudget | None = None
+    key_spend: float = 0.0
+    calls: tuple[Call, ...] = ()
+
+    def _record(self, name: str, payload: object = None) -> None:
+        self.calls = (*self.calls, (name, payload))
 
     async def resolve_end_user(self, prefetch_key: bool) -> None:
-        self.calls = (*self.calls, ("resolve_end_user", prefetch_key))
+        self._record("resolve_end_user", prefetch_key)
 
     async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None:
-        self.calls = (*self.calls, ("resolve_presented_key", hashed_token))
+        self._record("resolve_presented_key", hashed_token)
         return self.presented
 
     async def evict_presented_key(self, hashed_token: str) -> None:
-        self.calls = (*self.calls, ("evict_presented_key", hashed_token))
+        self._record("evict_presented_key", hashed_token)
 
     async def load_key(self, hashed_token: str) -> None:
-        self.calls = (*self.calls, ("load_key", hashed_token))
+        self._record("load_key", hashed_token)
+
+    async def prepare_key(self, loaded_from_db: bool) -> None:
+        self._record("prepare_key", loaded_from_db)
+
+    async def inspect_key(self) -> KeyFacts:
+        self._record("inspect_key")
+        return self.key
+
+    async def enforce_model_access(self) -> None:
+        self._record("enforce_model_access")
+
+    async def load_key_owner(self, user_id: str) -> OwnerFacts | None:
+        self._record("load_key_owner", user_id)
+        return self.owner
+
+    async def check_budget_exemption(self) -> bool:
+        self._record("check_budget_exemption")
+        return self.exempt
+
+    async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None:
+        self._record("read_team_member_budget", (user_id, team_id))
+        return self.team_member_budget
+
+    async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None:
+        self._record("alert_team_member_budget", budget)
+
+    async def send_key_budget_alerts(self) -> None:
+        self._record("send_key_budget_alerts")
+
+    async def read_key_spend(self, max_budget: float) -> float:
+        self._record("read_key_spend", max_budget)
+        return self.key_spend
+
+    async def alert_key_budget(self, spend: float) -> None:
+        self._record("alert_key_budget", spend)
+
+
+def _facts(token: str | None, route: str = "/chat/completions", is_llm_api_route: bool = True) -> RequestFacts:
+    return RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route)
+
+
+async def _run(step: Step, handler: RecordingHandler) -> tuple[Outcome, tuple[str, ...]]:
+    outcome: Final = await drive(step, handler)
+    return outcome, tuple(name for name, _ in handler.calls)
 
 
 async def _decide(
@@ -64,11 +143,34 @@ async def _decide(
     *,
     route: str = "/chat/completions",
     settings: Settings = SETTINGS,
-    presented: PresentedKey | None = None,
-) -> tuple[Outcome, tuple[tuple[str, object], ...]]:
-    handler: Final = RecordingHandler(presented=presented)
-    outcome: Final = await drive(PythonCore().start(RequestFacts(route=route, token=token), settings, NOW), handler)
-    return outcome, handler.calls
+    handler: RecordingHandler | None = None,
+) -> tuple[Outcome, tuple[str, ...]]:
+    return await _run(PythonCore().start(_facts(token, route), settings, NOW), handler or RecordingHandler())
+
+
+async def _check(
+    key: KeyFacts,
+    *,
+    is_llm_api_route: bool = True,
+    settings: Settings = SETTINGS,
+    owner: OwnerFacts | None = None,
+    exempt: bool = False,
+    team_member_budget: TeamMemberBudget | None = None,
+    key_spend: float = 0.0,
+) -> tuple[Outcome, tuple[str, ...]]:
+    handler: Final = RecordingHandler(
+        key=key, owner=owner, exempt=exempt, team_member_budget=team_member_budget, key_spend=key_spend
+    )
+    facts: Final = _facts("hash", is_llm_api_route=is_llm_api_route)
+    return await _run(PythonCore().start_resolved_key(facts, settings, NOW), handler)
+
+
+def _allowed(budget_checks_skipped: bool = False, budget_throttle_pct: float | None = None) -> Outcome:
+    return Allow(
+        VirtualKey(
+            loaded_from_db=True, budget_checks_skipped=budget_checks_skipped, budget_throttle_pct=budget_throttle_pct
+        )
+    )
 
 
 @pytest.mark.asyncio
@@ -92,71 +194,206 @@ async def test_credential_decisions_do_no_io(route: str, token: str | None, sett
 @pytest.mark.asyncio
 @pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"])
 async def test_master_key_resolves_end_user_without_key_prefetch(route: str):
-    assert await _decide(MASTER, route=route) == (Allow(MasterKey()), (("resolve_end_user", False),))
+    handler: Final = RecordingHandler()
+    assert await _decide(MASTER, route=route, handler=handler) == (Allow(MasterKey()), ("resolve_end_user",))
+    assert handler.calls == (("resolve_end_user", False),)
 
 
 @pytest.mark.asyncio
-async def test_unknown_virtual_key_is_loaded_from_the_db_after_a_cache_miss():
-    hashed: Final = hash_token("sk-x")
-    assert await _decide("sk-x") == (
-        Allow(VirtualKey(loaded_from_db=True)),
-        (("resolve_end_user", True), ("resolve_presented_key", hashed), ("load_key", hashed)),
-    )
+async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked():
+    handler: Final = RecordingHandler()
+    outcome, calls = await _decide("sk-x", handler=handler)
+    assert outcome == Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None))
+    assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", *CHECKED)
+    assert handler.calls[2] == ("load_key", hash_token("sk-x"))
+    assert handler.calls[3] == ("prepare_key", True)
 
 
 @pytest.mark.asyncio
-async def test_cached_virtual_key_is_not_loaded_again():
-    assert await _decide("sk-x", presented=NonAdminKey()) == (
-        Allow(VirtualKey(loaded_from_db=False)),
-        (("resolve_end_user", True), ("resolve_presented_key", hash_token("sk-x"))),
-    )
+async def test_cached_virtual_key_is_checked_without_a_db_load():
+    handler: Final = RecordingHandler(presented=NonAdminKey())
+    outcome, calls = await _decide("sk-x", handler=handler)
+    assert outcome == Allow(VirtualKey(loaded_from_db=False, budget_checks_skipped=False, budget_throttle_pct=None))
+    assert calls == ("resolve_end_user", "resolve_presented_key", *CHECKED)
+    assert handler.calls[2] == ("prepare_key", False)
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW])
-async def test_unexpired_cached_admin_key_skips_route_and_db_checks(expires: datetime | None):
+async def test_unexpired_cached_admin_key_skips_route_db_and_key_checks(expires: datetime | None):
     no_db: Final = Settings(MASTER, False, False, frozenset({"/global/spend/reset"}))
-    outcome, _ = await _decide(
-        "sk-admin", route="/global/spend/reset", settings=no_db, presented=AdminKey(expires=expires)
+    outcome, calls = await _decide(
+        "sk-admin",
+        route="/global/spend/reset",
+        settings=no_db,
+        handler=RecordingHandler(presented=AdminKey(expires=expires)),
     )
-    assert outcome == Allow(CachedAdminKey())
+    assert (outcome, calls) == (Allow(CachedAdminKey()), ("resolve_end_user", "resolve_presented_key"))
 
 
 @pytest.mark.asyncio
-async def test_expired_cached_admin_key_is_evicted_then_denied():
+async def test_expired_cached_admin_key_is_evicted_then_denied_with_the_presented_key():
     expires: Final = NOW - timedelta(seconds=1)
-    hashed: Final = hash_token("sk-admin")
-    assert await _decide("sk-admin", presented=AdminKey(expires=expires)) == (
-        Deny(ExpiredKey(expires=expires, now=NOW)),
-        (("resolve_end_user", True), ("resolve_presented_key", hashed), ("evict_presented_key", hashed)),
-    )
+    handler: Final = RecordingHandler(presented=AdminKey(expires=expires))
+    outcome, calls = await _decide("sk-admin", handler=handler)
+    assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label="sk-admin"))
+    assert calls == ("resolve_end_user", "resolve_presented_key", "evict_presented_key")
+    assert handler.calls[-1] == ("evict_presented_key", hash_token("sk-admin"))
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("presented", [None, NonAdminKey()])
 async def test_non_admin_key_cannot_use_master_key_only_routes(presented: PresentedKey | None):
-    outcome, calls = await _decide("sk-x", route="/global/spend/reset", presented=presented)
-    assert outcome == Deny(MasterKeyOnlyRoute())
-    assert "load_key" not in {name for name, _ in calls}
+    outcome, calls = await _decide("sk-x", route="/global/spend/reset", handler=RecordingHandler(presented=presented))
+    assert (outcome, calls) == (Deny(MasterKeyOnlyRoute()), ("resolve_end_user", "resolve_presented_key"))
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("presented", [None, NonAdminKey()])
 async def test_virtual_key_needs_a_database(presented: PresentedKey | None):
     no_db: Final = Settings(MASTER, False, False, frozenset())
-    outcome, _ = await _decide("sk-x", settings=no_db, presented=presented)
+    outcome, _ = await _decide("sk-x", settings=no_db, handler=RecordingHandler(presented=presented))
     assert outcome == Deny(NoDatabase())
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("token", ["eyJhbGciOiJIUzI1NiJ9.e30.sig", "a" * 64, "sk"])
 async def test_uncached_token_without_virtual_key_prefix_is_rejected_before_the_db(token: str):
-    outcome, calls = await _decide(token)
-    assert outcome == Deny(NotAVirtualKey())
-    assert [name for name, _ in calls] == ["resolve_end_user", "resolve_presented_key"]
+    assert await _decide(token) == (Deny(NotAVirtualKey()), ("resolve_end_user", "resolve_presented_key"))
 
 
 @pytest.mark.asyncio
-async def test_cached_token_without_virtual_key_prefix_is_accepted():
-    outcome, _ = await _decide("ui-login-blob", presented=NonAdminKey())
-    assert outcome == Allow(VirtualKey(loaded_from_db=False))
+async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with():
+    expires: Final = NOW - timedelta(seconds=1)
+    outcome, _ = await _decide("sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires)))
+    assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x")))
+
+
+@pytest.mark.asyncio
+async def test_blocked_key_is_denied_before_model_access():
+    assert await _check(replace(PLAIN_KEY, blocked=True)) == (Deny(KeyBlocked()), ("inspect_key",))
+
+
+@pytest.mark.asyncio
+async def test_deactivated_owner_is_denied_after_model_access():
+    outcome, calls = await _check(replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True))
+    assert (outcome, calls) == (
+        Deny(OwnerDeactivated(user_id="u1")),
+        ("inspect_key", "enforce_model_access", "load_key_owner"),
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("owner", [None, OwnerFacts(deactivated=False)])
+async def test_active_or_missing_owner_continues(owner: OwnerFacts | None):
+    outcome, _ = await _check(replace(PLAIN_KEY, user_id="u1"), owner=owner)
+    assert outcome == _allowed()
+
+
+@pytest.mark.asyncio
+async def test_exempt_request_skips_every_budget_check_but_not_expiry():
+    key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=99.0, max_budget=1.0)
+    assert await _check(key, exempt=True) == (
+        _allowed(budget_checks_skipped=True),
+        ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption"),
+    )
+    expired: Final = replace(key, expires=NOW - timedelta(seconds=1))
+    outcome, _ = await _check(expired, exempt=True)
+    assert outcome == Deny(ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash"))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("budget", "denied"),
+    [
+        (TeamMemberBudget(spend=10.0, max_budget=10.0), True),
+        (TeamMemberBudget(spend=10.1, max_budget=10.0), True),
+        (TeamMemberBudget(spend=9.9, max_budget=10.0), False),
+        (None, False),
+    ],
+)
+async def test_team_member_budget_alerts_then_denies_at_the_limit(budget: TeamMemberBudget | None, denied: bool):
+    key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0)
+    outcome, calls = await _check(key, team_member_budget=budget)
+    if denied:
+        assert budget is not None
+        assert outcome == Deny(
+            TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=budget.spend, max_budget=budget.max_budget)
+        )
+        assert calls[-2:] == ("read_team_member_budget", "alert_team_member_budget")
+    else:
+        assert outcome == _allowed()
+        assert "alert_team_member_budget" not in calls
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "key",
+    [
+        replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=None),
+        replace(PLAIN_KEY, user_id=None, team_id="t1", team_member_spend=0.0),
+        replace(PLAIN_KEY, user_id="u1", team_id=None, team_member_spend=0.0),
+    ],
+)
+async def test_team_member_budget_is_read_only_for_a_member_with_tracked_spend(key: KeyFacts):
+    _, calls = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
+    assert "read_team_member_budget" not in calls
+
+
+@pytest.mark.asyncio
+async def test_team_member_budget_needs_a_database():
+    key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0)
+    no_db: Final = Settings(MASTER, False, False, frozenset())
+    _, calls = await _check(key, settings=no_db, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
+    assert "read_team_member_budget" not in calls
+
+
+@pytest.mark.asyncio
+async def test_team_member_budget_is_checked_before_expiry():
+    key: Final = replace(
+        PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0, expires=NOW - timedelta(seconds=1)
+    )
+    outcome, _ = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
+    assert outcome == Deny(TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=1.0, max_budget=0.5))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("spend", "max_budget", "throttle", "expected"),
+    [
+        (5.0, 5.0, None, Deny(KeyBudgetExceeded(spend=5.0, max_budget=5.0))),
+        (6.0, 5.0, None, Deny(KeyBudgetExceeded(spend=6.0, max_budget=5.0))),
+        (4.9, 5.0, None, Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None))),
+        (6.0, 5.0, 0.5, Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=0.5))),
+        (4.9, 5.0, 0.5, Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None))),
+        (
+            1e9,
+            math.inf,
+            None,
+            Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)),
+        ),
+        (
+            1e9,
+            math.nan,
+            None,
+            Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)),
+        ),
+    ],
+)
+async def test_key_budget_denies_or_throttles_at_the_limit(
+    spend: float, max_budget: float, throttle: float | None, expected: Outcome
+):
+    key: Final = replace(PLAIN_KEY, max_budget=max_budget, budget_throttle_pct=throttle)
+    handler: Final = RecordingHandler(key=key, key_spend=spend)
+    outcome, calls = await _run(PythonCore().start_resolved_key(_facts("hash"), SETTINGS, NOW), handler)
+    assert outcome == expected
+    assert calls[-3:] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget")
+    assert handler.calls[-1] == ("alert_key_budget", spend)
+
+
+@pytest.mark.asyncio
+async def test_key_budget_is_enforced_only_on_llm_routes():
+    key: Final = replace(PLAIN_KEY, max_budget=1.0)
+    outcome, calls = await _check(key, is_llm_api_route=False, key_spend=2.0)
+    assert outcome == _allowed()
+    assert calls[-1] == "send_key_budget_alerts"

tests/unit/proxy/auth/test_user_api_key_auth.py

@@ -1600,7 +1600,9 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning():
     tree = ast.parse(
         textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))
         + "\n"
-        + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key))
+        + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth))
+        + "\n"
+        + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects))
     )
 
     def calls_before_each_return(node):
@@ -1645,7 +1647,9 @@ def test_every_jwt_branch_carries_the_user_model_budget():
     tree = ast.parse(
         textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))
         + "\n"
-        + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key))
+        + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth))
+        + "\n"
+        + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects))
     )
 
     assignments = [
@@ -1663,7 +1667,7 @@ def test_every_jwt_branch_carries_the_user_model_budget():
     assert "auto_registered" in targets, (
         f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}"
     )
-    assert "valid_token" in targets, "the virtual-key path must carry the user's model budget"
+    assert "key" in targets, "the virtual-key path must carry the user's model budget"
 
 
 @pytest.mark.asyncio
@@ -1782,7 +1786,9 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path():
     tree = ast.parse(
         textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))
         + "\n"
-        + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key))
+        + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth))
+        + "\n"
+        + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects))
     )
 
     def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool:
@@ -1803,7 +1809,9 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path():
     calls = [
         node
         for node in ast.walk(tree)
-        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget"
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id in {"_check_user_model_budget", "_check_soft_and_model_budgets"}
     ]
     assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}"
 
@@ -1927,7 +1935,9 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach():
     tree = ast.parse(
         textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))
         + "\n"
-        + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key))
+        + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth))
+        + "\n"
+        + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects))
     )
 
     # Half one: the shared block copies the user row's budget onto the token.