From d92438c4ceac681046ec17ff534571bd1076d881 Mon Sep 17 00:00:00 2001 From: Ivana Dunisijevic Date: Wed, 7 Oct 2026 22:22:42 +0200 Subject: [PATCH 1/4] fix(proxy): retain deferred stream logging --- litellm/proxy/utils.py | 4 ++- .../test_deferred_guardrail_logging.py | 34 ++++++++----------- .../proxy_logging/test_streaming_hooks.py | 20 ++++++----- 3 files changed, 29 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3e73e2dee40..61c08bd54e9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4184,6 +4184,8 @@ class ProxyLogging: logging_obj: Final = request_data.get("litellm_logging_obj") if logging_obj is None: return + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( logging_obj, "_on_deferred_stream_complete", None ) @@ -4191,7 +4193,7 @@ class ProxyLogging: if _deferred_cb is not None and _args is not None: logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - asyncio.create_task(_deferred_cb(*_args)) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_deferred_cb(*_args)) _fire_deferred_stream_logging = fire_deferred_stream_logging diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 17427e839b9..d1369f92c9f 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -36,6 +36,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import ProxyLogging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.guardrails import GuardrailEventHooks # --------------------------------------------------------------------------- @@ -736,9 +737,7 @@ class TestDeferredStreamingClosure: # CSW stored args; now simulate what ProxyLogging does request_data = {"litellm_logging_obj": mock_logging_obj} ProxyLogging._fire_deferred_stream_logging(request_data) - - await asyncio.sleep(0) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert guardrail_called is True, "Guardrail hook should be called" assert ( @@ -948,9 +947,7 @@ class TestDeferredStreamingClosure: # CSW stored args; now simulate what ProxyLogging does request_data = {"litellm_logging_obj": mock_logging_obj} ProxyLogging._fire_deferred_stream_logging(request_data) - - await asyncio.sleep(0) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert hook_called is True, "Production closure must call guardrail hook" assert ( @@ -1320,8 +1317,7 @@ class TestFireDeferredStreamLogging: request_data = {"litellm_logging_obj": mock_logging_obj} ProxyLogging._fire_deferred_stream_logging(request_data) - - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert callback_called is True assert callback_args["response"] == "test_response" @@ -1521,7 +1517,7 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, False) ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert recorded["result"] is assembled assert recorded["cache_hit"] is False @@ -1555,7 +1551,7 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, False) ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert recorded["result"] is assembled assert recorded["cache_hit"] is False @@ -1563,8 +1559,6 @@ class TestArmDeferredStreamDispatch: @pytest.mark.asyncio async def test_native_stream_closure_enqueues_single_coroutine(self): - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): @@ -1587,7 +1581,9 @@ class TestArmDeferredStreamDispatch: GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" ) as mock_enqueue: ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + deferred = mock_enqueue.call_args.kwargs["async_coroutine"] + mock_enqueue.reset_mock() + await deferred mock_enqueue.assert_called_once_with(async_coroutine=coro) assert recorded == {} coro.close() @@ -1600,8 +1596,6 @@ class TestArmDeferredStreamDispatch: so stores (assembled_response, cache_hit). The closure armed for a raw generator must accept that shape too, or _fire_deferred_stream_logging raises TypeError and the request loses its spend log and callbacks.""" - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): @@ -1620,7 +1614,9 @@ class TestArmDeferredStreamDispatch: GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" ) as mock_enqueue: ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + deferred = mock_enqueue.call_args.kwargs["async_coroutine"] + mock_enqueue.reset_mock() + await deferred mock_enqueue.assert_not_called() assert recorded["result"] is assembled @@ -1630,8 +1626,6 @@ class TestArmDeferredStreamDispatch: @pytest.mark.asyncio @pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())]) async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog): - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): @@ -1652,7 +1646,9 @@ class TestArmDeferredStreamDispatch: caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + deferred = mock_enqueue.call_args.kwargs["async_coroutine"] + mock_enqueue.reset_mock() + await deferred mock_enqueue.assert_not_called() assert recorded == {} diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index 50f50478ad3..f462340a808 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -13,7 +13,7 @@ import asyncio from collections.abc import AsyncGenerator, AsyncIterator from datetime import datetime from typing import Any, Dict, Final, List -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException @@ -23,6 +23,7 @@ from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) @@ -563,8 +564,7 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea request_data=request_data, ): pass - await asyncio.sleep(0) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)] @@ -599,8 +599,7 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ request_data=request_data, ): pass - await asyncio.sleep(0) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() assert [event[0] for event in events] == ["logging_dispatched"] assert logging_obj._deferred_stream_complete_args is None @@ -724,8 +723,7 @@ async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes request_data=request_data, ): pass - await asyncio.sleep(0) - await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() snapshot = { "events": events, @@ -751,8 +749,12 @@ async def test_fire_deferred_stream_logging_fires_callback(): logging_obj._on_deferred_stream_complete = deferred logging_obj._deferred_stream_complete_args = ("payload",) - ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue: + ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) + + enqueue.assert_called_once() + assert captured == {} + await enqueue.call_args.kwargs["async_coroutine"] snapshot = { "arg": captured["arg"], "callback_cleared": logging_obj._on_deferred_stream_complete is None, From 307bb06a0514506d3ac8d489c02a8b5e186eb4e6 Mon Sep 17 00:00:00 2001 From: Ivana Dunisijevic Date: Thu, 8 Oct 2026 17:08:32 +0200 Subject: [PATCH 2/4] fix(proxy): retain deferred stream callbacks --- litellm/proxy/common_request_processing.py | 14 ++- litellm/proxy/utils.py | 23 ++++- .../test_deferred_guardrail_logging.py | 92 +++++++++++++------ .../proxy_logging/test_streaming_hooks.py | 26 +++--- 4 files changed, 104 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 317203cc2d5..3b0888b9c0e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3630,14 +3630,12 @@ class ProxyBaseLLMRequestProcessing: # callbacks fire regardless of the call-type heuristic in # _is_sync_litellm_request (which only recognizes a subset of # async markers stored in litellm_params). - asyncio.create_task( - _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers( - _response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + await _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers( + _response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, ) except Exception as e: verbose_proxy_logger.exception( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 61c08bd54e9..185b7d4f51d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1103,6 +1103,21 @@ def _call_type_for_route(route: str | None) -> str | None: _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data")) +_DEFERRED_STREAM_CALLBACK_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: retains background tasks + + +async def _run_deferred_stream_callback(callback: Coroutine[object, object, object]) -> None: + try: + await callback + except Exception as e: + verbose_proxy_logger.exception("Error in deferred stream callback: %s", e) + + +def _retain_deferred_stream_callback(callback: Coroutine[object, object, object]) -> asyncio.Task[None]: + task: Final = asyncio.create_task(_run_deferred_stream_callback(callback)) + _DEFERRED_STREAM_CALLBACK_TASKS.add(task) + task.add_done_callback(_DEFERRED_STREAM_CALLBACK_TASKS.discard) + return task def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: @@ -4172,7 +4187,7 @@ class ProxyLogging: record_served_output_texts(logging_obj.model_call_details, served_stream_output_texts(served_chunks)) @staticmethod - def fire_deferred_stream_logging(request_data: dict) -> None: + def fire_deferred_stream_logging(request_data: dict) -> asyncio.Task[None] | None: """ Fire the deferred streaming logging callback after the full streaming pipeline (including guardrail end-of-stream blocks) has completed. @@ -4183,8 +4198,7 @@ class ProxyLogging: """ logging_obj: Final = request_data.get("litellm_logging_obj") if logging_obj is None: - return - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + return None _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( logging_obj, "_on_deferred_stream_complete", None @@ -4193,7 +4207,8 @@ class ProxyLogging: if _deferred_cb is not None and _args is not None: logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_deferred_cb(*_args)) + return _retain_deferred_stream_callback(_deferred_cb(*_args)) + return None _fire_deferred_stream_logging = fire_deferred_stream_logging diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index d1369f92c9f..e4e1f6e3151 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -736,8 +736,9 @@ class TestDeferredStreamingClosure: # CSW stored args; now simulate what ProxyLogging does request_data = {"litellm_logging_obj": mock_logging_obj} - ProxyLogging._fire_deferred_stream_logging(request_data) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data) + assert deferred_task is not None + await deferred_task assert guardrail_called is True, "Guardrail hook should be called" assert ( @@ -946,8 +947,9 @@ class TestDeferredStreamingClosure: # CSW stored args; now simulate what ProxyLogging does request_data = {"litellm_logging_obj": mock_logging_obj} - ProxyLogging._fire_deferred_stream_logging(request_data) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data) + assert deferred_task is not None + await deferred_task assert hook_called is True, "Production closure must call guardrail hook" assert ( @@ -1316,9 +1318,12 @@ class TestFireDeferredStreamLogging: mock_logging_obj._deferred_stream_complete_args = ("test_response", True) request_data = {"litellm_logging_obj": mock_logging_obj} - ProxyLogging._fire_deferred_stream_logging(request_data) - await GLOBAL_LOGGING_WORKER.flush() + with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue: + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data) + assert deferred_task is not None + await deferred_task + enqueue.assert_not_called() assert callback_called is True assert callback_args["response"] == "test_response" assert callback_args["cache_hit"] is True @@ -1326,6 +1331,34 @@ class TestFireDeferredStreamLogging: assert mock_logging_obj._on_deferred_stream_complete is None assert mock_logging_obj._deferred_stream_complete_args is None + @pytest.mark.asyncio + async def test_deferred_guardrails_wait_for_success_logging(self): + logging_started = asyncio.Event() + release_logging = asyncio.Event() + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + logging_started.set() + await release_logging.wait() + + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) + + with patch("litellm.callbacks", []): + deferred_task = asyncio.create_task( + ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + ) + await logging_started.wait() + assert deferred_task.done() is False + release_logging.set() + await deferred_task + @pytest.mark.asyncio async def test_noop_when_no_deferred_args(self): """_fire_deferred_stream_logging should be a no-op when no deferred @@ -1516,8 +1549,9 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, False) - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task assert recorded["result"] is assembled assert recorded["cache_hit"] is False @@ -1550,8 +1584,9 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, False) - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task assert recorded["result"] is assembled assert recorded["cache_hit"] is False @@ -1577,13 +1612,14 @@ class TestArmDeferredStreamDispatch: coro = _logging_coroutine() logging_obj._deferred_stream_complete_args = (coro,) - with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam - GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" - ) as mock_enqueue: - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - deferred = mock_enqueue.call_args.kwargs["async_coroutine"] - mock_enqueue.reset_mock() - await deferred + with ( + patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue + ): + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task mock_enqueue.assert_called_once_with(async_coroutine=coro) assert recorded == {} coro.close() @@ -1610,13 +1646,14 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, True) - with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam - GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" - ) as mock_enqueue: - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - deferred = mock_enqueue.call_args.kwargs["async_coroutine"] - mock_enqueue.reset_mock() - await deferred + with ( + patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue + ): + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task mock_enqueue.assert_not_called() assert recorded["result"] is assembled @@ -1645,10 +1682,9 @@ class TestArmDeferredStreamDispatch: ) as mock_enqueue, caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - deferred = mock_enqueue.call_args.kwargs["async_coroutine"] - mock_enqueue.reset_mock() - await deferred + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task mock_enqueue.assert_not_called() assert recorded == {} diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index f462340a808..86429080997 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,10 +10,10 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime from typing import Any, Dict, Final, List -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException @@ -23,7 +23,6 @@ from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) @@ -180,6 +179,13 @@ async def _one_chunk() -> AsyncGenerator[object, None]: yield "chunk" +async def _wait_until(condition: Callable[[], bool]) -> None: + for _ in range(200): + if condition(): + return + await asyncio.sleep(0.01) + + class _AttributeStream: _hidden_params = {"model_id": "m-1"} model = "gpt-x" @@ -564,7 +570,7 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea request_data=request_data, ): pass - await GLOBAL_LOGGING_WORKER.flush() + await _wait_until(lambda: len(events) == 2) assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)] @@ -599,7 +605,7 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ request_data=request_data, ): pass - await GLOBAL_LOGGING_WORKER.flush() + await _wait_until(lambda: len(events) == 1) assert [event[0] for event in events] == ["logging_dispatched"] assert logging_obj._deferred_stream_complete_args is None @@ -723,7 +729,7 @@ async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes request_data=request_data, ): pass - await GLOBAL_LOGGING_WORKER.flush() + await _wait_until(lambda: len(events) == 1) snapshot = { "events": events, @@ -749,12 +755,10 @@ async def test_fire_deferred_stream_logging_fires_callback(): logging_obj._on_deferred_stream_complete = deferred logging_obj._deferred_stream_complete_args = ("payload",) - with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue: - ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) - - enqueue.assert_called_once() + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) + assert deferred_task is not None assert captured == {} - await enqueue.call_args.kwargs["async_coroutine"] + await deferred_task snapshot = { "arg": captured["arg"], "callback_cleared": logging_obj._on_deferred_stream_complete is None, From 8bee832c4ee4fb53f8730aaa20a8cc69b6442287 Mon Sep 17 00:00:00 2001 From: Ivana Dunisijevic Date: Thu, 8 Oct 2026 17:15:24 +0200 Subject: [PATCH 3/4] fix(proxy): annotate deferred callback boundary --- litellm/proxy/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 185b7d4f51d..97d01743b7f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1109,7 +1109,7 @@ _DEFERRED_STREAM_CALLBACK_TASKS: Final[set[asyncio.Task[None]]] = set() # mutab async def _run_deferred_stream_callback(callback: Coroutine[object, object, object]) -> None: try: await callback - except Exception as e: + except Exception as e: # noqa: BLE001 # background callback failures must not escape verbose_proxy_logger.exception("Error in deferred stream callback: %s", e) From aedc0e3aecd1201054fb0ed5b50a716d16acf2eb Mon Sep 17 00:00:00 2001 From: Ivana Dunisijevic Date: Thu, 8 Oct 2026 17:47:49 +0200 Subject: [PATCH 4/4] test(proxy): cover deferred callback failures --- .../test_deferred_guardrail_logging.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index e4e1f6e3151..b08a5b310fa 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1331,6 +1331,25 @@ class TestFireDeferredStreamLogging: assert mock_logging_obj._on_deferred_stream_complete is None assert mock_logging_obj._deferred_stream_complete_args is None + @pytest.mark.asyncio + async def test_callback_failure_is_logged_without_escaping(self): + """A failed background callback must not surface as an unhandled task error.""" + callback_error = RuntimeError("deferred callback failed") + + async def failing_callback(): + raise callback_error + + mock_logging_obj = MagicMock() + mock_logging_obj._on_deferred_stream_complete = failing_callback + mock_logging_obj._deferred_stream_complete_args = () + + with patch("litellm.proxy.utils.verbose_proxy_logger.exception") as log_exception: + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": mock_logging_obj}) + assert deferred_task is not None + await deferred_task + + log_exception.assert_called_once_with("Error in deferred stream callback: %s", callback_error) + @pytest.mark.asyncio async def test_deferred_guardrails_wait_for_success_logging(self): logging_started = asyncio.Event()