From 9912e4b2712e84580b146bf6476cf9fe54b753b1 Mon Sep 17 00:00:00 2001 From: aidanasingh Date: Mon, 6 Jul 2026 22:29:05 -0400 Subject: [PATCH 1/5] fix(router): apply weighted-failover exclusion before deployment-affinity callback --- litellm/router.py | 18 ++--- .../test_router_weighted_failover.py | 65 +++++++++++++++++++ 2 files changed, 74 insertions(+), 9 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 12b96430334..02e3de07c52 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10297,6 +10297,15 @@ class Router: healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in + ## this request via weighted-failover. Always honored, regardless of the + ## router-level flag, so a stale exclusion key on kwargs cannot escape. + _excluded_deployment_ids = (request_kwargs or {}).pop("_excluded_deployment_ids", None) + healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + cast(List[Dict], healthy_deployments), + excluded_deployment_ids=_excluded_deployment_ids, + ) + healthy_deployments = await self.async_callback_filter_deployments( model=model, healthy_deployments=healthy_deployments, @@ -10327,15 +10336,6 @@ class Router: cast(List[Dict], healthy_deployments), target_order=_target_order ) - ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in - ## this request via weighted-failover. Always honored, regardless of the - ## router-level flag, so a stale exclusion key on kwargs cannot escape. - _excluded_deployment_ids = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( - cast(List[Dict], healthy_deployments), - excluded_deployment_ids=_excluded_deployment_ids, - ) - if len(healthy_deployments) == 0: exception = await async_raise_no_deployment_exception( litellm_router_instance=self, diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 8faf6bcd9cf..fbd715e90b7 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -13,7 +13,11 @@ from unittest.mock import AsyncMock, patch import pytest +import litellm from litellm import Router +from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, +) from litellm.utils import _get_excluded_filtered_deployments @@ -769,3 +773,64 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo ) assert response._hidden_params["model_id"] == "fallback" + + +@pytest.mark.asyncio +async def test_weighted_failover_exclusion_runs_before_affinity_callback(): + """Regression: the weighted-failover exclusion must be applied before the + deployment-affinity callback in async_get_healthy_deployments. Otherwise a + session pinned to the just-failed deployment collapses the candidate set to + that one deployment, the exclusion then empties it, and the request escapes + to cross-group fallback even though a healthy sibling in the same group is + available.""" + callbacks_snapshot = list(litellm.callbacks) + try: + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key"}, + "model_info": {"id": "dep-b"}, + }, + ], + routing_strategy="simple-shuffle", + enable_weighted_failover=True, + optional_pre_call_checks=["session_affinity"], + ) + + affinity_check = next( + c + for c in (router.optional_callbacks or []) + if isinstance(c, DeploymentAffinityCheck) + ) + session_id = "sess-abc" + await affinity_check.cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="test-model", session_id=session_id + ), + {"model_id": "dep-a"}, + ) + + pinned = await router.async_get_healthy_deployments( + model="test-model", + request_kwargs={"metadata": {"session_id": session_id}}, + messages=[{"role": "user", "content": "hi"}], + ) + assert {d["model_info"]["id"] for d in pinned} == {"dep-a"} + + rerouted = await router.async_get_healthy_deployments( + model="test-model", + request_kwargs={ + "metadata": {"session_id": session_id}, + "_excluded_deployment_ids": ["dep-a"], + }, + messages=[{"role": "user", "content": "hi"}], + ) + assert {d["model_info"]["id"] for d in rerouted} == {"dep-b"} + finally: + litellm.callbacks[:] = callbacks_snapshot From c55e0ba1db0a6b33477bb4a7825523fb9aecab56 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:15:14 -0700 Subject: [PATCH 2/5] fix(router): let affinity pins yield to deployments that already failed this request Moving the weighted failover exclusion ahead of the affinity check also changed what tag, order, plugin and managed-file filters saw. Restore the original filter order and have the affinity check ignore session and API key pins that point at a deployment already excluded or retry-skipped in this request. previous_response_id pins still win. --- litellm/router.py | 18 +++--- .../deployment_affinity_check.py | 28 ++++++++- .../test_deployment_affinity_check.py | 57 +++++++++++++++++++ tests/unit/test_router_weighted_failover.py | 10 ++-- 4 files changed, 98 insertions(+), 15 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 36459b6c105..74facf200d2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14703,15 +14703,6 @@ class Router: healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in - ## this request via weighted-failover. Always honored, regardless of the - ## router-level flag, so a stale exclusion key on kwargs cannot escape. - _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils.get_excluded_filtered_deployments( - healthy_deployments, - excluded_deployment_ids=_excluded_deployment_ids, - ) - healthy_deployments = await self.async_callback_filter_deployments( model=model, healthy_deployments=healthy_deployments, @@ -14791,6 +14782,15 @@ class Router: target_order=_target_order, ) + ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in + ## this request via weighted-failover. Always honored, regardless of the + ## router-level flag, so a stale exclusion key on kwargs cannot escape. + _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( + cast(list[dict], healthy_deployments), + excluded_deployment_ids=_excluded_deployment_ids, + ) + ## RETRY SKIP ## -> drop deployments that already refused this request with a ## non-retryable status, unless that leaves nothing, so the caller still gets ## the provider's own error instead of a no-deployments error. diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 88fb8876175..342a2720f12 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -16,6 +16,7 @@ import hashlib from collections.abc import Callable, Mapping, Sequence from typing import Final, cast +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target @@ -38,6 +39,8 @@ class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] +_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) + VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( { "deployment_affinity", @@ -327,6 +330,24 @@ class DeploymentAffinityCheck(CustomLogger): return stored return None + @staticmethod + def _deployment_ids(value: object) -> tuple[str, ...]: + if value is None: + return () + try: + return _DEPLOYMENT_IDS_ADAPTER.validate_python(value) + except ValidationError: + return () + + @classmethod + def _deployment_ids_failed_in_request(cls, request_kwargs: Mapping[str, object]) -> frozenset[str]: + return frozenset( + ( + *cls._deployment_ids(request_kwargs.get("_excluded_deployment_ids")), + *cls._deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), + ) + ) + def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) @@ -405,6 +426,9 @@ class DeploymentAffinityCheck(CustomLogger): session_affinity_active: Final = ( enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None ) + failed_deployment_ids: Final = self._deployment_ids_failed_in_request( + cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names + ) user_key: Final = ( self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if (session_affinity_active or enable_user_key) @@ -426,7 +450,7 @@ class DeploymentAffinityCheck(CustomLogger): elif isinstance(session_cache_result, str): session_model_id = session_cache_result - if session_model_id: + if session_model_id and session_model_id not in failed_deployment_ids: session_deployment: Final = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, model_id=session_model_id, @@ -461,7 +485,7 @@ class DeploymentAffinityCheck(CustomLogger): # Backwards / safety: allow raw string values. model_id = cache_result - if not model_id: + if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments deployment = self._find_deployment_by_model_id( diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index d8bc4c45ab8..f1393a17cbe 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1209,3 +1209,60 @@ async def test_virtual_key_hash_wins_over_user_id_for_affinity(): parent_otel_span=None, ) assert other_key_same_user == healthy_deployments + + +class _ForbiddenError(Exception): + status_code: Final = 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pre_call_check", "metadata"), + [ + ("session_affinity", {"session_id": "sess-retry"}), + ("deployment_affinity", {"user_api_key_hash": "hashed-key-retry"}), + ], +) +async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_refused( + pre_call_check: str, metadata: dict[str, str] +): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": _ForbiddenError("A refuses"), + "weight": 1, + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from B", + "weight": 0, + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=1, + retry_after=0, + disable_cooldowns=True, + optional_pre_call_checks=[pre_call_check], + ) + + responses: Final = [ + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + for _ in range(2) + ] + + assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index 8bec010e3d0..adc4dd587d6 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -8,7 +8,7 @@ cross-group fallback runs. """ from collections import Counter -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, patch import pytest @@ -954,9 +954,11 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo ("deployment_affinity", {"user_api_key_hash": "hashed-key-1"}), ], ) -async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deployment(pre_call_check, metadata): +async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deployment( + pre_call_check: str, metadata: dict[str, str] +): """A pin on the failed deployment must not hide its healthy sibling from weighted failover.""" - router = Router( + router: Final = Router( model_list=[ { "model_name": "test-model", @@ -996,7 +998,7 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy fallbacks=[{"test-model": ["fallback-model"]}], ) - responses = [ + responses: Final = [ await router.acompletion( model="test-model", messages=[{"role": "user", "content": "hi"}], From 00cc4b67cdec9890243b0aaeea6c46937e17f673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:46:35 -0700 Subject: [PATCH 3/5] fix(router): release affinity pins on deployments that failed this request Delete a session or API key pin whose deployment already failed the request, so the deployment that serves it claims the pin. Prompt caching pins skip failed deployments too, and the proxy strips the internal failed-id fields from client bodies. --- litellm/proxy/litellm_pre_call_utils.py | 2 ++ litellm/router.py | 11 +++---- litellm/router_utils/common_utils.py | 28 +++++++++++++++- .../deployment_affinity_check.py | 30 ++++------------- .../prompt_caching_deployment_check.py | 5 ++- .../unit/proxy/test_litellm_pre_call_utils.py | 2 ++ .../test_deployment_affinity_check.py | 33 +++++++++++++++++++ .../test_prompt_caching_deployment_check.py | 21 ++++++++++++ tests/unit/test_router/test_router.py | 4 +-- tests/unit/test_router_weighted_failover.py | 4 +++ 10 files changed, 106 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index f8e6498502b..96bef31512b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -355,6 +355,8 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "fallback_depth", "_target_order", "attempted_targets", + "_excluded_deployment_ids", + "_retry_skipped_deployment_ids", "proxy_server_request", "standard_logging_object", "secret_fields", diff --git a/litellm/router.py b/litellm/router.py index 74facf200d2..c02b9a0b15d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -209,6 +209,7 @@ from litellm.router_utils.clientside_credential_handler import ( ) from litellm.router_utils.common_utils import ( MID_STREAM_CONTINUATION_KWARG, + as_retry_skipped_deployment_ids, deployments_the_caller_may_use, filter_team_based_models, filter_web_search_deployments, @@ -668,10 +669,6 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) -def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: - return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else () - - def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: if silent_model is None: return () @@ -8730,7 +8727,7 @@ class Router: return () if litellm.should_retry(status_code): return () - already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped) + already_skipped_ids: Final = as_retry_skipped_deployment_ids(already_skipped) skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) verbose_router_logger.debug( "Retry skips deployments that already answered %s to this request: %s", status_code, skipped @@ -14794,7 +14791,7 @@ class Router: ## RETRY SKIP ## -> drop deployments that already refused this request with a ## non-retryable status, unless that leaves nothing, so the caller still gets ## the provider's own error instead of a no-deployments error. - _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None ) healthy_deployments = ( @@ -16028,7 +16025,7 @@ class Router: ) ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. - _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None ) healthy_deployments = ( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index a00d84babf3..a84ea67a2ec 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -4,7 +4,7 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject @@ -38,6 +38,8 @@ _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapt _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) +_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_OBJECT_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment: @@ -57,6 +59,30 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]: return {key: value for key, value in kwargs.items() if key not in ROUTER_ONLY_CALL_KWARGS} +def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: + if not isinstance(value, tuple): + return () + return tuple(item for item in _OBJECT_TUPLE_ADAPTER.validate_python(value) if isinstance(item, str)) + + +def _as_excluded_deployment_ids(value: object) -> tuple[str, ...]: + if value is None: + return () + try: + return _DEPLOYMENT_IDS_ADAPTER.validate_python(value) + except ValidationError: + return () + + +def deployment_ids_failed_in_request(request_kwargs: Mapping[str, object]) -> frozenset[str]: + return frozenset( + ( + *_as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), + *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), + ) + ) + + def request_user_api_key_auth(request_kwargs: Mapping[str, object]) -> object: """The caller's auth object, from whichever metadata bucket this surface writes to.""" metadata_value: Final = request_kwargs.get("metadata") diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 342a2720f12..927f1aa7e16 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -16,7 +16,6 @@ import hashlib from collections.abc import Callable, Mapping, Sequence from typing import Final, cast -from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target @@ -31,6 +30,7 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -39,8 +39,6 @@ class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] -_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) - VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( { "deployment_affinity", @@ -330,24 +328,6 @@ class DeploymentAffinityCheck(CustomLogger): return stored return None - @staticmethod - def _deployment_ids(value: object) -> tuple[str, ...]: - if value is None: - return () - try: - return _DEPLOYMENT_IDS_ADAPTER.validate_python(value) - except ValidationError: - return () - - @classmethod - def _deployment_ids_failed_in_request(cls, request_kwargs: Mapping[str, object]) -> frozenset[str]: - return frozenset( - ( - *cls._deployment_ids(request_kwargs.get("_excluded_deployment_ids")), - *cls._deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), - ) - ) - def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) @@ -426,7 +406,7 @@ class DeploymentAffinityCheck(CustomLogger): session_affinity_active: Final = ( enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None ) - failed_deployment_ids: Final = self._deployment_ids_failed_in_request( + failed_deployment_ids: Final = deployment_ids_failed_in_request( cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names ) user_key: Final = ( @@ -450,7 +430,9 @@ class DeploymentAffinityCheck(CustomLogger): elif isinstance(session_cache_result, str): session_model_id = session_cache_result - if session_model_id and session_model_id not in failed_deployment_ids: + if session_model_id in failed_deployment_ids: + await self.cache.async_delete_cache(key=session_cache_key) + elif session_model_id: session_deployment: Final = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, model_id=session_model_id, @@ -485,6 +467,8 @@ class DeploymentAffinityCheck(CustomLogger): # Backwards / safety: allow raw string values. model_id = cache_result + if model_id in failed_deployment_ids: + await self.cache.async_delete_cache(key=cache_key) if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index bb3475cdcbe..c04bab4d360 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -19,6 +19,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model +from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import ( @@ -152,7 +153,9 @@ class PromptCachingDeploymentCheck(CustomLogger): request_tools, ), ) - if model_id_dict is not None: + if model_id_dict is not None and model_id_dict["model_id"] not in deployment_ids_failed_in_request( + request_data or {} + ): model_id: Final = model_id_dict["model_id"] for deployment in healthy_deployments: if deployment["model_info"]["id"] == model_id: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c2a5b01fe4e..28d62fc3084 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1109,6 +1109,8 @@ async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection ("fallback_depth", True), ("_target_order", 2), ("attempted_targets", ["forged"]), + ("_excluded_deployment_ids", ["pinned-deployment"]), + ("_retry_skipped_deployment_ids", ["pinned-deployment"]), ], ) async def test_add_litellm_data_to_request_strips_forged_fallback_hop_state( diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index f1393a17cbe..a960f9ef5be 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -1266,3 +1267,35 @@ async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_ ] assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] + assert [response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] for response in responses] == [ + 1, + 0, + ] + + +@pytest.mark.asyncio +async def test_previous_response_id_pin_holds_even_after_its_deployment_failed(): + callback: Final = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=True, + ) + healthy_deployments: Final = [ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "A"}}, + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "B"}}, + ] + + assert await callback.async_filter_deployments( + model="gpt-4o", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={ + "previous_response_id": ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="A", response_id="resp_upstream" + ), + "_excluded_deployment_ids": ["A"], + "metadata": {"user_api_key_hash": "hashed-key"}, + }, + parent_otel_span=None, + ) == [healthy_deployments[0]] diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 8d893160780..9487e4d6321 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -313,6 +313,27 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): assert filtered == deployments +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failed_in_request", + [{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}], +) +async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request( + failed_in_request: dict[str, object], +): + cache: Final = DualCache() + check: Final = PromptCachingDeploymentCheck(cache=cache) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + assert await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"), + messages=messages, + request_kwargs=failed_in_request, + ) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d645ab4566d..e3a12b4eb63 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -20980,9 +20980,9 @@ def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( ], ) def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): - from litellm.router import _as_retry_skipped_deployment_ids + from litellm.router_utils.common_utils import as_retry_skipped_deployment_ids - assert _as_retry_skipped_deployment_ids(value) == expected + assert as_retry_skipped_deployment_ids(value) == expected @pytest.mark.parametrize( diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index adc4dd587d6..1808541666f 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -1008,3 +1008,7 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy ] assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] + assert [response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] for response in responses] == [ + 1, + 0, + ] From 6fe7f1bf84f20a5a36e6e03edc41c75b59a9d26e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:13:27 -0700 Subject: [PATCH 4/5] fix(router): keep affinity pins in place and skip failed origins in encrypted content affinity A pin now only yields for the request whose pinned deployment failed. Deleting it moved pins after one 429 or a request-caused 400 and added a Redis call on the routing path. Encrypted content affinity also ignores deployments that failed this request, so a failed origin takes the strip-and-dispatch path. --- litellm/router_utils/common_utils.py | 4 +- .../deployment_affinity_check.py | 6 +-- .../encrypted_content_affinity_check.py | 20 +++++-- .../test_deployment_affinity_check.py | 2 +- .../test_encrypted_content_affinity_check.py | 53 +++++++++++++++++++ tests/unit/test_router_weighted_failover.py | 33 +++++++----- 6 files changed, 94 insertions(+), 24 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index a84ea67a2ec..17fc1c9d39b 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -39,7 +39,6 @@ _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = tuple[Mapping[str, object], ...] ) _DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) -_OBJECT_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment: @@ -62,7 +61,8 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]: def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: if not isinstance(value, tuple): return () - return tuple(item for item in _OBJECT_TUPLE_ADAPTER.validate_python(value) if isinstance(item, str)) + items: Final[tuple[object, ...]] = value + return tuple(item for item in items if isinstance(item, str)) def _as_excluded_deployment_ids(value: object) -> tuple[str, ...]: diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 927f1aa7e16..91f1dabde40 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -430,9 +430,7 @@ class DeploymentAffinityCheck(CustomLogger): elif isinstance(session_cache_result, str): session_model_id = session_cache_result - if session_model_id in failed_deployment_ids: - await self.cache.async_delete_cache(key=session_cache_key) - elif session_model_id: + if session_model_id and session_model_id not in failed_deployment_ids: session_deployment: Final = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, model_id=session_model_id, @@ -467,8 +465,6 @@ class DeploymentAffinityCheck(CustomLogger): # Backwards / safety: allow raw string values. model_id = cache_result - if model_id in failed_deployment_ids: - await self.cache.async_delete_cache(key=cache_key) if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments 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 5d47bf63a38..b376f358f7f 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 @@ -51,6 +51,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_encrypted_reasoning_from_messages, ) from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues from litellm.types.router import Deployment from litellm.utils import get_order_filtered_deployments @@ -182,6 +183,14 @@ class EncryptedContentAffinityCheck(CustomLogger): None, ) + @staticmethod + def _deployment_model_id(deployment: Mapping[str, object]) -> str | None: + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping): + return None + deployment_id: Final = cast(Mapping[str, object], model_info).get("id") # cast-ok: narrowed by isinstance + return None if deployment_id is None else str(deployment_id) + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -345,7 +354,8 @@ class EncryptedContentAffinityCheck(CustomLogger): typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments) if not self._is_enabled_for_model_group(model): return typed_healthy_deployments - target_order: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs).get("_target_order") + typed_request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs) + target_order: Final = typed_request_kwargs.get("_target_order") candidates: Final = ( get_order_filtered_deployments(typed_healthy_deployments, target_order=target_order) if isinstance(target_order, int) @@ -375,8 +385,12 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id, ) + failed_deployment_ids: Final = deployment_ids_failed_in_request(typed_request_kwargs) + untried_candidates: Final = [ + candidate for candidate in candidates if self._deployment_model_id(candidate) not in failed_deployment_ids + ] deployment: Final = self._find_deployment_by_model_id( - healthy_deployments=candidates, + healthy_deployments=untried_candidates, model_id=model_id, ) if deployment is not None: @@ -392,7 +406,7 @@ class EncryptedContentAffinityCheck(CustomLogger): # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary( - healthy_deployments=candidates, + healthy_deployments=untried_candidates, model_id=model_id, ) if boundary_matches: diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index a960f9ef5be..fa75f980cc0 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1269,7 +1269,7 @@ async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_ assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] assert [response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] for response in responses] == [ 1, - 0, + 1, ] diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 87fbf4433a1..da654394dfb 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -2606,3 +2606,56 @@ async def test_affinity_strips_for_team_and_pattern_routes(): assert result is sibling_pool assert not request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failed_in_request", + [{"_excluded_deployment_ids": ["region-a"]}, {"_retry_skipped_deployment_ids": ("region-a",)}], +) +async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request(failed_in_request): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock("https://region-a.example.com/v1", "key-a") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["region-a", "region-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", "region-a") + group = [ + { + "model_info": {"id": "region-a"}, + "model_name": "gpt-5.4", + "litellm_params": {"api_base": "https://region-a.example.com/v1", "api_key": "key-a"}, + }, + { + "model_info": {"id": "region-b"}, + "model_name": "gpt-5.4", + "litellm_params": {"api_base": "https://region-b.example.com/v1", "api_key": "key-b"}, + }, + ] + request_kwargs = { + "litellm_metadata": {}, + "input": [ + { + "type": "reasoning", + "id": ResponsesAPIRequestUtils._build_encrypted_item_id("region-a", "rs_test"), + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + **failed_in_request, + } + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=group, + messages=None, + request_kwargs=request_kwargs, + ) + + assert result is group + assert request_kwargs["input"][0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index 1808541666f..8b7fc81dbde 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -15,6 +15,7 @@ import pytest import litellm from litellm import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.utils import get_excluded_filtered_deployments @@ -957,7 +958,6 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deployment( pre_call_check: str, metadata: dict[str, str] ): - """A pin on the failed deployment must not hide its healthy sibling from weighted failover.""" router: Final = Router( model_list=[ { @@ -998,17 +998,24 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy fallbacks=[{"test-model": ["fallback-model"]}], ) - responses: Final = [ - await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "hi"}], - metadata=dict(metadata), + while_a_is_down: Final = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) + router.upsert_deployment( + Deployment( + model_name="test-model", + litellm_params=LiteLLM_Params(model="gpt-4o", api_key="bad", mock_response="ok from A", weight=1), + model_info=ModelInfo(id="A"), ) - for _ in range(2) - ] + ) + after_a_recovers: Final = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + ) - assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] - assert [response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] for response in responses] == [ - 1, - 0, - ] + assert while_a_is_down._hidden_params["model_id"] == "B" + assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 + assert after_a_recovers._hidden_params["model_id"] == "A" From b1f66dbcb54af702041b2b569ec8e7ca5f71cfeb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:34:21 -0700 Subject: [PATCH 5/5] fix(router): keep encrypted reasoning when a retry can only go back to its origin The encrypted content check now drops failed deployments the way the router does: a weighted failover exclusion always holds, and a retry skip that would leave nothing falls back to the whole group. The proxy also drops client-sent metadata._failover_excluded_ids, and the encrypted content tests use a real Router --- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/router_utils/common_utils.py | 4 +- .../encrypted_content_affinity_check.py | 18 +++- .../unit/proxy/test_litellm_pre_call_utils.py | 2 + .../test_encrypted_content_affinity_check.py | 98 ++++++++++++------- 5 files changed, 80 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 96bef31512b..c1db9951592 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -437,6 +437,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "_pipeline_managed_guardrails", "client_disconnected", "error_information", + "_failover_excluded_ids", PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 17fc1c9d39b..91b9e5cb114 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -65,7 +65,7 @@ def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: return tuple(item for item in items if isinstance(item, str)) -def _as_excluded_deployment_ids(value: object) -> tuple[str, ...]: +def as_excluded_deployment_ids(value: object) -> tuple[str, ...]: if value is None: return () try: @@ -77,7 +77,7 @@ def _as_excluded_deployment_ids(value: object) -> tuple[str, ...]: def deployment_ids_failed_in_request(request_kwargs: Mapping[str, object]) -> frozenset[str]: return frozenset( ( - *_as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), + *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), ) ) 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 b376f358f7f..39f72d01341 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 @@ -51,7 +51,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_encrypted_reasoning_from_messages, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.common_utils import deployment_ids_failed_in_request +from litellm.router_utils.common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids from litellm.types.llms.openai import AllMessageValues from litellm.types.router import Deployment from litellm.utils import get_order_filtered_deployments @@ -385,10 +385,20 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id, ) - failed_deployment_ids: Final = deployment_ids_failed_in_request(typed_request_kwargs) - untried_candidates: Final = [ - candidate for candidate in candidates if self._deployment_model_id(candidate) not in failed_deployment_ids + excluded_ids: Final = frozenset( + as_excluded_deployment_ids(typed_request_kwargs.get("_excluded_deployment_ids")) + ) + retry_skipped_ids: Final = frozenset( + as_retry_skipped_deployment_ids(typed_request_kwargs.get("_retry_skipped_deployment_ids")) + ) + not_excluded_candidates: Final = [ + candidate for candidate in candidates if self._deployment_model_id(candidate) not in excluded_ids ] + untried_candidates: Final = [ + candidate + for candidate in not_excluded_candidates + if self._deployment_model_id(candidate) not in retry_skipped_ids + ] or not_excluded_candidates deployment: Final = self._find_deployment_by_model_id( healthy_deployments=untried_candidates, model_id=model_id, diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 28d62fc3084..d3580c7e3dc 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1195,6 +1195,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "_managed_background_response": True, "_guardrail_pipelines": [{"name": "spoofed"}], "_pipeline_managed_guardrails": ["evaded"], + "_failover_excluded_ids": ["pinned-deployment"], "safe_user_metadata": "kept", } data = { @@ -1253,6 +1254,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "_managed_background_response", "_guardrail_pipelines", "_pipeline_managed_guardrails", + "_failover_excluded_ids", } assert "litellm_metadata" not in updated for stripped_key in stripped_keys: diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index da654394dfb..4c804ccdc09 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -16,6 +16,7 @@ The mechanism works without any cache and supports two encoding strategies: """ import time +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -2608,54 +2609,77 @@ async def test_affinity_strips_for_team_and_pattern_routes(): assert not request_kwargs["input"][1].get("encrypted_content") -@pytest.mark.asyncio -@pytest.mark.parametrize( +def _region_deployment(region: str) -> dict[str, object]: + return { + "model_name": "gpt-5.4", + "litellm_params": { + "model": "openai/gpt-5.4", + "api_base": f"https://{region}.example.com/v1", + "api_key": f"key-{region}", + }, + "model_info": {"id": region}, + } + + +def _reasoning_input_minted_by(origin: str) -> list[dict[str, object]]: + return [ + { + "type": "reasoning", + "id": ResponsesAPIRequestUtils._build_encrypted_item_id(origin, "rs_test"), + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", origin), + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ] + + +_ORIGIN_FAILED_THIS_REQUEST: Final = pytest.mark.parametrize( "failed_in_request", [{"_excluded_deployment_ids": ["region-a"]}, {"_retry_skipped_deployment_ids": ("region-a",)}], ) -async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request(failed_in_request): + + +@pytest.mark.asyncio +@_ORIGIN_FAILED_THIS_REQUEST +async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request( + failed_in_request: dict[str, object], +): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) - originating = _make_originating_mock("https://region-a.example.com/v1", "key-a") - mock_router = _make_router_mock_with_cooldown( - originating, cooldown_entries=[], routed_group_model_ids=["region-a", "region-b"] - ) - check = EncryptedContentAffinityCheck(router=mock_router) - wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", "region-a") - group = [ - { - "model_info": {"id": "region-a"}, - "model_name": "gpt-5.4", - "litellm_params": {"api_base": "https://region-a.example.com/v1", "api_key": "key-a"}, - }, - { - "model_info": {"id": "region-b"}, - "model_name": "gpt-5.4", - "litellm_params": {"api_base": "https://region-b.example.com/v1", "api_key": "key-b"}, - }, - ] - request_kwargs = { - "litellm_metadata": {}, - "input": [ - { - "type": "reasoning", - "id": ResponsesAPIRequestUtils._build_encrypted_item_id("region-a", "rs_test"), - "encrypted_content": wrapped, - "summary": [{"type": "summary_text", "text": "scattering"}], - }, - {"role": "user", "content": "and sunsets?"}, - ], - **failed_in_request, - } + router: Final = litellm.Router(model_list=[_region_deployment("region-a"), _region_deployment("region-b")]) + group: Final = router.model_list + reasoning_input: Final = _reasoning_input_minted_by("region-a") - result = await check.async_filter_deployments( + result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments( model="gpt-5.4", healthy_deployments=group, messages=None, - request_kwargs=request_kwargs, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input, **failed_in_request}, ) assert result is group - assert request_kwargs["input"][0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} + assert reasoning_input[0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("regions", [("region-a",), ("region-a", "region-b")]) +async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_the_request(regions: tuple[str, ...]): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router(model_list=[_region_deployment(region) for region in regions]) + reasoning_input: Final = _reasoning_input_minted_by("region-a") + original_reasoning: Final = dict(reasoning_input[0]) + + result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments( + model="gpt-5.4", + healthy_deployments=router.model_list, + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input, "_retry_skipped_deployment_ids": regions}, + ) + + assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"] + assert reasoning_input[0] == original_reasoning