perf(proxy): one MGET and one pipeline for post-call spend counters, no team/user/org refetch on the response path (#40841)

* perf(auth): prefetch user, team, membership, org and project in one MGET, one query and one pipeline

Auth read each object with its own Redis GET and, on a miss, its own DB
query, then the admission spend counters with one GET each. The prefetch
warms every entry the checks read with one MGET, one raw query for the
Redis misses and one pipeline write, and a per-request batch serves the
spend counter reads from one MGET. The per-object getters stay the
readers and the fallback, so enforcement does not depend on the prefetch

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(auth): keep prefetch and spend batch collections immutable

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* perf(auth): let the cold spend-counter reseed reuse the admission MGET instead of one GET per counter

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* perf(auth): prefetch referenced auth objects only after the key's model access check passes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(auth): give the prefetch-ordering test's patches their test-quality reasons

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(auth): move the real-Postgres prefetch join test to the proxy_behavior shard

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(auth): read NULL nested permission and budget lists as [] in the prefetch join

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* perf(proxy): batch post-call spend counter reads and carry budget state through the request

Post-call warm checks, reservation reads and reconcile reads for one request now go through a task-local spend counter batch: one MGET answers every counter, successful increments write their result back into the batch so no second Redis read follows, and invalidation forgets the key. RedisCache.async_increment sends INCRBYFLOAT and its TTL command in one pipeline round trip.

Auth pins frozen team, user and org budget snapshots on UserAPIKeyAuth, the pre-call setup writes them into the request metadata, and Prometheus reads them back instead of calling get_key_object, get_team_object, get_user_object and get_org_object on the response path. The getters stay as the fallback for requests that carried nothing (custom auth, unauthenticated routes, skipped checks).

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* perf(proxy): reconcile the budget reservation and the post-call warm checks from one MGET and one pipeline

A scope opened inside an open spend counter batch binds into it instead of starting its own, so the reservation reconcile and the post-call warm checks share the request's single MGET. The reconcile reads every reserved counter concurrently, sends the consistent adjustments in one INCRBYFLOAT+EXPIRE pipeline and settles a flushed or reseeded counter on its own afterwards, keeping the pre-call resize fail-closed. PendingSpendIncrement moves to spend_counter_batch so budget_reservation can build a pipeline without importing a private name

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* chore(proxy): drop the dataclass import left behind by the PendingSpendIncrement move

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(types): import Self from typing_extensions so the proxy imports on Python 3.10

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): use a neutral organization alias in the carried budget state tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): cover recorded and forgotten spend counter values in the request batch

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(caching): assert async_set_cache_pipeline_with_ttls keeps per-entry TTLs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(proxy): type the reservation entry carried through reconcile adjustments

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(auth): map the model table's aliases column to model_aliases in the prefetch join and read user memberships the way get_user_object does

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/caching/redis_cache.py

@@ -1273,6 +1273,24 @@ class RedisCache(BaseCache):
         if len(self.redis_batch_writing_buffer) >= self.redis_flush_size:
             await self.flush_cache_buffer()  # logging done in here
 
+    @staticmethod
+    async def _incrbyfloat_with_ttl(
+        _redis_client: "Redis", key: str, value: float, ttl: int | None, refresh_ttl: bool
+    ) -> float:
+        """INCRBYFLOAT plus its TTL command in one round trip; a third only when an unexpiring key needs an EXPIRE."""
+        if ttl is None:
+            return await _redis_client.incrbyfloat(name=key, amount=value)
+        async with _redis_client.pipeline(transaction=False) as pipe:
+            pipe.incrbyfloat(name=key, amount=value)
+            if refresh_ttl:
+                pipe.expire(key, ttl)
+            else:
+                pipe.ttl(key)
+            result, ttl_or_expire = await pipe.execute()
+        if not refresh_ttl and ttl_or_expire == -1:
+            await _redis_client.expire(key, ttl)
+        return float(result)
+
     @_redis_circuit_breaker_guard
     async def async_increment(
         self,
@@ -1289,14 +1307,9 @@ class RedisCache(BaseCache):
         _used_ttl: Final = self.get_ttl(ttl=ttl)
         key = self.check_and_fix_namespace(key=key)
         try:
-            result: Final = await _redis_client.incrbyfloat(name=key, amount=value)
-            if _used_ttl is not None:
-                if refresh_ttl:
-                    await _redis_client.expire(key, _used_ttl)
-                else:
-                    current_ttl: Final = await _redis_client.ttl(key)
-                    if current_ttl == -1:
-                        await _redis_client.expire(key, _used_ttl)
+            result: Final = await self._incrbyfloat_with_ttl(
+                _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl
+            )
 
             ## LOGGING ##
             end_time = time.time()

litellm/integrations/prometheus.py

@@ -52,6 +52,12 @@ from litellm.types.integrations.prometheus import (
     _sanitize_prometheus_label_value,
     validate_prometheus_deployment_and_latency_caller_identity,
 )
+from litellm.types.proxy.carried_budget_state import (
+    KeyBudgetSnapshot,
+    OrgBudgetSnapshot,
+    TeamBudgetSnapshot,
+    UserBudgetSnapshot,
+)
 from litellm.types.utils import (
     StandardLoggingGuardrailInformation,
     StandardLoggingPayload,
@@ -1941,6 +1947,8 @@ class PrometheusLogger(CustomLogger):
 
         _user_spend: Final = _metadata.get("user_api_key_user_spend", None)
         _user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None)
+        _user_email: Final = _metadata.get("user_api_key_user_email", None)
+        _org_alias: Final = _metadata.get("user_api_key_org_alias", None)
 
         # Bound the per-request budget-metric emission so that slow Redis/DB
         # lookups under load cannot consume the whole LoggingWorker watchdog
@@ -1957,6 +1965,7 @@ class PrometheusLogger(CustomLogger):
                 response_cost=response_cost,
                 key_max_budget=_api_key_max_budget,
                 key_spend=_api_key_spend,
+                carried=KeyBudgetSnapshot.from_metadata(_metadata),
             ),
             self._set_team_budget_metrics_after_api_request(
                 user_api_team=user_api_team,
@@ -1964,16 +1973,21 @@ class PrometheusLogger(CustomLogger):
                 team_spend=_team_spend,
                 team_max_budget=_team_max_budget,
                 response_cost=response_cost,
+                carried=TeamBudgetSnapshot.from_metadata(_metadata),
             ),
             self._set_user_budget_metrics_after_api_request(
                 user_id=user_id,
                 user_spend=_user_spend,
                 user_max_budget=_user_max_budget,
                 response_cost=response_cost,
+                carried=UserBudgetSnapshot.from_metadata(_metadata),
+                user_email=_user_email if isinstance(_user_email, str) else None,
             ),
             self._set_org_budget_metrics_after_api_request(
                 org_id=user_api_key_org_id,
                 response_cost=response_cost,
+                carried=OrgBudgetSnapshot.from_metadata(_metadata),
+                org_alias=_org_alias if isinstance(_org_alias, str) else None,
             ),
             return_exceptions=True,
         )
@@ -3821,6 +3835,7 @@ class PrometheusLogger(CustomLogger):
         team_spend: float | None,
         team_max_budget: float | None,
         response_cost: float,
+        carried: TeamBudgetSnapshot | None = None,
     ):
         """
         Set team budget metrics after an LLM API request
@@ -3839,6 +3854,7 @@ class PrometheusLogger(CustomLogger):
                 spend=team_spend,
                 max_budget=team_max_budget,
                 response_cost=response_cost,
+                carried=carried,
             )
 
             self._set_team_budget_metrics(team_object)
@@ -3850,18 +3866,26 @@ class PrometheusLogger(CustomLogger):
         spend: float | None,
         max_budget: float | None,
         response_cost: float,
+        carried: TeamBudgetSnapshot | None = None,
     ) -> LiteLLM_TeamTable:
         """
         Assemble a LiteLLM_TeamTable object
 
-        for fields not available in metadata, we fetch from db
-        Fields not available in metadata:
-        - `budget_reset_at`
+        ``budget_reset_at`` comes from the auth-carried snapshot when the request has one,
+        otherwise from the team lookup
         """
         from litellm.proxy.auth.auth_checks import get_team_object
         from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
 
         _total_team_spend: Final = (spend or 0) + response_cost
+        if carried is not None:
+            return LiteLLM_TeamTable(
+                team_id=team_id,
+                team_alias=team_alias,
+                spend=_total_team_spend,
+                max_budget=max_budget if max_budget is not None else carried.max_budget,
+                budget_reset_at=carried.budget_reset_at,
+            )
         team_object: Final = LiteLLM_TeamTable(
             team_id=team_id,
             team_alias=team_alias,
@@ -3946,11 +3970,13 @@ class PrometheusLogger(CustomLogger):
         self,
         org_id: str | None,
         response_cost: float,
+        carried: OrgBudgetSnapshot | None = None,
+        org_alias: str | None = None,
     ):
         """
         Set org budget metrics after an LLM API request
 
-        - Fetches org info via cache (get_org_object)
+        - Uses the auth-carried org budget when the request has one, else fetches via get_org_object
         - Sets org budget metrics
         """
         if isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric):
@@ -3959,6 +3985,16 @@ class PrometheusLogger(CustomLogger):
         if not org_id:
             return
 
+        if carried is not None:
+            self._set_org_budget_metrics(
+                org_id=org_id,
+                org_alias=org_alias or "",
+                spend=carried.spend + response_cost,
+                max_budget=carried.max_budget,
+                budget_reset_at=None,
+            )
+            return
+
         from litellm.proxy.auth.auth_checks import get_org_object
         from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
 
@@ -3979,7 +4015,6 @@ class PrometheusLogger(CustomLogger):
         if org_info is None:
             return
 
-        org_alias: Final = org_info.organization_alias or ""
         _total_org_spend: Final = (org_info.spend or 0.0) + response_cost
         budget_table: Final = org_info.litellm_budget_table
         max_budget: Final = budget_table.max_budget if budget_table else None
@@ -3987,7 +4022,7 @@ class PrometheusLogger(CustomLogger):
 
         self._set_org_budget_metrics(
             org_id=org_id,
-            org_alias=org_alias,
+            org_alias=org_info.organization_alias or "",
             spend=_total_org_spend,
             max_budget=max_budget,
             budget_reset_at=budget_reset_at,
@@ -4084,6 +4119,7 @@ class PrometheusLogger(CustomLogger):
         response_cost: float,
         key_max_budget: float | None,
         key_spend: float | None,
+        carried: KeyBudgetSnapshot | None = None,
     ):
         if isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric):
             return
@@ -4095,6 +4131,7 @@ class PrometheusLogger(CustomLogger):
                 key_max_budget=key_max_budget,
                 key_spend=key_spend,
                 response_cost=response_cost,
+                carried=carried,
             )
             self._set_key_budget_metrics(user_api_key_dict)
 
@@ -4105,6 +4142,7 @@ class PrometheusLogger(CustomLogger):
         key_max_budget: float | None,
         key_spend: float | None,
         response_cost: float,
+        carried: KeyBudgetSnapshot | None = None,
     ) -> UserAPIKeyAuth:
         """
         Assemble a UserAPIKeyAuth object
