fix(responses): run stream failure and success hooks on the iterating loop instead of blocking it

litellm/responses/streaming_iterator.py

@@ -6,7 +6,7 @@
 import time
 import traceback
 import uuid
-from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -169,6 +169,26 @@ def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str)
         verbose_logger.error("%s failed: %s", task_name, exception)
 
 
+_PENDING_LOGGING_TASKS: Final[set[asyncio.Task[object]]] = set()  # mutable-ok: strong refs to pending logging tasks
+
+
+def _running_loop() -> asyncio.AbstractEventLoop | None:
+    try:
+        return asyncio.get_running_loop()
+    except RuntimeError:
+        return None
+
+
+def _spawn_logging_task(
+    running_loop: asyncio.AbstractEventLoop, coroutine: Coroutine[object, object, object], *, task_name: str
+) -> asyncio.Task[object]:
+    task: Final = running_loop.create_task(coroutine)
+    _PENDING_LOGGING_TASKS.add(task)
+    task.add_done_callback(_PENDING_LOGGING_TASKS.discard)
+    task.add_done_callback(lambda done: _log_background_task_failure(done, task_name=task_name))
+    return task
+
+
 _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(
     {
         "server_error": 500,
@@ -275,6 +295,8 @@ class BaseResponsesAPIStreamingIterator:
     This class contains shared logic for both synchronous and asynchronous iterators.
     """
 
+    _pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
+
     def __init__(
         self,
         response: httpx.Response,
@@ -839,8 +861,21 @@ def _run_post_success_hooks(self, end_time: datetime):
             except Exception:
                 typed_call_type = None
 
+        running_loop: Final = _running_loop()
+        if running_loop is not None:
+            self._record_pending_logging_task(
+                _spawn_logging_task(
+                    running_loop,
+                    async_post_call_success_deployment_hook(
+                        request_data=request_payload,
+                        response=self.completed_response,
+                        call_type=typed_call_type,
+                    ),
+                    task_name="Responses stream post-call success hook",
+                )
+            )
+            return
         try:
-            # Call synchronously; async hook will be executed via asyncio.run in a new loop
             run_async_function(
                 async_function=async_post_call_success_deployment_hook,
                 request_data=request_payload,
@@ -861,28 +896,63 @@ def _handle_failure(self, exception: Exception):
         self._failure_handled = True
 
         traceback_exception: Final = traceback.format_exc()
+        end_time: Final = datetime.now()
+        running_loop: Final = _running_loop()
+        if running_loop is not None:
+            self._record_pending_logging_task(
+                _spawn_logging_task(
+                    running_loop,
+                    self._run_failure_handlers_in_order(exception, traceback_exception, end_time),
+                    task_name="Responses stream failure logging",
+                )
+            )
+            return
         try:
             run_async_function(
                 async_function=self.logging_obj.async_failure_handler,
                 exception=exception,
                 traceback_exception=traceback_exception,
                 start_time=self.start_time,
-                end_time=datetime.now(),
+                end_time=end_time,
             )
         except Exception:
             pass
+        self._submit_sync_failure_handler(exception, traceback_exception, end_time)
+
+    async def _run_failure_handlers_in_order(
+        self, exception: Exception, traceback_exception: str, end_time: datetime
+    ) -> None:
+        try:
+            await self.logging_obj.async_failure_handler(
+                exception=exception,
+                traceback_exception=traceback_exception,
+                start_time=self.start_time,
+                end_time=end_time,
+            )
+        finally:
+            self._submit_sync_failure_handler(exception, traceback_exception, end_time)
 
+    def _submit_sync_failure_handler(self, exception: Exception, traceback_exception: str, end_time: datetime) -> None:
         try:
             executor.submit(
                 self.logging_obj.failure_handler,
                 exception,
                 traceback_exception,
                 self.start_time,
-                datetime.now(),
+                end_time,
             )
         except Exception:
             pass
 
+    def _record_pending_logging_task(self, task: asyncio.Task[object]) -> None:
+        self._pending_logging_tasks = (*self._pending_logging_tasks, task)
+
+    async def _await_pending_logging(self) -> None:
+        pending: Final = self._pending_logging_tasks
+        self._pending_logging_tasks = ()
+        if pending:
+            await asyncio.wait(pending)
+
     def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None:
         self._yielded_first_chunk = True
         if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES:
@@ -970,6 +1040,13 @@ def __aiter__(self):
         return self
 
     async def __anext__(self) -> ResponsesAPIStreamingResponse:
+        try:
+            return await self._next_event()
+        except Exception:
+            await self._await_pending_logging()
+            raise
+
+    async def _next_event(self) -> ResponsesAPIStreamingResponse:
         try:
             self._check_max_streaming_duration()
             while True:

tests/unit/responses/test_streaming_iterator.py

@@ -3,7 +3,9 @@
 (Prometheus, OTEL, SpendLogs completionStartTime) do not fall back to
 completion_start_time = end_time."""
 
+import asyncio
 import json
+from collections.abc import Callable
 from datetime import datetime
 from typing import Final, Optional
 from unittest.mock import AsyncMock, Mock, patch
@@ -14,6 +16,7 @@
 
 import litellm
 from litellm.exceptions import MidStreamFallbackError
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
 from litellm.responses.streaming_iterator import (
@@ -417,6 +420,175 @@ def test_sync_complete_stream_still_ends_normally(trailer):
     assert logging_obj.async_failure_handler.await_count == 0
 
 
+class _LoopRecordingLogger(CustomLogger):
+    def __init__(self) -> None:
+        super().__init__()
+        self.failure_loop: asyncio.AbstractEventLoop | None = None
+        self.failure_deployment_id: str | None = None
+        self.failure_finished = False
+        self.hook_loop: asyncio.AbstractEventLoop | None = None
+        self.hook_finished = False
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self.failure_loop = asyncio.get_running_loop()
+        self.failure_deployment_id = kwargs["litellm_params"].get("model_info", {}).get("id")
+        await asyncio.sleep(0.05)
+        self.failure_finished = True
+
+    async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
+        self.hook_loop = asyncio.get_running_loop()
+        await asyncio.sleep(0.05)
+        self.hook_finished = True
+        return None
+
+
+class _SyncOnlyRecordingLogger(CustomLogger):
+    def __init__(self) -> None:
+        super().__init__()
+        self.sync_failure_finished = False
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self.sync_failure_finished = True
+
+
+class _OrderRecordingSyncLogger(CustomLogger):
+    def __init__(self, async_recorder: _LoopRecordingLogger) -> None:
+        super().__init__()
+        self._async_recorder: Final = async_recorder
+        self.async_failure_finished_first: bool | None = None
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self.async_failure_finished_first = self._async_recorder.failure_finished
+
+
+def _real_logging_obj(
+    *, call_type: str = "aresponses", litellm_params: dict[str, object] | None = None
+) -> LiteLLMLoggingObj:
+    logging_obj: Final = LiteLLMLoggingObj(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type=call_type,
+        start_time=datetime.now(),
+        litellm_call_id="lit-8678-test",
+        function_id="lit-8678-test",
+    )
+    logging_obj.model_call_details["litellm_params"] = (
+        dict(litellm_params) if litellm_params is not None else {"aresponses": True}
+    )
+    return logging_obj
+
+
+async def _wait_until(condition: Callable[[], bool]) -> None:
+    for _ in range(200):
+        if condition():
+            return
+        await asyncio.sleep(0.01)
+    raise AssertionError("condition never became true")
+
+
+@pytest.mark.asyncio
+async def test_transport_error_failure_logging_runs_on_the_iterating_loop(monkeypatch):
+    """LIT-8678: a stream failure used to run async_failure_handler on a helper loop in a
+    worker thread and block the iterating loop until it finished, so a callback waiting on
+    state bound to that loop (a batch logger's flush lock) stalled the whole proxy."""
+    recorder: Final = _LoopRecordingLogger()
+    monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
+    monkeypatch.setattr(litellm, "failure_callback", [])
+    iterator: Final = _make_iterator(
+        sse_events=_PARTIAL_OUTPUT_EVENTS,
+        logging_obj=_real_logging_obj(),
+        trailing_error=httpx.ReadError("Response payload is not completed"),
+    )
+
+    with pytest.raises(httpx.ReadError):
+        async for _ in iterator:
+            pass
+
+    assert recorder.failure_finished is True
+    assert recorder.failure_loop is asyncio.get_running_loop()
+
+
+@pytest.mark.asyncio
+async def test_failure_logging_finishes_before_the_error_reaches_the_consumer(monkeypatch):
+    """The router's mid-stream fallback re-enters the same logging object for the next
+    deployment as soon as it catches the error, so failure logging that still runs after
+    the raise reads the fallback deployment's params and cools down the wrong deployment."""
+    recorder: Final = _LoopRecordingLogger()
+    monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
+    monkeypatch.setattr(litellm, "failure_callback", [])
+    logging_obj: Final = _real_logging_obj(
+        litellm_params={"aresponses": True, "model_info": {"id": "primary-deployment"}}
+    )
+    iterator: Final = _make_iterator(
+        sse_events=[],
+        logging_obj=logging_obj,
+        trailing_error=httpx.ReadError("Response payload is not completed"),
+    )
+
+    with pytest.raises(MidStreamFallbackError):
+        async for _ in iterator:
+            pass
+    logging_obj.model_call_details["litellm_params"]["model_info"] = {"id": "fallback-deployment"}
+    await _wait_until(lambda: recorder.failure_finished)
+
+    assert recorder.failure_deployment_id == "primary-deployment"
+
+
+@pytest.mark.asyncio
+async def test_sync_stream_failure_inside_a_running_loop_still_runs_sync_only_callbacks(monkeypatch):
+    recorder: Final = _SyncOnlyRecordingLogger()
+    monkeypatch.setattr(litellm, "failure_callback", [recorder])
+    monkeypatch.setattr(litellm, "_async_failure_callback", [])
+    iterator: Final = _make_sync_iterator(
+        sse_events=_PARTIAL_OUTPUT_EVENTS,
+        logging_obj=_real_logging_obj(call_type="responses", litellm_params={}),
+        trailing_error=httpx.ReadError("Response payload is not completed"),
+    )
+
+    with pytest.raises(httpx.ReadError):
+        for _ in iterator:
+            pass
+
+    await _wait_until(lambda: recorder.sync_failure_finished)
+
+
+@pytest.mark.asyncio
+async def test_sync_failure_callbacks_run_after_async_failure_logging_finishes(monkeypatch):
+    """Both handlers read the same logging object, so the sync one must not start while the
+    async one is still running, which is the ordering the blocking dispatch used to give."""
+    async_recorder: Final = _LoopRecordingLogger()
+    sync_recorder: Final = _OrderRecordingSyncLogger(async_recorder)
+    monkeypatch.setattr(litellm, "_async_failure_callback", [async_recorder])
+    monkeypatch.setattr(litellm, "failure_callback", [sync_recorder])
+    iterator: Final = _make_sync_iterator(
+        sse_events=_PARTIAL_OUTPUT_EVENTS,
+        logging_obj=_real_logging_obj(call_type="responses", litellm_params={}),
+        trailing_error=httpx.ReadError("Response payload is not completed"),
+    )
+
+    with pytest.raises(httpx.ReadError):
+        for _ in iterator:
+            pass
+
+    await _wait_until(lambda: sync_recorder.async_failure_finished_first is not None)
+
+    assert sync_recorder.async_failure_finished_first is True
+
+
+@pytest.mark.asyncio
+async def test_completed_stream_success_deployment_hook_runs_on_the_iterating_loop(monkeypatch):
+    recorder: Final = _LoopRecordingLogger()
+    monkeypatch.setattr(litellm, "callbacks", [recorder])
+    iterator: Final = _make_iterator(sse_events=_COMPLETE_STREAM_EVENTS, logging_obj=_logging_obj_stub())
+
+    async for _ in iterator:
+        pass
+
+    assert recorder.hook_finished is True
+    assert recorder.hook_loop is asyncio.get_running_loop()
+
+
 def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
     """
     Regression test for LIT-6184 on the /v1/responses streaming surface: the