fix(messages): bill a native stream the client abandoned like python (#50425)

When the client disconnects mid-stream, the native adapter now logs success for the chunks it delivered before teardown, so spend, SpendLogs and the key's budget and tpm reservations settle the same way they do on the Python path

The proxy records a client disconnect, with its 499 error information, in the stream generator's cleanup. The native teardown billed the stream before that cleanup, so the row had empty error information. A proxy stream torn down by cancellation now bills from the task's done callback, after the disconnect is recorded

litellm-rust/crates/callbacks-legacy-python/python_contract.json

@@ -130,7 +130,8 @@
     "chunks",
     "start",
     "end",
-    "first_chunk"
+    "first_chunk",
+    "stream_teardown"
   ],
   "stream_failure": [
     "logger",

litellm-rust/crates/callbacks-legacy-python/src/adapter.rs

@@ -329,7 +329,12 @@ impl LegacyLogging {
         logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
     }
 
-    fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
+    fn stream_success(
+        &self,
+        py: Python<'_>,
+        stream: &DeliveredStream,
+        teardown: bool,
+    ) -> PyResult<()> {
         let logger = self.logger()?;
         let billing = self.stream_billing().ok_or(HostError::MissingState)?;
         let billed = Streaming::Success.call(
@@ -346,6 +351,7 @@ impl LegacyLogging {
                 &self.start,
                 &self.end,
                 &stream.first_chunk,
+                teardown,
             ),
         );
         match billed {
@@ -667,7 +673,7 @@ impl LegacyLogging {
         self.end = Some(datetime(py, timing.end_time)?);
         self.response = Some(response.clone_ref(py));
         match &self.stream {
-            Some(stream) => self.stream_success(py, stream)?,
+            Some(stream) => self.stream_success(py, stream, false)?,
             None => self.dispatch_success(py)?,
         }
         Ok(HookStep::Ready(()))
@@ -737,8 +743,28 @@ impl LegacyLogging {
     }
 }
 
+impl LegacyLogging {
+    // A caller that goes away mid-stream cancels the call before it ends: Python's Messages
+    // stream iterator still bills the chunks it read, straight away rather than deferred.
+    fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> {
+        if !self.asynchronous || self.end.is_some() {
+            return Ok(());
+        }
+        let Some(stream) = self.stream.take() else {
+            return Ok(());
+        };
+        self.end = Some(datetime(py, epoch_seconds())?);
+        self.stream_success(py, &stream, true)
+    }
+}
+
 impl PythonOwned for LegacyLogging {
     fn close(&mut self, py: Python<'_>) {
+        if self.logger.is_some()
+            && let Err(error) = self.bill_abandoned_stream(py)
+        {
+            error.write_unraisable(py, None);
+        }
         if let Some(logger) = self.logger.take()
             && let Err(error) = logger.restore_context(py)
         {
@@ -2186,6 +2212,45 @@ assert chunks[1] is last
         });
     }
 
+    #[rstest]
+    #[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
+    #[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
+    #[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])]
+    fn closing_an_unfinished_stream_bills_its_delivered_chunks_once(
+        #[case] ended: Option<bool>,
+        #[case] expected: &[&str],
+    ) {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = namespace(
+                py,
+                c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}",
+            );
+            let mut logging = LegacyLogging {
+                operation: crate::LoggingOperation::Messages,
+                ..logged(py, &locals, true)
+            };
+            logging
+                .on_stream_open(py, &local(&locals, "head").unbind())
+                .unwrap();
+            logging
+                .on_stream_chunk(py, &local(&locals, "first").unbind())
+                .unwrap();
+            match ended {
+                Some(true) => drop(succeed(py, &locals, &mut logging)),
+                Some(false) => drop(fail(py, &locals, &mut logging)),
+                None => {}
+            }
+            litellm_host_python::PythonOwned::close(&mut logging, py);
+            let names: Vec<String> = local(&locals, "logger")
+                .call_method0("names")
+                .unwrap()
+                .extract()
+                .unwrap();
+            assert_eq!(names, expected);
+        });
+    }
+
     #[rstest]
     #[case::synchronous(false, &["failure_handler"])]
     #[case::asynchronous(true, &[])]

litellm-rust/crates/callbacks-legacy-python/src/test_support.rs

@@ -98,7 +98,7 @@ FAKES = {
     'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
     'stream_opened': lambda logger, head: logger.record('stream_opened', head),
     'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head),
-    'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
+    'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
         'stream_success', list(chunks)
     ),
     'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),

litellm/rust_bridge/host/callbacks_legacy_python.py

@@ -11,6 +11,7 @@ import contextvars
 import datetime
 import inspect
 import json
+import sys
 import traceback
 import uuid
 import weakref
