Merge nate_messages_native_disconnect_error_info into integration

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)
 
 

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"]