mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
165dbc2243
commit
e450f2d54c
6 changed files with 471 additions and 246 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue