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 9badee854fb..1564c0da634 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 @@ -38,7 +38,7 @@ Safe to enable globally: import hashlib from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import cache, lru_cache, partial from time import monotonic from typing import TYPE_CHECKING, Final, Optional, cast @@ -62,7 +62,7 @@ from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.secret_managers.main import secret_manager_would_be_consulted from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import Deployment, GenericLiteLLMParams +from litellm.types.router import Deployment, GenericLiteLLMParams, ModelInfo from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential @@ -91,8 +91,11 @@ class _EncryptionBoundary: model: Callable[[], str | None] = field(compare=False) account_scope: frozenset[tuple[str, object]] explicit_pair: bool + declared: str | None = None def is_shared_with(self, other: "_EncryptionBoundary") -> bool: + if self.declared is not None and other.declared is not None: + return self.declared == other.declared if (self.api_base, self.api_key) != (other.api_base, other.api_key): return False if self.explicit_pair and other.explicit_pair: @@ -290,6 +293,37 @@ class EncryptedContentAffinityCheck(CustomLogger): ) ) + @staticmethod + def _deployment_encryption_boundary( + litellm_params: object, + model_info: object, + clock: Callable[[], float] = monotonic, + ) -> _EncryptionBoundary | None: + """ + ``_encryption_boundary_key``, overridden by an operator-declared + ``model_info.encrypted_content_boundary``. Two deployments that both declare one are peers exactly + when the names match, which lets endpoints of one provider organization with different ``api_base`` + values share encrypted reasoning. + """ + declared: Final = EncryptedContentAffinityCheck._declared_encryption_boundary(model_info) + resolved: Final = EncryptedContentAffinityCheck._encryption_boundary_key(litellm_params, clock) + if declared is None: + return resolved + if resolved is None: + return _EncryptionBoundary(None, None, lambda: None, frozenset(), explicit_pair=False, declared=declared) + return replace(resolved, declared=declared) + + @staticmethod + def _declared_encryption_boundary(model_info: object) -> str | None: + if isinstance(model_info, ModelInfo): + return model_info.encrypted_content_boundary + if not isinstance(model_info, Mapping): + return None + declared: Final = cast(Mapping[str, object], model_info).get( # cast-ok: narrowed by isinstance + "encrypted_content_boundary" + ) + return declared if isinstance(declared, str) and declared else None + @staticmethod @lru_cache(maxsize=1024) def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None: @@ -408,9 +442,9 @@ class EncryptedContentAffinityCheck(CustomLogger): def _find_deployments_on_same_encryption_boundary( self, - healthy_deployments: list[dict], + healthy_deployments: list[dict[str, object]], model_id: str, - ) -> tuple[list[dict], Deployment | None]: + ) -> tuple[list[dict[str, object]], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating deployment's encryption boundary, alongside the originating @@ -421,15 +455,20 @@ class EncryptedContentAffinityCheck(CustomLogger): originating: Final = self.ctx.registry().deployment(model_id) if originating is None: return [], None - boundary: Final = self._encryption_boundary_key( - originating.litellm_params.model_dump(exclude_none=True), self._clock + boundary: Final = self._deployment_encryption_boundary( + originating.litellm_params.model_dump(exclude_none=True), originating.model_info, self._clock ) if boundary is None: return [], originating matches: Final = [ d for d in healthy_deployments - if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None + if ( + candidate := self._deployment_encryption_boundary( + d.get("litellm_params", {}), d.get("model_info"), self._clock + ) + ) + is not None and boundary.is_shared_with(candidate) ] return matches, originating @@ -464,8 +503,8 @@ class EncryptedContentAffinityCheck(CustomLogger): boundary for target in target_deployments if ( - boundary := EncryptedContentAffinityCheck._encryption_boundary_key( - target.get("litellm_params"), clock + boundary := EncryptedContentAffinityCheck._deployment_encryption_boundary( + target.get("litellm_params"), target.get("model_info"), clock ) ) is not None @@ -482,8 +521,8 @@ class EncryptedContentAffinityCheck(CustomLogger): return False origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( - EncryptedContentAffinityCheck._encryption_boundary_key( - origin.litellm_params.model_dump(exclude_none=True), clock + EncryptedContentAffinityCheck._deployment_encryption_boundary( + origin.litellm_params.model_dump(exclude_none=True), origin.model_info, clock ) if origin is not None else None diff --git a/litellm/types/router.py b/litellm/types/router.py index c0b182f0f20..6f5652db771 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -353,6 +353,10 @@ class ModelInfo(MirroredPricingParams): max_embedding_batch_size: Annotated[int, Field(gt=0)] | None = None + # deployments that set the same value share one provider encryption boundary, so encrypted reasoning + # minted on one can be replayed to another (e.g. regional endpoints of one OpenAI organization) + encrypted_content_boundary: Annotated[str, Field(min_length=1)] | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided 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 76a574bae24..adf08e43c95 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 @@ -3293,3 +3293,194 @@ async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manage assert pin_before_the_change == [peer] assert pin_after_the_change == [other_org] assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0] + + +# --------------------------------------------------------------------------- +# Operator-declared encryption boundary (model_info.encrypted_content_boundary) +# --------------------------------------------------------------------------- + + +def _regional_deployment(model_name: str, region: str, boundary: str | None) -> dict: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_base": f"https://{region}.example.invalid/v1", + "api_key": "sk-one-org", + }, + "model_info": {"id": region} | ({} if boundary is None else {"encrypted_content_boundary": boundary}), + } + + +async def _forwarded_reasoning_after_model_switch(origin_boundary: str | None, peer_boundary: str | None) -> tuple: + router: Final = litellm.Router( + model_list=[ + _regional_deployment("reasoning-a", "region-1", origin_boundary), + _regional_deployment("reasoning-b", "region-2", peer_boundary), + ], + optional_pre_call_checks=["encrypted_content_affinity"], + num_retries=0, + ) + handler: Final = AsyncMock(return_value=_build_mock_response(output_items=[], response_id="resp_second")) + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new=handler, + ): + response: Final = await router.aresponses( + model="reasoning-b", + input=[ + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + }, + {"role": "user", "content": "follow-up"}, + ], + litellm_metadata={}, + ) + forwarded_reasoning: Final = handler.await_args.kwargs["input"][0] + return response._hidden_params["model_id"], forwarded_reasoning.get("encrypted_content") + + +@pytest.mark.asyncio +async def test_declared_boundary_keeps_reasoning_across_endpoints_with_different_api_bases(): + assert await _forwarded_reasoning_after_model_switch("org-1", "org-1") == ("region-2", "blob-from-region-1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("origin_boundary", "peer_boundary"), + [(None, None), ("org-1", "org-2"), ("org-1", None), (None, "org-1")], + ids=["undeclared", "different-boundaries", "only-origin-declared", "only-peer-declared"], +) +async def test_reasoning_is_stripped_unless_both_deployments_declare_the_same_boundary( + origin_boundary, peer_boundary +): + assert await _forwarded_reasoning_after_model_switch(origin_boundary, peer_boundary) == ("region-2", None) + + +def test_declared_boundary_overrides_a_shared_api_base_and_key(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + other_org: Final = { + "model_info": {"id": "region-1-org-2", "encrypted_content_boundary": "org-2"}, + "litellm_params": {"api_base": "https://region-1.example.invalid/v1", "api_key": "sk-one-org"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([other_org], "region-1") + + assert matches == [] + + +@pytest.mark.asyncio +async def test_boundary_pin_only_selects_deployments_that_declare_the_origins_boundary(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + peer: Final = _regional_deployment("reasoning-b", "region-2", "org-1") + stranger: Final = _regional_deployment("reasoning-b", "region-3", "org-2") + request_kwargs: Final = { + "input": [ + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + } + ] + } + + result: Final = await check.async_filter_deployments( + model="reasoning-b", + healthy_deployments=[stranger, peer], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [peer] + assert request_kwargs["input"][0]["encrypted_content"] == ( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob-from-region-1", "region-1") + ) + + +def test_a_boundary_declared_on_one_side_keeps_the_shared_api_base_and_key_match(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + undeclared_twin: Final = { + "model_info": {"id": "region-1-twin"}, + "litellm_params": {"api_base": "https://region-1.example.invalid/v1", "api_key": "sk-one-org"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([undeclared_twin], "region-1") + + assert matches == [undeclared_twin] + + +def test_declared_boundary_matches_deployments_whose_credentials_cannot_be_resolved(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "") + origin.litellm_params.model_dump.return_value = {"api_base": "https://region-1.example.invalid/v1"} + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + peer: Final = { + "model_info": {"id": "region-2", "encrypted_content_boundary": "org-1"}, + "litellm_params": {"api_base": "https://region-2.example.invalid/v1"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([peer], "region-1") + + assert matches == [peer] + + +async def _forwarded_reasoning_after_order_fallback(boundary: str | None) -> tuple: + primary: Final = _regional_deployment("reasoning", "region-1", boundary) + backup: Final = _regional_deployment("reasoning", "region-3", boundary) + router: Final = litellm.Router( + model_list=[ + primary | {"litellm_params": primary["litellm_params"] | {"order": 1}}, + backup | {"litellm_params": backup["litellm_params"] | {"order": 2}}, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + disable_cooldowns=True, + num_retries=0, + ) + + async def upstream(*, litellm_params, **_kwargs): + if "region-1" in litellm_params.api_base: + raise litellm.InternalServerError("upstream 500", llm_provider="openai", model="gpt-5.6") + return _build_mock_response(output_items=[], response_id="resp_second") + + handler: Final = AsyncMock(side_effect=upstream) + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new=handler, + ): + response: Final = await router.aresponses( + model="reasoning", + input=[ + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + }, + {"role": "user", "content": "follow-up"}, + ], + litellm_metadata={}, + ) + forwarded_reasoning: Final = handler.await_args.kwargs["input"][0] + return response._hidden_params["model_id"], forwarded_reasoning.get("encrypted_content") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("boundary", "expected_reasoning"), + [("org-1", "blob-from-region-1"), (None, None)], + ids=["declared", "undeclared"], +) +async def test_order_fallback_keeps_reasoning_only_on_a_declared_boundary_peer(boundary, expected_reasoning): + assert await _forwarded_reasoning_after_order_fallback(boundary) == ("region-3", expected_reasoning) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6f8c49cc0f9..b171a72dcf9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -52372,6 +52372,8 @@ export interface components { discoverable?: boolean | null; /** Enable Tag Filtering */ enable_tag_filtering?: boolean | null; + /** Encrypted Content Boundary */ + encrypted_content_boundary?: string | null; /** Enforce Supported Endpoints */ enforce_supported_endpoints?: boolean | null; /** Explicit Cache Read Input Token Cost */