fix(prometheus): move the committed counters for charges with no auth on the row

A billed background response is charged from its creator's stored attribution, after the request's auth is gone. Its charge moved the live counters but not the committed ones, so every gauge read one charge behind. It now moves the committed counters that already exist, and its own gauge adds its cost to them.

litellm/integrations/prometheus.py

@@ -2284,6 +2284,7 @@ class PrometheusLogger(CustomLogger):
         _org_alias: Final = _metadata.get("user_api_key_org_alias", None)
         _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata)
         _committed_at_auth: Final = committed_budget_spend_from_metadata(_carried_metadata)
+        _row_carries_auth: Final = "user_api_key_auth" in _carried_metadata
 
         # Bound the per-request budget-metric emission so that slow Redis/DB
         # lookups under load cannot consume the whole LoggingWorker watchdog
@@ -2295,9 +2296,20 @@ class PrometheusLogger(CustomLogger):
         budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout()
 
         async def emit_budget_metrics() -> tuple[object, ...]:
+            committed_now: Final = await _committed_spend_now(
+                _committed_at_auth.counter_keys
+                if _row_carries_auth
+                else EntitySpendCounterKeys.of(
+                    token=user_api_key,
+                    team_id=user_api_team,
+                    user_id=user_id,
+                    org_id=user_api_key_org_id,
+                    end_user_id=end_user_id,
+                )
+            )
             committed: Final = _committed_spend_before_request(
-                at_auth=_committed_at_auth,
-                now=await _committed_spend_now(_committed_at_auth.counter_keys),
+                at_auth=_committed_at_auth if _row_carries_auth else committed_now,
+                now=committed_now,
                 response_cost=response_cost,
             )
             return await asyncio.gather(

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -39,6 +39,7 @@ from litellm.proxy.db.db_spend_update_writer import (
 from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route
 from litellm.proxy.route_llm_request import ProxyModelNotFoundError
+from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys
 from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
 from litellm.proxy.spend_tracking.spend_event import (
     ObjectMapping,
@@ -1022,6 +1023,30 @@ class _IncrementSpendCounters(Protocol):
     ) -> None: ...
 
 
+async def _committed_spend_scope_of_row(
+    kwargs: Mapping[str, object],
+    response_cost: float,
+    token: str | None,
+    team_id: str | None,
+    user_id: str | None,
+    org_id: str | None,
+    end_user_id: str | None,
+) -> EntitySpendCounterKeys:
+    metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs)))
+    if "user_api_key_auth" in metadata:
+        return committed_budget_spend_from_metadata(metadata).counter_keys
+    if response_cost <= 0 or not committed_spend_tracked():
+        return EntitySpendCounterKeys()
+    from litellm.proxy.proxy_server import spend_counter_cache
+
+    return await warm_committed_counter_keys(
+        spend_counter_cache,
+        EntitySpendCounterKeys.of(
+            token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id
+        ),
+    )
+
+
 async def _update_database_and_spend_counters(
     proxy_logging_obj: "ProxyLogging",
     increment_spend_counters: _IncrementSpendCounters,
@@ -1153,9 +1178,15 @@ async def _update_database_and_spend_counters_in_batch(
             request_started_at=start_time,
             model_access_groups=model_access_groups,
             project_id=project_id,
-            committed_spend_scope=committed_budget_spend_from_metadata(
-                _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=kwargs))
-            ).counter_keys,
+            committed_spend_scope=await _committed_spend_scope_of_row(
+                kwargs=kwargs,
+                response_cost=response_cost,
+                token=user_api_key,
+                team_id=team_id,
+                user_id=user_id,
+                org_id=org_id,
+                end_user_id=end_user_id,
+            ),
         )
     except Exception:
         if budget_reservation is not None:

litellm/proxy/spend_tracking/committed_spend.py

@@ -108,6 +108,20 @@ async def read_committed_budget_spend(
     return _committed_budget_spend(counter_keys=counter_keys, committed=committed)
 
 
+async def warm_committed_counter_keys(
+    spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys
+) -> EntitySpendCounterKeys:
+    """For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists."""
+    committed: Final = await read_committed_budget_spend(spend_counter_cache, counter_keys)
+    return EntitySpendCounterKeys(
+        key=counter_keys.key if committed.key is not None else None,
+        team=counter_keys.team if committed.team is not None else None,
+        user=counter_keys.user if committed.user is not None else None,
+        org=counter_keys.org if committed.org is not None else None,
+        end_user=counter_keys.end_user if committed.end_user is not None else None,
+    )
+
+
 async def read_admission_committed_spend(
     spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys
 ) -> CommittedBudgetSpend:

tests/unit/integrations/test_prometheus_carried_budget_state.py

@@ -311,6 +311,44 @@ async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted(
     }
 
 
+def _billed_row_metadata() -> dict:
+    return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"}
+
+
+@pytest.mark.asyncio
+async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to(
+    prometheus_logger, getters, spend_counter_cache
+):
+    """Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet."""
+    _stub_gauges(prometheus_logger)
+    _store_committed(spend_counter_cache, key=6.0, team=33.0, user=11.0, org=63.0)
+
+    await _emit(prometheus_logger, _billed_row_metadata())
+
+    assert _remaining(prometheus_logger) == {
+        "key": pytest.approx(2.0),
+        "team": pytest.approx(265.0),
+        "user": pytest.approx(37.0),
+        "org": 435.0,
+    }
+
+
+@pytest.mark.asyncio
+async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend(
+    prometheus_logger, getters, spend_counter_cache
+):
+    _stub_gauges(prometheus_logger)
+
+    await _emit(prometheus_logger, _billed_row_metadata())
+
+    assert _remaining(prometheus_logger) == {
+        "key": pytest.approx(7.0),
+        "team": pytest.approx(278.0),
+        "user": pytest.approx(43.0),
+        "org": 458.0,
+    }
+
+
 @pytest.mark.asyncio
 async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend(
     prometheus_logger, getters, monkeypatch

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -1070,6 +1070,54 @@ async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(m
     assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("tracked", "expected_key", "expected_org"),
+    [(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)],
+)
+async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist(
+    tracked, expected_key, expected_org
+):
+    """A background response is billed from its creator's stored attribution, after the request's auth is gone."""
+    from litellm.proxy.proxy_server import increment_spend_counters
+
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
+    spend_counter_cache = DualCache()
+    spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:key:test_api_key", value=1.0)
+    spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:org:test_org_id", value=2.0)
+
+    with (
+        patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked),
+        patch("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache),
+        patch(
+            "litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db",
+            AsyncMock(return_value=50.0),
+        ),
+    ):
+        await _update_database_and_spend_counters(
+            proxy_logging_obj=proxy_logging_obj,
+            increment_spend_counters=increment_spend_counters,
+            user_api_key="test_api_key",
+            user_id="test_user_id",
+            end_user_id=None,
+            team_id="test_team_id",
+            org_id="test_org_id",
+            kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}},
+            completion_response=None,
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+            response_cost=0.2,
+            budget_reservation=None,
+        )
+
+    cached = spend_counter_cache.in_memory_cache.get_cache
+    assert cached(key="committed_spend:key:test_api_key") == expected_key
+    assert cached(key="committed_spend:org:test_org_id") == expected_org
+    assert cached(key="committed_spend:team:test_team_id") is None
+    assert cached(key="committed_spend:user:test_user_id") is None
+
+
 @pytest.mark.asyncio
 async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails():
     proxy_logging_obj = MagicMock()