fix(budgets): write the model access group spend counter after each call

Two problems, both caught in review.

The new table only landed in the root schema.prisma. Client generation reads
litellm/proxy/schema.prisma and packaging reads the copy under
litellm-proxy-extras, so the generated client had no
litellm_modelaccessgroupbudgettable and every budget read and write against
it would have failed at runtime. The root is the source of truth; both
copies are now byte-identical to it.

Nothing incremented spend:model_access_group:{group} after a call. Only the
reservation path ever wrote it, so with disable_budget_reservation the
read-time check was reading a counter nobody maintained and falling back to
the row's spend, which is cached for up to DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL.
A caller could run well past the pool inside that window, which is precisely
the case the read-time check exists to cover.

increment_spend_counters now takes the matched groups and charges them
through _init_and_increment_unreserved_spend_counter, so a group already
covered by a reservation is skipped rather than counted twice. The cost
callback sources the names with get_request_model_access_groups, the same
reader the spend writer uses.

litellm/proxy/common_utils/user_api_key_cache.py

@@ -201,11 +201,12 @@ def model_access_group_registry_cache_key() -> str:
 
 
 def model_access_group_spend_counter_key(access_group_name: str) -> str:
-    """Spend counter key for one model access group; shared so its three owners cannot drift.
+    """Spend counter key for one model access group; shared so its four owners cannot drift.
 
-    The reservation path writes it, auth reads it to enforce ``max_budget``, and the reset job
-    clears it on rollover. A copy that drifts in any one of them silently resets or reads a
-    counter nobody else touches, which shows up as a budget that never trips or never resets.
+    The reservation path writes it up front, the cost callback writes it after the call, auth
+    reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that
+    drifts in any one of them silently resets or reads a counter nobody else touches, which shows
+    up as a budget that never trips or never resets.
     """
     return f"spend:model_access_group:{access_group_name}"
 

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -1,5 +1,6 @@
 import asyncio
 import traceback
+from collections.abc import Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, cast
 
@@ -27,6 +28,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import (
 )
 from litellm.proxy.spend_tracking.spend_tracking_utils import (
     _sanitize_error_information_for_spend_logs,
+    get_request_model_access_groups,
 )
 from litellm.proxy.utils import ProxyUpdateSpend
 from litellm.types.utils import (
@@ -258,6 +260,7 @@ class _ProxyDBLogger(CustomLogger):
                 sl_object=sl_object,
                 metadata=metadata,
             )
+            model_access_groups: Final = get_request_model_access_groups(kwargs)
 
             if response_cost is not None:
                 user_api_key: Final = metadata.get("user_api_key", None)
@@ -296,6 +299,7 @@ class _ProxyDBLogger(CustomLogger):
                         response_cost=response_cost,
                         budget_reservation=budget_reservation,
                         request_tags=tags,
+                        model_access_groups=model_access_groups,
                     )
 
                     # update cache (fire-and-forget for backward compat:
@@ -572,6 +576,7 @@ async def _update_database_and_spend_counters(
     response_cost: float,
     budget_reservation: dict | None,
     request_tags: list[str] | None = None,
+    model_access_groups: Sequence[str] | None = None,
 ) -> None:
     try:
         await proxy_logging_obj.db_spend_update_writer.update_database(
@@ -610,6 +615,7 @@ async def _update_database_and_spend_counters(
             budget_reservation=budget_reservation,
             end_user_id=end_user_id,
             tags=request_tags,
+            model_access_groups=model_access_groups,
         )
     except Exception:
         if budget_reservation is not None:

litellm/proxy/proxy_server.py

@@ -382,6 +382,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
     UserApiKeyCache,
     end_user_cache_key,
     get_management_object_ttl,
+    model_access_group_cache_key,
+    model_access_group_spend_counter_key,
     tag_cache_key,
 )
 from litellm.proxy.config_resolvers import resolve_fields
@@ -2648,6 +2650,7 @@ async def increment_spend_counters(
     budget_reservation: dict | None = None,
     end_user_id: str | None = None,
     tags: list[str] | None = None,
+    model_access_groups: Sequence[str] | None = None,
 ):
     """
     Atomically increment spend counters for budget enforcement.
@@ -2777,6 +2780,13 @@ async def increment_spend_counters(
             )
             if end_user_id is not None or tags is not None
             else None,
