[Feat]Cancel upstream on client disconnect (#14295)

* cancel upstream on client disconnect

* add comments

* add test

* set timeout in constraints.py

* Guard against missing 'type' key

* update dependency to fix uvicorn bugs

litellm/constants.py

@@ -893,6 +893,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

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import time
 import traceback
 from datetime import datetime
 from typing import (
@@ -24,6 +25,7 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
     DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
+    DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS,
     STREAM_SSE_DATA_PREFIX,
 )
 from litellm.litellm_core_utils.dd_tracing import tracer
@@ -175,6 +177,29 @@ async def create_streaming_response(
     )
 
 
+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
@@ -425,12 +450,24 @@ 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:
+            raise HTTPException(
+                status_code=499,
+                detail="Client disconnected the request",
+            )
 
         response = responses[1]
 

litellm/proxy/proxy_server.py

@@ -997,33 +997,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."""

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -6122,15 +6122,15 @@ zstd = ["zstandard (>=0.18.0)"]
 
 [[package]]
 name = "uvicorn"
-version = "0.29.0"
+version = "0.32.1"
 description = "The lightning-fast ASGI server."
 optional = true
 python-versions = ">=3.8"
 groups = ["main"]
 markers = "python_version >= \"3.10\" and (extra == \"mlflow\" or extra == \"proxy\") or extra == \"proxy\""
 files = [
-    {file = "uvicorn-0.29.0-py3-none-any.whl", hash = "sha256:2c2aac7ff4f4365c206fd773a39bf4ebd1047c238f8b8268ad996829323473de"},
-    {file = "uvicorn-0.29.0.tar.gz", hash = "sha256:6a69214c0b6a087462412670b3ef21224fa48cae0e452b5883e8e8bdfdd11dd0"},
+    {file = "uvicorn-0.32.1-py3-none-any.whl", hash = "sha256:82ad92fd58da0d12af7482ecdb5f2470a04c9c9a53ced65b9bbb4a205377602e"},
+    {file = "uvicorn-0.32.1.tar.gz", hash = "sha256:ee9519c246a72b1c084cea8d3b44ed6026e78a4a309cbedae9c37e4cb9fbb175"},
 ]
 
 [package.dependencies]
@@ -6139,7 +6139,7 @@ h11 = ">=0.8"
 typing-extensions = {version = ">=4.0", markers = "python_version < \"3.11\""}
 
 [package.extras]
-standard = ["colorama (>=0.4) ; sys_platform == \"win32\"", "httptools (>=0.5.0)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1) ; sys_platform != \"win32\" and sys_platform != \"cygwin\" and platform_python_implementation != \"PyPy\"", "watchfiles (>=0.13)", "websockets (>=10.4)"]
+standard = ["colorama (>=0.4) ; sys_platform == \"win32\"", "httptools (>=0.6.3)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1) ; sys_platform != \"win32\" and sys_platform != \"cygwin\" and platform_python_implementation != \"PyPy\"", "watchfiles (>=0.13)", "websockets (>=10.4)"]
 
 [[package]]
 name = "uvloop"
@@ -6576,4 +6576,4 @@ utils = ["numpydoc"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.8.1,<4.0, !=3.9.7"
-content-hash = "f41e6359109c5c52dba2a28f301b04030d865265f408974082b390bf45568a01"
+content-hash = "e48cc445bc012e020a9e311942e46833dda587b70a630a04bfc08b629746fe56"

pyproject.toml

@@ -34,7 +34,7 @@ pydantic = "^2.5.0"
 jsonschema = "^4.22.0"
 numpydoc = {version = "*", optional = true} # used in utils.py
 
-uvicorn = {version = "^0.29.0", optional = true}
+uvicorn = {version = "^0.32.0", optional = true}
 uvloop = {version = "^0.21.0", optional = true, markers="sys_platform != 'win32'"}
 gunicorn = {version = "^23.0.0", optional = true}
 fastapi = {version = "^0.115.5", optional = true}

requirements.txt

@@ -5,7 +5,7 @@ openai==1.99.5  # openai req.
 fastapi==0.115.5 # server dep
 backoff==2.2.1 # server dep
 pyyaml==6.0.2 # server dep
-uvicorn==0.29.0 # server dep
+uvicorn==0.32.0 # server dep
 gunicorn==23.0.0 # server dep
 fastuuid==0.12.0 # for uuid4
 uvloop==0.21.0 # uvicorn dep, gives us much better performance under load

tests/proxy_unit_tests/test_client_disconnection.py

@@ -0,0 +1,47 @@
+"""
+Test client disconnection detection functionality.
+"""
+import asyncio
+import pytest
+from unittest.mock import AsyncMock
+
+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 = AsyncMock()
+
+    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 = AsyncMock()
+
+    # 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()