From 2e5ebf826fd02aecb4e4b2d005465cffe859a170 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 13 May 2026 09:03:13 -0700 Subject: [PATCH] fix(responses): register cooldowns on failure + fail fast on stale encrypted_content (#27820) --- litellm/responses/main.py | 7 +- .../encrypted_content_affinity_check.py | 144 ++++++++- .../test_responses_router_cooldown.py | 88 +++++ .../test_encrypted_content_affinity_check.py | 303 +++++++++++++++++- 4 files changed, 522 insertions(+), 20 deletions(-) create mode 100644 tests/test_litellm/responses/test_responses_router_cooldown.py diff --git a/litellm/responses/main.py b/litellm/responses/main.py index b6dc5afb94..4ee9235af7 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1128,7 +1128,6 @@ def responses( ) ) - # Pre Call logging litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, @@ -1138,6 +1137,12 @@ def responses( **responses_api_request_params, "aresponses": _is_async, "litellm_call_id": litellm_call_id, + "model_info": kwargs.get("model_info"), + "metadata": ( + kwargs["litellm_metadata"] + if "litellm_metadata" in kwargs + else kwargs.get("metadata") + ), }, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 866073c84d..4ed19c5cd2 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -36,11 +36,20 @@ Safe to enable globally: - No cache required. """ +import time from typing import TYPE_CHECKING, Any, List, Optional, cast +import httpx + from litellm._logging import verbose_router_logger +from litellm.exceptions import ( + BadRequestError, + RateLimitError, + ServiceUnavailableError, +) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues if TYPE_CHECKING: @@ -153,27 +162,31 @@ class EncryptedContentAffinityCheck(CustomLogger): self, healthy_deployments: List[dict], model_id: str, - ) -> List[dict]: + ) -> tuple[List[dict], Any]: """ Deployments in ``healthy_deployments`` sharing the originating - deployment's ``(api_base, api_key)``. Returns ``[]`` if router isn't - wired in, the originating deployment was removed, or no boundary match. + deployment's ``(api_base, api_key)``, alongside the originating + deployment object (or ``None`` if it was removed / router unavailable). + Returns ``([], originating_or_None)`` when no boundary match exists, + so the caller can reuse the looked-up ``originating`` rather than + re-querying the router. """ if self.router is None: - return [] + return [], None originating = self.router.get_deployment(model_id=model_id) if originating is None: - return [] + return [], None boundary = self._encryption_boundary_key( originating.litellm_params.model_dump(exclude_none=True) ) if boundary is None: - return [] - return [ + return [], originating + matches = [ d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary ] + return matches, originating # ------------------------------------------------------------------ # Request routing (pre-call filter) @@ -189,7 +202,14 @@ class EncryptedContentAffinityCheck(CustomLogger): ) -> List[dict]: """ If the request ``input`` contains litellm-encoded item IDs, decode the - embedded ``model_id`` and pin the request to that deployment. + embedded ``model_id`` and pin the request to that deployment. Raises + ``RateLimitError`` / ``ServiceUnavailableError`` / ``BadRequestError`` + when the originating deployment is unavailable and no encryption-boundary + peer exists, rather than dispatching a doomed request to a non-peer + deployment. The 429/503 split mirrors the originating cooldown's status: + a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the + remaining cooldown window) so OpenAI-compatible clients back off and + retry after the deployment is eligible again. """ request_kwargs = request_kwargs or {} typed_healthy_deployments = cast(List[dict], healthy_deployments) @@ -229,9 +249,11 @@ class EncryptedContentAffinityCheck(CustomLogger): return [deployment] # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. - boundary_matches = self._find_deployments_on_same_encryption_boundary( - healthy_deployments=typed_healthy_deployments, - model_id=model_id, + boundary_matches, originating = ( + self._find_deployments_on_same_encryption_boundary( + healthy_deployments=typed_healthy_deployments, + model_id=model_id, + ) ) if boundary_matches: verbose_router_logger.debug( @@ -243,10 +265,98 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["_encrypted_content_affinity_pinned"] = True return boundary_matches - verbose_router_logger.error( - "EncryptedContentAffinityCheck: decoded deployment=%s not found in " - "healthy_deployments and no boundary match available; falling back to " - "full deployment pool (encrypted_content may be rejected upstream)", - model_id, + # Dispatching to a non-peer would guarantee an upstream + # `invalid_encrypted_content` 400, so fail fast with a clearer error. + raise await self._unavailable_origin_error( + model=model, + model_id=model_id, + originating=originating, + parent_otel_span=parent_otel_span, ) - return typed_healthy_deployments + + async def _unavailable_origin_error( + self, + model: str, + model_id: str, + originating: Any, + parent_otel_span: Optional[Span], + ) -> Exception: + # Public error messages intentionally omit the originating ``model_id`` so + # an authenticated caller forging encrypted-content markers cannot use the + # error surface to enumerate which deployment IDs exist on this router. + if originating is None: + return BadRequestError( + message=( + "The deployment that produced this encrypted_content is no " + "longer configured on this router, and no deployment on the " + "same encryption boundary is available. Re-issue the request " + "without the stale encrypted_content items, or restore the " + "originating deployment." + ), + model=model, + llm_provider="", + ) + + cooldown = await self._get_origin_cooldown( + model_id=model_id, parent_otel_span=parent_otel_span + ) + + if cooldown is not None and str(cooldown.get("status_code")) == "429": + retry_after = self._cooldown_seconds_remaining(cooldown) + return RateLimitError( + message=( + "The deployment that produced this encrypted_content is " + f"rate-limited (cooling down for ~{retry_after}s), and no " + "deployment on the same encryption boundary is configured. " + "Retry after the Retry-After window or configure a deployment " + "with the same (api_base, api_key)." + ), + llm_provider="", + model=model, + response=httpx.Response( + status_code=429, + headers={"retry-after": str(retry_after)}, + request=httpx.Request("POST", "https://litellm.ai/"), + ), + ) + + return ServiceUnavailableError( + message=( + "The deployment that produced this encrypted_content is " + "currently unavailable (likely cooled down), and no deployment " + "on the same encryption boundary is configured. Retry later or " + "configure a deployment with the same (api_base, api_key)." + ), + llm_provider="", + model=model, + ) + + async def _get_origin_cooldown( + self, + model_id: str, + parent_otel_span: Optional[Span], + ) -> Optional[CooldownCacheValue]: + if self.router is None: + return None + cooldown_cache = getattr(self.router, "cooldown_cache", None) + if cooldown_cache is None: + return None + try: + active = await cooldown_cache.async_get_active_cooldowns( + model_ids=[model_id], parent_otel_span=parent_otel_span + ) + except Exception: + return None + for cached_model_id, value in active: + if cached_model_id == model_id: + return value + return None + + @staticmethod + def _cooldown_seconds_remaining(cooldown: CooldownCacheValue) -> int: + remaining = ( + float(cooldown.get("timestamp", 0.0)) + + float(cooldown.get("cooldown_time", 0.0)) + - time.time() + ) + return max(1, int(remaining)) diff --git a/tests/test_litellm/responses/test_responses_router_cooldown.py b/tests/test_litellm/responses/test_responses_router_cooldown.py new file mode 100644 index 0000000000..48e2d2455e --- /dev/null +++ b/tests/test_litellm/responses/test_responses_router_cooldown.py @@ -0,0 +1,88 @@ +""" +Regression: Responses API router must register cooldowns on deployment +failures. Previously the Responses API path built ``litellm_params`` without +``model_info``, so ``Router.deployment_callback_on_failure`` exited early via +the "No model_info found" branch and the failing deployment was never added +to the cooldown set. +""" + +import os +import sys +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments + + +@pytest.mark.asyncio +async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): + failing_deployment_id = "deployment-rate-limited" + + router = litellm.Router( + model_list=[ + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-1", + }, + "model_info": {"id": failing_deployment_id}, + }, + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-2", + }, + "model_info": {"id": "deployment-healthy"}, + }, + ], + num_retries=0, + cooldown_time=60, + ) + + rate_limit_error = litellm.RateLimitError( + message="upstream throttled", + llm_provider="openai", + model="openai/gpt-5.1-codex", + response=httpx.Response( + status_code=429, + request=httpx.Request("POST", "https://api.openai.com/v1/responses"), + ), + ) + + def pin_to_failing_deployment(seq): + for d in seq: + if d["model_info"]["id"] == failing_deployment_id: + return d + return seq[0] + + with ( + patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new_callable=AsyncMock, + side_effect=rate_limit_error, + ), + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=pin_to_failing_deployment, + ), + ): + with pytest.raises(litellm.RateLimitError): + await router.aresponses( + model="openai.gpt-5.1-codex", + input="hi", + ) + + cooldown_ids = await _async_get_cooldown_deployments( + litellm_router_instance=router, parent_otel_span=None + ) + assert failing_deployment_id in cooldown_ids, ( + f"Responses API failure callback did not register cooldown for " + f"{failing_deployment_id!r}; cooldown set was {cooldown_ids}" + ) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 1f6848d3b5..07d894d040 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -17,6 +17,8 @@ The mechanism works without any cache and supports two encoding strategies: import os import sys +import time +from typing import List, Optional from unittest.mock import AsyncMock, patch import pytest @@ -1067,11 +1069,12 @@ def test_boundary_fallback_no_router_ref_returns_empty(): "litellm_params": {"api_base": "https://x", "api_key": "k"}, } ] - matches = check._find_deployments_on_same_encryption_boundary( + matches, originating = check._find_deployments_on_same_encryption_boundary( healthy_deployments=healthy, model_id="dep-2", ) assert matches == [] + assert originating is None def test_boundary_fallback_originating_deployment_removed_returns_empty(): @@ -1096,11 +1099,12 @@ def test_boundary_fallback_originating_deployment_removed_returns_empty(): "litellm_params": {"api_base": "https://x", "api_key": "k"}, } ] - matches = check._find_deployments_on_same_encryption_boundary( + matches, originating = check._find_deployments_on_same_encryption_boundary( healthy_deployments=healthy, model_id="dep-removed", ) assert matches == [] + assert originating is None mock_router.get_deployment.assert_called_once_with(model_id="dep-removed") @@ -1172,3 +1176,298 @@ def test_boundary_key_rejects_non_dict_like_inputs(): ) is None ) + + +# --------------------------------------------------------------------------- +# Fail-fast when originating deployment is unavailable and no boundary peer +# --------------------------------------------------------------------------- + + +def _make_originating_mock(api_base: str, api_key: str): + from unittest.mock import MagicMock + + originating = MagicMock() + originating.litellm_params.model_dump.return_value = { + "api_base": api_base, + "api_key": api_key, + } + return originating + + +def _make_router_mock_with_cooldown( + originating, cooldown_entries: Optional[List[tuple]] = None +): + """ + Build a MagicMock router whose ``cooldown_cache.async_get_active_cooldowns`` + returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown). + """ + from unittest.mock import AsyncMock, MagicMock + + mock_router = MagicMock() + mock_router.get_deployment.return_value = originating + mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock( + return_value=list(cooldown_entries or []) + ) + return mock_router + + +@pytest.mark.asyncio +async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_429(): + """ + Originating deployment is in the router config, in cooldown for a non-429 + cause (e.g. a 500), and no boundary peer is configured. The check must + surface this as a 503 (transient, but not rate-limit-specific) rather than + dispatching to a non-peer deployment. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") + mock_router = _make_router_mock_with_cooldown( + originating, + cooldown_entries=[ + ( + "deployment-a-cooled", + { + "exception_received": "boom", + "status_code": "500", + "timestamp": time.time(), + "cooldown_time": 60.0, + }, + ) + ], + ) + + check = EncryptedContentAffinityCheck(router=mock_router) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-a-cooled", "rs_test" + ) + healthy_only_b = [ + { + "model_info": {"id": "deployment-b"}, + "litellm_params": { + "api_base": "https://account-b.openai.azure.com/", + "api_key": "key-b", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = { + "input": [{"id": encoded_id, "type": "reasoning"}], + } + + with pytest.raises(ServiceUnavailableError) as excinfo: + await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) + + # Public error message intentionally omits the originating model_id to + # avoid an authenticated-caller probing oracle. + assert "deployment-a-cooled" not in str(excinfo.value) + assert excinfo.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_for_429(): + """ + Originating deployment is in cooldown specifically because of a 429. + The check must surface this as a 429 RateLimitError with a Retry-After + header derived from the cooldown's remaining window, so OpenAI-compatible + clients respect the backoff instead of giving up on a 503. + """ + from litellm.exceptions import RateLimitError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") + cooldown_started = time.time() - 5.0 + mock_router = _make_router_mock_with_cooldown( + originating, + cooldown_entries=[ + ( + "deployment-a-cooled-429", + { + "exception_received": "rate limited", + "status_code": "429", + "timestamp": cooldown_started, + "cooldown_time": 60.0, + }, + ) + ], + ) + + check = EncryptedContentAffinityCheck(router=mock_router) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-a-cooled-429", "rs_test" + ) + healthy_only_b = [ + { + "model_info": {"id": "deployment-b"}, + "litellm_params": { + "api_base": "https://account-b.openai.azure.com/", + "api_key": "key-b", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = { + "input": [{"id": encoded_id, "type": "reasoning"}], + } + + with pytest.raises(RateLimitError) as excinfo: + await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) + + assert "deployment-a-cooled-429" not in str(excinfo.value) + assert excinfo.value.status_code == 429 + retry_after = excinfo.value.response.headers.get("retry-after") + assert retry_after is not None + assert 1 <= int(retry_after) <= 60 + + +@pytest.mark.asyncio +async def test_affinity_raises_service_unavailable_when_origin_filtered_without_cooldown_entry(): + """ + Originating deployment is configured but absent from healthy_deployments + with no active cooldown entry. Surface as 503 (we cannot prove the cause + was rate-limiting) rather than guessing 429. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") + mock_router = _make_router_mock_with_cooldown(originating, cooldown_entries=[]) + + check = EncryptedContentAffinityCheck(router=mock_router) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-a-filtered", "rs_test" + ) + healthy_only_b = [ + { + "model_info": {"id": "deployment-b"}, + "litellm_params": { + "api_base": "https://account-b.openai.azure.com/", + "api_key": "key-b", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = { + "input": [{"id": encoded_id, "type": "reasoning"}], + } + + with pytest.raises(ServiceUnavailableError) as excinfo: + await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) + + assert excinfo.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_affinity_raises_bad_request_when_origin_removed(): + """ + Originating deployment was removed from the router config and no boundary + peer is available. This is permanent (the stale encrypted_content cannot + be honored), so surface a 400 with actionable text. + """ + from unittest.mock import MagicMock + + from litellm.exceptions import BadRequestError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + mock_router = MagicMock() + mock_router.get_deployment.return_value = None + + check = EncryptedContentAffinityCheck(router=mock_router) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-removed", "rs_test" + ) + healthy_only_b = [ + { + "model_info": {"id": "deployment-b"}, + "litellm_params": { + "api_base": "https://account-b.openai.azure.com/", + "api_key": "key-b", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = { + "input": [{"id": encoded_id, "type": "reasoning"}], + } + + with pytest.raises(BadRequestError) as excinfo: + await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) + + assert "deployment-removed" not in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_affinity_does_not_raise_when_boundary_peer_available(): + """ + Even when the originating deployment is filtered out, if a peer on the + same (api_base, api_key) is in healthy_deployments, the boundary-match + path must succeed silently — no exception. + """ + from unittest.mock import MagicMock + + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = MagicMock() + originating.litellm_params.model_dump.return_value = { + "api_base": "https://account-a.openai.azure.com/", + "api_key": "key-a", + } + mock_router = MagicMock() + mock_router.get_deployment.return_value = originating + + check = EncryptedContentAffinityCheck(router=mock_router) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-a", "rs_test" + ) + peer = { + "model_info": {"id": "deployment-a-peer"}, + "litellm_params": { + "api_base": "https://account-a.openai.azure.com/", + "api_key": "key-a", + "model": "azure/gpt-5.4", + }, + } + request_kwargs = { + "input": [{"id": encoded_id, "type": "reasoning"}], + } + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=[peer], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [peer] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True