diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 23faa964d5f..12cffd9eabe 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -57,6 +57,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, @@ -1859,6 +1860,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]", @@ -2939,7 +2951,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 diff --git a/litellm/types/router.py b/litellm/types/router.py index cd1389aea9d..f151eab88d9 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -267,6 +267,10 @@ def _whole_float_as_int(value: object) -> object: class ModelInfo(MirroredPricingParams): 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 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 92bf2b45fab..c4d2e605a47 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import anthropic import httpx import pytest +import respx from fastapi import HTTPException, Request, Response, status from fastapi.responses import JSONResponse, StreamingResponse @@ -5634,6 +5635,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( [ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bd08b743c7d..c5842d432e8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50234,6 +50234,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 */