@@ -457,6 +458,14 @@ def stream_opened(logger: Logging, head: Mapping[str, object]) -> None:
     )
 
 
+def _cancelling(task: asyncio.Task[object]) -> bool:
+    return task.cancelling() > 0 if sys.version_info >= (3, 11) else True
+
+
+def _proxy_request(logger: LoggingSurface) -> object:
+    return logger.litellm_params.get("proxy_server_request")
+
+
 def stream_success(
     logger: Logging,
     url_route: str,
@@ -466,6 +475,7 @@ def stream_success(
     start: datetime.datetime,
     end: datetime.datetime,
     first_chunk: datetime.datetime | None,
+    stream_teardown: bool,
 ) -> None:
     from litellm.llms.anthropic.messages.native.streaming_iterator import (
         GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
@@ -490,7 +500,7 @@ def stream_success(
         raw_bytes=chunks,
         end_time=end,
     )
-    if getattr(logger, "_on_deferred_stream_complete", None) is not None:
+    if not stream_teardown and getattr(logger, "_on_deferred_stream_complete", None) is not None:
         logger._deferred_stream_complete_args = (coroutine,)  # pyright: ignore[reportAttributeAccessIssue]  # the proxy's deferred stream release reads this slot
         return
     try:
@@ -500,6 +510,11 @@ def stream_success(
 
         executor.submit(contextvars.copy_context().run, asyncio.run, coroutine)
         return
+    cancelled: Final = asyncio.current_task()
+    if stream_teardown and cancelled is not None and _cancelling(cancelled) and _proxy_request(logger):
+        # the proxy's disconnect cleanup runs later in this task and stamps the 499 error_information this row carries
+        cancelled.add_done_callback(lambda _: enqueue_logging(coroutine))
+        return
     enqueue_logging(coroutine)
 
 

tests/test_litellm_rust/messages/test_streaming.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from collections.abc import AsyncIterator, Iterator
 from typing import Final
@@ -73,6 +74,25 @@ async def test_stream_closed_early_logs_success_once(server: RecordingServer) ->
         await anext(stream)
 
 
+@pytest.mark.asyncio
+async def test_stream_cancelled_while_waiting_on_the_provider_logs_success_for_what_it_read(
+    server: RecordingServer,
+) -> None:
+    server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS, chunk_delay=1))
+    recorder: Final = RecordingLogger()
+
+    stream: Final = await call(server, stream=True, callbacks=[recorder])
+    assert isinstance(stream, AsyncIterator)
+    await anext(stream)
+    with pytest.raises(TimeoutError):
+        await asyncio.wait_for(anext(stream), timeout=0.2)
+
+    success: Final = await recorder.wait_for_async("async_log_success_event")
+    assert len(success) == 1
+    assert success[0].response.usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"]
+    assert "log_failure_event" not in recorder.names
+
+
 async def _streamed_pre_call_input(server: RecordingServer) -> object:
     server.enqueue(STREAM)
     recorder: Final = RecordingLogger()

tests/unit/rust_bridge/host/test_callbacks_legacy_python.py

@@ -1,3 +1,4 @@
+import asyncio
 import datetime
 import inspect
 from collections.abc import Mapping
@@ -362,3 +363,43 @@ async def test_a_native_stream_failure_logs_like_the_python_chunk_processor(monk
     assert logged == [("Error in chunk_processor: %s", cut)]
     assert "cache_hit" not in logger.model_call_details
     assert scheduled[0]["exception"] is cut
+
+
+@pytest.mark.asyncio
+async def test_a_proxy_stream_abandoned_mid_cancel_bills_after_the_disconnect_cleanup(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
+
+    logger: Final = _messages_logger()
+    logger.update_environment_variables(
+        litellm_params={"metadata": {}, "proxy_server_request": {"url": "/v1/messages"}}, optional_params={}
+    )
+
+    async def billing() -> None:
+        return None
+
+    events: Final[list[str]] = []
+    monkeypatch.setattr(PassThroughStreamingHandler, "_route_streaming_logging_to_handler", lambda **_: billing())
+    monkeypatch.setattr(legacy, "enqueue_logging", lambda coroutine: (events.append("billed"), coroutine.close()))
+    started: Final = asyncio.Event()
+    now: Final = datetime.datetime.now()
+
+    async def relay() -> None:
+        try:
+            started.set()
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            legacy.stream_success(logger, "/v1/messages", "anthropic", {}, [b"event: ping\n\n"], now, now, None, True)
+            await asyncio.sleep(0)
+            events.append("disconnect recorded")
+            raise
+
+    task: Final = asyncio.create_task(relay())
+    await started.wait()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    await asyncio.sleep(0)
+
+    assert events == ["disconnect recorded", "billed"]