+            _increment_model_access_group_spend_counters(
+                model_access_groups=model_access_groups,
+                response_cost=cost,
+                reserved_counter_keys=reserved_counter_keys,
+            )
+            if model_access_groups
+            else None,
             _increment_org_spend_counter(
                 org_id=org_id,
                 response_cost=cost,
@@ -2865,6 +2875,33 @@ async def _increment_end_user_and_tag_spend_counters(
         )
 
 
+async def _increment_model_access_group_spend_counters(
+    model_access_groups: Sequence[object],
+    response_cost: float,
+    reserved_counter_keys: set[str],
+) -> None:
+    """Charge the model access groups that authorized this request.
+
+    Without this the counter auth reads is written only by the reservation path, so
+    ``disable_budget_reservation`` would leave ``_model_access_group_max_budget_check`` enforcing
+    against the DB row's spend, which lags by up to the cache TTL.
+
+    Typed ``object`` rather than ``str`` because the names reach the cost callback out of request
+    metadata, which the coercion upstream filters to a list but not to strings. A non-string that
+    slipped through would build a counter key nothing else ever reads.
+    """
+    unique_groups: Final = tuple(
+        dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str))
+    )
+    for group in unique_groups:
+        await _init_and_increment_unreserved_spend_counter(
+            counter_key=model_access_group_spend_counter_key(group),
+            source_cache_key=model_access_group_cache_key(group),
+            increment=response_cost,
+            reserved_counter_keys=reserved_counter_keys,
+        )
+
+
 async def _increment_org_spend_counter(
     org_id: str | None,
     response_cost: float,

tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py

@@ -1,14 +1,14 @@
 
-import pytest
-
-
 from datetime import datetime
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import pytest
+
+from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.hooks.proxy_track_cost_callback import (
-    _ProxyDBLogger,
     _get_budget_reservation_from_metadata,
+    _ProxyDBLogger,
     _should_track_cost_callback,
     _update_database_and_spend_counters,
 )
@@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
         response_cost=0.2,
         budget_reservation=budget_reservation,
         request_tags=["tag-a"],
+        model_access_groups=("premium",),
     )
 
     proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
@@ -598,6 +599,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
         budget_reservation=budget_reservation,
         end_user_id="test_end_user_id",
         tags=["tag-a"],
+        model_access_groups=("premium",),
     )
 
 
@@ -1875,3 +1877,82 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
         assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (
             1 if expect_spend_log else 0
         )
+
+
+@pytest.mark.asyncio
+async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped():
+    """Auth stamps the matched groups onto request metadata; the callback has to carry them through.
+
+    Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with
+    reservations disabled the budget check reads a counter no one maintains.
+    """
+    logger = _ProxyDBLogger()
+    kwargs = {
+        "model": "gpt-4",
+        "call_type": "acompletion",
+        "litellm_params": {
+            "metadata": {
+                "user_api_key": "hashed-key",
+                "user_api_key_user_id": "user-1",
+                MODEL_ACCESS_GROUP_METADATA_KEY: ["premium", "starter"],
+            },
+        },
+        "standard_logging_object": {"response_cost": 0.25, "request_tags": None},
+        "stream": False,
+    }
+
+    with (
+        patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
+        patch(
+            "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
+        ) as mock_increment_spend_counters,
+        patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock),
+    ):
+        mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
+        mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
+        await logger._PROXY_track_cost_callback(
+            kwargs=kwargs,
+            completion_response=None,
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+        )
+
+    mock_increment_spend_counters.assert_awaited_once()
+    assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == (
+        "premium",
+        "starter",
+    )
+
+
+@pytest.mark.asyncio
+async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped():
+    """A request no budgeted group authorized must not debit anything."""
+    logger = _ProxyDBLogger()
+    kwargs = {
+        "model": "gpt-4",
+        "call_type": "acompletion",
+        "litellm_params": {
+            "metadata": {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"},
+        },
+        "standard_logging_object": {"response_cost": 0.25, "request_tags": None},
+        "stream": False,
+    }
+
+    with (
+        patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
+        patch(
+            "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
+        ) as mock_increment_spend_counters,
+        patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock),
+    ):
+        mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
+        mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
+        await logger._PROXY_track_cost_callback(
+            kwargs=kwargs,
+            completion_response=None,
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+        )
+
+    mock_increment_spend_counters.assert_awaited_once()
+    assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == ()

tests/test_litellm/proxy/test_budget_reservation.py

@@ -33,6 +33,7 @@ from litellm.proxy.common_utils.reset_budget_job import _model_access_group_coun
 from litellm.proxy.common_utils.user_api_key_cache import (
     UserApiKeyCache,
     model_access_group_cache_key,
+    model_access_group_spend_counter_key,
 )
 from litellm.proxy.spend_tracking.budget_reservation import (
     TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS,
@@ -47,6 +48,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
 )
 from litellm.proxy.utils import ProxyLogging
 from litellm.router import Router
