From 4716b46c243b90d833ea08de8bebcc9a20c963d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:31:08 -0700 Subject: [PATCH] fix(router): read the affinity pin from the messages argument and strip bridge reasoning on a cross-group route The encrypted_content_affinity check only read the Anthropic history from request_kwargs["messages"], so a caller that passes it through the callback's messages argument alone skipped the pin. Read the argument first and fall back to the kwargs. When the minting deployment is not a candidate of the routed group, the base already strips the Responses input's encrypted reasoning; do the same for the bridge-tagged thinking blocks in Anthropic messages so the routed deployment gets the readable thinking text instead of ciphertext it cannot decrypt. --- .../prompt_templates/common_utils.py | 43 ++++++++++- .../encrypted_content_affinity_check.py | 14 +++- ...ore_utils_prompt_templates_common_utils.py | 60 ++++++++++++++- .../test_encrypted_content_affinity_check.py | 73 +++++++++++++++---- 4 files changed, 166 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ae9b7f6bc4b..7ca2a8aad6a 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import io import json import mimetypes import re -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from itertools import groupby from os import PathLike from pathlib import Path @@ -1889,6 +1889,47 @@ def is_encrypted_reasoning_block(block: object) -> bool: return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping)) +def strip_encrypted_reasoning_from_messages(messages: object) -> None: + """Drop the encrypted reasoning a routed deployment cannot decrypt from Anthropic-shaped + history, keeping the readable thinking text. + + Mutates the content lists in place: the router's fallback snapshot shares these + message objects, so a rebound list would replay the stripped blocks on the fallback hop. + """ + if not isinstance(messages, list): + return + for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json + _strip_encrypted_reasoning_from_blocks(content) + + +def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: + return ( + cast(list[object], content) # cast-ok: narrowed by isinstance + for message in messages + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + ) + + +def _strip_encrypted_reasoning_from_blocks(content: object) -> None: + blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance + stripped: Final = tuple(_without_encrypted_reasoning_block(block) for block in blocks) + blocks[:] = (block for block in stripped if block is not None) # rebind-ok: list shared with fallback snapshot + + +def _without_encrypted_reasoning_block(block: object) -> object | None: + if not is_encrypted_reasoning_block(block): + return block + mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by is_encrypted_reasoning_block + if mapping.get("type") != "thinking" or not mapping.get("thinking"): + return None + kept: Final[dict[str, object]] = { # mutable-ok: thinking block rebuilt without the undecryptable signature + key: value for key, value in mapping.items() if key != "signature" + } + return kept + + def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: index, block = indexed_block return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary" 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 b2edb13a4bf..cdd70e6baf2 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 @@ -48,7 +48,10 @@ from litellm.exceptions import ( ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_content_of_block +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_content_of_block, + 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 @@ -274,8 +277,9 @@ class EncryptedContentAffinityCheck(CustomLogger): parent_otel_span: Span | None = None, ) -> list[dict]: """ - If the request ``input`` contains litellm-encoded item IDs, decode the - embedded ``model_id`` and pin the request to that deployment. Raises + 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 @@ -304,9 +308,10 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True request_input: Final = request_kwargs.get("input") + anthropic_messages: Final = messages or request_kwargs.get("messages") model_id: Final = self._extract_model_id_from_input( request_input - ) or self._extract_model_id_from_anthropic_messages(request_kwargs.get("messages")) + ) or self._extract_model_id_from_anthropic_messages(anthropic_messages) if not model_id: return typed_healthy_deployments @@ -363,6 +368,7 @@ class EncryptedContentAffinityCheck(CustomLogger): 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 diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index ecb07d8c937..71b71d55fa7 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1,3 +1,4 @@ +import copy import functools import json import os @@ -20,6 +21,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_encrypted_reasoning_block, responses_reasoning_items_from_thinking_blocks, split_concatenated_json_objects, + strip_encrypted_reasoning_from_messages, update_messages_with_model_file_ids, ) @@ -1567,7 +1569,9 @@ class TestEncryptedReasoningReplay: def test_signature_round_trips_the_encrypted_content(self): assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes" - @pytest.mark.parametrize("signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7]) + @pytest.mark.parametrize( + "signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7] + ) def test_anything_else_is_not_encrypted_content(self, signature): assert encrypted_content_from_signature(signature) is None @@ -1576,7 +1580,11 @@ class TestEncryptedReasoningReplay: [{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}] ) assert items == ( - {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Plan."}], "encrypted_content": "gAAAA_1"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Plan."}], + "encrypted_content": "gAAAA_1", + }, ) def test_encrypted_redacted_block_replays_with_an_empty_summary(self): @@ -1596,7 +1604,10 @@ class TestEncryptedReasoningReplay: ] ) assert items == ( - {"type": "reasoning", "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}]}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}], + }, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"}, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]}, ) @@ -1621,3 +1632,46 @@ class TestEncryptedReasoningReplay: ) def test_is_encrypted_reasoning_block(self, block, expected): assert is_encrypted_reasoning_block(block) is expected + + def test_strip_keeps_the_readable_thinking_and_drops_the_undecryptable_bytes(self): + assistant_content = [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("g2")}, + {"type": "thinking", "thinking": "", "signature": encrypted_reasoning_signature("g3")}, + {"type": "text", "text": "answer"}, + ] + messages = [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": assistant_content}, + {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}, + ] + + strip_encrypted_reasoning_from_messages(messages) + + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "thinking", "thinking": "packed by the bridge"}, + {"type": "text", "text": "answer"}, + ] + assert messages[0] == {"role": "user", "content": "question"} + assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]} + + @pytest.mark.parametrize( + "messages", + [ + "not a list", + None, + [{"role": "user", "content": None}], + [{"role": "user", "content": "plain string"}], + ["not a message"], + [{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}]}], + ], + ) + def test_strip_leaves_history_without_bridge_reasoning_untouched(self, messages): + before = copy.deepcopy(messages) + + strip_encrypted_reasoning_from_messages(messages) + + assert messages == before diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 99a650436b8..57c02e43fc1 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1597,26 +1597,12 @@ async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_throu {"model_info": {"id": "openai-org-a"}, "litellm_params": {"model": "openai/gpt-5.1"}}, {"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5.1"}}, ] - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", "openai-org-b") - request_kwargs = { - "messages": [ - {"role": "user", "content": "Solve the zebra puzzle"}, - { - "role": "assistant", - "content": [ - {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, - {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, - {"type": "text", "text": "The zebra owner lives in the green house."}, - ], - }, - {"role": "user", "content": "And who drinks water?"}, - ], - } + request_kwargs = {"model": "gpt-5.1"} pinned = await check.async_filter_deployments( model="gpt-5.1", healthy_deployments=deployments, - messages=None, + messages=_bridge_replayed_anthropic_messages(minted_by="openai-org-b"), request_kwargs=request_kwargs, ) @@ -1624,6 +1610,61 @@ async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_throu assert request_kwargs["_encrypted_content_affinity_pinned"] is True +def _bridge_replayed_anthropic_messages(minted_by: str) -> list: + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", minted_by) + return [ + {"role": "user", "content": "Solve the zebra puzzle"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + { + "type": "thinking", + "thinking": "The bridge packed this one", + "signature": f"litellm_encrypted_reasoning:{wrapped}", + }, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ], + }, + {"role": "user", "content": "And who drinks water?"}, + ] + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_routed_to_another_group(): + """ + The /v1/messages twin of the tier-change case: the routed group holds no deployment + of the org that minted the reasoning, so the bridge-tagged blocks are stripped down + to their readable thinking text and the request dispatches to the routed pool. + """ + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["openai-org-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5-nano"}}] + messages = _bridge_replayed_anthropic_messages(minted_by="openai-org-a") + assistant_content = messages[1]["content"] + request_kwargs = {"model": "gpt-5.1"} + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "thinking", "thinking": "The bridge packed this one"}, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ] + + class TestStripEncryptedReasoningFromInput: def test_keeps_summary_and_drops_encrypted_content_and_id(self): wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a")