Revert "[Feat]Cancel upstream on client disconnect"

litellm/constants.py

@@ -910,9 +910,6 @@
 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,7 +1,6 @@
 import asyncio
 import json
 import logging
-import time
 import traceback
 from datetime import datetime
 from typing import (
@@ -25,7 +24,6 @@
 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
@@ -177,29 +175,6 @@ async def combined_generator() -> AsyncGenerator[str, None]:
     )
 
 
-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
@@ -450,24 +425,12 @@ async def base_process_llm_request(
         )
         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
 
-        # 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",
-            )
+        responses = await llm_responses
 
         response = responses[1]
 

litellm/proxy/proxy_server.py

@@ -997,6 +997,33 @@ def swagger_monkey_patch(*args, **kwargs):
 ### 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."""

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.32.0", optional = true}
+uvicorn = {version = "^0.29.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.32.0 # server dep
+uvicorn==0.29.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