From e20706a830b4988f2fb0556c6594cb449192645a Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 29 Oct 2026 20:02:14 -0700 Subject: [PATCH] fix(proxy): cancel disconnected requests waiting in pre-call hooks --- litellm/proxy/common_request_processing.py | 32 ++++---- .../proxy/test_common_request_processing.py | 74 ++++++++++++++++++- 2 files changed, 90 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7cd752c2297..5962cd05a25 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1921,7 +1921,7 @@ def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None: async def _cancel_llm_call_on_client_disconnect( request: Request, - llm_api_call: "asyncio.Future[_LlmCallT]", + llm_api_call: "asyncio.Future[_LlmCallT] | anyio.CancelScope", disconnect_event: asyncio.Event, ) -> None: try: @@ -1940,21 +1940,18 @@ async def _cancel_llm_call_on_client_disconnect( async def _await_llm_call_cancelling_on_disconnect( request: Request, - llm_api_call: "asyncio.Future[_LlmCallT]", + llm_api_call: Awaitable[_LlmCallT], ) -> _LlmCallT: disconnect_event: Final = asyncio.Event() - monitor: Final = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event)) - try: - return await llm_api_call - except asyncio.CancelledError: - if disconnect_event.is_set(): - raise HTTPException( - status_code=499, - detail=_CLIENT_DISCONNECT_DETAIL, - ) - raise - finally: - monitor.cancel() + with anyio.CancelScope() as cancel_scope: + monitor: Final = asyncio.create_task( + _cancel_llm_call_on_client_disconnect(request, cancel_scope, disconnect_event) + ) + try: + return await llm_api_call + finally: + monitor.cancel() + raise HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL) def _timing_values( @@ -2587,12 +2584,17 @@ class ProxyBaseLLMRequestProcessing: prefetch_model = self.data.get("model") if llm_router is not None and isinstance(prefetch_model, str): RoutingContext(llm_router).arm_read_prefetch(prefetch_model, self.data) - self.data = await proxy_logging_obj.pre_call_hook( + pre_call: Final = proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type, skip_guardrails=skip_guardrails, ) + self.data = ( + await _await_llm_call_cancelling_on_disconnect(request, pre_call) + if general_settings.get("cancel_on_disconnect", False) + else await pre_call + ) await _enforce_guardrail_added_tag_budgets( data=self.data, tags_before_guardrails=self._tags_before_guardrails, diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index aeab7a3d348..dc61adb64be 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -17,6 +17,7 @@ import pytest import respx from fastapi import HTTPException, Request, Response, status from fastapi.responses import JSONResponse, StreamingResponse +from starlette.types import Message import litellm from litellm._uuid import uuid @@ -5733,6 +5734,74 @@ class TestCancelOnDisconnect: return Request(scope={"type": "http", "headers": []}, receive=receive) + @pytest.mark.parametrize("enabled", (True, False)) + async def test_disconnect_cancels_a_waiting_pre_call_hook( + self, monkeypatch: pytest.MonkeyPatch, enabled: bool + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + messages: Final[asyncio.Queue[Message]] = asyncio.Queue() + + class WaitingHook(CustomLogger): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> dict[str, object]: + entered.set() + try: + await release.wait() + except asyncio.CancelledError: + cancelled.set() + raise + return data + + monkeypatch.setattr(litellm, "callbacks", [WaitingHook()]) + request: Final = Request( + scope={"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}, + receive=messages.get, + ) + processor: Final = ProxyBaseLLMRequestProcessing( + data={"model": "openai/test-model", "messages": [{"role": "user", "content": "hello"}]} + ) + task: Final = asyncio.create_task( + processor.common_processing_pre_call_logic( + request=request, + general_settings={"cancel_on_disconnect": enabled}, + user_api_key_dict=ProxyUserAPIKeyAuth(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=RealProxyConfig(), + route_type="acompletion", + ) + ) + try: + await asyncio.wait_for(entered.wait(), 5) + messages.put_nowait({"type": "http.disconnect"}) + if enabled: + with pytest.raises(HTTPException) as exc_info: + await asyncio.wait_for(task, 5) + assert exc_info.value.status_code == 499 + assert cancelled.is_set() + else: + release.set() + data, _ = await asyncio.wait_for(task, 5) + assert data["model"] == "openai/test-model" + assert not cancelled.is_set() + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + async def test_disconnect_watcher_preserves_hook_context_changes(self) -> None: + value: Final = contextvars.ContextVar("hook-context", default="before") + + async def pre_call() -> str: + value.set("after") + return "completed" + + assert await _await_llm_call_cancelling_on_disconnect(self._request([]), pre_call()) == "completed" + assert value.get() == "after" + async def test_monitor_cancels_llm_call_and_sets_event_on_disconnect(self): request = self._request( [ @@ -7667,7 +7736,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-synthetic"}, "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, }, - {"model_name": "chat-model", "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "sk-synthetic"}}, + { + "model_name": "chat-model", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "sk-synthetic"}, + }, ], fallbacks=[{"rate-limited": ["embed-alias", "chat-model"]}], )