diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index c0fe1c7b52f..8a212223523 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/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", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 28bd5f2cefb..f68541c5c7b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/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, + #[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 c88a6496423..ee16e0c0059 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/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), diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 5a368573f3d..ba557f769eb 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/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) diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index 491d257b482..99b6101da4e 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 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() diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index b00df6547d2..cb594767908 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/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"]