diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..c70a9708d49 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/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", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..59d528c0fa4 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/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, + #[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 = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + #[rstest] #[case::synchronous(false, &["failure_handler"])] #[case::asynchronous(true, &[])] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 8acfa7a898b..a74121c4b56 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/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), diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92590531cdc..b4c7a6a4ee5 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/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: diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index e57631f48b9..0674d77e52c 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/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