From 5fc869c91583af627a60d31657695afa3dbac76b Mon Sep 17 00:00:00 2001 From: Aakash Pydi Date: Sun, 27 Sep 2026 16:46:06 -0700 Subject: [PATCH] fix(streaming): await failure logging before re-raising a failed text stream Await the failure handlers instead of scheduling them, matching the non-streaming path, so the failure is logged even if the event loop shuts down right after. The test now checks the failure payload and no longer waits on a timer --- litellm/llms/openai/completion/handler.py | 5 +---- .../completion/test_completion_handler.py | 21 ++++++++++--------- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 6dc738afe83..52bc64e92ec 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -1,4 +1,3 @@ -import asyncio import json import traceback from collections.abc import Callable @@ -303,9 +302,7 @@ class OpenAITextCompletion(BaseLLM): try: raw_response: Final = await openai_client.completions.with_raw_response.create(**data) except Exception as e: - asyncio.create_task( - logging_obj.dispatch_failure_handlers(e, traceback.format_exc(), prefer_async_handlers=True) - ) + await 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 6850c6a9cca..647495121f1 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -6,8 +6,6 @@ Regression tests for https://github.com/BerriAI/litellm/issues/27410 """ -import asyncio - import httpx import pytest import respx @@ -93,12 +91,12 @@ async def test_acompletion_forwards_client_headers_to_provider( request_headers = mock_completions_endpoint.calls.last.request.headers assert request_headers["x-mycorp-llmcall-id"] == "abc-123" -class _FailureCounter(CustomLogger): +class _FailureRecorder(CustomLogger): def __init__(self): - self.failures = 0 + self.payloads = [] async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - self.failures += 1 + self.payloads.append(kwargs["standard_logging_object"]) @pytest.mark.parametrize( @@ -111,10 +109,10 @@ class _FailureCounter(CustomLogger): @respx.mock async def test_astream_failing_before_first_byte_logs_one_failure(provider_response, expected_error, monkeypatch): respx.post("https://api.openai.com/v1/completions").mock(side_effect=provider_response) - counter = _FailureCounter() + recorder = _FailureRecorder() monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setattr(litellm, "callbacks", [counter]) - monkeypatch.setattr(litellm, "_async_failure_callback", [counter]) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) response = await atext_completion( model="gpt-3.5-turbo-instruct", prompt="hello", max_tokens=5, stream=True, max_retries=0 @@ -122,6 +120,9 @@ 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.sleep(0.5) - assert counter.failures == 1 + assert len(recorder.payloads) == 1 + payload = recorder.payloads[0] + assert payload["status"] == "failure" + assert payload["custom_llm_provider"] == "text-completion-openai" + assert payload["model"] == "gpt-3.5-turbo-instruct"