+from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
 
 
 @pytest.fixture()
@@ -3067,3 +3069,117 @@ async def test_model_access_group_counter_blocks_a_request_over_the_group_budget
 
     assert exc_info.value.entity_id == "premium"
     assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value
+
+
+async def _cache_model_access_group_budget(key_cache, group, spend, max_budget=None):
+    await key_cache.async_set_cache(
+        key=model_access_group_cache_key(group),
+        value=ModelAccessGroupBudget(access_group_name=group, spend=spend, max_budget=max_budget),
+        model_type=ModelAccessGroupBudget,
+    )
+
+
+async def _reserve_for_model_access_groups(key_cache, groups, estimate):
+    """Reserve against the given groups, whose rows are already cached, so nothing hits the DB."""
+    with patch(
+        "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
+        return_value=estimate,
+    ):
+        return await reserve_budget_for_request(
+            request_body=_request_body(),
+            route="/chat/completions",
+            llm_router=None,
+            valid_token=UserAPIKeyAuth(
+                api_key="hashed", token="tok-mag-counter", matched_model_access_groups=list(groups)
+            ),
+            team_object=None,
+            user_object=None,
+            prisma_client=_ModelAccessGroupBudgetPrisma(),
+            user_api_key_cache=key_cache,
+            proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache),
+        )
+
+
+@pytest.mark.asyncio
+async def test_model_access_group_counter_accumulates_across_calls_without_a_reservation(spend_counter_state):
+    """With reservations disabled nothing writes the counter up front, so the cost callback must.
+
+    Otherwise the read-time budget check enforces against the DB row's spend, which the cache
+    holds for the full TTL, and a caller runs past the ceiling for that whole window.
+    """
+    counter_cache, key_cache = spend_counter_state
+    await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0)
+
+    from litellm.proxy.proxy_server import increment_spend_counters
+
+    counter_key = model_access_group_spend_counter_key("premium")
+
+    await increment_spend_counters(
+        token=None, team_id=None, user_id=None, response_cost=0.25, model_access_groups=["premium"]
+    )
+    assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.25)
+
+    await increment_spend_counters(
+        token=None, team_id=None, user_id=None, response_cost=0.75, model_access_groups=["premium", "premium", ""]
+    )
+    assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.0)
+    assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None
+
+
+@pytest.mark.asyncio
+async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state):
+    """The reservation already wrote this counter, so the post-call pass has to skip it."""
+    counter_cache, key_cache = spend_counter_state
+    await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0)
+
+    reservation = await _reserve_for_model_access_groups(key_cache, ["premium"], estimate=0.6)
+    counter_key = model_access_group_spend_counter_key("premium")
+    assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.6)
+
+    from litellm.proxy.proxy_server import increment_spend_counters
+
+    await increment_spend_counters(
+        token=None,
+        team_id=None,
+        user_id=None,
+        response_cost=0.2,
+        budget_reservation=reservation,
+        model_access_groups=["premium"],
+    )
+
+    # 1.0 recorded + the reservation reconciled down to the 0.2 actually spent. A second
+    # increment would land at 1.4.
+    assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.2)
+
+
+@pytest.mark.asyncio
+async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one(spend_counter_state):
+    """A budgetless group reserves nothing, so only the post-call pass can charge it.
+
+    Both groups authorized the request and both get debited, each exactly once, whether or not
+    the reservation path happened to hold a counter for them.
+    """
+    counter_cache, key_cache = spend_counter_state
+    await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0)
+    await _cache_model_access_group_budget(key_cache, "starter", spend=4.0)
+
+    reservation = await _reserve_for_model_access_groups(key_cache, ["premium", "starter"], estimate=0.6)
+    assert [entry["entity_id"] for entry in reservation["entries"]] == ["premium"]
+
+    from litellm.proxy.proxy_server import increment_spend_counters
+
+    await increment_spend_counters(
+        token=None,
+        team_id=None,
+        user_id=None,
+        response_cost=0.2,
+        budget_reservation=reservation,
+        model_access_groups=["premium", "starter", "starter", "premium"],
+    )
+
+    assert counter_cache.in_memory_cache.get_cache(
+        key=model_access_group_spend_counter_key("premium")
+    ) == pytest.approx(1.2)
+    assert counter_cache.in_memory_cache.get_cache(
+        key=model_access_group_spend_counter_key("starter")
+    ) == pytest.approx(4.2)