fix(router): strip encrypted reasoning the pinned deployment cannot decrypt

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -2015,7 +2015,11 @@ def is_unsignable_thinking_block(block: object) -> bool:
     return not (isinstance(thinking_text, str) and len(thinking_text.strip()) > 0)
 
 
-def strip_encrypted_reasoning_from_messages(messages: object) -> None:
+def strip_encrypted_reasoning_from_messages(
+    messages: object,
+    *,
+    should_strip: Callable[[Mapping[str, object]], bool] | None = None,
+) -> None:
     """Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from
     Anthropic-shaped history.
 
@@ -2030,7 +2034,7 @@ def strip_encrypted_reasoning_from_messages(messages: object) -> None:
     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)
+        _strip_encrypted_reasoning_from_blocks(content, should_strip=should_strip)
 
 
 def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]:
@@ -2043,9 +2047,18 @@ def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]:
     )
 
 
-def _strip_encrypted_reasoning_from_blocks(content: object) -> None:
+def _strip_encrypted_reasoning_from_blocks(
+    content: object,
+    *,
+    should_strip: Callable[[Mapping[str, object]], bool] | None = None,
+) -> None:
     blocks: Final = cast(list[object], content)  # cast-ok: narrowed by the caller's isinstance
-    kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block))
+    kept: Final = tuple(
+        block
+        for block in blocks
+        if not is_encrypted_reasoning_block(block)
+        or (should_strip is not None and not should_strip(cast(Mapping[str, object], block)))
+    )
     blocks[:] = kept
 
 

litellm/responses/utils.py

@@ -1,6 +1,6 @@
 import base64
 import re
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from functools import reduce
 from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload
 
@@ -556,7 +556,11 @@ def _restore_encrypted_content_item_ids_in_input(request_input: _RequestInputT)
         return request_input
 
     @staticmethod
-    def strip_encrypted_reasoning_from_input(request_input: object) -> None:
+    def strip_encrypted_reasoning_from_input(
+        request_input: object,
+        *,
+        should_strip: Callable[[Mapping[str, object]], bool] | None = None,
+    ) -> None:
         """Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary.
 
         Mutates ``request_input`` in place: the router's fallback snapshot shares this
@@ -565,7 +569,12 @@ def strip_encrypted_reasoning_from_input(request_input: object) -> None:
         if not isinstance(request_input, list):
             return
         items: Final = cast(list[object], request_input)  # cast-ok: untyped client json
-        stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items)
+        stripped: Final = tuple(
+            ResponsesAPIRequestUtils._without_encrypted_reasoning(item)
+            if should_strip is None or (isinstance(item, Mapping) and should_strip(cast(Mapping[str, object], item)))
+            else item
+            for item in items
+        )
         items[:] = (item for item in stripped if item is not None)
 
     @staticmethod

litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py

