fix(proxy): address review comments on client disconnect feature

litellm/proxy/common_request_processing.py

@@ -487,14 +487,22 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
     return False
 
 
-async def _check_request_disconnection(request: Request, llm_api_call_task):
+async def _check_request_disconnection(
+    request: Request,
+    llm_api_call_task,
+    disconnect_event: asyncio.Event,
+):
     """
     Asynchronously checks if the request is disconnected at regular intervals.
-    If the request is disconnected
-    - cancel the litellm.router task
+    If the request is disconnected:
+    - cancels the litellm.router task (effective for non-streaming requests)
+    - sets disconnect_event so the caller can distinguish a client disconnect
+      from other sources of CancelledError (e.g. server shutdown)
 
     Parameters:
-    - request: Request: The request object to check for disconnection.
+    - request: The request object to check for disconnection.
+    - llm_api_call_task: The asyncio gather future to cancel on disconnect.
+    - disconnect_event: Event set when the client disconnects.
     Returns:
     - None
     """
@@ -508,6 +516,7 @@ async def _check_request_disconnection(request: Request, llm_api_call_task):
             # cancel the LLM API Call task if any passed - this is passed from individual providers
             # Example OpenAI, Azure, VertexAI etc
             llm_api_call_task.cancel()
+            disconnect_event.set()
             return
 
 
@@ -1087,16 +1096,25 @@ class ProxyBaseLLMRequestProcessing:
         )  # run the moderation check in parallel to the actual llm api call
 
         # Execute the task to detect disconnection
-        disconnect_task = asyncio.create_task(_check_request_disconnection(request, llm_responses))
+        disconnect_event = asyncio.Event()
+        disconnect_task = asyncio.create_task(
+            _check_request_disconnection(request, llm_responses, disconnect_event)
+        )
 
         try:
             # wait for call to end
-            # Note: In the case of streaming, processing does not wait here, so disconnection detection is performed in StreamingResponse.
+            # Note: for streaming requests llm_responses resolves quickly once the
+            # upstream connection is established; the ASGI transport layer handles
+            # cancellation of the upstream when the client disconnects mid-stream.
             responses = await llm_responses
             disconnect_task.cancel()
-
         except asyncio.CancelledError:
-            verbose_proxy_logger.info("Client disconnected, cancelled upstream LLM request")
+            disconnect_task.cancel()
+            if disconnect_event.is_set():
+                raise HTTPException(
+                    status_code=499,
+                    detail="Client disconnected the request",
+                )
             raise
 
         response = responses[1]
@@ -1602,9 +1620,12 @@ class ProxyBaseLLMRequestProcessing:
         version: Optional[str] = None,
     ):
         """Raises ProxyException (OpenAI API compatible) if an exception is raised"""
-        verbose_proxy_logger.exception(
-            f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}"
-        )
+        if isinstance(e, HTTPException) and e.status_code == 499:
+            verbose_proxy_logger.info("Client disconnected the request (499)")
+        else:
+            verbose_proxy_logger.exception(
+                f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}"
+            )
         # Allow callbacks to transform the error response
         transformed_exception = await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict,

tests/test_litellm/proxy/test_client_disconnection.py

@@ -3,14 +3,14 @@ Test client disconnection detection functionality.
 """
 import asyncio
 import pytest
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from litellm.proxy.common_request_processing import _check_request_disconnection
 
 
 @pytest.mark.asyncio
 async def test_check_request_disconnection_with_disconnect():
-    """Test that _check_request_disconnection cancels task when client disconnects."""
+    """Test that _check_request_disconnection cancels task and sets event when client disconnects."""
     mock_request = AsyncMock()
     mock_request.receive.side_effect = [
         {"type": "http.request"},   # First call
@@ -18,23 +18,27 @@ async def test_check_request_disconnection_with_disconnect():
     ]
 
     mock_llm_task = MagicMock()  # sync mock so .cancel() doesn't return a coroutine
+    disconnect_event = asyncio.Event()
 
-    await _check_request_disconnection(mock_request, mock_llm_task)
+    with patch("litellm.proxy.common_request_processing.asyncio.sleep", new_callable=AsyncMock):
+        await _check_request_disconnection(mock_request, mock_llm_task, disconnect_event)
 
     mock_llm_task.cancel.assert_called_once()
+    assert disconnect_event.is_set()
 
 
 @pytest.mark.asyncio
 async def test_check_request_disconnection_no_disconnect():
-    """Test that _check_request_disconnection handles normal requests."""
+    """Test that _check_request_disconnection does not cancel task during normal operation."""
     mock_request = AsyncMock()
     mock_request.receive.return_value = {"type": "http.request"}
 
     mock_llm_task = MagicMock()  # sync mock so .cancel() doesn't return a coroutine
+    disconnect_event = asyncio.Event()
 
-    # This will timeout after 600 seconds, but we don't need to wait
-    # Just test that it doesn't crash immediately
-    task = asyncio.create_task(_check_request_disconnection(mock_request, mock_llm_task))
+    task = asyncio.create_task(
+        _check_request_disconnection(mock_request, mock_llm_task, disconnect_event)
+    )
     await asyncio.sleep(0.1)  # Let it run briefly
     task.cancel()
 
@@ -43,5 +47,5 @@ async def test_check_request_disconnection_no_disconnect():
     except asyncio.CancelledError:
         pass
 
-    # Task should not be cancelled during normal operation
     mock_llm_task.cancel.assert_not_called()
+    assert not disconnect_event.is_set()