From 26330c88978d1ec637986069aa5d7a6e7f3899cb Mon Sep 17 00:00:00 2001 From: CreateRandom Date: Wed, 15 Apr 2026 14:30:01 +0200 Subject: [PATCH 1/4] feat(proxy): cancel upstream LLM request on client disconnect --- litellm/constants.py | 3 ++ litellm/proxy/common_request_processing.py | 39 ++++++++++++++- litellm/proxy/proxy_server.py | 28 ----------- .../proxy/test_client_disconnection.py | 47 +++++++++++++++++++ 4 files changed, 87 insertions(+), 30 deletions(-) create mode 100644 tests/test_litellm/proxy/test_client_disconnection.py diff --git a/litellm/constants.py b/litellm/constants.py index d0596bed684..055f6bd0b68 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1287,6 +1287,9 @@ 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" diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 037f913ad07..fbc8ef6cced 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -26,6 +26,7 @@ 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 @@ async def base_process_llm_request( # noqa: PLR0915 ) 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] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9981c049c18..cf837d65a9b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1608,34 +1608,6 @@ async def root_redirect(): ### 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 diff --git a/tests/test_litellm/proxy/test_client_disconnection.py b/tests/test_litellm/proxy/test_client_disconnection.py new file mode 100644 index 00000000000..1d85bb97032 --- /dev/null +++ b/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() From c5f48ccfabec8a86e652049302719f1de2f1c835 Mon Sep 17 00:00:00 2001 From: CreateRandom Date: Wed, 15 Apr 2026 16:17:04 +0200 Subject: [PATCH 2/4] fix(proxy): address review comments on client disconnect feature --- litellm/proxy/common_request_processing.py | 43 ++++++++++++++----- .../proxy/test_client_disconnection.py | 20 +++++---- 2 files changed, 44 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fbc8ef6cced..7f372df7969 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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 @@ async def base_process_llm_request( # noqa: PLR0915 ) # 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 @@ async def _handle_llm_api_exception( 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, diff --git a/tests/test_litellm/proxy/test_client_disconnection.py b/tests/test_litellm/proxy/test_client_disconnection.py index 1d85bb97032..380ac7f9119 100644 --- a/tests/test_litellm/proxy/test_client_disconnection.py +++ b/tests/test_litellm/proxy/test_client_disconnection.py @@ -3,14 +3,14 @@ """ 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() From 65bed54f2824dcc7c8a367c9c41e9652c46e9b8b Mon Sep 17 00:00:00 2001 From: CreateRandom Date: Thu, 16 Apr 2026 09:46:03 +0200 Subject: [PATCH 3/4] fix(proxy): skip post_call_failure_hook for 499 client disconnect --- litellm/proxy/common_request_processing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7f372df7969..2ac792b04f7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1622,6 +1622,7 @@ async def _handle_llm_api_exception( """Raises ProxyException (OpenAI API compatible) if an exception is raised""" if isinstance(e, HTTPException) and e.status_code == 499: verbose_proxy_logger.info("Client disconnected the request (499)") + raise e else: verbose_proxy_logger.exception( f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}" From c9c03ebe9ac408ef698593628a006bfd4bf32ed4 Mon Sep 17 00:00:00 2001 From: CreateRandom Date: Thu, 16 Apr 2026 09:54:49 +0200 Subject: [PATCH 4/4] style: apply black formatting --- litellm/proxy/common_request_processing.py | 24 +++++++++++-------- .../proxy/test_client_disconnection.py | 13 ++++++---- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2ac792b04f7..33a69000223 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -900,9 +900,11 @@ def _debug_log_request_payload(self) -> None: "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Keys: %s", len(_payload_str), MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, - list(self.data.keys()) - if isinstance(self.data, dict) - else type(self.data).__name__, + ( + list(self.data.keys()) + if isinstance(self.data, dict) + else type(self.data).__name__ + ), ) else: verbose_proxy_logger.debug( @@ -1181,9 +1183,9 @@ async def base_process_llm_request( # noqa: PLR0915 # aliasing/routing, but the OpenAI-compatible response `model` field should reflect # what the client sent. if requested_model_from_client: - self.data[ - "_litellm_client_requested_model" - ] = requested_model_from_client + self.data["_litellm_client_requested_model"] = ( + requested_model_from_client + ) # Streaming: attach a closure that fires after all guardrail # end-of-stream blocks complete. CSW.__anext__ stores the @@ -1788,7 +1790,9 @@ async def async_streaming_data_generator( verbose_proxy_logger.debug("inside generator") try: str_so_far = "" - async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( + async for ( + chunk + ) in proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, request_data=request_data, @@ -2016,9 +2020,9 @@ def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs[ - "cache_creation_input_tokens" - ] = cache_creation_input_tokens + usage_kwargs["cache_creation_input_tokens"] = ( + cache_creation_input_tokens + ) if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/tests/test_litellm/proxy/test_client_disconnection.py b/tests/test_litellm/proxy/test_client_disconnection.py index 380ac7f9119..bd0b3d5042e 100644 --- a/tests/test_litellm/proxy/test_client_disconnection.py +++ b/tests/test_litellm/proxy/test_client_disconnection.py @@ -1,6 +1,7 @@ """ Test client disconnection detection functionality. """ + import asyncio import pytest from unittest.mock import AsyncMock, MagicMock, patch @@ -13,15 +14,19 @@ async def test_check_request_disconnection_with_disconnect(): """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 - {"type": "http.disconnect"} # Second call - disconnect + {"type": "http.request"}, # First call + {"type": "http.disconnect"}, # Second call - disconnect ] mock_llm_task = MagicMock() # sync mock so .cancel() doesn't return a coroutine disconnect_event = asyncio.Event() - with patch("litellm.proxy.common_request_processing.asyncio.sleep", new_callable=AsyncMock): - await _check_request_disconnection(mock_request, mock_llm_task, disconnect_event) + 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()