diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 23680b89f91..3fc754e8bb8 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -482,6 +482,16 @@ 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 cancelled.cancelling() + and logger.litellm_params.get("proxy_server_request") + ): + # 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/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index a9ec6a4f6ff..a6557adfbfc 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 @@ -346,3 +347,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"]