feat(proxy): cancel upstream LLM request on client disconnect

litellm/constants.py

@@ -1287,6 +1287,9 @@ MAX_SPENDLOG_ROWS_TO_QUERY = int(
 DEFAULT_SOFT_BUDGET = float(
     os.getenv("DEFAULT_SOFT_BUDGET", 50.0)
 )  # by default all litellm proxy keys have a soft budget of 50.0
+DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS = int(
+    os.getenv("DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS", 600)
+)  # 10 minutes timeout for client disconnect checking in proxy
 # makes it clear this is a rate limit error for a litellm virtual key
 RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY = "LiteLLM Virtual Key user_api_key_hash"
 

litellm/proxy/common_request_processing.py

@@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.constants import (
     DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
+    DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS,
     DEFAULT_MAX_RECURSE_DEPTH,
     LITELLM_DETAILED_TIMING,
     MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
@@ -486,6 +487,30 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
     return False
 
 
+async def _check_request_disconnection(request: Request, llm_api_call_task):
+    """
+    Asynchronously checks if the request is disconnected at regular intervals.
+    If the request is disconnected
+    - cancel the litellm.router task
+
+    Parameters:
+    - request: Request: The request object to check for disconnection.
+    Returns:
+    - None
+    """
+
+    # only run this function for configured timeout -> if these don't get cancelled -> we don't want the server to have many while loops
+    start_time = time.time()
+    while time.time() - start_time < DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS:
+        await asyncio.sleep(1)
+        message = await request.receive()
+        if message.get("type") == "http.disconnect":
+            # 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()
+            return
+
+
 class ProxyBaseLLMRequestProcessing:
     def __init__(self, data: dict):
         self.data = data
@@ -1057,12 +1082,22 @@ class ProxyBaseLLMRequestProcessing:
         )
         tasks.append(llm_call)
 
-        # wait for call to end
         llm_responses = asyncio.gather(
             *tasks
         )  # run the moderation check in parallel to the actual llm api call
 
-        responses = await llm_responses
+        # Execute the task to detect disconnection
+        disconnect_task = asyncio.create_task(_check_request_disconnection(request, llm_responses))
+
+        try:
+            # wait for call to end
+            # Note: In the case of streaming, processing does not wait here, so disconnection detection is performed in StreamingResponse.
+            responses = await llm_responses
+            disconnect_task.cancel()
+
+        except asyncio.CancelledError:
+            verbose_proxy_logger.info("Client disconnected, cancelled upstream LLM request")
+            raise
 
         response = responses[1]
 

litellm/proxy/proxy_server.py

@@ -1608,34 +1608,6 @@ db_writer_client: Optional[AsyncHTTPHandler] = None
 ### logger ###
 
 
-async def check_request_disconnection(request: Request, llm_api_call_task):
-    """
-    Asynchronously checks if the request is disconnected at regular intervals.
-    If the request is disconnected
-    - cancel the litellm.router task
-    - raises an HTTPException with status code 499 and detail "Client disconnected the request".
-
-    Parameters:
-    - request: Request: The request object to check for disconnection.
-    Returns:
-    - None
-    """
-
-    # only run this function for 10 mins -> if these don't get cancelled -> we don't want the server to have many while loops
-    start_time = time.time()
-    while time.time() - start_time < 600:
-        await asyncio.sleep(1)
-        if await request.is_disconnected():
-            # 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()
-
-            raise HTTPException(
-                status_code=499,
-                detail="Client disconnected the request",
-            )
-
-
 def _resolve_typed_dict_type(typ):
     """Resolve the actual TypedDict class from a potentially wrapped type."""
     from typing_extensions import _TypedDictMeta  # type: ignore

tests/test_litellm/proxy/test_client_disconnection.py

@@ -0,0 +1,47 @@
+"""
+Test client disconnection detection functionality.
+"""
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+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."""
+    mock_request = AsyncMock()
+    mock_request.receive.side_effect = [
+        {"type": "http.request"},   # First call
+        {"type": "http.disconnect"} # Second call - disconnect
+    ]
+
+    mock_llm_task = MagicMock()  # sync mock so .cancel() doesn't return a coroutine
+
+    await _check_request_disconnection(mock_request, mock_llm_task)
+
+    mock_llm_task.cancel.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_check_request_disconnection_no_disconnect():
+    """Test that _check_request_disconnection handles normal requests."""
+    mock_request = AsyncMock()
+    mock_request.receive.return_value = {"type": "http.request"}
+
+    mock_llm_task = MagicMock()  # sync mock so .cancel() doesn't return a coroutine
+
+    # 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))
+    await asyncio.sleep(0.1)  # Let it run briefly
+    task.cancel()
+
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+
+    # Task should not be cancelled during normal operation
+    mock_llm_task.cancel.assert_not_called()