diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 74ee87c772b..b1872ad8212 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4055,14 +4055,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 29749afb429..1a6d1f6684c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1159,6 +1159,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: # noqa: BLE001 # background callback failures must not escape + 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]: @@ -4420,7 +4435,7 @@ class ProxyLogging: record_served_output_texts(logging_obj.model_call_details, texts) @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. @@ -4431,7 +4446,8 @@ class ProxyLogging: """ logging_obj: Final = request_data.get("litellm_logging_obj") if logging_obj is None: - return + return None + _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( logging_obj, "_on_deferred_stream_complete", None ) @@ -4439,7 +4455,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 - asyncio.create_task(_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 17427e839b9..b08a5b310fa 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 # --------------------------------------------------------------------------- @@ -735,10 +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 asyncio.sleep(0) - await asyncio.sleep(0) + 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 ( @@ -947,10 +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 asyncio.sleep(0) - await asyncio.sleep(0) + 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 ( @@ -1319,10 +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 asyncio.sleep(0) + 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 @@ -1330,6 +1331,53 @@ 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() + 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 @@ -1520,8 +1568,9 @@ 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) + 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 @@ -1554,8 +1603,9 @@ 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) + 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 @@ -1563,8 +1613,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(): @@ -1583,11 +1631,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}) - await asyncio.sleep(0) + 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() @@ -1600,8 +1651,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(): @@ -1616,11 +1665,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}) - await asyncio.sleep(0) + 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 @@ -1630,8 +1682,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(): @@ -1651,8 +1701,9 @@ class TestArmDeferredStreamDispatch: ) as mock_enqueue, caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + 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 50f50478ad3..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,7 +10,7 @@ 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 @@ -179,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" @@ -563,8 +570,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 _wait_until(lambda: len(events) == 2) assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)] @@ -599,8 +605,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 _wait_until(lambda: len(events) == 1) assert [event[0] for event in events] == ["logging_dispatched"] assert logging_obj._deferred_stream_complete_args is None @@ -724,8 +729,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 _wait_until(lambda: len(events) == 1) snapshot = { "events": events, @@ -751,8 +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",) - ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) - await asyncio.sleep(0) + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + assert captured == {} + await deferred_task snapshot = { "arg": captured["arg"], "callback_cleared": logging_obj._on_deferred_stream_complete is None,