Merge 32026a6272 into b6f91a221b

litellm/integrations/prometheus.py

@@ -4302,12 +4302,6 @@ class PrometheusLogger(CustomLogger):
                 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,
-            spend=_total_team_spend,
-            max_budget=max_budget,
-        )
         try:
             team_info: Final = await get_team_object(
                 team_id=team_id,
@@ -4316,14 +4310,27 @@ class PrometheusLogger(CustomLogger):
             )
         except Exception as e:
             verbose_logger.debug("[Non-Blocking] Prometheus: Error getting team info: %s", e)
-            return team_object
+            return LiteLLM_TeamTable(
+                team_id=team_id,
+                team_alias=team_alias,
+                spend=_total_team_spend,
+                max_budget=max_budget,
+            )
 
-        if team_info:
-            team_object.budget_reset_at = team_info.budget_reset_at
-            if team_object.max_budget is None and team_info.max_budget is not None:
-                team_object.max_budget = team_info.max_budget
-
-        return team_object
+        # metadata spend of None means "unknown", not "nothing spent" -- fall back to the DB value
+        _base_spend: Final = (
+            spend if spend is not None else (team_info.spend if team_info and team_info.spend is not None else 0)
+        )
+        _resolved_max_budget: Final = (
+            max_budget if max_budget is not None else (team_info.max_budget if team_info else None)
+        )
+        return LiteLLM_TeamTable(
+            team_id=team_id,
+            team_alias=team_alias,
+            spend=_base_spend + response_cost,
+            max_budget=_resolved_max_budget,
+            budget_reset_at=team_info.budget_reset_at if team_info else None,
+        )
 
     def _set_team_budget_metrics(
         self,

tests/unit/integrations/test_prometheus_user_team_metrics.py

@@ -338,6 +338,61 @@ async def test_assemble_team_object_does_not_override_metadata_max_budget(
     ), "max_budget from metadata must not be replaced by the DB value"
 
 
+async def test_assemble_team_object_uses_db_spend_when_metadata_is_none(
+    prometheus_logger,
+):
+    """
+    When spend is None in request metadata (cached team object not hydrated),
+    _assemble_team_object must backfill it from get_team_object, otherwise the
+    remaining-budget gauge reports the full max_budget as if nothing was spent.
+    """
+    db_team = MagicMock()
+    db_team.max_budget = 1000.0
+    db_team.spend = 902.60
+    db_team.budget_reset_at = None
+
+    with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team:  # test-quality-ok: stubbing the DB fetch, same pattern as the sibling max_budget fallback tests above
+        mock_get_team.return_value = db_team
+        team_object = await prometheus_logger._assemble_team_object(
+            team_id="team-1",
+            team_alias="team-a",
+            spend=None,
+            max_budget=1000.0,
+            response_cost=0.5,
+        )
+
+    assert (
+        team_object.spend == 902.60 + 0.5
+    ), "spend should be populated from DB (plus request cost) when metadata value is None"
+
+
+async def test_assemble_team_object_does_not_override_metadata_spend(
+    prometheus_logger,
+):
+    """
+    When spend IS present in request metadata, it must not be replaced by the
+    DB value.
+    """
+    db_team = MagicMock()
+    db_team.max_budget = 1000.0
+    db_team.spend = 9999.0
+    db_team.budget_reset_at = None
+
+    with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team:  # test-quality-ok: stubbing the DB fetch, same pattern as the sibling max_budget fallback tests above
+        mock_get_team.return_value = db_team
+        team_object = await prometheus_logger._assemble_team_object(
+            team_id="team-1",
+            team_alias="team-a",
+            spend=50.0,
+            max_budget=1000.0,
+            response_cost=1.0,
+        )
+
+    assert (
+        team_object.spend == 50.0 + 1.0
+    ), "spend from metadata must not be replaced by the DB value"
+
+
 async def test_set_team_budget_metrics_after_api_request_no_inf_when_metadata_budget_none(
     prometheus_logger,
 ):