From 5ac640e49d414f33b9d7a59be4a20f150271bc03 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:22:22 -0700 Subject: [PATCH] fix(responses): run stream failure and success hooks on the iterating loop instead of blocking it (#43270) * fix(responses): run stream failure and success hooks on the iterating loop instead of blocking it A dropped provider stream on native /v1/responses ran the failure logging through run_async_function from inside the async iterator, which parks the event loop thread on a helper-thread future until every failure callback returns, and never returns when a callback waits on state only that loop can advance. With a running loop the failure handlers (and the completed-stream success deployment hook) are now scheduled as tasks on it, the way chat streaming already does; the sync iterator keeps its blocking path * fix(responses): await stream failure and success logging on the iterating loop before propagating Keep the merge-base hook set for the native Responses stream: async_failure_handler plus the executor-thread failure_handler on failure, and the post-call success deployment hook on completion. Inside a running loop the async handler is scheduled as a task on that loop and the async iterator awaits it before re-raising, so the loop is never blocked on a foreign-loop future and the failure is attributed before the router's fallback wrapper re-enters the same logging object. The sync iterator inside a running loop keeps the task fire-and-forget with a strong reference. * fix(responses): submit the sync failure handler only after the async one finishes --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 85 ++++++++- .../unit/responses/test_streaming_iterator.py | 172 ++++++++++++++++++ 2 files changed, 253 insertions(+), 4 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f537d24eaa..12bc9adbac8 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -6,7 +6,7 @@ import json 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 @@ class BaseResponsesAPIStreamingIterator: 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 @@ class BaseResponsesAPIStreamingIterator: 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 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): 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: diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 9dbbc20591e..2f6dccb37f3 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -3,7 +3,9 @@ completion_start_time on the first chunk so downstream TTFT consumers (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 @@ from pydantic_core import PydanticSerializationError 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