From ba550f397092346b1ff72c76591d3066a71ba177 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:58:08 -0800 Subject: [PATCH] fix(streaming): log the streamed part's usage when a stream is closed early close(), aclose() and both context managers on CustomStreamWrapper now log one partial success built from the chunks already streamed, so an SDK app that closes a stream when its client hangs up keeps that usage and cost. Nothing extra logs when the stream already ended, failed, was never read, when streaming logging is disabled, or when the proxy already logged a failure for it. --- litellm/litellm_core_utils/litellm_logging.py | 4 + .../litellm_core_utils/streaming_handler.py | 100 +++- .../test_streaming_handler.py | 465 +++++++++++++++++- 3 files changed, 552 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ce57a9bd4a5..7456a72478f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2803,6 +2803,10 @@ class Logging(LiteLLMLoggingBaseClass): if self.cost_breakdown is not None: self.cost_breakdown = _with_failed_attempts_cost(self.cost_breakdown, own_cost, failed_attempts_cost) + @property + def stream_failure_claimed(self) -> bool: + return self._stream_failure_claimed + def claim_stream_failure_report(self) -> bool: if self._stream_failure_claimed: return False diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index aa6ce6ea99e..4a956d7226e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -341,6 +341,21 @@ class _VertexChunkLike(Protocol): candidates: Sequence[_VertexCandidateLike] +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _DispatchesSuccessHandlers: + return logging_obj + + class _ParsedChunkHiddenParams(LiteLLMBaseModel): provider_specific_fields: Mapping[str, object] | None = None upstream_provider: str | None = None @@ -429,6 +444,8 @@ class CustomStreamWrapper: self.sent_last_chunk = False self._openai_handler_receives_chat_chunks = False self._failure_logging: asyncio.Task[None] | None = None + self._outcome_logged = False + self._iterated_provider_stream = False self._closed = False self._stream_created_time: float = time.time() @@ -613,6 +630,54 @@ class CustomStreamWrapper: self._restore_consumer_correlation_context(guarded=True) def close(self) -> None: + self._log_abandoned_stream() + self._release() + + async def aclose(self) -> None: + with anyio.CancelScope(shield=True): + await self._alog_abandoned_stream() + await self._arelease() + + def _claim_abandoned_stream_log(self) -> bool: + if ( + self._outcome_logged + or not self._iterated_provider_stream + or litellm.disable_streaming_logging is True + or self.logging_obj.stream_failure_claimed + ): + return False + self._outcome_logged = True + return True + + def _log_abandoned_stream(self) -> None: + if self._claim_abandoned_stream_log(): + executor.submit(self._log_streamed_part) + + def _log_streamed_part(self) -> None: + partial: Final = self.partial_response() + if partial is not None: + self.logging_obj.success_handler(partial, None, None, self._is_cache_hit()) + + async def _alog_abandoned_stream(self) -> None: + if not self._claim_abandoned_stream_log(): + return + try: + partial: Final = await asyncify(self.partial_response)() + if partial is not None: + await _success_dispatcher(self.logging_obj).dispatch_success_handlers( + partial, + cache_hit=self._is_cache_hit(), + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + except Exception as e: # noqa: BLE001 # best-effort usage logging; closing must not raise into the caller + verbose_logger.debug("CustomStreamWrapper.aclose: could not log the abandoned stream's usage: %s", e) + + def _is_cache_hit(self) -> bool: + return self.custom_llm_provider == "cached_response" + + def _release(self) -> None: stream_to_close: Final = cast(object, self.completion_stream) # cast-ok: completion_stream is untyped self.completion_stream = None self._closed = True @@ -626,7 +691,7 @@ class CustomStreamWrapper: verbose_logger.debug("CustomStreamWrapper.close: error closing completion_stream: %s", e) self._restore_consumer_correlation_context() - async def aclose(self): + async def _arelease(self) -> None: # Restore the consumer's outer context only after the underlying # provider stream's own close (and its diagnostic logging below, if # closing fails) completes - not before - so those log lines still @@ -2313,9 +2378,8 @@ class CustomStreamWrapper: def __next__(self) -> "ModelResponseStream": if self._closed: raise StopIteration - cache_hit = False - if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": - cache_hit = True + self._iterated_provider_stream = True + cache_hit: Final = self._is_cache_hit() self._check_max_streaming_duration() try: if self.completion_stream is None: @@ -2401,12 +2465,14 @@ class CustomStreamWrapper: error: Final = self._stream_eof_error() eof_traceback: Final = traceback.format_exc() self._record_partial_usage_for_failure() + self._outcome_logged = True threading.Thread(target=self.logging_obj.failure_handler, args=(error, eof_traceback)).start() - self.close() + self._release() self._handle_stream_fallback_error(error) if self.sent_last_chunk is True: complete_streaming_response: Final = self._build_complete_streaming_response() response = self._new_usage_chunk() + self._outcome_logged = True if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params( complete_streaming_response, self.custom_llm_provider, self.chunks @@ -2463,7 +2529,7 @@ class CustomStreamWrapper: if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response - self.close() + self._release() raise # Re-raise StopIteration else: self.sent_last_chunk = True @@ -2496,8 +2562,9 @@ class CustomStreamWrapper: except Exception as e: traceback_exception: Final = traceback.format_exc() # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated + self._outcome_logged = True threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start() - self.close() + self._release() self._handle_stream_fallback_error(e) def fetch_sync_stream(self): @@ -2519,9 +2586,8 @@ class CustomStreamWrapper: async def __anext__(self) -> "ModelResponseStream": if self._closed: raise StopAsyncIteration - cache_hit = False - if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": - cache_hit = True + self._iterated_provider_stream = True + cache_hit: Final = self._is_cache_hit() try: # Inside the try (not before it) so a raised litellm.Timeout flows # through the same except Exception -> _handle_stream_fallback_error @@ -2608,7 +2674,7 @@ class CustomStreamWrapper: return usage_free_chunk except (StopAsyncIteration, StopIteration): if self._should_fail_on_clean_eof(): - await self.aclose() + await self._arelease() self._log_stream_failure_and_raise(self._stream_eof_error()) return await self._finalize_completed_stream(cache_hit=cache_hit) except httpx.TimeoutException as e: # if httpx read timeout error occues @@ -2616,15 +2682,15 @@ class CustomStreamWrapper: ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}" self._dispatch_stream_failure(e, traceback_exception) - await self.aclose() + await self._arelease() self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: - await self.aclose() + await self._arelease() self._log_stream_failure_and_raise(e) return await self._finalize_completed_stream(cache_hit=cache_hit) except Exception as e: - await self.aclose() + await self._arelease() self._log_stream_failure_and_raise(e) async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": @@ -2668,6 +2734,7 @@ class CustomStreamWrapper: self.sent_stream_usage = True return response + self._outcome_logged = True _deferred_cb: Final = getattr( self.logging_obj, "_on_deferred_stream_complete", @@ -2690,7 +2757,7 @@ class CustomStreamWrapper: # when consumers use ``async for`` on sync-SDK streams. Legacy string # callbacks still run via executor.submit inside dispatch_success_handlers. asyncio.create_task( - self.logging_obj.dispatch_success_handlers( + _success_dispatcher(self.logging_obj).dispatch_success_handlers( complete_streaming_response, cache_hit=cache_hit, start_time=None, @@ -2699,7 +2766,7 @@ class CustomStreamWrapper: ) ) - await self.aclose() + await self._arelease() raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True @@ -2722,6 +2789,7 @@ class CustomStreamWrapper: self._handle_stream_fallback_error(e) def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None: + self._outcome_logged = True self._record_partial_usage_for_failure() self._failure_logging = asyncio.create_task( self.logging_obj.dispatch_attempt_failure_handlers(e, traceback_exception) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cf8de9c40f5..7a40d0ab31f 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -13,14 +13,16 @@ from dataclasses import dataclass from itertools import chain from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType -from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast +from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, TypeAlias, TypedDict, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +import respx from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, TypeAdapter +from typing_extensions import ReadOnly import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var @@ -47,6 +49,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( ChoiceLogprobs, + Choices, CompletionTokensDetailsWrapper, Delta, ModelResponse, @@ -9571,3 +9574,463 @@ def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( assert usage.prompt_tokens_details.cached_tokens == 80 assert usage.completion_tokens_details is not None assert usage.completion_tokens_details.reasoning_tokens == 5 + + +_CUT_STREAM_URL: Final = "https://api.openai.com/v1/chat/completions" +_CUT_STREAM_MODEL: Final = "gpt-5.6" +_CUT_STREAM_PROMPT: Final = "Write a story about a lighthouse keeper." +_CUT_STREAM_DELTAS: Final = ("Once", " upon", " a", " time", " a", " keeper", " lit", " the", " lamp", ".") +_CUT_STREAM_PROMPT_TOKENS: Final = 431 +_CUT_STREAM_COMPLETION_TOKENS: Final = 57 +_CloseWith: TypeAlias = Literal["aclose", "async_with", "close", "with"] + + +def _cut_stream_event(fields: Mapping[str, object]) -> str: + chunk: Final = { + "id": "chatcmpl-cut-stream", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": _CUT_STREAM_MODEL, + **fields, + } + return f"data: {json.dumps(chunk)}\n\n" + + +def _cut_stream_content(position: int, text: str) -> str: + delta: Final = {"role": "assistant", "content": text} if position == 0 else {"content": text} + return _cut_stream_event({"choices": [{"index": 0, "delta": delta, "finish_reason": None}]}) + + +def _cut_stream_content_sse(deltas: Sequence[str]) -> str: + return "".join(_cut_stream_content(position, text) for position, text in enumerate(deltas)) + + +def _cut_stream_sse(deltas: Sequence[str]) -> str: + return ( + _cut_stream_content_sse(deltas) + + _cut_stream_event({"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}) + + _cut_stream_event( + { + "choices": [], + "usage": { + "prompt_tokens": _CUT_STREAM_PROMPT_TOKENS, + "completion_tokens": _CUT_STREAM_COMPLETION_TOKENS, + "total_tokens": _CUT_STREAM_PROMPT_TOKENS + _CUT_STREAM_COMPLETION_TOKENS, + }, + } + ) + + "data: [DONE]\n\n" + ) + + +_CUT_STREAM_SSE: Final = _cut_stream_sse(_CUT_STREAM_DELTAS) + + +class _ConnectionResetMidStream(httpx.SyncByteStream, httpx.AsyncByteStream): + def __iter__(self) -> Iterator[bytes]: + yield _cut_stream_content_sse(_CUT_STREAM_DELTAS).encode() + raise httpx.ReadError("connection reset by peer") + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _cut_stream_content_sse(_CUT_STREAM_DELTAS).encode() + raise httpx.ReadError("connection reset by peer") + + +class _LoggedStreamOutcome(TypedDict): + model: ReadOnly[str] + response_cost: ReadOnly[float] + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + + +@dataclass(frozen=True, slots=True) +class _LoggedStreamSuccess: + outcome: _LoggedStreamOutcome + text: str + + +def _logged_text(response_obj: object) -> str: + choices: Final = response_obj.choices if isinstance(response_obj, ModelResponse) else [] + first: Final = choices[0] if choices else None + return (first.message.content or "") if isinstance(first, Choices) else "" + + +class _StreamOutcomeRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.successes: tuple[_LoggedStreamSuccess, ...] = () + self.failures: tuple[str, ...] = () + self._recording: Final = threading.Lock() + self.logged: Final = threading.Event() + + def _record_success(self, kwargs: Mapping[str, object], response_obj: object) -> None: + success: Final = _LoggedStreamSuccess( + outcome=TypeAdapter(_LoggedStreamOutcome).validate_python(kwargs["standard_logging_object"]), + text=_logged_text(response_obj), + ) + with self._recording: + self.successes = (*self.successes, success) + self.logged.set() + + def _record_failure(self, mode: str) -> None: + with self._recording: + self.failures = (*self.failures, mode) + self.logged.set() + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record_success(kwargs, response_obj) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record_success(kwargs, response_obj) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record_failure("sync") + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record_failure("async") + + async def settled(self) -> None: + assert await asyncio.to_thread(self.logged.wait, 5), "no success or failure was logged" + await asyncio.sleep(0.5) + + +@pytest.fixture +def stream_outcomes(monkeypatch: pytest.MonkeyPatch) -> _StreamOutcomeRecorder: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _StreamOutcomeRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + return recorder + + +def _serve_openai_stream(respx_mock: respx.MockRouter, stream: str | _ConnectionResetMidStream) -> None: + headers: Final = {"content-type": "text/event-stream"} + response: Final = ( + httpx.Response(200, text=stream, headers=headers) + if isinstance(stream, str) + else httpx.Response(200, stream=stream, headers=headers) + ) + respx_mock.post(_CUT_STREAM_URL).mock(return_value=response) + + +def _open_sync_stream() -> CustomStreamWrapper: + stream: Final = litellm.completion( + model=_CUT_STREAM_MODEL, + messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}], + stream=True, + stream_options={"include_usage": True}, + api_key="sk-unit-test", + ) + assert isinstance(stream, CustomStreamWrapper) + return stream + + +async def _open_async_stream() -> CustomStreamWrapper: + stream: Final = await litellm.acompletion( + model=_CUT_STREAM_MODEL, + messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}], + stream=True, + stream_options={"include_usage": True}, + api_key="sk-unit-test", + ) + assert isinstance(stream, CustomStreamWrapper) + return stream + + +def _text_of(chunks: Sequence[ModelResponseStream]) -> str: + return "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + + +async def _read_async(stream: CustomStreamWrapper, chunks: int) -> str: + return _text_of([await stream.__anext__() for _ in range(chunks)]) + + +async def _read_two_chunks_then_close(close_with: _CloseWith) -> str: + if close_with == "close": + sync_stream: Final = _open_sync_stream() + sync_seen: Final = _text_of(tuple(itertools.islice(sync_stream, 2))) + sync_stream.close() + return sync_seen + if close_with == "with": + with _open_sync_stream() as managed_sync_stream: + return _text_of(tuple(itertools.islice(managed_sync_stream, 2))) + if close_with == "async_with": + async with await _open_async_stream() as managed_stream: + return await _read_async(managed_stream, 2) + stream: Final = await _open_async_stream() + seen: Final = await _read_async(stream, 2) + await stream.aclose() + return seen + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_with", ["aclose", "async_with", "close", "with"]) +async def test_stream_closed_by_its_consumer_before_the_end_logs_the_streamed_part_once( + close_with: _CloseWith, stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter +) -> None: + _serve_openai_stream(respx_mock, _CUT_STREAM_SSE) + + seen: Final = await _read_two_chunks_then_close(close_with) + await stream_outcomes.settled() + + assert seen == "".join(_CUT_STREAM_DELTAS[:2]) + assert stream_outcomes.failures == () + assert len(stream_outcomes.successes) == 1 + logged: Final = stream_outcomes.successes[0] + assert logged.text == seen + assert logged.outcome["prompt_tokens"] == litellm.token_counter( + model=_CUT_STREAM_MODEL, messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}] + ) + assert logged.outcome["completion_tokens"] == litellm.token_counter( + model=_CUT_STREAM_MODEL, text=seen, count_response_tokens=True + ) + assert logged.outcome["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_async_stream_closed_right_after_its_usage_chunk_logs_the_provider_usage( + stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter +) -> None: + _serve_openai_stream(respx_mock, _CUT_STREAM_SSE) + stream: Final = await _open_async_stream() + + received: Final = [] + async for chunk in stream: + received.append(chunk) + if isinstance(getattr(chunk, "usage", None), Usage): + break + await stream.aclose() + await stream_outcomes.settled() + + assert _text_of(received) == "".join(_CUT_STREAM_DELTAS) + assert len(stream_outcomes.successes) == 1 + logged: Final = stream_outcomes.successes[0] + assert logged.text == "".join(_CUT_STREAM_DELTAS) + assert (logged.outcome["prompt_tokens"], logged.outcome["completion_tokens"]) == ( + _CUT_STREAM_PROMPT_TOKENS, + _CUT_STREAM_COMPLETION_TOKENS, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_stream_read_to_its_end_then_closed_logs_the_provider_usage_once( + sync_mode: bool, stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter +) -> None: + _serve_openai_stream(respx_mock, _CUT_STREAM_SSE) + + if sync_mode: + sync_stream: Final = _open_sync_stream() + sync_seen: Final = _text_of(tuple(sync_stream)) + sync_stream.close() + else: + stream: Final = await _open_async_stream() + seen: Final = _text_of([chunk async for chunk in stream]) + await stream.aclose() + await stream_outcomes.settled() + + assert (sync_seen if sync_mode else seen) == "".join(_CUT_STREAM_DELTAS) + assert len(stream_outcomes.successes) == 1 + logged: Final = stream_outcomes.successes[0] + assert (logged.outcome["prompt_tokens"], logged.outcome["completion_tokens"]) == ( + _CUT_STREAM_PROMPT_TOKENS, + _CUT_STREAM_COMPLETION_TOKENS, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "broken_stream", + [ + pytest.param(_ConnectionResetMidStream(), id="connection_reset"), + pytest.param(_cut_stream_content_sse(_CUT_STREAM_DELTAS), id="body_ended_before_finish_reason"), + ], +) +async def test_stream_that_broke_mid_flight_logs_no_success_when_closed( + sync_mode: bool, + broken_stream: str | _ConnectionResetMidStream, + stream_outcomes: _StreamOutcomeRecorder, + respx_mock: respx.MockRouter, +) -> None: + _serve_openai_stream(respx_mock, broken_stream) + + if sync_mode: + sync_stream: Final = _open_sync_stream() + with pytest.raises(MidStreamFallbackError): + tuple(sync_stream) + sync_stream.close() + else: + stream: Final = await _open_async_stream() + with pytest.raises(MidStreamFallbackError): + [chunk async for chunk in stream] + await stream.aclose() + await stream_outcomes.settled() + + assert stream_outcomes.successes == () + assert len(stream_outcomes.failures) >= 1 + + +@pytest.mark.asyncio +async def test_stream_whose_failure_was_already_reported_logs_no_success_when_closed( + stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter +) -> None: + _serve_openai_stream(respx_mock, _CUT_STREAM_SSE) + stream: Final = await _open_async_stream() + await _read_async(stream, 2) + + assert stream.logging_obj.claim_stream_failure_report() is True + await stream.aclose() + await asyncio.sleep(0.5) + + assert stream_outcomes.successes == () + + +@pytest.mark.asyncio +async def test_router_stream_closed_by_its_consumer_before_the_end_logs_the_streamed_part_once( + stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter +) -> None: + _serve_openai_stream(respx_mock, _CUT_STREAM_SSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "story-writer", + "litellm_params": {"model": f"openai/{_CUT_STREAM_MODEL}", "api_key": "sk-unit-test"}, + } + ] + ) + stream: Final = await router.acompletion( + model="story-writer", + messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}], + stream=True, + stream_options={"include_usage": True}, + ) + assert isinstance(stream, CustomStreamWrapper) + + seen: Final = await _read_async(stream, 2) + await stream.aclose() + await stream_outcomes.settled() + + assert seen == "".join(_CUT_STREAM_DELTAS[:2]) + assert len(stream_outcomes.successes) == 1 + logged: Final = stream_outcomes.successes[0] + assert logged.text == seen + assert logged.outcome["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_router_stream_closed_after_a_mid_stream_fallback_logs_only_the_fallback_part( + stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter +) -> None: + backup_api_base: Final = "https://backup.openai.test/v1" + backup_deltas: Final = (" Then", " the", " storm", " came", ".") + _serve_openai_stream(respx_mock, _ConnectionResetMidStream()) + respx_mock.post(f"{backup_api_base}/chat/completions").mock( + return_value=httpx.Response( + 200, text=_cut_stream_sse(backup_deltas), headers={"content-type": "text/event-stream"} + ) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "story-writer", + "litellm_params": {"model": f"openai/{_CUT_STREAM_MODEL}", "api_key": "sk-unit-test"}, + }, + { + "model_name": "story-writer-backup", + "litellm_params": { + "model": f"openai/{_CUT_STREAM_MODEL}", + "api_key": "sk-unit-test", + "api_base": backup_api_base, + }, + "model_info": {"supports_assistant_prefill": True}, + }, + ], + fallbacks=[{"story-writer": ["story-writer-backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + stream: Final = await router.acompletion( + model="story-writer", + messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}], + stream=True, + stream_options={"include_usage": True}, + ) + assert isinstance(stream, CustomStreamWrapper) + + primary_part: Final = await _read_async(stream, len(_CUT_STREAM_DELTAS)) + fallback_part: Final = await _read_async(stream, 2) + await stream.aclose() + await stream_outcomes.settled() + + assert primary_part == "".join(_CUT_STREAM_DELTAS) + assert fallback_part == "".join(backup_deltas[:2]) + assert [logged.text for logged in stream_outcomes.successes] == [fallback_part] + + +@pytest.mark.asyncio +async def test_stream_closed_early_logs_nothing_while_streaming_logging_is_disabled( + stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_streaming_logging", True) + _serve_openai_stream(respx_mock, _CUT_STREAM_SSE) + stream: Final = await _open_async_stream() + await _read_async(stream, 2) + + await stream.aclose() + await asyncio.sleep(0.5) + + assert stream_outcomes.successes == () + + +class _SuccessLoggingDown(Logging): + async def dispatch_success_handlers(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("success logging is down") + + +class _ProviderChunks: + def __init__(self, chunks: Sequence[ModelResponseStream]) -> None: + self._chunks: Final = iter(chunks) + self.closed = False + + def __aiter__(self) -> "_ProviderChunks": + return self + + async def __anext__(self) -> ModelResponseStream: + chunk: Final = next(self._chunks, None) + if chunk is None: + raise StopAsyncIteration + return chunk + + async def aclose(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +async def test_async_stream_closed_early_still_closes_the_provider_stream_when_logging_its_part_fails() -> None: + provider_stream: Final = _ProviderChunks(tuple(_stream_chunk({"content": text}) for text in _CUT_STREAM_DELTAS)) + stream: Final = CustomStreamWrapper( + completion_stream=provider_stream, + model=_CUT_STREAM_MODEL, + logging_obj=_SuccessLoggingDown( + model=_CUT_STREAM_MODEL, + messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="cut-stream-logging-down", + function_id="cut-stream-logging-down", + ), + ) + seen: Final = await _read_async(stream, 2) + + await stream.aclose() + + assert seen == "".join(_CUT_STREAM_DELTAS[:2]) + assert provider_stream.closed