diff --git a/.github/workflows/test-terraform-modules.yml b/.github/workflows/test-terraform-modules.yml index 0e3e5330453..52006d9b578 100644 --- a/.github/workflows/test-terraform-modules.yml +++ b/.github/workflows/test-terraform-modules.yml @@ -4,6 +4,7 @@ on: push: paths: - "terraform/litellm/aws/**" + - "terraform/litellm/gcp/**" - ".github/workflows/test-terraform-modules.yml" pull_request: branches: @@ -13,6 +14,7 @@ on: - "litellm_**" paths: - "terraform/litellm/aws/**" + - "terraform/litellm/gcp/**" - ".github/workflows/test-terraform-modules.yml" permissions: @@ -52,3 +54,32 @@ jobs: # Plan-only, mock_provider-backed: no AWS credentials, no API calls. - name: test run: terraform test + + gcp-module: + name: fmt, validate, test (gcp) + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: terraform/litellm/gcp + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2 + with: + terraform_version: 1.13.3 + terraform_wrapper: false + + - name: fmt + run: terraform fmt -recursive -check -diff + + - name: init + run: terraform init -backend=false -input=false + + - name: validate + run: terraform validate + + - name: test + run: terraform test diff --git a/litellm/constants.py b/litellm/constants.py index e4fb2a00297..ce744e9c58a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -89,6 +89,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = ( # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)) +BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: Final = 256 * 1024 REDACTED_BY_LITELLM: Final = "redacted-by-litellm" # in-memory stand-in handed to provider converters for redacted arguments; never stored REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}" diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 1e2fa318786..2213c5fe275 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -97,7 +97,11 @@ class CloudZeroStreamer: continue # Convert lists back to DataFrames - return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records} + return { + date_key: pl.DataFrame(records, infer_schema_length=None) + for date_key, records in daily_batches.items() + if records + } def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime: """Parse timestamp string and convert to UTC.""" diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index ffc8fe1c1f5..12a0ee55fad 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -95,7 +95,7 @@ class CBFTransformer: if len(cbf_data) > 0: console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]") - return pl.DataFrame(cbf_data) + return pl.DataFrame(cbf_data, infer_schema_length=None) def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord: """Create a single CBF record from LiteLLM daily spend row.""" diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c462b1edb98..2d66a280663 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -948,6 +948,23 @@ class CustomGuardrail(CustomLogger): """ return False + def _suppressed_by_auto_router_compression(self) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. + + Reads request-scoped state set by `arm_pre_call`, never request metadata. The + caller controls metadata, and metadata reaches spend logs the caller can read, + so a suppression list carried there would be one a request could replay to + switch off a PII or content-filter guardrail for itself. + """ + name: Final = self.guardrail_name + if not name: + return False + from litellm.proxy.guardrails.auto_router_compression import ( + suppressed_compression_guardrails, + ) + + return name in suppressed_compression_guardrails() + def should_run_guardrail( self, data, @@ -956,6 +973,9 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ + if self._suppressed_by_auto_router_compression(): + return False + requested_guardrails: Final = self.get_guardrail_from_metadata(data) disable_global_guardrail: Final = self.get_disable_global_guardrail(data) opted_out_global_guardrails: Final = self.get_opted_out_global_guardrails_from_metadata(data) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8f1b2ce1cc2..83e0b4d84f1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -78,7 +78,10 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( InteractionsUsageObjectTransformation, ) -from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages +from litellm.litellm_core_utils.logging_utils import ( + truncate_base64_in_messages, + truncate_base64_in_messages_async, +) from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, @@ -538,6 +541,7 @@ class Logging(LiteLLMLoggingBaseClass): self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( self.initialize_standard_built_in_tools_params(kwargs) ) + self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: datetime.datetime | None = None self._llm_caching_handler: LLMCachingHandler | None = None @@ -1914,7 +1918,9 @@ class Logging(LiteLLMLoggingBaseClass): two paths cannot mutate it at the same time. ``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from ``completion()``); legacy string callbacks still run via - ``executor.submit(failure_handler)`` when configured. + ``executor.submit(failure_handler)`` when configured, and still get submitted + when the awaiting task is cancelled (e.g. the event loop shuts down right after + the request failed). """ litellm_params: Final = self.model_call_details.get("litellm_params", {}) or {} sync_sdk: Final = self._is_sync_litellm_request(litellm_params) @@ -1923,12 +1929,11 @@ class Logging(LiteLLMLoggingBaseClass): self.failure_handler(exception, traceback_exception) return - await self.async_failure_handler(exception, traceback_exception) - - if not self._should_run_sync_failure_callbacks_for_async_calls(): - return - - executor.submit(self.failure_handler, exception, traceback_exception) + try: + await self.async_failure_handler(exception, traceback_exception) + finally: + if self._should_run_sync_failure_callbacks_for_async_calls(): + executor.submit(self.failure_handler, exception, traceback_exception) def should_run_logging( self, @@ -2933,6 +2938,11 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage + self.truncated_messages_for_logging = await truncate_base64_in_messages_async( + StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=self.model_call_details, messages=self.model_call_details.get("messages") + ) + ) start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, end_time=end_time, @@ -3225,8 +3235,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details = {} if ( - self.model_call_details.get("log_event_type") == "failed_api_call" - and self.model_call_details.get("exception") is exception + self.model_call_details.get("exception") is exception and self.model_call_details.get("standard_logging_object") is not None ): return start_time, self.model_call_details["end_time"] @@ -6202,9 +6211,13 @@ def get_standard_logging_object_payload( model_id=_model_id, requester_ip_address=clean_metadata.get("requester_ip_address", None), user_agent=clean_metadata.get("user_agent", None), - messages=truncate_base64_in_messages( - StandardLoggingPayloadSetup.append_system_prompt_messages( - kwargs=kwargs, messages=kwargs.get("messages") + messages=( + logging_obj.truncated_messages_for_logging + if logging_obj.truncated_messages_for_logging is not None + else truncate_base64_in_messages( + StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=kwargs.get("messages") + ) ) ), response=final_response_obj, diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index f3b1b29a9ad..44daef42e14 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -3,12 +3,15 @@ import functools import inspect import re import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger -from litellm.constants import MAX_BASE64_LENGTH_FOR_LOGGING +from litellm.constants import ( + BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS, + MAX_BASE64_LENGTH_FOR_LOGGING, +) from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -141,6 +144,39 @@ def truncate_base64_in_messages( return messages +_StringTree = str | Sequence["_StringTree"] | Mapping[str, "_StringTree"] | None + + +def _iter_string_leaves(value: _StringTree) -> Iterator[str]: + stack: Final[list[_StringTree]] = [value] # mutable-ok: explicit stack, recursive functions are banned in litellm/ + while stack: + match stack.pop(): + case str() as text: + yield text + case Mapping() as mapping: + stack.extend(mapping.values()) + case Sequence() as items: + stack.extend(items) + case None: + pass + + +async def truncate_base64_in_messages_async( + messages: str | list | dict | None, # mutable-ok: same contract as truncate_base64_in_messages +) -> str | list | dict | None: # mutable-ok: same contract as truncate_base64_in_messages + """ + Same result as truncate_base64_in_messages, but payloads whose string content + reaches BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are scanned in a worker + thread so the regex pass over multi-MB base64 images does not block the event loop. + """ + if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: + return messages + total_chars: Final = sum(len(leaf) for leaf in _iter_string_leaves(messages)) + if total_chars < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: + return truncate_base64_in_messages(messages) + return await asyncio.to_thread(truncate_base64_in_messages, messages) + + # Global service logger instance to avoid recreating it _service_logger = None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bcbcc6bc579..dc1e8db1628 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -159,6 +159,9 @@ from litellm.proxy._types import ( from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + id_jag_assertion_capture_gap_at_startup, +) from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider @@ -1382,6 +1385,20 @@ def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str ) +def _warn_config_id_jag_server_outruns_sso(server: MCPServer) -> None: + if server.auth_type != MCPAuth.oauth2_id_jag: + return + gap: Final = id_jag_assertion_capture_gap_at_startup() + if gap is None: + return + verbose_logger.warning( + "MCP server %r (id=%s, source=config) is declared with auth_type=oauth2_id_jag, but %s.", + get_server_prefix(server), + server.server_id, + gap, + ) + + def _deserialize_json_dict(data: str | _StringMap | None) -> dict[str, str] | None: """ Deserialize optional JSON mappings stored in the database. @@ -2393,6 +2410,7 @@ class MCPServerManager: ) self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") + _warn_config_id_jag_server_outruns_sso(new_server) self.config_mcp_servers[server_id] = new_server self._set_oauth_discovery_deferred( server_id, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 9273ddda9cf..5e28396dcfb 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -31,11 +31,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto TokenCacheBackend, TokenStoreUnavailable, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( - RedisDistributedLock, -) -from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( - RedisRefreshCoordinator, +from litellm.proxy._experimental.mcp_server.outbound_credentials.runtime_refresh_coordinator import ( + runtime_refresh_coordinator, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import ( OAuthTokenCacheCodec, @@ -131,23 +128,17 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres ) from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 - redis_cache: Final = user_api_key_cache.redis_cache - if redis_cache is None: + coordinator: Final = runtime_refresh_coordinator() + if coordinator is None: return None, None, False codec: Final = OAuthTokenCacheCodec( encrypt_value_helper, lambda blob: decrypt_value_helper(blob, "mcp_per_user_token", exception_type="debug"), ) - # user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) and the - # Redis client from init_async_client() is partially typed - both are untyped-boundary casts. + # user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) - an + # untyped-boundary cast. cache: Final[AsyncCache] = user_api_key_cache # pyright: ignore - redis_client: Final = redis_cache.init_async_client() # pyright: ignore - lock: Final = RedisDistributedLock( - redis_client, # pyright: ignore - namespace_key=redis_cache.check_and_fix_namespace, - ) backend: Final = DualCacheTokenCacheBackend(cache, codec) - coordinator: Final = RedisRefreshCoordinator(lock) return backend, coordinator, True diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 404baa14350..8328aae01ab 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -46,11 +46,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + default_sso_assertion_store, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( AssertionStoreUnavailable, - DbSSOAssertionStore, SSOAssertionStore, - SSOIdentityAssertion, + assertion_expired, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, @@ -129,7 +131,7 @@ class UpstreamCredentialProvider: self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient() self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache() self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource() - self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or DbSSOAssertionStore() + self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or default_sso_assertion_store() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -246,7 +248,7 @@ class UpstreamCredentialProvider: "Sign in through LiteLLM SSO so the gateway captures one." ) ) - if _assertion_expired(assertion, datetime.now(timezone.utc)): + if assertion_expired(assertion, datetime.now(timezone.utc)): return Error( CredError.of_precondition_required( "The stored IdP identity assertion for this user has expired. Sign in through " @@ -405,19 +407,6 @@ def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str: return hashlib.sha256(material.encode()).hexdigest() -def _assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: - """Whether the stored assertion's ``exp`` has passed. An assertion carrying no expiry is - treated as usable and left for the IdP to reject, since the store records what the id_token - claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a - stored value that lost its offset compares instead of raising. - """ - expires_at: Final = assertion.expires_at - if expires_at is None: - return False - normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) - return normalized <= now - - def _id_jag_fingerprint(subject_token: str, server_id: str, config: IdJagConfig) -> str: """What the cached leg-2 bearer was minted from: the subject token, the server, and the config. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py new file mode 100644 index 00000000000..e799838b5b7 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -0,0 +1,41 @@ +"""The runtime ``RefreshCoordinator``: cross-replica single-flight when Redis is wired. + +Builds ``RedisRefreshCoordinator`` over the proxy's shared Redis so one refresh runs per key +across the fleet, or returns ``None`` when Redis is absent so the caller keeps the foundation's +in-process default (correct for a single replica). The proxy globals it reads are not ready at +import time, so this is called per composition rather than held as module state. + +Shared by every credential arm that renews a stored grant: a rotating refresh token must be +redeemed once across all workers, so each arm electing its own winner with its own lock shape +would be a bug waiting to differ. +""" + +from __future__ import annotations + +from typing import Final + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + RefreshCoordinator, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( + RedisDistributedLock, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( + RedisRefreshCoordinator, +) + + +def runtime_refresh_coordinator() -> RefreshCoordinator | None: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # runtime global + + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is None: + return None + # The Redis client from init_async_client() is only partially typed; the lock validates every + # reply it depends on, so the untyped boundary is contained here. + redis_client: Final = redis_cache.init_async_client() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm redis wrapper is untyped + lock: Final = RedisDistributedLock( + redis_client, # pyright: ignore[reportArgumentType,reportUnknownArgumentType] # litellm redis wrapper is untyped + namespace_key=redis_cache.check_and_fix_namespace, + ) + return RedisRefreshCoordinator(lock) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py new file mode 100644 index 00000000000..7600fd7ab8a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py @@ -0,0 +1,469 @@ +"""Renew the stored SSO identity assertion so an ID-JAG agent outlives one id_token. + +The ``oauth2_id_jag`` arm asserts the id_token captured at the user's last interactive sign-in, so +without renewal an agent holding a brokered LiteLLM key can act for that user only until that token's +``exp``, typically an hour, and the sole recovery is another interactive login. The assertion already +carries the IdP refresh token beside it; this module is what redeems it. + +``RefreshingSSOAssertionStore`` wraps any ``SSOAssertionStore`` and satisfies the same protocol, so +the egress arm is unchanged: it still reads one assertion and still judges expiry itself. Renewal is +lazy (only a read that finds a near-expiry assertion triggers one, so IdP traffic tracks actual use, +not the size of the user table) and single-flighted per user through the same ``RefreshCoordinator`` +the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent +redemptions of one token as replay and can revoke the whole grant chain. + +The refresh is redeemed against the generic-OIDC client the login itself used +(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy +reconciles from the stored SSO row into the process environment at startup), authenticated the way +that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the +credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept +only one of the two. An assertion can only exist if that client minted it, so no other client could +redeem its refresh token, and no other method is known to be accepted. A deployment whose +``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss +logs the scope by name rather than failing silently. + +Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse +onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged +so the reader's own guard challenges the user to sign in again, while a transient IdP failure raises +``AssertionStoreUnavailable`` so the reader answers 503 instead of blaming the user for an outage. + +One ambiguity remains under Redis-coordinated renewal across replicas. A cross-replica loser that +finds the row still expiring after the holder finished cannot tell a refused refresh from a renewal +that could not be recorded. Redeeming itself could consume a refresh token the holder may already +have rotated, so it answers retryable 503 rather than guessing a sign-in challenge. The next +uncontended read settles the outcome itself: a refusal challenges, and a successful refresh persists. +If the holder rotated the token but its write failed, that rotation is lost and the next uncontended +read's refusal challenges, which is the only honest answer because the rotated token was never +recorded. On the refusal path, the loser pays for one retry before that challenge. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final, Literal, Protocol + +import httpx +from pydantic import SecretStr, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import Timeout +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped +) +from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( + build_token_endpoint_client_auth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InProcessRefreshCoordinator, + RefreshCoordinator, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.runtime_refresh_coordinator import ( + runtime_refresh_coordinator, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + DbSSOAssertionStore, + SSOAssertionStore, + SSOIdentityAssertion, + assertion_expired, + assertion_from_sso_login, + fetch_sso_identity_assertion, + persist_sso_identity_assertion, +) +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.mcp import MCPTokenEndpointAuthMethod + +_BODY_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(dict[str, object]) + +_REFRESH_GRANT_TYPE: Final = "refresh_token" +# The lock namespace for the one assertion row a user has; the sibling arm keys the same lock by +# server_id, and no server_id can collide with this literal. +_SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion" +# Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of +# the exchange is replaced first. Matches the sibling per-user token store's skew. +_DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0 + + +class AssertionRead(Protocol): + """Reads the user's stored assertion row.""" + + async def __call__(self, user_id: str) -> SSOIdentityAssertion | None: ... + + +class AssertionWrite(Protocol): + """Replaces the user's stored assertion row.""" + + async def __call__(self, user_id: str, assertion: SSOIdentityAssertion) -> None: ... + + +class CoordinatorFactory(Protocol): + """Builds the cross-replica coordinator, or ``None`` when there is no shared lock to build on.""" + + def __call__(self) -> RefreshCoordinator | None: ... + + +class FormPost(Protocol): + """POSTs an OAuth form and hands back the raw response.""" + + async def __call__( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> httpx.Response | None: ... + + +@dataclass(frozen=True, slots=True) +class SSOClientConfig: + """The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how.""" + + token_endpoint: str + client_id: str + client_secret: SecretStr + auth_method: MCPTokenEndpointAuthMethod + + +def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None: + """The configured generic-OIDC client, or ``None`` when the deployment has none. + + Read from the process environment because that is where the login path reads it + (``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row + at startup, so this resolves to the same client that minted the assertion. ``None`` is an + ordinary state, not an error: a deployment signing in through a provider that captures no + assertion has nothing here to renew, and a client with no secret is not a confidential client + that could redeem one. + + ``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because + the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while + the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so + a constant here would authenticate the renewal differently from the sign-in that produced the + refresh token and 401 against an IdP application registered for only one of the two. + """ + token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT") + client_id: Final = env.get("GENERIC_CLIENT_ID") + client_secret: Final = env.get("GENERIC_CLIENT_SECRET") + if not token_endpoint or not client_id or not client_secret: + return None + includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" + return SSOClientConfig( + token_endpoint=token_endpoint, + client_id=client_id, + client_secret=SecretStr(client_secret), + auth_method="client_secret_post" if includes_client_id else "client_secret_basic", + ) + + +@dataclass(frozen=True, slots=True) +class RefreshFailure: + """Why a renewal produced nothing, split by what the caller can do about it. + + ``rejected`` is settled: this refresh token will never work again, so the user has to sign in. + ``unavailable`` is transient: the same attempt may succeed in a minute, so telling the user to + sign in again would be a lie about whose problem it is. Both arms carry the same payload, so + this is a ``Literal`` discriminant rather than a ``tagged_union``; consumers still ``match`` on + ``kind`` with an ``assert_never`` tail. + """ + + kind: Literal["rejected", "unavailable"] + detail: str + + @staticmethod + def of_rejected(detail: str) -> RefreshFailure: + return RefreshFailure(kind="rejected", detail=detail) + + @staticmethod + def of_unavailable(detail: str) -> RefreshFailure: + return RefreshFailure(kind="unavailable", detail=detail) + + +class TokenEndpointTransport(Protocol): + """One form POST to the IdP token endpoint, with the refusal/outage split preserved. + + That split is the whole reason this is not the resolver's ``TokenEndpointClient``: that + collaborator maps every non-2xx to ``upstream_unavailable``, which is right for an exchange leg + and wrong here, where a 400 ``invalid_grant`` means the stored refresh token is dead and the user + must act. + """ + + async def post( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> Result[Mapping[str, object], RefreshFailure]: ... + + +async def post_form(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None: + # litellm's httpx handler is only partially typed; nothing but the response object crosses back, + # and the transport below validates its body, so the untyped boundary is contained here. + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped + return await client.post(url, data=form, headers=headers) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportReturnType,reportArgumentType] # litellm http handler is untyped and its stub narrows data=/headers= to dict, which httpx itself does not require + + +class HttpxTokenEndpointTransport: + """The live transport. 4xx is the IdP refusing this grant; anything else is an outage. + + The POST itself is injected so that split, which decides whether the user is challenged or told + to wait, is testable without a live IdP. + """ + + def __init__(self, post: FormPost = post_form) -> None: + self._post = post + + async def post( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> Result[Mapping[str, object], RefreshFailure]: + try: + response: Final = await self._post(url, form, headers) + if response is None: + return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned no response")) + response.raise_for_status() + body: Final = _BODY_ADAPTER.validate_python(response.json()) # pyright: ignore[reportAny] # untyped JSON; the adapter is the type gate + except httpx.HTTPStatusError as exc: + status: Final = exc.response.status_code + if 400 <= status < 500: + return Error(RefreshFailure.of_rejected(f"the IdP refused the refresh with status {status}")) + return Error(RefreshFailure.of_unavailable(f"the IdP token endpoint answered with status {status}")) + except (httpx.RequestError, Timeout) as exc: + return Error(RefreshFailure.of_unavailable(f"the IdP token endpoint is unreachable ({type(exc).__name__})")) + except json.JSONDecodeError: + return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned a non-JSON response")) + except ValidationError: + return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned a non-object response")) + return Ok(body) + + +class SSOAssertionRefresher: + """Redeems the stored refresh token for a current id_token and writes the rotation back. + + Collaborators are injected so the orchestration, the untyped response parsing and the + write-back race are all testable without an IdP or a database. + """ + + def __init__( + self, + transport: TokenEndpointTransport, + *, + client_config: Callable[[], SSOClientConfig | None] = lambda: sso_client_config(os.environ), + read: AssertionRead = fetch_sso_identity_assertion, + write: AssertionWrite = persist_sso_identity_assertion, + ) -> None: + self._transport = transport + self._client_config = client_config + self._read = read + self._write = write + + async def refresh( + self, user_id: str, assertion: SSOIdentityAssertion + ) -> Result[SSOIdentityAssertion, RefreshFailure]: + if assertion.refresh_token is None: + verbose_proxy_logger.warning( + "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was " + "captured with it, so it cannot be renewed without another interactive sign-in. Add " + "'offline_access' to GENERIC_SCOPE so the SSO login captures one.", + user_id, + ) + return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in")) + config: Final = self._client_config() + if config is None: + verbose_proxy_logger.warning( + "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed " + "because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, " + "GENERIC_CLIENT_SECRET).", + user_id, + ) + return Error(RefreshFailure.of_rejected("the generic SSO client is not configured")) + + carried_refresh_token: Final = assertion.refresh_token.get_secret_value() + # Whichever method the SSO login used for this client, since that is the one the IdP + # application is known to accept: an assertion only exists to renew because a sign-in already + # authenticated this client that way. + client_auth: Final = build_token_endpoint_client_auth( + auth_method=config.auth_method, + client_id=config.client_id, + client_secret=config.client_secret.get_secret_value(), + ) + form: Final = { # mutable-ok: the RFC 6749 form body is a wire format the HTTP client takes as a mapping + "grant_type": _REFRESH_GRANT_TYPE, + "refresh_token": carried_refresh_token, + **client_auth.body, + } + match await self._transport.post(config.token_endpoint, form, client_auth.headers): + case Error(failure): + return Error(failure) + case Ok(body): + return await self._renewed_from(user_id, assertion, body, carried_refresh_token) + + async def _renewed_from( + self, + user_id: str, + previous: SSOIdentityAssertion, + body: Mapping[str, object], + carried_refresh_token: str, + ) -> Result[SSOIdentityAssertion, RefreshFailure]: + """The renewed assertion, built by the same validator the login path uses. + + A rotated refresh token replaces the stored one; an omitted one carries forward, since an + IdP that does not rotate expects the original to keep working. + """ + rotated: Final = body.get("refresh_token") + renewed: Final = assertion_from_sso_login( + body.get("id_token"), + rotated if isinstance(rotated, str) and rotated else carried_refresh_token, + ) + if renewed is None: + verbose_proxy_logger.warning( + "ID-JAG: the IdP accepted the refresh for user_id=%s but returned no usable id_token, so there " + "is nothing to assert upstream. The SSO client's grant needs the 'openid' scope for the token " + "endpoint to return one on a refresh.", + user_id, + ) + return Error(RefreshFailure.of_rejected("the IdP's refresh response carried no usable id_token")) + failure: Final = await self._store_renewal(user_id, previous, renewed) + if failure is not None: + return Error(failure) + return Ok(renewed) + + async def _store_renewal( + self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion + ) -> RefreshFailure | None: + """Write the renewal back, unless the row moved on while this renewal was in flight. + + The row is one per user and last-write-wins, so an interactive sign-in landing mid-renewal + would otherwise be overwritten with a refresh token the IdP has already rotated away, costing + that user a sign-in later. Comparing against the id_token this renewal started from is what + detects that; skipping is safe because the newer row is the one the reader wants anyway. + + A failed write is transient, not settled. The store, not this return value, is what every + caller reads, so a renewal that could not be recorded is a renewal nobody will see; saying so + keeps a database problem answering 503 rather than telling the user to sign in again over it. + """ + try: + current: Final = await self._read(user_id) + if current is not None and current.id_token.get_secret_value() != previous.id_token.get_secret_value(): + verbose_proxy_logger.info( + "ID-JAG: a newer IdP identity assertion for user_id=%s was stored while this renewal was in " + "flight; keeping the stored one.", + user_id, + ) + return None + await self._write(user_id, renewed) + except Exception as exc: # noqa: BLE001 # any storage failure is transient here, never the user's fault + verbose_proxy_logger.warning( + "ID-JAG: could not persist the renewed IdP identity assertion for user_id=%s, so the rotated " + "refresh token is lost and this user will have to sign in again once the renewed token expires: %s", + user_id, + exc, + ) + return RefreshFailure.of_unavailable("the renewed IdP identity assertion could not be persisted") + return None + + +class RefreshingSSOAssertionStore: + """An ``SSOAssertionStore`` that renews a near-expiry assertion before handing it back. + + Reads the inner store; an assertion still comfortably inside its lifetime is returned untouched, + so the common path costs exactly what it did before. Otherwise one renewal runs per user through + the injected ``RefreshCoordinator`` and every caller then re-reads the inner store, which is the + authority: the winner's write is what they all observe, and a renewal the write-back guard + skipped yields the newer assertion that displaced it rather than a private copy. + + A refusal leaves the expired assertion in place for the reader's own guard to reject, so the user + sees the same sign-in-again challenge as before this store existed. A transient IdP failure + raises ``AssertionStoreUnavailable``, the protocol's existing signal for "this is not the user's + fault"; concurrent in-process callers share that outcome, while a cross-replica loser answers 503 + when its re-read still finds the row expiring. On the refusal path that costs the loser one retry, + which then challenges. If the holder rotated the token but its write failed, the rotation is lost + and the next uncontended read's refusal challenges, the only honest answer because that token was + never recorded. + """ + + def __init__( + self, + inner: SSOAssertionStore, + refresher: SSOAssertionRefresher, + *, + fresh_read: AssertionRead, + coordinator_factory: CoordinatorFactory = runtime_refresh_coordinator, + expiry_skew_seconds: float = _DEFAULT_EXPIRY_SKEW_SECONDS, + clock: Callable[[], datetime] = lambda: datetime.now(timezone.utc), + ) -> None: + self._inner = inner + self._refresher = refresher + self._fresh_read = fresh_read + self._coordinator_factory = coordinator_factory + self._in_process_coordinator = InProcessRefreshCoordinator() + self._distributed_coordinator: RefreshCoordinator | None = None + self._skew = timedelta(seconds=expiry_skew_seconds) + self._clock = clock + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + assertion: Final = await self._inner.fetch(user_id) + if not self._expiring(assertion): + return assertion + await self._coordinator().run( + user_id, + _SINGLE_FLIGHT_KEY, + refresh=lambda: self._renew(user_id), + reread=lambda: self._reread_renewed(user_id), + ) + return await self._fresh_read(user_id) + + def _expiring(self, assertion: SSOIdentityAssertion | None) -> bool: + return assertion is not None and assertion_expired(assertion, self._clock() + self._skew) + + def _coordinator(self) -> RefreshCoordinator: + """The cross-replica coordinator once Redis is reachable, else the in-process one. + + Built on first use and kept, because the proxy's Redis client is not wired at import time; + retried while it is absent so a proxy that gains Redis later stops electing per-worker. + """ + if self._distributed_coordinator is None: + self._distributed_coordinator = self._coordinator_factory() + return self._distributed_coordinator or self._in_process_coordinator + + async def _renew(self, user_id: str) -> None: + """The elected renewal, judged from a fresh read so a rotation another replica just landed is + never redeemed again. Returns nothing: the inner store, not this return value, is what every + caller reads afterwards, so the winner and the losers cannot disagree.""" + latest: Final = await self._fresh_read(user_id) + if latest is None or not self._expiring(latest): + return + match await self._refresher.refresh(user_id, latest): + case Ok(_): + return + case Error(failure): + match failure.kind: + case "rejected": + return + case "unavailable": + raise AssertionStoreUnavailable(failure.detail) + assert_never(failure.kind) + + async def _reread_renewed(self, user_id: str) -> None: + """A loser cannot distinguish refusal from an unrecorded renewal without risking token replay. + + It answers retryable 503 instead of guessing a sign-in challenge; the retry runs uncontended + and settles the outcome itself. + """ + latest: Final = await self._fresh_read(user_id) + if self._expiring(latest): + raise AssertionStoreUnavailable( + f"the IdP identity assertion for user_id={user_id} was being renewed by another replica " + "and is not yet current; retry shortly" + ) + + +def default_sso_assertion_store() -> SSOAssertionStore: + """The live read seam for the ``id_jag`` arm: the stored assertion, renewed when it is stale.""" + db_store: Final = DbSSOAssertionStore() + fresh_read: Final = db_store.fetch_uncached + return RefreshingSSOAssertionStore( + db_store, + SSOAssertionRefresher(HttpxTokenEndpointTransport(), read=fresh_read), + fresh_read=fresh_read, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 6552008ca54..f7b92df5ba3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -127,6 +127,23 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden ) +def assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: + """Whether the assertion's ``exp`` has passed at ``now``. An assertion carrying no expiry is + treated as usable and left for the IdP to reject, since the store records what the id_token + claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a + stored value that lost its offset compares instead of raising. + + Lives beside the model rather than in either reader so the egress guard and the renewal + trigger judge the same field the same way; passing a ``now`` in the future is how a caller + asks "is this about to expire" without a second, driftable predicate. + """ + expires_at: Final = assertion.expires_at + if expires_at is None: + return False + normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) + return normalized <= now + + async def ema_assertion_retention_enabled() -> bool: """Whether any MCP server uses ``oauth2_id_jag``, evaluated per login so the gateway only retains bearer material while an EMA upstream exists to spend it on. Judged against the two @@ -146,7 +163,9 @@ async def ema_assertion_retention_enabled() -> bool: return True if prisma_client is None: return False - row = await prisma_client.db.litellm_mcpservertable.find_first(where={"auth_type": MCPAuth.oauth2_id_jag.value}) + row: Final = await prisma_client.db.litellm_mcpservertable.find_first( + where={"auth_type": MCPAuth.oauth2_id_jag.value} + ) return row is not None @@ -158,7 +177,7 @@ async def persist_sso_identity_assertion( if prisma_client is None: return - payload: Final[dict[str, str]] = { + payload: Final = { "id_token": assertion.id_token.get_secret_value(), **({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}), **({"issuer": assertion.issuer} if assertion.issuer else {}), @@ -220,11 +239,13 @@ async def fetch_sso_identity_assertion( class AssertionStoreUnavailable(Exception): - """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down). + """Raised by ``fetch`` when the assertion cannot be read for a transient reason: the DB is + down, or the IdP behind a renewing store could not be reached. Distinct from returning ``None`` for "this user has no captured assertion": an outage must not read as a definite absence, which would tell the user to sign in again over a transient failure, - and it must not escape as an unhandled error on the egress or retry path. Mirrors + and it must not escape as an unhandled error on the egress or retry path. The message names the + real component for the operator log; callers get the reader's generic 503. Mirrors ``TokenStoreUnavailable`` on the sibling per-user OAuth store. """ @@ -257,6 +278,12 @@ class DbSSOAssertionStore: except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence raise AssertionStoreUnavailable(str(exc)) from exc + async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None: + try: + return await _read_assertion_from_db(user_id) + except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence + raise AssertionStoreUnavailable(str(exc)) from exc + async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, new_master_key: str) -> None: """Re-encrypt every stored assertion under ``new_master_key`` during a salt-key rotation, @@ -280,7 +307,9 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, row.user_id, ) return False - re_encrypted = _STR_ADAPTER.validate_python(encrypt_value_helper(plaintext, new_encryption_key=new_master_key)) + re_encrypted: Final = _STR_ADAPTER.validate_python( + encrypt_value_helper(plaintext, new_encryption_key=new_master_key) + ) await prisma_client.db.litellm_ssoidentityassertion.update( where={"user_id": row.user_id}, data={"assertion_b64": re_encrypted}, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f25fa46197e..5e6c9b34332 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router @@ -2004,6 +2005,12 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + # An auto router with its own compression policy is authoritative for this + # request: suppress every other compression guardrail and arm whichever one + # the policy names for the model call, before those guardrails get a chance + # to run below. + await _arm_auto_router_compression(data=self.data, llm_router=llm_router) + self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py new file mode 100644 index 00000000000..98707e7ddca --- /dev/null +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -0,0 +1,267 @@ +""" +Decouples prompt compression between an auto router's routing decision and the model +it routes to, via ``auto_router_routing_compression`` / ``auto_router_model_compression`` +on the marker deployment: a guardrail name, or ``"none"``. + +Neither key set inherits today's behaviour. Either key set makes the auto router +authoritative and suppresses every other compression guardrail for that request. +""" + +import contextvars +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.router import Router + +COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) +_NO_COMPRESSION: Final = "none" + +# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a +# suppression list they can read is one they can replay to disable any guardrail. +_suppressed_compression_guardrails: Final[contextvars.ContextVar[frozenset[str]]] = contextvars.ContextVar( + "litellm_auto_router_suppressed_compression_guardrails", default=frozenset() +) + + +def suppressed_compression_guardrails() -> frozenset[str]: + """Names of the compression guardrails this request's auto router suppresses.""" + return _suppressed_compression_guardrails.get() + + +# Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and nothing +# compresses; the router must not assume the model hop already ran. +_model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( + "litellm_auto_router_model_hop_armed", default=False +) + + +def model_hop_compression_armed() -> bool: + """True when this request's model-side compression guardrail was actually armed.""" + return _model_hop_armed.get() + + +@dataclass(frozen=True, slots=True) +class AutoRouterCompressionPolicy: + """An auto router's compression choice for each hop. ``None`` means no compression.""" + + routing: str | None + model: str | None + + @property + def is_same(self) -> bool: + return self.routing == self.model + + +def _normalized_compression_choice(raw: object) -> str | None: + if not isinstance(raw, str) or not raw: + return None + return None if raw.strip().lower() == _NO_COMPRESSION else raw + + +def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRouterCompressionPolicy | None: + raw_routing: Final = litellm_params.get("auto_router_routing_compression") + raw_model: Final = litellm_params.get("auto_router_model_compression") + if raw_routing is None and raw_model is None: + return None + return AutoRouterCompressionPolicy( + routing=_normalized_compression_choice(raw_routing), + model=_normalized_compression_choice(raw_model), + ) + + +def policy_for_model( + llm_router: "Router | None", + model_alias: str, + team_id: str | None, + request_tags: Sequence[str], +) -> AutoRouterCompressionPolicy | None: + """The compression policy of the auto router marker `model_alias` resolves to. + + Pre-call arming and the routing hook both resolve through here, so an alias with + several tag-scoped markers cannot suppress under one and then route under another. + """ + if llm_router is None: + return None + deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () + markers: Final = tuple( + litellm_params + for deployment in deployments + if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) # pyright: ignore[reportUnnecessaryIsInstance] # filters out non-Mapping + and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + ) + requested: Final = frozenset(request_tags) + tag_matched: Final = tuple( + params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) + ) + # Untagged only: a marker scoped to tags this request lacks describes other traffic. + untagged: Final = tuple(params for params in markers if not params.get("tags")) + # Lazy, so the first marker carrying a policy wins and the rest are never read. + candidates: Final = (policy_from_litellm_params(params) for params in (*tag_matched, *untagged)) + return next((policy for policy in candidates if policy is not None), None) + + +def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + for meta_key in ("metadata", "litellm_metadata"): + meta = request_kwargs.get(meta_key) + if isinstance(meta, Mapping): + team_id = meta.get("user_api_key_team_id") + if isinstance(team_id, str): + return team_id + return None + + +def _compression_guardrail_classes() -> tuple[type, ...]: + """The registered guardrail classes whose provider compresses prompts.""" + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + return tuple(cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS) + + +def is_compression_guardrail(guardrail: object) -> bool: + """Whether `guardrail` is an instance of a compression guardrail provider. + + Both hops validate through here: the policy fields are operator-supplied names, and + an unvalidated one would get handed the conversation and invoked. + """ + classes: Final = _compression_guardrail_classes() + return bool(classes) and isinstance(guardrail, classes) + + +def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: + """Every currently-active guardrail whose type is a compression guardrail.""" + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + + if not _compression_guardrail_classes(): + return () + active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name) + + +async def arm_pre_call( + data: dict[str, object], # mutable-ok: arms the live request dict in place + llm_router: "Router | None", +) -> None: + """Apply an auto router's compression policy, if any, before guardrails run. + + Suppresses every other compression guardrail and re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``. + """ + _suppressed_compression_guardrails.set(frozenset()) + _model_hop_armed.set(False) + if llm_router is None: + return + + model_alias: Final = data.get("model") + if not isinstance(model_alias, str) or not model_alias: + return + + from litellm.router_strategy.tag_based_routing import ( + _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # used in router.py and budget_limiter.py too + ) + + policy: Final = policy_for_model( + llm_router=llm_router, + model_alias=model_alias, + team_id=team_id_from_request(data), + request_tags=_get_tags_from_request_kwargs(data), + ) + if policy is None: + return + + _suppressed_compression_guardrails.set( + frozenset( + name + for guardrail in _active_compression_guardrails() + if (name := guardrail.guardrail_name) and name != policy.model + ) + ) + + # Arming adds the name to `metadata["guardrails"]`, which runs it even if not default_on. + armed_model_hop: Final = policy.model is not None and any( + guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails() + ) + if policy.model is not None and not armed_model_hop: + verbose_proxy_logger.warning( + "AutoRouter compression: '%s' is not an active compression guardrail; the model hop is uncompressed", + policy.model, + ) + + if armed_model_hop: + _model_hop_armed.set(True) + _, metadata = get_or_create_metadata_bucket(data) + requested: Final = metadata.get("guardrails") + existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () + if policy.model not in existing: + # A list: litellm_pre_call_utils isinstance-checks this key and drops a tuple. + metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list + + +def _as_routing_messages( + messages: Iterable[Mapping[str, object]], +) -> list[dict[str, object]]: # mutable-ok: shape fixed by the pre-routing hook protocol + """A fresh, independently mutable copy, the shape the pre-routing hook takes.""" + return [dict(message) for message in messages] # mutable-ok: shape fixed by the pre-routing hook protocol + + +async def messages_for_routing( + policy: AutoRouterCompressionPolicy | None, + # list[dict], not Sequence[Mapping]: fixed by the async_pre_routing_hook protocol. + messages: list[dict[str, object]] | None, # mutable-ok: shape fixed by the pre-routing hook protocol + request_kwargs: Mapping[str, object], +) -> list[dict[str, object]] | None: # mutable-ok: shape fixed by the pre-routing hook protocol + """Messages to use for a routing decision, per `policy.routing`. None means the + caller should route on whatever it already has. + + Reads the live messages, never a pre-guardrail copy: this compresses through a real + guardrail that POSTs the text out, so routing on a pre-masking snapshot would leak + what the masking guardrail stripped. When the model hop already compressed and the + hops differ, routing therefore reads the compressed text rather than the original. + """ + if policy is None or policy.routing is None: + return None + + if not messages: + return None + + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) + if guardrail is None: + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing + ) + return _as_routing_messages(messages) + + if not is_compression_guardrail(guardrail): + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages", + policy.routing, + ) + return _as_routing_messages(messages) + + inputs: Final[GenericGuardrailAPIInputs] = { + "structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape + } + model: Final = request_kwargs.get("model") + # Throwaway: apply_guardrail writes stats here, so routing never double-counts into + # extract_compression_saved_tokens. + stats_sink: Final = {"messages": messages, "model": model} # mutable-ok: apply_guardrail writes its stats here + result: Final = await guardrail.apply_guardrail( + inputs=inputs, + request_data=stats_sink, + input_type="request", + ) + compressed: Final = result.get("structured_messages") + return compressed if isinstance(compressed, list) else _as_routing_messages(messages) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 40cc2e57932..ae266792391 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -64,6 +64,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + id_jag_assertion_capture_gap, +) from litellm.proxy.management_helpers.audit_logs import ( get_audit_log_changed_by, is_audit_logging_enabled, @@ -272,6 +275,22 @@ if MCP_AVAILABLE: _validate_mcp_server_name_fields(payload) _validate_upstream_token_header(payload) + def warn_if_id_jag_server_outruns_sso(server_id: str | None, auth_type: MCPAuth | str | None) -> None: + """Registering an ``oauth2_id_jag`` server under an SSO provider that captures no IdP + identity assertion is a dead configuration: nothing here fails, and then every ID-JAG call + fails for every user with a message that only ever tells them to sign in again. Say it once, + at the moment the admin can still act on it.""" + if auth_type != MCPAuth.oauth2_id_jag: + return + gap = id_jag_assertion_capture_gap() + if gap is None: + return + verbose_proxy_logger.warning( + "MCP server %s is registered with auth_type=oauth2_id_jag, but %s.", + server_id, + gap, + ) + def stamp_omitted_oauth2_flow(payload: NewMCPServerRequest) -> None: """Fallback only: fill in oauth2_flow when an oauth2 create omits it. @@ -1623,6 +1642,8 @@ if MCP_AVAILABLE: detail={"error": f"Error creating mcp server: {e}"}, ) + warn_if_id_jag_server_outruns_sso(new_mcp_server.server_id, new_mcp_server.auth_type) + # Registry refresh is best-effort: the row is already committed, so a # failure here (e.g. an unrelated malformed row in the table) must not # surface as a 500 and orphan the created server, which would push the @@ -2726,6 +2747,7 @@ if MCP_AVAILABLE: status_code=status.HTTP_404_NOT_FOUND, detail={"error": f"MCP Server not found, passed server_id={payload.server_id}"}, ) + warn_if_id_jag_server_outruns_sso(mcp_server_record_updated.server_id, mcp_server_record_updated.auth_type) await global_mcp_server_manager.update_server(mcp_server_record_updated) # Ensure registry is up to date by reloading from database diff --git a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py new file mode 100644 index 00000000000..404fdfc83a9 --- /dev/null +++ b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py @@ -0,0 +1,81 @@ +"""Whether the SSO provider the login callback dispatches to can capture an IdP identity assertion. + +An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693 +subject token. Only the generic OIDC login path reaches a token response the gateway retains one +from, so a deployment whose SSO runs through Google, Microsoft or SAML never stores an assertion +and every store-sourced ID-JAG exchange fails for every user, however many times they sign in. +Neither side can see that alone: the MCP registration knows nothing about SSO and the login knows +nothing about MCP. This module is the one shared answer both warn from. +""" + +from __future__ import annotations + +import os +from enum import Enum + +from typing_extensions import assert_never + +from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler + +_GENERIC_OIDC_REMEDY = ( + "Point SSO at the generic OIDC provider (GENERIC_CLIENT_ID), the one login path whose token " + "response the gateway retains an id_token from" +) + + +class ActiveSSOProvider(str, Enum): + google = "google" + microsoft = "microsoft" + generic = "generic" + saml = "saml" + none = "none" + + +def active_sso_provider() -> ActiveSSOProvider: + """The provider the SSO callback will dispatch to. + + Mirrors the callback's precedence rather than reporting everything configured: an environment + carrying both GOOGLE_CLIENT_ID and GENERIC_CLIENT_ID runs the Google branch, so it must report + Google. Presence is judged the way the callback judges it, so a client id set to the empty + string still selects that branch here. + """ + if os.getenv("GOOGLE_CLIENT_ID") is not None: + return ActiveSSOProvider.google + if os.getenv("MICROSOFT_CLIENT_ID") is not None: + return ActiveSSOProvider.microsoft + if os.getenv("GENERIC_CLIENT_ID") is not None: + return ActiveSSOProvider.generic + if SAMLAuthHandler.is_saml_configured(): + return ActiveSSOProvider.saml + return ActiveSSOProvider.none + + +def id_jag_assertion_capture_gap() -> str | None: + """Why ID-JAG cannot work under the active SSO provider, phrased for an operator reading a log, + or ``None`` when that provider does capture an assertion.""" + provider = active_sso_provider() + match provider: + case ActiveSSOProvider.generic: + return None + case ActiveSSOProvider.none: + return ( + "no SSO provider is configured, so no IdP identity assertion is ever captured and " + f"ID-JAG credential resolution fails for every user. {_GENERIC_OIDC_REMEDY}" + ) + case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.saml: + return ( + f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no " + "IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter " + f"how often they sign in. {_GENERIC_OIDC_REMEDY}" + ) + case _: + assert_never(provider) + + +def id_jag_assertion_capture_gap_at_startup() -> str | None: + """Config load runs before SSO settings stored in the database are reconciled into the process + environment, so an unresolved provider at that point is not yet a gap; the SSO callback reports it + once a login happens.""" + if active_sso_provider() is ActiveSSOProvider.none: + return None + return id_jag_assertion_capture_gap() diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 84150ef7935..3e6434a5afd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,7 +16,7 @@ import json import os import re import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from html import escape from types import MappingProxyType @@ -29,6 +29,7 @@ from typing import ( NoReturn, Optional, Protocol, + TypeAlias, Union, cast, overload, @@ -70,6 +71,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SSOIdentityAssertion, assertion_from_sso_login, + ema_assertion_retention_enabled, retain_sso_identity_assertion_for_ema, ) from litellm.proxy._types import ( @@ -105,6 +107,9 @@ from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + id_jag_assertion_capture_gap, +) from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler from litellm.proxy.management_endpoints.sso_helper_utils import ( check_is_admin_only_access, @@ -1677,6 +1682,46 @@ async def get_generic_sso_response( return result or {}, received_response, access_token_payload, sso_assertion +RetentionCheck: TypeAlias = Callable[[], Awaitable[bool]] # mutable-ok: Callable parameter syntax + + +async def warn_if_id_jag_assertion_uncaptured( + assertion: SSOIdentityAssertion | None, *, retention_enabled: RetentionCheck | None = None +) -> None: + """Say, at the one moment it is knowable, that this login gave an ``oauth2_id_jag`` server + nothing to spend. Without it the operator only ever sees the per-request failure, which cannot + tell a user who has never signed in from a provider that will never capture. Kept strictly + diagnostic: a store outage is swallowed, since a login must not fail over a log line.""" + if assertion is not None: + return + try: + check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled + if not await check(): + return + except Exception as exc: # noqa: BLE001 # diagnostics must never break the login + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers after SSO login: %s", exc) + return + gap: Final = id_jag_assertion_capture_gap() + verbose_proxy_logger.warning( + "SSO login captured no IdP identity assertion while an oauth2_id_jag MCP server is registered: %s", + gap if gap is not None else "the identity provider's token response carried no usable id_token", + ) + + +async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None = None) -> None: + gap: Final = id_jag_assertion_capture_gap() + if gap is None: + return + try: + check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled + if not await check(): + return + except Exception as exc: # noqa: BLE001 # diagnostics must never break the page they annotate + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers: %s", exc) + return + verbose_proxy_logger.warning("SSO debug callback ran with an oauth2_id_jag capture gap: %s", gap) + + async def create_team_member_add_task(team_id, user_info): """Create a task for adding a member to a team.""" try: @@ -2269,6 +2314,7 @@ async def _complete_cli_sso_callback_session( raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") await retain_sso_identity_assertion_for_ema(user_id=user_info.user_id, assertion=sso_assertion) + await warn_if_id_jag_assertion_uncaptured(sso_assertion) teams: list[str] = [] if hasattr(user_info, "teams") and user_info.teams: @@ -3599,6 +3645,7 @@ class SSOAuthenticationHandler: if isinstance(user_id, str) and user_id: await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=sso_assertion) + await warn_if_id_jag_assertion_uncaptured(sso_assertion) disabled_non_admin_personal_key_creation: Final = get_disabled_non_admin_personal_key_creation() litellm_dashboard_ui = get_custom_url(request_base_url=str(request.base_url), route="ui/") @@ -4733,6 +4780,7 @@ async def debug_sso_callback(request: Request): safe_raw_claims: Final = {k: v for k, v in (received_response or {}).items() if k not in _OAUTH_TOKEN_FIELDS} safe_access_token_claims = {k: v for k, v in (access_token_payload or {}).items() if k not in _OAUTH_TOKEN_FIELDS} + await warn_if_id_jag_capture_gap() sso_payload: Final = { "parsed_by_proxy": filtered_result, "raw_claims": safe_raw_claims, diff --git a/litellm/router.py b/litellm/router.py index cf65a06d604..172c1f8f31c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8381,17 +8381,12 @@ class Router: ## LOG FAILURE EVENT if logging_obj is not None: asyncio.create_task( - logging_obj.async_failure_handler( + logging_obj.dispatch_failure_handlers( exception=e, traceback_exception=traceback.format_exc(), - end_time=time.time(), + prefer_async_handlers=True, ) ) - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback.format_exc()), - ).start() # log response _set_cooldown_deployments( litellm_router_instance=self, exception_status=e.status_code, @@ -8404,17 +8399,12 @@ class Router: ## LOG FAILURE EVENT if logging_obj is not None: asyncio.create_task( - logging_obj.async_failure_handler( + logging_obj.dispatch_failure_handlers( exception=e, traceback_exception=traceback.format_exc(), - end_time=time.time(), + prefer_async_handlers=True, ) ) - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback.format_exc()), - ).start() # log response raise e async def async_callback_filter_deployments( @@ -8452,17 +8442,12 @@ class Router: ## LOG FAILURE EVENT if logging_obj is not None: asyncio.create_task( - logging_obj.async_failure_handler( + logging_obj.dispatch_failure_handlers( exception=e, traceback_exception=traceback.format_exc(), - end_time=time.time(), + prefer_async_handlers=True, ) ) - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback.format_exc()), - ).start() # log response raise e return returned_healthy_deployments @@ -12641,13 +12626,13 @@ class Router: logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) if logging_obj is not None: - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming - asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception=e, + traceback_exception=traceback_exception, + prefer_async_handlers=True, + ) + ) raise e async def async_get_available_deployment_for_pass_through( @@ -12775,11 +12760,13 @@ class Router: if request_kwargs is not None: logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) if logging_obj is not None: - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() - asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception=e, + traceback_exception=traceback_exception, + prefer_async_handlers=True, + ) + ) raise e async def _run_routing_plugins( @@ -13044,13 +13031,48 @@ class Router: ) return None - pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( + from litellm.proxy.guardrails.auto_router_compression import ( + messages_for_routing, + model_hop_compression_armed, + policy_for_model, + team_id_from_request, + ) + + # Same tag-aware lookup the proxy's pre-call arming used, so an alias with + # several tag-scoped markers cannot suppress under one and route under another. + compression_policy: Final = policy_for_model( + llm_router=self, + model_alias=registered_model_name, + team_id=team_id_from_request(request_kwargs), + request_tags=_get_tags_from_request_kwargs(request_kwargs), + ) + # Shared compression already ran in the pre-call hook, so reuse it rather than + # compressing twice. Conditional on arming having actually happened: only the + # proxy arms, and on the SDK path the shortcut would skip both hops entirely. + needs_independent_routing_compression: Final = compression_policy is not None and not ( + compression_policy.is_same and compression_policy.model is not None and model_hop_compression_armed() + ) + routing_messages: Final = ( + await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs) + if needs_independent_routing_compression + else None + ) + + routed: Final = await selected_strategy.strategy.async_pre_routing_hook( model=registered_model_name, request_kwargs=request_kwargs, - messages=messages, + messages=routing_messages if routing_messages is not None else messages, input=input, specific_deployment=specific_deployment, ) + # Routing-only compression must not leak into the response: the model call and + # deployment-context filtering key off this field. Compared by value, since + # pydantic rebuilds the list rather than keeping the object passed in. + pre_routing_hook_response: Final = ( + routed.model_copy(update={"messages": messages}) # mutable-ok: pydantic's model_copy takes a dict + if routed is not None and routing_messages is not None and routed.messages == routing_messages + else routed + ) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 0775e0a4039..d5234e27ec6 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -56,8 +56,9 @@ class PatternMatchRouter: This class will store a mapping for regex pattern: List[Deployments] """ - def __init__(self): + def __init__(self, pattern_utils: type[PatternUtils] = PatternUtils): self.patterns: dict[str, list] = {} + self._pattern_utils: Final = pattern_utils def add_pattern(self, pattern: str, llm_deployment: dict): """ @@ -69,9 +70,10 @@ class PatternMatchRouter: """ # Convert the pattern to a regex regex: Final = self._pattern_to_regex(pattern) - if regex not in self.patterns: - self.patterns[regex] = [] - self.patterns[regex].append(llm_deployment) + if regex in self.patterns: + self.patterns[regex].append(llm_deployment) + return + self.patterns = dict(self._pattern_utils.sorted_patterns({**self.patterns, regex: [llm_deployment]})) def remove_deployment(self, model_id: str) -> None: """ @@ -138,11 +140,12 @@ class PatternMatchRouter: if request is None: return None - sorted_patterns: Final = PatternUtils.sorted_patterns(self.patterns) regex_filtered_model_names: Final = ( - [self._pattern_to_regex(m) for m in filtered_model_names] if filtered_model_names is not None else [] + tuple(self._pattern_to_regex(m) for m in filtered_model_names) + if filtered_model_names is not None + else () ) - for pattern, llm_deployments in sorted_patterns: + for pattern, llm_deployments in self.patterns.items(): if filtered_model_names is not None and pattern not in regex_filtered_model_names: continue pattern_match = re.match(pattern, request) diff --git a/litellm/types/router.py b/litellm/types/router.py index 832399f6dbb..fb55a5f35b7 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -428,6 +428,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): auto_router_default_model: str | None = None auto_router_embedding_model: str | None = None auto_router_max_input_chars: int | None = None + # Compression policy for the two hops of a routed request. Both unset means the + # request's own compression guardrails apply to both, as they always have. + auto_router_routing_compression: str | None = None + auto_router_model_compression: str | None = None # complexity-router params complexity_router_config: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5f5a0646414..1142b8263f4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3787,6 +3787,8 @@ all_litellm_params = ( "auto_router_default_model", "auto_router_embedding_model", "auto_router_max_input_chars", + "auto_router_routing_compression", + "auto_router_model_compression", "complexity_router_config", "complexity_router_default_model", "adaptive_router_config", diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 88e9979148f..c93e5f6b303 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -392,6 +392,63 @@ with its own provider config (one `examples/default`-style root per project), or fork the module to add `configuration_aliases` and pass per-instance `providers = { ... }`. +## Dependencies only (run LiteLLM on GKE) + +Set `create_runtime = false` to provision Cloud SQL, Memorystore, GCS, +Secret Manager, and the runtime service account without Cloud Run or the +load balancer. For a Shared VPC, set the full host-project network ID and +skip PSA creation after the host project has configured it: + +```hcl +create_runtime = false +network_id = "projects//global/networks/" +create_psa_connection = false +``` + +The host project must already have Private Services Access configured on +that network and the Service Networking API enabled; the module cannot set +PSA up from a service project. GKE nodes must sit on the same Shared VPC so +the Cloud SQL and Memorystore private IPs are routable from the pods. Run +the root with its provider pointed at the project that should own the +dependencies. `create_runtime = true` with `network_id` set is also allowed, +but the Serverless VPC Access connector has to live in the same project as +the network, so that combination only works when the VPC is in the +deployment project + +Map the outputs into the Helm values as follows: + +```yaml +database: + writer: + host: + dbname: + passwordSecret: + name: + reader: + host: + dbname: + passwordSecret: + name: +redis: + host: + port: +masterKey: + secretName: +``` + +Create the database Secret with keys `username` (the `db_username` output) +and `password` (read it with `gcloud secrets versions access latest +--secret=`), and the master key Secret from +`master_key_secret_id` the same way. Memorystore only accepts TLS by +default, so store the `redis_server_ca_pem` output in a third Secret, +mount it into the gateway and backend pods via `volumes` / `volumeMounts`, +and add `REDIS_SSL=true` and `REDIS_SSL_CA_CERTS=` to each +component's `extraEnv`. Setting `redis_transit_encryption = false` removes +the CA plumbing at the cost of plaintext Redis traffic inside the VPC + +The chart's pre-install/pre-upgrade migration hook runs the Prisma +migration, so nothing replaces the Cloud Run migrations Job in this mode + ## Storage and database retention Two opt-in tripwires guard against accidental data loss on @@ -409,14 +466,15 @@ Flip `cloudsql_deletion_protection` to `false` or `gcs_force_destroy` to ## Redis encryption -Memorystore runs with `transit_encryption_mode = "SERVER_AUTHENTICATION"`, -so the proxy connects via `rediss://`. The instance's self-signed CA cert -(`server_ca_certs[0].cert`) is shipped to gateway + backend as -`REDIS_CA_PEM_B64`; their entrypoint shell decodes it to `/tmp/redis-ca.pem` -before uvicorn starts and points `REDIS_SSL_CA_CERTS` at that path. No -extra config needed — but if you ever swap Memorystore for an external -Redis, override `REDIS_HOST`/`REDIS_PORT` and either drop these env vars -or point them at your own CA. +By default, Memorystore runs with +`transit_encryption_mode = "SERVER_AUTHENTICATION"`, so Cloud Run connects +via `rediss://`. The instance's self-signed CA cert +(`server_ca_certs[0].cert`) is shipped to gateway and backend as +`REDIS_CA_PEM_B64`; their entrypoint shell decodes it to +`/tmp/redis-ca.pem` before uvicorn starts and points `REDIS_SSL_CA_CERTS` at +that path. Set `redis_transit_encryption = false` to use plaintext Redis. +For GKE, use `redis_server_ca_pem` as described in the dependencies-only +section, or accept the security tradeoff of disabling transit encryption ## Files @@ -434,4 +492,5 @@ or point them at your own CA. | `iam.tf` | Runtime SA + Cloud SQL client + Secret Manager accessor | | `cloudrun.tf` | 3 Cloud Run services + Cloud Run Job for migrations | | `load_balancer.tf`| External HTTPS LB, serverless NEGs, URL map for path routing | -| `outputs.tf` | LB IP, service URLs, secret IDs, migration `execute` command | +| `outputs.tf` | LB IP, service URLs, dependency endpoints, secret IDs, migration command | +| `tests/` | Plan-only mock-provider coverage for deployment modes and Redis encryption | diff --git a/terraform/litellm/gcp/bootstrap.tf b/terraform/litellm/gcp/bootstrap.tf index b929c4d76f3..dead5c41f6b 100644 --- a/terraform/litellm/gcp/bootstrap.tf +++ b/terraform/litellm/gcp/bootstrap.tf @@ -15,15 +15,17 @@ # enough to invoke Cloud Run admin APIs (`gcloud auth login`). resource "terraform_data" "migration" { + count = var.create_runtime ? 1 : 0 + triggers_replace = { - job_id = google_cloud_run_v2_job.migrations.id + job_id = google_cloud_run_v2_job.migrations[0].id job_image = local.migrations_image } provisioner "local-exec" { interpreter = ["bash", "-c"] environment = { - JOB = google_cloud_run_v2_job.migrations.name + JOB = google_cloud_run_v2_job.migrations[0].name REGION = var.region PROJECT = var.project_id } diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 5a5c361b832..84ae8b9247f 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -6,25 +6,28 @@ locals { # Memorystore exposes a self-signed CA cert per instance; we ship it as # a base64 env var and decode it to a file at container startup so the # rediss:// connection can validate. Public cert, not sensitive. - redis_ca_pem_b64 = base64encode(google_redis_instance.this.server_ca_certs[0].cert) + redis_ca_pem_b64 = var.redis_transit_encryption ? base64encode(google_redis_instance.this.server_ca_certs[0].cert) : "" - shared_env_kv = [ - { name = "DATABASE_HOST", value = google_sql_database_instance.writer.private_ip_address }, - { name = "DATABASE_PORT", value = "5432" }, - { name = "DATABASE_USER", value = var.db_username }, - { name = "DATABASE_NAME", value = var.db_name }, - { name = "DATABASE_HOST_READ_REPLICA", value = google_sql_database_instance.reader.private_ip_address }, - { name = "DATABASE_PORT_READ_REPLICA", value = "5432" }, - { name = "REDIS_HOST", value = google_redis_instance.this.host }, - { name = "REDIS_PORT", value = tostring(google_redis_instance.this.port) }, - # _redis.get_redis_url_from_environment honors REDIS_SSL to flip the - # scheme to rediss://; REDIS_SSL_CA_CERTS is mapped via - # _get_redis_env_kwarg_mapping → ssl_ca_certs on the redis-py client. - { name = "REDIS_SSL", value = "true" }, - { name = "REDIS_SSL_CA_CERTS", value = "/tmp/redis-ca.pem" }, - { name = "REDIS_CA_PEM_B64", value = local.redis_ca_pem_b64 }, - { name = "GCS_BUCKET_NAME", value = google_storage_bucket.this.name }, - ] + shared_env_kv = concat( + [ + { name = "DATABASE_HOST", value = google_sql_database_instance.writer.private_ip_address }, + { name = "DATABASE_PORT", value = "5432" }, + { name = "DATABASE_USER", value = var.db_username }, + { name = "DATABASE_NAME", value = var.db_name }, + { name = "DATABASE_HOST_READ_REPLICA", value = google_sql_database_instance.reader.private_ip_address }, + { name = "DATABASE_PORT_READ_REPLICA", value = "5432" }, + { name = "REDIS_HOST", value = google_redis_instance.this.host }, + { name = "REDIS_PORT", value = tostring(google_redis_instance.this.port) }, + ], + var.redis_transit_encryption ? [ + { name = "REDIS_SSL", value = "true" }, + { name = "REDIS_SSL_CA_CERTS", value = "/tmp/redis-ca.pem" }, + { name = "REDIS_CA_PEM_B64", value = local.redis_ca_pem_b64 }, + ] : [], + [ + { name = "GCS_BUCKET_NAME", value = google_storage_bucket.this.name }, + ], + ) # OTel v2 is opt-in and gated on otel_endpoint, matching the AWS stack — # nothing OTel-related is added to the container env until an endpoint is @@ -126,9 +129,9 @@ locals { # Decode the Memorystore CA cert (passed as REDIS_CA_PEM_B64) to the # path REDIS_SSL_CA_CERTS points at, so the redis-py client can validate # the rediss:// handshake. - redis_ca_fragment = [ + redis_ca_fragment = var.redis_transit_encryption ? [ "python -c \"import os, base64, pathlib; pathlib.Path(os.environ['REDIS_SSL_CA_CERTS']).write_bytes(base64.b64decode(os.environ['REDIS_CA_PEM_B64']))\"" - ] + ] : [] database_url_fragment = [ "export DATABASE_URL=\"postgresql://$${DATABASE_USER}:$${DATABASE_PASSWORD}@$${DATABASE_HOST}:$${DATABASE_PORT}/$${DATABASE_NAME}\"", @@ -171,29 +174,7 @@ locals { # ---------- Gateway ---------- resource "google_cloud_run_v2_service" "gateway" { - # Metering needs a client certificate AND its key. Each secret is created only - # when its own PEM is supplied, so an endpoint set with a missing key would - # otherwise apply cleanly and leave the proxy logging "missing config" and - # never exporting. ca_cert_pem stays optional: empty means fall back to the - # system trust store. - # - # The guard lives here, on an unconditional resource, rather than on the cert - # secret: that secret is count-gated on the cert itself, so it has zero - # instances in exactly the case this must catch. Adding count or for_each to - # this resource would silently stop the guard from evaluating. - # - # endpoint cert key -> result - # "" any any -> metering off, no secrets created - # set set set -> metering on - # set any-missing -> plan fails here - lifecycle { - precondition { - condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" - ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." - } - } + count = var.create_runtime ? 1 : 0 name = "${local.name}-gateway" location = var.region @@ -206,7 +187,7 @@ resource "google_cloud_run_v2_service" "gateway" { max_instance_request_concurrency = var.gateway_max_instance_request_concurrency vpc_access { - connector = google_vpc_access_connector.this.id + connector = google_vpc_access_connector.this[0].id egress = "PRIVATE_RANGES_ONLY" } @@ -312,17 +293,7 @@ resource "google_cloud_run_v2_service" "gateway" { # ---------- Backend ---------- resource "google_cloud_run_v2_service" "backend" { - # Same guard as the gateway: the backend meters too (it serves the named-server - # MCP transport), and a targeted apply of just this resource must not slip a - # billing endpoint through without the credentials to use it. - lifecycle { - precondition { - condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" - ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." - } - } + count = var.create_runtime ? 1 : 0 name = "${local.name}-backend" location = var.region @@ -335,7 +306,7 @@ resource "google_cloud_run_v2_service" "backend" { max_instance_request_concurrency = var.backend_max_instance_request_concurrency vpc_access { - connector = google_vpc_access_connector.this.id + connector = google_vpc_access_connector.this[0].id egress = "PRIVATE_RANGES_ONLY" } @@ -443,6 +414,8 @@ resource "google_cloud_run_v2_service" "backend" { # with zero IAM bindings, so a compromised UI container can't pivot to # Secret Manager / Cloud SQL via the metadata service. resource "google_cloud_run_v2_service" "ui" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-ui" location = var.region ingress = "INGRESS_TRAFFIC_INTERNAL_LOAD_BALANCER" @@ -450,7 +423,7 @@ resource "google_cloud_run_v2_service" "ui" { deletion_protection = false template { - service_account = google_service_account.ui_runtime.email + service_account = google_service_account.ui_runtime[0].email max_instance_request_concurrency = var.ui_max_instance_request_concurrency scaling { @@ -491,25 +464,31 @@ resource "google_cloud_run_v2_service" "ui" { # (LITELLM_MASTER_KEY); these IAM bindings just open up Cloud Run's invoker # gate so the LB request makes it to the container. resource "google_cloud_run_v2_service_iam_member" "gateway_allusers" { + count = var.create_runtime ? 1 : 0 + project = var.project_id - location = google_cloud_run_v2_service.gateway.location - name = google_cloud_run_v2_service.gateway.name + location = google_cloud_run_v2_service.gateway[0].location + name = google_cloud_run_v2_service.gateway[0].name role = "roles/run.invoker" member = "allUsers" } resource "google_cloud_run_v2_service_iam_member" "backend_allusers" { + count = var.create_runtime ? 1 : 0 + project = var.project_id - location = google_cloud_run_v2_service.backend.location - name = google_cloud_run_v2_service.backend.name + location = google_cloud_run_v2_service.backend[0].location + name = google_cloud_run_v2_service.backend[0].name role = "roles/run.invoker" member = "allUsers" } resource "google_cloud_run_v2_service_iam_member" "ui_allusers" { + count = var.create_runtime ? 1 : 0 + project = var.project_id - location = google_cloud_run_v2_service.ui.location - name = google_cloud_run_v2_service.ui.name + location = google_cloud_run_v2_service.ui[0].location + name = google_cloud_run_v2_service.ui[0].name role = "roles/run.invoker" member = "allUsers" } @@ -519,6 +498,8 @@ resource "google_cloud_run_v2_service_iam_member" "ui_allusers" { # assembles DATABASE_URL from the DATABASE_* env vars and runs `prisma # migrate deploy`. No proxy_config, no master key, no shell wrapper. resource "google_cloud_run_v2_job" "migrations" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-migrations" location = var.region labels = local.labels @@ -529,7 +510,7 @@ resource "google_cloud_run_v2_job" "migrations" { service_account = google_service_account.runtime.email vpc_access { - connector = google_vpc_access_connector.this.id + connector = google_vpc_access_connector.this[0].id egress = "PRIVATE_RANGES_ONLY" } diff --git a/terraform/litellm/gcp/cloudsql.tf b/terraform/litellm/gcp/cloudsql.tf index c9c2d03b2de..777434b4727 100644 --- a/terraform/litellm/gcp/cloudsql.tf +++ b/terraform/litellm/gcp/cloudsql.tf @@ -36,7 +36,7 @@ resource "google_sql_database_instance" "writer" { ip_configuration { ipv4_enabled = false - private_network = google_compute_network.this.id + private_network = local.network_id } insights_config { @@ -55,6 +55,11 @@ resource "google_sql_database_instance" "writer" { # (full data loss). Set the initial size only; let Cloud SQL own it # thereafter. ignore_changes = [settings[0].disk_size] + + precondition { + condition = var.create_psa_connection || var.network_id != "" + error_message = "create_psa_connection must be true unless network_id references an existing VPC with Private Services Access configured." + } } } @@ -76,7 +81,7 @@ resource "google_sql_database_instance" "reader" { ip_configuration { ipv4_enabled = false - private_network = google_compute_network.this.id + private_network = local.network_id } } diff --git a/terraform/litellm/gcp/examples/default/main.tf b/terraform/litellm/gcp/examples/default/main.tf index 8760d445f0c..f44b2a9a001 100644 --- a/terraform/litellm/gcp/examples/default/main.tf +++ b/terraform/litellm/gcp/examples/default/main.tf @@ -31,6 +31,11 @@ module "litellm" { tenant = var.tenant env = var.env + create_runtime = var.create_runtime + network_id = var.network_id + create_psa_connection = var.create_psa_connection + redis_transit_encryption = var.redis_transit_encryption + litellm_master_key = var.litellm_master_key litellm_license = var.litellm_license ui_password = var.ui_password diff --git a/terraform/litellm/gcp/examples/default/outputs.tf b/terraform/litellm/gcp/examples/default/outputs.tf index 3a9343c4850..48cdc1af66e 100644 --- a/terraform/litellm/gcp/examples/default/outputs.tf +++ b/terraform/litellm/gcp/examples/default/outputs.tf @@ -38,6 +38,31 @@ output "redis_endpoint" { value = module.litellm.redis_endpoint } +output "redis_host" { + description = "Memorystore Redis host." + value = module.litellm.redis_host +} + +output "redis_port" { + description = "Memorystore Redis port." + value = module.litellm.redis_port +} + +output "redis_server_ca_pem" { + description = "Memorystore server CA PEM." + value = module.litellm.redis_server_ca_pem +} + +output "db_username" { + description = "Cloud SQL application username." + value = module.litellm.db_username +} + +output "db_name" { + description = "Cloud SQL database name." + value = module.litellm.db_name +} + output "gcs_bucket" { description = "GCS bucket name." value = module.litellm.gcs_bucket @@ -53,6 +78,11 @@ output "db_password_secret_id" { value = module.litellm.db_password_secret_id } +output "runtime_service_account_email" { + description = "Runtime service account email." + value = module.litellm.runtime_service_account_email +} + output "migration_run_command" { description = "Break-glass command to re-run the one-off migration job." value = module.litellm.migration_run_command diff --git a/terraform/litellm/gcp/examples/default/terraform.tfvars.example b/terraform/litellm/gcp/examples/default/terraform.tfvars.example index 4416cf0ee5d..c35206503bb 100644 --- a/terraform/litellm/gcp/examples/default/terraform.tfvars.example +++ b/terraform/litellm/gcp/examples/default/terraform.tfvars.example @@ -8,6 +8,14 @@ region = "us-central1" tenant = "acme" env = "stage" +# Deployment mode. For dependencies only on a Shared VPC, set +# create_runtime = false, network_id to the full host-project network ID, and +# create_psa_connection = false after configuring PSA on that network. +# create_runtime = true +# network_id = "" +# create_psa_connection = true +# redis_transit_encryption = true + # Tenant-supplied secrets. Prefer TF_VAR_litellm_master_key / # TF_VAR_litellm_license / TF_VAR_ui_password env vars so the values don't # end up in a committed tfvars file. All three are optional — when diff --git a/terraform/litellm/gcp/examples/default/variables.tf b/terraform/litellm/gcp/examples/default/variables.tf index 56e5ec88ef8..88b57ce27eb 100644 --- a/terraform/litellm/gcp/examples/default/variables.tf +++ b/terraform/litellm/gcp/examples/default/variables.tf @@ -26,6 +26,30 @@ variable "env" { type = string } +variable "create_runtime" { + description = "Create Cloud Run and load balancer resources." + type = bool + default = true +} + +variable "network_id" { + description = "Existing VPC network resource ID. Empty creates a VPC." + type = string + default = "" +} + +variable "create_psa_connection" { + description = "Create Private Services Access resources." + type = bool + default = true +} + +variable "redis_transit_encryption" { + description = "Enable Memorystore transit encryption." + type = bool + default = true +} + # Sensitive — prefer TF_VAR_litellm_master_key / TF_VAR_litellm_license / # TF_VAR_ui_password so values stay out of any committed tfvars file. variable "litellm_master_key" { diff --git a/terraform/litellm/gcp/iam.tf b/terraform/litellm/gcp/iam.tf index 09df5e7dff0..509e6d48ffd 100644 --- a/terraform/litellm/gcp/iam.tf +++ b/terraform/litellm/gcp/iam.tf @@ -6,6 +6,15 @@ resource "google_service_account" "runtime" { account_id = "${local.name}-runtime" display_name = "LiteLLM Cloud Run runtime" + + lifecycle { + precondition { + condition = !var.create_runtime || var.billing_metrics_endpoint == "" || ( + var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + ) + error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set and create_runtime is true." + } + } } # UI runtime SA — no role bindings. The UI is static nginx with no DB, @@ -14,6 +23,8 @@ resource "google_service_account" "runtime" { # project's serverless service agent (not this SA), so it doesn't need # artifactregistry.reader either. resource "google_service_account" "ui_runtime" { + count = var.create_runtime ? 1 : 0 + account_id = "${local.name}-ui-runtime" display_name = "LiteLLM Cloud Run UI runtime (no data-plane access)" } diff --git a/terraform/litellm/gcp/load_balancer.tf b/terraform/litellm/gcp/load_balancer.tf index 11f30d0f944..57e8af8210f 100644 --- a/terraform/litellm/gcp/load_balancer.tf +++ b/terraform/litellm/gcp/load_balancer.tf @@ -14,77 +14,93 @@ locals { } resource "google_compute_global_address" "lb" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-lb-ip" labels = local.labels } # Serverless NEGs — one per Cloud Run service. resource "google_compute_region_network_endpoint_group" "gateway" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-gateway-neg" region = var.region network_endpoint_type = "SERVERLESS" cloud_run { - service = google_cloud_run_v2_service.gateway.name + service = google_cloud_run_v2_service.gateway[0].name } } resource "google_compute_region_network_endpoint_group" "backend" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-backend-neg" region = var.region network_endpoint_type = "SERVERLESS" cloud_run { - service = google_cloud_run_v2_service.backend.name + service = google_cloud_run_v2_service.backend[0].name } } resource "google_compute_region_network_endpoint_group" "ui" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-ui-neg" region = var.region network_endpoint_type = "SERVERLESS" cloud_run { - service = google_cloud_run_v2_service.ui.name + service = google_cloud_run_v2_service.ui[0].name } } # Backend services wrap each NEG. resource "google_compute_backend_service" "gateway" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-gateway-bs" protocol = "HTTP" load_balancing_scheme = "EXTERNAL_MANAGED" backend { - group = google_compute_region_network_endpoint_group.gateway.id + group = google_compute_region_network_endpoint_group.gateway[0].id } } resource "google_compute_backend_service" "backend" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-backend-bs" protocol = "HTTP" load_balancing_scheme = "EXTERNAL_MANAGED" backend { - group = google_compute_region_network_endpoint_group.backend.id + group = google_compute_region_network_endpoint_group.backend[0].id } } resource "google_compute_backend_service" "ui" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-ui-bs" protocol = "HTTP" load_balancing_scheme = "EXTERNAL_MANAGED" backend { - group = google_compute_region_network_endpoint_group.ui.id + group = google_compute_region_network_endpoint_group.ui[0].id } } # URL map. Default → backend (management API). Path matchers route the # gateway and UI prefixes elsewhere. resource "google_compute_url_map" "this" { + count = var.create_runtime ? 1 : 0 + name = local.name - default_service = google_compute_backend_service.backend.id + default_service = google_compute_backend_service.backend[0].id host_rule { hosts = ["*"] @@ -93,13 +109,13 @@ resource "google_compute_url_map" "this" { path_matcher { name = "main" - default_service = google_compute_backend_service.backend.id + default_service = google_compute_backend_service.backend[0].id # UI paths (catch them before any /v1/* gateway rules so /favicon.ico # and / take precedence). path_rule { paths = local.ui_path_prefixes - service = google_compute_backend_service.ui.id + service = google_compute_backend_service.ui[0].id } # Gateway path prefixes. GCP URL maps cap a path_rule at 10 path globs, @@ -108,7 +124,7 @@ resource "google_compute_url_map" "this" { for_each = { for idx, chunk in chunklist(local.gateway_path_prefixes, 10) : idx => chunk } content { paths = path_rule.value - service = google_compute_backend_service.gateway.id + service = google_compute_backend_service.gateway[0].id } } } @@ -118,7 +134,7 @@ resource "google_compute_url_map" "this" { # target proxy when TLS is enabled; otherwise the regular path-routing # URL map is attached to the HTTP proxy and everything stays plaintext. resource "google_compute_url_map" "https_redirect" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 name = "${local.name}-redirect" default_url_redirect { @@ -129,8 +145,10 @@ resource "google_compute_url_map" "https_redirect" { } resource "google_compute_target_http_proxy" "this" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-http" - url_map = local.tls_enabled ? google_compute_url_map.https_redirect[0].id : google_compute_url_map.this.id + url_map = local.tls_enabled ? google_compute_url_map.https_redirect[0].id : google_compute_url_map.this[0].id # Default-deny on the HTTP-only path: TLS is the supported posture. # Operators must either supply DNS names or explicitly opt in. @@ -143,12 +161,14 @@ resource "google_compute_target_http_proxy" "this" { } resource "google_compute_global_forwarding_rule" "http" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-http" ip_protocol = "TCP" port_range = "80" load_balancing_scheme = "EXTERNAL_MANAGED" - ip_address = google_compute_global_address.lb.address - target = google_compute_target_http_proxy.this.id + ip_address = google_compute_global_address.lb[0].address + target = google_compute_target_http_proxy.this[0].id labels = local.labels } @@ -161,7 +181,7 @@ resource "google_compute_global_forwarding_rule" "http" { # transitions to ACTIVE. resource "google_compute_managed_ssl_certificate" "this" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 # A managed cert's `domains` is immutable, so changing var.lb_domains # forces replacement, and the cert is referenced by the HTTPS target @@ -181,19 +201,19 @@ resource "google_compute_managed_ssl_certificate" "this" { } resource "google_compute_target_https_proxy" "this" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 name = "${local.name}-https" - url_map = google_compute_url_map.this.id + url_map = google_compute_url_map.this[0].id ssl_certificates = [google_compute_managed_ssl_certificate.this[0].id] } resource "google_compute_global_forwarding_rule" "https" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 name = "${local.name}-https" ip_protocol = "TCP" port_range = "443" load_balancing_scheme = "EXTERNAL_MANAGED" - ip_address = google_compute_global_address.lb.address + ip_address = google_compute_global_address.lb[0].address target = google_compute_target_https_proxy.this[0].id labels = local.labels } diff --git a/terraform/litellm/gcp/locals.tf b/terraform/litellm/gcp/locals.tf index 9a817eba605..3861413d496 100644 --- a/terraform/litellm/gcp/locals.tf +++ b/terraform/litellm/gcp/locals.tf @@ -21,6 +21,9 @@ locals { var.labels, ) + create_network = var.network_id == "" + network_id = local.create_network ? google_compute_network.this[0].id : var.network_id + gateway_path_prefixes = [ "/v1/chat/*", "/chat/*", "/v1/completions*", "/completions*", @@ -74,7 +77,7 @@ locals { "/ui/*", ] - proxy_config_enabled = length(keys(var.proxy_config)) > 0 + proxy_config_enabled = var.create_runtime && length(keys(var.proxy_config)) > 0 proxy_config_yaml = local.proxy_config_enabled ? yamlencode(var.proxy_config) : "" proxy_config_mount_path = "/etc/litellm" diff --git a/terraform/litellm/gcp/network.tf b/terraform/litellm/gcp/network.tf index a1ccaed02f9..47c7bf94a2b 100644 --- a/terraform/litellm/gcp/network.tf +++ b/terraform/litellm/gcp/network.tf @@ -1,13 +1,17 @@ resource "google_compute_network" "this" { + count = local.create_network ? 1 : 0 + name = local.name auto_create_subnetworks = false routing_mode = "REGIONAL" } resource "google_compute_subnetwork" "this" { + count = local.create_network ? 1 : 0 + name = "${local.name}-${var.region}" region = var.region - network = google_compute_network.this.id + network = google_compute_network.this[0].id ip_cidr_range = var.subnet_cidr private_ip_google_access = true } @@ -16,17 +20,21 @@ resource "google_compute_subnetwork" "this" { # managed services peer with the VPC over the connection below using # addresses from this range. resource "google_compute_global_address" "psa" { + count = var.create_psa_connection ? 1 : 0 + name = "${local.name}-psa" purpose = "VPC_PEERING" address_type = "INTERNAL" prefix_length = 16 - network = google_compute_network.this.id + network = local.network_id } resource "google_service_networking_connection" "psa" { - network = google_compute_network.this.id + count = var.create_psa_connection ? 1 : 0 + + network = local.network_id service = "servicenetworking.googleapis.com" - reserved_peering_ranges = [google_compute_global_address.psa.name] + reserved_peering_ranges = [google_compute_global_address.psa[0].name] } # Serverless VPC Access connector — required so Cloud Run can reach @@ -37,9 +45,11 @@ resource "google_service_networking_connection" "psa" { # for low-to-moderate Cloud Run egress; bump max if your services push # heavy private-network traffic. resource "google_vpc_access_connector" "this" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-conn" region = var.region - network = google_compute_network.this.name + network = local.network_id ip_cidr_range = var.vpc_connector_cidr min_instances = 2 max_instances = 3 diff --git a/terraform/litellm/gcp/outputs.tf b/terraform/litellm/gcp/outputs.tf index 6f1f1d5ccf4..2a4742f42cf 100644 --- a/terraform/litellm/gcp/outputs.tf +++ b/terraform/litellm/gcp/outputs.tf @@ -1,26 +1,26 @@ output "lb_ip" { - description = "Global anycast IP of the external HTTPS load balancer." - value = google_compute_global_address.lb.address + description = "Global anycast IP of the external HTTPS load balancer. Null when create_runtime is false." + value = var.create_runtime ? one(google_compute_global_address.lb[*].address) : null } output "lb_url" { - description = "Proxy URL. Switches scheme based on whether lb_domains is set; when TLS is enabled the URL points at the first listed domain (since managed certs are tied to the hostname, not the anycast IP). The dashboard is served at /, the API at /v1/*." - value = local.tls_enabled ? "https://${var.lb_domains[0]}" : "http://${google_compute_global_address.lb.address}" + description = "Proxy URL, or null when create_runtime is false. Switches scheme based on whether lb_domains is set." + value = var.create_runtime ? (local.tls_enabled ? "https://${var.lb_domains[0]}" : "http://${one(google_compute_global_address.lb[*].address)}") : null } output "gateway_service_url" { - description = "Default Cloud Run URL for the gateway (bypasses the LB)." - value = google_cloud_run_v2_service.gateway.uri + description = "Default Cloud Run URL for the gateway, or null when create_runtime is false." + value = var.create_runtime ? one(google_cloud_run_v2_service.gateway[*].uri) : null } output "backend_service_url" { - description = "Default Cloud Run URL for the backend (bypasses the LB)." - value = google_cloud_run_v2_service.backend.uri + description = "Default Cloud Run URL for the backend, or null when create_runtime is false." + value = var.create_runtime ? one(google_cloud_run_v2_service.backend[*].uri) : null } output "ui_service_url" { - description = "Default Cloud Run URL for the UI (bypasses the LB)." - value = google_cloud_run_v2_service.ui.uri + description = "Default Cloud Run URL for the UI, or null when create_runtime is false." + value = var.create_runtime ? one(google_cloud_run_v2_service.ui[*].uri) : null } output "cloudsql_writer_ip" { @@ -38,6 +38,36 @@ output "redis_endpoint" { value = "${google_redis_instance.this.host}:${google_redis_instance.this.port}" } +output "runtime_service_account_email" { + description = "Runtime service account email for Cloud Run or GKE Workload Identity." + value = google_service_account.runtime.email +} + +output "redis_host" { + description = "Memorystore Redis host." + value = google_redis_instance.this.host +} + +output "redis_port" { + description = "Memorystore Redis port." + value = google_redis_instance.this.port +} + +output "redis_server_ca_pem" { + description = "Memorystore server CA PEM. Mount it in the pod and set REDIS_SSL=true and REDIS_SSL_CA_CERTS= via extraEnv when transit encryption is enabled." + value = var.redis_transit_encryption ? google_redis_instance.this.server_ca_certs[0].cert : null +} + +output "db_username" { + description = "Cloud SQL application username." + value = var.db_username +} + +output "db_name" { + description = "Cloud SQL database name." + value = var.db_name +} + output "gcs_bucket" { description = "GCS bucket name. Exposed to gateway + backend as GCS_BUCKET_NAME. Reference from proxy_config via `os.environ/GCS_BUCKET_NAME`." value = google_storage_bucket.this.name @@ -54,11 +84,11 @@ output "db_password_secret_id" { } output "migration_run_command" { - description = "Shell command that executes the one-off migration job against Cloud SQL. Run this once after the first apply." - value = format( + description = "Shell command that executes the one-off migration job against Cloud SQL, or null when create_runtime is false." + value = var.create_runtime ? format( "gcloud run jobs execute %s --region %s --project %s --wait", - google_cloud_run_v2_job.migrations.name, + one(google_cloud_run_v2_job.migrations[*].name), var.region, var.project_id, - ) + ) : null } diff --git a/terraform/litellm/gcp/redis.tf b/terraform/litellm/gcp/redis.tf index 0e07c416e85..0602758f090 100644 --- a/terraform/litellm/gcp/redis.tf +++ b/terraform/litellm/gcp/redis.tf @@ -4,7 +4,7 @@ resource "google_redis_instance" "this" { memory_size_gb = var.redis_memory_size_gb region = var.region - authorized_network = google_compute_network.this.id + authorized_network = local.network_id connect_mode = "PRIVATE_SERVICE_ACCESS" redis_version = "REDIS_7_0" @@ -16,7 +16,7 @@ resource "google_redis_instance" "this" { # and passed to the proxy as REDIS_CA_PEM_B64); the proxy decodes it to # /tmp/redis-ca.pem at startup and uses it to validate the rediss:// # handshake. Mirrors `transit_encryption_enabled = true` on AWS. - transit_encryption_mode = "SERVER_AUTHENTICATION" + transit_encryption_mode = var.redis_transit_encryption ? "SERVER_AUTHENTICATION" : "DISABLED" depends_on = [google_service_networking_connection.psa] } diff --git a/terraform/litellm/gcp/tests/deps_only.tftest.hcl b/terraform/litellm/gcp/tests/deps_only.tftest.hcl new file mode 100644 index 00000000000..610c49d5b52 --- /dev/null +++ b/terraform/litellm/gcp/tests/deps_only.tftest.hcl @@ -0,0 +1,175 @@ +mock_provider "google" { + mock_resource "google_redis_instance" { + defaults = { + host = "10.0.0.4" + port = 6379 + server_ca_certs = [{ + cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----" + }] + } + } +} + +mock_provider "google-beta" {} +mock_provider "random" {} + +variables { + project_id = "test-project" + tenant = "tenant" + env = "test" + allow_plaintext_lb = true + image_registry = "us-central1-docker.pkg.dev/test-project/litellm" +} + +run "default_creates_everything" { + command = plan + + assert { + condition = alltrue([ + length(google_compute_network.this) == 1, + length(google_compute_subnetwork.this) == 1, + length(google_compute_global_address.psa) == 1, + length(google_service_networking_connection.psa) == 1, + length(google_vpc_access_connector.this) == 1, + length(google_cloud_run_v2_service.gateway) == 1, + length(google_cloud_run_v2_service.backend) == 1, + length(google_cloud_run_v2_service.ui) == 1, + length(google_cloud_run_v2_job.migrations) == 1, + length(google_compute_global_address.lb) == 1, + length(terraform_data.migration) == 1, + ]) + error_message = "The default mode must create networking, runtime services, the load balancer, and migrations." + } + + assert { + condition = google_redis_instance.this.transit_encryption_mode == "SERVER_AUTHENTICATION" + error_message = "Redis transit encryption must remain enabled by default." + } + + assert { + condition = length(local.shared_env_kv) == 12 + error_message = "The default runtime environment must include GCS and the three Redis TLS entries." + } +} + +run "deps_only_creates_no_runtime" { + command = plan + + variables { + create_runtime = false + proxy_config = { + model_list = [] + } + } + + assert { + condition = alltrue([ + length(google_cloud_run_v2_service.gateway) == 0, + length(google_cloud_run_v2_service.backend) == 0, + length(google_cloud_run_v2_service.ui) == 0, + length(google_cloud_run_v2_job.migrations) == 0, + length(google_cloud_run_v2_service_iam_member.gateway_allusers) == 0, + length(google_cloud_run_v2_service_iam_member.backend_allusers) == 0, + length(google_cloud_run_v2_service_iam_member.ui_allusers) == 0, + length(google_compute_global_address.lb) == 0, + length(google_compute_region_network_endpoint_group.gateway) == 0, + length(google_compute_region_network_endpoint_group.backend) == 0, + length(google_compute_region_network_endpoint_group.ui) == 0, + length(google_compute_backend_service.gateway) == 0, + length(google_compute_backend_service.backend) == 0, + length(google_compute_backend_service.ui) == 0, + length(google_compute_url_map.this) == 0, + length(google_compute_url_map.https_redirect) == 0, + length(google_compute_target_http_proxy.this) == 0, + length(google_compute_global_forwarding_rule.http) == 0, + length(google_compute_managed_ssl_certificate.this) == 0, + length(google_compute_target_https_proxy.this) == 0, + length(google_compute_global_forwarding_rule.https) == 0, + length(terraform_data.migration) == 0, + length(google_vpc_access_connector.this) == 0, + length(google_service_account.ui_runtime) == 0, + length(google_storage_bucket.proxy_config) == 0, + ]) + error_message = "Dependencies-only mode must omit all runtime, load balancer, connector, UI identity, and proxy config resources." + } + + assert { + condition = alltrue([ + google_sql_database_instance.writer.name == "tenant-litellm-test", + google_sql_database_instance.reader.name == "tenant-litellm-test-reader", + google_redis_instance.this.name == "tenant-litellm-test", + google_storage_bucket.this.force_destroy == false, + google_secret_manager_secret.master_key.secret_id == "tenant-litellm-test-master-key", + google_secret_manager_secret.db_password.secret_id == "tenant-litellm-test-db-password", + google_service_account.runtime.account_id == "tenant-litellm-test-runtime", + ]) + error_message = "Dependencies-only mode must retain data stores, secrets, and the runtime service account." + } + + assert { + condition = output.lb_url == null && output.migration_run_command == null + error_message = "Runtime outputs must be null while dependency outputs remain available." + } +} + +run "existing_network_attaches_data_stores" { + command = plan + + variables { + network_id = "projects/host-proj/global/networks/shared" + create_psa_connection = false + create_runtime = false + } + + assert { + condition = alltrue([ + length(google_compute_network.this) == 0, + length(google_compute_subnetwork.this) == 0, + length(google_compute_global_address.psa) == 0, + length(google_service_networking_connection.psa) == 0, + google_sql_database_instance.writer.settings[0].ip_configuration[0].private_network == var.network_id, + google_redis_instance.this.authorized_network == var.network_id, + ]) + error_message = "An existing VPC must receive the Cloud SQL and Memorystore private-network attachments." + } +} + +run "psa_required_without_existing_network" { + command = plan + + variables { + create_psa_connection = false + } + + expect_failures = [ + google_sql_database_instance.writer, + ] +} + +run "redis_plaintext_drops_tls_env" { + command = plan + + variables { + redis_transit_encryption = false + } + + assert { + condition = google_redis_instance.this.transit_encryption_mode == "DISABLED" + error_message = "Redis transit encryption must be disabled when requested." + } + + assert { + condition = length(local.shared_env_kv) == 9 + error_message = "Plaintext Redis mode must include GCS and omit the three Redis TLS entries." + } + + assert { + condition = length([for env in local.shared_env_kv : env if env.name == "REDIS_SSL"]) == 0 + error_message = "Plaintext Redis mode must not set REDIS_SSL." + } + + assert { + condition = length(local.redis_ca_fragment) == 0 + error_message = "Plaintext Redis mode must not decode a Redis CA at startup." + } +} diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 1162e100bb2..9c68ed3db76 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -79,16 +79,42 @@ variable "ui_password" { sensitive = true } +# ---------- Deployment mode ---------- + +variable "create_runtime" { + description = "Create Cloud Run, load balancer, VPC connector, runtime support resources, and the migration job. Set false for GKE or another external runtime." + type = bool + default = true +} + +variable "network_id" { + description = "Existing VPC network resource ID (`projects//global/networks/`). When set, no VPC or subnet is created. A VPC connector requires this network to be in the deployment project when create_runtime is true." + type = string + default = "" +} + +variable "create_psa_connection" { + description = "Create the Private Services Access range and connection for Cloud SQL and Memorystore. Set false when the existing network already has PSA configured." + type = bool + default = true +} + +variable "redis_transit_encryption" { + description = "Enable Memorystore transit encryption and inject Redis TLS settings into Cloud Run. Set false to use plaintext Redis." + type = bool + default = true +} + # ---------- Networking ---------- variable "subnet_cidr" { - description = "Primary CIDR block for the LiteLLM subnet." + description = "Primary CIDR block for the LiteLLM subnet. Unused when network_id is set." type = string default = "10.40.0.0/16" } variable "vpc_connector_cidr" { - description = "CIDR for the Serverless VPC Access connector. /28 required." + description = "CIDR for the Serverless VPC Access connector. /28 required. Unused when create_runtime is false." type = string default = "10.41.0.0/28" } diff --git a/tests/e2e/access_control/test_access_control_e2e.py b/tests/e2e/access_control/test_access_control_e2e.py index af7e9a099fd..9d01f2915e7 100644 --- a/tests/e2e/access_control/test_access_control_e2e.py +++ b/tests/e2e/access_control/test_access_control_e2e.py @@ -24,7 +24,7 @@ from access_control_client import ( from e2e_config import unique_marker from e2e_http import Success, UnauthorizedError, UnknownApiError, unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, LiteLLMParamsBody from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -32,6 +32,7 @@ pytestmark = pytest.mark.e2e ALLOWED_MODEL = "gemini-2.5-flash" DISALLOWED_MODEL = "gpt-5.5" VIRTUAL_KEY_BACKEND = "anthropic/claude-haiku-4-5-20251001" +EMBEDDING_MODEL = "openai-text-embedding-3-small" class TestAccessControl: @@ -71,6 +72,31 @@ class TestAccessControl: f"403 body must be a model-access denial, got: {result.body[:300]}" ) + @pytest.mark.covers("other.auth.virtual_key.route_group_allowed") + def test_llm_api_routes_group_grants_every_llm_endpoint( + self, client: AccessControlClient, resources: ResourceManager + ) -> None: + key = client.llm_only_key() + resources.defer(lambda: client.delete_key(key)) + + chat = client.chat_status(key, ALLOWED_MODEL, f"capital of France? {unique_marker()}") + assert chat.status_code == 200, ( + f"llm_api_routes key must reach /chat/completions, got {chat.status_code}: {chat.body[:300]}" + ) + assert ChatResponse.model_validate_json(chat.body).choices, ( + f"200 must carry a real completion, not an error envelope: {chat.body[:300]}" + ) + + embedding = unwrap( + client.proxy.embed(key, EmbedBody(model=EMBEDDING_MODEL, input=f"route group {unique_marker()}")) + ) + assert embedding.model, f"llm_api_routes key reached /embeddings but got no model back: {embedding}" + + denied = client.create_model_status(key, f"e2e-route-group-{unique_marker()}") + assert denied.status_code == 403 and ROUTE_NOT_ALLOWED_MARKER in denied.body, ( + f"the same key must still be shut out of /model/new, got {denied.status_code}: {denied.body[:300]}" + ) + def test_llm_only_key_forbidden_from_management_route_403( self, client: AccessControlClient, resources: ResourceManager ) -> None: diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index a2c17b0fb66..0d224061381 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -10,6 +10,11 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix intentionally not covered here. - Vertex (gemini-2.5-flash): prompt caching via ``cache_control`` context caching; the second identical call must report cached prompt tokens > 0. +- Anthropic (claude-haiku-4-5, direct): the same ``cache_control`` prefix over + the OpenAI-compatible route; the second call must report cache-read tokens > 0. +- OpenAI (gpt-5.6): automatic prompt caching needs no request marker, so the + cacheable prefix goes out as a plain system string with a ``prompt_cache_key`` + and the second call must report ``prompt_tokens_details.cached_tokens`` > 0. service_tier lives in test_provider_features_e2e.py. @@ -21,6 +26,7 @@ built from the typed content blocks shared in ``endpoints_client.py``. from __future__ import annotations import time +from collections.abc import Callable import pytest from pydantic import BaseModel @@ -29,7 +35,7 @@ from e2e_config import unique_marker from e2e_http import Result, unwrap from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager -from models import ChatResponse, LiteLLMParamsBody, Usage +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage from passthrough_client import PassthroughClient import os @@ -37,6 +43,8 @@ pytestmark = pytest.mark.e2e BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" VERTEX_MODEL = "vertex_ai/gemini-2.5-flash" +ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001" +OPENAI_MODEL = "openai/gpt-5.6" class CacheChatBody(BaseModel): @@ -89,17 +97,36 @@ def _cache_chat( ) +def _plain_cache_chat( + client: PassthroughClient, key: str, model: str, prefix: str, cache_key: str +) -> Result[ChatResponse]: + """The same cacheable prefix as a plain system string, for providers that cache + automatically and take no per-block marker (OpenAI).""" + return client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="system", content=prefix), + ChatMessage(role="user", content="Reply with one word."), + ], + max_tokens=64, + prompt_cache_key=cache_key, + ), + ) + + def _assert_cache_read_on_second_call( - client: PassthroughClient, key: str, model: str + model: str, send: Callable[[str], Result[ChatResponse]] ) -> None: prefix = _cacheable_prefix() - first = unwrap(_cache_chat(client, key, model, prefix)) + first = unwrap(send(prefix)) assert first.choices, f"{model}: first cache-priming call returned no choices: {first}" deadline = time.monotonic() + 30.0 while True: - second = unwrap(_cache_chat(client, key, model, prefix)) + second = unwrap(send(prefix)) read_tokens = _cached_read_tokens(second.usage) if read_tokens > 0 or time.monotonic() >= deadline: break @@ -125,7 +152,8 @@ class TestCacheControl: LiteLLMParamsBody(model=BEDROCK_MODEL, aws_region_name="us-east-1"), ) resources.defer(lambda: client.proxy.delete_model(model_id)) - _assert_cache_read_on_second_call(client, resources.key(), model) + key = resources.key() + _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix)) @pytest.mark.covers( "llm.chat_completions.vertex.prompt_cache_5m.nonstream.works", @@ -145,4 +173,40 @@ class TestCacheControl: ), ) resources.defer(lambda: client.proxy.delete_model(model_id)) - _assert_cache_read_on_second_call(client, resources.key(), model) + key = resources.key() + _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix)) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works", + exercised_on=[], + ) + def test_anthropic_prompt_caching_reads_cache( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-anthropic-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=ANTHROPIC_MODEL, api_key="os.environ/ANTHROPIC_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix)) + + @pytest.mark.covers( + "llm.chat_completions.openai.prompt_cache_5m.nonstream.works", + exercised_on=[], + ) + def test_openai_prompt_caching_reads_cache( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=OPENAI_MODEL, api_key="os.environ/OPENAI_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + cache_key = f"e2e-openai-cache-{unique_marker()}" + _assert_cache_read_on_second_call( + model, lambda prefix: _plain_cache_chat(client, key, model, prefix, cache_key) + ) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 68c0dfab897..87bd32d8dab 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -11,7 +11,7 @@ fails that provider's row here. The per-provider classes below cover the OpenAI-compatible /chat/completions translation for providers customers reach by registering their own deployment -via /model/new (Cohere, Gemini, hosted_vllm), each deleted on teardown. +via /model/new (Cohere, Gemini, hosted_vllm, Anthropic), each deleted on teardown. """ from __future__ import annotations @@ -46,6 +46,7 @@ pytestmark = pytest.mark.e2e COHERE_BACKEND = "cohere/command-r-08-2024" GEMINI_BACKEND = "gemini/gemini-2.5-flash" OPENAI_BACKEND = "openai/gpt-5.6" +ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5-20251001" BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" @@ -746,3 +747,98 @@ class TestBedrockConverseChatCompletions: response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) _assert_describes_cat(response) + + +class TestAnthropicChatCompletions: + """Anthropic via the OpenAI-compatible /chat/completions path, the translation + customers on the OpenAI SDK rely on when they route to Claude. The streamed call + must deliver real content deltas, and a tool-forced call must come back as a + well-formed tool_call on both the non-streamed and streamed paths. + """ + + def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.anthropic.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-tool") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.tool_use.stream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_streams_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-tool-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + stream=True, + ), + ) + assert result.ok and result.is_streaming, f"tool stream was not established: {result}" + assert result.stream_error is None, f"tool stream carried an error event: {result.stream_error}" + name, arguments = _streamed_tool_call(result.stream_events) + assert name == "get_weather", f"streamed tool call named {name!r}: {result.stream_events[:5]}" + args = _WeatherArgs.model_validate_json(arguments) + assert args.location.strip(), f"streamed tool call arguments missing location: {arguments!r}" diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index f951eb328f5..5520ca0cee5 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,4 +1,4 @@ -"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex. +"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex, Cohere. Each test registers the deployment it needs at runtime (deleted on teardown) and asserts a non-empty, non-zero vector came back. The LIT-3167 guard in @@ -86,6 +86,26 @@ class TestEmbeddingsEndpoint: f"embedding vector is all zeros: {result.body[:300]}" ) + @pytest.mark.covers("llm.embeddings.cohere.basic.nonstream.works") + def test_cohere_embeddings_returns_vector( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-embeddings-cohere-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody(model="cohere/embed-v4.0", api_key="os.environ/COHERE_API_KEY"), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.embeddings(key, model, "Say this is a test!") + require_successful_call(result) + parsed = EmbeddingsResult.model_validate_json(result.body) + assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" + assert any(component != 0.0 for component in parsed.first_vector), ( + f"embedding vector is all zeros: {result.body[:300]}" + ) + @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") def test_vertex_embeddings_returns_vector( self, endpoints_client: EndpointsClient, resources: ResourceManager diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index 50ea8f4b4df..447fe7d30d9 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -17,7 +17,7 @@ import pytest from e2e_config import CHEAP_OPENAI_MODEL, unique_marker from e2e_http import require_successful_call, unwrap from lifecycle import ResourceManager -from models import KeyGenerateBody, SpendLogRow +from models import ChatResponse, KeyGenerateBody, SpendLogRow from passthrough_client import ( AnthropicTool, GeminiFunctionDeclaration, @@ -344,6 +344,44 @@ class TestOpenAIPassthroughSpend: ) +class TestOpenAIProviderPrefixChat: + """OpenAI-format chat through the raw `/openai/{endpoint}` passthrough (LIT-4752). + + The body goes to OpenAI untranslated with the proxy's own OPENAI_API_KEY swapped + in, so the customer gets OpenAI's real completion back, and the gateway must + still write a costed pass_through_endpoint row for it. + """ + + @pytest.mark.covers("llm.chat_completions.openai.passthrough.nonstream.cost_logged") + def test_openai_prefix_chat_returns_completion_and_logs_its_cost( + self, client: PassthroughClient, scoped_key: str + ) -> None: + result = client.openai_chat(scoped_key, CHEAP_OPENAI_MODEL, f"Say hi in one word. {unique_marker()}") + require_successful_call(result) + + completion = ChatResponse.model_validate_json(result.body) + assert completion.id, f"/openai/v1/chat/completions relayed no completion id: {result.body[:300]}" + content = ( + completion.choices[0].message.content + if completion.choices and completion.choices[0].message + else None + ) + assert content and content.strip(), ( + f"/openai/v1/chat/completions relayed an empty completion: {result.body[:300]}" + ) + assert completion.usage is not None, f"the completion carried no usage to price from: {completion}" + + row = _fetch_cost_breakdown(client, completion.id) + assert row.prompt_tokens == completion.usage.prompt_tokens, ( + f"logged {row.prompt_tokens} prompt tokens, the completion the customer read " + f"reported {completion.usage.prompt_tokens}" + ) + assert row.completion_tokens == completion.usage.completion_tokens, ( + f"logged {row.completion_tokens} completion tokens, the completion the customer read " + f"reported {completion.usage.completion_tokens}" + ) + + class TestOpenAIPassthroughWebsocket: """The OpenAI passthrough prefixes must answer a websocket upgrade, not only a POST. diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 387280c8023..2b897f5f07f 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -40,6 +40,8 @@ from models import ( KeyListParams, KeyListResponse, KeyRegenerateBody, + KeyResetSpendBody, + KeyResetSpendResponse, KeyUpdateBody, ModelDeleteBody, OrgDeleteBody, @@ -191,16 +193,26 @@ class ManagementClient: response_type=NoBody, ) ) - def regenerate_key(self, key: str) -> str: + def regenerate_key(self, key: str, *, grace_period: str | None = None) -> str: return unwrap( self.proxy.transport.post( "/key/regenerate", headers=self.proxy.transport.master, - json=KeyRegenerateBody(key=key), + json=KeyRegenerateBody(key=key, grace_period=grace_period), response_type=KeyGenerateResponse, ) ).key + def reset_key_spend(self, key: str, reset_to: float) -> KeyResetSpendResponse: + return unwrap( + self.proxy.transport.post( + f"/key/{key}/reset_spend", + headers=self.proxy.transport.master, + json=KeyResetSpendBody(reset_to=reset_to), + response_type=KeyResetSpendResponse, + ) + ) + def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]: """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is who is asking: the master key by default, or a virtual key.""" diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index 711175abb0d..8b7d5f0eb6f 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -18,7 +18,7 @@ from typing import Literal import pytest from e2e_config import unique_marker -from e2e_http import NoBody, unwrap +from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient from models import KeyDeleteBody, KeyGenerateBody, KeyUpdateBody @@ -26,6 +26,9 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e +TINY_BUDGET = 3e-6 +SPEND_MODEL = "claude-haiku-4-5" + class KeyToggleBlockBody(BaseModel): key: str @@ -82,6 +85,30 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke return key +def _is_budget_block(outcome: StreamingResponse) -> bool: + return not outcome.ok and "budget_exceeded" in outcome.body + + +def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None: + for _ in range(40): + outcome = client.chat_status(key, SPEND_MODEL, f"spend {unique_marker()}") + if _is_budget_block(outcome): + assert outcome.status_code == 429, ( + f"budget refusal must be 429, got {outcome.status_code}: {outcome.body[:200]}" + ) + return + assert outcome.ok, f"paid call failed before the budget tripped ({outcome.status_code}): {outcome.body[:300]}" + time.sleep(2) + pytest.fail(f"max_budget={TINY_BUDGET} never blocked a call on the key") + + +def _settled_spend(client: ManagementClient, key: str) -> float | None: + first = client.proxy.key_info(key).spend or 0.0 + time.sleep(client.proxy.poll_interval) + second = client.proxy.key_info(key).spend or 0.0 + return second if first > 0 and first == second else None + + def _block(client: ManagementClient, key: str) -> None: _ = unwrap( client.proxy.transport.post( @@ -197,6 +224,32 @@ class TestKeyManagementRoutes: "/key/info never reported max_budget 42.0 after /key/bulk_update before the deadline", ) + @pytest.mark.covers("other.key_mgmt.spend_reset.resets_to_value") + def test_reset_spend_zeroes_recorded_spend_and_lifts_the_budget_block( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + key = _generate_key(client, resources, KeyGenerateBody(models=[SPEND_MODEL], max_budget=TINY_BUDGET)) + _spend_until_budget_blocks(client, key) + recorded = _poll( + client, lambda: _settled_spend(client, key), "key spend never landed in /key/info before the deadline" + ) + + reset = client.reset_key_spend(key, reset_to=0.0) + assert reset.previous_spend == recorded, ( + f"reset_spend reported previous_spend {reset.previous_spend}, /key/info had recorded {recorded}" + ) + assert reset.spend == 0.0, f"reset_spend to 0 reported spend {reset.spend}" + assert client.proxy.key_info(key).spend == 0.0, "/key/info still reports spend after the reset to 0" + + def call_allowed_again() -> bool | None: + outcome = client.chat_status(key, SPEND_MODEL, f"after reset {unique_marker()}") + if _is_budget_block(outcome): + return None + assert outcome.ok, f"post-reset call failed ({outcome.status_code}): {outcome.body[:300]}" + return True + + _ = _poll(client, call_allowed_again, "the key stayed budget-blocked after its spend was reset to 0") + @pytest.mark.covers("mgmt.key.generate.admin_only") def test_generate_forbidden_for_non_admin_key( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index a56eb853823..476165b715d 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -12,6 +12,7 @@ from __future__ import annotations import math import time from collections.abc import Callable +from typing import Final import pytest @@ -42,6 +43,10 @@ from models import ( pytestmark = pytest.mark.e2e +REGENERATE_GRACE_PERIOD = "15s" +REGENERATE_GRACE_SECONDS = 15.0 + + def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: @@ -365,6 +370,36 @@ class TestKeyRegeneration: client, old_rejected, "old key was still accepted after regeneration (never rejected 401) at the deadline" ) + @pytest.mark.covers("other.key_mgmt.regenerate.grace_period_honored") + def test_regenerate_with_grace_period_keeps_old_key_until_revoked( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + old_key = _generate_key(client, resources, KeyGenerateBody(models=["gpt-5.5"])) + + new_key = client.regenerate_key(old_key, grace_period=REGENERATE_GRACE_PERIOD) + resources.defer(lambda: client.proxy.delete_key(new_key)) + revoke_at: Final = time.monotonic() + REGENERATE_GRACE_SECONDS + assert new_key != old_key, "regenerate returned the same key string, so no rotation happened" + + def old_accepted() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.ok else None + + _ = _poll(client, old_accepted, "old key was rejected 401 inside its grace period at the deadline") + assert time.monotonic() < revoke_at, ( + f"old key was only accepted after its {REGENERATE_GRACE_PERIOD} grace period had elapsed" + ) + + def old_rejected() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code == 401 else None + + _ = _poll( + client, + old_rejected, + f"old key was still accepted past its {REGENERATE_GRACE_PERIOD} grace period (never 401) at the deadline", + ) + class TestTeamRoutes: @pytest.mark.covers("mgmt.team.new.persists") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b5229744d6f..5de49ead3ed 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -83,6 +83,16 @@ class KeyGenerateResponse(BaseModel): class KeyRegenerateBody(BaseModel): key: str + grace_period: str | None = None + + +class KeyResetSpendBody(BaseModel): + reset_to: float + + +class KeyResetSpendResponse(BaseModel): + spend: float + previous_spend: float class KeyDeleteBody(BaseModel): diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts index 5c7438cfd11..0537ba193be 100644 --- a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts +++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts @@ -184,6 +184,7 @@ test.describe("Admin tables scroll inside the page", () => { ); try { await navigateToPage(page, Page.TagManagement); + await setRowsPerPage(page, "50"); await expectRowsAtLeast(page, SEED_ROWS); expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); } finally { @@ -207,6 +208,7 @@ test.describe("Admin tables scroll inside the page", () => { ); try { await navigateToPage(page, Page.ModelHubTable); + await setRowsPerPage(page, "50"); await expectRowsAtLeast(page, SEED_ROWS); expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); } finally { diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py index 35d295d581a..291badd91b4 100644 --- a/tests/router_unit_tests/test_router_index_management.py +++ b/tests/router_unit_tests/test_router_index_management.py @@ -241,7 +241,7 @@ class TestRouterIndexManagement: "_get_deployment_by_litellm_model": "lookup by litellm_params.model, which is not indexed", "_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match', "config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)", - "heuristic_v2_router_limit_violation": "counts heuristic_v2 routers across the whole list; admin path only (auto-router init/upsert)", + "auto_router_capability_violation": "counts gated auto-routers across the whole list; admin path only (auto-router init/upsert)", } # Get path to router.py diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py index 1a95e45b2d5..d4e49a1252f 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py +++ b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py @@ -69,6 +69,30 @@ class TestCloudZeroStreamer: assert "2025-01-19" in result assert len(result["2025-01-19"]) == 1 + def test_group_by_date_infers_schema_from_every_row(self): + """Test daily batches retain optional string columns that are null for thousands of leading rows.""" + streamer = CloudZeroStreamer("test-key", "test-connection") + leading_nulls = 10_000 + rows = [ + {"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": None} + for _ in range(leading_nulls) + ] + rows.append( + {"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": "team-alias"} + ) + data = pl.DataFrame( + rows, + schema={"time/usage_start": pl.String, "resource/tag:team_alias": pl.String}, + ) + + result = streamer._group_by_date(data) + + batch = result["2025-01-19"] + assert len(batch) == leading_nulls + 1 + assert batch.schema["resource/tag:team_alias"] == pl.String + assert batch["resource/tag:team_alias"].null_count() == leading_nulls + assert batch.tail(1).item(0, "resource/tag:team_alias") == "team-alias" + def test_parse_and_convert_timestamp_utc(self): """Test _parse_and_convert_timestamp method with UTC timestamp.""" streamer = CloudZeroStreamer("test-key", "test-connection") diff --git a/tests/test_litellm/integrations/cloudzero/test_transform.py b/tests/test_litellm/integrations/cloudzero/test_transform.py index 3ec2fe6779e..cf8d70702f9 100644 --- a/tests/test_litellm/integrations/cloudzero/test_transform.py +++ b/tests/test_litellm/integrations/cloudzero/test_transform.py @@ -86,6 +86,33 @@ class TestCBFTransformer: assert result.is_empty() + def test_transform_keeps_tags_first_seen_after_row_100(self): + transformer = CBFTransformer() + teamless_rows = 101 + team_rows = 2 + total_rows = teamless_rows + team_rows + data = pl.DataFrame( + { + "date": ["2025-01-19"] * total_rows, + "successful_requests": [1] * total_rows, + "spend": [0.5] * total_rows, + "prompt_tokens": [10] * total_rows, + "completion_tokens": [5] * total_rows, + "model": ["gpt-4"] * total_rows, + "custom_llm_provider": ["openai"] * total_rows, + "api_key": ["sk-late-team"] * total_rows, + "team_id": pl.Series([None] * teamless_rows + ["team-late"] * team_rows, dtype=pl.String), + "team_alias": pl.Series([None] * teamless_rows + ["Late Team"] * team_rows, dtype=pl.String), + } + ) + + result = transformer.transform(data) + + assert len(result) == total_rows + assert "resource/tag:team_alias" in result.columns + assert result["resource/tag:team_alias"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows + assert result["resource/tag:entity_id"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows + def test_create_cbf_record(self): """Test _create_cbf_record method with valid row data.""" transformer = CBFTransformer() diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 49a52157c8a..885bd1d4d72 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -17,7 +17,6 @@ if TYPE_CHECKING: class TestCustomGuardrailDeploymentHook: - @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_guardrails(self): """Test that method returns kwargs unchanged when no guardrails are present""" @@ -30,18 +29,14 @@ class TestCustomGuardrailDeploymentHook: "guardrails": None, } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs # Test with guardrails as non-list kwargs["guardrails"] = "not_a_list" - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs @@ -68,9 +63,7 @@ class TestCustomGuardrailDeploymentHook: "user_api_key_request_route": "test_route", } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) # Verify async_pre_call_hook was called with correct parameters custom_guardrail.async_pre_call_hook.assert_called_once() @@ -103,9 +96,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -118,9 +109,7 @@ class TestCustomGuardrailDeploymentHook: } guardrail.mark_pre_call_hook_ran(kwargs) - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 0 @@ -134,9 +123,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -148,9 +135,7 @@ class TestCustomGuardrailDeploymentHook: "metadata": {}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 @@ -179,9 +164,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -193,15 +176,12 @@ class TestCustomGuardrailDeploymentHook: "metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 class TestCustomGuardrailShouldRunGuardrail: - def test_should_run_guardrail_with_litellm_metadata(self): """Test that should_run_guardrail works with litellm_metadata pattern""" from litellm.types.guardrails import GuardrailEventHooks @@ -218,9 +198,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -240,9 +218,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -259,9 +235,7 @@ class TestCustomGuardrailShouldRunGuardrail: # Test with guardrails at root level data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]} - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -281,9 +255,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["different_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is False @@ -302,9 +274,7 @@ class TestCustomGuardrailShouldRunGuardrail: "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True, "Global guardrail should run when default_on=True" # Test 2: User-injected disable at root level is IGNORED @@ -316,9 +286,7 @@ class TestCustomGuardrailShouldRunGuardrail: result = custom_guardrail.should_run_guardrail( data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call ) - assert ( - result is True - ), "User-injected disable_global_guardrails should be ignored" + assert result is True, "User-injected disable_global_guardrails should be ignored" # Test 3: User-injected disable in metadata is IGNORED data_with_disable_metadata = { @@ -349,12 +317,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {"request_tags": ["user-supplied"]}, } - result = custom_guardrail.should_run_guardrail( - data=data_cross_key, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" + result = custom_guardrail.should_run_guardrail(data=data_cross_key, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" # Test 6: After the pre-call strip runs, user-injected # user_api_key_metadata in the non-authoritative metadata key is gone. @@ -365,12 +329,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {}, # post-strip: attacker payload removed } - result = custom_guardrail.should_run_guardrail( - data=data_post_strip, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must be respected when other metadata key is empty" + result = custom_guardrail.should_run_guardrail(data=data_post_strip, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must be respected when other metadata key is empty" def test_should_run_guardrail_key_disable_global_not_overruled_by_team_guardrail_list( self, @@ -436,12 +396,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "opted_out_global_guardrails": ["global_guardrail"], } - assert ( - custom_guardrail.should_run_guardrail( - data=data_root, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_root, event_type=GuardrailEventHooks.pre_call) is True # Test 2: User-injected opt-out in metadata is IGNORED data_metadata = { @@ -450,10 +405,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": ["global_guardrail"]}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_metadata, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_metadata, event_type=GuardrailEventHooks.pre_call) is True ) # Test 4: a different guardrail in the opt-out list → still runs @@ -462,12 +414,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_other, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_other, event_type=GuardrailEventHooks.pre_call) is True # Test 5: empty opt-out list → still runs data_empty = { @@ -475,12 +422,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": []}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_empty, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_empty, event_type=GuardrailEventHooks.pre_call) is True # Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs data_malformed = { @@ -489,10 +431,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": True}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_malformed, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_malformed, event_type=GuardrailEventHooks.pre_call) is True ) def test_should_run_guardrail_opt_out_does_not_affect_non_global(self): @@ -515,12 +454,69 @@ class TestCustomGuardrailShouldRunGuardrail: "guardrails": ["opt_in_guardrail"], }, } - assert ( - non_global.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) - is True + assert non_global.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True + + def test_should_run_guardrail_suppressed_by_auto_router_compression(self): + """An auto router's own compression policy can suppress an otherwise-eligible + guardrail, even one that is default_on and explicitly requested.""" + from litellm.proxy.guardrails import auto_router_compression + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, ) + token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"headroom-default"})) + try: + assert ( + always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call) + is False + ) + finally: + auto_router_compression._suppressed_compression_guardrails.reset(token) + + def test_should_run_guardrail_suppression_does_not_affect_other_names(self): + from litellm.proxy.guardrails import auto_router_compression + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"some-other-guardrail"})) + try: + assert ( + always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call) + is True + ) + finally: + auto_router_compression._suppressed_compression_guardrails.reset(token) + + def test_request_metadata_can_never_suppress_a_guardrail(self): + """Regression (security): suppression state is request-scoped and server-set, + never read from metadata. Metadata reaches spend logs the caller can read, so + anything honored from there is something a later request could replay to switch + off a PII or content-filter guardrail for itself.""" + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + forged = { + "model": "smart-router", + "metadata": { + "_auto_router_suppressed_compression_guardrails": [ + "headroom-default", + "any-token:headroom-default", + ], + }, + } + + assert always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) is True class TestApplyGuardrailCheck: @@ -559,35 +555,33 @@ class TestApplyGuardrailCheck: child_with_override = ChildGuardrailWithOverride() # Test: CustomGuardrail itself has apply_guardrail in its __dict__ - assert ( - "apply_guardrail" in type(CustomGuardrail()).__dict__ - ), "CustomGuardrail should have apply_guardrail in its own __dict__" + assert "apply_guardrail" in type(CustomGuardrail()).__dict__, ( + "CustomGuardrail should have apply_guardrail in its own __dict__" + ) # Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(parent_instance).__dict__ - ), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(parent_instance).__dict__, ( + "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(child_without_override).__dict__ - ), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(child_without_override).__dict__, ( + "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__ - assert ( - "apply_guardrail" in type(child_with_override).__dict__ - ), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + assert "apply_guardrail" in type(child_with_override).__dict__, ( + "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + ) # Verify that all instances still have the method via inheritance (hasattr) - assert hasattr( - parent_instance, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_without_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_with_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" + assert hasattr(parent_instance, "apply_guardrail"), "All instances should have apply_guardrail via inheritance" + assert hasattr(child_without_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) + assert hasattr(child_with_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) class TestGuardrailLoggingAggregation: @@ -614,11 +608,7 @@ class TestGuardrailLoggingAggregation: def test_appends_to_existing_metadata_list(self): request_data = { - "metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "existing_guardrail"} - ] - } + "metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "existing_guardrail"}]} } self._invoke_add_log(request_data) @@ -630,11 +620,7 @@ class TestGuardrailLoggingAggregation: assert info[1]["guardrail_name"] == "test_guardrail" def test_converts_existing_metadata_dict_to_list(self): - request_data = { - "metadata": { - "standard_logging_guardrail_information": {"guardrail_name": "legacy"} - } - } + request_data = {"metadata": {"standard_logging_guardrail_information": {"guardrail_name": "legacy"}}} self._invoke_add_log(request_data) @@ -646,18 +632,12 @@ class TestGuardrailLoggingAggregation: def test_appends_to_litellm_metadata(self): request_data = { - "litellm_metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "litellm_existing"} - ] - } + "litellm_metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "litellm_existing"}]} } self._invoke_add_log(request_data) - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert isinstance(info, list) assert len(info) == 2 assert info[1]["guardrail_name"] == "test_guardrail" @@ -674,12 +654,10 @@ class TestGuardrailLoggingAggregation: self._invoke_add_log(request_data) - assert ( - "standard_logging_guardrail_information" not in request_data["metadata"] - ), "entry landed in the caller's metadata, where the spend log does not read it" - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + assert "standard_logging_guardrail_information" not in request_data["metadata"], ( + "entry landed in the caller's metadata, where the spend log does not read it" + ) + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert len(info) == 1 assert info[0]["guardrail_name"] == "test_guardrail" @@ -697,9 +675,7 @@ class TestGuardrailLoggingAggregation: } self._invoke_add_log(request_data) - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name="test_guardrail" - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name="test_guardrail") buckets = { key @@ -745,9 +721,7 @@ class TestGuardrailOtelSpanEmission: assert len(captured) == 1 emitted = captured[0] - recorded = request_data["metadata"]["standard_logging_guardrail_information"][ - -1 - ] + recorded = request_data["metadata"]["standard_logging_guardrail_information"][-1] assert emitted is recorded assert emitted["guardrail_name"] == "emit_guard" assert emitted["start_time"] == 1.0 @@ -757,9 +731,7 @@ class TestGuardrailOtelSpanEmission: def _boom(_entry): raise RuntimeError("otel exporter down") - monkeypatch.setattr( - "litellm.integrations.otel.logger.emit_guardrail_span", _boom - ) + monkeypatch.setattr("litellm.integrations.otel.logger.emit_guardrail_span", _boom) request_data = {"metadata": {}} self._record(self._make_guardrail(), request_data) @@ -856,9 +828,7 @@ class TestGuardrailSensitiveFieldStripping: duration=1.0, ) - logged_response = request_data["metadata"][ - "standard_logging_guardrail_information" - ][0]["guardrail_response"] + logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert "secret_fields" not in logged_response assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response) @@ -871,9 +841,7 @@ class TestGuardrailSensitiveFieldStripping: guardrail_json_response=[ { "result": "ok", - "secret_fields": { - "raw_headers": {"authorization": "Bearer sk-secret"} - }, + "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}, }, {"result": "also_ok"}, ], @@ -927,9 +895,7 @@ class TestGuardrailResponseCredentialMasking: duration=1.0, ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] masked_key = logged["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] assert masked_key != plaintext_key @@ -938,10 +904,7 @@ class TestGuardrailResponseCredentialMasking: assert logged["model"] == "gpt-4o-mini" assert logged["messages"] == [{"role": "user", "content": "hi"}] - assert ( - logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] - == "proj-name" - ) + assert logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] == "proj-name" def test_nested_user_api_key_auth_metadata_is_masked(self): import json @@ -1000,9 +963,7 @@ class TestGuardrailResponseCredentialMasking: request_data: dict = {"metadata": {}} guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}] - }, + guardrail_json_response={"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}, request_data=request_data, guardrail_status="success", ) @@ -1025,9 +986,7 @@ class TestGuardrailResponseCredentialMasking: guardrail_status="success", ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert logged["flagged"] is True assert logged["score"] == 0.94 assert logged["tokens_used"] == 42 @@ -1039,18 +998,14 @@ class TestGuardrailResponseCredentialMasking: plaintext = "lsv2_pt_abcdef1234567890" guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "metadata_snapshot": { - "callback_vars": {"langsmith_api_key": plaintext} - } - }, + guardrail_json_response={"metadata_snapshot": {"callback_vars": {"langsmith_api_key": plaintext}}}, request_data=request_data, guardrail_status="success", ) - masked = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ]["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] + masked = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"][ + "metadata_snapshot" + ]["callback_vars"]["langsmith_api_key"] assert masked != plaintext assert masked.startswith(plaintext[:4]) assert masked.endswith(plaintext[-4:]) @@ -1544,9 +1499,7 @@ class TestEventTypeLogging: guardrail = TestGuardrail() request_data = {"metadata": {}} - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1, ( @@ -1588,9 +1541,7 @@ class TestEventTypeLogging: request_data = {"metadata": {}} with pytest.raises(ValueError, match="blocked"): - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1 @@ -1719,9 +1670,7 @@ class TestTracingFieldsPopulation: guardrail_json_response="blocked", request_data=request_data, guardrail_status="guardrail_intervened", - tracing_detail=GuardrailTracingDetail( - policy_template="EU AI Act Article 5" - ), + tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"), ) slg_list = request_data["metadata"]["standard_logging_guardrail_information"] @@ -1763,13 +1712,7 @@ class TestCustomGuardrailSpendLogMatchRedaction: cg = CustomGuardrail(guardrail_name="test-rail") raw = { "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ] } request_data: dict = {"metadata": {}} @@ -1780,17 +1723,10 @@ class TestCustomGuardrailSpendLogMatchRedaction: ) slg = request_data["metadata"]["standard_logging_guardrail_information"][0] assert ( - slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) - assert ( - raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "GG" - ) + assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "GG" def test_add_standard_logging_redacts_regex_field(self): cg = CustomGuardrail(guardrail_name="test-rail") diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index af75691eb10..1991170707d 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1114,6 +1114,56 @@ async def test_logging_non_streaming_request(): litellm.callbacks = original_callbacks +@pytest.mark.asyncio +async def test_async_success_handler_truncates_large_base64_off_the_event_loop(monkeypatch): + """The standard logging payload's base64 scan of a large multimodal request must not run on the loop thread.""" + import threading + + from litellm.litellm_core_utils import logging_utils + + loop_thread = threading.get_ident() + scan_threads: list[int] = [] + original_scan = logging_utils._truncate_base64_in_string + + def recording_scan(value: str) -> str: + scan_threads.append(threading.get_ident()) + return original_scan(value) + + monkeypatch.setattr(logging_utils, "_truncate_base64_in_string", recording_scan) + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + + logged = asyncio.Event() + captured: dict = {} + + class CaptureLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + captured["standard_logging_object"] = kwargs["standard_logging_object"] + logged.set() + + monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()]) + payload = "L" * 20_000 + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{payload}"}}, + ], + } + ], + mock_response="ok", + ) + await asyncio.wait_for(logged.wait(), timeout=10) + + logged_url = captured["standard_logging_object"]["messages"][0]["content"][1]["image_url"]["url"] + assert "base64_data truncated" in logged_url + assert payload not in logged_url + assert scan_threads + assert loop_thread not in scan_threads + + @pytest.mark.parametrize( "async_flag", [ @@ -1496,6 +1546,62 @@ async def test_dispatch_failure_handlers_async_completes_before_sync_submit( assert events == ["async_start", "async_end", "sync_submit"] +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_submits_sync_handler_when_task_is_cancelled( + logging_obj, +): + """Cancelling the dispatch task mid-await still submits the sync failure_handler. + + Router failure paths fire the dispatcher with ``asyncio.create_task`` and raise + right away. When the event loop is torn down before the task finishes (a short + ``asyncio.run`` in the SDK), the cancelled task must still hand the sync callbacks + to the executor, as the old raw-thread path did, and only once the async handler + has stopped. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + events: list[str] = [] + async_started = asyncio.Event() + + async def _async_failure(exc, tb, **kwargs): + events.append("async_start") + async_started.set() + await asyncio.sleep(10) + events.append("async_end") + + def _submit(*args, **kwargs): + events.append("sync_submit") + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", new_callable=MagicMock), + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=True, + ), + patch( # test-quality-ok: the executor submit is the observable + "litellm.litellm_core_utils.litellm_logging.executor.submit", + side_effect=_submit, + ), + ): + task = asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + ) + await async_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert events == ["async_start", "sync_submit"] + + @pytest.mark.asyncio async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks( logging_obj, @@ -6027,6 +6133,34 @@ def test_failure_handler_helper_fn_builds_payload_once_per_exception(): assert obj.model_call_details["standard_logging_object"] is not first_payload +@pytest.mark.asyncio +async def test_sync_failure_handler_reuses_payload_after_callable_async_callback(): + """Regression for LIT-6886: the proxy runs async_failure_handler, then the threaded + failure_handler, for every rejected request. A plain-function async callback (the + Router registers one) is dispatched through CustomLogger.async_log_event, which + restamps log_event_type on the shared model_call_details; the sync handler then + rebuilt the standardized payload, doubling the redaction and payload cost of a 403.""" + router_style_callback = AsyncMock() + obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6886-1", + function_id="f", + dynamic_async_failure_callbacks=[router_style_callback], + ) + exc = _raise_and_catch(_ClientError(status_code=403, message="key not allowed to access model")) + await obj.async_failure_handler(exception=exc, traceback_exception="") + first_payload = obj.model_call_details["standard_logging_object"] + assert first_payload is not None + assert router_style_callback.await_count == 1 + + obj.failure_handler(exc, "") + assert obj.model_call_details["standard_logging_object"] is first_payload + + @pytest.mark.asyncio async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_obj): """The savings gate reads litellm_gateway_injected_cache from the request's diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/test_litellm/litellm_core_utils/test_logging_utils.py index b0dad0bf228..f9913f1935d 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_utils.py @@ -2,12 +2,16 @@ Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers. """ +import threading + import pytest +from litellm.litellm_core_utils import logging_utils from litellm.litellm_core_utils.logging_utils import ( _format_base64_size, _truncate_base64_in_string, truncate_base64_in_messages, + truncate_base64_in_messages_async, ) # --------------------------------------------------------------------------- @@ -157,3 +161,70 @@ class TestTruncateBase64InMessages: result[0]["content"][0]["image_url"]["url"] == f"data:image/png;base64,{short}" ) + + +# --------------------------------------------------------------------------- +# truncate_base64_in_messages_async +# --------------------------------------------------------------------------- + + +def _image_messages(payload: str) -> list: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{payload}"}}, + ], + } + ] + + +@pytest.fixture +def scan_threads(monkeypatch): + """Record the thread that runs every base64 regex scan.""" + threads: list[int] = [] + original = logging_utils._truncate_base64_in_string + + def recording_scan(value: str) -> str: + threads.append(threading.get_ident()) + return original(value) + + monkeypatch.setattr(logging_utils, "_truncate_base64_in_string", recording_scan) + return threads + + +class TestTruncateBase64InMessagesAsync: + @pytest.mark.asyncio + async def test_large_payload_is_scanned_off_the_event_loop(self, monkeypatch, scan_threads): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + payload = "I" * 20_000 + messages = _image_messages(payload) + + result = await truncate_base64_in_messages_async(messages) + offload_threads = tuple(scan_threads) + + assert result == truncate_base64_in_messages(messages) + assert payload not in result[0]["content"][1]["image_url"]["url"] + assert payload in messages[0]["content"][1]["image_url"]["url"] + assert offload_threads + assert threading.get_ident() not in offload_threads + + @pytest.mark.asyncio + async def test_small_payload_stays_on_the_calling_thread(self, monkeypatch, scan_threads): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + messages = _image_messages("J" * 200) + + result = await truncate_base64_in_messages_async(messages) + + assert result == truncate_base64_in_messages(messages) + assert scan_threads + assert set(scan_threads) == {threading.get_ident()} + + @pytest.mark.asyncio + async def test_none_and_disabled_truncation_short_circuit(self, monkeypatch, scan_threads): + assert await truncate_base64_in_messages_async(None) is None + monkeypatch.setattr(logging_utils, "MAX_BASE64_LENGTH_FOR_LOGGING", 0) + messages = _image_messages("K" * 20_000) + assert await truncate_base64_in_messages_async(messages) is messages + assert scan_threads == [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 5e2f2cf97d7..6b3098d9e60 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -9,9 +9,11 @@ returning the stub. import asyncio import logging +import time from datetime import datetime, timedelta, timezone import httpx +import jwt as pyjwt import pytest from pydantic import SecretStr @@ -42,6 +44,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto OAuthToken, TokenStoreUnavailable, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + RefreshingSSOAssertionStore, + SSOAssertionRefresher, + SSOClientConfig, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( AssertionStoreUnavailable, SSOIdentityAssertion, @@ -589,6 +596,73 @@ async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_id assert endpoint.calls == [] +@pytest.mark.asyncio +async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging(): + """The unattended-agent case end to end: the user last signed in more than an id_token lifetime + ago, so without renewal this is the 412 above. With the renewing store wired the arm resolves, + and leg 1 asserts the renewed token rather than the one that ran out.""" + renewed_id_token = pyjwt.encode( + {"iss": "https://idp.example.com", "sub": "alice", "exp": int(time.time()) + 3600}, + "test-idp-signing-key-32-bytes-long-xxxx", + algorithm="HS256", + ) + expired = SSOIdentityAssertion( + id_token=SecretStr("stale-id-token"), + refresh_token=SecretStr("rt_1"), + expires_at=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + rows = {"alice": expired} + + async def _read(user_id: str) -> SSOIdentityAssertion | None: + return rows.get(user_id) + + async def _write(user_id: str, assertion: SSOIdentityAssertion) -> None: + rows[user_id] = assertion + + class _Inner: + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + return await _read(user_id) + + class _Transport: + async def post(self, url, form, headers): + return Ok({"access_token": "at", "id_token": renewed_id_token}) + + refresher = SSOAssertionRefresher( + _Transport(), + client_config=lambda: SSOClientConfig( + token_endpoint="https://idp.example.com/token", + client_id="litellm", + client_secret=SecretStr("s"), + auth_method="client_secret_basic", + ), + read=_read, + write=_write, + ) + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, + sso_assertion_store=RefreshingSSOAssertionStore( + _Inner(), refresher, fresh_read=_read, coordinator_factory=lambda: None + ), + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == renewed_id_token + + +def test_the_resolver_defaults_to_the_renewing_assertion_store(): + """A resolver built without collaborators is what production gets, so the default has to renew; + the plain database reader would strand every agent an id_token lifetime after its user's login.""" + provider = UpstreamCredentialProvider() + + assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # noqa: SLF001 # the wiring is the assertion + + @pytest.mark.asyncio async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry(): endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py new file mode 100644 index 00000000000..d80913f8d33 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py @@ -0,0 +1,794 @@ +"""Tests for renewing the stored SSO identity assertion behind the ID-JAG arm. + +Pins the contract an unattended agent depends on: an assertion that has run out is renewed from the +refresh token captured beside it instead of stranding the agent until its user signs in again, the +IdP sees one redemption per user no matter how many tool calls arrive at once, a rotation is written +back without overwriting a sign-in that landed mid-renewal, and the two failure kinds stay +distinguishable - a dead refresh token still challenges the user, an unreachable IdP does not. +""" + +import asyncio +import base64 +import itertools +import logging +import time +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime, timedelta, timezone + +import httpx +import jwt as pyjwt +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + HttpxTokenEndpointTransport, + RefreshFailure, + RefreshingSSOAssertionStore, + SSOAssertionRefresher, + SSOClientConfig, + default_sso_assertion_store, + sso_client_config, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + SSOIdentityAssertion, +) + +SIGNING_KEY = "test-idp-signing-key-32-bytes-long-xxxx" +ISSUER = "https://idp.example.com" +TOKEN_ENDPOINT = "https://idp.example.com/token" + +_CLIENT = SSOClientConfig( + token_endpoint=TOKEN_ENDPOINT, + client_id="litellm", + client_secret=SecretStr("s3cret"), + auth_method="client_secret_basic", +) +_POST_CLIENT = SSOClientConfig( + token_endpoint=TOKEN_ENDPOINT, + client_id="litellm", + client_secret=SecretStr("s3cret"), + auth_method="client_secret_post", +) + + +_MINTED = itertools.count() + + +def _id_token(subject: str = "u1", exp_offset: int = 3600) -> str: + """A distinct token per call. Two mints with the same claims in the same second would encode + identically, which would let a test that means "the renewed token replaced the old one" pass + while comparing a value to itself.""" + return pyjwt.encode( + {"iss": ISSUER, "sub": subject, "exp": int(time.time()) + exp_offset, "jti": f"t{next(_MINTED)}"}, + SIGNING_KEY, + algorithm="HS256", + ) + + +def _stored(id_token: str, *, expires_in: int, refresh_token: str | None = "rt_1") -> SSOIdentityAssertion: + """A row as the SSO callback wrote it: ``expires_in`` seconds from now, mirroring the id_token.""" + return SSOIdentityAssertion( + id_token=SecretStr(id_token), + refresh_token=SecretStr(refresh_token) if refresh_token else None, + issuer=ISSUER, + expires_at=datetime.now(timezone.utc) + timedelta(seconds=expires_in), + ) + + +class _FakeRows: + """The one assertion row per user: the inner read seam and the refresher's read/write pair.""" + + def __init__(self, rows: dict[str, SSOIdentityAssertion] | None = None) -> None: + self.rows: dict[str, SSOIdentityAssertion] = dict(rows or {}) + self.cached_rows: dict[str, SSOIdentityAssertion] = {} + self.reads: list[str] = [] + self.writes: list[tuple[str, SSOIdentityAssertion]] = [] + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + self.reads.append(user_id) + # A real suspension point, so concurrent callers interleave here instead of running to + # completion one at a time and never actually racing. + await asyncio.sleep(0) + return self.cached_rows.get(user_id, self.rows.get(user_id)) + + async def fetch_fresh(self, user_id: str) -> SSOIdentityAssertion | None: + self.reads.append(user_id) + await asyncio.sleep(0) + return self.rows.get(user_id) + + async def write(self, user_id: str, assertion: SSOIdentityAssertion) -> None: + self.writes.append((user_id, assertion)) + self.rows[user_id] = assertion + + +class _FakeTransport: + """Answers every refresh with the same canned result, optionally holding until ``gate`` opens.""" + + def __init__( + self, + response: Result[Mapping[str, object], RefreshFailure], + *, + gate: asyncio.Event | None = None, + on_call: Callable[[], None] | None = None, + ) -> None: + self._response = response + self._gate = gate + self._on_call = on_call + self.calls: list[tuple[str, dict[str, str]]] = [] + self.headers: list[dict[str, str]] = [] + + async def post( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> Result[Mapping[str, object], RefreshFailure]: + self.calls.append((url, dict(form))) + self.headers.append(dict(headers)) + if self._on_call is not None: + self._on_call() + if self._gate is not None: + await self._gate.wait() + return self._response + + +def _renewal(id_token: str, refresh_token: str | None = None) -> Result[Mapping[str, object], RefreshFailure]: + body: dict[str, object] = {"access_token": "at", "id_token": id_token, "token_type": "Bearer"} + return Ok({**body, "refresh_token": refresh_token} if refresh_token else body) + + +def _store( + rows: _FakeRows, + transport: _FakeTransport, + *, + client_config: Callable[[], SSOClientConfig | None] = lambda: _CLIENT, + coordinator_factory: Callable[[], object] = lambda: None, +) -> RefreshingSSOAssertionStore: + refresher = SSOAssertionRefresher(transport, client_config=client_config, read=rows.fetch, write=rows.write) + return RefreshingSSOAssertionStore( + rows, + refresher, + fresh_read=rows.fetch_fresh, + coordinator_factory=coordinator_factory, # pyright: ignore[reportArgumentType] # test doubles stand in for the runtime factory + ) + + +async def _until(predicate: Callable[[], bool]) -> None: + for _ in range(2000): + if predicate(): + return + await asyncio.sleep(0) + raise AssertionError("condition never became true") + + +@pytest.mark.asyncio +async def test_an_expiring_assertion_is_renewed_and_the_renewal_is_what_the_reader_gets(): + """The whole point: an agent calling after its user's id_token ran out keeps working.""" + stale, fresh = _id_token(exp_offset=-1), _id_token() + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(_renewal(fresh)) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == fresh + assert len(transport.calls) == 1 + url, form = transport.calls[0] + assert url == TOKEN_ENDPOINT + assert form["grant_type"] == "refresh_token" + assert form["refresh_token"] == "rt_1" + + +@pytest.mark.asyncio +async def test_a_basic_auth_login_gets_a_basic_auth_refresh(): + """The non-PKCE login always sends HTTP Basic, so the renewal must too; credentials in the body + would 401 against an IdP application registered for Basic.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + await _store(rows, transport).fetch("alice") + + expected = base64.b64encode(b"litellm:s3cret").decode() + assert transport.headers[0]["Authorization"] == f"Basic {expected}" + _url, form = transport.calls[0] + assert "client_secret" not in form + assert "client_id" not in form + + +@pytest.mark.asyncio +async def test_a_body_credential_login_gets_a_body_credential_refresh(): + """The mirror case. A PKCE deployment with GENERIC_INCLUDE_CLIENT_ID set signs in with the + credentials in the body, so Basic here would 401 against an application registered for post; the + renewal has to follow the login rather than a constant.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + await _store(rows, transport, client_config=lambda: _POST_CLIENT).fetch("alice") + + assert "Authorization" not in transport.headers[0] + _url, form = transport.calls[0] + assert form["client_id"] == "litellm" + assert form["client_secret"] == "s3cret" + + +@pytest.mark.parametrize( + ("include_client_id", "expected"), + [ + (None, "client_secret_basic"), + ("false", "client_secret_basic"), + ("TRUE", "client_secret_post"), + ("true", "client_secret_post"), + ], +) +def test_the_auth_method_follows_the_flag_the_login_reads(include_client_id, expected): + """``GENERIC_INCLUDE_CLIENT_ID`` is what the PKCE login branches on, parsed the same way it + parses it, so the renewal cannot pick a method the sign-in did not use.""" + env = { + "GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, + "GENERIC_CLIENT_ID": "litellm", + "GENERIC_CLIENT_SECRET": "s3cret", + **({"GENERIC_INCLUDE_CLIENT_ID": include_client_id} if include_client_id is not None else {}), + } + + config = sso_client_config(env) + + assert config is not None + assert config.auth_method == expected + + +@pytest.mark.asyncio +async def test_an_assertion_well_inside_its_lifetime_never_reaches_the_idp(): + """The common path must cost exactly what it did before this store existed.""" + current = _id_token() + rows = _FakeRows({"alice": _stored(current, expires_in=1800)}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == current + assert transport.calls == [] + assert rows.writes == [] + + +@pytest.mark.asyncio +async def test_renewal_starts_inside_the_skew_rather_than_after_expiry(): + """A token that would die between resolution and the second exchange leg is replaced first.""" + about_to_expire, fresh = _id_token(), _id_token() + assert about_to_expire != fresh + rows = _FakeRows({"alice": _stored(about_to_expire, expires_in=30)}) + transport = _FakeTransport(_renewal(fresh)) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == fresh + + +@pytest.mark.asyncio +async def test_a_user_with_no_stored_assertion_is_still_absent(): + rows = _FakeRows() + transport = _FakeTransport(_renewal(_id_token())) + + assert await _store(rows, transport).fetch("nobody") is None + assert transport.calls == [] + + +@pytest.mark.asyncio +async def test_a_refused_refresh_leaves_the_expired_assertion_for_the_reader_to_reject(): + """A dead refresh token is the user's problem, and the reader's expiry guard is what tells them; + swapping in a renewed-looking value or hiding the row would break that challenge.""" + stale = _id_token(exp_offset=-1) + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(Error(RefreshFailure.of_rejected("the IdP refused the refresh with status 400"))) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == stale + assert rows.writes == [] + + +@pytest.mark.asyncio +async def test_an_unreachable_idp_is_a_store_outage_not_a_sign_in_again_challenge(): + """503, not 412: the user has nothing to fix by signing in again while the IdP is down.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(Error(RefreshFailure.of_unavailable("the IdP token endpoint is unreachable"))) + + with pytest.raises(AssertionStoreUnavailable): + await _store(rows, transport).fetch("alice") + + +@pytest.mark.asyncio +async def test_a_missing_refresh_token_names_the_scope_the_operator_has_to_set(caplog): + """Nothing to redeem is the default state of a deployment, so the log has to say what to change + or the feature stays silently inert.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1, refresh_token=None)}) + transport = _FakeTransport(_renewal(_id_token())) + + with caplog.at_level(logging.WARNING): + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert transport.calls == [] + assert "GENERIC_SCOPE" in caplog.text + assert "offline_access" in caplog.text + + +@pytest.mark.asyncio +async def test_an_unconfigured_sso_client_never_calls_the_idp(caplog): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + with caplog.at_level(logging.WARNING): + served = await _store(rows, transport, client_config=lambda: None).fetch("alice") + + assert served is not None + assert transport.calls == [] + assert "GENERIC_TOKEN_ENDPOINT" in caplog.text + + +@pytest.mark.asyncio +async def test_a_refresh_response_carrying_no_id_token_is_refused(caplog): + """An access token is not an identity assertion, so there is nothing to assert upstream.""" + stale = _id_token(exp_offset=-1) + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(Ok({"access_token": "at", "token_type": "Bearer"})) + + with caplog.at_level(logging.WARNING): + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == stale + assert rows.writes == [] + assert "openid" in caplog.text + + +@pytest.mark.asyncio +async def test_a_rotated_refresh_token_replaces_the_stored_one(): + """An IdP that rotates invalidates the old token, so keeping it would cost a sign-in next time.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(), refresh_token="rt_2")) + + await _store(rows, transport).fetch("alice") + + stored = rows.rows["alice"] + assert stored.refresh_token is not None + assert stored.refresh_token.get_secret_value() == "rt_2" + + +@pytest.mark.asyncio +async def test_an_omitted_refresh_token_carries_the_previous_one_forward(): + """An IdP that does not rotate expects the original to keep working; dropping it would strand + the user after exactly one renewal.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + await _store(rows, transport).fetch("alice") + + stored = rows.rows["alice"] + assert stored.refresh_token is not None + assert stored.refresh_token.get_secret_value() == "rt_1" + + +@pytest.mark.asyncio +async def test_the_renewed_expiry_moves_forward_so_the_next_read_does_not_refresh_again(): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(exp_offset=3600))) + store = _store(rows, transport) + + await store.fetch("alice") + await store.fetch("alice") + + assert len(transport.calls) == 1 + + +async def _explode(user_id: str, assertion: SSOIdentityAssertion) -> None: + raise RuntimeError("write failed") + + +@pytest.mark.asyncio +async def test_a_renewal_that_cannot_be_recorded_is_reported_as_transient(): + """The store is what every caller reads, so a renewal nobody can see is not a success. Calling it + one would hand back a token the gateway failed to record.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + refresher = SSOAssertionRefresher( + _FakeTransport(_renewal(_id_token())), client_config=lambda: _CLIENT, read=rows.fetch, write=_explode + ) + + outcome = await refresher.refresh("alice", rows.rows["alice"]) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_failed_write_does_not_tell_the_user_to_sign_in_again(): + """A database that cannot take the write is not something signing in again fixes, so the reader + has to see an outage rather than the stale row's expiry.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode) + store = RefreshingSSOAssertionStore( + rows, + refresher, + fresh_read=rows.fetch_fresh, + coordinator_factory=lambda: None, # pyright: ignore[reportArgumentType] # test double stands in for the runtime factory + ) + + with pytest.raises(AssertionStoreUnavailable): + await store.fetch("alice") + + assert len(transport.calls) == 1 + + +@pytest.mark.asyncio +async def test_concurrent_reads_for_one_user_redeem_the_refresh_token_once(): + """A burst of tool calls must not replay one refresh token N times: an IdP that rotates reads + that as reuse and can revoke the whole grant chain.""" + gate = asyncio.Event() + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + fresh = _id_token() + transport = _FakeTransport(_renewal(fresh), gate=gate) + store = _store(rows, transport) + + callers = [asyncio.create_task(store.fetch("alice")) for _ in range(8)] + await _until(lambda: len(transport.calls) >= 1 and len(rows.reads) >= 8) + # Guards against a vacuous pass: every caller must have read the expired row and entered the + # renewal branch while the winner is still blocked, otherwise they never raced at all. + assert len(rows.reads) >= 8 + assert not any(task.done() for task in callers) + + gate.set() + served = await asyncio.gather(*callers) + + assert len(transport.calls) == 1 + assert {assertion.id_token.get_secret_value() for assertion in served if assertion is not None} == {fresh} + + +@pytest.mark.asyncio +async def test_concurrent_reads_for_different_users_each_get_their_own_refresh(): + """Single-flight is per user; collapsing across users would leave everyone but one stranded.""" + gate = asyncio.Event() + rows = _FakeRows( + { + "alice": _stored(_id_token("alice", exp_offset=-1), expires_in=-1), + "bob": _stored(_id_token("bob", exp_offset=-1), expires_in=-1), + } + ) + transport = _FakeTransport(_renewal(_id_token()), gate=gate) + store = _store(rows, transport) + + callers = [asyncio.create_task(store.fetch(user)) for user in ("alice", "bob")] + await _until(lambda: len(transport.calls) >= 2) + gate.set() + await asyncio.gather(*callers) + + assert len(transport.calls) == 2 + assert {form["refresh_token"] for _url, form in transport.calls} == {"rt_1"} + + +@pytest.mark.asyncio +async def test_a_renewal_writes_back_when_the_row_did_not_move(): + """The refresh-then-sign-in ordering: nothing displaced the row, so the rotation must land.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + fresh = _id_token() + transport = _FakeTransport(_renewal(fresh, refresh_token="rt_2")) + + served = await _store(rows, transport).fetch("alice") + + assert [user_id for user_id, _assertion in rows.writes] == ["alice"] + assert rows.rows["alice"].id_token.get_secret_value() == fresh + assert served is not None + assert served.id_token.get_secret_value() == fresh + + +@pytest.mark.asyncio +async def test_a_sign_in_landing_mid_renewal_is_not_overwritten(): + """The sign-in-then-refresh ordering. The login wrote a newer assertion while the IdP call was in + flight; overwriting it would put back a refresh token the IdP has already rotated away, costing + that user a sign-in later.""" + from_login = _stored(_id_token("alice"), expires_in=3600, refresh_token="rt_from_login") + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + + def _login_lands() -> None: + rows.rows["alice"] = from_login + + transport = _FakeTransport(_renewal(_id_token(), refresh_token="rt_2"), on_call=_login_lands) + + served = await _store(rows, transport).fetch("alice") + + assert rows.writes == [] + stored = rows.rows["alice"] + assert stored.refresh_token is not None + assert stored.refresh_token.get_secret_value() == "rt_from_login" + assert served is not None + assert served.id_token.get_secret_value() == from_login.id_token.get_secret_value() + + +class _RecordingCoordinator: + """Stands in for the cross-replica coordinator, running the winner's refresh inline.""" + + def __init__(self) -> None: + self.runs: list[tuple[str, str]] = [] + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[None]], + reread: Callable[[], Awaitable[None]], + ) -> None: + self.runs.append((user_id, server_id)) + return await refresh() + + +class _ReplaceThenRefreshCoordinator: + """Replaces the row before running the elected refresh.""" + + def __init__(self, replace: Callable[[], None]) -> None: + self._replace = replace + self.runs: list[tuple[str, str]] = [] + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[None]], + reread: Callable[[], Awaitable[None]], + ) -> None: + self.runs.append((user_id, server_id)) + self._replace() + return await refresh() + + +class _HeldCoordinator: + """Emulates a cross-replica holder finishing before the loser re-reads.""" + + def __init__(self, before_reread: Callable[[], None] | None = None) -> None: + self._before_reread = before_reread + self.runs: list[tuple[str, str]] = [] + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[None]], + reread: Callable[[], Awaitable[None]], + ) -> None: + self.runs.append((user_id, server_id)) + if self._before_reread is not None: + self._before_reread() + return await reread() + + +@pytest.mark.asyncio +async def test_an_elected_renewal_redeems_the_row_it_re_reads_not_the_one_it_entered_with(): + stale = _id_token(exp_offset=-1) + fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2") + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + coordinator = _ReplaceThenRefreshCoordinator(lambda: rows.rows.__setitem__("alice", fresh)) + + served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert transport.calls == [] + assert served is not None + assert served.id_token.get_secret_value() == fresh.id_token.get_secret_value() + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_whose_winner_renewed_reads_the_renewal_without_redeeming(): + fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2") + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + coordinator = _HeldCoordinator(before_reread=lambda: rows.rows.__setitem__("alice", fresh)) + + served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert served is fresh + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_rereads_past_a_stale_process_local_cache(): + stale = _stored(_id_token(exp_offset=-1), expires_in=-1) + fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2") + rows = _FakeRows({"alice": fresh}) + rows.cached_rows["alice"] = stale + transport = _FakeTransport(_renewal(_id_token())) + coordinator = _HeldCoordinator() + + served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert served is fresh + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_in_challenge(): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode) + coordinator = _HeldCoordinator() + store = RefreshingSSOAssertionStore( + rows, + refresher, + fresh_read=rows.fetch_fresh, + coordinator_factory=lambda: coordinator, # pyright: ignore[reportArgumentType] # test double stands in for the runtime factory + ) + + with pytest.raises(AssertionStoreUnavailable): + await store.fetch("alice") + + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_never_redeems_the_token_the_holder_may_have_rotated(): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(Error(RefreshFailure.of_rejected("dead"))) + coordinator = _HeldCoordinator() + + with pytest.raises(AssertionStoreUnavailable): + await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_the_cross_replica_coordinator_is_used_and_built_once(): + """Redis elects one refresher across the fleet; rebuilding its client per renewal would open a + connection every time.""" + coordinator = _RecordingCoordinator() + builds: list[int] = [] + + def _factory() -> object: + builds.append(1) + return coordinator + + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(exp_offset=-1))) + store = _store(rows, transport, coordinator_factory=_factory) + + await store.fetch("alice") + await store.fetch("alice") + + assert len(builds) == 1 + assert coordinator.runs == [("alice", "sso_identity_assertion"), ("alice", "sso_identity_assertion")] + + +@pytest.mark.asyncio +async def test_the_in_process_coordinator_is_retried_until_redis_appears(): + """A proxy that gains Redis after boot must stop electing a winner per worker.""" + coordinator = _RecordingCoordinator() + available: list[bool] = [False] + + def _factory() -> object | None: + return coordinator if available[0] else None + + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(exp_offset=-1))) + store = _store(rows, transport, coordinator_factory=_factory) + + await store.fetch("alice") + assert coordinator.runs == [] + + available[0] = True + await store.fetch("alice") + assert coordinator.runs == [("alice", "sso_identity_assertion")] + + +@pytest.mark.parametrize( + "env", + [ + {}, + {"GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"}, + {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_SECRET": "s"}, + {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_ID": "litellm"}, + {"GENERIC_TOKEN_ENDPOINT": "", "GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"}, + ], +) +def test_a_partial_sso_client_is_no_client(env): + """Redeeming against a half-configured client would post credentials nowhere useful; the arm + treats it as "cannot renew" and falls back to the sign-in challenge.""" + assert sso_client_config(env) is None + + +def test_the_configured_sso_client_is_the_one_the_login_used(): + config = sso_client_config( + { + "GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, + "GENERIC_CLIENT_ID": "litellm", + "GENERIC_CLIENT_SECRET": "s3cret", + } + ) + + assert config is not None + assert config.token_endpoint == TOKEN_ENDPOINT + assert config.client_id == "litellm" + assert config.client_secret.get_secret_value() == "s3cret" + + +def test_the_live_store_renews_over_the_database_reader(): + """The composition root has to produce a renewing store, or none of this runs in production.""" + assert isinstance(default_sso_assertion_store(), RefreshingSSOAssertionStore) + + +def _responding(response: httpx.Response | None) -> HttpxTokenEndpointTransport: + async def _post(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None: + return response + + return HttpxTokenEndpointTransport(_post) + + +def _json_response(status: int, payload: dict[str, object]) -> httpx.Response: + return httpx.Response(status, json=payload, request=httpx.Request("POST", TOKEN_ENDPOINT)) + + +@pytest.mark.parametrize("status", [400, 401, 403]) +@pytest.mark.asyncio +async def test_the_idp_declining_the_grant_is_a_refusal_the_user_must_act_on(status): + """A 4xx means this refresh token is finished; calling that an outage would sit the user behind a + 503 forever instead of telling them to sign in.""" + outcome = await _responding(_json_response(status, {"error": "invalid_grant"})).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "rejected" + + +@pytest.mark.parametrize("status", [500, 502, 503]) +@pytest.mark.asyncio +async def test_a_failing_idp_is_an_outage_not_a_refusal(status): + """The refresh token is probably fine; telling the user to sign in again would blame them for + someone else's outage, and would burn their session for nothing.""" + outcome = await _responding(_json_response(status, {})).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_an_unreachable_endpoint_is_an_outage(): + async def _post(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None: + raise httpx.ConnectError("connection refused") + + outcome = await HttpxTokenEndpointTransport(_post).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_non_json_body_is_an_outage(): + response = httpx.Response(200, text="maintenance", request=httpx.Request("POST", TOKEN_ENDPOINT)) + + outcome = await _responding(response).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_missing_response_is_an_outage(): + outcome = await _responding(None).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_successful_grant_is_handed_back_as_the_parsed_body(): + outcome = await _responding(_json_response(200, {"access_token": "at", "id_token": "idt"})).post( + TOKEN_ENDPOINT, {"grant_type": "refresh_token"}, {} + ) + + assert isinstance(outcome, Ok) + assert outcome.ok["id_token"] == "idt" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 34dc067e7a3..e3ed48713aa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -461,6 +461,30 @@ class TestMCPServerManager: base.update(overrides) return {"m2mserver": base} + def _id_jag_config(self): + return { + "idjag_server": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2_id_jag, + "client_id": "cid", + "client_secret": "csec", + "token_exchange_endpoint": "https://idp.example.com/token", + "id_jag_resource_token_endpoint": "https://resource.example.com/token", + "id_jag_resource": "https://resource.example.com", + } + } + + def _clear_sso_env(self, monkeypatch): + for env_var in ( + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + ): + monkeypatch.delenv(env_var, raising=False) + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on"]) def test_mcp_oauth_discovery_on_startup_true_values(self, value): with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": value}): @@ -1130,6 +1154,72 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.oauth2_flow is None + @pytest.mark.asyncio + async def test_load_servers_from_config_warns_for_id_jag_with_google_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + warnings = [message for message in caplog.messages if "oauth2_id_jag" in message] + assert len(warnings) == 1 + assert "idjag_server" in warnings[0] + assert "GENERIC_CLIENT_ID" in warnings[0] + + @pytest.mark.asyncio + async def test_load_servers_from_config_does_not_warn_for_id_jag_without_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) + + @pytest.mark.asyncio + async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + manager = MCPServerManager() + config = { + "api_key_server": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.api_key, + "auth_value": "upstream-secret", + } + } + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(config) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) + + @pytest.mark.asyncio + async def test_load_servers_from_config_does_not_warn_for_id_jag_with_generic_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) + def _client_forwarded_config(self, auth_type, **overrides): base = { "url": "https://example.com/mcp", diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py new file mode 100644 index 00000000000..db2f94306fb --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -0,0 +1,381 @@ +"""Unit tests for litellm.proxy.guardrails.auto_router_compression.""" + +import json +from typing import Any + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails import auto_router_compression +from litellm.proxy.guardrails.auto_router_compression import ( + AutoRouterCompressionPolicy, + arm_pre_call, + messages_for_routing, + policy_for_model, + policy_from_litellm_params, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + + +class TestPolicyFromLitellmParams: + def test_neither_key_set_is_no_policy(self): + assert policy_from_litellm_params({}) is None + + def test_routing_only(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "headroom-a"}) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_normalizes_to_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"} + ) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_is_case_insensitive(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "NONE"}) + assert policy == AutoRouterCompressionPolicy(routing=None, model=None) + + def test_is_same_true_for_matching_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "x"} + ) + assert policy.is_same is True + + def test_is_same_false_for_different_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "y"} + ) + assert policy.is_same is False + + def test_is_same_true_when_both_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"} + ) + assert policy.is_same is True + + +class _FakeRouter: + """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + + def __init__(self, deployments: list[dict[str, Any]]): + self._deployments = deployments + + def get_model_list(self, model_name, team_id=None): + return [d for d in self._deployments if d.get("model_name") == model_name] + + +def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: + return { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + **compression, + **({"tags": tags} if tags is not None else {}), + }, + } + + +class TestPolicyForModel: + def test_no_router_returns_none(self): + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None + + def test_no_marker_deployment_returns_none(self): + router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + + def test_marker_deployment_without_policy_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + + def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_picks_the_marker_whose_tags_the_request_carries(self): + """Regression: an alias with several tag-scoped markers must not suppress one + marker's guardrail and then route under a different marker's policy.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-us"}, tags=["us"]), + ] + ) + + eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + + assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) + assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) + + def test_untagged_marker_matches_any_request(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + policy = policy_for_model( + llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + ) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self): + """Regression: a "us" request must not fall back to an "eu" marker's policy.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-default"}), + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) + + def test_no_untagged_fallback_means_no_policy(self): + """No matching marker means no policy, not an unrelated slice's compression.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) + assert ( + policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + ) + + def test_tag_scoped_marker_takes_precedence_over_untagged(self): + """Regression: when multiple markers exist, the tag-scoped one the request + actually matches should be used, not the first untagged one.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-untagged"}), + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) + + +class _RecordingCompressionGuardrail(CustomGuardrail): + """A guardrail whose apply_guardrail marks every text message as compressed.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.request_data_seen: list[dict] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.request_data_seen.append(request_data) + structured_messages = inputs.get("structured_messages") or [] + compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages] + return {**inputs, "structured_messages": compressed} + + +@pytest.fixture +def registered_guardrail(monkeypatch): + import litellm + from litellm.proxy.guardrails import guardrail_registry + + # Registered under a compression provider name: both hops refuse a name that does + # not resolve to one, so a bare callback would (correctly) never be used. + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail) + guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + +class _NonCompressionGuardrail(CustomGuardrail): + """A guardrail that is not a compression provider, e.g. a PII or content filter.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.called = False + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.called = True + return inputs + + +class TestArmPreCall: + @pytest.mark.asyncio + async def test_no_router_is_noop(self): + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=None) + assert "metadata" not in data + + @pytest.mark.asyncio + async def test_no_policy_does_not_create_metadata_bucket(self): + router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=router) + assert "metadata" not in data + assert "litellm_metadata" not in data + + @pytest.mark.asyncio + async def test_policy_suppresses_active_compression_guardrails(self, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem( + guardrail_registry.guardrail_class_registry, "fake-provider", _RecordingCompressionGuardrail + ) + monkeypatch.setattr( + "litellm.proxy.guardrails.auto_router_compression.COMPRESSION_GUARDRAIL_PROVIDERS", + frozenset({"fake-provider"}), + ) + import litellm + + always_on = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + litellm.logging_callback_manager.add_litellm_callback(always_on) + try: + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=router) + assert auto_router_compression.suppressed_compression_guardrails() == frozenset({"always-on-compression"}) + # Suppression state must never ride along in metadata: that reaches spend + # logs the caller can read, and anything there is replayable. + assert "always-on-compression" not in json.dumps(data.get("metadata", {})) + assert always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is False + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + + @pytest.mark.asyncio + async def test_suppression_state_never_enters_request_metadata(self): + """Regression (security): metadata reaches spend logs, so a suppression list + there is one a caller could read back and replay to disable a guardrail.""" + guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + import litellm + + litellm.logging_callback_manager.add_litellm_callback(guardrail) + try: + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=router) + assert "suppress" not in json.dumps(data).lower() + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + @pytest.mark.asyncio + async def test_model_side_guardrail_is_requested_even_when_not_default_on(self, monkeypatch): + import litellm + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail) + active = _RecordingCompressionGuardrail(guardrail_name="headroom-b") + litellm.logging_callback_manager.add_litellm_callback(active) + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-b", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + try: + await arm_pre_call(data=data, llm_router=router) + assert data["metadata"]["guardrails"] == ["headroom-b"] + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(active) + + @pytest.mark.asyncio + async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): + """Regression (security): arm_pre_call runs before the guardrails, so any copy it + kept would be pre-masking text that routing then POSTs to an external service.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} + + await arm_pre_call(data=data, llm_router=router) + + assert "123-45-6789" not in json.dumps(data.get("metadata", {})) + assert not hasattr(auto_router_compression, "_routing_messages_snapshot") + + +class TestMessagesForRouting: + @pytest.mark.asyncio + async def test_no_policy_returns_none(self): + assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_routing_none_with_no_model_compression_returns_none(self): + """Nothing compressed either hop, so the caller's own messages are already right.""" + policy = AutoRouterCompressionPolicy(routing=None, model=None) + assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self): + """No uncompressed copy survives the model hop, and keeping one would mean + retaining the pre-masking text. Routing reads what it has.""" + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] + + assert await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_routes_on_the_uncompressed_messages(self): + policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None) + messages = [{"role": "user", "content": "hi"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result == messages + + @pytest.mark.asyncio + async def test_compresses_via_the_named_guardrail(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hello world"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}] + + @pytest.mark.asyncio + async def test_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail): + """Regression (security): routing POSTs its input out, so it must read what the + earlier guardrails left behind, not a pre-masking copy.""" + policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") + masked = [{"role": "user", "content": "my ssn is [REDACTED]"}] + + result = await messages_for_routing(policy=policy, messages=masked, request_kwargs={}) + + assert result == [{"role": "user", "content": "[COMPRESSED] my ssn is [REDACTED]"}] + assert registered_guardrail.request_data_seen[0]["messages"] == masked + + @pytest.mark.asyncio + async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch): + """Regression (security): naming an ordinary guardrail must not turn the routing + hop into a way to ship prompts to whatever service backs it.""" + import litellm + + other = _NonCompressionGuardrail(guardrail_name="pii-filter") + litellm.logging_callback_manager.add_litellm_callback(other) + try: + policy = AutoRouterCompressionPolicy(routing="pii-filter", model=None) + messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] + + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + + assert other.called is False + assert result == messages + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(other) + + @pytest.mark.asyncio + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): + """Regression: a guardrail writes stats onto the request_data it is given, so + passing the caller's own would double-count into extract_compression_saved_tokens.""" + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hi"}] + request_kwargs = {"metadata": {}} + await messages_for_routing(policy=policy, messages=messages, request_kwargs=request_kwargs) + assert registered_guardrail.request_data_seen[0] is not request_kwargs + assert request_kwargs == {"metadata": {}} diff --git a/tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py b/tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py new file mode 100644 index 00000000000..ff4fbbfb695 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py @@ -0,0 +1,117 @@ +import pytest + +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + ActiveSSOProvider, + active_sso_provider, + id_jag_assertion_capture_gap, + id_jag_assertion_capture_gap_at_startup, +) + +_SSO_ENV_VARS = ( + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", +) + + +@pytest.fixture(autouse=True) +def _isolated_sso_env(monkeypatch): + """Every SSO selector is read from the process environment, so a value left behind by + another test would silently decide this one's answer.""" + for name in _SSO_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +class TestActiveSSOProviderMirrorsTheCallback: + """The gap warning is only as good as its agreement with the branch the login callback + actually takes, so provider selection is asserted branch by branch, including the + precedence that makes a co-configured generic client unreachable.""" + + def test_google_client_id_selects_google(self, monkeypatch): + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + assert active_sso_provider() is ActiveSSOProvider.google + + def test_microsoft_client_id_selects_microsoft(self, monkeypatch): + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid") + assert active_sso_provider() is ActiveSSOProvider.microsoft + + def test_generic_client_id_selects_generic(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert active_sso_provider() is ActiveSSOProvider.generic + + def test_saml_metadata_selects_saml(self, monkeypatch): + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert active_sso_provider() is ActiveSSOProvider.saml + + def test_nothing_configured_selects_none(self): + assert active_sso_provider() is ActiveSSOProvider.none + + def test_google_outranks_a_co_configured_generic_client(self, monkeypatch): + """The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and + no assertion is captured; reporting generic would clear a gap that is still open.""" + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert active_sso_provider() is ActiveSSOProvider.google + + def test_microsoft_outranks_a_co_configured_generic_client(self, monkeypatch): + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid") + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert active_sso_provider() is ActiveSSOProvider.microsoft + + def test_generic_outranks_saml(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert active_sso_provider() is ActiveSSOProvider.generic + + +class TestIdJagAssertionCaptureGap: + def test_generic_oidc_has_no_gap(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert id_jag_assertion_capture_gap() is None + + @pytest.mark.parametrize( + "env_var, provider_label", + [ + ("GOOGLE_CLIENT_ID", "google"), + ("MICROSOFT_CLIENT_ID", "microsoft"), + ("SAML_IDP_METADATA_URL", "saml"), + ], + ) + def test_non_capturing_provider_is_named_with_the_remedy(self, monkeypatch, env_var, provider_label): + monkeypatch.setenv(env_var, "configured") + gap = id_jag_assertion_capture_gap() + assert gap is not None + assert provider_label in gap + assert "GENERIC_CLIENT_ID" in gap + + def test_no_sso_configured_reports_a_gap(self): + gap = id_jag_assertion_capture_gap() + assert gap is not None + assert "no SSO provider is configured" in gap + + def test_google_beside_generic_still_reports_a_gap(self, monkeypatch): + """The precedence trap in operator terms: adding a generic client id without removing + GOOGLE_CLIENT_ID does not fix the deployment, so the gap must not clear.""" + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + gap = id_jag_assertion_capture_gap() + assert gap is not None + assert "google" in gap + + +class TestIdJagAssertionCaptureGapAtStartup: + def test_no_provider_at_startup_is_not_yet_a_gap(self): + assert id_jag_assertion_capture_gap_at_startup() is None + + def test_google_provider_at_startup_reports_the_capture_gap(self, monkeypatch): + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + startup_gap = id_jag_assertion_capture_gap_at_startup() + callback_gap = id_jag_assertion_capture_gap() + assert startup_gap is not None + assert startup_gap == callback_gap + + def test_generic_provider_at_startup_has_no_gap(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert id_jag_assertion_capture_gap_at_startup() is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index adab3538b58..71ff7de89b0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2,6 +2,7 @@ import os import sys import types import json +import logging from contextlib import ExitStack from datetime import datetime, timedelta from types import SimpleNamespace @@ -3840,6 +3841,155 @@ class TestAddMCPServerAtomicity: mock_manager.reload_servers_from_database.assert_not_awaited() +class TestIdJagRegistrationWarnsAboutTheSSOGap: + """An `oauth2_id_jag` server only ever works when the login path captures an IdP identity + assertion, and only the generic OIDC arm does. Registering one under Google or Microsoft + succeeds and then fails for every user on every call, so the mismatch has to be said at + registration time, while the admin is still looking at the configuration.""" + + @staticmethod + def _clear_sso_env(monkeypatch): + for name in ( + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + ): + monkeypatch.delenv(name, raising=False) + + @staticmethod + def _id_jag_warnings(caplog) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "oauth2_id_jag" in record.getMessage() + ] + + @staticmethod + def _server_record(auth_type) -> LiteLLM_MCPServerTable: + record = generate_mock_mcp_server_db_record(server_id="ema-1", alias="ema") + record.auth_type = auth_type + return record + + async def _run_create(self, monkeypatch, provider_env, auth_type, caplog): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + self._clear_sso_env(monkeypatch) + for name, value in provider_env.items(): + monkeypatch.setenv(name, value) + + mock_manager = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ( + patch( # test-quality-ok: endpoint test stubs the Prisma client lookup + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint test stubs MCP server creation + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=self._server_record(auth_type)), + ), + patch( # test-quality-ok: endpoint reads the global MCP manager + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await add_mcp_server( + payload=NewMCPServerRequest( + alias="ema", + url="https://ema.example.com/mcp", + transport=MCPTransport.http, + ), + user_api_key_dict=generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ), + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "provider_env, expected_fragment", + [ + ({"GOOGLE_CLIENT_ID": "cid"}, "google"), + ({"MICROSOFT_CLIENT_ID": "cid"}, "microsoft"), + ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), + ({}, "no SSO provider is configured"), + ], + ) + async def test_create_warns_under_a_provider_that_captures_nothing( + self, monkeypatch, caplog, provider_env, expected_fragment + ): + await self._run_create(monkeypatch, provider_env, MCPAuth.oauth2_id_jag, caplog) + warnings = self._id_jag_warnings(caplog) + assert len(warnings) == 1 + assert expected_fragment in str(warnings[0]) + assert "ema-1" in str(warnings[0]) + + @pytest.mark.asyncio + async def test_create_is_silent_under_generic_oidc(self, monkeypatch, caplog): + await self._run_create(monkeypatch, {"GENERIC_CLIENT_ID": "cid"}, MCPAuth.oauth2_id_jag, caplog) + assert self._id_jag_warnings(caplog) == [] + + @pytest.mark.asyncio + async def test_create_is_silent_for_other_auth_types(self, monkeypatch, caplog): + """Nothing but the id_jag arm sources credentials from a stored SSO assertion, so no + other server registered under Google has anything to warn about.""" + await self._run_create(monkeypatch, {"GOOGLE_CLIENT_ID": "cid"}, MCPAuth.api_key, caplog) + assert self._id_jag_warnings(caplog) == [] + + @pytest.mark.asyncio + async def test_update_to_id_jag_warns(self, monkeypatch, caplog): + """Switching an existing server onto id_jag opens the same gap a create does.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + edit_mcp_server, + ) + + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + + mock_manager = MagicMock() + mock_manager.update_server = AsyncMock() + mock_manager.reload_servers_from_database = AsyncMock() + with ( + patch( # test-quality-ok: endpoint test stubs the Prisma client lookup + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint test stubs the MCP server lookup + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=self._server_record(MCPAuth.api_key)), + ), + patch( # test-quality-ok: endpoint test stubs MCP server updates + "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", + AsyncMock(return_value=self._server_record(MCPAuth.oauth2_id_jag)), + ), + patch( # test-quality-ok: endpoint test stubs credential cleanup + "litellm.proxy.management_endpoints.mcp_management_endpoints.purge_user_oauth_credentials_for_server", + AsyncMock(return_value=0), + ), + patch( # test-quality-ok: endpoint reads the global MCP manager + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await edit_mcp_server( + payload=UpdateMCPServerRequest(server_id="ema-1", auth_type=MCPAuth.oauth2_id_jag), + user_api_key_dict=generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ), + ) + + warnings = self._id_jag_warnings(caplog) + assert len(warnings) == 1 + assert "google" in str(warnings[0]) + + class TestHealthCheckServers: """Test suite for health check servers endpoint""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 5dfff53f7c3..8d8bc15f9be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1,17 +1,16 @@ import asyncio import json +import logging import os -from contextlib import asynccontextmanager +from contextlib import ExitStack, asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException, Request -from litellm._uuid import uuid - - import litellm +from litellm._uuid import uuid from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO @@ -1615,8 +1614,8 @@ async def test_get_generic_sso_response_with_empty_headers(): async def test_get_generic_sso_response_includes_token_claims_when_enabled(monkeypatch): import jwt as pyjwt - from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) mock_jwt_handler = MagicMock(spec=JWTHandler) @@ -2321,10 +2320,10 @@ class TestCustomUISSO: async def test_handle_custom_ui_sso_sign_in_success(self): """Test successful custom UI SSO sign-in with valid headers""" from fastapi_sso.sso.base import OpenID - from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler # Mock request with custom headers @@ -2400,6 +2399,7 @@ class TestCustomUISSO: from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler mock_request = MagicMock(spec=Request) @@ -2436,10 +2436,10 @@ class TestCustomUISSO: and its methods are called with the correct parameters """ from fastapi_sso.sso.base import OpenID - from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler # Create a real custom handler class instance @@ -8167,6 +8167,128 @@ async def test_debug_sso_callback_handles_missing_raw_response(): assert "user@example.com" in body +# ── The debug page is where an operator lands when ID-JAG is failing ────────── + +_GOOGLE_DEBUG_CLIENT_ID = "debug-google-client-id" +_GENERIC_DEBUG_CLIENT_ID = "debug-generic-client-id" + + +async def _render_debug_page(provider_env, id_jag_registered, force_inert=False): + """Drive /sso/debug/callback and return the raw response body.""" + from litellm.proxy.management_endpoints.ui_sso import GoogleSSOHandler, debug_sso_callback + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://proxy.example.com/" + mock_request.cookies = {} + mock_request.query_params = {} + + parsed = {"sub": "user_123", "email": "u@example.com"} + + async def fake_generic(**kwargs): + return parsed, {"sub": "user_123"}, {"scope": "openid"}, None + + async def fake_google(**kwargs): + return parsed + + stack = [ + patch.dict(os.environ, provider_env, clear=False), + patch( # test-quality-ok: endpoint test stubs the upstream generic IdP boundary + "litellm.proxy.management_endpoints.ui_sso.get_generic_sso_response", side_effect=fake_generic + ), + patch.object( # test-quality-ok: endpoint test stubs the upstream Google IdP boundary + GoogleSSOHandler, "get_google_callback_response", side_effect=fake_google + ), + patch( # test-quality-ok: debug endpoint reads this module global without an injection seam + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=id_jag_registered), + ), + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals + ] + if force_inert: + stack.append( + patch( # test-quality-ok: force-inert reference isolates the endpoint's pre-change response + "litellm.proxy.management_endpoints.ui_sso.warn_if_id_jag_capture_gap", + AsyncMock(return_value=None), + ) + ) + + with ExitStack() as es: + for ctx in stack: + es.enter_context(ctx) + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + if var not in provider_env: + os.environ.pop(var, None) + response = await debug_sso_callback(mock_request) + + return response.body.decode() + + +@pytest.mark.asyncio +async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "google" in warnings[0] + assert "GENERIC_CLIENT_ID" in warnings[0] + assert "id_jag" not in body + assert "GENERIC_CLIENT_ID" not in body + + +@pytest.mark.asyncio +async def test_debug_page_is_byte_identical_when_the_provider_captures(): + """A deployment with no gap must get the page it got before this change, to the byte. The + comparison is against the endpoint with the diagnostic forced inert, not against a guess.""" + with_feature = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True + ) + pre_change = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, + id_jag_registered=True, + force_inert=True, + ) + + assert with_feature == pre_change + assert "id_jag" not in with_feature + + +@pytest.mark.asyncio +async def test_debug_page_is_byte_identical_when_no_id_jag_server_is_registered(): + """Most deployments run Google SSO and no id_jag server at all; their debug page must not + grow an ID-JAG section about a feature they do not use.""" + with_feature = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False + ) + pre_change = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, + id_jag_registered=False, + force_inert=True, + ) + + assert with_feature == pre_change + assert "id_jag" not in with_feature + + +@pytest.mark.asyncio +async def test_debug_page_survives_a_store_outage(monkeypatch, caplog): + """The page's job is to render claims; an unreachable MCP table must cost it the annotation, + not the page.""" + from litellm.proxy.management_endpoints.ui_sso import warn_if_id_jag_capture_gap + + monkeypatch.setenv("GOOGLE_CLIENT_ID", _GOOGLE_DEBUG_CLIENT_ID) + retention_check = AsyncMock(side_effect=Exception("db down")) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + assert await warn_if_id_jag_capture_gap(retention_enabled=retention_check) is None + + retention_check.assert_awaited_once() + + assert _id_jag_gap_warnings(caplog) == [] + + async def _render_legacy_login_page(env_overrides, general_settings): from litellm.proxy.management_endpoints.ui_sso import google_login @@ -8261,8 +8383,8 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): that /sso/key/generate enforces; the ACS re-checks it after validating the assertion, so the entitlement DB query never runs on unvalidated input.""" from litellm.proxy._types import ProxyException - from litellm.proxy.management_endpoints.ui_sso import saml_callback from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import saml_callback call_order: list[str] = [] @@ -8681,6 +8803,248 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): assert response.status_code == 200 +def _id_jag_gap_warnings(caplog) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "oauth2_id_jag" in record.getMessage() + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_env, expected_fragment", + [ + ({"GOOGLE_CLIENT_ID": "cid"}, "google"), + ({"MICROSOFT_CLIENT_ID": "cid", "MICROSOFT_TENANT": "t"}, "microsoft"), + ({}, "no SSO provider is configured"), + ], +) +async def test_uncaptured_assertion_warns_when_an_id_jag_server_is_registered( + monkeypatch, caplog, provider_env, expected_fragment +): + """A provider with no capture path leaves ID-JAG permanently broken, and the only place + that is knowable is the login itself; without this line the operator sees nothing at all.""" + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(name, raising=False) + for name, value in provider_env.items(): + monkeypatch.setenv(name, value) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True)) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert expected_fragment in str(warnings[0]) + + +@pytest.mark.asyncio +async def test_generic_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog): + """Generic OIDC has a capture path, so there is no configuration gap to report; the login + still handed the id_jag arm nothing, and that must not pass silently.""" + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("GENERIC_CLIENT_ID", "cid") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True)) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "no usable id_token" in str(warnings[0]) + + +@pytest.mark.asyncio +async def test_no_warning_when_the_assertion_was_captured(monkeypatch, caplog): + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + assertion_from_sso_login, + ) + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + assertion = assertion_from_sso_login(_ema_id_token(), None) + assert assertion is not None + + retention_mock = AsyncMock(return_value=True) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(assertion, retention_enabled=retention_mock) + + assert _id_jag_gap_warnings(caplog) == [] + retention_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_no_warning_when_no_id_jag_server_is_registered(monkeypatch, caplog): + """Most deployments never register one; a warning about ID-JAG on every login there would + be pure noise and would train operators to ignore it.""" + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=False)) + + assert _id_jag_gap_warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_store_outage_does_not_break_the_login(monkeypatch, caplog): + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + assert ( + await warn_if_id_jag_assertion_uncaptured( + None, retention_enabled=AsyncMock(side_effect=Exception("db down")) + ) + is None + ) + + assert _id_jag_gap_warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog): + """Wiring: the browser login path must reach the diagnostic, not just define it.""" + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.cookies = {} + + with ( + patch( # test-quality-ok: endpoint test stubs the Prisma client lookup + "litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock() + ), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.user_custom_sso", None), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: endpoint reads proxy globals + patch( # test-quality-ok: endpoint test stubs key generation at its module boundary + "litellm.proxy.proxy_server.generate_key_helper_fn", + AsyncMock(return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}), + ), + patch( # test-quality-ok: endpoint test stubs the user database lookup + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=None), + ), + patch( # test-quality-ok: endpoint test stubs the admin database lookup + "litellm.proxy.management_endpoints.ui_sso.check_and_update_if_proxy_admin_id", + AsyncMock(return_value="internal_user"), + ), + patch( # test-quality-ok: endpoint test stubs assertion persistence + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + patch( # test-quality-ok: endpoint reads this module global without an injection seam + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=True), + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await SSOAuthenticationHandler.get_redirect_response_from_openid( + result=CustomOpenID( + id="raw-idp-subject", + email="u@example.com", + first_name="U", + last_name="Ser", + display_name="U Ser", + provider="google", + team_ids=[], + user_role=None, + ), + request=mock_request, + received_response=None, + generic_client_id=None, + ui_access_mode=None, + access_token_payload=None, + jwt_handler=None, + sso_assertion=None, + ) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "google" in str(warnings[0]) + + +@pytest.mark.asyncio +async def test_cli_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog): + """Wiring: the CLI login path shares the gap, so it must share the diagnostic.""" + from litellm.proxy.management_endpoints.ui_sso import ( + _complete_cli_sso_callback_session, + ) + + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "cid") + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + + user_info = MagicMock() + user_info.user_id = "cli-user-id" + user_info.user_role = "internal_user" + user_info.models = [] + user_info.teams = [] + + with ( + patch( # test-quality-ok: endpoint test stubs the user database lookup + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=user_info), + ), + patch( # test-quality-ok: endpoint test stubs CLI team lookup + "litellm.proxy.management_endpoints.ui_sso.fetch_cli_sso_team_details", + AsyncMock(return_value=[]), + ), + patch( # test-quality-ok: endpoint test stubs attribution metadata + "litellm.proxy.management_endpoints.ui_sso.build_cli_sso_attribution_metadata", + return_value={}, + ), + patch( # test-quality-ok: endpoint test stubs assertion persistence + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + patch( # test-quality-ok: endpoint reads this module global without an injection seam + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=True), + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await _complete_cli_sso_callback_session( + request=mock_request, + key="cli-login-id", + flow={}, + result={"sub": "raw-idp-subject"}, + parsed_openid_result={ + "user_id": "raw-idp-subject", + "user_email": "u@example.com", + "user_role": None, + }, + user_defined_values=None, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + cli_sso_session_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + sso_assertion=None, + ) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "microsoft" in str(warnings[0]) + + def _cli_callback_kwargs(flow): return { "request": _cli_callback_request(), diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f7fe6ad9d39..96d9b0c5a26 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -376,6 +376,75 @@ class TestProxyBaseLLMRequestProcessing: assert "litellm_logging_obj" not in persisted_body json.dumps(persisted_body) + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): + """arm_pre_call must run before pre_call_hook: an auto router's own compression + policy has to be in `data["metadata"]` (naming the model-side guardrail so it + runs even if it isn't default_on) by the time guardrails see the request.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails import guardrail_registry + + # The model hop is only armed for a name that resolves to an active compression + # guardrail, so arming it has to have a real one to resolve to. + class _FakeCompressionGuardrail(CustomGuardrail): + pass + + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _FakeCompressionGuardrail) + active_guardrail = _FakeCompressionGuardrail(guardrail_name="headroom-model") + litellm.logging_callback_manager.add_litellm_callback(active_guardrail) + + processing_obj = ProxyBaseLLMRequestProcessing(data={}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + + seen_metadata: dict = {} + + async def mock_pre_call_hook(user_api_key_dict, data, call_type): + seen_metadata.update(data.get("metadata") or {}) + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + + fake_llm_router = MagicMock() + fake_llm_router.get_model_list.return_value = [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-model", + }, + } + ] + mock_proxy_config = MagicMock(spec=ProxyConfig) + mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + + try: + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=fake_llm_router, + ) + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(active_guardrail) + + assert seen_metadata.get("guardrails") == ["headroom-model"] + def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): mock_set_active_span_tag = MagicMock(return_value=True) import litellm.proxy.dd_span_tagger diff --git a/tests/test_litellm/router_utils/test_pattern_match_deployments.py b/tests/test_litellm/router_utils/test_pattern_match_deployments.py index 795d448ef5f..f9d9345cd26 100644 --- a/tests/test_litellm/router_utils/test_pattern_match_deployments.py +++ b/tests/test_litellm/router_utils/test_pattern_match_deployments.py @@ -2,8 +2,10 @@ from __future__ import annotations +from unittest.mock import Mock + from litellm.router_utils import pattern_match_deployments -from litellm.router_utils.pattern_match_deployments import PatternMatchRouter +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils def _wildcard_deployment(model_name: str) -> dict: @@ -76,3 +78,31 @@ def test_get_pattern_still_resolves_unqualified_names(monkeypatch): router = PatternMatchRouter() router.add_pattern("openai/*", _wildcard_deployment("openai/*")) assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"] + + +class _CountingPatternUtils(PatternUtils): + sorted_patterns = staticmethod(Mock(wraps=PatternUtils.sorted_patterns)) + + +def test_route_never_sorts_and_the_most_specific_pattern_still_wins_after_registry_changes(): + """Regression for LIT-6886: the auth layer walks the wildcard registry for every request, so an + unmatched model name (an invalid-model 403) re-sorted every pattern by specificity per request and + a burst of rejections saturated the worker CPU. Lookups must not sort; adding a pattern or removing + a deployment must still leave the most specific pattern winning.""" + router = PatternMatchRouter(pattern_utils=_CountingPatternUtils) + router.add_pattern("openai/*", _wildcard_deployment("openai/*")) + router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*")) + router.add_pattern("openai/gpt-*", {"model_name": "openai/gpt-*", "litellm_params": {"model": "azure/gpt-*"}}) + sorts_after_setup = _CountingPatternUtils.sorted_patterns.call_count + + for _ in range(3): + assert router.route("does-not-exist") is None + assert _matched_models(router.route("openai/gpt-4o")) == ["azure/gpt-4o"] + assert _matched_models(router.route("openai/o3")) == ["openai/o3"] + assert _CountingPatternUtils.sorted_patterns.call_count == sorts_after_setup + + router.add_pattern("openai/*", {**_wildcard_deployment("openai/*"), "model_info": {"id": "id-1"}}) + assert len(_matched_models(router.route("openai/o3"))) == 2 + router.remove_deployment("id-1") + assert _matched_models(router.route("openai/gpt-4o")) == ["azure/gpt-4o"] + assert _matched_models(router.route("openai/o3")) == ["openai/o3"] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f779abaf870..2df3c2d9339 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5,6 +5,7 @@ import json import logging import os import threading +from datetime import datetime from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -19,7 +20,9 @@ import respx import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, @@ -10028,6 +10031,215 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestAutoRouterCompressionDecoupling: + """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` + decouple what the routing decision sees from what the model call sees. The one + assertion that must hold under any mutation: the strategy can be routed on + compressed text while the caller's own `messages` list - the one that would reach + the model - is never touched.""" + + class _RecordingStrategy: + """Echoes back whatever `messages` it was handed, like every real strategy does.""" + + def __init__(self): + self.received_messages: list[dict] | None = None + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + self.received_messages = messages + return PreRoutingHookResponse(model="gemini-flash", messages=messages) + + class _CompressingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.call_count = 0 + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.call_count += 1 + structured_messages = inputs.get("structured_messages") or [] + compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages] + return {**inputs, "structured_messages": compressed} + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + def _router(self, marker_litellm_params: dict) -> tuple[litellm.Router, "_RecordingStrategy"]: + from litellm.types.router import TaggedPreRoutingStrategy + + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + router = litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + **marker_litellm_params, + }, + }, + { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + }, + ], + ) + for name in ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers"): + setattr(router, name, {}) + strategy = self._RecordingStrategy() + router.complexity_routers = {"smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=strategy)]} + return router, strategy + + @pytest.fixture + def registered_guardrail(self, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + # Registered under a compression provider name: both hops refuse a name that + # does not resolve to one, so a bare callback would never be used. + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", self._CompressingGuardrail) + guardrail = self._CompressingGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + @pytest.mark.asyncio + async def test_routing_side_compression_never_reaches_the_caller_messages(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "none", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + assert response.messages == original_messages + + @pytest.mark.asyncio + async def test_model_side_compression_alone_leaves_routing_uncompressed(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy(self, registered_guardrail): + """Routing asked for no compression while the model hop compressed, so the only + messages left are that guardrail's output and the strategy classifies on them. + + Keeping a pre-compression copy to classify on instead is what this deliberately + gives up: that copy is taken before the pre-call guardrails run, so it still + holds whatever a masking guardrail exists to strip, and routing-side compression + POSTs its input to an external service.""" + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) + + assert strategy.received_messages == model_compressed + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + @pytest.mark.asyncio + async def test_same_compression_still_compresses_routing_when_nothing_armed_it(self, registered_guardrail): + """Regression: only the proxy calls arm_pre_call. Used through the SDK, nothing + arms the model-side guardrail and nothing has compressed anything, so reusing a + model-hop result that was never produced would serve the request with no + compression on either hop, silently ignoring the configuration.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + uncompressed = self._messages() + assert auto_router_compression.model_hop_compression_armed() is False + + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=uncompressed + ) + + assert strategy.received_messages != uncompressed + assert registered_guardrail.call_count == 1 + + async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): + """The same/different distinction exists so a shared choice does not pay for + compression twice: by the time the router runs, `messages` already reflects + whatever the ordinary pre-call guardrail pipeline did for the model call, so + the routing decision must reuse it rather than calling the guardrail again.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + # Stands in for what the proxy's ordinary pre-call guardrail pipeline would + # have already produced for the model call, since `auto_router_model_compression` + # names a guardrail: the router never triggers that pipeline itself. + already_compressed_messages = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + # arm_pre_call is what would have armed that guardrail, and only the proxy calls + # it; the reuse below is conditional on it having run. + armed = auto_router_compression._model_hop_armed.set(True) + + try: + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages + ) + finally: + auto_router_compression._model_hop_armed.reset(armed) + + assert strategy.received_messages == already_compressed_messages + assert response.messages == already_compressed_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_no_policy_is_fully_unaffected(self, registered_guardrail): + router, strategy = self._router({}) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages is original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + +@pytest.mark.usefixtures("local_model_cost_map") + @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -12999,6 +13211,136 @@ async def test_router_retry_policy_controls_upstream_attempt_count( assert upstream.call_count == expected_upstream_calls +def _make_failure_logging_obj(): + return LiteLLMLogging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="lit-6960", + function_id="f", + ) + + +async def _assert_router_failure_logging_is_coordinated(logging_obj, trigger, expected_exception): + """The sync failure_handler must not start until async_failure_handler has finished on the shared logging_obj.""" + events: list[str] = [] + sync_done = threading.Event() + + async def _async_failure(*args, **kwargs): + events.append("async_start") + await asyncio.sleep(0.05) + events.append("async_end") + + def _sync_failure(*args, **kwargs): + events.append("sync_start") + sync_done.set() + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", side_effect=_sync_failure), + patch.object(logging_obj, "_should_run_sync_failure_callbacks_for_async_calls", return_value=True), + ): + with pytest.raises(expected_exception): + await trigger() + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + await asyncio.gather(*pending) + assert await asyncio.to_thread(sync_done.wait, 5), "failure_handler never ran" + + assert events == ["async_start", "async_end", "sync_start"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hook_error", + [ + litellm.RateLimitError(message="rpm exceeded", llm_provider="openai", model="gpt-5.6"), + RuntimeError("pre call check blew up"), + ], +) +async def test_async_routing_strategy_pre_call_checks_failure_logging_is_coordinated(hook_error): + class _RaisingPreCallCheck(CustomLogger): + async def async_pre_call_check(self, deployment, parent_otel_span): + raise hook_error + + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + deployment = router.model_list[0] + logging_obj = _make_failure_logging_obj() + + with patch.object(litellm, "callbacks", [_RaisingPreCallCheck()]): # test-quality-ok: router reads this global + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_routing_strategy_pre_call_checks( + deployment=deployment, parent_otel_span=None, logging_obj=logging_obj + ), + type(hook_error), + ) + + +@pytest.mark.asyncio +async def test_async_callback_filter_deployments_failure_logging_is_coordinated(): + class _RaisingFilter(CustomLogger): + async def async_filter_deployments(self, *args, **kwargs): + raise RuntimeError("filter blew up") + + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + logging_obj = _make_failure_logging_obj() + + with patch.object(litellm, "callbacks", [_RaisingFilter()]): # test-quality-ok: router reads this global + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_callback_filter_deployments( + model="gpt-5.6", + healthy_deployments=router.model_list, + messages=None, + parent_otel_span=None, + request_kwargs={}, + logging_obj=logging_obj, + ), + RuntimeError, + ) + + +@pytest.mark.asyncio +async def test_async_get_available_deployment_failure_logging_is_coordinated(): + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + logging_obj = _make_failure_logging_obj() + + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_get_available_deployment( + model="model-that-is-not-configured", + request_kwargs={"litellm_logging_obj": logging_obj}, + messages=[{"role": "user", "content": "hi"}], + ), + litellm.BadRequestError, + ) + + +@pytest.mark.asyncio +async def test_async_get_available_deployment_for_pass_through_failure_logging_is_coordinated(): + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + logging_obj = _make_failure_logging_obj() + + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_get_available_deployment_for_pass_through( + model="gpt-5.6", + request_kwargs={"litellm_logging_obj": logging_obj}, + ), + litellm.BadRequestError, + ) + + class _InFlightTracker: def __init__(self) -> None: self.current = 0 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 481e3591ce9..225134b4e2b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22181 + "limit": 22180 }, "LIT002": { "limit": 26745 @@ -9,7 +9,7 @@ "limit": 261 }, "LIT004": { - "limit": 40 + "limit": 38 }, "LIT005": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16464 + "limit": 16462 }, "LIT011": { "limit": 5506 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx index 9f27daab6b3..4e0590df72d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx @@ -1,36 +1,61 @@ -import { render, screen, waitFor } from "@testing-library/react"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { describe, expect, it, vi } from "vitest"; +import { type UrlUpdateEvent } from "nuqs/adapters/testing"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; import GuardrailsMonitorView from "./GuardrailsMonitorView"; import * as networking from "@/components/networking"; +import { renderWithProviders, screen, testQueryClient, waitFor } from "@/../tests/test-utils"; vi.mock("@/components/networking", () => ({ getGuardrailsUsageOverview: vi.fn(), + getGuardrailsUsageDetail: vi.fn(), + getGuardrailsUsageLogs: vi.fn(), formatDate: vi.fn((d: Date) => d.toISOString().slice(0, 10)), })); -const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview); +vi.mock("@/components/GuardrailsMonitor/LogViewer", () => ({ + LogViewer: ({ guardrailName }: { guardrailName: string }) =>
{guardrailName}
, +})); -function wrapper({ children }: { children: React.ReactNode }) { - const queryClient = new QueryClient({ - defaultOptions: { - queries: { retry: false }, - }, - }); - return {children}; -} +const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview); +const mockGetGuardrailsUsageDetail = vi.mocked(networking.getGuardrailsUsageDetail); +const mockGetGuardrailsUsageLogs = vi.mocked(networking.getGuardrailsUsageLogs); + +const emptyOverview = { rows: [], chart: [], totalRequests: 0, totalBlocked: 0, passRate: 100 }; + +const piiRow = { + id: "gr-pii", + name: "PII Guard", + type: "pii", + provider: "LiteLLM", + requestsEvaluated: 10, + failRate: 10, + status: "healthy" as const, + trend: "stable" as const, +}; + +const piiDetail = { + guardrail_name: "PII Guard", + description: "", + status: "healthy", + provider: "LiteLLM", + type: "pii", + requestsEvaluated: 10, + failRate: 10, + avgScore: 0.5, + avgLatency: 20, +}; describe("GuardrailsMonitorView", () => { - it("should render overview and fetch guardrails usage when accessToken is provided", async () => { - mockGetGuardrailsUsageOverview.mockResolvedValue({ - rows: [], - chart: [], - totalRequests: 0, - totalBlocked: 0, - passRate: 100, - }); + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + mockGetGuardrailsUsageOverview.mockResolvedValue(emptyOverview); + mockGetGuardrailsUsageDetail.mockResolvedValue(piiDetail); + mockGetGuardrailsUsageLogs.mockResolvedValue({ logs: [], total: 0 }); + }); - render(, { wrapper }); + it("should render overview and fetch guardrails usage when accessToken is provided", async () => { + renderWithProviders(); expect(await screen.findByRole("heading", { name: /Guardrails Monitor/i })).toBeInTheDocument(); await waitFor(() => { @@ -39,7 +64,54 @@ describe("GuardrailsMonitorView", () => { }); it("should render without crashing when accessToken is null", async () => { - render(, { wrapper }); + renderWithProviders(); expect(await screen.findByRole("heading", { name: /Guardrails Monitor/i })).toBeInTheDocument(); }); + + describe("guardrail detail deep link (?guardrail=)", () => { + it("should open the detail view directly from a ?guardrail= deep link", async () => { + renderWithProviders(, { searchParams: "?guardrail=gr-pii" }); + + expect(await screen.findByRole("heading", { name: "PII Guard" })).toBeInTheDocument(); + expect(mockGetGuardrailsUsageDetail).toHaveBeenCalledWith( + "test-token", + "gr-pii", + expect.any(String), + expect.any(String), + ); + expect(screen.queryByRole("heading", { name: /Guardrails Monitor/i })).not.toBeInTheDocument(); + }); + + it("should push ?guardrail= as a new history entry when a guardrail is selected", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + mockGetGuardrailsUsageOverview.mockResolvedValue({ ...emptyOverview, rows: [piiRow] }); + renderWithProviders(, { onUrlUpdate }); + + await user.click(await screen.findByRole("button", { name: "PII Guard" })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.get("guardrail")).toBe("gr-pii"); + expect(lastUpdate.options.history).toBe("push"); + expect(await screen.findByRole("heading", { name: "PII Guard" })).toBeInTheDocument(); + }); + + it("should clear ?guardrail= by replacing history when going back to the overview", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + renderWithProviders(, { + searchParams: "?guardrail=gr-pii", + onUrlUpdate, + }); + + await user.click(await screen.findByRole("button", { name: /back to overview/i })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.has("guardrail")).toBe(false); + expect(lastUpdate.options.history).toBe("replace"); + expect(await screen.findByRole("heading", { name: /Guardrails Monitor/i })).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx index a9acf3e6377..f90a46e19e4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx @@ -1,12 +1,11 @@ import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; +import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useMemo, useState } from "react"; import { formatDate } from "@/components/networking"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { GuardrailDetail } from "./GuardrailDetail"; import { GuardrailsOverview } from "./GuardrailsOverview"; -type View = { type: "overview" } | { type: "detail"; guardrailId: string }; - interface GuardrailsMonitorViewProps { accessToken?: string | null; } @@ -16,7 +15,10 @@ const defaultStart = new Date(); defaultStart.setDate(defaultStart.getDate() - 7); export default function GuardrailsMonitorView({ accessToken = null }: GuardrailsMonitorViewProps) { - const [view, setView] = useState({ type: "overview" }); + const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( + "guardrail", + parseAsString.withOptions({ history: "push" }), + ); const initialFrom = useMemo(() => new Date(defaultStart), []); const initialTo = useMemo(() => new Date(defaultEnd), []); @@ -34,11 +36,11 @@ export default function GuardrailsMonitorView({ accessToken = null }: Guardrails }, []); const handleSelectGuardrail = (id: string) => { - setView({ type: "detail", guardrailId: id }); + void setSelectedGuardrailId(id); }; const handleBack = () => { - setView({ type: "overview" }); + void setSelectedGuardrailId(null, { history: "replace" }); }; const dateRangeControl = ( @@ -47,7 +49,7 @@ export default function GuardrailsMonitorView({ accessToken = null }: Guardrails return (
- {view.type === "overview" ? ( + {!selectedGuardrailId ? (
{dateRangeControl}
({ getGuardrailsList: vi.fn(), @@ -15,16 +16,21 @@ vi.mock("./add_guardrail_form", () => ({ vi.mock("./guardrail_table", () => ({ __esModule: true, - default: ({ guardrailsList, onDeleteClick }: any) => ( + default: ({ guardrailsList, onDeleteClick, onGuardrailClick }: any) => (
Mock Guardrail Table
{guardrailsList.length > 0 && ( - + <> + + + )}
), @@ -32,7 +38,12 @@ vi.mock("./guardrail_table", () => ({ vi.mock("./guardrail_info", () => ({ __esModule: true, - default: () =>
Mock Guardrail Info View
, + default: ({ guardrailId, onClose }: { guardrailId: string; onClose: () => void }) => ( +
+
Mock Guardrail Info View {guardrailId}
+ +
+ ), })); vi.mock("./GuardrailTestPlayground", async () => { @@ -112,7 +123,7 @@ describe("GuardrailsPanel", () => { }); it("should render the component", async () => { - render(); + renderWithProviders(); expect(screen.getByText("Guardrails")).toBeInTheDocument(); // Activate the Guardrails tab so its content (including the Add button) is rendered fireEvent.click(screen.getByText("Guardrails")); @@ -120,7 +131,7 @@ describe("GuardrailsPanel", () => { }); it("should delete the clicked guardrail after confirming in the modal", async () => { - render(); + renderWithProviders(); fireEvent.click(screen.getByText("Guardrails")); fireEvent.click(await screen.findByTestId("delete-button")); @@ -139,14 +150,14 @@ describe("GuardrailsPanel", () => { }); it("should mount every tab panel up front so panel state survives tab switches", async () => { - render(); + renderWithProviders(); expect(await screen.findByLabelText("playground draft")).toBeInTheDocument(); expect(screen.getByText("Mock Team Guardrails Tab")).toBeInTheDocument(); }); it("should keep test playground state when switching tabs away and back", async () => { - render(); + renderWithProviders(); fireEvent.click(screen.getByText("Test Playground")); @@ -161,7 +172,7 @@ describe("GuardrailsPanel", () => { }); it("should not delete anything when the modal is cancelled", async () => { - render(); + renderWithProviders(); fireEvent.click(screen.getByText("Guardrails")); fireEvent.click(await screen.findByTestId("delete-button")); @@ -171,4 +182,42 @@ describe("GuardrailsPanel", () => { expect(mockDeleteGuardrailCall).not.toHaveBeenCalled(); }); + + describe("guardrail detail deep link (?guardrail=)", () => { + it("should open the guardrail info view directly from a ?guardrail= deep link", async () => { + renderWithProviders(, { searchParams: "?guardrail=test-guardrail-1" }); + + expect(await screen.findByTestId("guardrail-info-view")).toHaveTextContent("test-guardrail-1"); + expect(screen.queryByText("Mock Guardrail Table")).not.toBeInTheDocument(); + }); + + it("should push ?guardrail= as a new history entry when a guardrail row is clicked", async () => { + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + renderWithProviders(, { onUrlUpdate }); + + fireEvent.click(await screen.findByTestId("open-button")); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.get("guardrail")).toBe("test-guardrail-1"); + expect(lastUpdate.options.history).toBe("push"); + expect(await screen.findByTestId("guardrail-info-view")).toHaveTextContent("test-guardrail-1"); + }); + + it("should clear ?guardrail= by replacing history when the info view is closed", async () => { + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + renderWithProviders(, { + searchParams: "?guardrail=test-guardrail-1", + onUrlUpdate, + }); + + fireEvent.click(await screen.findByRole("button", { name: "Close Guardrail Info" })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.has("guardrail")).toBe(false); + expect(lastUpdate.options.history).toBe("replace"); + expect(await screen.findByText("Mock Guardrail Table")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index 7e59abf8e3d..901e39004f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -1,3 +1,4 @@ +import { parseAsString, useQueryState } from "nuqs"; import React, { useState, useEffect } from "react"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { ChevronDown, Code, Plus } from "lucide-react"; @@ -40,7 +41,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [selectedGuardrailId, setSelectedGuardrailId] = useState(null); + const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( + "guardrail", + parseAsString.withOptions({ history: "push" }), + ); const isAdmin = userRole ? isAdminRole(userRole) : false; const fetchGuardrails = async () => { @@ -63,16 +67,20 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole fetchGuardrails(); }, [accessToken]); + const closeGuardrailDetail = () => { + void setSelectedGuardrailId(null, { history: "replace" }); + }; + const handleAddGuardrail = () => { if (selectedGuardrailId) { - setSelectedGuardrailId(null); + closeGuardrailDetail(); } setIsAddModalVisible(true); }; const handleAddCustomCodeGuardrail = () => { if (selectedGuardrailId) { - setSelectedGuardrailId(null); + closeGuardrailDetail(); } setIsCustomCodeModalVisible(true); }; @@ -175,7 +183,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole {selectedGuardrailId ? ( setSelectedGuardrailId(null)} + onClose={closeGuardrailDetail} accessToken={accessToken} isAdmin={isAdmin} /> @@ -184,7 +192,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole guardrailsList={guardrailsList} isLoading={isLoading} onDeleteClick={handleDeleteClick} - onGuardrailClick={(id) => setSelectedGuardrailId(id)} + onGuardrailClick={(id) => void setSelectedGuardrailId(id)} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index fcffc2122e7..836921104ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -338,3 +338,27 @@ describe("Guardrail Info", () => { expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); }); + +describe("Guardrail Info when the guardrail cannot be loaded", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should keep Back to Guardrails reachable so a stale ?guardrail= link is not a dead end", async () => { + vi.mocked(networking.getGuardrailInfo).mockRejectedValue(new Error("Guardrail stale-id not found")); + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: [], + }); + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + const onClose = vi.fn(); + + render(); + + expect(await screen.findByText("Guardrail not found")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: /back to guardrails/i })); + expect(onClose).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index aaa656d15bb..c9162d99934 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -481,16 +481,18 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, return
Loading...
; } + const backButton = ( + + ); + if (!guardrailData) { - return
Guardrail not found
; + return
{backButton}Guardrail not found
; } - // Format date helper function - const formatDate = (dateString?: string) => { - if (!dateString) return "-"; - const date = new Date(dateString); - return date.toLocaleString(); - }; + const formatDate = (dateString?: string) => (dateString ? new Date(dateString).toLocaleString() : "-"); // Format the provider display name and logo const { logo, displayName } = getGuardrailLogoAndName(guardrailData.litellm_params?.guardrail || ""); @@ -510,10 +512,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, return (
- + {backButton}

{guardrailData.guardrail_name || "Unnamed Guardrail"}

{guardrailData.guardrail_id}

diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index cad41557100..d7df6ce33bb 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,11 +1,11 @@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { SearchSelect } from "@/components/shared/SearchSelect"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react"; import { Switch } from "@/components/ui/switch"; import { AffinityControls } from "./AffinityControls"; +import TierRowSelect from "./TierRowSelect"; import { ModalityRoutingControls } from "./ModalityRoutingControls"; import { Card, CardContent } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; @@ -50,6 +50,8 @@ import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; +import CompressionControls from "./CompressionControls"; +import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; export type { DimensionWeights, TierBoundaries, TokenThresholds }; export type { CustomTierSet, TierRow } from "./tier_rows"; @@ -370,27 +372,6 @@ const TierRowEditFields: React.FC<{ ); -const TierRowSelect: React.FC<{ - label: string; - options: { value: string; label: string }[]; - value: string | null; - onValueChange: (rowId: string) => void; - placeholder?: string; -}> = ({ label, options, value, onValueChange, placeholder }) => ( - -); - export type AdaptiveEligible = "all" | "classified_tier"; export type ComplexityTierLabels = Partial>; @@ -502,6 +483,10 @@ interface ComplexityRouterConfigProps { onMatchThresholdChange?: (threshold: number) => void; escalationKeywords?: string[]; onEscalationKeywordsChange?: (keywords: string[]) => void; + // Optional: not part of complexity_router_config, since it applies to every + // pre-routing strategy, not just the complexity router. + autoRouterCompression?: AutoRouterCompressionState; + onAutoRouterCompressionChange?: (state: AutoRouterCompressionState) => void; showValidationErrors?: boolean; } @@ -604,6 +589,8 @@ const ComplexityRouterConfig: React.FC = ({ onMatchThresholdChange = () => {}, escalationKeywords = [], onEscalationKeywordsChange, + autoRouterCompression = DEFAULT_AUTO_ROUTER_COMPRESSION, + onAutoRouterCompressionChange, showValidationErrors = false, }) => { const customTierSet = value.custom_tier_set; @@ -877,6 +864,17 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), + ...(onAutoRouterCompressionChange + ? [ + { + key: "compression", + label: Advanced: Compression, + children: ( + + ), + }, + ] + : []), ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange ? [ { diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx new file mode 100644 index 00000000000..c1817918f60 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -0,0 +1,91 @@ +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { SearchSelect, SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { Info } from "lucide-react"; +import React from "react"; +import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { + AutoRouterCompressionState, + isCompressionGuardrailProvider, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; + +interface CompressionControlsProps { + value: AutoRouterCompressionState; + onChange: (state: AutoRouterCompressionState) => void; +} + +const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; + +const CompressionControls: React.FC = ({ value, onChange }) => { + const { routing, sameAsRouting, model } = value; + const onRoutingChange = (newRouting: string | undefined) => + onChange({ ...value, routing: newRouting, sameAsRouting: newRouting === undefined ? true : sameAsRouting }); + const onSameAsRoutingChange = (newSameAsRouting: boolean) => onChange({ ...value, sameAsRouting: newSameAsRouting }); + const onModelChange = (newModel: string | undefined) => onChange({ ...value, model: newModel }); + + const { data } = useGuardrails(); + const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) + .filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail)) + .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); + const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; + + return ( +
+
+
+ Routing decision + + + +
+ onRoutingChange(value === "" ? undefined : value)} + placeholder="Inherit from the request's own compression guardrails" + emptyText="No compression guardrails found" + aria-label="Routing decision compression" + /> +
+ + {routing !== undefined && ( +
+ Model call + onSameAsRoutingChange(value === "same")} + className="w-full" + > +
+ + +
+ + + {!sameAsRouting && ( +
+ onModelChange(value === "" ? undefined : value)} + placeholder="None (no compression)" + emptyText="No compression guardrails found" + aria-label="Model call compression" + /> +
+ )} +
+ )} +
+ ); +}; + +export default CompressionControls; diff --git a/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx b/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx new file mode 100644 index 00000000000..ad7d53f9eae --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx @@ -0,0 +1,25 @@ +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import React from "react"; + +const TierRowSelect: React.FC<{ + label: string; + options: { value: string; label: string }[]; + value: string | null; + onValueChange: (rowId: string) => void; + placeholder?: string; +}> = ({ label, options, value, onValueChange, placeholder }) => ( + +); + +export default TierRowSelect; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d2f6b10c3a6..5605a993ded 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1,4 +1,12 @@ -import { renderWithProviders, screen, waitFor, within, fireEvent, testQueryClient } from "../../../tests/test-utils"; +import { + renderWithProviders, + screen, + waitFor, + within, + fireEvent, + testQueryClient, + chooseSelectOption, +} from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { vi } from "vitest"; import AddAutoRouterTab from "./add_auto_router_tab"; @@ -522,6 +530,71 @@ describe("AddAutoRouterTab", () => { ); }); + describe("prompt compression", () => { + it("leaves both compression keys out of the create payload when the section is untouched", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-router"); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted).not.toHaveProperty("auto_router_routing_compression"); + expect(submitted).not.toHaveProperty("auto_router_model_compression"); + }); + + it("mirrors an explicit no-compression routing choice onto the model call by default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-explicit-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + + it("defaults the model call to none when different is chosen but nothing is picked there", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "different-compression-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + await user.click(screen.getByText("Use a different compression")); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + }); + // The scalar floor is the one scorer knob with no group dict behind it, so its wiring into the create // payload is only proven end to end. 0 is the case a truthy check would silently drop. it("carries a reasoning override floor of 0 through to the create payload", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index ff71f161800..60780f9cbe0 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -32,6 +32,11 @@ import ComplexityRouterConfig, { } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { DEFAULT_ESCALATION_KEYWORDS } from "./EscalationKeywords"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, +} from "./buildAutoRouterCompression"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { BuildComplexityRouterConfigParams, @@ -194,6 +199,9 @@ const AddAutoRouterTab: React.FC = ({ const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); const [escalationKeywords, setEscalationKeywords] = useState(DEFAULT_ESCALATION_KEYWORDS); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); const [isSubmitting, setIsSubmitting] = useState(false); @@ -465,6 +473,7 @@ const AddAutoRouterTab: React.FC = ({ model_type: "complexity_router", complexity_router_config: complexityRouterConfigPayload, model_access_group: form.getValues("model_access_group"), + ...buildAutoRouterCompressionParams(autoRouterCompression), }; await handleAddAutoRouterSubmit(submitValues, accessToken, () => form.reset(EMPTY_FORM_VALUES), handleOk); @@ -670,6 +679,8 @@ const AddAutoRouterTab: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} showValidationErrors={showValidationErrors} />
diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts new file mode 100644 index 00000000000..20d6af50d18 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -0,0 +1,119 @@ +import { + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; + +describe("buildAutoRouterCompressionParams", () => { + it("omits both keys when routing was never configured", () => { + expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({}); + }); + + it("mirrors routing onto model when same-as-routing is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + }); + + it("uses the explicit model choice when different is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: "headroom-b", + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + }); + + it("defaults the model side to none when different is chosen but nothing is picked", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: NO_COMPRESSION, + }); + }); + + it("sends the none sentinel when routing itself is explicitly turned off", () => { + const params = buildAutoRouterCompressionParams({ + routing: NO_COMPRESSION, + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: NO_COMPRESSION, + auto_router_model_compression: NO_COMPRESSION, + }); + }); +}); + +describe("hydrateAutoRouterCompression", () => { + it("returns the default state when neither key is set", () => { + expect(hydrateAutoRouterCompression({})).toEqual(DEFAULT_AUTO_ROUTER_COMPRESSION); + }); + + it("is same-as-routing when the model value matches routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("is different when the model value diverges from routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" }); + }); + + it("treats a missing model key as no model-hop compression, not same-as-routing", () => { + const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "none" }); + }); + + it("re-saving a routing-only config leaves the model hop uncompressed", () => { + // Regression: the backend reads an absent model key as no model-hop compression. + // Hydrating it as same-as-routing made opening the router and saving any unrelated + // edit write the routing guardrail onto the model hop, so the model call silently + // started receiving compressed messages. + const stored = { auto_router_routing_compression: "headroom-a" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored)); + expect(rebuilt.auto_router_model_compression).toBe("none"); + expect(rebuilt.auto_router_model_compression).not.toBe("headroom-a"); + }); + + it("surfaces a stored model-only policy instead of reading as untouched", () => { + // Regression: the backend treats either key alone as an authoritative policy, so a + // model-only config that hydrated to the inherit state was invisible in the form, + // and the next save overwrote the stored model hop with the routing value. + const state = hydrateAutoRouterCompression({ auto_router_model_compression: "headroom-b" }); + expect(state).toEqual({ routing: "none", sameAsRouting: false, model: "headroom-b" }); + }); + + it("round-trips a model-only policy without changing either hop", () => { + const stored = { auto_router_model_compression: "headroom-b" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored)); + expect(rebuilt.auto_router_model_compression).toBe("headroom-b"); + expect(rebuilt.auto_router_routing_compression).toBe("none"); + }); + + it("round-trips through buildAutoRouterCompressionParams", () => { + const original = { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "none" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(original)); + expect(rebuilt).toEqual(original); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts new file mode 100644 index 00000000000..6f401a12865 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -0,0 +1,69 @@ +/** + * Maps the auto router's compression form state to the two flat litellm_params keys + * the backend reads (litellm.proxy.guardrails.auto_router_compression), and back. + * + * `routing` being undefined means the section was never touched: both keys are + * omitted from the payload, and the request's own compression guardrails apply to + * both hops unchanged. Once `routing` has a value (a guardrail name, or the "none" + * sentinel for explicit no-compression), the auto router is authoritative and the + * model side always gets a concrete value too, mirroring `routing` when same-as + * is chosen and defaulting to "none" otherwise. + */ + +export const NO_COMPRESSION = "none"; + +/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in + * litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */ +export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr"]; + +export const isCompressionGuardrailProvider = (provider: unknown): boolean => + typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase()); + +export interface AutoRouterCompressionState { + routing: string | undefined; + sameAsRouting: boolean; + model: string | undefined; +} + +export interface AutoRouterCompressionLitellmParams { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; +} + +export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = { + routing: undefined, + sameAsRouting: true, + model: undefined, +}; + +export const buildAutoRouterCompressionParams = ( + state: AutoRouterCompressionState, +): AutoRouterCompressionLitellmParams => { + if (state.routing === undefined) return {}; + return { + auto_router_routing_compression: state.routing, + auto_router_model_compression: state.sameAsRouting ? state.routing : state.model ?? NO_COMPRESSION, + }; +}; + +export const hydrateAutoRouterCompression = (litellmParams: { + auto_router_routing_compression?: string | null; + auto_router_model_compression?: string | null; +}): AutoRouterCompressionState => { + const storedRouting = litellmParams.auto_router_routing_compression ?? undefined; + const storedModel = litellmParams.auto_router_model_compression ?? undefined; + + // Only neither key set means the section was never touched. The backend treats + // either key on its own as an authoritative policy (policy_from_litellm_params), so + // reading a model-only config as untouched would hide it from the form and let the + // next save overwrite the stored model hop. + if (storedRouting === undefined && storedModel === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + + // An absent key on either hop is no compression for that hop, not same-as-the-other: + // the backend reads it as None. Hydrating it as same-as-routing would make re-saving + // an unrelated edit write one hop's guardrail onto the other. + const routing = storedRouting ?? NO_COMPRESSION; + const model = storedModel ?? NO_COMPRESSION; + const sameAsRouting = model === routing; + return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; +}; diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx index 9385836ce1a..59d9ecf205e 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx @@ -1,8 +1,9 @@ import { modelCreateCall } from "../networking"; import { toast } from "@/lib/toast"; import type { ComplexityRouterConfigPayload } from "./build_complexity_router_config"; +import type { AutoRouterCompressionLitellmParams } from "./buildAutoRouterCompression"; -export interface AddAutoRouterValues { +export interface AddAutoRouterValues extends AutoRouterCompressionLitellmParams { auto_router_name: string; auto_router_default_model: string | undefined; model_type: "complexity_router"; @@ -24,6 +25,8 @@ export const handleAddAutoRouterSubmit = async ( model: "auto_router/complexity_router", complexity_router_config: values.complexity_router_config, complexity_router_default_model: values.auto_router_default_model, + auto_router_routing_compression: values.auto_router_routing_compression, + auto_router_model_compression: values.auto_router_model_compression, }, model_info: { ...(values.team_id ? { team_id: values.team_id } : {}), diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index e45f84fa646..3367c810991 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -1036,6 +1036,87 @@ describe("EditAutoRouterModal with a stored custom tier set", () => { expect(savedConfig().tier_model_configs).toEqual(CUSTOM_STORED.tier_model_configs); }); }); +describe("EditAutoRouterModal prompt compression", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const savedLitellmParams = () => { + const [, payload] = modelPatchUpdateCall.mock.calls.at(-1) ?? []; + return payload?.litellm_params; + }; + + const renderWithStoredCompression = (compression?: { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; + }) => + renderWithProviders( + , + ); + + it("leaves both compression keys out of an untouched save when none were stored", async () => { + const user = userEvent.setup(); + renderWithStoredCompression(); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_routing_compression"); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression"); + }); + + it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); + expect(savedLitellmParams()?.auto_router_model_compression).toBe("headroom-a"); + }); + + it("shows a stored different-compression choice as Use a different compression, not Same", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByText("Advanced: Compression")); + + expect(await screen.findByRole("combobox", { name: "Routing decision compression" })).toHaveValue("headroom-a"); + expect(screen.getByRole("radio", { name: "Use a different compression" })).toBeChecked(); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toHaveValue("None (no compression)"); + }); + + it("preserves a stored different-compression choice through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); + expect(savedLitellmParams()?.auto_router_model_compression).toBe("none"); + }); +}); describe("EditAutoRouterModal classifier vision", () => { beforeEach(() => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index dd3ff1241ea..3c0013e267e 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -41,6 +41,12 @@ import { } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, +} from "../add_model/buildAutoRouterCompression"; import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords"; import { hydrateDimensionWeights, @@ -447,6 +453,9 @@ const EditAutoRouterModal: React.FC = ({ const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", @@ -539,6 +548,12 @@ const EditAutoRouterModal: React.FC = ({ setMatchThreshold( typeof parsedConfig.match_threshold === "number" ? parsedConfig.match_threshold : DEFAULT_MATCH_THRESHOLD, ); + setAutoRouterCompression( + hydrateAutoRouterCompression({ + auto_router_routing_compression: modelData.litellm_params?.auto_router_routing_compression, + auto_router_model_compression: modelData.litellm_params?.auto_router_model_compression, + }), + ); form.reset({ ...EMPTY_FORM_VALUES, @@ -651,6 +666,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, + ...buildAutoRouterCompressionParams(autoRouterCompression), }; const updatedModelInfo = { ...modelData.model_info, @@ -772,6 +788,8 @@ const EditAutoRouterModal: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} />
) : ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 676fb4c45ef..d3b583daec4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29356,6 +29356,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */ @@ -39528,6 +39532,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */