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 @@ 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, 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 @@ 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()