From d1687110fe6e0cb73518788619dab1de99f8ff8f Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Fri, 6 Nov 2026 12:43:48 +0000 Subject: [PATCH] fix(messages): bill an abandoned native proxy stream after the disconnect is recorded The proxy records a client disconnect, with its 499 error information, in the stream generator's cleanup, and the logged payload reads it from metadata. The native teardown billed the stream during cancel propagation, before that cleanup, so the logged row had empty error information --- .../host/callbacks_legacy_python.py | 10 +++++ .../host/test_callbacks_legacy_python.py | 41 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index b4c7a6a4ee5..9beed71a5bf 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -454,6 +454,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 1ba0188e76d..2fa6f3b74b2 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 @@ -298,3 +299,43 @@ def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None: assert head["additional_headers"] == {"request-id": "req_native", "llm_provider-request-id": "req_native"} assert head["litellm_call_id"] == "messages" assert "_response_ms" in head + + +@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"]