Merge 132f3cbfe0 into 7f1b0a4633

litellm/proxy/common_request_processing.py

@@ -63,6 +63,7 @@ from litellm.litellm_core_utils.bug_report import (
     strip_bug_report_notice,
 )
 from litellm.litellm_core_utils.core_helpers import (
+    as_str_mapping,
     finite_cost,
     get_or_create_metadata_bucket,
     independent_snapshot,
@@ -2019,6 +2020,17 @@ def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
     )
 
 
+def _model_group_cancels_on_disconnect(llm_router: Router, model: object) -> bool:
+    if not isinstance(model, str):
+        return False
+    group: Final = llm_router.get_model_from_alias(model=model) or model
+    return any(
+        info.get("cancel_on_disconnect") is True
+        for deployment in llm_router.get_model_list_of_routed_group(group)
+        if (info := as_str_mapping(_REQUEST_FIELDS.validate_python(deployment).get("model_info"))) is not None
+    )
+
+
 async def _cancel_llm_call_on_client_disconnect(
     request: Request,
     llm_api_call: "asyncio.Future[_LlmCallT]",
@@ -3138,7 +3150,13 @@ class ProxyBaseLLMRequestProcessing:
         llm_responses: Final = asyncio.gather(*tasks)  # run the moderation check in parallel to the actual llm api call
 
         try:
-            if general_settings.get("cancel_on_disconnect", False):
+            if general_settings.get("cancel_on_disconnect", False) or (
+                llm_router is not None
+                and _model_group_cancels_on_disconnect(
+                    llm_router,
+                    _REQUEST_FIELDS.validate_python(self.data).get("model"),  # pyright: ignore[reportUnknownMemberType]  # legacy request state is validated by the adapter
+                )
+            ):
                 responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses)
             else:
                 responses = await llm_responses

litellm/types/router.py

@@ -294,6 +294,10 @@ class ModelInfo(MirroredPricingParams):
     enforce_supported_endpoints: bool | None = None
     id: str | None  # Allow id to be optional on input, but it will always be present as a str in the model instance
     db_model: bool = False  # used for proxy - to separate models which are stored in the db vs. config.
+    cancel_on_disconnect: bool = Field(
+        False,
+        description="Opt this model group into cancelling non-streaming upstream calls when the client disconnects. Enabled if any deployment in the group opts in. A global cancel_on_disconnect setting still applies to every group.",
+    )
     updated_at: datetime.datetime | None = None
     updated_by: str | None = None
 

tests/unit/proxy/test_common_request_processing.py

@@ -5954,6 +5954,129 @@ class TestCancelOnDisconnect:
 
         return Request(scope={"type": "http", "headers": []}, receive=receive)
 
+    @pytest.mark.parametrize(
+        "model,group_enabled,global_enabled,expect_cancel",
+        [
+            ("interactive", True, False, True),
+            ("batch", True, False, False),
+            ("interactive", False, False, False),
+            ("batch", False, True, True),
+            ("alias", True, False, True),
+            ("wild/selected", True, False, True),
+        ],
+    )
+    async def test_model_group_disconnect_cancels_only_opted_in_work(
+        self,
+        model: str,
+        group_enabled: bool,
+        global_enabled: bool,
+        expect_cancel: bool,
+        respx_mock: respx.MockRouter,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.proxy.proxy_server import ProxyConfig as RuntimeProxyConfig
+
+        monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+        started: Final = asyncio.Event()
+        cancelled: Final = asyncio.Event()
+        messages: Final = [{"role": "user", "content": "disconnect control"}]
+        parameters: Final = {
+            "model": "openai/gpt-6.1-sol",
+            "api_key": "synthetic-key",
+            "api_base": "https://disconnect.example/v1",
+        }
+        router: Final = Router(
+            model_list=[
+                {
+                    "model_name": "interactive",
+                    "litellm_params": parameters,
+                    "model_info": {"id": "interactive-passive", "cancel_on_disconnect": False},
+                },
+                {
+                    "model_name": "interactive",
+                    "litellm_params": parameters,
+                    "model_info": {"cancel_on_disconnect": group_enabled},
+                },
+                {"model_name": "batch", "litellm_params": parameters, "model_info": {"cancel_on_disconnect": False}},
+                {
+                    "model_name": "wild/*",
+                    "litellm_params": parameters,
+                    "model_info": {"cancel_on_disconnect": group_enabled},
+                },
+            ],
+            model_group_alias={"alias": "interactive"},
+            num_retries=0,
+        )
+
+        async def upstream(_: httpx.Request) -> httpx.Response:
+            started.set()
+            try:
+                await asyncio.sleep(0)
+                return httpx.Response(
+                    200,
+                    json={
+                        "id": "completion-disconnect",
+                        "object": "chat.completion",
+                        "created": 1,
+                        "model": "gpt-6.1-sol",
+                        "choices": [
+                            {
+                                "index": 0,
+                                "message": {"role": "assistant", "content": "completed"},
+                                "finish_reason": "stop",
+                            }
+                        ],
+                        "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
+                    },
+                )
+            except asyncio.CancelledError:
+                cancelled.set()
+                raise
+
+        async def receive() -> Mapping[str, object]:
+            await started.wait()
+            return {"type": "http.disconnect"}
+
+        respx_mock.post("https://disconnect.example/v1/chat/completions").mock(side_effect=upstream)
+        logging_obj: Final = Logging(
+            model=model,
+            messages=messages,
+            stream=False,
+            call_type="acompletion",
+            start_time=datetime.datetime(2026, 10, 1),
+            litellm_call_id="disconnect-control",
+            function_id="disconnect-control",
+        )
+        processor: Final = ProxyBaseLLMRequestProcessing(
+            data={
+                "model": model,
+                "messages": messages,
+                "litellm_logging_obj": logging_obj,
+                "metadata": {"cancel_on_disconnect": True, "model_info": {"cancel_on_disconnect": True}},
+            }
+        )
+        call: Final = processor.base_process_llm_request(
+            request=Request(scope={"type": "http", "headers": []}, receive=receive),
+            fastapi_response=Response(),
+            user_api_key_dict=ProxyUserAPIKeyAuth(api_key="synthetic-key"),
+            route_type="acompletion",
+            proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
+            general_settings={"cancel_on_disconnect": global_enabled},
+            proxy_config=RuntimeProxyConfig(),
+            llm_router=router,
+            skip_pre_call_logic=True,
+        )
+        if expect_cancel:
+            with pytest.raises(HTTPException) as error:
+                await call
+            assert error.value.status_code == 499
+        else:
+            response: Final = await call
+            assert response.choices[0].message.content == "completed"
+        assert cancelled.is_set() is expect_cancel
+
+
     async def test_monitor_cancels_llm_call_and_sets_event_on_disconnect(self):
         request = self._request(
             [

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -52171,6 +52171,12 @@ export interface components {
             cache_creation_input_token_cost_above_1hr?: number | null;
             /** Cache Read Input Token Cost */
             cache_read_input_token_cost?: number | null;
+            /**
+             * Cancel On Disconnect
+             * @description Opt this model group into cancelling non-streaming upstream calls when the client disconnects. Enabled if any deployment in the group opts in. A global cancel_on_disconnect setting still applies to every group.
+             * @default false
+             */
+            cancel_on_disconnect: boolean;
             /** Cost Per Ptu Per Hour */
             cost_per_ptu_per_hour?: number | null;
             /** Created At */