fix(router): pin bridge-replayed encrypted reasoning to the deployment that minted it

The encrypted_content_affinity check only read the pin from the Responses
input, which /v1/messages builds after the router has picked a deployment,
so a model group spread across OpenAI orgs sent follow-up turns to the
wrong org and got invalid_encrypted_content back. The check now also
decodes the pin from bridge-tagged thinking and redacted_thinking blocks
in the Anthropic messages. The bridge also keeps a deployment's own
include list next to reasoning.encrypted_content instead of replacing it.
This commit is contained in:
mateo-berri 2026-09-09 16:29:52 -07:00
parent f49ebc93ea
commit e27a018a6c
5 changed files with 104 additions and 10 deletions

View file

@ -1873,7 +1873,7 @@ def _encrypted_reasoning_field(block: Mapping[str, object]) -> object:
return None
def _encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
def encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
return encrypted_content_from_signature(_encrypted_reasoning_field(block))
@ -1900,7 +1900,7 @@ def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) ->
for block in group
if (text := _readable_thinking_text(block))
]
encrypted_content: Final = _encrypted_content_of_block(group[0])
encrypted_content: Final = encrypted_content_of_block(group[0])
if encrypted_content is not None:
return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content)
if not summary:

View file

@ -126,7 +126,7 @@ def _build_responses_kwargs(
responses_kwargs["stream"] = True
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
excluded: Final = {"anthropic_messages"}
excluded: Final = frozenset(("anthropic_messages",))
for key, value in forwarded_kwargs.items():
if key == "litellm_logging_obj" and value is not None:
from litellm.litellm_core_utils.litellm_logging import (
@ -147,6 +147,14 @@ def _build_responses_kwargs(
if explicit_prompt_cache_key is not None:
responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key
deployment_include: Final = forwarded_kwargs.get("include")
bridge_include: Final = responses_kwargs.get("include")
if isinstance(deployment_include, list) and isinstance(bridge_include, list):
responses_kwargs["include"] = [
*bridge_include,
*(item for item in deployment_include if item not in bridge_include),
]
return responses_kwargs

View file

@ -37,6 +37,7 @@ Safe to enable globally:
"""
import time
from collections.abc import Iterator, Mapping
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
import httpx
@ -48,6 +49,7 @@ 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.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.cooldown_cache import CooldownCacheValue
from litellm.types.llms.openai import AllMessageValues
@ -138,15 +140,48 @@ class EncryptedContentAffinityCheck(CustomLogger):
# 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,
_,
) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content)
if model_id:
return model_id
return None
@staticmethod
def _anthropic_content_blocks(messages: object) -> Iterator[Mapping[str, object]]:
if not isinstance(messages, list):
return iter(())
return (
cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance
for message in cast(list[object], messages) # cast-ok: narrowed by isinstance
if isinstance(message, Mapping)
for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance
if isinstance(content, list)
for block in cast(list[object], content) # cast-ok: narrowed by isinstance
if isinstance(block, Mapping)
)
@staticmethod
def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | None:
model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
return model_id or None
@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
),
None,
)
@staticmethod
def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
for deployment in healthy_deployments:
@ -248,13 +283,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
if "litellm_metadata" in request_kwargs:
request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True
request_input: Final = request_kwargs.get("input")
model_id: Final = self._extract_model_id_from_input(request_input)
model_id: Final = self._extract_model_id_from_input(
request_kwargs.get("input")
) or self._extract_model_id_from_anthropic_messages(request_kwargs.get("messages"))
if not model_id:
return typed_healthy_deployments
verbose_router_logger.debug(
"EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs",
"EncryptedContentAffinityCheck: decoded model_id=%s from the request's encrypted content markers",
model_id,
)

View file

@ -77,6 +77,16 @@ def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_reje
assert "reasoning" in responses_kwargs
def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning():
responses_kwargs = _build_responses_kwargs(
max_tokens=1024,
messages=MESSAGES,
model="openai/gpt-5.6-luna",
extra_kwargs={"custom_llm_provider": "openai", "include": ["file_search_call.results"]},
)
assert responses_kwargs["include"] == ["reasoning.encrypted_content", "file_search_call.results"]
def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key():
responses_kwargs = _build_responses_kwargs(
max_tokens=1024,

View file

@ -1656,3 +1656,43 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
finally:
router.discard()
@pytest.mark.asyncio
async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_through_the_bridge():
"""
Claude Code behind /v1/messages replays the encrypted reasoning the bridge packed
into a thinking block's signature (or a redacted block's data). The pin has to be
read from those blocks because the bridge builds the Responses `input` only after
the router has picked a deployment.
"""
check = EncryptedContentAffinityCheck()
deployments = [
{"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?"},
],
}
pinned = await check.async_filter_deployments(
model="gpt-5.1",
healthy_deployments=deployments,
messages=None,
request_kwargs=request_kwargs,
)
assert [d["model_info"]["id"] for d in pinned] == ["openai-org-b"]
assert request_kwargs["_encrypted_content_affinity_pinned"] is True