@@ -4113,6 +4151,14 @@ class PrometheusLogger(CustomLogger):
         from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
 
         _total_key_spend: Final = (key_spend or 0) + response_cost
+        if carried is not None:
+            return UserAPIKeyAuth(
+                token=user_api_key,
+                key_alias=user_api_key_alias,
+                max_budget=key_max_budget,
+                spend=_total_key_spend,
+                budget_reset_at=carried.budget_reset_at,
+            )
         user_api_key_dict: Final = UserAPIKeyAuth(
             token=user_api_key,
             key_alias=user_api_key_alias,
@@ -4140,6 +4186,8 @@ class PrometheusLogger(CustomLogger):
         user_spend: float | None,
         user_max_budget: float | None,
         response_cost: float,
+        carried: UserBudgetSnapshot | None = None,
+        user_email: str | None = None,
     ):
         """
         Set user budget metrics after an LLM API request
@@ -4157,6 +4205,8 @@ class PrometheusLogger(CustomLogger):
                 spend=user_spend,
                 max_budget=user_max_budget,
                 response_cost=response_cost,
+                carried=carried,
+                user_email=user_email,
             )
 
             self._set_user_budget_metrics(user_object)
@@ -4167,18 +4217,28 @@ class PrometheusLogger(CustomLogger):
         spend: float | None,
         max_budget: float | None,
         response_cost: float,
+        carried: UserBudgetSnapshot | None = None,
+        user_email: str | None = None,
     ) -> LiteLLM_UserTable:
         """
         Assemble a LiteLLM_UserTable object
 
-        for fields not available in metadata, we fetch from db
-        Fields not available in metadata:
-        - `budget_reset_at`
+        ``budget_reset_at`` and ``user_alias`` come from the auth-carried snapshot when the
+        request has one, otherwise from the user lookup
         """
         from litellm.proxy.auth.auth_checks import get_user_object
         from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
 
         _total_user_spend: Final = (spend or 0) + response_cost
+        if carried is not None:
+            return LiteLLM_UserTable(
+                user_id=user_id,
+                spend=_total_user_spend,
+                max_budget=max_budget if max_budget is not None else carried.max_budget,
+                budget_reset_at=carried.budget_reset_at,
+                user_alias=carried.user_alias,
+                user_email=user_email,
+            )
         user_object: Final = LiteLLM_UserTable(
             user_id=user_id,
             spend=_total_user_spend,

litellm/proxy/_types.py

@@ -40,6 +40,11 @@ from litellm.types.mcp import (
     MCPTransportType,
 )
 from litellm.types.mcp_server.mcp_server_manager import MCPInfo
+from litellm.types.proxy.carried_budget_state import (
+    OrgBudgetSnapshot,
+    TeamBudgetSnapshot,
+    UserBudgetSnapshot,
+)
 from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
 from litellm.types.router import RouterErrors, UpdateRouterConfig
 from litellm.types.secret_managers.main import KeyManagementSystem
@@ -3106,6 +3111,9 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView):  # the expected response ob
         ),
     )
     budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True)
+    team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
+    user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
+    org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True)
     matched_model_access_groups: list[str] | None = Field(default=None, exclude=True)
     budget_throttle_pct: float | None = Field(default=None, exclude=True)
     user: Any | None = None  # Expanded user object when expand=user is used

litellm/proxy/auth/auth_checks.py

@@ -102,6 +102,7 @@ from litellm.proxy.guardrails.tool_name_extraction import (
 )
 from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
+from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
 from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
 from litellm.repositories.budget_repository import BudgetRepository
 from litellm.repositories.object_permission_repository import ObjectPermissionRepository
@@ -5635,6 +5636,8 @@ async def _organization_max_budget_check(
     if org_table is None:
         return
 
+    carry_organization_budget_state(valid_token=valid_token, org_table=org_table)
+
     # Get max_budget from organization's budget table
     org_max_budget: float | None = None
     if org_table.litellm_budget_table is not None:

litellm/proxy/auth/user_api_key_auth.py

@@ -103,6 +103,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
 )
 from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state
 from litellm.proxy.spend_tracking.spend_counter_batch import (
     bind_admission_counter_keys,
     release_spend_counter_batch,
@@ -2631,6 +2632,11 @@ async def _run_centralized_common_checks(
         None if isinstance(end_user_result, BaseException) else end_user_result
     )
     global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result
+    carry_team_and_user_budget_state(
+        valid_token=user_api_key_auth_obj,
+        team_object=team_object,
+        user_object=user_object,
+    )
 
     if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
         user_api_key_auth_obj.org_id = team_object.organization_id

litellm/proxy/db/spend_counter_reseed.py

@@ -24,7 +24,7 @@ from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import Litellm_EntityType
 from litellm.proxy.db.db_lookup_gate import db_lookup_gate
-from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch
+from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value
 from litellm.repositories.organization_repository import OrganizationRepository
 from litellm.repositories.table_repositories import (
     BudgetWindowSpendRepository,
@@ -197,9 +197,8 @@ class SpendCounterReseed:
 
     @staticmethod
     async def _read_active_batch(counter_key: str) -> tuple[float | None, bool] | None:
-        """The request's admission MGET already answered for this counter; a Redis miss there is authoritative."""
-        batch: Final = active_spend_counter_batch()
-        return None if batch is None else await batch.read(counter_key)
+        """The request's MGET answers for this counter; a Redis miss there is authoritative."""
+        return await read_batched_spend_counter(counter_key)
 
     @staticmethod
     async def coalesced(
@@ -263,6 +262,7 @@ class SpendCounterReseed:
                         key=counter_key,
                         value=current_value,
                     )
+                    record_spend_counter_value(counter_key, current_value)
                 else:
                     cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
                     seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend
@@ -414,8 +414,11 @@ class SpendCounterReseed:
     ) -> float | None:
         lock: Final = await SpendCounterReseed._get_lock(counter_key)
         async with lock:
-            redis_clean_miss = False
-            if spend_counter_cache.redis_cache is not None:
+            batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
+            if batched is not None and batched[0] is not None:
+                return batched[0]
+            redis_clean_miss = batched is not None
+            if spend_counter_cache.redis_cache is not None and not redis_clean_miss:
                 try:
                     val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
                     if val is not None:
@@ -459,6 +462,7 @@ class SpendCounterReseed:
                         key=counter_key,
                         value=current_value,
                     )
+                    record_spend_counter_value(counter_key, float(current_value))
                 else:
                     cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
                     seeded_spend: Final = (

litellm/proxy/litellm_pre_call_utils.py

@@ -64,6 +64,7 @@ from litellm.proxy.common_utils.callback_utils import (
     strip_callback_config,
 )
 from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
+from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
 from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
 
 # Cache special headers as a frozenset for O(1) lookup performance
@@ -2300,6 +2301,7 @@ async def add_litellm_data_to_request(
     data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget
     user_model_budget: Final = user_api_key_dict.user_model_max_budget
     data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget  # rebind-ok: out-param
+    data[_metadata_variable_name].update(carried_budget_metadata(user_api_key_dict))
 
     data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
     data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata)

litellm/proxy/proxy_server.py

@@ -26,7 +26,6 @@ from collections.abc import (
     MutableMapping,
     Sequence,
 )
-from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from types import MappingProxyType, UnionType
 from typing import (
@@ -659,7 +658,15 @@ from litellm.proxy.route_priority import hot_routes_first
 from litellm.proxy.search_endpoints.endpoints import router as search_router
 from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
 from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
-from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch
+from litellm.proxy.spend_tracking.spend_counter_batch import (
+    PendingSpendIncrement,
+    active_spend_counter_batch,
+    forget_spend_counter,
+    post_call_counter_keys,
+    read_batched_spend_counter,
+    record_spend_counter_value,
+    spend_counter_batch_scope,
+)
 from litellm.proxy.spend_tracking.spend_management_endpoints import (
     router as spend_management_router,
 )
@@ -2628,6 +2635,7 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
     if needs_update:
         spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
     if spend_counter_cache.redis_cache is not None:
+        forget_spend_counter(counter_key)
         try:
             await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend)
         except Exception:
@@ -2768,12 +2776,6 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float)
     return fallback_spend, False
 
 
