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] 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