diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index e0155f88de1..f0233f5f040 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -254,6 +254,7 @@ class CustomStreamWrapper: self.sent_last_chunk = False self._failure_logging: asyncio.Task[None] | None = None self._closed = False + self._abandonment_logged: bool = False self._stream_created_time: float = time.time() litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate( @@ -355,7 +356,7 @@ class CustomStreamWrapper: return self def __exit__(self, *exc_info: object) -> None: - self.close() + self.close_stream() async def __aenter__(self) -> "CustomStreamWrapper": return self @@ -363,6 +364,81 @@ class CustomStreamWrapper: async def __aexit__(self, *exc_info: object) -> None: await self.aclose() + def close_stream(self) -> None: + """Deterministically end a stream the consumer stopped reading early. + + A stream abandoned without reaching its final chunk would otherwise + never emit its terminal success log (spans, spend, custom callbacks + all miss the call). ``close_stream()`` flushes that log from the chunks + received so far and then closes the underlying provider stream; safe to + call after the stream has been fully consumed (it is a no-op then). + + Deliberately *not* the same entry point as ``close()``: ``Router`` + closes every response in a ``finally`` block, so treating that + internal teardown as "the consumer abandoned this stream" would + assemble a response and submit a success log on behalf of every + fallback and shadow stream. ``close()`` therefore stays a pure + teardown, and a caller that abandons a stream has to say so, through + this method or the context manager above. + """ + self._flush_abandoned_stream_logging() + self.close() + + def _flush_abandoned_stream_logging(self) -> None: + """Emit a terminal success log for a stream the consumer abandoned. + + A stream the caller stops reading early (break, drop, aclose) never + reaches the final-chunk branch that assembles + ``complete_streaming_response`` and submits the terminal success log, + so every success callback (tracing, spend, ...) silently misses the + call even though tokens were served. Assemble the response from the + chunks received so far and submit the same logging once, mirroring how + the proxy bills a partial stream on client disconnect. + + Only reachable through ``close_stream()`` / the context manager. It is + deliberately not wired into either of the other two hooks this class + already had: + + - ``__del__`` runs at an arbitrary point on an arbitrary thread, + possibly during interpreter shutdown or after the event loop it + needs is gone, so submitting logging from there is neither reliable + nor testable. It also races the success callbacks of a stream that is + still being drained. + - ``aclose()`` is litellm's own cleanup path. `Router` calls it on + fallback and shutdown paths (`hasattr(model_response, "aclose")`), so + treating it as "the consumer abandoned this stream" would assemble a + response and submit a success log on behalf of every internal + teardown. Measured cost: with the flush in ``aclose()`` the + ``tests/unit/test_router`` suite went from 150s to over 25 minutes, + because each one runs `stream_chunk_builder` (token counting) + synchronously. + + So a caller that abandons a stream has to say so, through + ``close_stream()``. + """ + if self._abandonment_logged or self.sent_last_chunk: + return + if not self.chunks: # pyright: ignore[reportUnknownMemberType] # chunks is a loosely-typed accumulator + return + self._abandonment_logged = True + try: + complete_streaming_response: Final = litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # stream_chunk_builder is exposed on the package root + chunks=self.chunks, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # chunks is a loosely-typed accumulator + messages=self.messages, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + except Exception as e: # noqa: BLE001 # an abandoned stream must still emit whatever telemetry was built; the failure is logged, not raised + verbose_logger.warning( + "stream_chunk_builder raised for an abandoned stream (%s); the partial call is not logged.", + str(e), + ) + return + if complete_streaming_response is None: + return + self.logging_obj.model_call_details["async_complete_streaming_response"] = complete_streaming_response + self.run_success_logging_and_cache_storage(complete_streaming_response, cache_hit=False) # pyright: ignore[reportUnknownMemberType] # resolved dynamically on the logging object + def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None: """Restore trace_id/session_id in the *consuming* thread/task/context. diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 50bcc14b681..ae4334dea80 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6872,3 +6872,152 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): async with _wrapper_over(stream): pass assert stream.close_calls == 1 + + +class TestCloseFlushesAbandonedStream: + """A stream abandoned before its final chunk must still emit its terminal + success log when closed (issue #45736): tracing, spend and custom + callbacks otherwise silently miss the call.""" + + class _RecordingLoggingObj: + """Logging object with one *real* attribute for the tests to assert on. + + `check_test_quality.py` flags a test whose every assertion inspects a + patched mock (TQ002): it restates the call instead of checking the + outcome. The flush writes the assembled response into + `model_call_details` before submitting the log, so a genuine dict there + is something a caller observes. Everything else stays a mock because + the wrapper's own success-logging path reaches far into the object. + """ + + def __init__(self) -> None: + self.model_call_details: dict = {} + self._rest = MagicMock() + self._rest.async_success_handler = AsyncMock() + + def __getattr__(self, name: str): + return getattr(self._rest, name) + + def _wrapper_with_chunk(self): + """Return a wrapper plus the logging object it was constructed with.""" + logging_obj = self._RecordingLoggingObj() + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider=None, + ) + wrapper.sent_last_chunk = False + wrapper.chunks.append(ModelResponseStream(id="c1", choices=[])) + return wrapper, logging_obj + + def test_close_flushes_terminal_log_for_abandoned_stream(self): + wrapper, logging_obj = self._wrapper_with_chunk() + + wrapper.close_stream() + + # the assembled response reached the logging object, and the wrapper + # recorded that it has now logged the stream + assembled = logging_obj.model_call_details["async_complete_streaming_response"] + assert assembled is not None + assert wrapper._abandonment_logged is True + + def test_close_is_noop_after_natural_completion(self): + wrapper, logging_obj = self._wrapper_with_chunk() + wrapper.sent_last_chunk = True + + wrapper.close_stream() + + # the natural completion already handled this stream, so closing must + # neither assemble a response nor claim the stream was abandoned + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is False + + def test_double_close_flushes_only_once(self): + wrapper, logging_obj = self._wrapper_with_chunk() + + wrapper.close_stream() + first = logging_obj.model_call_details["async_complete_streaming_response"] + wrapper.close_stream() + + # the second close neither replaced nor re-assembled anything + assert logging_obj.model_call_details["async_complete_streaming_response"] is first + + def test_context_manager_exit_flushes(self): + wrapper, logging_obj = self._wrapper_with_chunk() + + with wrapper: + pass + + assert "async_complete_streaming_response" in logging_obj.model_call_details + + def test_router_close_duck_typing_contract(self): + """`Router.stream_with_fallbacks` calls close() on anything that has it. + + It does so from a `finally`, on every response including shadow and + fallback streams, so `close()` has to stay a pure teardown: it must + not assemble a response nor submit the terminal success log. The + flushing entry point is `close_stream`, and this pins that split. + """ + wrapper, logging_obj = self._wrapper_with_chunk() + + wrapper.close() + + # close() is the teardown Router performs; it must not bill the stream + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is False + assert callable(wrapper.close_stream) + + def test_close_releases_underlying_stream(self): + wrapper, _ = self._wrapper_with_chunk() + released = [] + + class _Stream: + def close(inner_self): + released.append(True) + + wrapper.completion_stream = _Stream() + + wrapper.close_stream() + + # the provider stream was closed and the reference dropped + assert released == [True] + assert wrapper.completion_stream is None + + def test_no_flush_when_nothing_was_streamed(self): + wrapper, logging_obj = self._wrapper_with_chunk() + wrapper.chunks = [] + + wrapper.close_stream() + + # nothing was served, so there is nothing to bill or log + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is False + + def test_flush_survives_stream_chunk_builder_raising(self): + wrapper, logging_obj = self._wrapper_with_chunk() + + def _boom(**_kwargs): + raise RuntimeError("builder exploded") + + with patch("litellm.stream_chunk_builder", _boom): + wrapper.close_stream() + + # the failure is logged, not raised, and nothing half-built is recorded + assert "async_complete_streaming_response" not in logging_obj.model_call_details + assert wrapper._abandonment_logged is True + + def test_close_survives_provider_stream_raising(self): + wrapper, _ = self._wrapper_with_chunk() + + class _Angry: + def close(inner_self): + raise OSError("already closed") + + wrapper.completion_stream = _Angry() + + wrapper.close_stream() + + # a provider that fails to close must not mask the caller's own error, + # and the reference is still dropped + assert wrapper.completion_stream is None