-@dataclass(frozen=True, slots=True)
-class _PendingSpendIncrement:
-    counter_key: str
-    increment: float
-
-
 async def increment_spend_counters(
     token: str | None,
     team_id: str | None,
@@ -2796,6 +2798,45 @@ async def increment_spend_counters(
     Awaited (not create_task) in the cost callback, so the counter is
     updated before the next request's auth check runs.
     """
+    with spend_counter_batch_scope(
+        spend_counter_cache.redis_cache,
+        counter_keys=post_call_counter_keys(
+            token=token,
+            team_id=team_id,
+            user_id=user_id,
+            org_id=org_id,
+            end_user_id=end_user_id,
+            tags=tags,
+            model_access_groups=model_access_groups,
+        ),
+    ):
+        await _increment_spend_counters_batched(
+            token=token,
+            team_id=team_id,
+            user_id=user_id,
+            response_cost=response_cost,
+            org_id=org_id,
+            budget_reservation=budget_reservation,
+            end_user_id=end_user_id,
+            tags=tags,
+            request_started_at=request_started_at,
+            model_access_groups=model_access_groups,
+        )
+
+
+async def _increment_spend_counters_batched(
+    token: str | None,
+    team_id: str | None,
+    user_id: str | None,
+    response_cost: float | None,
+    org_id: str | None,
+    budget_reservation: dict | None,
+    end_user_id: str | None,
+    tags: list[str] | None,
+    request_started_at: datetime | None,
+    model_access_groups: Sequence[str] | None,
+):
+    """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET."""
     reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update(
         budget_reservation=budget_reservation,
         response_cost=response_cost,
@@ -2808,7 +2849,7 @@ async def increment_spend_counters(
 
     cost: Final[float] = response_cost
 
-    async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
+    async def _key_scope(key_token: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
         # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow
         # hashes raw "sk-..." keys before they reach the callback). The
         # startswith("sk-") check is a safety net matching update_cache —
@@ -2819,7 +2860,7 @@ async def increment_spend_counters(
             hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token
         )
         key_counter_key: Final = f"spend:key:{hashed_token}"
-        key_pending: Final[tuple[_PendingSpendIncrement, ...]] = (
+        key_pending: Final[tuple[PendingSpendIncrement, ...]] = (
             ()
             if key_counter_key in reserved_counter_keys
             else (
@@ -2831,7 +2872,7 @@ async def increment_spend_counters(
             )
         )
 
-        async def _key_window_increment(window: object) -> _PendingSpendIncrement | None:
+        async def _key_window_increment(window: object) -> PendingSpendIncrement | None:
             duration = (
                 window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
             )
@@ -2878,9 +2919,9 @@ async def increment_spend_counters(
         )
         return key_pending + tuple(item for item in window_pending if item is not None)
 
-    async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
+    async def _team_scope(scope_team_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
         team_counter_key: Final = f"spend:team:{scope_team_id}"
-        team_pending: Final[tuple[_PendingSpendIncrement, ...]] = (
+        team_pending: Final[tuple[PendingSpendIncrement, ...]] = (
             ()
             if team_counter_key in reserved_counter_keys
             else (
@@ -2892,7 +2933,7 @@ async def increment_spend_counters(
             )
         )
 
-        async def _team_window_increment(window: object) -> _PendingSpendIncrement | None:
+        async def _team_window_increment(window: object) -> PendingSpendIncrement | None:
             duration = (
                 window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
             )
@@ -2941,7 +2982,7 @@ async def increment_spend_counters(
 
     async def _team_member_scope(
         scope_user_id: str, scope_team_id: str
-    ) -> tuple[_PendingSpendIncrement | BaseException, ...]:
+    ) -> tuple[PendingSpendIncrement | BaseException, ...]:
         team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}"
         if team_member_counter_key in reserved_counter_keys:
             return ()
@@ -2953,7 +2994,7 @@ async def increment_spend_counters(
             ),
         )
 
-    async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
+    async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
         user_counter_key: Final = f"spend:user:{scope_user_id}"
         if user_counter_key in reserved_counter_keys:
             return ()
@@ -3063,7 +3104,7 @@ async def _prepare_end_user_and_tag_spend_increments(
     tags: list[str] | None,
     response_cost: float,
     reserved_counter_keys: set[str],
-) -> tuple[_PendingSpendIncrement | BaseException, ...]:
+) -> tuple[PendingSpendIncrement | BaseException, ...]:
     unique_tags: Final = (
         tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else ()
     )
@@ -3100,7 +3141,7 @@ async def _prepare_model_access_group_spend_increments(
     model_access_groups: Sequence[object],
     response_cost: float,
     reserved_counter_keys: set[str],
-) -> tuple[_PendingSpendIncrement | BaseException, ...]:
+) -> tuple[PendingSpendIncrement | BaseException, ...]:
     """Charge the model access groups that authorized this request.
 
     Without this the counter auth reads is written only by the reservation path, so
@@ -3133,7 +3174,7 @@ async def _prepare_org_spend_increment(
     org_id: str | None,
     response_cost: float,
     reserved_counter_keys: set[str],
-) -> tuple[_PendingSpendIncrement, ...]:
+) -> tuple[PendingSpendIncrement, ...]:
     if org_id is None:
         return ()
 
@@ -3151,7 +3192,7 @@ async def _prepare_unreserved_spend_counter_increment(
     source_cache_key: str | list[str],
     increment: float,
     reserved_counter_keys: set[str],
-) -> _PendingSpendIncrement | None:
+) -> PendingSpendIncrement | None:
     if counter_key in reserved_counter_keys:
         return None
 
