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

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.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013FDwtkaNTifWAP4W7JZavE

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

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

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

@@ -274,7 +274,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(
@@ -291,6 +296,7 @@ impl LegacyLogging {
                 &self.start,
                 &self.end,
                 &stream.first_chunk,
+                teardown,
             ),
         );
         match billed {
@@ -592,7 +598,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(()))
@@ -662,8 +668,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)
         {
@@ -1981,6 +2007,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

@@ -92,7 +92,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

@@ -419,6 +419,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,
@@ -443,7 +444,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:

tests/test_litellm_rust/messages/test_streaming.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import AsyncIterator, Iterator
 from typing import Final
 
@@ -68,3 +69,22 @@ async def test_stream_closed_early_logs_success_once(server: RecordingServer) ->
     assert len(await recorder.wait_for_async("async_log_success_event")) == 1
     with pytest.raises(StopAsyncIteration):
         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