From 876bbecf3fe31f2954a5a3d856b92b45efa14200 Mon Sep 17 00:00:00 2001 From: Abhimanyu Kapur <38531241+akapur99@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:02:15 -0700 Subject: [PATCH] fix: move the budget gate off the success callback into the detached task MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit greptile: _key_or_team_is_over_budget awaited get_current_spend inline in async_log_success_event, and that read can fall back to an authoritative DB query — the same request-path DB touch the job-snapshot rework just removed. The gate now runs first inside _run_shadow_eval, where a detached task absorbs the latency and the paid shadow/judge calls still never fire for an over-budget key. The callback is back to zero awaits beyond its own bookkeeping; a regression test asserts get_current_spend is never awaited before the callback returns. Co-Authored-By: Claude --- litellm/integrations/shadow_eval_logger.py | 13 +++- .../integrations/test_shadow_eval_logger.py | 70 +++++++++++-------- 2 files changed, 49 insertions(+), 34 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index d8292574ac5..bd8e53c9892 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -330,8 +330,6 @@ class ShadowEvalLogger(CustomLogger): return # only chat-shaped traffic is comparable if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS: return - if await _key_or_team_is_over_budget(metadata): - return # the shadowed key/team has no budget left for the extra calls raw_messages: Final = kwargs.get("messages") self._inflight_shadow_tasks += 1 task: Final = asyncio.create_task( @@ -347,6 +345,7 @@ class ShadowEvalLogger(CustomLogger): dict(payload.get("model_parameters") or {}) # mutable-ok: frozen snapshot ), parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot + budget_metadata=MappingProxyType(dict(metadata)), # mutable-ok: frozen snapshot ) ) task.add_done_callback(lambda _: setattr(self, "_inflight_shadow_tasks", self._inflight_shadow_tasks - 1)) @@ -451,13 +450,21 @@ class ShadowEvalLogger(CustomLogger): real_model: str, model_parameters: Mapping[str, object], parent_metadata: Mapping[str, object], + budget_metadata: Mapping[str, object] = _EMPTY_METADATA, ) -> None: - """Detached background task: shadow call -> blind judge -> verdict row.""" + """Detached background task: budget gate -> shadow call -> blind judge -> verdict row. + + The budget read lives here, not in the success hook: get_current_spend can fall + back to an authoritative DB read, which a detached task absorbs and the + production callback must not. + """ prisma: Final = self._prisma_provider() try: real_text: Final = self._extract_response_text(response_obj) if not real_text or not messages: return + if await _key_or_team_is_over_budget(budget_metadata): + return shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters, parent_metadata) if shadow is None: diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 78ac8ccb0f4..73022ab15ca 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -528,12 +528,13 @@ class TestKeyOrTeamIsOverBudget: @pytest.mark.asyncio -class TestSuccessHookSkipsWhenOverBudget: - async def test_over_budget_key_is_skipped_before_scheduling_the_shadow_task(self): - job = ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running") - logger, _, router = _logger_with_mocks(job) - logger._run_shadow_eval = AsyncMock() - kwargs = { +class TestBudgetGateRunsInTheDetachedTaskNotTheCallback: + """get_current_spend can fall back to an authoritative DB read; the production + success callback must never await that, so the gate lives in the detached task.""" + + @staticmethod + def _kwargs(spend: float): + return { "standard_logging_object": { "id": "req-1", "model": "gpt-4o", @@ -541,44 +542,51 @@ class TestSuccessHookSkipsWhenOverBudget: "metadata": { "user_api_key_hash": "key-hash", "user_api_key_max_budget": 10.0, - "user_api_key_spend": 10.0, + "user_api_key_spend": spend, }, }, "litellm_params": {"metadata": {}}, "messages": [{"role": "user", "content": "hi"}], } + @staticmethod + def _job(): + return ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running") + + async def test_over_budget_key_never_fires_a_shadow_call(self): + logger, _, router = _logger_with_mocks(self._job()) + response = {"choices": [{"message": {"content": "real answer"}}]} + with patch("litellm.proxy.proxy_server.get_current_spend", AsyncMock(return_value=10.0)): - await logger.async_log_success_event(kwargs, MagicMock(), None, None) - await asyncio.sleep(0) + await logger.async_log_success_event(self._kwargs(spend=10.0), response, None, None) + await asyncio.sleep(0.05) - logger._run_shadow_eval.assert_not_awaited() router.acompletion.assert_not_called() - async def test_under_budget_key_still_schedules_the_shadow_task(self): - job = ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running") - logger, _, router = _logger_with_mocks(job) - logger._run_shadow_eval = AsyncMock() - kwargs = { - "standard_logging_object": { - "id": "req-1", - "model": "gpt-4o", - "call_type": "acompletion", - "metadata": { - "user_api_key_hash": "key-hash", - "user_api_key_max_budget": 10.0, - "user_api_key_spend": 4.0, - }, - }, - "litellm_params": {"metadata": {}}, - "messages": [{"role": "user", "content": "hi"}], - } + async def test_under_budget_key_fires_the_shadow_call(self): + logger, _, router = _logger_with_mocks(self._job()) + router.acompletion = AsyncMock(return_value={"choices": [{"message": {"content": "shadow"}}]}) + response = {"choices": [{"message": {"content": "real answer"}}]} with patch("litellm.proxy.proxy_server.get_current_spend", AsyncMock(return_value=4.0)): - await logger.async_log_success_event(kwargs, MagicMock(), None, None) - await asyncio.sleep(0) + await logger.async_log_success_event(self._kwargs(spend=4.0), response, None, None) + await asyncio.sleep(0.05) - logger._run_shadow_eval.assert_awaited_once() + router.acompletion.assert_awaited_once() + + async def test_the_callback_itself_never_awaits_the_spend_read(self): + """The spend read must happen after the callback returns, inside the task.""" + logger, _, _ = _logger_with_mocks(self._job()) + get_current_spend = AsyncMock(return_value=0.0) + response = {"choices": [{"message": {"content": "real answer"}}]} + + with patch("litellm.proxy.proxy_server.get_current_spend", get_current_spend): + await logger.async_log_success_event(self._kwargs(spend=0.0), response, None, None) + spend_reads_when_callback_returned = get_current_spend.await_count + await asyncio.sleep(0.05) + + assert spend_reads_when_callback_returned == 0 + assert get_current_spend.await_count > 0 @pytest.mark.asyncio