@@ -3166,7 +3207,7 @@ async def _prepare_spend_counter_increment(
     counter_key: str,
     source_cache_key: str | list[str],
     increment: float,
-) -> _PendingSpendIncrement:
+) -> PendingSpendIncrement:
     """
     Initialize counter from the authoritative DB spend value if not yet
     set, then return the pending increment for the caller to apply in one
@@ -3188,7 +3229,7 @@ async def _prepare_spend_counter_increment(
         counter_key=counter_key,
         source_cache_key=source_cache_key,
     )
-    return _PendingSpendIncrement(counter_key=counter_key, increment=increment)
+    return PendingSpendIncrement(counter_key=counter_key, increment=increment)
 
 
 async def _enqueue_window_spend_row_update(
@@ -3247,7 +3288,7 @@ async def _prepare_window_spend_counter_increment(
     window_duration: str | None,
     window_start: datetime | None,
     increment: float,
-) -> _PendingSpendIncrement | None:
+) -> PendingSpendIncrement | None:
     if window_start is None:
         verbose_proxy_logger.warning(
             "Skipping spend counter increment for invalid budget window %s",
@@ -3264,7 +3305,7 @@ async def _prepare_window_spend_counter_increment(
     )
     if initialized is False:
         return None
-    return _PendingSpendIncrement(counter_key=counter_key, increment=increment)
+    return PendingSpendIncrement(counter_key=counter_key, increment=increment)
 
 
 async def _ensure_spend_counter_initialized(
@@ -3331,6 +3372,14 @@ async def _ensure_window_spend_counter_initialized(
 
 
 async def _is_spend_counter_cache_warm(counter_key: str) -> bool:
+    batched: Final = await read_batched_spend_counter(counter_key)
+    if batched is not None:
+        batched_value, _ = batched
+        if batched_value is None:
+            return False
+        spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value)
+        return True
+
     if spend_counter_cache.redis_cache is not None:
         try:
             current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache(
@@ -3375,6 +3424,7 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float):
             key=counter_key,
             value=current_value,
         )
+        record_spend_counter_value(counter_key, float(current_value))
         return current_value
 
     return await SpendCounterReseed.increment_in_memory(
@@ -3383,6 +3433,7 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float):
 
 
 async def _invalidate_spend_counter(counter_key: str):
+    forget_spend_counter(counter_key)
     spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
     if spend_counter_cache.redis_cache is not None:
         try:
@@ -3395,7 +3446,16 @@ async def _invalidate_spend_counter(counter_key: str):
             )
 
 
-async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None:
+async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None:
+    try:
+        await increment_spend_counters_pipeline(pending=pending)
+    except RedisCircuitBreakerOpenError:
+        return
+
+
+async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> None:
+    """One INCRBYFLOAT+EXPIRE pipeline for every pending counter; on failure every counter is invalidated
+    before the error propagates, so no caller can read a half-applied batch."""
     if not pending:
         return
     redis_cache: Final = spend_counter_cache.redis_cache
@@ -3412,13 +3472,12 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme
     ]
     try:
         results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list)
-    except Exception as e:
+    except Exception:
         await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending))
-        if isinstance(e, RedisCircuitBreakerOpenError):
-            return
         raise
     for item, current_value in zip(pending, results or ()):
         spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
+        record_spend_counter_value(item.counter_key, float(current_value))
 
 
 async def update_cache(

litellm/proxy/spend_tracking/budget_reservation.py

@@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
     tag_cache_key,
     team_membership_reservation_cache_key,
 )
+from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope
 from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm.router import Router
 from litellm.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer
@@ -257,46 +258,47 @@ async def reserve_budget_for_request(
 
     applied_entries: Final[list[dict[str, float | str]]] = []
     try:
-        for counter in counters:
-            entry = _counter_to_reservation_entry(
-                counter=counter,
-                reserved_cost=reservation_cost,
-            )
-            applied_entries.append(entry)
-            try:
-                reserved_value = await _reserve_counter(
+        with _counters_batch_scope(frozenset(counter.counter_key for counter in counters)):
+            for counter in counters:
+                entry = _counter_to_reservation_entry(
                     counter=counter,
-                    reservation_cost=reservation_cost,
+                    reserved_cost=reservation_cost,
                 )
-            except _CounterReservationUnavailable as exc:
-                if exc.touched_counter and not exc.counter_invalidated:
-                    await _release_applied_entries_best_effort(
-                        entries=[entry],
-                        default_reserved_cost=reservation_cost,
+                applied_entries.append(entry)
+                try:
+                    reserved_value = await _reserve_counter(
+                        counter=counter,
+                        reservation_cost=reservation_cost,
                     )
-                applied_entries.remove(entry)
-                if fail_closed_budget_enforcement:
-                    _raise_reservation_unavailable(counter_key=counter.counter_key)
-                continue
+                except _CounterReservationUnavailable as exc:
+                    if exc.touched_counter and not exc.counter_invalidated:
+                        await _release_applied_entries_best_effort(
+                            entries=[entry],
+                            default_reserved_cost=reservation_cost,
+                        )
+                    applied_entries.remove(entry)
+                    if fail_closed_budget_enforcement:
+                        _raise_reservation_unavailable(counter_key=counter.counter_key)
+                    continue
 
-            if reserved_value is not None:
-                current_spend = reserved_value
-            else:
-                cached_spend = current_spend_by_counter_key.get(counter.counter_key)
-                if cached_spend is None:
-                    cached_spend = await _get_current_counter_value(counter=counter)
-                current_spend = cached_spend + reservation_cost
-            if current_spend > counter.max_budget:
-                reservation_cost = await _apply_over_budget_reservation_policy(
-                    counter=counter,
-                    valid_token=valid_token,
-                    entry=entry,
-                    applied_entries=applied_entries,
-                    reservation_cost=reservation_cost,
-                    current_spend=current_spend,
-                    fail_closed_budget_enforcement=fail_closed_budget_enforcement,
-                )
-                continue
+                if reserved_value is not None:
+                    current_spend = reserved_value
+                else:
+                    cached_spend = current_spend_by_counter_key.get(counter.counter_key)
+                    if cached_spend is None:
+                        cached_spend = await _get_current_counter_value(counter=counter)
+                    current_spend = cached_spend + reservation_cost
+                if current_spend > counter.max_budget:
+                    reservation_cost = await _apply_over_budget_reservation_policy(
+                        counter=counter,
+                        valid_token=valid_token,
+                        entry=entry,
+                        applied_entries=applied_entries,
+                        reservation_cost=reservation_cost,
+                        current_spend=current_spend,
+                        fail_closed_budget_enforcement=fail_closed_budget_enforcement,
+                    )
+                    continue
     except Exception:
         await _release_applied_entries_best_effort(
             entries=applied_entries,
@@ -878,67 +880,92 @@ async def _get_current_counter_value(counter: _BudgetCounter) -> float:
     )
 
 
+def _counters_batch_scope(counter_keys: frozenset[str]) -> spend_counter_batch_scope:
+    """Each counter is read once, then written, so one MGET up front serves every read in the loop."""
+    from litellm.proxy.proxy_server import spend_counter_cache
+
+    return spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys)
+
+
+@dataclass(frozen=True, slots=True)
+class _EntryAdjustment:
+    entry: dict[str, float | str]
+    counter_key: str
+    target_adjustment: float
+    adjustment: float
+
+
+def _entry_adjustment(
+    entry: dict[str, float | str], actual_cost: float, default_reserved_cost: float
+) -> _EntryAdjustment | None:
+    counter_key: Final = entry.get("counter_key")
+    if counter_key is None:
+        return None
+    target_adjustment: Final = actual_cost - _get_entry_reserved_cost(
+        entry=entry, default_reserved_cost=default_reserved_cost
+    )
+    adjustment: Final = target_adjustment - float(entry.get("applied_adjustment") or 0.0)
+    if adjustment == 0:
+        return None
+    return _EntryAdjustment(
+        entry=entry, counter_key=str(counter_key), target_adjustment=target_adjustment, adjustment=adjustment
+    )
+
+
 async def _set_reserved_entries_actual_cost(
     entries: list[dict],
     actual_cost: float,
     default_reserved_cost: float,
     reseed_on_inconsistent: bool = True,
 ) -> None:
-    for entry in entries:
-        await _set_reserved_entry_actual_cost(
-            entry=entry,
-            actual_cost=actual_cost,
-            default_reserved_cost=default_reserved_cost,
-            reseed_on_inconsistent=reseed_on_inconsistent,
+    """Every reserved counter is read from one MGET and the consistent adjustments go out in one pipeline.
+    A counter that was flushed or reseeded since reservation is settled on its own after the pipeline."""
+    from litellm.proxy.proxy_server import increment_spend_counters_pipeline
+
+    with _counters_batch_scope(frozenset(str(entry["counter_key"]) for entry in entries if "counter_key" in entry)):
+        adjustments: Final = tuple(
+            adjustment
+            for entry in entries
+            if (adjustment := _entry_adjustment(entry, actual_cost, default_reserved_cost)) is not None
         )
-
-
-async def _set_reserved_entry_actual_cost(
-    entry: dict,
-    actual_cost: float,
-    default_reserved_cost: float,
-    reseed_on_inconsistent: bool = True,
-) -> None:
-    from litellm.proxy.proxy_server import (
-        _increment_spend_counter_cache,
-        reseed_spend_counter_from_db,
-    )
-
-    counter_key: Final = entry.get("counter_key")
-    if counter_key is None:
-        return
-    reserved_cost: Final = _get_entry_reserved_cost(
-        entry=entry,
-        default_reserved_cost=default_reserved_cost,
-    )
-    target_adjustment: Final = actual_cost - reserved_cost
-    applied_adjustment: Final = float(entry.get("applied_adjustment") or 0.0)
-    adjustment: Final = target_adjustment - applied_adjustment
-    if adjustment == 0:
-        return
-    if await _counter_can_apply_adjustment(
-        counter_key=counter_key,
-        adjustment=adjustment,
-    ):
-        await _increment_spend_counter_cache(
-            counter_key=counter_key,
-            increment=adjustment,
+        consistent: Final = tuple(
+            await asyncio.gather(
+                *(
+                    _counter_can_apply_adjustment(counter_key=item.counter_key, adjustment=item.adjustment)
+                    for item in adjustments
+                )
+            )
         )
-    elif reseed_on_inconsistent:
-        # Post-call reconcile / release: the counter was flushed, expired or reseeded
-        # between reservation and reconcile, so the optimistic delta no longer applies.
-        # Reseed from the DB floor (which cannot include this request's cost yet) and
-        # add the settled cost, since increment_spend_counters skips reserved keys.
-        reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key)
-        if reseeded and actual_cost > 0:
-            await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost)
-    else:
-        # Pre-call admission resize: the in-flight reservation cost is not yet
-        # persisted, so the DB floor would discard it. Keep the original
-        # fail-closed behavior (raise -> reserve_budget_for_request releases and
-        # denies) rather than admitting against an inconsistent counter.
-        raise RuntimeError(f"Cannot resize budget reservation against inconsistent counter {counter_key}")
-    entry["applied_adjustment"] = target_adjustment
+        inconsistent: Final = tuple(item for item, ok in zip(adjustments, consistent) if not ok)
+        if inconsistent and not reseed_on_inconsistent:
+            # Pre-call admission resize: the in-flight reservation cost is not yet
+            # persisted, so the DB floor would discard it. Keep the original
+            # fail-closed behavior (raise -> reserve_budget_for_request releases and
+            # denies) rather than admitting against an inconsistent counter.
+            raise RuntimeError(
+                f"Cannot resize budget reservation against inconsistent counter {inconsistent[0].counter_key}"
+            )
+        applicable: Final = tuple(item for item, ok in zip(adjustments, consistent) if ok)
+        await increment_spend_counters_pipeline(
+            pending=tuple(
+                PendingSpendIncrement(counter_key=item.counter_key, increment=item.adjustment) for item in applicable
+            )
+        )
+        for item in inconsistent:
+            await _reseed_reserved_entry(item=item, actual_cost=actual_cost)
+        for item in adjustments:
+            item.entry["applied_adjustment"] = item.target_adjustment
+
+
+async def _reseed_reserved_entry(item: _EntryAdjustment, actual_cost: float) -> None:
+    """Post-call reconcile / release of a counter that was flushed, expired or reseeded between reservation and
+    reconcile: the optimistic delta no longer applies, so reseed from the DB floor (which cannot include this
+    request's cost yet) and add the settled cost, since increment_spend_counters skips reserved keys."""
+    from litellm.proxy.proxy_server import _increment_spend_counter_cache, reseed_spend_counter_from_db
+
+    reseeded: Final = await reseed_spend_counter_from_db(counter_key=item.counter_key)
+    if reseeded and actual_cost > 0:
+        await _increment_spend_counter_cache(counter_key=item.counter_key, increment=actual_cost)
 
 
 async def _counter_can_apply_adjustment(
@@ -963,8 +990,8 @@ async def _release_applied_entries_best_effort(
 ) -> None:
     for entry in entries:
         try:
-            await _set_reserved_entry_actual_cost(
-                entry=entry,
+            await _set_reserved_entries_actual_cost(
+                entries=[entry],  # mutable-ok: the reconcile takes the reservation's list of entries
                 actual_cost=0.0,
                 default_reserved_cost=default_reserved_cost,
             )

litellm/proxy/spend_tracking/carried_budget_state.py

@@ -0,0 +1,60 @@
+"""Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata."""
+
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import Final
+
+from litellm.models.organization import LiteLLM_OrganizationTable
+from litellm.models.team import LiteLLM_TeamTable
+from litellm.models.user import LiteLLM_UserTable
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.proxy.carried_budget_state import (
+    OrgBudgetSnapshot,
+    TeamBudgetSnapshot,
+    UserBudgetSnapshot,
+)
+
+
+def carry_team_and_user_budget_state(
+    valid_token: UserAPIKeyAuth,
+    team_object: LiteLLM_TeamTable | None,
+    user_object: LiteLLM_UserTable | None,
+) -> None:
+    if team_object is not None:
+        valid_token.team_budget_snapshot = TeamBudgetSnapshot(  # rebind-ok: the request credential is pinned in place
+            budget_reset_at=team_object.budget_reset_at,
+            max_budget=team_object.max_budget,
+        )
+    if user_object is not None:
+        valid_token.user_budget_snapshot = UserBudgetSnapshot(  # rebind-ok: same object the caller keeps using
+            budget_reset_at=user_object.budget_reset_at,
+            max_budget=user_object.max_budget,
+            user_alias=user_object.user_alias,
+        )
+
+
+def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None:
+    budget_table: Final = org_table.litellm_budget_table
+    valid_token.organization_alias = (
+        org_table.organization_alias
+    )  # rebind-ok: the request credential is pinned in place
+    valid_token.org_budget_snapshot = OrgBudgetSnapshot(  # rebind-ok: same object the caller keeps using
+        spend=org_table.spend,
+        max_budget=budget_table.max_budget if budget_table is not None else None,
+    )
+
+
+def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object]:
+    snapshots: Final = (
+        valid_token.team_budget_snapshot,
+        valid_token.user_budget_snapshot,
+        valid_token.org_budget_snapshot,
+    )
+    return MappingProxyType(
+        {
+            key: value
+            for snapshot in snapshots
+            if snapshot is not None
+            for key, value in snapshot.metadata_entries().items()
+        }
+    )

litellm/proxy/spend_tracking/spend_counter_batch.py

@@ -1,8 +1,9 @@
-"""One Redis MGET for every spend counter the admission checks read, instead of one GET per counter."""
+"""One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each."""
 
 import asyncio
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator, Mapping, Sequence
 from contextvars import ContextVar, Token
+from dataclasses import dataclass
 from types import MappingProxyType, TracebackType
 from typing import Final
 
@@ -11,11 +12,18 @@ from pydantic import TypeAdapter
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.redis_cache import RedisCache
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key
 
 _CounterValues: Final = TypeAdapter(dict[str, float | None])
 _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
 
 
+@dataclass(frozen=True, slots=True)
+class PendingSpendIncrement:
+    counter_key: str
+    increment: float
+
+
 class SpendCounterBatch:
     """Bound counters are read with one MGET on first use; counters bound later join the next MGET.
     ``async_batch_get_cache`` maps a clean miss to ``None`` and drops keys only when Redis failed, so an absent
@@ -35,6 +43,10 @@ class SpendCounterBatch:
     def counter_keys(self) -> frozenset[str]:
         return self._keys
 
+    @property
+    def is_open(self) -> bool:
+        return self._open
+
     def bind(self, counter_keys: frozenset[str]) -> None:
         if self._open:
             self._keys = self._keys | counter_keys
@@ -52,12 +64,29 @@ class SpendCounterBatch:
             return None
         return loaded[counter_key], True
 
+    def record(self, counter_key: str, value: float) -> None:
+        """A write returned the counter's new value; later reads in this scope see it instead of the MGET value."""
+        if not self._open:
+            return
+        key: Final = frozenset((counter_key,))
+        self._keys = self._keys | key
+        self._fetched = self._fetched | key
+        self._loaded = MappingProxyType({**self._loaded, counter_key: value})
+
+    def forget(self, counter_key: str) -> None:
+        """A write left the counter's value unknown; later reads in this scope go to Redis."""
+        key: Final = frozenset((counter_key,))
+        self._keys = self._keys - key
+        self._fetched = self._fetched - key
+        self._loaded = MappingProxyType({k: v for k, v in self._loaded.items() if k != counter_key})
+
     async def _load(self) -> Mapping[str, float | None]:
         async with self._lock:
             pending: Final = self._keys - self._fetched
             if pending:
                 self._fetched = self._fetched | pending
-                self._loaded = MappingProxyType({**self._loaded, **await self._fetch(pending)})
+                fetched: Final = await self._fetch(pending)
+                self._loaded = MappingProxyType({**fetched, **self._loaded})
             return self._loaded
 
     async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]:
@@ -78,17 +107,26 @@ def active_spend_counter_batch() -> SpendCounterBatch | None:
 
 
 class spend_counter_batch_scope:
-    """Reads inside the scope share one MGET once ``bind_admission_counter_keys`` has run."""
+    """Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it.
+    Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET."""
 
-    __slots__ = ("_redis_cache", "_token")
+    __slots__ = ("_counter_keys", "_redis_cache", "_token")
 
-    def __init__(self, redis_cache: RedisCache | None) -> None:
+    def __init__(self, redis_cache: RedisCache | None, counter_keys: frozenset[str] = frozenset()) -> None:
         self._redis_cache: Final = redis_cache
+        self._counter_keys: Final = counter_keys
         self._token: Token[SpendCounterBatch | None] | None = None
 
     def __enter__(self) -> None:
-        if self._redis_cache is not None:
-            self._token = _active_batch.set(SpendCounterBatch(self._redis_cache))
+        if self._redis_cache is None:
+            return
+        outer: Final = _active_batch.get()
+        if outer is not None and outer.is_open:
+            outer.bind(self._counter_keys)
+            return
+        batch: Final = SpendCounterBatch(self._redis_cache)
+        batch.bind(self._counter_keys)
+        self._token = _active_batch.set(batch)
 
     def __exit__(
         self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
@@ -122,9 +160,58 @@ def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> fr
     return frozenset(_iter_admission_counter_keys(token, end_user_id))
 
 
+def post_call_counter_keys(
+    token: str | None,
+    team_id: str | None,
+    user_id: str | None,
+    org_id: str | None,
+    end_user_id: str | None,
+    tags: Sequence[object] | None,
+    model_access_groups: Sequence[object] | None,
+) -> frozenset[str]:
+    """Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read."""
+    entity_keys: Final = admission_counter_keys(
+        UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id
+    )
+    tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str))
+    group_keys: Final = frozenset(
+        model_access_group_spend_counter_key(group)
+        for group in model_access_groups or ()
+        if group and isinstance(group, str)
+    )
+    return entity_keys | tag_keys | group_keys
+
+
 def bind_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> None:
     """Idempotent: call again after the token gains ids (end user, team org) so those counters join the MGET."""
