From d980bbf4c93ea48d98465b83c9bf0219c2ad7d36 Mon Sep 17 00:00:00 2001 From: Aakash Pydi Date: Sun, 27 Sep 2026 16:58:43 -0700 Subject: [PATCH] fix(streaming): schedule text stream failure logging like the stream wrapper Awaiting the handlers let a slow or raising callback delay or replace the provider error. Schedule them the way CustomStreamWrapper already logs stream failures, and have the test wait on the callback instead of a timer --- litellm/llms/openai/completion/handler.py | 5 ++++- tests/unit/llms/openai/completion/test_completion_handler.py | 5 +++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 52bc64e92ec..6dc738afe83 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -1,3 +1,4 @@ +import asyncio import json import traceback from collections.abc import Callable @@ -302,7 +303,9 @@ class OpenAITextCompletion(BaseLLM): try: raw_response: Final = await openai_client.completions.with_raw_response.create(**data) except Exception as e: - await logging_obj.dispatch_failure_handlers(e, traceback.format_exc(), prefer_async_handlers=True) + asyncio.create_task( + logging_obj.dispatch_failure_handlers(e, traceback.format_exc(), prefer_async_handlers=True) + ) raise response: Final = raw_response.parse() streamwrapper: Final = CustomStreamWrapper( diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index 647495121f1..4af8c3ad3c4 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -6,6 +6,8 @@ Regression tests for https://github.com/BerriAI/litellm/issues/27410 """ +import asyncio + import httpx import pytest import respx @@ -94,9 +96,11 @@ async def test_acompletion_forwards_client_headers_to_provider( class _FailureRecorder(CustomLogger): def __init__(self): self.payloads = [] + self.logged = asyncio.Event() async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): self.payloads.append(kwargs["standard_logging_object"]) + self.logged.set() @pytest.mark.parametrize( @@ -120,6 +124,7 @@ async def test_astream_failing_before_first_byte_logs_one_failure(provider_respo with pytest.raises(expected_error): async for _ in response: pass + await asyncio.wait_for(recorder.logged.wait(), timeout=5) assert len(recorder.payloads) == 1 payload = recorder.payloads[0]