From f43b4db2c1cf33d9482c5ba1ce68942a380ed806 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:53:02 -0700 Subject: [PATCH] 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 | 16 ++++++- .../proxy/hooks/proxy_track_cost_callback.py | 37 ++++++++++++-- .../proxy/spend_tracking/committed_spend.py | 14 ++++++ .../test_prometheus_carried_budget_state.py | 38 +++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 48 +++++++++++++++++++ 5 files changed, 148 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2636d92a7cf..d71d0f8d403 100644 --- a/litellm/integrations/prometheus.py +++ b/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( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 97018b74b95..7f95121ea23 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/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: diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py index 671662fa8c0..0f542040d8e 100644 --- a/litellm/proxy/spend_tracking/committed_spend.py +++ b/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: diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 74525da4c3e..1a746de588e 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/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 diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c3f4c4e9d97..c736f92da6d 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/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()