+    bind_spend_counter_keys(admission_counter_keys(token, end_user_id))
+
+
+def bind_spend_counter_keys(counter_keys: frozenset[str]) -> None:
     batch: Final = _active_batch.get()
     if batch is None:
         return
-    batch.bind(admission_counter_keys(token, end_user_id))
+    batch.bind(counter_keys)
+
+
+def record_spend_counter_value(counter_key: str, value: float) -> None:
+    batch: Final = _active_batch.get()
+    if batch is None:
+        return
+    batch.record(counter_key, value)
+
+
+def forget_spend_counter(counter_key: str) -> None:
+    batch: Final = _active_batch.get()
+    if batch is None:
+        return
+    batch.forget(counter_key)
+
+
+async def read_batched_spend_counter(counter_key: str) -> tuple[float | None, bool] | None:
+    """Bind-on-read for counters only known at read time (budget windows); the first reader pays the MGET."""
+    batch: Final = _active_batch.get()
+    if batch is None:
+        return None
+    batch.bind(frozenset((counter_key,)))
+    return await batch.read(counter_key)

litellm/types/proxy/carried_budget_state.py

@@ -0,0 +1,64 @@
+"""Budget fields auth already resolved, carried on the request so success logging does no object lookups.
+
+Auth pins one snapshot per entity on ``UserAPIKeyAuth`` (request-scoped, never cached), the pre-call
+setup writes them into the request metadata under the aliased ``user_api_key_*`` names, and a logger
+reads them back with ``from_metadata``. ``None`` means this request never carried that entity
+(unauthenticated route, custom auth, budget check skipped) and the logger keeps its own lookup.
+"""
+
+from collections.abc import Mapping
+from datetime import datetime
+from types import MappingProxyType
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from typing_extensions import Self
+
+
+class _BudgetSnapshot(BaseModel):
+    model_config = ConfigDict(frozen=True, populate_by_name=True, extra="ignore")
+
+    def metadata_entries(self) -> Mapping[str, object]:
+        return MappingProxyType(self.model_dump(by_alias=True, mode="json"))
+
+    @classmethod
+    def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None:
+        try:
+            return cls.model_validate(metadata)
+        except ValidationError:
+            return None
+
+
+class KeyBudgetSnapshot(_BudgetSnapshot):
+    """Read-only view of the ``user_api_key_budget_reset_at`` entry ``add_user_api_key_auth_to_request_metadata`` writes."""
+
+    budget_reset_at: datetime | None = Field(
+        validation_alias="user_api_key_budget_reset_at", serialization_alias="user_api_key_budget_reset_at"
+    )
+
+
+class TeamBudgetSnapshot(_BudgetSnapshot):
+    budget_reset_at: datetime | None = Field(
+        validation_alias="user_api_key_team_budget_reset_at", serialization_alias="user_api_key_team_budget_reset_at"
+    )
+    max_budget: float | None = Field(
+        validation_alias="user_api_key_team_table_max_budget", serialization_alias="user_api_key_team_table_max_budget"
+    )
+
+
+class UserBudgetSnapshot(_BudgetSnapshot):
+    budget_reset_at: datetime | None = Field(
+        validation_alias="user_api_key_user_budget_reset_at", serialization_alias="user_api_key_user_budget_reset_at"
+    )
+    max_budget: float | None = Field(
+        validation_alias="user_api_key_user_table_max_budget", serialization_alias="user_api_key_user_table_max_budget"
+    )
+    user_alias: str | None = Field(
+        validation_alias="user_api_key_user_alias", serialization_alias="user_api_key_user_alias"
+    )
+
+
+class OrgBudgetSnapshot(_BudgetSnapshot):
+    spend: float = Field(validation_alias="user_api_key_org_spend", serialization_alias="user_api_key_org_spend")
+    max_budget: float | None = Field(
+        validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget"
+    )

tests/test_litellm/caching/test_redis_cache.py

@@ -1205,6 +1205,92 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new
     assert breaker._state == breaker.CLOSED
 
 
+class _RoundTripCountingRedis:
+    """Fake redis.asyncio client: one round trip per awaited command or pipeline execute."""
+
+    def __init__(self, ttl: int) -> None:
+        self.values: dict[str, float] = {}
+        self.ttls: dict[str, int] = {}
+        self.round_trips = 0
+        self._initial_ttl = ttl
+
+    async def incrbyfloat(self, name: str, amount: float) -> float:
+        self.round_trips += 1
+        return self._incr(name, amount)
+
+    async def expire(self, name: str, time: int) -> bool:
+        self.round_trips += 1
+        self.ttls[name] = time
+        return True
+
+    def _incr(self, name: str, amount: float) -> float:
+        self.values[name] = self.values.get(name, 0.0) + amount
+        self.ttls.setdefault(name, self._initial_ttl)
+        return self.values[name]
+
+    def pipeline(self, transaction: bool) -> "_RoundTripCountingRedis._Pipeline":
+        return _RoundTripCountingRedis._Pipeline(self)
+
+    class _Pipeline:
+        def __init__(self, client: "_RoundTripCountingRedis") -> None:
+            self._client = client
+            self._commands: list[tuple[str, tuple[object, ...]]] = []
+
+        async def __aenter__(self) -> "_RoundTripCountingRedis._Pipeline":
+            return self
+
+        async def __aexit__(self, *exc_info: object) -> None:
+            return None
+
+        def incrbyfloat(self, name: str, amount: float) -> None:
+            self._commands.append(("incrbyfloat", (name, amount)))
+
+        def expire(self, name: str, time: int) -> None:
+            self._commands.append(("expire", (name, time)))
+
+        def ttl(self, name: str) -> None:
+            self._commands.append(("ttl", (name,)))
+
+        async def execute(self) -> list[object]:
+            self._client.round_trips += 1
+            results: list[object] = []
+            for command, args in self._commands:
+                if command == "incrbyfloat":
+                    results.append(self._client._incr(str(args[0]), float(args[1])))  # pyright: ignore[reportArgumentType]  # fake stores str/float
+                elif command == "expire":
+                    self._client.ttls[str(args[0])] = int(args[1])  # pyright: ignore[reportArgumentType]  # fake stores int
+                    results.append(True)
+                else:
+                    results.append(self._client.ttls.get(str(args[0]), -2))
+            return results
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("refresh_ttl", "existing_ttl", "expected_round_trips", "expected_ttl"),
+    [
+        pytest.param(True, 100, 2, 60, id="refresh_ttl: INCRBYFLOAT+EXPIRE in one round trip each"),
+        pytest.param(False, 100, 2, 100, id="keep ttl: INCRBYFLOAT+TTL in one round trip each, no EXPIRE"),
+        pytest.param(False, -1, 3, 60, id="unexpiring key: INCRBYFLOAT+TTL then EXPIRE once, 1 trip after"),
+    ],
+)
+async def test_async_increment_pipelines_the_ttl_command(
+    monkeypatch, redis_no_ping, refresh_ttl, existing_ttl, expected_round_trips, expected_ttl
+):
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    redis_cache = RedisCache(namespace="ns")
+    client = _RoundTripCountingRedis(ttl=existing_ttl)
+
+    with patch.object(redis_cache, "init_async_client", return_value=client):
+        first = await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=refresh_ttl)
+        second = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=refresh_ttl)
+
+    assert (first, second) == (1.5, 3.5)
+    assert client.values == {"ns:spend:key:k": 3.5}
+    assert client.ttls == {"ns:spend:key:k": expected_ttl}
+    assert client.round_trips == expected_round_trips
+
+
 class _SetRecordingPipeline:
     def __init__(self) -> None:
         self.sets: list[tuple[str, str, timedelta | None]] = []

tests/test_litellm/integrations/test_prometheus_carried_budget_state.py

