scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

chore: merge litellm_internal_staging into litellm_anthropic_wif_backend

Browse source
This commit is contained in:
mateo-berri 2026-09-05 12:58:17 -07:00
commit cc5c6c059a
82 changed files with 5618 additions and 522 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
.github/workflows/test-terraform-modules.yml vendored
View file

@ -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

1
litellm/constants.py
View file

@ -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 = "{}"

6
litellm/integrations/cloudzero/cz_stream_api.py
View file

@ -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."""

2
litellm/integrations/cloudzero/transform.py
View file

@ -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."""

20
litellm/integrations/custom_guardrail.py
View file

@ -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)

39
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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,

40
litellm/litellm_core_utils/logging_utils.py
View file

@ -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

18
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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,

21
litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py
View file

@ -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

23
litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py
View file

@ -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.

41
litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py Normal file
View file

@ -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)

469
litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py Normal file
View file

@ -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,
)

39
litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py
View file

@ -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},

7
litellm/proxy/common_request_processing.py
View file

@ -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,

267
litellm/proxy/guardrails/auto_router_compression.py Normal file
View file

@ -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)

22
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -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

81
litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py Normal file
View file

@ -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()

50
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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,

92
litellm/router.py
View file

@ -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),

17
litellm/router_utils/pattern_match_deployments.py
View file

@ -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)

4
litellm/types/router.py
View file

@ -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

2
litellm/types/utils.py
View file

@ -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",

77
terraform/litellm/gcp/README.md
View file

@ -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/<host>/global/networks/<vpc>"
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: <cloudsql_writer_ip>
dbname: <db_name>
passwordSecret:
name: <kubernetes-secret-with-db-credentials>
reader:
host: <cloudsql_reader_ip>
dbname: <db_name>
passwordSecret:
name: <kubernetes-secret-with-db-credentials>
redis:
host: <redis_host>
port: <redis_port>
masterKey:
secretName: <kubernetes-secret-with-master-key>
```
Create the database Secret with keys `username` (the `db_username` output)
and `password` (read it with `gcloud secrets versions access latest
--secret=<db_password_secret_id>`), 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=<mount path>` 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 |

6
terraform/litellm/gcp/bootstrap.tf
View file

@ -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
}

109
terraform/litellm/gcp/cloudrun.tf
View file

@ -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"
}

9
terraform/litellm/gcp/cloudsql.tf
View file

@ -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
}
}

5
terraform/litellm/gcp/examples/default/main.tf
View file

@ -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

30
terraform/litellm/gcp/examples/default/outputs.tf
View file

@ -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

8
terraform/litellm/gcp/examples/default/terraform.tfvars.example
View file

@ -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

24
terraform/litellm/gcp/examples/default/variables.tf
View file

@ -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" {

11
terraform/litellm/gcp/iam.tf
View file

@ -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)"
}

58
terraform/litellm/gcp/load_balancer.tf
View file

@ -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
}

5
terraform/litellm/gcp/locals.tf
View file

@ -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"

20
terraform/litellm/gcp/network.tf
View file

@ -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

58
terraform/litellm/gcp/outputs.tf
View file

@ -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=<path> 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
}

4
terraform/litellm/gcp/redis.tf
View file

@ -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]
}

175
terraform/litellm/gcp/tests/deps_only.tftest.hcl Normal file
View file

@ -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."
}
}

30
terraform/litellm/gcp/variables.tf
View file

