From ff91b89b8d57771edd4d8ddb8b020093d45fcf5d Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:47:04 -0800 Subject: [PATCH] fix(router): keep encrypted reasoning when a fallback hops between same-key deployments with no api_base (#49396) * fix(router): keep encrypted reasoning on a hop between same-key deployments on the provider's default endpoint * test(router): type the default-endpoint boundary test helpers * fix(guardrails): restore the google_text_moderation embeddings fix an earlier main merge dropped * fix(router): match encryption boundaries on the key, account, and region the request really uses A deployment's named credential now fills only the values the deployment leaves out, through the same helper the request path uses, so the boundary names the key and base the request sends. A deployment with no api_key takes the provider's default key, the boundary carries the OpenAI organization and project a deployment sends, and a Bedrock Mantle region named in the model routes through the provider resolver before the Responses URL builder * fix(router): keep header-signed accounts apart on a shared boundary and skip boundary work no reasoning needs Headers the request logger masks as secrets (Authorization, api-key, x-api-key) now join the boundary's account scope, with extra_headers overriding headers the way the request merges them, so deployments that sign in through a header on one base never share encrypted reasoning. Target boundaries are built only when a reasoning item from another deployment needs the comparison, and the Responses input strip asks only about encrypted reasoning items, so a pinned turn or fallback hop with nothing to compare makes no secret-manager reads. * fix(router): resolve default boundaries only for openai and drop the unsent organization param Aggregators and Bedrock Mantle serve several vendors on one key, so a missing key or base there no longer yields a shared boundary. The Responses request never sends the organization param, so it no longer splits the boundary, while a Cookie header now does. * fix(router): match explicit pairs as on main and key default boundaries by model An explicit api_base plus api_key pair is the whole boundary again, so tag headers no longer split deployments that main kept together. The default OpenAI boundary now carries the model, since one OPENAI_BASE_URL can front several vendors, and a deployment that declares another provider gets no boundary before any secret lookup. * fix(router): match an explicit key and the environment key on one api_base A deployment with an explicit api_base and api_key now shares a boundary with one on the same api_base that takes the same key from OPENAI_API_KEY, when the model and account headers also match. Two explicit pairs still match on base and key alone, as on main. * fix(router): skip a deployment whose params cannot be read in boundary checks A named credential with a value of the wrong type, or an extra header whose name is not a string, made the boundary check raise and fail the whole request. That deployment now has no boundary, so the others still keep the reasoning they can decrypt. Two explicit api_base plus api_key pairs also compare without resolving the provider, so they read no secrets, as on main. * fix(router): ask a secret manager for default boundaries at most once a minute per deployment * fix(router): resolve a changed named credential at once in the cached boundary and take an injected clock --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../get_llm_provider_logic.py | 6 +- litellm/responses/utils.py | 15 +- .../encrypted_content_affinity_check.py | 285 ++++++-- ...allback_hop_foreign_encrypted_reasoning.py | 36 +- .../test_encrypted_content_affinity_check.py | 681 ++++++++++++++++-- 5 files changed, 917 insertions(+), 106 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 38242291b71..94ba02bfbd7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -146,6 +146,10 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider +def declared_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: + return custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + + def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: """The authenticating provider this pair already names, or None. @@ -154,7 +158,7 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str and for a declared pair the resolver's answer is the declaration itself, so metadata callers adopt the declaration instead of resolving. """ - declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + declared: Final = declared_provider(model, custom_llm_provider) return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 1a97c1cb48c..43ed9a3b96f 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -678,20 +678,25 @@ class ResponsesAPIRequestUtils: return items: Final = cast(list[object], request_input) # cast-ok: untyped client json 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))) + ResponsesAPIRequestUtils._without_encrypted_reasoning(reasoning) + if (reasoning := ResponsesAPIRequestUtils._encrypted_reasoning_item(item)) is not None + and (should_strip is None or should_strip(reasoning)) else item for item in items ) items[:] = (item for item in stripped if item is not None) @staticmethod - def _without_encrypted_reasoning(item: object) -> object | None: + def _encrypted_reasoning_item(item: object) -> Mapping[str, object] | None: if not isinstance(item, dict): - return item + return None reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"): - return reasoning + return None + return reasoning + + @staticmethod + def _without_encrypted_reasoning(reasoning: Mapping[str, object]) -> object | None: readable: Final = any( ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") ) 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 6f3f0794371..f3425596afa 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 @@ -36,31 +36,75 @@ Safe to enable globally: - No cache required. """ -from collections.abc import Iterator, Mapping, Sequence -from functools import cache +import hashlib +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from functools import cache, lru_cache, partial +from time import monotonic from typing import TYPE_CHECKING, Final, Optional, cast -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError +import litellm from litellm._logging import verbose_router_logger +from litellm.exceptions import AuthenticationError, BadRequestError from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, get_llm_provider +from litellm.litellm_core_utils.litellm_logging import has_sensitive_keyword from litellm.litellm_core_utils.prompt_templates.common_utils import ( anthropic_content_lists, encrypted_content_of_block, strip_encrypted_reasoning_from_messages, ) from litellm.litellm_core_utils.routing_common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids +from litellm.llms.base_llm.chat.transformation import BaseLLMException 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 -from litellm.utils import get_order_filtered_deployments +from litellm.types.router import Deployment, GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential if TYPE_CHECKING: from litellm.router_api import RoutingContext _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[object, object]) +_ACCOUNT_SCOPE_HEADERS: Final = frozenset({"openai-organization", "openai-project", "cookie"}) +_SINGLE_VENDOR_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value}) +_DEFAULT_ROUTE_SECRETS: Final = ("USE_LITELLM_PROXY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE") +_SECRET_MANAGER_REFRESH_SECONDS: Final = 60 + + +@dataclass(frozen=True, slots=True) +class _ProviderRoute: + model: str + provider: str + api_base: str | None + + +@dataclass(frozen=True, slots=True) +class _EncryptionBoundary: + api_base: object + api_key: object + model: Callable[[], str | None] = field(compare=False) + account_scope: frozenset[tuple[str, object]] + explicit_pair: bool + + def is_shared_with(self, other: "_EncryptionBoundary") -> bool: + if (self.api_base, self.api_key) != (other.api_base, other.api_key): + return False + if self.explicit_pair and other.explicit_pair: + return True + return self.account_scope == other.account_scope and self.model() == other.model() + + +@dataclass(frozen=True, slots=True) +class _ParamsInRefreshWindow: + digest: str + window: int + params: Mapping[str, object] = field(compare=False) class EncryptedContentAffinityCheck(CustomLogger): @@ -83,9 +127,11 @@ class EncryptedContentAffinityCheck(CustomLogger): ctx: Optional["RoutingContext"] = None, enable_global_affinity: bool = True, model_group_affinity_config: dict[str, list[str]] | None = None, + clock: Callable[[], float] = monotonic, ) -> None: super().__init__() self.ctx = ctx + self._clock = clock self.enable_global_affinity = enable_global_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} @@ -209,44 +255,156 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _encryption_boundary_key( litellm_params: object, - ) -> tuple[object, object] | None: + clock: Callable[[], float] = monotonic, + ) -> _EncryptionBoundary | None: """ - ``(api_base, api_key)`` identifies an upstream encryption boundary. - The values are resolved from the deployment and its named credential - without modifying the deployment. + The upstream encryption boundary a deployment's encrypted reasoning belongs to. The values + are resolved the way the request resolves them, the deployment's own values winning over its + named credential's, without modifying the deployment. Two explicit ``(api_base, api_key)`` + pairs match on those values alone. On a provider that serves one vendor's models, a deployment + with no ``api_key`` takes the provider's default key, one with no ``api_base`` takes the endpoint + its Responses URL builder resolves, and a match involving either also needs the same model and + the same headers that pick an account. - Accepts any object exposing dict-style ``.get(key, default)``: plain - dicts (the common case in ``healthy_deployments``) as well as - ``LiteLLM_Params``-style Pydantic instances, which define a custom - ``.get()``. A stricter ``isinstance(dict)`` guard would silently drop - the latter from boundary matching and fall back to the full pool — - i.e. trigger the exact ``invalid_encrypted_content`` failure this - check exists to prevent. + Accepts plain dicts (the common case in ``healthy_deployments``) as + well as ``LiteLLM_Params`` Pydantic instances. A stricter + ``isinstance(dict)`` guard would silently drop the latter from + boundary matching and fall back to the full pool — i.e. trigger the + exact ``invalid_encrypted_content`` failure this check exists to prevent. + + When a secret manager would be asked for the provider defaults, the boundary of one set of + deployment and named credential values is reused for up to a minute, so a run of pin misses + does not ask it again for every candidate, while a changed credential resolves at once. """ - getter: Final = getattr(litellm_params, "get", None) - if not callable(getter): + params: Final = EncryptedContentAffinityCheck._params_mapping(litellm_params) + if params is None: return None - api_base: Final = getter("api_base") - api_key: Final = getter("api_key") - credential_name: Final = getter("litellm_credential_name") - credential_values: Final[Mapping[str, object] | None] = ( - CredentialAccessor.get_credential_values(credential_name) - if isinstance(credential_name, str) and credential_name - else None + effective_params: Final = with_named_credential(params) + if not any(secret_manager_would_be_consulted(name) for name in _DEFAULT_ROUTE_SECRETS): + return EncryptedContentAffinityCheck._resolve_encryption_boundary(effective_params) + return EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window( + _ParamsInRefreshWindow( + hashlib.sha256(repr(effective_params).encode()).hexdigest(), + int(clock() // _SECRET_MANAGER_REFRESH_SECONDS), + effective_params, + ) ) - effective_api_base: Final = ( - credential_values.get("api_base") - if credential_values is not None and "api_base" in credential_values - else api_base - ) - effective_api_key: Final = ( - credential_values.get("api_key") - if credential_values is not None and "api_key" in credential_values - else api_key - ) - if not effective_api_base or not effective_api_key: + + @staticmethod + @lru_cache(maxsize=1024) + def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None: + return EncryptedContentAffinityCheck._resolve_encryption_boundary(key.params) + + @staticmethod + def _resolve_encryption_boundary(effective_params: Mapping[str, object]) -> _EncryptionBoundary | None: + explicit_api_key: Final = effective_params.get("api_key") + explicit_api_base: Final = effective_params.get("api_base") + account_scope: Final = EncryptedContentAffinityCheck._account_scope(effective_params) + if explicit_api_key and explicit_api_base: + return _EncryptionBoundary( + explicit_api_base, + explicit_api_key, + cache(partial(EncryptedContentAffinityCheck._single_vendor_model, effective_params)), + account_scope, + explicit_pair=True, + ) + route: Final = EncryptedContentAffinityCheck._single_vendor_route(effective_params) + if route is None: + return None + api_key: Final = explicit_api_key or EncryptedContentAffinityCheck._default_api_key(route) + if not api_key: + return None + api_base: Final = explicit_api_base or EncryptedContentAffinityCheck._default_responses_endpoint( + route, effective_params + ) + if not api_base: + return None + return _EncryptionBoundary(api_base, api_key, lambda: route.model, account_scope, explicit_pair=False) + + @staticmethod + def _params_mapping(litellm_params: object) -> Mapping[str, object] | None: + if isinstance(litellm_params, BaseModel): + return litellm_params.model_dump(exclude_none=True) + if isinstance(litellm_params, Mapping): + return _REQUEST_KWARGS_ADAPTER.validate_python(litellm_params) + return None + + @staticmethod + def _account_scope(litellm_params: Mapping[str, object]) -> frozenset[tuple[str, object]]: + return frozenset( + (name.lower(), str(value)) + for name, value in EncryptedContentAffinityCheck._header_items(litellm_params) + if name.lower() in _ACCOUNT_SCOPE_HEADERS or has_sensitive_keyword(name) + ) + + @staticmethod + def _header_items(litellm_params: Mapping[str, object]) -> Iterator[tuple[str, object]]: + headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("headers")) + extra_headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("extra_headers")) + overridden: Final = frozenset(name.lower() for name in extra_headers) + yield from ((name, value) for name, value in headers.items() if name.lower() not in overridden) + yield from extra_headers.items() + + @staticmethod + def _header_mapping(headers: object) -> Mapping[str, object]: + if isinstance(headers, Mapping): + return {str(name): value for name, value in _HEADERS_ADAPTER.validate_python(headers).items()} + return {} + + @staticmethod + def _single_vendor_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None: + route: Final = EncryptedContentAffinityCheck._provider_route(litellm_params) + return route if route is not None and route.provider in _SINGLE_VENDOR_PROVIDERS else None + + @staticmethod + def _single_vendor_model(litellm_params: Mapping[str, object]) -> str | None: + route: Final = EncryptedContentAffinityCheck._single_vendor_route(litellm_params) + return route.model if route is not None else None + + @staticmethod + def _provider_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None: + model: Final = litellm_params.get("model") + if not isinstance(model, str): + return None + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + declared: Final = declared_provider( + model, custom_llm_provider if isinstance(custom_llm_provider, str) else None + ) + if declared in litellm.provider_list and declared not in _SINGLE_VENDOR_PROVIDERS: + return None + try: + provider_model, provider, _, provider_api_base = get_llm_provider( + model=model, litellm_params=GenericLiteLLMParams.model_validate(dict(litellm_params)) + ) + except (BadRequestError, ValidationError): + return None + return _ProviderRoute(provider_model, provider, provider_api_base) + + @staticmethod + def _default_api_key(route: _ProviderRoute) -> str | None: + model_info: Final = ProviderConfigManager.get_provider_model_info( + model=route.model, provider=LlmProviders(route.provider) + ) + if model_info is None: + return None + try: + return model_info.get_api_key() + except NotImplementedError: + return None + + @staticmethod + def _default_responses_endpoint(route: _ProviderRoute, litellm_params: Mapping[str, object]) -> str | None: + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + provider=route.provider, model=route.model + ) + if responses_config is None: + return None + try: + return responses_config.get_complete_url( # pyright: ignore[reportUnknownMemberType] # the base types litellm_params as a bare dict + api_base=route.api_base, litellm_params=dict(litellm_params) + ) + except (ValueError, BaseLLMException, AuthenticationError): return None - return (effective_api_base, effective_api_key) def _find_deployments_on_same_encryption_boundary( self, @@ -255,7 +413,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) -> tuple[list[dict], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating - deployment's ``(api_base, api_key)``, alongside the originating + deployment's encryption boundary, alongside the originating deployment object (or ``None`` if it was removed / router unavailable). """ if self.ctx is None: @@ -263,11 +421,16 @@ 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)) + boundary: Final = self._encryption_boundary_key( + originating.litellm_params.model_dump(exclude_none=True), self._clock + ) if boundary is None: return [], originating matches: Final = [ - d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary + d + for d in healthy_deployments + if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None + and boundary.is_shared_with(candidate) ] return matches, originating @@ -280,6 +443,7 @@ class EncryptedContentAffinityCheck(CustomLogger): target_deployments: Sequence[Mapping[str, object]], *, unmarked_origin: str | None, + clock: Callable[[], float] = monotonic, ) -> None: """ Drop the encrypted reasoning that none of ``target_deployments`` minted or shares an @@ -293,12 +457,19 @@ class EncryptedContentAffinityCheck(CustomLogger): 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 := EncryptedContentAffinityCheck._encryption_boundary_key(target.get("litellm_params"))) - is not None - ) + + @cache + def target_boundaries() -> tuple[_EncryptionBoundary, ...]: + return tuple( + boundary + for target in target_deployments + if ( + boundary := EncryptedContentAffinityCheck._encryption_boundary_key( + target.get("litellm_params"), clock + ) + ) + is not None + ) @cache def target_can_decrypt(marked_origin: str | None) -> bool: @@ -312,12 +483,14 @@ class EncryptedContentAffinityCheck(CustomLogger): origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( EncryptedContentAffinityCheck._encryption_boundary_key( - origin.litellm_params.model_dump(exclude_none=True) + origin.litellm_params.model_dump(exclude_none=True), clock ) if origin is not None else None ) - return origin_boundary is not None and origin_boundary in target_boundaries + return origin_boundary is not None and any( + origin_boundary.is_shared_with(target_boundary) for target_boundary in target_boundaries() + ) def should_strip_input_item(item: Mapping[str, object]) -> bool: return not target_can_decrypt(EncryptedContentAffinityCheck._model_id_of_input_item(item)) @@ -424,7 +597,12 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.ctx, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id + self.ctx, + request_input, + anthropic_messages, + (deployment,), + unmarked_origin=model_id, + clock=self._clock, ) return [deployment] @@ -442,7 +620,12 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.ctx, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id + self.ctx, + request_input, + anthropic_messages, + boundary_matches, + unmarked_origin=model_id, + clock=self._clock, ) return boundary_matches diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index e6ec3444369..ef2a9e5b72f 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -303,7 +303,7 @@ def scenario_api_base(scenario: Scenario, scenario_id: str, response: StoredResp def deployment( scenario: Scenario, group: str, - api_base: str, + api_base: str | None, *, order: int | None, model: str = RESPONSES_MODEL, @@ -316,7 +316,7 @@ def deployment( "litellm_params": { "model": model, "api_key": PROVIDER_KEY, - "api_base": api_base, + **({} if api_base is None else {"api_base": api_base}), **({} if order is None else {"order": order}), **extra_params, }, @@ -863,6 +863,38 @@ def test_hop_inside_one_encryption_boundary_keeps_the_reasoning(gateway: Gateway assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies +def test_hop_between_same_key_deployments_on_the_providers_default_endpoint_keeps_the_reasoning( + gateway: Gateway, tmp_path: Path +) -> None: + shared: Final = f"default-{uuid.uuid4().hex[:12]}" + with ( + scratch_database() as database, + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database, "OPENAI_BASE_URL": f"{gateway.upstream_url}/{shared}"}, + remove_environment=("DIRECT_URL", "DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE"), + ) as owned, + owned.gateway.scenario() as scenario, + ): + owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}}) + one: Final = scenario_api_base(scenario, f"{shared}-o1", failing()).rsplit("/", 1)[-1] + two: Final = scenario_api_base(scenario, f"{shared}-o2", healthy_json()).rsplit("/", 1)[-1] + order_one: Final = deployment( + scenario, shared, None, order=1, extra_params={"extra_headers": {"x-scripted-scenario": one}} + ) + order_two: Final = deployment( + scenario, shared, None, order=2, extra_params={"extra_headers": {"x-scripted-scenario": two}} + ) + mark: Final = marker() + answer: Final = responses_answer( + owned.gateway, HTTPX, Target(shared, frozenset({order_one, order_two})), history(mark) + ) + assert_served_by(answer, order_two) + bodies: Final = bodies_for(observed(owned.gateway), shared) + assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies + + def assert_cross_group_hop_stripped(gateway: Gateway, source_scenario: str, target_scenario: str, mark: str) -> None: records: Final = observed(gateway) assert [body.get("input") for body in bodies_for(records, source_scenario)] == [history(mark)], records 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 d184bf84035..a4d279e32e6 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 @@ -15,17 +15,21 @@ The mechanism works without any cache and supports two encoding strategies: encrypted_content back to their original forms before sending to the upstream provider. """ +import datetime import time +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, patch import pytest import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_api import RoutingContext from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem @pytest.fixture(autouse=True) @@ -1100,14 +1104,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params) assert pydantic_key is not None - assert ( - pydantic_key - == plain_key - == ( - "https://mateo-resource.openai.azure.com", - "fake-azure-resource-key-a", - ) - ) + assert pydantic_key == plain_key def test_boundary_key_resolves_missing_values_from_named_credential(): @@ -1131,14 +1128,35 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) - assert boundary == ("https://account-a.example.com", "credential-key-a") + assert boundary is not None + assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key( + {"api_base": "https://account-a.example.com", "api_key": "credential-key-a"} + ) + assert boundary != EncryptedContentAffinityCheck._encryption_boundary_key( + {"api_base": "https://account-a.example.com", "api_key": "credential-key-b"} + ) -def test_boundary_key_matches_named_credential_precedence(): +@pytest.mark.parametrize( + "deployment_values", + [ + pytest.param({"api_base": "https://deployment.example.com", "api_key": "deployment-key"}, id="set values"), + pytest.param({"api_base": "", "api_key": ""}, id="empty values"), + ], +) +@pytest.mark.usefixtures("no_process_wide_api_base") +def test_boundary_key_resolves_named_credential_values_the_way_the_request_does( + monkeypatch: pytest.MonkeyPatch, deployment_values: dict[str, str] +): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) + from litellm.utils import with_named_credential + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + deployment: Final = {"model": "openai/gpt-5.5", **deployment_values, "litellm_credential_name": "account-a"} with patch.object( litellm, "credential_list", @@ -1153,46 +1171,14 @@ def test_boundary_key_matches_named_credential_precedence(): ) ], ): - boundary = EncryptedContentAffinityCheck._encryption_boundary_key( - { - "api_base": "https://deployment.example.com", - "api_key": "deployment-key", - "litellm_credential_name": "account-a", - } - ) + boundary: Final = EncryptedContentAffinityCheck._encryption_boundary_key(deployment) + sent_params: Final = with_named_credential(deployment) - assert boundary == ("https://credential.example.com", "credential-key-a") - - -def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, + assert boundary is not None + assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key( + {key: value for key, value in sent_params.items() if key != "litellm_credential_name"} ) - with patch.object( - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ): - boundary = EncryptedContentAffinityCheck._encryption_boundary_key( - { - "api_base": "", - "api_key": "", - "litellm_credential_name": "account-a", - } - ) - - assert boundary == ("https://credential.example.com", "credential-key-a") - def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( @@ -2728,3 +2714,604 @@ async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_th assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"] assert reasoning_input[0] == original_reasoning + + +def _default_endpoint_deployment(model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]: + return {"model_name": "gpt-5.5", "litellm_params": dict(litellm_params), "model_info": {"id": model_id}} + + +def _deployment_of(ctx: RoutingContext, model_id: str) -> dict[str, object]: + deployment: Final = ctx.registry().deployment(model_id) + assert deployment is not None + return deployment.model_dump(exclude_none=True) + + +SUMMARY_ONLY_REASONING: Final = {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} +UNMARKED_REASONING: Final = {**SUMMARY_ONLY_REASONING, "id": "rs_test", "encrypted_content": "gAAAAA-blob"} +OPENAI_SHARED_KEY: Final = {"model": "openai/gpt-5.5", "api_key": "shared-key"} + + +@pytest.fixture +def no_process_wide_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + + +@pytest.mark.usefixtures("no_process_wide_api_base") +@pytest.mark.parametrize("marked", [True, False], ids=["marked origin", "unmarked origin"]) +@pytest.mark.parametrize( + "origin_params, target_params, target_can_decrypt", + [ + pytest.param(OPENAI_SHARED_KEY, OPENAI_SHARED_KEY, True, id="same key on the provider default endpoint"), + pytest.param( + {"model": "gpt-5.5", "custom_llm_provider": "openai", "api_key": "shared-key"}, + OPENAI_SHARED_KEY, + True, + id="provider named apart from the model", + ), + pytest.param(OPENAI_SHARED_KEY, {**OPENAI_SHARED_KEY, "api_key": "other-key"}, False, id="different key"), + pytest.param( + OPENAI_SHARED_KEY, + {**OPENAI_SHARED_KEY, "model": "openai/gpt-5.5-mini"}, + False, + id="same key on the default endpoint for another model", + ), + pytest.param( + {"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"}, + {"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"}, + False, + id="same key on Bedrock Mantle, which serves several vendors", + ), + pytest.param( + {"model": "openrouter/openai/gpt-6.1-sol", "api_key": "shared-key"}, + {"model": "openrouter/anthropic/claude-haiku-5.5", "api_key": "shared-key"}, + False, + id="same key on an aggregator that serves several vendors", + ), + pytest.param( + {"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"}, + {"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"}, + False, + id="provider without a Responses URL builder", + ), + pytest.param( + {"model": "azure/gpt-5.5", "api_key": "shared-key"}, + {"model": "azure/gpt-5.5", "api_key": "shared-key"}, + False, + id="provider whose builder needs an api_base", + ), + pytest.param( + {"model": "github_copilot/gpt-5.5", "api_key": "shared-key"}, + {"model": "github_copilot/gpt-5.5", "api_key": "shared-key"}, + False, + id="same key on a provider that authenticates with its own token", + ), + pytest.param( + {"model": "groq/llama-4-scout"}, + {"model": "groq/llama-4-scout"}, + False, + id="provider with no default key", + ), + pytest.param( + {"model": "vertex_ai/gemini-3-pro"}, + {"model": "vertex_ai/gemini-3-pro"}, + False, + id="provider whose default key is a credentials file", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "api_base": "https://openai-proxy.example.com/v1"}, + OPENAI_SHARED_KEY, + False, + id="explicit api_base beside the default endpoint", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"openai-organization": "org-b"}}, + False, + id="same key sent to different OpenAI organizations", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Project": "proj-a"}}, + OPENAI_SHARED_KEY, + False, + id="same key with an OpenAI project beside one without", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"Cookie": "session=account-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"cookie": "session=account-b"}}, + False, + id="same key signed in to different accounts by cookie", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "headers": {"api-key": "account-a"}}, + {**OPENAI_SHARED_KEY, "headers": {"api-key": "account-b"}}, + False, + id="explicit key beside different api-key headers", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "organization": "org-a"}, + {**OPENAI_SHARED_KEY, "organization": "org-b"}, + True, + id="organization param the Responses request never sends", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "a", "OpenAI-Organization": "org-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "b", "OpenAI-Organization": "org-a"}}, + True, + id="same OpenAI organization beside unrelated headers", + ), + pytest.param( + { + **OPENAI_SHARED_KEY, + "headers": {"OpenAI-Organization": "org-a"}, + "extra_headers": {"openai-organization": "org-b"}, + }, + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-b"}}, + True, + id="extra headers override the organization the headers set", + ), + ], +) +def test_hop_strip_keeps_reasoning_between_same_key_deployments_on_the_default_endpoint( + origin_params: dict[str, object], target_params: dict[str, object], target_can_decrypt: bool, marked: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", origin_params), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] if marked else dict(UNMARKED_REASONING) + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, + request_input, + None, + (_deployment_of(ctx, "target"),), + unmarked_origin=None if marked else "origin", + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +@pytest.mark.usefixtures("no_process_wide_api_base") +@pytest.mark.parametrize( + "target_api_key, target_can_decrypt", + [ + pytest.param(None, True, id="both on the environment key"), + pytest.param("env-key", True, id="explicit key equal to the environment key"), + pytest.param("other-key", False, id="explicit key other than the environment key"), + ], +) +def test_hop_strip_resolves_a_missing_key_from_the_environment( + monkeypatch: pytest.MonkeyPatch, target_api_key: str | None, target_can_decrypt: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + target_params: Final = {"model": "openai/gpt-5.5", **({"api_key": target_api_key} if target_api_key else {})} + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", {"model": "openai/gpt-5.5"}), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +def test_hop_strip_leaves_input_items_that_are_not_objects_as_sent(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + request_input: Final = ["and sunsets?", *_reasoning_input_minted_by("origin")] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + None, request_input, None, (), unmarked_origin=None + ) + + assert request_input == ["and sunsets?", SUMMARY_ONLY_REASONING, {"role": "user", "content": "and sunsets?"}] + + +GATEWAY_BASE: Final = {"model": "openai/gpt-5.5", "api_base": "https://gateway.example.com/v1"} + + +@pytest.mark.parametrize( + "origin_params, target_params, target_can_decrypt", + [ + pytest.param( + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}}, + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-b"}}, + False, + id="environment key replaced by different Authorization headers", + ), + pytest.param( + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "a"}}, + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "b"}}, + True, + id="explicit pair beside different tag headers", + ), + pytest.param( + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}}, + {**GATEWAY_BASE, "extra_headers": {"authorization": "Bearer sk-account-a"}}, + True, + id="same Authorization header on both", + ), + ], +) +def test_hop_strip_keeps_accounts_that_sign_in_through_headers_apart( + monkeypatch: pytest.MonkeyPatch, + origin_params: dict[str, object], + target_params: dict[str, object], + target_can_decrypt: bool, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", origin_params), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +@pytest.mark.parametrize( + "environment_key, target_model, target_can_decrypt", + [ + pytest.param("shared-key", GATEWAY_BASE["model"], True, id="environment key equal to the explicit key"), + pytest.param("other-key", GATEWAY_BASE["model"], False, id="environment key other than the explicit key"), + pytest.param("shared-key", "openai/gpt-5.5-mini", False, id="equal keys serving another model"), + ], +) +def test_hop_strip_matches_an_explicit_key_and_the_environment_key_on_one_api_base( + monkeypatch: pytest.MonkeyPatch, environment_key: str, target_model: str, target_can_decrypt: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", environment_key) + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", {**GATEWAY_BASE, "api_key": "shared-key"}), + _default_endpoint_deployment("target", {**GATEWAY_BASE, "model": target_model}), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +class _SecretReadCounter(CustomSecretManager): + def __init__(self) -> None: + super().__init__(secret_manager_name="read-counter") + self.reads: int = 0 + + async def async_read_secret( + self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None + ) -> str | None: + self.reads += 1 + return None + + def sync_read_secret( + self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None + ) -> str | None: + self.reads += 1 + return None + + +@pytest.mark.parametrize( + "deployment_params, reasoning_minted_by", + [ + pytest.param(OPENAI_SHARED_KEY, None, id="no reasoning in the input"), + pytest.param(OPENAI_SHARED_KEY, "target", id="reasoning the hop's target minted"), + pytest.param( + {"model": "groq/llama-4-scout"}, "origin", id="reasoning from a provider with no default boundary" + ), + pytest.param({**GATEWAY_BASE, "api_key": "shared-key"}, "origin", id="reasoning between explicit pairs"), + ], +) +def test_hop_strip_reads_no_secrets_when_no_default_boundary_needs_resolving( + monkeypatch: pytest.MonkeyPatch, deployment_params: dict[str, object], reasoning_minted_by: str | None +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", deployment_params), + _default_endpoint_deployment("target", deployment_params), + ] + ) + ctx: Final = RoutingContext(router) + secret_reads: Final = _SecretReadCounter() + monkeypatch.setattr(litellm, "secret_manager_client", secret_reads) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + request_input: Final = ( + _reasoning_input_minted_by(reasoning_minted_by) + if reasoning_minted_by + else [{"role": "user", "content": "and sunsets?"}] + ) + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin="origin" + ) + + assert secret_reads.reads == 0 + + +@pytest.mark.parametrize( + "unreadable_params", + [ + pytest.param( + {"model": "openai/gpt-5.5", "litellm_credential_name": "mistyped"}, + id="named credential with a value of the wrong type", + ), + pytest.param( + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {1: "one"}}, + id="extra header whose name is not a string", + ), + ], +) +def test_hop_strip_keeps_reasoning_a_target_can_decrypt_beside_a_target_with_unreadable_params( + monkeypatch: pytest.MonkeyPatch, unreadable_params: dict[str, object] +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="mistyped", + credential_values={"api_key": "shared-key", "api_version": datetime.date(2024, 10, 21)}, + credential_info={}, + ) + ], + ) + shared_pair: Final = {**GATEWAY_BASE, "api_key": "shared-key"} + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", shared_pair), + _default_endpoint_deployment("unreadable", unreadable_params), + _default_endpoint_deployment("target", shared_pair), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, + request_input, + None, + (_deployment_of(ctx, "unreadable"), _deployment_of(ctx, "target")), + unmarked_origin=None, + ) + + assert request_input[0] == original_reasoning + + +def test_boundary_key_puts_pydantic_and_plain_params_on_one_default_endpoint_boundary(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + from litellm.types.router import LiteLLM_Params + + pydantic_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(LiteLLM_Params(**OPENAI_SHARED_KEY)) + plain_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(dict(OPENAI_SHARED_KEY)) + + assert pydantic_key is not None + assert plain_key is not None + assert pydantic_key.is_shared_with(plain_key) + assert plain_key.is_shared_with(pydantic_key) + + +def test_boundary_key_has_no_default_endpoint_for_a_model_no_provider_claims(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + unmapped: Final = {"model": "unmapped-model-name", "api_key": "shared-key"} + + assert EncryptedContentAffinityCheck._encryption_boundary_key(unmapped) is None + + +def test_boundary_key_has_no_default_key_for_a_custom_provider(monkeypatch: pytest.MonkeyPatch): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom"]) + monkeypatch.setattr(litellm, "_custom_providers", ["qa-custom"]) + + assert EncryptedContentAffinityCheck._encryption_boundary_key({"model": "qa-custom/some-model"}) is None + + +@pytest.mark.asyncio +async def test_affinity_pins_to_the_same_key_peer_on_the_default_endpoint_when_the_origin_cannot_serve(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", OPENAI_SHARED_KEY), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + reasoning_input: Final = _reasoning_input_minted_by("origin") + original_reasoning: Final = dict(reasoning_input[0]) + + result: Final[list[dict[str, object]]] = await EncryptedContentAffinityCheck(ctx=ctx).async_filter_deployments( + model="gpt-5.5", + healthy_deployments=[peer, _deployment_of(ctx, "other-org")], + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input}, + ) + + assert result == [peer] + assert reasoning_input[0] == original_reasoning + + +class _HandAdvancedClock: + def __init__(self) -> None: + self.now: float = 0.0 + + def __call__(self) -> float: + return self.now + + +def _count_secret_reads(monkeypatch: pytest.MonkeyPatch) -> _SecretReadCounter: + secret_reads: Final = _SecretReadCounter() + monkeypatch.setattr(litellm, "secret_manager_client", secret_reads) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + return secret_reads + + +async def _pin_a_turn_minted_by_the_origin( + affinity: object, healthy_deployments: list[dict[str, object]] +) -> tuple[list[dict[str, object]], Mapping[str, object]]: + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + assert isinstance(affinity, EncryptedContentAffinityCheck) + reasoning_input: Final = _reasoning_input_minted_by("origin") + result: Final[list[dict[str, object]]] = await affinity.async_filter_deployments( + model="gpt-5.5", + healthy_deployments=list(healthy_deployments), + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input}, + ) + return result, reasoning_input[0] + + +@pytest.mark.asyncio +async def test_repeated_pins_to_a_default_endpoint_peer_ask_the_secret_manager_once_a_minute( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", OPENAI_SHARED_KEY), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + healthy_deployments: Final = [peer, _deployment_of(ctx, "other-org")] + secret_reads: Final = _count_secret_reads(monkeypatch) + EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear() + clock: Final = _HandAdvancedClock() + affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=clock) + + first_pin, first_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + reads_for_the_first_pin: Final = secret_reads.reads + clock.now = 59.0 + second_pin, second_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + reads_within_the_minute: Final = secret_reads.reads + clock.now = 60.0 + third_pin, third_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + + assert first_pin == second_pin == third_pin == [peer] + assert first_reasoning == second_reasoning == third_reasoning == _reasoning_input_minted_by("origin")[0] + assert reads_for_the_first_pin > 0 + assert reads_within_the_minute == reads_for_the_first_pin + assert secret_reads.reads > reads_within_the_minute + + +@pytest.mark.asyncio +async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manager_refresh_minute( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + def origin_credential(api_key: str) -> list[CredentialItem]: + return [ + CredentialItem(credential_name="origin-account", credential_values={"api_key": api_key}, credential_info={}) + ] + + monkeypatch.setattr(litellm, "credential_list", origin_credential("shared-key")) + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment( + "origin", {"model": "openai/gpt-5.5", "litellm_credential_name": "origin-account"} + ), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + other_org: Final = _deployment_of(ctx, "other-org") + _count_secret_reads(monkeypatch) + EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear() + affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=_HandAdvancedClock()) + + pin_before_the_change, _ = await _pin_a_turn_minted_by_the_origin(affinity, [peer, other_org]) + monkeypatch.setattr(litellm, "credential_list", origin_credential("other-key")) + pin_after_the_change, reasoning_after_the_change = await _pin_a_turn_minted_by_the_origin( + affinity, [peer, other_org] + ) + + assert pin_before_the_change == [peer] + assert pin_after_the_change == [other_org] + assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0]