diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e812b03385d..5c8c275c607 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2021,7 +2021,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: @@ -2040,21 +2040,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( @@ -2670,12 +2667,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 5c27ad8ce77..58e6280cc71 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -19,6 +19,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 @@ -5954,6 +5955,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( [