From 0ce16cd68d55480125930f39c115a91ff5b07041 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:47:51 -0700 Subject: [PATCH] fix(prometheus): read the committed org spend on failed requests without the auth object --- litellm/integrations/prometheus.py | 14 ++++++++++---- .../test_prometheus_carried_budget_state.py | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1f3b208a460..2636d92a7cf 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -377,6 +377,12 @@ def _customer_spend_after_request(cached_spend: float, committed_spend: float | return max(committed_spend + response_cost, cached_spend) +def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys: + if org is None or org.max_budget is None: + return EntitySpendCounterKeys() + return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None) + + async def _committed_spend_now(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: if not counter_keys.tracked(): return CommittedBudgetSpend() @@ -2862,16 +2868,16 @@ class PrometheusLogger(CustomLogger): get_litellm_metadata_from_kwargs(kwargs) ) failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias") - failure_at_auth: Final = committed_budget_spend_from_metadata(failure_metadata) + failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata) failure_committed: Final = _committed_spend_before_request( - at_auth=failure_at_auth, - now=await _committed_spend_now(EntitySpendCounterKeys(org=failure_at_auth.counter_keys.org)), + at_auth=committed_budget_spend_from_metadata(failure_metadata), + now=await _committed_spend_now(_budgeted_org_counter(user_api_key_org_id, failure_org)), response_cost=0, ) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, response_cost=0, - carried=OrgBudgetSnapshot.from_metadata(failure_metadata), + carried=failure_org, org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None, committed_spend=failure_committed.org, ) diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index b7738a62a00..2886ff39586 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -434,6 +434,20 @@ async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_commit prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) +@pytest.mark.asyncio +async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=3.0, team=30.0, user=8.0, org=65.0) + metadata = _request_metadata(_token_committed_at_auth()) + del metadata["user_api_key_auth"] + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + @pytest.mark.asyncio async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): _stub_gauges(prometheus_logger)