Merge d0733bf98c into 24d92f77d5

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.
 

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