@@ -0,0 +1,289 @@
+"""
+Post-request budget gauges read the key/team/user/org state auth already resolved
+from request metadata. get_*_object only runs when that state is missing (custom
+auth, SDK callers, failure paths)
+"""
+
+from datetime import datetime, timezone
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from prometheus_client import REGISTRY
+
+from litellm.integrations.prometheus import PrometheusLogger
+from litellm.proxy._types import (
+    LiteLLM_BudgetTable,
+    LiteLLM_OrganizationTable,
+    LiteLLM_TeamTable,
+    LiteLLM_UserTable,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+from litellm.proxy.spend_tracking.carried_budget_state import (
+    carried_budget_metadata,
+    carry_organization_budget_state,
+    carry_team_and_user_budget_state,
+)
+from litellm.types.proxy.carried_budget_state import (
+    KeyBudgetSnapshot,
+    TeamBudgetSnapshot,
+    UserBudgetSnapshot,
+)
+
+TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc)
+USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc)
+KEY_RESET_AT = datetime(2026, 12, 1, tzinfo=timezone.utc)
+
+
+@pytest.fixture(autouse=True)
+def cleanup_prometheus_registry():
+    for collector in list(REGISTRY._collector_to_names.keys()):
+        try:
+            REGISTRY.unregister(collector)
+        except Exception:
+            pass
+    yield
+    for collector in list(REGISTRY._collector_to_names.keys()):
+        try:
+            REGISTRY.unregister(collector)
+        except Exception:
+            pass
+
+
+@pytest.fixture
+def prometheus_logger():
+    return PrometheusLogger()
+
+
+@pytest.fixture
+def getters():
+    """Every response-path object getter, patched where prometheus imports them from."""
+    mocks = {
+        "get_key_object": AsyncMock(return_value=UserAPIKeyAuth(token="hashed", budget_reset_at=KEY_RESET_AT)),
+        "get_team_object": AsyncMock(
+            return_value=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT, max_budget=300.0)
+        ),
+        "get_user_object": AsyncMock(
+            return_value=LiteLLM_UserTable(
+                user_id="u1",
+                budget_reset_at=USER_RESET_AT,
+                user_email="alice@example.com",
+                user_alias="Alice",
+                max_budget=50.0,
+            )
+        ),
+        "get_org_object": AsyncMock(
+            return_value=LiteLLM_OrganizationTable(
+                organization_id="o1",
+                organization_alias="platform-org",
+                budget_id="b1",
+                created_by="admin",
+                updated_by="admin",
+                spend=40.0,
+                litellm_budget_table=LiteLLM_BudgetTable(max_budget=500.0),
+            )
+        ),
+    }
+    with (
+        patch.multiple(  # test-quality-ok: prometheus reads these proxy_server globals at call time, no injection seam
+            "litellm.proxy.proxy_server", prisma_client=MagicMock(), user_api_key_cache=MagicMock()
+        ),
+        patch.multiple(  # test-quality-ok: the getters are the DB boundary this test counts calls to
+            "litellm.proxy.auth.auth_checks", **mocks
+        ),
+    ):
+        yield mocks
+
+
+def _authed_token() -> UserAPIKeyAuth:
+    token = UserAPIKeyAuth(
+        token="hashed",
+        key_alias="key-alias",
+        team_id="t1",
+        team_alias="team-alias",
+        user_id="u1",
+        user_email="alice@example.com",
+        org_id="o1",
+        spend=1.0,
+        max_budget=10.0,
+        team_spend=20.0,
+        team_max_budget=300.0,
+        user_spend=5.0,
+        user_max_budget=50.0,
+        budget_reset_at=KEY_RESET_AT,
+    )
+    carry_team_and_user_budget_state(
+        valid_token=token,
+        team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT, max_budget=300.0),
+        user_object=LiteLLM_UserTable(user_id="u1", budget_reset_at=USER_RESET_AT, user_alias="Alice", max_budget=50.0),
+    )
+    carry_organization_budget_state(
+        valid_token=token,
+        org_table=LiteLLM_OrganizationTable(
+            organization_id="o1",
+            organization_alias="platform-org",
+            budget_id="b1",
+            created_by="admin",
+            updated_by="admin",
+            spend=40.0,
+            litellm_budget_table=LiteLLM_BudgetTable(max_budget=500.0),
+        ),
+    )
+    return token
+
+
+def _request_metadata(token: UserAPIKeyAuth) -> dict:
+    """What add_user_api_key_auth_to_request_metadata leaves in litellm_params["metadata"]."""
+    return {
+        **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token),
+        **carried_budget_metadata(token),
+    }
+
+
+def _stub_gauges(prometheus_logger: PrometheusLogger) -> None:
+    for name in (
+        "litellm_remaining_api_key_budget_metric",
+        "litellm_api_key_max_budget_metric",
+        "litellm_api_key_budget_remaining_hours_metric",
+        "litellm_remaining_team_budget_metric",
+        "litellm_team_max_budget_metric",
+        "litellm_team_budget_remaining_hours_metric",
+        "litellm_remaining_user_budget_metric",
+        "litellm_user_max_budget_metric",
+        "litellm_user_budget_remaining_hours_metric",
+        "litellm_remaining_org_budget_metric",
+        "litellm_org_max_budget_metric",
+        "litellm_org_budget_remaining_hours_metric",
+    ):
+        setattr(prometheus_logger, name, MagicMock())
+
+
+async def _emit(prometheus_logger: PrometheusLogger, metadata: dict) -> None:
+    await prometheus_logger._increment_remaining_budget_metrics(
+        user_api_team="t1",
+        user_api_team_alias="team-alias",
+        user_api_key="hashed",
+        user_api_key_alias="key-alias",
+        litellm_params={"metadata": metadata},
+        response_cost=2.0,
+        user_id="u1",
+        user_api_key_org_id="o1",
+    )
+
+
+@pytest.mark.asyncio
+async def test_authed_request_sets_every_gauge_without_any_object_getter(prometheus_logger, getters):
+    _stub_gauges(prometheus_logger)
+
+    await _emit(prometheus_logger, _request_metadata(_authed_token()))
+
+    assert all(getter.await_count == 0 for getter in getters.values()), {
+        name: getter.await_count for name, getter in getters.items()
+    }
+    remaining = {
+        "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0],
+        "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0],
+        "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0],
+        "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0],
+    }
+    assert remaining == {
+        "key": pytest.approx(7.0),
+        "team": pytest.approx(278.0),
+        "user": pytest.approx(43.0),
+        "org": 458.0,
+    }
+    prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(500.0)
+    prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.assert_called_once()
+    prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.assert_called_once()
+    prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters):
+    _stub_gauges(prometheus_logger)
+
+    await _emit(prometheus_logger, {"user_api_key_team_spend": 20.0, "user_api_key_team_max_budget": 300.0})
+
+    assert {name: getter.await_count for name, getter in getters.items()} == {
+        "get_key_object": 1,
+        "get_team_object": 1,
+        "get_user_object": 1,
+        "get_org_object": 1,
+    }
+    prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0)
+    prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_partial_carried_state_only_skips_the_carried_objects(prometheus_logger, getters):
+    _stub_gauges(prometheus_logger)
+    token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1")
+    carry_team_and_user_budget_state(
+        valid_token=token,
+        team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT),
+        user_object=None,
+    )
+
+    await _emit(prometheus_logger, dict(carried_budget_metadata(token)))
+
+    assert {name: getter.await_count for name, getter in getters.items()} == {
+        "get_key_object": 1,
+        "get_team_object": 0,
+        "get_user_object": 1,
+        "get_org_object": 1,
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("metadata_max_budget", [300.0, None], ids=["metadata has max_budget", "filled from object"])
+async def test_carried_objects_match_what_the_getters_would_have_produced(
+    prometheus_logger, getters, metadata_max_budget
+):
+    metadata = _request_metadata(_authed_token())
+    user_max_budget = 50.0 if metadata_max_budget is not None else None
+
+    carried_team = await prometheus_logger._assemble_team_object(
+        team_id="t1",
+        team_alias="team-alias",
+        spend=20.0,
+        max_budget=metadata_max_budget,
+        response_cost=2.0,
+        carried=TeamBudgetSnapshot.from_metadata(metadata),
+    )
+    fetched_team = await prometheus_logger._assemble_team_object(
+        team_id="t1", team_alias="team-alias", spend=20.0, max_budget=metadata_max_budget, response_cost=2.0
+    )
+    carried_user = await prometheus_logger._assemble_user_object(
+        user_id="u1",
+        spend=5.0,
+        max_budget=user_max_budget,
+        response_cost=2.0,
+        carried=UserBudgetSnapshot.from_metadata(metadata),
+        user_email="alice@example.com",
+    )
+    fetched_user = await prometheus_logger._assemble_user_object(
+        user_id="u1", spend=5.0, max_budget=user_max_budget, response_cost=2.0
+    )
+    carried_key = await prometheus_logger._assemble_key_object(
+        user_api_key="hashed",
+        user_api_key_alias="key-alias",
+        key_max_budget=10.0,
+        key_spend=1.0,
+        response_cost=2.0,
+        carried=KeyBudgetSnapshot.from_metadata(metadata),
+    )
+    fetched_key = await prometheus_logger._assemble_key_object(
+        user_api_key="hashed", user_api_key_alias="key-alias", key_max_budget=10.0, key_spend=1.0, response_cost=2.0
+    )
+
+    assert carried_team == fetched_team
+    assert carried_team.max_budget == 300.0
+    assert carried_user == fetched_user
+    assert carried_user.max_budget == 50.0
+    assert carried_key == fetched_key
+    assert {name: getter.await_count for name, getter in getters.items()} == {
+        "get_key_object": 1,
+        "get_team_object": 1,
+        "get_user_object": 1,
+        "get_org_object": 0,
+    }

tests/test_litellm/proxy/auth/test_auth_checks.py

@@ -5791,6 +5791,42 @@ async def test_budget_checks_only_run_on_llm_api_routes(scope, route, expect_blo
             assert await _run() is True
 
 
+@pytest.mark.asyncio
+async def test_organization_budget_check_carries_org_state_on_the_token():
+    """The org row auth already fetched is pinned on the token so the response path
+    (Prometheus org budget gauges) reads it from request metadata instead of calling
+    get_org_object again."""
+    from litellm.proxy._types import LiteLLM_OrganizationTable
+    from litellm.proxy.auth.auth_checks import _organization_max_budget_check
+    from litellm.types.proxy.carried_budget_state import OrgBudgetSnapshot
+
+    org_table = LiteLLM_OrganizationTable(
+        organization_id="o1",
+        organization_alias="platform-org",
+        budget_id="b1",
+        created_by="admin",
+        updated_by="admin",
+        spend=12.5,
+        litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0),
+    )
+    token = UserAPIKeyAuth(token="k1", org_id="o1")
+    user_api_key_cache = UserApiKeyCache()
+    await user_api_key_cache.async_set_cache(
+        key="org_id:o1:with_budget", value=org_table, model_type=LiteLLM_OrganizationTable
+    )
+
+    await _organization_max_budget_check(
+        valid_token=token,
+        team_object=None,
+        prisma_client=MagicMock(),
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=MagicMock(),
+    )
+
+    assert token.organization_alias == "platform-org"
+    assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0)
+
+
 @pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"])
 @pytest.mark.asyncio
 async def test_spend_capable_non_llm_routes_still_enforce_budget(route):