@@ -36,7 +36,8 @@
 - No cache required.
 """
 
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator, Mapping, Sequence
+from functools import cache
 from typing import TYPE_CHECKING, Final, Optional, cast
 
 from litellm._logging import verbose_router_logger
@@ -114,23 +115,31 @@ def _extract_model_id_from_input(request_input: object) -> str | None:
         if not isinstance(request_input, list):
             return None
 
-        for item in request_input:
-            if not isinstance(item, dict):
-                continue
+        return next(
+            (
+                model_id
+                for item in request_input
+                if (model_id := EncryptedContentAffinityCheck._model_id_of_input_item(item)) is not None
+            ),
+            None,
+        )
+
+    @staticmethod
+    def _model_id_of_input_item(item: object) -> str | None:
+        if not isinstance(item, dict):
+            return None
 
-            # First, try to decode from item ID (if present)
-            item_id = item.get("id")
-            if item_id and isinstance(item_id, str):
-                decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id)
-                if decoded:
-                    return decoded.get("model_id")
+        item_id: Final = item.get("id")
+        if item_id and isinstance(item_id, str):
+            decoded: Final = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id)
+            if decoded:
+                return decoded.get("model_id")
 
-            # If no encoded ID, check if encrypted_content itself is wrapped
-            encrypted_content = item.get("encrypted_content")
-            if encrypted_content and isinstance(encrypted_content, str):
-                model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content)
-                if model_id:
-                    return model_id
+        encrypted_content: Final = item.get("encrypted_content")
+        if encrypted_content and isinstance(encrypted_content, str):
+            model_id: Final = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content)
+            if model_id:
+                return model_id
 
         return None
 
@@ -150,19 +159,20 @@ def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | No
         model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
         return model_id or None
 
+    @staticmethod
+    def _model_id_of_anthropic_block(block: Mapping[str, object]) -> str | None:
+        encrypted_content: Final = encrypted_content_of_block(block)
+        if encrypted_content is None:
+            return None
+        return EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content)
+
     @staticmethod
     def _extract_model_id_from_anthropic_messages(messages: object) -> str | None:
         return next(
             (
                 model_id
                 for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages)
-                if (encrypted_content := encrypted_content_of_block(block)) is not None
-                if (
-                    model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(
-                        encrypted_content
-                    )
-                )
-                is not None
+                if (model_id := EncryptedContentAffinityCheck._model_id_of_anthropic_block(block)) is not None
             ),
             None,
         )
@@ -243,6 +253,50 @@ def _find_deployments_on_same_encryption_boundary(
         ]
         return matches, originating
 
+    def _strip_reasoning_the_target_cannot_decrypt(
+        self,
+        request_input: object,
+        anthropic_messages: object,
+        target_deployments: Sequence[Mapping[str, object]],
+    ) -> None:
+        target_ids: Final = frozenset(
+            str(model_info["id"])
+            for target in target_deployments
+            if isinstance((model_info := target.get("model_info")), Mapping) and model_info.get("id") is not None
+        )
+        target_boundaries: Final = frozenset(
+            boundary
+            for target in target_deployments
+            if (boundary := self._encryption_boundary_key(target.get("litellm_params"))) is not None
+        )
+
+        @cache
+        def target_can_decrypt(origin_model_id: str) -> bool:
+            if origin_model_id in target_ids:
+                return True
+            if self.router is None:
+                return False
+            origin: Final = self.router.get_deployment(model_id=origin_model_id)
+            origin_boundary: Final = (
+                self._encryption_boundary_key(origin.litellm_params.model_dump(exclude_none=True))
+                if origin is not None
+                else None
+            )
+            return origin_boundary is not None and origin_boundary in target_boundaries
+
+        def should_strip_input_item(item: Mapping[str, object]) -> bool:
+            origin_model_id: Final = self._model_id_of_input_item(item)
+            return origin_model_id is not None and not target_can_decrypt(origin_model_id)
+
+        def should_strip_anthropic_block(block: Mapping[str, object]) -> bool:
+            origin_model_id: Final = self._model_id_of_anthropic_block(block)
+            return origin_model_id is not None and not target_can_decrypt(origin_model_id)
+
+        ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(
+            request_input, should_strip=should_strip_input_item
+        )
+        strip_encrypted_reasoning_from_messages(anthropic_messages, should_strip=should_strip_anthropic_block)
+
     # ------------------------------------------------------------------
     # Request routing  (pre-call filter)
     # ------------------------------------------------------------------
@@ -303,6 +357,7 @@ async def async_filter_deployments(
                 model_id,
             )
             request_kwargs["_encrypted_content_affinity_pinned"] = True
+            self._strip_reasoning_the_target_cannot_decrypt(request_input, anthropic_messages, (deployment,))
             return [deployment]
 
         # Follow-up switched model_name (LIT-2531): pin by Azure resource instead.
@@ -318,6 +373,7 @@ async def async_filter_deployments(
                 len(boundary_matches),
             )
             request_kwargs["_encrypted_content_affinity_pinned"] = True
+            self._strip_reasoning_the_target_cannot_decrypt(request_input, anthropic_messages, boundary_matches)
             return boundary_matches
 
         # The origin cannot serve this turn and no peer shares its encryption boundary, so its

tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py

@@ -1879,6 +1879,20 @@ def test_strip_drops_every_bridge_tagged_block_and_leaves_no_unsigned_thinking_b
         assert messages[0] == {"role": "user", "content": "question"}
         assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}
 
+    def test_strip_uses_predicate_to_keep_selected_encrypted_blocks(self):
+        kept_signature = encrypted_reasoning_signature("keep")
+        stripped_signature = encrypted_reasoning_signature("strip")
+        content = [
+            {"type": "thinking", "thinking": "keep", "signature": kept_signature},
+            {"type": "thinking", "thinking": "strip", "signature": stripped_signature},
+        ]
+        messages = [{"role": "assistant", "content": content}]
+
+        strip_encrypted_reasoning_from_messages(messages, should_strip=lambda block: block.get("thinking") == "strip")
+
+        assert messages[0]["content"] is content
+        assert content == [{"type": "thinking", "thinking": "keep", "signature": kept_signature}]
+
     @pytest.mark.parametrize(
         "messages",
         [