From e450f2d54c113a962c38e32cf3d0d29341991c3e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:31:43 -0700 Subject: [PATCH] fix(router): serve Responses turns from a sibling when the encrypted content origin has no boundary peer (#43015) * test(integration): reproduce encrypted_content_affinity 503 when origin has no boundary peer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): keep encrypted content affinity turn one out of the response cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): degrade encrypted_content_affinity when the origin has no encryption-boundary peer A routed-group candidate that is currently unavailable and shares its (api_base, api_key) with no healthy deployment used to raise a proxy-level 503/429 from _unavailable_origin_error, even though healthy siblings in the same model group could still serve the turn. Strip the encrypted reasoning and dispatch to the healthy pool instead, matching the existing cross-group behavior, and log a warning naming the origin model_id and routed group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): bound the degraded-affinity log marker and assert the strip on the wire Address review findings: restore num_retries alongside optional_pre_call_checks in the integration test teardown, record scenario request bodies on the scripted upstream so the tests can assert no encrypted reasoning reaches the sibling, and truncate the client-supplied model_id in the degraded-dispatch warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tests): only record JSON bodies on the scripted upstream Multipart uploads to scripted POST routes have no JSON body, so gate the observation recording on the request content-type Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): make encrypted_content_affinity runtime-toggleable so /config/update can turn it off --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router.py | 8 +- .../encrypted_content_affinity_check.py | 166 ++---------- tests/integration/_support/upstream.py | 8 + ...test_encrypted_content_affinity_no_peer.py | 240 ++++++++++++++++++ .../test_router_retry_policy_update.py | 58 +++++ .../test_encrypted_content_affinity_check.py | 237 ++++++++++------- 6 files changed, 471 insertions(+), 246 deletions(-) create mode 100644 tests/integration/routing/test_encrypted_content_affinity_no_peer.py diff --git a/litellm/router.py b/litellm/router.py index da92eef9102..56d27fedf28 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -227,6 +227,9 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, ) +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, +) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( build_io_token_rate_limit_headers, deployment_has_io_token_limits, @@ -438,6 +441,7 @@ _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = M { "prompt_caching": PromptCachingDeploymentCheck, "enforce_model_rate_limits": ModelRateLimitingCheck, + "encrypted_content_affinity": EncryptedContentAffinityCheck, } ) @@ -2207,10 +2211,6 @@ class Router: ) def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, - ) - def _move_before_deployment_affinity( callback_list: list[Any], callback_to_move: EncryptedContentAffinityCheck, 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 e5e40d7d6f5..cf58f3b3d3c 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,17 +36,10 @@ Safe to enable globally: - No cache required. """ -import time from collections.abc import Iterator, Mapping -from typing import TYPE_CHECKING, Final, Optional, Protocol, cast - -import httpx +from typing import TYPE_CHECKING, Final, Optional, cast from litellm._logging import verbose_router_logger -from litellm.exceptions import ( - RateLimitError, - ServiceUnavailableError, -) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -55,7 +48,6 @@ 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.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues from litellm.types.router import Deployment @@ -63,14 +55,6 @@ if TYPE_CHECKING: from litellm.router import Router -class _SupportsActiveCooldowns(Protocol): - """Cooldown-cache handle: this check only reads back the currently active cooldowns.""" - - async def async_get_active_cooldowns( - self, model_ids: list[str], parent_otel_span: Span | None - ) -> list[tuple[str, CooldownCacheValue]]: ... - - class EncryptedContentAffinityCheck(CustomLogger): """ Routes follow-up Responses API requests to the deployment that produced @@ -194,23 +178,6 @@ class EncryptedContentAffinityCheck(CustomLogger): return deployment return None - @staticmethod - def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None: - containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping)) - return next((tid for tid in team_ids if isinstance(tid, str)), None) - - def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]: - """ - Deployment ids that could serve this turn's routed ``model``, as the router - resolves a route (model_group_alias / routing group / model_name / team / - pattern). Delegates to the router so the full precedence is not re-derived here - and no deployment ids are written into request kwargs bound for the provider. - """ - if self.router is None: - return frozenset() - return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs)) - @staticmethod def _encryption_boundary_key( litellm_params: object, @@ -262,9 +229,6 @@ class EncryptedContentAffinityCheck(CustomLogger): Deployments in ``healthy_deployments`` sharing the originating 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 [], None @@ -294,18 +258,12 @@ class EncryptedContentAffinityCheck(CustomLogger): """ If the request ``input`` contains litellm-encoded item IDs, or its Anthropic ``messages`` replay a bridge-tagged thinking block, decode the embedded - ``model_id`` and pin the request to that deployment. Raises - ``RateLimitError`` / ``ServiceUnavailableError`` when the originating - deployment is a member of the routed model group but currently unavailable - and no encryption-boundary peer exists, rather than dispatching a doomed - request to a non-peer deployment. When the origin is not a member of the - routed group (an auto-router tier change, a model switch with no peer, a - removed deployment, or an unknown/forged marker), the encrypted reasoning is - stripped and the request dispatches with its readable history instead. 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. + ``model_id`` and pin the request to that deployment. When the origin cannot + serve this turn and no encryption-boundary peer is configured (it is + unhealthy, the request was routed to a different group by an auto-router tier + change or model switch, or the marker is removed/unknown/forged), the + encrypted reasoning is stripped and the request dispatches to the healthy + pool with its readable history instead of failing. """ request_kwargs = request_kwargs or {} typed_healthy_deployments: Final = cast(list[dict], healthy_deployments) @@ -348,7 +306,7 @@ class EncryptedContentAffinityCheck(CustomLogger): return [deployment] # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. - boundary_matches, originating = self._find_deployments_on_same_encryption_boundary( + boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary( healthy_deployments=typed_healthy_deployments, model_id=model_id, ) @@ -362,101 +320,17 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["_encrypted_content_affinity_pinned"] = True return boundary_matches - # The origin cannot serve this turn's routed group and no peer shares the boundary, so its + # The origin cannot serve this turn and no peer shares its encryption boundary, so its # encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch - # to the routed group instead of failing. Membership is tested by deployment id against the set - # the router actually resolved for this route, not by model-group name, so an alias, a - # provider-qualified spelling, a team-public name, or a pattern route of the same group is not - # mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is - # treated the same as a cross-group one, which also denies an authenticated caller a - # deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and - # dispatch rather than returning distinguishable responses. Only a genuine same-group member - # that is currently unavailable falls through to the fail-fast, preserving the cooldown contract. - routed_group_model_ids: Final = ( - self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset() + # to the healthy pool instead of failing the request. This also denies an authenticated caller a + # deployment-id existence oracle: a same-group id, a cross-group id, a removed id and a forged + # marker all strip and dispatch rather than returning distinguishable responses. + verbose_router_logger.warning( + "EncryptedContentAffinityCheck: model_id=%s cannot serve group %s and no deployment on the same " + "encryption boundary is configured; forwarding without its encrypted reasoning", + model_id[:64], + model, ) - if str(model_id) not in routed_group_model_ids: - verbose_router_logger.debug( - "EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; " - "forwarding without its encrypted reasoning", - model_id, - model, - ) - ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) - strip_encrypted_reasoning_from_messages(anthropic_messages) - return typed_healthy_deployments - - # The origin is a member of the routed group but currently unavailable (cooled down); fail fast - # rather than dispatching to a non-peer, which would guarantee an upstream 400. - raise await self._unavailable_origin_error( - model=model, - model_id=model_id, - parent_otel_span=parent_otel_span, - ) - - async def _unavailable_origin_error( - self, - model: str, - model_id: str, - parent_otel_span: Span | None, - ) -> 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. - cooldown: Final = 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: Final = 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: Span | None, - ) -> CooldownCacheValue | None: - if self.router is None: - return None - cooldown_cache: Final[_SupportsActiveCooldowns | None] = getattr(self.router, "cooldown_cache", None) - if cooldown_cache is None: - return None - try: - active: Final = 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)) + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + strip_encrypted_reasoning_from_messages(anthropic_messages) + return typed_healthy_deployments diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index eea539643d4..e645126032b 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -239,6 +239,14 @@ class Provider: response: Final = self.scenario_store.get(scenario_id) if response is None: return JSONResponse({"error": "Unknown scenario"}, status_code=404) + if request.method == "POST" and "json" in request.headers.get("content-type", ""): + raw_body: Final = await request.body() + if raw_body: + body: Final = JSON_OBJECT.validate_json(raw_body) + if isinstance(body, dict): + self.observations.put( + Observation(request.url.path, request.headers.get("authorization", ""), body) + ) if isinstance(response, RoutedResponse): route_key: Final = f"{request.method} /{'/'.join(segments[1:])}" route: Final = next( diff --git a/tests/integration/routing/test_encrypted_content_affinity_no_peer.py b/tests/integration/routing/test_encrypted_content_affinity_no_peer.py new file mode 100644 index 00000000000..1025278aa10 --- /dev/null +++ b/tests/integration/routing/test_encrypted_content_affinity_no_peer.py @@ -0,0 +1,240 @@ +"""Encrypted-content affinity when the origin deployment has no encryption-boundary peer. + +A multi-region model group has several deployments sharing one upstream api_key but a +distinct api_base each, with ``optional_pre_call_checks: [encrypted_content_affinity]`` +and ``disable_cooldowns: true`` (the integration proxy config). A follow-up +``POST /v1/responses`` that replays a reasoning item must never fail at the proxy while +sibling deployments in the same group are healthy: when the origin cannot serve the turn +its encrypted reasoning should be stripped and the request dispatched to a sibling. +""" + +import json +import uuid +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse +from pydantic import JsonValue + +AFFINITY_CHECK: Final = "encrypted_content_affinity" +PROVIDER_MODEL: Final = "openai/gpt-5" +PROVIDER_KEY: Final = "integration-provider-key" +DEPLOYMENT_COUNT: Final = 3 + + +def _responses_payload() -> dict[str, JsonValue]: + return { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5-scripted", + "output": [ + { + "type": "reasoning", + "id": "rs_$REQUEST_ID", + "summary": [], + "encrypted_content": "ZHNra2RrZA==", + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted answer"}], + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12}, + } + + +def _enable_affinity_check(scenario: Scenario) -> None: + gateway: Final = scenario.gateway + settings: Final = object_value(gateway.get("/router/settings")["current_values"]) + current: Final = settings.get("optional_pre_call_checks") + original: Final = list(current) if isinstance(current, list) else [] + original_retries: Final = settings.get("num_retries") + scenario.cleanups.callback( + lambda: gateway.post( + "/config/update", + {"router_settings": {"optional_pre_call_checks": original, "num_retries": original_retries}}, + ) + ) + gateway.post( + "/config/update", + {"router_settings": {"optional_pre_call_checks": [*original, AFFINITY_CHECK], "num_retries": 0}}, + ) + + +def _last_group_request_body(gateway: Gateway, group: str) -> dict[str, JsonValue]: + """The request body the upstream last saw for this group, to prove the strip reached the wire.""" + with httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream: + requests: Final = object_value(upstream.get("/__observations").json()).get("requests") + assert isinstance(requests, list) + bodies: Final = [ + object_value(request).get("body") + for request in requests + if str(object_value(request).get("path")).startswith(f"/{group}-") + ] + assert bodies, f"upstream saw no requests for {group}: {requests}" + body: Final = bodies[-1] + assert isinstance(body, dict), f"upstream request body is not an object: {body}" + return body + + +def _assert_no_encrypted_reasoning_reached_upstream(gateway: Gateway, group: str) -> None: + body: Final = _last_group_request_body(gateway, group) + items: Final = body.get("input") + assert isinstance(items, list), f"upstream request carried no input list: {body}" + assert not any( + isinstance(item, dict) and ("encrypted_content" in item or "litellm_enc" in json.dumps(item)) for item in items + ), f"undecryptable reasoning reached the sibling: {body}" + + +def _multi_region_group(scenario: Scenario) -> tuple[str, tuple[str, ...]]: + """Three deployments in one model group: one shared api_key, a distinct api_base each.""" + gateway: Final = scenario.gateway + group: Final = f"enc-affinity-{uuid.uuid4().hex}" + handles: Final = tuple( + register_scenario( + f"{group}-{index}", + JsonResponse(content_type="application/json", body=_responses_payload()), + ) + for index in range(DEPLOYMENT_COUNT) + ) + for handle in handles: + scenario.cleanups.callback(delete_scenario, handle) + + def delete_model_if_present(model_id: str) -> None: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + if any(object_value(object_value(entry)["model_info"])["id"] == model_id for entry in entries): + scenario.delete_model(model_id) + + deployment_ids: Final = tuple( + string_value( + object_value( + gateway.post( + "/model/new", + { + "model_name": group, + "litellm_params": { + "model": PROVIDER_MODEL, + "api_key": PROVIDER_KEY, + "api_base": handle.api_base(), + }, + "model_info": {}, + }, + )["model_info"] + )["id"] + ) + for handle in handles + ) + for model_id in deployment_ids: + scenario.cleanups.callback(delete_model_if_present, model_id) + return group, deployment_ids + + +def _user_message(text: str) -> dict[str, JsonValue]: + return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]} + + +def _responses_turn(gateway: Gateway, group: str, request_input: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/responses", + {"model": group, "input": request_input, "store": False, "include": ["reasoning.encrypted_content"]}, + ) + + +def _turn_one(gateway: Gateway, group: str, deployment_ids: tuple[str, ...]) -> tuple[str, list[JsonValue]]: + # the proxy caches responses, so the turn-one prompt needs a unique marker or a + # stale body carrying another run's encoded model_id would replay instead + response: Final = _responses_turn(gateway, group, f"hello affinity {uuid.uuid4().hex}") + assert response.status_code == 200, response.text + origin: Final = str(response.headers["x-litellm-model-id"]) + assert origin in deployment_ids, f"turn one served by unknown deployment {origin}: {response.text}" + output: Final = object_value(response.json()).get("output") + assert isinstance(output, list), f"turn one returned no output items: {response.text}" + reasoning: Final = next((item for item in output if object_value(item).get("type") == "reasoning"), None) + assert isinstance(reasoning, dict), f"turn one returned no reasoning item: {response.text}" + assert str(object_value(reasoning)["id"]).startswith("encitem_"), ( + f"affinity encoding did not run on turn one: {reasoning}" + ) + assert isinstance(object_value(reasoning).get("encrypted_content"), str), ( + f"turn one reasoning item has no encrypted_content: {reasoning}" + ) + message: Final = next((item for item in output if object_value(item).get("type") == "message"), None) + assert isinstance(message, dict), f"turn one returned no message item: {response.text}" + return origin, [reasoning, message] + + +def _replay(gateway: Gateway, group: str, items: list[JsonValue]) -> httpx.Response: + return _responses_turn( + gateway, + group, + [_user_message("hello affinity"), *items, _user_message("continue the conversation")], + ) + + +def _model_blocked(gateway: Gateway, model_id: str) -> bool: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + entry: Final = next( + (entry for entry in entries if object_value(object_value(entry)["model_info"])["id"] == model_id), + None, + ) + return entry is not None and object_value(object_value(entry)["model_info"]).get("blocked") is True + + +def test_replayed_encrypted_content_serves_from_sibling_when_origin_blocked(gateway: Gateway) -> None: + """Origin excluded from healthy deployments (admin-blocked, no cooldown) must not 503. + + On unfixed code the origin is a routed-group candidate with no (api_base, api_key) + peer, so the affinity check raises a proxy-level 503 instead of stripping the + encrypted reasoning and dispatching to a healthy sibling. + """ + with gateway.scenario() as scenario: + _enable_affinity_check(scenario) + group, deployment_ids = _multi_region_group(scenario) + origin, items = _turn_one(gateway, group, deployment_ids) + + gateway.post("/model/block", {"model_id": origin}) + eventually(lambda: _model_blocked(gateway, origin), lambda blocked: blocked) + + response: Final = _replay(gateway, group, items) + assert response.status_code == 200, response.text + siblings: Final = tuple(model_id for model_id in deployment_ids if model_id != origin) + assert response.headers.get("x-litellm-model-id") in siblings, ( + f"turn two served by {response.headers.get('x-litellm-model-id')}, " + f"expected a sibling of blocked origin {origin}: {response.text}" + ) + _assert_no_encrypted_reasoning_reached_upstream(gateway, group) + + +def test_replayed_encrypted_content_serves_from_sibling_when_origin_deleted(gateway: Gateway) -> None: + """Origin permanently removed: strip-and-dispatch to a sibling, the shipped behavior.""" + with gateway.scenario() as scenario: + _enable_affinity_check(scenario) + group, deployment_ids = _multi_region_group(scenario) + origin, items = _turn_one(gateway, group, deployment_ids) + + gateway.post("/model/delete", {"id": origin}) + eventually( + lambda: gateway.get("/model/info")["data"], + lambda entries: ( + isinstance(entries, list) + and all(object_value(object_value(entry)["model_info"])["id"] != origin for entry in entries) + ), + ) + + response: Final = _replay(gateway, group, items) + assert response.status_code == 200, response.text + siblings: Final = tuple(model_id for model_id in deployment_ids if model_id != origin) + assert response.headers.get("x-litellm-model-id") in siblings, ( + f"turn two served by {response.headers.get('x-litellm-model-id')}, " + f"expected a sibling of deleted origin {origin}: {response.text}" + ) + _assert_no_encrypted_reasoning_reached_upstream(gateway, group) diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 0a3dcba325a..b8c4e5fe084 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -30,6 +30,7 @@ from pydantic import ValidationError import litellm from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck from litellm.types.router import RetryPolicy, UpdateRouterConfig @@ -203,6 +204,63 @@ def test_update_settings_replaces_toggleable_pre_call_checks(): assert any(isinstance(callback, ModelRateLimitingCheck) for callback in (router.optional_callbacks or [])) +def test_update_settings_clears_omitted_encrypted_content_affinity_check(): + router = _build_router() + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + router.update_settings(optional_pre_call_checks=[]) + + assert not any( + isinstance(callback, EncryptedContentAffinityCheck) for callback in (router.optional_callbacks or []) + ) + assert not any(isinstance(callback, EncryptedContentAffinityCheck) for callback in litellm.callbacks) + + +@pytest.mark.asyncio +async def test_update_settings_turning_off_encrypted_content_affinity_stops_flagging_requests(): + router = _build_router() + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + enabled_request: Final = {"litellm_metadata": {}, "input": "hello"} + await router.async_get_available_deployment(model="test-model", request_kwargs=enabled_request) + assert enabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + + router.update_settings(optional_pre_call_checks=[]) + disabled_request: Final = {"litellm_metadata": {}, "input": "hello"} + await router.async_get_available_deployment(model="test-model", request_kwargs=disabled_request) + assert "encrypted_content_affinity_enabled" not in disabled_request["litellm_metadata"] + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + reenabled_request: Final = {"litellm_metadata": {}, "input": "hello"} + await router.async_get_available_deployment(model="test-model", request_kwargs=reenabled_request) + assert reenabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + + +def test_update_settings_keeps_per_group_encrypted_content_affinity_when_global_toggle_is_omitted(): + router = litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-fake", + "api_base": "http://localhost:9999", + }, + } + ], + model_group_affinity_config={"test-model": ["encrypted_content_affinity"]}, + ) + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + router.update_settings(optional_pre_call_checks=[]) + + affinity_checks: Final = [ + callback for callback in (router.optional_callbacks or []) if isinstance(callback, EncryptedContentAffinityCheck) + ] + assert len(affinity_checks) == 1 + assert affinity_checks[0].enable_global_affinity is False + + @pytest.mark.asyncio async def test_update_settings_preserves_router_budget_limiting_when_omitted(monkeypatch): async def _disable_periodic_sync(*args, **kwargs): 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 aa34fbd6bf7..836049c88a2 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 @@ -1285,7 +1285,7 @@ def test_boundary_key_rejects_non_dict_like_inputs(): # --------------------------------------------------------------------------- -# Fail-fast when originating deployment is unavailable and no boundary peer +# Degraded dispatch when the originating deployment is unavailable and no boundary peer # --------------------------------------------------------------------------- @@ -1323,14 +1323,13 @@ def _make_router_mock_with_cooldown( @pytest.mark.asyncio -async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_429(): +async def test_affinity_strips_and_dispatches_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. + degrade: strip the encrypted reasoning and dispatch to the healthy pool + rather than failing the request. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1366,32 +1365,36 @@ async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_42 } ] request_kwargs = { - "input": [{"id": encoded_id, "type": "reasoning"}], + "input": [ + { + "id": encoded_id, + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + "gAAAAA-blob", "deployment-a-cooled" + ), + } + ], } - 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, - ) + result = 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 + assert result is healthy_only_b + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"]) @pytest.mark.asyncio -async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_for_429(): +async def test_affinity_strips_and_dispatches_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. + Originating deployment is in cooldown specifically because of a 429 and no + boundary peer is configured. The check degrades the same way: strip the + encrypted reasoning and dispatch to the healthy pool rather than surfacing + a rate-limit error to the caller. """ - from litellm.exceptions import RateLimitError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1431,29 +1434,23 @@ async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_fo "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, - ) + result = 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 + assert result is healthy_only_b @pytest.mark.asyncio -async def test_affinity_raises_service_unavailable_when_origin_filtered_without_cooldown_entry(): +async def test_affinity_strips_and_dispatches_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. + with no active cooldown entry and no boundary peer. The check degrades: + strip the encrypted reasoning and dispatch to the healthy pool. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1480,15 +1477,72 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ "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, - ) + result = 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 + assert result is healthy_only_b + + +@pytest.mark.asyncio +async def test_affinity_serves_sibling_when_candidate_origin_has_no_boundary_peer(): + """ + Regression for the multi-region group where every deployment has a distinct + api_base: the origin is still a routed-group candidate but absent from + healthy_deployments, and no (api_base, api_key) peer exists. The turn must + degrade on a sibling with its encrypted reasoning stripped, not fail. + """ + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock("https://region-a.example.com/v1", "shared-key") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["region-a", "region-b", "region-c"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "region-a") + siblings = [ + { + "model_info": {"id": "region-b"}, + "model_name": "gpt-5.4", + "litellm_params": {"api_base": "https://region-b.example.com/v1", "api_key": "shared-key"}, + }, + { + "model_info": {"id": "region-c"}, + "model_name": "gpt-5.4", + "litellm_params": {"api_base": "https://region-c.example.com/v1", "api_key": "shared-key"}, + }, + ] + request_kwargs = { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "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?"}, + ], + } + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=siblings, + messages=None, + request_kwargs=request_kwargs, + ) + + assert result is siblings + assert request_kwargs["input"][1] == { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "scattering"}], + } @pytest.mark.asyncio @@ -1981,15 +2035,14 @@ async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_ @pytest.mark.asyncio -async def test_affinity_fails_fast_within_the_origins_own_group(): +async def test_affinity_degrades_within_the_origins_own_group(): """ - Negative class for the tier-change discriminator: the routed group IS the - origin's group (a same-group cooldown, not a tier change), so even with a - healthy non-origin sibling that cannot decrypt the content, the request - still fails fast and the encrypted reasoning is left intact rather than - stripped. Preserves the LIT-3051 cooldown contract. + The routed group IS the origin's group (a same-group cooldown), and the + healthy sibling sits on a different encryption boundary, so it cannot + decrypt the replayed reasoning. The check still degrades instead of + failing: the encrypted reasoning is stripped and the request dispatches + to the sibling. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2014,29 +2067,26 @@ async def test_affinity_fails_fast_within_the_origins_own_group(): ] request_kwargs = _cross_group_request_kwargs() - with pytest.raises(ServiceUnavailableError): - await check.async_filter_deployments( - model="gpt-reasoning-tier", - healthy_deployments=sibling_pool, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-reasoning-tier", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert request_kwargs["input"][1].get("encrypted_content") + assert result is sibling_pool + assert not request_kwargs["input"][1].get("encrypted_content") @pytest.mark.asyncio -async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_same_by_id(): +async def test_affinity_strips_when_group_is_spelled_differently_but_same_by_id(): """ - The discriminator must key on deployment-id membership, not on the model-group - name string. Here the origin's configured group is spelled ``openai/gpt-5.4-mini`` - while the routed group is the canonical ``gpt-5.4-mini``: same group, different - spelling. A name compare (``originating.model_name != model``) would read this as - a tier change and strip the reasoning it did not have to. Because the origin's id - is a member of the routed group, this is a same-group cooldown instead: the request - fails fast and the encrypted reasoning is left intact. + The origin's configured group is spelled ``openai/gpt-5.4-mini`` while the + routed group is the canonical ``gpt-5.4-mini``: same group, different + spelling, with the origin a member but currently unavailable. There is no + boundary peer, so the check degrades: strip the encrypted reasoning and + dispatch to the sibling pool. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2071,28 +2121,24 @@ async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_sam ], } - with pytest.raises(ServiceUnavailableError): - await check.async_filter_deployments( - model="gpt-5.4-mini", - healthy_deployments=sibling_pool, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4-mini", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert request_kwargs["input"][1].get("encrypted_content") + assert result is sibling_pool + assert not request_kwargs["input"][1].get("encrypted_content") @pytest.mark.asyncio -async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes(): +async def test_affinity_strips_for_team_and_pattern_routes(): """ - The exact `model_name` index does not include team-public or pattern routes, so a - same-group cooldown reached only through one of those would be misread as a tier change - and stripped. The check asks the router for the candidate ids it resolves for the route - (`get_candidate_model_ids_for_route`), which covers those paths, rather than the bare - index. Here that set marks the origin as a candidate, so the request fails fast with its - reasoning intact, and the routed group and team are passed through to the router. + Team-public or pattern routes resolve the same routed group as the origin's, + so an unavailable origin there also degrades rather than failing: the + encrypted reasoning is stripped and the request dispatches to the sibling. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2127,13 +2173,12 @@ async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes( ], } - with pytest.raises(ServiceUnavailableError): - await check.async_filter_deployments( - model="team-public-model", - healthy_deployments=sibling_pool, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="team-public-model", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert request_kwargs["input"][1].get("encrypted_content") - mock_router.get_candidate_model_ids_for_route.assert_called_once_with(model="team-public-model", team_id="teamA") + assert result is sibling_pool + assert not request_kwargs["input"][1].get("encrypted_content")