tests/test_litellm/proxy/auth/test_user_api_key_auth.py

@@ -5,7 +5,7 @@ import os
 import subprocess
 import sys
 from contextlib import contextmanager
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from textwrap import dedent
 from types import SimpleNamespace
@@ -23,6 +23,7 @@ from litellm.proxy._types import (
     LiteLLM_JWTAuth,
     LiteLLM_BudgetTable,
     LiteLLM_EndUserTable,
+    LiteLLM_TeamTableCachedObj,
     LiteLLM_UserTable,
     LitellmUserRoles,
     ProxyErrorTypes,
@@ -48,6 +49,7 @@ from litellm.proxy.auth.user_api_key_auth import (
     get_api_key,
     user_api_key_auth,
 )
+from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
 
 
 class _RoutingRequest:
@@ -4019,6 +4021,65 @@ async def test_centralized_common_checks_routes_header_tags_to_litellm_metadata(
     assert "metadata" not in request_data
 
 
+@pytest.mark.asyncio
+async def test_centralized_common_checks_carries_team_and_user_budget_state_on_the_token():
+    """The team and user objects auth resolves are pinned on the token so the
+    response path (Prometheus budget gauges) reads them from request metadata
+    instead of calling get_team_object / get_user_object again."""
+    from fastapi import Request
+    from starlette.datastructures import URL
+
+    import litellm.proxy.proxy_server as _proxy_server_mod
+
+    reset_at = datetime(2026, 10, 1, tzinfo=timezone.utc)
+    token = UserAPIKeyAuth(api_key="sk-test", token="hashed", team_id="t1", user_id="u1")
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    user_api_key_cache = DualCache()
+    await user_api_key_cache.async_set_cache(
+        key="team_id:t1",
+        value=LiteLLM_TeamTableCachedObj(team_id="t1", budget_reset_at=reset_at, max_budget=300.0),
+    )
+    await user_api_key_cache.async_set_cache(
+        key="u1",
+        value=LiteLLM_UserTable(user_id="u1", user_alias="Alice", budget_reset_at=None, max_budget=None),
+    )
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
+    attrs = {
+        **_proxy_attrs_for_centralized_checks(user_custom_auth=None),
+        "prisma_client": MagicMock(),
+        "user_api_key_cache": user_api_key_cache,
+        "proxy_logging_obj": proxy_logging_obj,
+    }
+    originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
+    try:
+        for k, v in attrs.items():
+            setattr(_proxy_server_mod, k, v)
+        with patch(  # test-quality-ok: the authz gate has its own tests above; this one checks the carry step before it
+            "litellm.proxy.auth.user_api_key_auth.common_checks",
+            new_callable=AsyncMock,
+        ):
+            await _run_centralized_common_checks(
+                user_api_key_auth_obj=token,
+                request=request,
+                request_data={"model": "gpt-5.4-mini"},
+                route="/chat/completions",
+            )
+    finally:
+        for k, v in originals.items():
+            setattr(_proxy_server_mod, k, v)
+
+    assert dict(carried_budget_metadata(token)) == {
+        "user_api_key_team_budget_reset_at": "2026-10-01T00:00:00Z",
+        "user_api_key_team_table_max_budget": 300.0,
+        "user_api_key_user_budget_reset_at": None,
+        "user_api_key_user_table_max_budget": None,
+        "user_api_key_user_alias": "Alice",
+    }
+
+
 @pytest.mark.asyncio
 async def test_centralized_common_checks_skipped_for_custom_auth_without_flag():
     """Existing RPS guarantee: custom-auth deployments without

tests/test_litellm/proxy/db/test_spend_counter_reseed.py

@@ -350,7 +350,7 @@ async def test_cold_reseed_preserves_concurrent_local_increment(
 
     increment_task: Final = asyncio.create_task(
         proxy_server._apply_spend_counter_increments(
-            pending=(proxy_server._PendingSpendIncrement(counter_key=counter_key, increment=increment),)
+            pending=(proxy_server.PendingSpendIncrement(counter_key=counter_key, increment=increment),)
         )
         if batch
         else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment)

tests/test_litellm/proxy/proxy_server/test_spend_counters.py

@@ -1151,10 +1151,10 @@ async def test_prepare_window_spend_counter_increment_missing_window_start_inval
 # ---------------------------------------------------------------------------
 
 
-def _two_pending_increments() -> tuple[ps._PendingSpendIncrement, ...]:
+def _two_pending_increments() -> tuple[ps.PendingSpendIncrement, ...]:
     return (
-        ps._PendingSpendIncrement(counter_key="spend:key:k", increment=1.5),
-        ps._PendingSpendIncrement(counter_key="spend:team:t", increment=1.5),
+        ps.PendingSpendIncrement(counter_key="spend:key:k", increment=1.5),
+        ps.PendingSpendIncrement(counter_key="spend:team:t", increment=1.5),
     )
 
 

tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py

@@ -0,0 +1,119 @@
+"""Auth-resolved budget state rides on ``UserAPIKeyAuth`` and round-trips through request metadata."""
+
+from datetime import datetime, timezone
+
+from litellm.proxy._types import (
+    LiteLLM_BudgetTable,
+    LiteLLM_OrganizationTable,
+    LiteLLM_TeamTable,
+    LiteLLM_UserTable,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.spend_tracking.carried_budget_state import (
+    carried_budget_metadata,
+    carry_organization_budget_state,
+    carry_team_and_user_budget_state,
+)
+from litellm.types.proxy.carried_budget_state import (
+    KeyBudgetSnapshot,
+    OrgBudgetSnapshot,
+    TeamBudgetSnapshot,
+    UserBudgetSnapshot,
+)
+
+RESET_AT = datetime(2026, 10, 1, 12, 30, tzinfo=timezone.utc)
+
+
+def test_team_and_user_state_round_trips_through_metadata():
+    token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1")
+    carry_team_and_user_budget_state(
+        valid_token=token,
+        team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=RESET_AT, max_budget=300.0),
+        user_object=LiteLLM_UserTable(user_id="u1", budget_reset_at=None, max_budget=None, user_alias="Alice"),
+    )
+
+    metadata = dict(carried_budget_metadata(token))
+
+    assert metadata == {
+        "user_api_key_team_budget_reset_at": RESET_AT.isoformat().replace("+00:00", "Z"),
+        "user_api_key_team_table_max_budget": 300.0,
+        "user_api_key_user_budget_reset_at": None,
+        "user_api_key_user_table_max_budget": None,
+        "user_api_key_user_alias": "Alice",
+    }
+    assert TeamBudgetSnapshot.from_metadata(metadata) == TeamBudgetSnapshot(budget_reset_at=RESET_AT, max_budget=300.0)
+    assert UserBudgetSnapshot.from_metadata(metadata) == UserBudgetSnapshot(
+        budget_reset_at=None, max_budget=None, user_alias="Alice"
+    )
+
+
+def test_missing_objects_leave_no_metadata_and_no_snapshot():
+    token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1")
+    carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None)
+
+    assert dict(carried_budget_metadata(token)) == {}
+    assert TeamBudgetSnapshot.from_metadata({}) is None
+    assert UserBudgetSnapshot.from_metadata({"user_api_key_user_alias": "Alice"}) is None
+    assert OrgBudgetSnapshot.from_metadata({"user_api_key_org_spend": 1.0}) is None
+    assert KeyBudgetSnapshot.from_metadata({}) is None
+
+
+def test_organization_state_carries_alias_spend_and_max_budget():
+    token = UserAPIKeyAuth(token="hashed", org_id="o1")
+    org = LiteLLM_OrganizationTable(
+        organization_id="o1",
+        organization_alias="platform-org",
+        budget_id="b1",
+        created_by="admin",
+        updated_by="admin",
+        spend=12.5,
+        litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0),
+    )
+
+    carry_organization_budget_state(valid_token=token, org_table=org)
+
+    assert token.organization_alias == "platform-org"
+    assert OrgBudgetSnapshot.from_metadata(carried_budget_metadata(token)) == OrgBudgetSnapshot(
+        spend=12.5, max_budget=100.0
+    )
+
+
+def test_organization_without_budget_table_carries_no_cap():
+    token = UserAPIKeyAuth(token="hashed", org_id="o1")
+    org = LiteLLM_OrganizationTable(
+        organization_id="o1",
+        budget_id="b1",
+        created_by="admin",
+        updated_by="admin",
+        spend=3.0,
+    )
+
+    carry_organization_budget_state(valid_token=token, org_table=org)
+
+    assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=3.0, max_budget=None)
+
+
+def test_key_snapshot_parses_the_iso_string_auth_metadata_writes():
+    assert KeyBudgetSnapshot.from_metadata({"user_api_key_budget_reset_at": RESET_AT.isoformat()}) == KeyBudgetSnapshot(
+        budget_reset_at=RESET_AT
+    )
+    assert KeyBudgetSnapshot.from_metadata({"user_api_key_budget_reset_at": None}) == KeyBudgetSnapshot(
+        budget_reset_at=None
+    )
+
+
+def test_snapshots_never_reach_the_serialized_token():
+    token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1")
+    carry_team_and_user_budget_state(
+        valid_token=token,
+        team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=RESET_AT),
+        user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"),
+    )
+    token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0)
+
+    dumped = token.model_dump()
+
+    assert "team_budget_snapshot" not in dumped
+    assert "user_budget_snapshot" not in dumped
+    assert "org_budget_snapshot" not in dumped
+    assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None

tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -52,6 +52,24 @@ class CountingRedis(RedisCache):
         self.commands.append(f"MGET {' '.join(key_list)}")
         return {key: self.store.get(key) for key in key_list}
 
+    def get_ttl(self, **kwargs: object) -> int | None:
+        return None
+
+    async def async_increment(self, key: str, value: float, **kwargs: object) -> float:
+        self.commands.append(f"INCRBYFLOAT {key} {value}")
+        return self._incr(key, value)
+
+    async def async_increment_pipeline(
+        self, increment_list: Sequence[Mapping[str, object]], **kwargs: object
+    ) -> list[float]:
+        self.commands.append(f"PIPELINE {' '.join(str(op['key']) for op in increment_list)}")
+        return [self._incr(str(op["key"]), float(str(op["increment_value"]))) for op in increment_list]
+
+    def _incr(self, key: str, value: float) -> float:
+        total = float(str(self.store.get(key, 0.0))) + value
+        self.store[key] = total
+        return total
+
 
 def _spend_counter_cache(redis: RedisCache | None, in_memory: dict[str, float] | None = None) -> MagicMock:
     cache = MagicMock()
@@ -119,6 +137,37 @@ async def test_keys_bound_after_the_first_read_join_one_more_mget_for_only_the_n
     assert redis.commands == ["MGET spend:key:hashed", "MGET spend:org:org"]
 
 
+@pytest.mark.asyncio
+async def test_a_recorded_write_result_answers_later_reads_without_another_redis_read():
+    redis = CountingRedis({"spend:key:hashed": 1.0})
+    batch = SpendCounterBatch(redis)
+    batch.bind(frozenset({"spend:key:hashed"}))
+    assert await batch.read("spend:key:hashed") == (1.0, True)
+
+    batch.record("spend:key:hashed", 3.5)
+    batch.record("spend:org:org", 7.0)
+
+    assert await batch.read("spend:key:hashed") == (3.5, True)
+    assert await batch.read("spend:org:org") == (7.0, True)
+    assert redis.commands == ["MGET spend:key:hashed"]
+
+
+@pytest.mark.asyncio
+async def test_a_forgotten_counter_is_read_fresh_from_redis_when_it_is_bound_again():
+    redis = CountingRedis({"spend:key:hashed": 1.0})
+    batch = SpendCounterBatch(redis)
+    batch.bind(frozenset({"spend:key:hashed"}))
+    batch.record("spend:key:hashed", 3.5)
+
+    batch.forget("spend:key:hashed")
+    assert await batch.read("spend:key:hashed") is None
+
+    redis.store["spend:key:hashed"] = 9.0
+    batch.bind(frozenset({"spend:key:hashed"}))
+    assert await batch.read("spend:key:hashed") == (9.0, True)
+    assert redis.commands == ["MGET spend:key:hashed"]
+
+
 @pytest.mark.asyncio
 async def test_failed_mget_hands_every_counter_back_to_the_caller():
     batch = SpendCounterBatch(CountingRedis(fail=True))
@@ -298,3 +347,181 @@ async def test_reseed_outside_the_scope_still_re_checks_redis_itself():
 
     assert value == 4.0
     assert redis.commands == ["GET spend:key:hashed"]
+
+
+POST_CALL_KEYS = TOKEN_KEYS | {"spend:tag:prod", "spend:model_access_group:premium"}
+
+
+@pytest.mark.asyncio
+async def test_post_call_increment_for_every_entity_costs_one_mget_and_one_pipeline(monkeypatch):
+    redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
+    monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
+    monkeypatch.setattr(ps, "prisma_client", None)
+
+    await ps.increment_spend_counters(
+        token="hashed",
+        team_id="team",
+        user_id="user",
+        org_id="org",
+        end_user_id="eu",
+        tags=["prod"],
+        model_access_groups=["premium"],
+        response_cost=0.5,
+    )
+
+    assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
+    assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS
+    assert set(redis.commands[1].split()[1:]) == POST_CALL_KEYS
+    assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.5 for key in POST_CALL_KEYS}
+
+
+@pytest.mark.asyncio
+async def test_post_call_cold_counters_seed_from_the_mget_miss_without_a_second_read(monkeypatch):
+    redis = CountingRedis({"spend:key:hashed": 1.0})
+    redis.async_set_cache = AsyncMock(return_value=True)
+    prisma = MagicMock()
+    prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0))
+    monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+
+    await ps.increment_spend_counters(token="hashed", team_id="team", user_id=None, response_cost=0.5)
+
+    assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
+    redis.async_set_cache.assert_awaited_once_with(key="spend:team:team", value=4.0, nx=True)
+    assert redis.store["spend:key:hashed"] == 1.5
+
+
+RESERVED_KEYS = frozenset(
+    {"spend:key:hashed", "spend:team:team", "spend:team_member:user:team", "spend:end_user:eu", "spend:org:org"}
+)
+
+
+def _reservation(reserved_cost: float, counter_keys: frozenset[str] = RESERVED_KEYS) -> dict[str, object]:
+    return {
+        "reserved_cost": reserved_cost,
+        "entries": [
+            {"counter_key": key, "entity_type": "Key", "entity_id": key, "reserved_cost": reserved_cost}
+            for key in sorted(counter_keys)
+        ],
+    }
+
+
+@pytest.mark.asyncio
+async def test_post_call_with_a_reservation_costs_one_mget_one_reconcile_pipeline_one_increment_pipeline(monkeypatch):
+    redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
+    monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
+    monkeypatch.setattr(ps, "prisma_client", None)
+    reservation = _reservation(reserved_cost=0.4)
+
+    await ps.increment_spend_counters(
+        token="hashed",
+        team_id="team",
+        user_id="user",
+        org_id="org",
+        end_user_id="eu",
+        tags=["prod"],
+        model_access_groups=["premium"],
+        response_cost=0.5,
+        budget_reservation=reservation,
+    )
+
+    assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "PIPELINE"], redis.commands
+    assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS, "reconcile and warm checks share the MGET"
+    assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS
+    assert set(redis.commands[2].split()[1:]) == POST_CALL_KEYS - RESERVED_KEYS
+    assert {key: round(redis.store[key], 6) for key in POST_CALL_KEYS} == {
+        key: (1.1 if key in RESERVED_KEYS else 1.5) for key in POST_CALL_KEYS
+    }
+    assert [round(entry["applied_adjustment"], 6) for entry in reservation["entries"]] == [0.1] * len(RESERVED_KEYS)
+    assert reservation["finalized"] is True
+
+
+@pytest.mark.asyncio
+async def test_reconcile_settles_a_flushed_counter_on_its_own_after_the_shared_pipeline(monkeypatch):
+    from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation
+
+    redis = CountingRedis({key: 1.0 for key in RESERVED_KEYS - {"spend:team:team"}})
+    redis.async_set_max = AsyncMock(return_value=4.0)
+    prisma = MagicMock()
+    prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0))
+    monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
+    monkeypatch.setattr(ps, "prisma_client", prisma)
+    reservation = _reservation(reserved_cost=0.4)
+
+    await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.5)
+
+    assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "INCRBYFLOAT"], redis.commands
+    assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS - {"spend:team:team"}
+    assert redis.commands[2] == "INCRBYFLOAT spend:team:team 0.5"
+    redis.async_set_max.assert_awaited_once()
+    assert redis.async_set_max.await_args.kwargs["key"] == "spend:team:team"
+    assert all(round(entry["applied_adjustment"], 6) == 0.1 for entry in reservation["entries"])
+
+
+@pytest.mark.asyncio
+async def test_pre_call_resize_against_an_inconsistent_counter_writes_nothing_and_denies(monkeypatch):
+    from litellm.proxy.spend_tracking.budget_reservation import _resize_applied_reservation
+
+    redis = CountingRedis({"spend:key:hashed": 1.0})
+    monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
+    monkeypatch.setattr(ps, "prisma_client", None)
+    entries = _reservation(reserved_cost=0.4, counter_keys=frozenset({"spend:key:hashed", "spend:team:team"}))[
+        "entries"
+    ]
+
+    with pytest.raises(RuntimeError, match="spend:team:team"):
+        await _resize_applied_reservation(entries=entries, current_reserved_cost=0.4, new_reserved_cost=0.9)
+
+    assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands
+    assert redis.store["spend:key:hashed"] == 1.0
+    assert all("applied_adjustment" not in entry for entry in entries)
+
+
+@pytest.mark.asyncio
+async def test_a_failed_reconcile_pipeline_invalidates_every_reserved_counter_and_falls_back(monkeypatch):
+    redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
+    redis.async_delete_cache = AsyncMock()
+    reconcile_pipeline_failed = False
+
+    async def _pipeline(increment_list: Sequence[Mapping[str, object]], **kwargs: object) -> list[float]:
+        nonlocal reconcile_pipeline_failed
+        if not reconcile_pipeline_failed:
+            reconcile_pipeline_failed = True
+            raise ConnectionError("redis down")
+        return await CountingRedis.async_increment_pipeline(redis, increment_list, **kwargs)
+
+    redis.async_increment_pipeline = _pipeline  # pyright: ignore[reportAttributeAccessIssue]  # instance override
+    monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
+    monkeypatch.setattr(ps, "prisma_client", None)
+    reservation = _reservation(reserved_cost=0.4)
+
+    await ps.increment_spend_counters(
+        token="hashed",
+        team_id="team",
+        user_id="user",
+        org_id="org",
+        end_user_id="eu",
+        response_cost=0.5,
+        budget_reservation=reservation,
+    )
+
+    assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == RESERVED_KEYS
+    assert all("applied_adjustment" not in entry for entry in reservation["entries"])
+    assert redis.commands[-1].split()[0] == "PIPELINE"
+    assert set(redis.commands[-1].split()[1:]) == RESERVED_KEYS | {"spend:user:user"}
+
+
+def test_a_scope_opened_inside_an_open_scope_joins_its_batch_and_a_closed_one_gets_its_own():
+    redis = CountingRedis()
+    with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:a"})):
+        outer = active_spend_counter_batch()
+        assert outer is not None
+        with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:b"})):
+            assert active_spend_counter_batch() is outer
+        assert outer.counter_keys == {"spend:key:a", "spend:key:b"}
+        release_spend_counter_batch()
+        with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:c"})):
+            inner = active_spend_counter_batch()
+            assert inner is not outer
+            assert inner is not None and inner.counter_keys == {"spend:key:c"}
+        assert active_spend_counter_batch() is outer

tests/test_litellm/proxy/test_proxy_server.py

@@ -8419,7 +8419,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments():
     async def assert_reservation_not_finalized_yet(**kwargs):
         assert budget_reservation["finalized"] is False
         incremented_counters.append(kwargs["counter_key"])
-        return ps._PendingSpendIncrement(
+        return ps.PendingSpendIncrement(
             counter_key=kwargs["counter_key"], increment=kwargs["increment"]
         )