scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

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

Browse source 
* 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>
This commit is contained in:
devin-ai-integration[bot] 2026-09-26 12:22:22 -07:00 • committed by GitHub
parent 9540f19e38
commit 5ac640e49d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 253 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

85
litellm/responses/streaming_iterator.py
View file

@ -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:

172
tests/unit/responses/test_streaming_iterator.py
View file

@ -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