Merge e20706a830 into aeede53112

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,

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(
             [