@ -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/<host-project>/global/networks/<name>`). 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"
}

28
tests/e2e/access_control/test_access_control_e2e.py
View file

@ -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:

76
tests/e2e/llm_translation/test_cache_control.py
View file

@ -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)
)

98
tests/e2e/llm_translation/test_chat_completions_regression_e2e.py
View file

@ -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}"

22
tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py
View file

@ -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

40
tests/e2e/llm_translation/test_passthrough_e2e.py
View file

@ -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.

16
tests/e2e/management/management_client.py
View file

@ -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."""

55
tests/e2e/management/test_key_management_e2e.py
View file

@ -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

35
tests/e2e/management/test_management_e2e.py
View file

@ -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")

10
tests/e2e/models.py
View file

@ -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):

2
tests/e2e/ui/tests/tables/tableScrolling.spec.ts
View file

@ -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 {

2
tests/router_unit_tests/test_router_index_management.py
View file

@ -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

24
tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py
View file

@ -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")

27
tests/test_litellm/integrations/cloudzero/test_transform.py
View file

@ -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()

328
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -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")

134
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -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

71
tests/test_litellm/litellm_core_utils/test_logging_utils.py
View file

@ -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 == []

74
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py
View file

@ -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"))

794
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py Normal file
View file

@ -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="<html>maintenance</html>", 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"

90
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -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",

381
tests/test_litellm/proxy/guardrails/test_auto_router_compression.py Normal file
View file

@ -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": {}}

117
tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py Normal file
View file

@ -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

150
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -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"""

380
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -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(),

69
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -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

32
tests/test_litellm/router_utils/test_pattern_match_deployments.py
View file

@ -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"]

342
tests/test_litellm/test_router.py
View file

@ -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

6
type-discipline-budget.json
View file

@ -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

116
ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx
View file

@ -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 }) => <div data-testid="log-viewer">{guardrailName}</div>,
}));
function wrapper({ children }: { children: React.ReactNode }) {
const queryClient = new QueryClient({
defaultOptions: {
queries: { retry: false },
},
});
return <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>;
}
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(<GuardrailsMonitorView accessToken="test-token" />, { wrapper });
it("should render overview and fetch guardrails usage when accessToken is provided", async () => {
renderWithProviders(<GuardrailsMonitorView accessToken="test-token" />);
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(<GuardrailsMonitorView accessToken={null} />, { wrapper });
renderWithProviders(<GuardrailsMonitorView accessToken={null} />);
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(<GuardrailsMonitorView accessToken="test-token" />, { 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(<GuardrailsMonitorView accessToken="test-token" />, { 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(<GuardrailsMonitorView accessToken="test-token" />, {
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();
});
});
});

16
ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx
View file

@ -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<View>({ 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 (
<main className="w-full min-w-0 flex-1 p-8">
{view.type === "overview" ? (
{!selectedGuardrailId ? (
<GuardrailsOverview
accessToken={accessToken}
startDate={startDate}
@ -59,7 +61,7 @@ export default function GuardrailsMonitorView({ accessToken = null }: Guardrails
<>
<div className="mb-4 flex items-center justify-end">{dateRangeControl}</div>
<GuardrailDetail
guardrailId={view.guardrailId}
guardrailId={selectedGuardrailId}
onBack={handleBack}
accessToken={accessToken}
startDate={startDate}

77
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
View file

@ -1,7 +1,8 @@
import { render, screen, fireEvent, waitFor, within } from "@testing-library/react";
import { type UrlUpdateEvent } from "nuqs/adapters/testing";
import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
import GuardrailsPanel from "./GuardrailsPanel";
import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking";
import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils";
vi.mock("@/components/networking", () => ({
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) => (
<div>
<div>Mock Guardrail Table</div>
{guardrailsList.length > 0 && (
<button
data-testid="delete-button"
onClick={() => onDeleteClick(guardrailsList[0].guardrail_id, guardrailsList[0].guardrail_name)}
>
Delete
</button>
<>
<button
data-testid="delete-button"
onClick={() => onDeleteClick(guardrailsList[0].guardrail_id, guardrailsList[0].guardrail_name)}
>
Delete
</button>
<button data-testid="open-button" onClick={() => onGuardrailClick(guardrailsList[0].guardrail_id)}>
Open
</button>
</>
)}
</div>
),
@ -32,7 +38,12 @@ vi.mock("./guardrail_table", () => ({
vi.mock("./guardrail_info", () => ({
__esModule: true,
default: () => <div>Mock Guardrail Info View</div>,
default: ({ guardrailId, onClose }: { guardrailId: string; onClose: () => void }) => (
<div>
<div data-testid="guardrail-info-view">Mock Guardrail Info View {guardrailId}</div>
<button onClick={onClose}>Close Guardrail Info</button>
</div>
),
}));
vi.mock("./GuardrailTestPlayground", async () => {
@ -112,7 +123,7 @@ describe("GuardrailsPanel", () => {
});
it("should render the component", async () => {
render(<GuardrailsPanel {...defaultProps} />);
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
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(<GuardrailsPanel {...defaultProps} />);
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
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(<GuardrailsPanel {...defaultProps} />);
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
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(<GuardrailsPanel {...defaultProps} />);
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
fireEvent.click(screen.getByText("Test Playground"));
@ -161,7 +172,7 @@ describe("GuardrailsPanel", () => {
});
it("should not delete anything when the modal is cancelled", async () => {
render(<GuardrailsPanel {...defaultProps} />);
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
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(<GuardrailsPanel {...defaultProps} />, { 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(<GuardrailsPanel {...defaultProps} />, { 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(<GuardrailsPanel {...defaultProps} />, {
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();
});
});
});

18
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx
View file

@ -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<GuardrailsPanelProps> = ({ accessToken, userRole
const [isDeleting, setIsDeleting] = useState(false);
const [guardrailToDelete, setGuardrailToDelete] = useState<Guardrail | null>(null);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [selectedGuardrailId, setSelectedGuardrailId] = useState<string | null>(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<GuardrailsPanelProps> = ({ 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<GuardrailsPanelProps> = ({ accessToken, userRole
{selectedGuardrailId ? (
<GuardrailInfoView
guardrailId={selectedGuardrailId}
onClose={() => setSelectedGuardrailId(null)}
onClose={closeGuardrailDetail}
accessToken={accessToken}
isAdmin={isAdmin}
/>
@ -184,7 +192,7 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
guardrailsList={guardrailsList}
isLoading={isLoading}
onDeleteClick={handleDeleteClick}
onGuardrailClick={(id) => setSelectedGuardrailId(id)}
onGuardrailClick={(id) => void setSelectedGuardrailId(id)}
/>
)}

24
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx
View file

@ -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(<GuardrailInfoView guardrailId="stale-id" onClose={onClose} accessToken="123" isAdmin={true} />);
expect(await screen.findByText("Guardrail not found")).toBeInTheDocument();
fireEvent.click(screen.getByRole("button", { name: /back to guardrails/i }));
expect(onClose).toHaveBeenCalledTimes(1);
});
});

21
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx
View file

@ -481,16 +481,18 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
return <div className="p-4">Loading...</div>;
}
const backButton = (
<Button variant="ghost" onClick={onClose} className="mb-4">
<ArrowLeft className="w-4 h-4" />
Back to Guardrails
</Button>
);
if (!guardrailData) {
return <div className="p-4">Guardrail not found</div>;
return <div className="p-4">{backButton}Guardrail not found</div>;
}
// 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<GuardrailInfoProps> = ({ guardrailId, onClose,
return (
<div className="p-4">
<div>
<Button variant="ghost" onClick={onClose} className="mb-4">
<ArrowLeft className="w-4 h-4" />
Back to Guardrails
</Button>
{backButton}
<h1 className="text-2xl font-semibold">{guardrailData.guardrail_name || "Unnamed Guardrail"}</h1>
<div className="flex items-center cursor-pointer">
<p className="text-muted-foreground font-mono">{guardrailData.guardrail_id}</p>

42
ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
View file

@ -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 }) => (
<Select items={options} value={value} onValueChange={(rowId: string | null) => rowId && onValueChange(rowId)}>
<SelectTrigger aria-label={label} className="w-full">
<SelectValue placeholder={placeholder} />
</SelectTrigger>
<SelectContent>
{options.map((option) => (
<SelectItem key={option.value} value={option.value}>
{option.label}
</SelectItem>
))}
</SelectContent>
</Select>
);
export type AdaptiveEligible = "all" | "classified_tier";
export type ComplexityTierLabels = Partial<Record<keyof ComplexityTiers, string>>;
@ -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<ComplexityRouterConfigProps> = ({
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<ComplexityRouterConfigProps> = ({
},
]
: []),
...(onAutoRouterCompressionChange
? [
{
key: "compression",
label: <strong className="text-foreground font-semibold">Advanced: Compression</strong>,
children: (
<CompressionControls value={autoRouterCompression} onChange={onAutoRouterCompressionChange} />
),
},
]
: []),
...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange
? [
{

91
ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx Normal file
View file

@ -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<CompressionControlsProps> = ({ 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 (
<div className="space-y-4">
<div>
<div className="mb-1 flex items-center gap-2">
<span className="text-sm font-medium">Routing decision</span>
<SimpleTooltip content="Compression applied to the classifier's own call that picks a tier, separate from the model the request routes to.">
<Info className="size-4 text-muted-foreground" />
</SimpleTooltip>
</div>
<SearchSelect
options={options}
value={routing ?? ""}
onValueChange={(value) => onRoutingChange(value === "" ? undefined : value)}
placeholder="Inherit from the request's own compression guardrails"
emptyText="No compression guardrails found"
aria-label="Routing decision compression"
/>
</div>
{routing !== undefined && (
<div>
<span className="mb-2 block text-sm font-medium">Model call</span>
<RadioGroup
value={sameAsRouting ? "same" : "different"}
onValueChange={(value: unknown) => onSameAsRoutingChange(value === "same")}
className="w-full"
>
<div className="flex w-full flex-col items-start gap-2">
<Label className="items-start font-normal leading-normal">
<RadioGroupItem value="same" className="mt-0.5" />
<span>Same as the routing decision</span>
</Label>
<Label className="items-start font-normal leading-normal">
<RadioGroupItem value="different" className="mt-0.5" />
<span>Use a different compression</span>
</Label>
</div>
</RadioGroup>
{!sameAsRouting && (
<div className="mt-3">
<SearchSelect
options={options}
value={model ?? ""}
onValueChange={(value) => onModelChange(value === "" ? undefined : value)}
placeholder="None (no compression)"
emptyText="No compression guardrails found"
aria-label="Model call compression"
/>
</div>
)}
</div>
)}
</div>
);
};
export default CompressionControls;

25
ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx Normal file
View file

@ -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 }) => (
<Select items={options} value={value} onValueChange={(rowId: string | null) => rowId && onValueChange(rowId)}>
<SelectTrigger aria-label={label} className="w-full">
<SelectValue placeholder={placeholder} />
</SelectTrigger>
<SelectContent>
{options.map((option) => (
<SelectItem key={option.value} value={option.value}>
{option.label}
</SelectItem>
))}
</SelectContent>
</Select>
);
export default TierRowSelect;

75
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
View file

@ -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(<Harness />);
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(<Harness />);
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(<Harness />);
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 () => {

11
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
View file

@ -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<AddAutoRouterTabProps> = ({
const [embeddingModel, setEmbeddingModel] = useState<string | undefined>(undefined);
const [matchThreshold, setMatchThreshold] = useState<number>(DEFAULT_MATCH_THRESHOLD);
const [escalationKeywords, setEscalationKeywords] = useState<string[]>(DEFAULT_ESCALATION_KEYWORDS);
const [autoRouterCompression, setAutoRouterCompression] = useState<AutoRouterCompressionState>(
DEFAULT_AUTO_ROUTER_COMPRESSION,
);
const [showValidationErrors, setShowValidationErrors] = useState<boolean>(false);
const [editingTiers, setEditingTiers] = useState(false);
const [isSubmitting, setIsSubmitting] = useState(false);
@ -465,6 +473,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
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<AddAutoRouterTabProps> = ({
onMatchThresholdChange={setMatchThreshold}
escalationKeywords={escalationKeywords}
onEscalationKeywordsChange={setEscalationKeywords}
autoRouterCompression={autoRouterCompression}
onAutoRouterCompressionChange={setAutoRouterCompression}
showValidationErrors={showValidationErrors}
/>
</div>

119
ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts Normal file
View file

@ -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);
});
});

69
ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts Normal file
View file

@ -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 };
};

5
ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx
View file

@ -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 } : {}),

81
ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
View file

@ -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(
<EditAutoRouterModal
isVisible
onCancel={vi.fn()}
onSuccess={vi.fn()}
modelData={{
...MODEL_DATA,
litellm_params: { ...MODEL_DATA.litellm_params, ...compression },
}}
accessToken="token"
userRole="Admin"
/>,
);
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(() => {

18
ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
View file

@ -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<EditAutoRouterModalProps> = ({
const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState<boolean>(false);
const [embeddingModel, setEmbeddingModel] = useState<string | undefined>(undefined);
const [matchThreshold, setMatchThreshold] = useState<number>(DEFAULT_MATCH_THRESHOLD);
const [autoRouterCompression, setAutoRouterCompression] = useState<AutoRouterCompressionState>(
DEFAULT_AUTO_ROUTER_COMPRESSION,
);
const [complexityRouterConfig, setComplexityRouterConfig] = useState<ComplexityRouterConfigValue>({
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] },
classifier_type: "heuristic",
@ -539,6 +548,12 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
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<EditAutoRouterModalProps> = ({
...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<EditAutoRouterModalProps> = ({
onMatchThresholdChange={setMatchThreshold}
escalationKeywords={escalationKeywords}
onEscalationKeywordsChange={setEscalationKeywords}
autoRouterCompression={autoRouterCompression}
onAutoRouterCompressionChange={setAutoRouterCompression}
/>
</div>
) : (

8
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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 */