Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_otel_v2_admin_owned_destinations

This commit is contained in:
Yucheng Zhu 2026-07-31 01:00:02 -07:00
commit a4bf719054
66 changed files with 4590 additions and 613 deletions

View file

@ -24,6 +24,8 @@ class S3Logger:
s3_aws_secret_access_key=None,
s3_aws_session_token=None,
s3_config=None,
s3_server_side_encryption: str | None = None,
s3_sse_kms_key_id: str | None = None,
**kwargs,
):
import boto3
@ -50,11 +52,16 @@ class S3Logger:
s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token")
s3_config = litellm.s3_callback_params.get("s3_config")
s3_path = litellm.s3_callback_params.get("s3_path")
s3_server_side_encryption = litellm.s3_callback_params.get("s3_server_side_encryption")
s3_sse_kms_key_id = litellm.s3_callback_params.get("s3_sse_kms_key_id")
# done reading litellm.s3_callback_params
s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False))
self.s3_use_team_prefix = s3_use_team_prefix
self.bucket_name = s3_bucket_name
self.s3_path = s3_path
self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
s3_server_side_encryption, s3_sse_kms_key_id
)
verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}")
# Create an S3 client with custom endpoint URL
self.s3_client = boto3.client(
@ -136,6 +143,15 @@ class S3Logger:
print_verbose(f"\ns3 Logger - Logging payload = {payload_str}")
sse_params = {
key: value
for key, value in {
"ServerSideEncryption": self.s3_server_side_encryption,
"SSEKMSKeyId": self.s3_sse_kms_key_id,
}.items()
if value
}
response = self.s3_client.put_object(
Bucket=self.bucket_name,
Key=s3_object_key,
@ -144,6 +160,7 @@ class S3Logger:
ContentLanguage="en",
ContentDisposition=f'inline; filename="{s3_object_download_filename}"',
CacheControl="private, immutable, max-age=31536000, s-maxage=0",
**sse_params,
)
print_verbose(f"Response from s3:{str(response)}")
@ -155,6 +172,33 @@ class S3Logger:
pass
def _validated_sse_value(name: str, value: str | None) -> str | None:
if value is None or isinstance(value, str):
return value
verbose_logger.warning(
f"s3 logging: ignoring {name} because it has invalid type {type(value).__name__}; expected a string"
)
return None
def resolve_sse_params(
server_side_encryption: str | None,
sse_kms_key_id: str | None,
) -> tuple[str | None, str | None]:
valid_sse = _validated_sse_value("s3_server_side_encryption", server_side_encryption)
valid_key_id = _validated_sse_value("s3_sse_kms_key_id", sse_kms_key_id)
algorithm = valid_sse or ("aws:kms" if valid_key_id else None)
if algorithm is None:
return None, None
if valid_key_id and not algorithm.startswith("aws:kms"):
verbose_logger.warning(
f"s3 logging: ignoring s3_sse_kms_key_id because s3_server_side_encryption is {algorithm}; "
"set it to aws:kms to encrypt with the KMS key"
)
return algorithm, None
return algorithm, valid_key_id
def get_s3_object_key(
s3_path: str,
prefix: str,

View file

@ -8,13 +8,14 @@ NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to uplo
import asyncio
import time
from collections.abc import Mapping
from datetime import datetime
from typing import List, Optional, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
from litellm.integrations.s3 import get_s3_object_key
from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@ -55,6 +56,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_use_key_prefix: bool = False,
s3_use_virtual_hosted_style: bool = False,
s3_server_side_encryption: Optional[str] = None,
s3_sse_kms_key_id: str | None = None,
s3_callback_params_override: Optional[dict] = None,
**kwargs,
):
@ -94,6 +96,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_use_key_prefix=s3_use_key_prefix,
s3_use_virtual_hosted_style=s3_use_virtual_hosted_style,
s3_server_side_encryption=s3_server_side_encryption,
s3_sse_kms_key_id=s3_sse_kms_key_id,
)
verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}")
@ -148,6 +151,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_use_key_prefix: bool = False,
s3_use_virtual_hosted_style: bool = False,
s3_server_side_encryption: Optional[str] = None,
s3_sse_kms_key_id: str | None = None,
params_source: Optional[dict] = None,
):
"""
@ -197,10 +201,20 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style
)
self.s3_server_side_encryption = params.get("s3_server_side_encryption") or s3_server_side_encryption
self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
params.get("s3_server_side_encryption") or s3_server_side_encryption,
params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id,
)
return
def _sse_headers(self) -> Mapping[str, str]:
candidates = {
"x-amz-server-side-encryption": self.s3_server_side_encryption,
"x-amz-server-side-encryption-aws-kms-key-id": self.s3_sse_kms_key_id,
}
return {key: value for key, value in candidates.items() if value}
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await self._async_log_event_base(
kwargs=kwargs,
@ -335,11 +349,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
"Content-Language": "en",
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
**(
{"x-amz-server-side-encryption": self.s3_server_side_encryption}
if self.s3_server_side_encryption
else {}
),
**self._sse_headers(),
}
req = requests.Request("PUT", url, data=json_string, headers=headers)
prepped = req.prepare()
@ -510,11 +520,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
"Content-Language": "en",
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
**(
{"x-amz-server-side-encryption": self.s3_server_side_encryption}
if self.s3_server_side_encryption
else {}
),
**self._sse_headers(),
}
req = requests.Request("PUT", url, data=json_string, headers=headers)
prepped = req.prepare()

View file

@ -2,7 +2,8 @@
## Helper utilities for cost_per_token()
from dataclasses import dataclass
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
from types import MappingProxyType
from typing import Any, Literal, Mapping, Optional, Tuple, TypedDict, cast
import litellm
from litellm._logging import verbose_logger
@ -39,6 +40,14 @@ _VALID_DATA_RESIDENCIES = frozenset(r.value for r in DataResidency)
# of being rebuilt for every model_info key on every call.
_SERVICE_TIER_SUFFIXES: tuple[str, ...] = tuple(f"_{st.value}" for st in ServiceTier)
_SERVICE_TIER_TO_COST_KEY_SUFFIX: Mapping[str, str] = MappingProxyType(
{
ServiceTier.FLEX.value: ServiceTier.FLEX.value,
ServiceTier.PRIORITY.value: ServiceTier.PRIORITY.value,
ServiceTier.FAST.value: ServiceTier.PRIORITY.value,
}
)
def _get_token_detail_value(details: object, key: str) -> Optional[int]:
if isinstance(details, dict):
@ -177,7 +186,7 @@ def _get_service_tier_cost_key(base_key: str, service_tier: Optional[str]) -> st
Args:
base_key: The base cost key (e.g., "input_cost_per_token")
service_tier: The service tier ("flex", "priority", or None for standard)
service_tier: The service tier ("flex", "priority", "fast", or None for standard)
Returns:
str: The cost key to use (e.g., "input_cost_per_token_flex" or "input_cost_per_token")
@ -185,12 +194,11 @@ def _get_service_tier_cost_key(base_key: str, service_tier: Optional[str]) -> st
if service_tier is None:
return base_key
# Only use service tier specific keys for "flex" and "priority"
if service_tier.lower() in [ServiceTier.FLEX.value, ServiceTier.PRIORITY.value]:
return f"{base_key}_{service_tier.lower()}"
suffix = _SERVICE_TIER_TO_COST_KEY_SUFFIX.get(service_tier.lower())
if suffix is None:
return base_key
# For any other service tier, use standard pricing
return base_key
return f"{base_key}_{suffix}"
def _parse_above_token_threshold(key: str) -> float:

View file

@ -23678,14 +23678,17 @@
"gpt-5.6": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06,
"cache_creation_input_token_cost_flex": 3.125e-06,
"cache_creation_input_token_cost_priority": 1.25e-05,
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 1e-06,
"input_cost_per_token": 5e-06,
"input_cost_per_token_above_272k_tokens": 1e-05,
"input_cost_per_token_above_272k_tokens_flex": 5e-06,
"input_cost_per_token_batches": 2.5e-06,
"input_cost_per_token_flex": 2.5e-06,
"input_cost_per_token_priority": 1e-05,
@ -23696,6 +23699,7 @@
"mode": "chat",
"output_cost_per_token": 3e-05,
"output_cost_per_token_above_272k_tokens": 4.5e-05,
"output_cost_per_token_above_272k_tokens_flex": 2.25e-05,
"output_cost_per_token_batches": 1.5e-05,
"output_cost_per_token_flex": 1.5e-05,
"output_cost_per_token_priority": 6e-05,
@ -23731,14 +23735,17 @@
"gpt-5.6-sol": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06,
"cache_creation_input_token_cost_flex": 3.125e-06,
"cache_creation_input_token_cost_priority": 1.25e-05,
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 1e-06,
"input_cost_per_token": 5e-06,
"input_cost_per_token_above_272k_tokens": 1e-05,
"input_cost_per_token_above_272k_tokens_flex": 5e-06,
"input_cost_per_token_batches": 2.5e-06,
"input_cost_per_token_flex": 2.5e-06,
"input_cost_per_token_priority": 1e-05,
@ -23749,6 +23756,7 @@
"mode": "chat",
"output_cost_per_token": 3e-05,
"output_cost_per_token_above_272k_tokens": 4.5e-05,
"output_cost_per_token_above_272k_tokens_flex": 2.25e-05,
"output_cost_per_token_batches": 1.5e-05,
"output_cost_per_token_flex": 1.5e-05,
"output_cost_per_token_priority": 6e-05,
@ -23782,29 +23790,33 @@
"supports_xhigh_reasoning_effort": true
},
"gpt-5.6-terra": {
"cache_creation_input_token_cost": 3.125e-06,
"cache_creation_input_token_cost_above_272k_tokens": 6.25e-06,
"cache_creation_input_token_cost_flex": 1.5625e-06,
"cache_creation_input_token_cost_priority": 6.25e-06,
"cache_read_input_token_cost": 2.5e-07,
"cache_read_input_token_cost_above_272k_tokens": 5e-07,
"cache_read_input_token_cost_flex": 1.25e-07,
"cache_read_input_token_cost_priority": 5e-07,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_above_272k_tokens": 5e-06,
"input_cost_per_token_batches": 1.25e-06,
"input_cost_per_token_flex": 1.25e-06,
"input_cost_per_token_priority": 5e-06,
"cache_creation_input_token_cost": 2.5e-06,
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
"cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06,
"cache_creation_input_token_cost_flex": 1.25e-06,
"cache_creation_input_token_cost_priority": 5e-06,
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_272k_tokens": 4e-07,
"cache_read_input_token_cost_above_272k_tokens_flex": 2e-07,
"cache_read_input_token_cost_flex": 1e-07,
"cache_read_input_token_cost_priority": 4e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"input_cost_per_token_above_272k_tokens_flex": 2e-06,
"input_cost_per_token_batches": 1e-06,
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 4e-06,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_above_272k_tokens": 2.25e-05,
"output_cost_per_token_batches": 7.5e-06,
"output_cost_per_token_flex": 7.5e-06,
"output_cost_per_token_priority": 3e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_272k_tokens": 1.8e-05,
"output_cost_per_token_above_272k_tokens_flex": 9e-06,
"output_cost_per_token_batches": 6e-06,
"output_cost_per_token_flex": 6e-06,
"output_cost_per_token_priority": 2.4e-05,
"regional_processing_uplift_multiplier_eu": 1.1,
"regional_processing_uplift_multiplier_us": 1.1,
"supported_endpoints": [
@ -23835,29 +23847,33 @@
"supports_xhigh_reasoning_effort": true
},
"gpt-5.6-luna": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_creation_input_token_cost_flex": 6.25e-07,
"cache_creation_input_token_cost_priority": 2.5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"cache_read_input_token_cost_flex": 5e-08,
"cache_read_input_token_cost_priority": 2e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"input_cost_per_token_batches": 5e-07,
"input_cost_per_token_flex": 5e-07,
"input_cost_per_token_priority": 2e-06,
"cache_creation_input_token_cost": 2.5e-07,
"cache_creation_input_token_cost_above_272k_tokens": 5e-07,
"cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07,
"cache_creation_input_token_cost_flex": 1.25e-07,
"cache_creation_input_token_cost_priority": 5e-07,
"cache_read_input_token_cost": 2e-08,
"cache_read_input_token_cost_above_272k_tokens": 4e-08,
"cache_read_input_token_cost_above_272k_tokens_flex": 2e-08,
"cache_read_input_token_cost_flex": 1e-08,
"cache_read_input_token_cost_priority": 4e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_272k_tokens": 4e-07,
"input_cost_per_token_above_272k_tokens_flex": 2e-07,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_token_flex": 1e-07,
"input_cost_per_token_priority": 4e-07,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"output_cost_per_token_above_272k_tokens": 9e-06,
"output_cost_per_token_batches": 3e-06,
"output_cost_per_token_flex": 3e-06,
"output_cost_per_token_priority": 1.2e-05,
"output_cost_per_token": 1.2e-06,
"output_cost_per_token_above_272k_tokens": 1.8e-06,
"output_cost_per_token_above_272k_tokens_flex": 9e-07,
"output_cost_per_token_batches": 6e-07,
"output_cost_per_token_flex": 6e-07,
"output_cost_per_token_priority": 2.4e-06,
"regional_processing_uplift_multiplier_eu": 1.1,
"regional_processing_uplift_multiplier_us": 1.1,
"supported_endpoints": [
@ -45120,10 +45136,10 @@
"supports_vision": true
},
"bedrock_mantle/openai.gpt-5.6-terra": {
"input_cost_per_token": 2.75e-06,
"cache_creation_input_token_cost": 3.4375e-06,
"cache_read_input_token_cost": 2.75e-07,
"output_cost_per_token": 1.65e-05,
"input_cost_per_token": 2.2e-06,
"cache_creation_input_token_cost": 2.75e-06,
"cache_read_input_token_cost": 2.2e-07,
"output_cost_per_token": 1.32e-05,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
@ -45148,10 +45164,10 @@
"supports_vision": true
},
"bedrock_mantle/openai.gpt-5.6-luna": {
"input_cost_per_token": 1.1e-06,
"cache_creation_input_token_cost": 1.375e-06,
"cache_read_input_token_cost": 1.1e-07,
"output_cost_per_token": 6.6e-06,
"input_cost_per_token": 2.2e-07,
"cache_creation_input_token_cost": 2.75e-07,
"cache_read_input_token_cost": 2.2e-08,
"output_cost_per_token": 1.32e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 272000,
"max_output_tokens": 128000,

View file

@ -102,6 +102,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
byok_description: List[str] = Field(default_factory=list)
byok_api_key_help_url: Optional[str] = None
has_user_credential: Optional[bool] = None
connected_app_reachable: bool | None = None
source_url: Optional[str] = None
timeout: Optional[float] = None
max_concurrent_requests: Optional[int] = None

View file

@ -1963,6 +1963,18 @@ class MCPRequestHandler:
return allowed_tools
@staticmethod
def tool_is_granted(bare_tool_name: str, allowed_tool_names: list[str] | None) -> bool:
"""Whether key/team tool permissions reach ``bare_tool_name`` on one server.
``None`` means no tool-level restriction; an empty list grants nothing. Entries
name a tool on a single server and every writer stores them bare, so the
comparison is exact against the bare name rather than against the spellings
routing accepts. Both the listing path and the call path answer through here, so
discovery cannot advertise a tool that ``tools/call`` then refuses.
"""
return allowed_tool_names is None or bare_tool_name in allowed_tool_names
@staticmethod
async def is_tool_allowed_for_server(
tool_name: str,
@ -1973,7 +1985,7 @@ class MCPRequestHandler:
Check if a specific tool is allowed for a server based on key/team permissions.
Args:
tool_name: Name of the tool to check
tool_name: Bare tool name, already resolved against the server's prefixes
server_id: Server ID
user_api_key_auth: User auth
@ -1984,17 +1996,7 @@ class MCPRequestHandler:
server_id=server_id,
user_api_key_auth=user_api_key_auth,
)
# None means no restrictions (allow all)
if allowed_tools is None:
return True
# Empty list means no tools allowed
if not allowed_tools:
return False
# Check if tool is in allowed list
return tool_name in allowed_tools
return MCPRequestHandler.tool_is_granted(tool_name, allowed_tools)
@staticmethod
def is_tool_allowed(

View file

@ -116,12 +116,15 @@ from litellm.proxy._experimental.mcp_server.utils import (
get_server_prefix,
interpolate_headers,
is_short_mcp_tool_prefix_enabled,
is_tool_name_prefixed,
iter_known_server_prefixes,
iter_known_tool_name_spellings,
match_known_server_prefix,
match_known_tool_name,
merge_mcp_headers,
normalize_server_name,
openapi_tool_name,
parse_admin_env_vars,
split_server_prefix_from_name,
strip_known_server_prefix,
validate_mcp_server_name,
)
from litellm.proxy._types import (
@ -1783,7 +1786,7 @@ class MCPServerManager:
# Generate tool name (without prefix initially)
operation_id = operation.get("operationId", f"{method}_{path.replace('/', '_')}")
base_tool_name = operation_id.replace(" ", "_").lower()
base_tool_name = openapi_tool_name(operation_id)
# Add server prefix to tool name
prefixed_tool_name = add_server_prefix_to_name(base_tool_name, server_prefix)
@ -4185,10 +4188,8 @@ class MCPServerManager:
# Register every known prefix form (alias, server_name, server_id,
# short ID) so call_tool can resolve regardless of which form a
# caller / cached client is using.
self.tool_name_to_mcp_server_name_mapping[original_name] = prefix
for known_prefix in iter_known_server_prefixes(server):
qualified = add_server_prefix_to_name(original_name, known_prefix)
self.tool_name_to_mcp_server_name_mapping[qualified] = prefix
for spelling in iter_known_tool_name_spellings(original_name, server):
self.tool_name_to_mcp_server_name_mapping[spelling] = prefix
verbose_logger.info(f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}")
return prefixed_tools
@ -4261,28 +4262,29 @@ class MCPServerManager:
def check_allowed_or_banned_tools(self, tool_name: str, server: MCPServer) -> bool:
"""
Check if the tool is allowed or banned for the given server
Check if the tool is allowed or banned for the given server.
``tool_name`` is bare: every caller resolves the boundary against the server's
registered prefixes before dispatch (``server.py``'s ``original_tool_name``, the
Responses handler's ``sanitized_tool_name``). Configured entries are matched by
deriving the spellings routing accepts, never by stripping the entry, which would
cut a second boundary out of a native name that opens with the server prefix.
"""
from litellm.proxy._experimental.mcp_server.utils import (
server_applies_tool_allowlist,
)
if server_applies_tool_allowlist(server):
if not server.allowed_tools:
return False
return tool_name in server.allowed_tools or f"{server.name}-{tool_name}" in server.allowed_tools
if server.disallowed_tools:
return (
tool_name not in server.disallowed_tools and f"{server.name}-{tool_name}" not in server.disallowed_tools
)
return True
return match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None
return match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is None
def validate_allowed_params(self, tool_name: str, arguments: dict[str, Any], server: MCPServer) -> None:
"""
Filter arguments to only include allowed parameters for the given tool.
Args:
tool_name: Name of the tool (with or without prefix)
tool_name: Bare tool name, already resolved against the server's
registered prefixes by the caller
arguments: Dictionary of arguments to filter
server: MCPServer configuration
@ -4292,23 +4294,12 @@ class MCPServerManager:
Raises:
HTTPException: If allowed_params is configured for this tool but arguments contain disallowed params
"""
from litellm.proxy._experimental.mcp_server.utils import (
split_server_prefix_from_name,
)
# If no allowed_params configured, return all arguments
if not server.allowed_params:
allowed_params = server.allowed_params or {}
matched = match_known_tool_name(tool_name, server, allowed_params)
if matched is None:
return
# Get the unprefixed tool name to match against config
unprefixed_tool_name, _ = split_server_prefix_from_name(tool_name)
# Check both prefixed and unprefixed tool names
allowed_params_list = server.allowed_params.get(tool_name) or server.allowed_params.get(unprefixed_tool_name)
# If this tool doesn't have allowed_params specified, allow all params
if allowed_params_list is None:
return None
allowed_params_list = allowed_params[matched]
# Filter arguments to only include allowed parameters
disallowed_params = [param for param in arguments.keys() if param not in allowed_params_list]
@ -4390,8 +4381,11 @@ class MCPServerManager:
global_mcp_tool_registry,
)
# Get the tool from the registry
tool = global_mcp_tool_registry.get_tool(f"{server.name}-{tool_name}")
# Registration used add_server_prefix_to_name(base, get_server_prefix(server)),
# and tool_name is the bare base name by the time call_tool reaches here, so
# rebuilding the key the same way reproduces it exactly
registry_key = add_server_prefix_to_name(tool_name, get_server_prefix(server))
tool = global_mcp_tool_registry.get_tool(registry_key)
if tool is None:
# Tool not found in registry
error_msg = f"OpenAPI tool {tool_name} not found in registry"
@ -5251,7 +5245,7 @@ class MCPServerManager:
for tool in tools:
# The tool.name here is already prefixed from _get_tools_from_server
# Extract original name for mapping
original_name, _ = split_server_prefix_from_name(tool.name)
original_name = strip_known_server_prefix(tool.name, server)
self.tool_name_to_mcp_server_name_mapping[original_name] = server.name
self.tool_name_to_mcp_server_name_mapping[tool.name] = server.name
@ -5288,13 +5282,10 @@ class MCPServerManager:
# If not found and tool name is prefixed, extract the prefix and
# match against any known form.
if is_tool_name_prefixed(tool_name, known_server_prefixes=set(prefix_to_server.keys())):
(
original_tool_name,
server_name_from_prefix,
) = split_server_prefix_from_name(tool_name)
normalised_prefix = normalize_server_name(server_name_from_prefix)
matched_server = prefix_to_server.get(normalised_prefix)
matched = match_known_server_prefix(tool_name, prefix_to_server.keys())
if matched is not None:
matched_prefix, original_tool_name = matched
matched_server = prefix_to_server.get(matched_prefix)
if matched_server is not None and (
original_tool_name in self.tool_name_to_mcp_server_name_mapping
or tool_name in self.tool_name_to_mcp_server_name_mapping

View file

@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
list_fault_http_status,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
acting_user_auth,
build_effective_auth_contexts,
)
from litellm.proxy._experimental.mcp_server.utils import (
@ -99,9 +100,9 @@ if MCP_AVAILABLE:
MCPServer,
_apply_toolset_scope,
_fire_mcp_tool_call_logging,
_tool_name_matches,
execute_mcp_tool,
filter_tools_by_allowed_tools,
filter_tools_by_key_team_permissions,
)
########################################################
@ -530,19 +531,17 @@ if MCP_AVAILABLE:
tools = filter_tools_by_allowed_tools(tools, server)
# Filter by the key's effective tool permissions through the same
# primitive the MCP protocol path uses (direct grants, toolset grants,
# and team/agent/org ceilings), so REST listing cannot drift from it
# function the MCP protocol path uses (direct grants, toolset grants,
# and team/agent/org ceilings), so REST listing cannot drift from it.
# Entries here are tool names on one server, written bare by every
# writer, and dispatch compares them bare; matching a wider set of
# spellings would advertise a tool that tools/call then refuses
if user_api_key_auth:
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
allowed_tools_for_server = await MCPRequestHandler.get_allowed_tools_for_server(
tools = await filter_tools_by_key_team_permissions(
tools=tools,
server_id=server.server_id,
user_api_key_auth=user_api_key_auth,
)
if allowed_tools_for_server is not None:
tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server)]
return _create_tool_response_objects(tools, server)
@ -667,13 +666,19 @@ if MCP_AVAILABLE:
"""Coerce an Optional[str] Query param to str|None, dropping unresolved FastAPI defaults."""
return value if isinstance(value, str) else None
async def _resolve_toolset_scope(
async def _resolve_acting_auth(
toolset_name: str | None,
user_api_key_dict: UserAPIKeyAuth,
) -> UserAPIKeyAuth:
"""Resolve ``toolset_name`` to its scoped ``UserAPIKeyAuth``, or return unchanged."""
"""The one credential this tools request acts as.
A toolset name narrows the caller's own credential to that toolset; otherwise a dashboard
session is swapped for its admitted subject. The two are mutually exclusive by construction,
which is why they share an owner: the admitted subject resolves per grant source and a team
source deliberately carries none of the caller's ``object_permission``, so a toolset
narrowing layered on top would evaporate on every team-granted server."""
if not toolset_name:
return user_api_key_dict
return await acting_user_auth(user_api_key_dict)
from litellm.proxy.utils import get_prisma_client_or_throw
@ -731,6 +736,7 @@ if MCP_AVAILABLE:
try:
mcp_server_name = _as_query_str(mcp_server_name)
toolset_name = _as_query_str(toolset_name)
user_api_key_dict = await _resolve_acting_auth(toolset_name, user_api_key_dict)
# The full catalog (allowlist filter skipped) is admin-only so the
# REST endpoint can't be used to enumerate deliberately-disabled tools.
@ -738,8 +744,6 @@ if MCP_AVAILABLE:
include_disabled_tools and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
)
user_api_key_dict = await _resolve_toolset_scope(toolset_name, user_api_key_dict)
if server_id is None:
server_id = mcp_server_name
@ -928,6 +932,7 @@ if MCP_AVAILABLE:
)
try:
user_api_key_dict = await acting_user_auth(user_api_key_dict)
data = await request.json()
tool_name = data.get("name")

View file

@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
extract_mcp_tool_result_error_message,
get_server_prefix,
iter_known_server_prefixes,
match_known_tool_name,
)
from litellm.proxy._types import (
ProxyException,
@ -1416,34 +1417,17 @@ if MCP_AVAILABLE:
return allowed_mcp_servers
def _tool_name_matches(tool_name: str, filter_list: list[str]) -> bool:
def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
"""
Check if a tool name matches any name in the filter list.
Checks both the full tool name and unprefixed version (without server prefix).
This allows users to configure simple tool names regardless of prefixing.
Comparison is case-insensitive to handle OpenAPI operationIds that may be in camelCase.
Args:
tool_name: The tool name to check (may be prefixed like "server-tool_name")
filter_list: List of tool names to match against
Returns:
True if the tool name (prefixed or unprefixed) is in the filter list
Reads the same owner the server-level permission checks use, so discovery hides
exactly what dispatch refuses. ``mcp_server`` is required: guessing the boundary
at the first separator mismatches every tool on a server whose prefix contains
the separator.
"""
from litellm.proxy._experimental.mcp_server.utils import (
split_server_prefix_from_name,
)
# Normalize filter list to lowercase for case-insensitive comparison
filter_list_lower = [f.lower() for f in filter_list]
if tool_name.lower() in filter_list_lower:
return True
# Check if the unprefixed name is in the list (case-insensitive)
unprefixed_name, _ = split_server_prefix_from_name(tool_name)
return unprefixed_name.lower() in filter_list_lower
bare_name = strip_known_server_prefix(tool_name, mcp_server)
return match_known_tool_name(bare_name, mcp_server, filter_list) is not None
def filter_tools_by_allowed_tools(
tools: list[MCPTool],
@ -1473,12 +1457,16 @@ if MCP_AVAILABLE:
if server_applies_tool_allowlist(mcp_server):
if not mcp_server.allowed_tools:
return []
tools_to_return = [tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools)]
tools_to_return = [
tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server)
]
# Filter by disallowed_tools (blacklist)
if mcp_server.disallowed_tools:
tools_to_return = [
tool for tool in tools_to_return if not _tool_name_matches(tool.name, mcp_server.disallowed_tools)
tool
for tool in tools_to_return
if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server)
]
return tools_to_return
@ -1498,7 +1486,7 @@ if MCP_AVAILABLE:
return tools
for tool in tools:
unprefixed, _ = split_server_prefix_from_name(tool.name)
unprefixed = strip_known_server_prefix(tool.name, mcp_server)
lookup_key = unprefixed or tool.name
if lookup_key in display_name_map:
tool.name = display_name_map[lookup_key]
@ -2317,14 +2305,16 @@ if MCP_AVAILABLE:
server_id=server_id,
user_api_key_auth=user_api_key_auth,
)
if allowed_tool_names is None:
return tools
# Tools arrive prefixed with the server's own prefix; strip exactly that
# prefix (resolved from the server) rather than the first separator, so a
# prefix containing the separator still reduces to the stored bare name.
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
return [t for t in tools if strip_known_server_prefix(t.name, server) in allowed_tool_names]
return [
t
for t in tools
if MCPRequestHandler.tool_is_granted(strip_known_server_prefix(t.name, server), allowed_tool_names)
]
async def _list_mcp_tools(
user_api_key_auth: UserAPIKeyAuth | None = None,
@ -2699,6 +2689,7 @@ if MCP_AVAILABLE:
break
if mcp_server is not None:
server_name = mcp_server.name
original_tool_name = strip_known_server_prefix(name, mcp_server)
if requested_server is not None:
if mcp_server is not None and mcp_server.server_id != requested_server.server_id:
@ -2716,6 +2707,7 @@ if MCP_AVAILABLE:
if mcp_server is None:
mcp_server = requested_server
server_name = requested_server.name
original_tool_name = strip_known_server_prefix(name, requested_server)
# Only enforce server-level permissions when we can resolve a server
if server_name:
@ -2887,13 +2879,14 @@ if MCP_AVAILABLE:
_request_resolved_auth_headers.reset(_resolved_token)
response = CallToolResult(content=cast(Any, local_content), isError=False)
# Try managed MCP server tool (pass the full prefixed name)
# Try managed MCP server tool (the name is bare; the prefix boundary was
# already resolved above against this server's registered prefixes)
# Primary and recommended way to use external MCP servers
#########################################################
elif mcp_server:
response = await _handle_managed_mcp_tool(
server_name=server_name,
name=original_tool_name, # Pass the full name (potentially prefixed)
name=original_tool_name,
arguments=arguments,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,

View file

@ -1,9 +1,11 @@
"""Helpers to resolve real team contexts for UI session tokens."""
"""Helpers to resolve the identity a dashboard UI session token acts as."""
from __future__ import annotations
from typing import List
from fastapi import HTTPException
from litellm._logging import verbose_logger
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import UserAPIKeyAuth
@ -23,12 +25,19 @@ def clone_user_api_key_auth_with_team(
return cloned_auth
def is_ui_session_credential(user_api_key_auth: UserAPIKeyAuth) -> bool:
"""Whether the caller is the dashboard's SSO-minted session token acting as its user,
the only credential shape allowed to widen a request to the owning user's identity."""
return user_api_key_auth.team_id == UI_SESSION_TOKEN_TEAM_ID and bool(user_api_key_auth.user_id)
async def resolve_ui_session_team_ids(
user_api_key_auth: UserAPIKeyAuth,
) -> List[str]:
"""Resolve the real team ids backing a UI session token."""
if user_api_key_auth.team_id != UI_SESSION_TOKEN_TEAM_ID or not user_api_key_auth.user_id:
if not is_ui_session_credential(user_api_key_auth):
return []
from litellm.proxy.auth.auth_checks import get_user_object
@ -68,12 +77,63 @@ async def resolve_ui_session_team_ids(
return resolved_team_ids
async def admitted_user_context(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None:
"""THE owner of "resolve this dashboard session's user identity": the same admitted-subject auth a
gateway OAuth session for this user resolves with, carrying the user row's own object permission,
on this request's tracing span. None for any other credential (a caller-passed key is never
widened) and on reload failure, which every caller reads as "no user-level identity available"."""
user_id = user_api_key_auth.user_id
if not is_ui_session_credential(user_api_key_auth) or user_id is None:
return None
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
try:
admitted = await MCPRequestHandler._reload_admitted_user(user_id)
except HTTPException as e:
verbose_logger.warning(f"MCP dashboard session: admitted-subject reload failed for {user_id}: {e.detail}")
return None
return admitted.model_copy(update={"parent_otel_span": user_api_key_auth.parent_otel_span})
async def acting_user_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth:
"""The principal acting-as-user MCP routes resolve permissions with. A non-admin dashboard
session acts as the admitted subject, the same identity a gateway session resolves with, so
server reachability, per-source tool ceilings, rate limits, and billing bind identically on
both surfaces. An admin session keeps its operator view and any caller-passed credential is
returned unchanged, never widened.
Do not combine this with a narrowing that rewrites a single credential's ``object_permission``
(toolset scope): the admitted subject resolves per grant source and a team source deliberately
carries none of the caller's own grants, so the narrowing would silently evaporate on every
team-granted server. A request carrying such a scope keeps the caller's own credential."""
if not is_ui_session_credential(user_api_key_auth):
return user_api_key_auth
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
if _user_has_admin_view(user_api_key_auth):
return user_api_key_auth
admitted = await admitted_user_context(user_api_key_auth)
return admitted if admitted is not None else user_api_key_auth
async def build_effective_auth_contexts(
user_api_key_auth: UserAPIKeyAuth,
) -> List[UserAPIKeyAuth]:
"""Return auth contexts that reflect the actual teams for UI session tokens."""
"""Every auth context a management or listing surface must resolve a UI session token through:
one per real team backing the session, plus the session user's own admitted identity, so a grant
made directly to the user row is as visible to the dashboard as it is to a gateway session."""
resolved_team_ids = await resolve_ui_session_team_ids(user_api_key_auth)
if resolved_team_ids:
return [clone_user_api_key_auth_with_team(user_api_key_auth, team_id) for team_id in resolved_team_ids]
return [user_api_key_auth]
team_contexts = (
[clone_user_api_key_auth_with_team(user_api_key_auth, team_id) for team_id in resolved_team_ids]
if resolved_team_ids
else [user_api_key_auth]
)
admitted_context = await admitted_user_context(user_api_key_auth)
if admitted_context is None:
return team_contexts
return [*team_contexts, admitted_context]

View file

@ -2,7 +2,10 @@
MCP Server Utilities
"""
import hashlib
import importlib
import json
import os
import re
from collections.abc import MutableMapping, MutableSequence
from typing import (
@ -17,12 +20,10 @@ from typing import (
Tuple,
Union,
)
import hashlib
import importlib
import os
from urllib.parse import quote
from litellm.types.mcp_server.mcp_server_manager import MCPServer
# Constants
#
# NOTE: The environment-backed values below are read once, when this module is
@ -326,8 +327,59 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]:
yield from _emit(server_id)
def iter_known_tool_name_spellings(tool_name: str, server: MCPServer) -> Iterator[str]:
"""Yield every name that denotes the bare ``tool_name`` on ``server``: the bare name,
then its wire spelling under each prefix ``iter_known_server_prefixes`` accepts.
``get_server_prefix`` covers only the currently published one, and that moves with the
alias and with ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``.
"""
yield tool_name
for prefix in iter_known_server_prefixes(server):
yield add_server_prefix_to_name(tool_name, prefix)
def openapi_tool_name(operation_id: str) -> str:
"""Return the tool name ``_register_openapi_tools`` registers ``operation_id`` under.
The single transform between a spec's operationId and the name the gateway serves.
Policy recovers the link by replaying this exact function, which is what keeps it from
deciding for a tool it does not name: two operationIds that register as two tools
necessarily normalize to two names here, because this is the map that registered them.
"""
return operation_id.replace(" ", "_").lower()
def match_known_tool_name(tool_name: str, server: MCPServer, names: Iterable[str]) -> str | None:
"""Return the entry of ``names`` that denotes ``tool_name`` on ``server``, else ``None``.
The single question every tool-name-keyed site asks: the allow list, the deny list,
``allowed_params`` and the discovery filter, so discovery hides exactly what dispatch
refuses. It spans every spelling routing accepts and no more, because a tool's identity
is the exact name routing dispatches; anything looser lets one policy decide two tools.
On an OpenAPI server the configured entry holds the spec's operationId while routing
holds :func:`openapi_tool_name` of it, so both sides go through that map first. Doing it
with the registering function rather than a lookalike is the whole safety argument: a
coarser one collapses operationIds that registration keeps apart.
Callers read the returned entry rather than testing a container's values, which is what
stops an explicitly empty ``allowed_params`` list from reading as "nothing configured".
"""
normalize = openapi_tool_name if getattr(server, "spec_path", None) else str
spellings = {normalize(spelling) for spelling in iter_known_tool_name_spellings(tool_name, server)}
return next((name for name in names if normalize(name) in spellings), None)
def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
"""Return the unprefixed name plus the server name used as prefix."""
"""Return the unprefixed name plus the server name used as prefix.
Cuts at the FIRST separator, so the two halves are only trustworthy as a
pair: they reassemble into ``prefixed_name`` exactly, which is what makes
this safe for routing. Reading one half on its own is a guess about where the
boundary fell, and that guess is wrong whenever the prefix itself contains
the separator. Callers that compare a half against configuration must use
:func:`match_known_server_prefix` or :func:`strip_known_server_prefix`.
"""
if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name:
parts = prefixed_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1)
if len(parts) == 2:
@ -335,6 +387,27 @@ def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
return prefixed_name, ""
def match_known_server_prefix(name: str, known_prefixes: Iterable[str]) -> tuple[str, str] | None:
"""Return ``(matched_prefix, bare_name)`` when ``name`` carries a known prefix.
Candidates are normalized and tried LONGEST first, so a prefix that itself
contains :data:`MCP_TOOL_PREFIX_SEPARATOR` (the UUID ``server_id`` used when
a server has no alias, or a legacy hyphenated alias) wins over a shorter
prefix that is merely its leading segment. Returns ``None`` when no candidate
matches, i.e. ``name`` carries none of these prefixes.
"""
candidates = sorted(
{normalize_server_name(prefix) for prefix in known_prefixes if prefix},
key=len,
reverse=True,
)
for prefix in candidates:
separator_suffixed = prefix + MCP_TOOL_PREFIX_SEPARATOR
if name.startswith(separator_suffixed):
return prefix, name[len(separator_suffixed) :]
return None
def strip_known_server_prefix(name: str, server: Optional[Any]) -> str:
"""Strip ``server``'s registered prefix from a prefixed tool/resource name.
@ -352,11 +425,8 @@ def strip_known_server_prefix(name: str, server: Optional[Any]) -> str:
"""
if server is None:
return split_server_prefix_from_name(name)[0]
for prefix in iter_known_server_prefixes(server):
candidate = normalize_server_name(prefix) + MCP_TOOL_PREFIX_SEPARATOR
if name.startswith(candidate):
return name[len(candidate) :]
return name
matched = match_known_server_prefix(name, iter_known_server_prefixes(server))
return name if matched is None else matched[1]
def is_tool_name_prefixed(
@ -367,15 +437,16 @@ def is_tool_name_prefixed(
Check if tool name has a known MCP server prefix.
When ``known_server_prefixes`` is provided the function verifies that the
substring before the first separator is an actual registered server
prefix. Without it the check falls back to the legacy heuristic
name actually starts with one of those prefixes followed by the separator,
matching the longest candidate first so a prefix containing the separator
still resolves. Without it the check falls back to the legacy heuristic
(separator present anywhere in the name), which can produce false
positives for non-MCP tools whose names contain hyphens
(e.g. ``text-to-speech``, ``code-review``).
Args:
tool_name: Tool name to check.
known_server_prefixes: Optional set of normalised server prefixes
known_server_prefixes: Optional set of normalized server prefixes
currently registered in the MCP manager. Pass this whenever
the caller has access to the server registry so that the check
is accurate.
@ -387,8 +458,7 @@ def is_tool_name_prefixed(
return False
if known_server_prefixes is not None:
candidate_prefix = tool_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1)[0]
return normalize_server_name(candidate_prefix) in known_server_prefixes
return match_known_server_prefix(tool_name, known_server_prefixes) is not None
# Legacy fallback – separator present somewhere in the name.
return True

View file

@ -20402,6 +20402,16 @@
"description": "Who created the attachment.",
"title": "Created By"
},
"definition_location": {
"default": "db",
"description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).",
"enum": [
"db",
"config"
],
"title": "Definition Location",
"type": "string"
},
"keys": {
"description": "Key patterns.",
"items": {
@ -20658,6 +20668,16 @@
"description": "Who created the policy.",
"title": "Created By"
},
"definition_location": {
"default": "db",
"description": "Where this policy is defined: 'db' (database) or 'config' (config.yaml).",
"enum": [
"db",
"config"
],
"title": "Definition Location",
"type": "string"
},
"description": {
"anyOf": [
{
@ -21129,12 +21149,45 @@
"title": "PolicyVersionStatusUpdateRequest",
"type": "object"
},
"UsageChartPoint": {
"properties": {
"blocked": {
"title": "Blocked",
"type": "integer"
},
"date": {
"title": "Date",
"type": "string"
},
"passed": {
"title": "Passed",
"type": "integer"
},
"score": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Score"
}
},
"required": [
"date",
"passed",
"blocked"
],
"title": "UsageChartPoint",
"type": "object"
},
"UsageOverviewResponse": {
"properties": {
"chart": {
"items": {
"additionalProperties": true,
"type": "object"
"$ref": "#/components/schemas/UsageChartPoint"
},
"title": "Chart",
"type": "array"
@ -21243,6 +21296,13 @@
},
"ValidationError": {
"properties": {
"ctx": {
"title": "Context",
"type": "object"
},
"input": {
"title": "Input"
},
"loc": {
"items": {
"anyOf": [
@ -21420,7 +21480,7 @@
},
"/policies/attachments/list": {
"get": {
"description": "List all policy attachments from the database.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"attachments\": [\n {\n \"attachment_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"scope\": \"*\",\n \"teams\": [],\n \"keys\": [],\n \"models\": [],\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
"description": "List all policy attachments from the database and config.yaml.\n\nConfig-defined attachments are returned with definition_location \"config\" and a\nsynthetic attachment_id (\"config-<index>\").\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"attachments\": [\n {\n \"attachment_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"scope\": \"*\",\n \"teams\": [],\n \"keys\": [],\n \"models\": [],\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
"operationId": "list_policy_attachments_policies_attachments_list_get",
"responses": {
"200": {
@ -21596,7 +21656,7 @@
},
"/policies/list": {
"get": {
"description": "List all policies from the database. Optionally filter by version_status.\n\nQuery params:\n- version_status: Optional. One of \"draft\", \"published\", \"production\".\n If omitted, all versions are returned.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"policies\": [\n {\n \"policy_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"version_number\": 1,\n \"version_status\": \"production\",\n \"inherit\": null,\n \"description\": \"Base guardrails for all requests\",\n \"guardrails_add\": [\"pii_masking\"],\n \"guardrails_remove\": [],\n \"condition\": null,\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
"description": "List all policies from the database and config.yaml. Optionally filter by version_status.\n\nConfig-defined policies are returned with definition_location \"config\" and are treated\nas production versions. On a name conflict with a DB policy, only the DB policy is returned.\n\nQuery params:\n- version_status: Optional. One of \"draft\", \"published\", \"production\".\n If omitted, all versions are returned.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"policies\": [\n {\n \"policy_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"version_number\": 1,\n \"version_status\": \"production\",\n \"inherit\": null,\n \"description\": \"Base guardrails for all requests\",\n \"guardrails_add\": [\"pii_masking\"],\n \"guardrails_remove\": [],\n \"condition\": null,\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
"operationId": "list_policies_policies_list_get",
"parameters": [
{

View file

@ -148,7 +148,9 @@ if MCP_AVAILABLE:
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
admitted_user_context,
build_effective_auth_contexts,
is_ui_session_credential,
)
from litellm.proxy._types import (
LiteLLM_MCPServerTable,
@ -939,6 +941,16 @@ if MCP_AVAILABLE:
aggregated.setdefault(server.server_id, server)
return list(aggregated.values())
async def _connected_app_reachable_server_ids(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]:
"""Server ids a connected app authorized by this dashboard user is served on the aggregate
MCP endpoint, resolved through the one owner of the admitted subject so the page and the
session cannot drift. Empty when that identity cannot be built, which is the true answer:
the same user cannot open a gateway session either."""
admitted = await admitted_user_context(user_api_key_dict)
if admitted is None:
return frozenset()
return frozenset(await global_mcp_server_manager.get_allowed_mcp_servers(admitted))
@router.get(
"/server",
description="Returns the mcp server list with associated teams",
@ -953,6 +965,12 @@ if MCP_AVAILABLE:
"servers the team has access to plus globally available (allow_all_keys) servers. "
"Used by the Create Key UI to show team-scoped MCP servers.",
),
connected_app_view: bool = Query(
False,
description="Annotate each returned server with connected_app_reachable: whether a "
"connected app authorized by the calling user (a gateway OAuth session) is served "
"this server on the aggregate MCP endpoint.",
),
):
"""
Get all of the configured mcp servers for the user in the db with their associated teams
@ -1009,6 +1027,11 @@ if MCP_AVAILABLE:
servers = await _resolve_accessible_mcp_servers(user_api_key_dict)
redacted_mcp_servers = _redact_mcp_credentials_list(servers)
if connected_app_view is True and is_ui_session_credential(user_api_key_dict):
reachable_ids = await _connected_app_reachable_server_ids(user_api_key_dict)
for server in redacted_mcp_servers:
server.connected_app_reachable = server.server_id in reachable_ids
# augment the mcp servers with public status
if litellm.public_mcp_servers is not None:
for server in redacted_mcp_servers:

View file

@ -42,6 +42,7 @@ class AttachmentRegistry:
def __init__(self):
self._attachments: List[PolicyAttachment] = []
self._config_attachments: tuple[PolicyAttachment, ...] = ()
self._initialized: bool = False
def load_attachments(self, attachments_config: List[Dict[str, Any]]) -> None:
@ -62,6 +63,7 @@ class AttachmentRegistry:
verbose_proxy_logger.error(f"Error loading attachment: {str(e)}")
raise ValueError(f"Invalid attachment: {str(e)}") from e
self._config_attachments = tuple(self._attachments)
self._initialized = True
verbose_proxy_logger.info(f"Loaded {len(self._attachments)} policy attachments")
@ -173,6 +175,15 @@ class AttachmentRegistry:
"""
return self._attachments.copy()
def get_config_attachments(self) -> tuple[PolicyAttachment, ...]:
"""
Get the attachments loaded from config.yaml.
Returns:
Tuple of config-defined PolicyAttachment objects
"""
return self._config_attachments
def get_attachments_for_policy(self, policy_name: str) -> List[PolicyAttachment]:
"""
Get all attachments for a specific policy.
@ -199,6 +210,7 @@ class AttachmentRegistry:
Clear all attachments from the registry.
"""
self._attachments = []
self._config_attachments = ()
self._initialized = False
def add_attachment(self, attachment: PolicyAttachment) -> None:
@ -428,6 +440,7 @@ class AttachmentRegistry:
) -> None:
"""
Sync policy attachments from the database to in-memory registry.
Config-loaded attachments are preserved.
Args:
prisma_client: The Prisma client instance
@ -435,11 +448,8 @@ class AttachmentRegistry:
try:
attachments = await self.get_all_attachments_from_db(prisma_client)
# Clear existing attachments and reload from DB
self._attachments = []
for attachment_response in attachments:
attachment = PolicyAttachment(
db_attachments = [
PolicyAttachment(
policy=attachment_response.policy_name,
scope=attachment_response.scope,
teams=(attachment_response.teams if attachment_response.teams else None),
@ -447,10 +457,15 @@ class AttachmentRegistry:
models=(attachment_response.models if attachment_response.models else None),
tags=attachment_response.tags if attachment_response.tags else None,
)
self._attachments.append(attachment)
for attachment_response in attachments
]
self._attachments = [*self._config_attachments, *db_attachments]
self._initialized = True
verbose_proxy_logger.info(f"Synced {len(attachments)} attachments from DB to in-memory registry")
verbose_proxy_logger.info(
f"Synced {len(attachments)} attachments from DB to in-memory registry "
f"({len(self._config_attachments)} config-defined attachments preserved)"
)
except Exception as e:
verbose_proxy_logger.exception(f"Error syncing attachments from DB: {e}")
raise Exception(f"Error syncing attachments from DB: {str(e)}")

View file

@ -17,6 +17,8 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import (
GuardrailPipeline,
PipelineTestRequest,
Policy,
PolicyAttachment,
PolicyAttachmentCreateRequest,
PolicyAttachmentDBResponse,
PolicyAttachmentListResponse,
@ -33,6 +35,35 @@ from litellm.types.proxy.policy_engine import (
router = APIRouter()
def _config_policy_to_db_response(policy_name: str, policy: Policy) -> PolicyDBResponse:
return PolicyDBResponse(
policy_id=policy_name,
policy_name=policy_name,
version_number=1,
version_status="production",
inherit=policy.inherit,
description=policy.description,
guardrails_add=policy.guardrails.get_add(),
guardrails_remove=policy.guardrails.get_remove(),
condition=policy.condition.model_dump() if policy.condition else None,
pipeline=policy.pipeline.model_dump() if policy.pipeline else None,
definition_location="config",
)
def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment) -> PolicyAttachmentDBResponse:
return PolicyAttachmentDBResponse(
attachment_id=f"config-{index}",
policy_name=attachment.policy,
scope=attachment.scope,
teams=attachment.teams or [],
keys=attachment.keys or [],
models=attachment.models or [],
tags=attachment.tags or [],
definition_location="config",
)
# ─────────────────────────────────────────────────────────────────────────────
# Policy CRUD Endpoints
# ─────────────────────────────────────────────────────────────────────────────
@ -46,7 +77,13 @@ router = APIRouter()
)
async def list_policies(version_status: Optional[str] = None):
"""
List all policies from the database. Optionally filter by version_status.
List all policies from the database and config.yaml. Optionally filter by version_status.
Config-defined policies are returned with definition_location "config" and are treated
as production versions. On a name conflict with a production DB policy, only the DB policy
is returned, mirroring runtime resolution where only production DB versions override config.
A draft or published DB version does not hide the config policy, since the config version
is still the one being enforced.
Query params:
- version_status: Optional. One of "draft", "published", "production".
@ -84,11 +121,27 @@ async def list_policies(version_status: Optional[str] = None):
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
policies = await get_policy_registry().get_all_policies_from_db(prisma_client, version_status=version_status)
registry = get_policy_registry()
db_policies = (
await registry.get_all_policies_from_db(prisma_client, version_status=version_status)
if prisma_client is not None
else []
)
db_policy_names = {
db_policy.policy_name for db_policy in db_policies if db_policy.version_status == "production"
}
include_config = version_status in (None, "production")
config_policies = (
[
_config_policy_to_db_response(policy_name, policy)
for policy_name, policy in registry.list_config_policies().items()
if policy_name not in db_policy_names
]
if include_config
else []
)
policies = db_policies + config_policies
return PolicyListDBResponse(policies=policies, total_count=len(policies))
except Exception as e:
verbose_proxy_logger.exception(f"Error listing policies: {e}")
@ -606,7 +659,10 @@ async def test_pipeline(
)
async def list_policy_attachments():
"""
List all policy attachments from the database.
List all policy attachments from the database and config.yaml.
Config-defined attachments are returned with definition_location "config" and a
synthetic attachment_id ("config-<index>").
Example Request:
```bash
@ -635,11 +691,14 @@ async def list_policy_attachments():
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
attachments = await get_attachment_registry().get_all_attachments_from_db(prisma_client)
registry = get_attachment_registry()
db_attachments = await registry.get_all_attachments_from_db(prisma_client) if prisma_client is not None else []
config_attachments = [
_config_attachment_to_db_response(index, attachment)
for index, attachment in enumerate(registry.get_config_attachments())
]
attachments = db_attachments + config_attachments
return PolicyAttachmentListResponse(attachments=attachments, total_count=len(attachments))
except Exception as e:
verbose_proxy_logger.exception(f"Error listing policy attachments: {e}")

View file

@ -13,6 +13,7 @@ from datetime import datetime, timezone
from typing import (
TYPE_CHECKING,
Any,
Literal,
Optional,
Protocol,
TypedDict,
@ -162,6 +163,8 @@ class PolicyRegistry:
def __init__(self):
self._policies: dict[str, Policy] = {}
self._config_policies: Mapping[str, Policy] = {}
self._sources: Mapping[str, Literal["db", "config"]] = {}
self._policies_by_id: dict[str, tuple[str, Policy]] = {}
self._initialized: bool = False
@ -174,6 +177,8 @@ class PolicyRegistry:
This is the raw config from the YAML file.
"""
self._policies = {}
self._config_policies = {}
self._sources = {}
self._policies_by_id = {}
for policy_name, policy_data in policies_config.items():
@ -185,6 +190,8 @@ class PolicyRegistry:
verbose_proxy_logger.error(f"Error loading policy '{policy_name}': {str(e)}")
raise ValueError(f"Invalid policy '{policy_name}': {str(e)}") from e
self._config_policies = dict(self._policies)
self._sources = {policy_name: "config" for policy_name in self._policies}
self._initialized = True
verbose_proxy_logger.info(f"Loaded {len(self._policies)} policies")
@ -299,23 +306,42 @@ class PolicyRegistry:
Clear all policies from the registry.
"""
self._policies = {}
self._config_policies = {}
self._sources = {}
self._initialized = False
def add_policy(self, policy_name: str, policy: Policy) -> None:
def get_source(self, policy_name: str) -> Optional[Literal["db", "config"]]:
"""
Return the provenance of an in-memory policy, or None if unknown.
"""
return self._sources.get(policy_name)
def list_config_policies(self) -> Mapping[str, Policy]:
"""
Return the policies loaded from config.yaml, keyed by policy name.
"""
return dict(self._config_policies)
def add_policy(self, policy_name: str, policy: Policy, source: Literal["db", "config"] = "db") -> None:
"""
Add or update a single policy.
Args:
policy_name: Name of the policy
policy: Policy object to add
source: Provenance of the policy ("db" or "config")
"""
self._policies[policy_name] = policy
self._sources = {**self._sources, policy_name: source}
if source == "config":
self._config_policies = {**self._config_policies, policy_name: policy}
self._initialized = True
verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}")
def remove_policy(self, policy_name: str) -> bool:
"""
Remove a policy by name.
Remove a policy by name. If a config-defined policy shares the name,
it is restored immediately instead of waiting for the next DB sync.
Args:
policy_name: Name of the policy to remove
@ -323,11 +349,18 @@ class PolicyRegistry:
Returns:
True if policy was removed, False if it didn't exist
"""
if policy_name in self._policies:
del self._policies[policy_name]
verbose_proxy_logger.debug(f"Removed policy: {policy_name}")
if policy_name not in self._policies:
return False
config_fallback = self._config_policies.get(policy_name)
if config_fallback is not None:
self._policies[policy_name] = config_fallback
self._sources = {**self._sources, policy_name: "config"}
verbose_proxy_logger.debug(f"Removed policy: {policy_name}; restored config-defined version")
return True
return False
del self._policies[policy_name]
self._sources = {name: source for name, source in self._sources.items() if name != policy_name}
verbose_proxy_logger.debug(f"Removed policy: {policy_name}")
return True
# ─────────────────────────────────────────────────────────────────────────
# Database CRUD Methods
@ -501,10 +534,15 @@ class PolicyRegistry:
# Remove from in-memory registry only if this was the production version
if version_status == "production":
self.remove_policy(policy_name)
result["warning"] = (
"Production version was deleted. No other version was promoted. "
"Promote another version to production if this policy should remain active."
)
if self.get_source(policy_name) == "config":
result["warning"] = (
"Production version was deleted. The config-defined policy with the same name is active again."
)
else:
result["warning"] = (
"Production version was deleted. No other version was promoted. "
"Promote another version to production if this policy should remain active."
)
return result
except Exception as e:
@ -591,14 +629,14 @@ class PolicyRegistry:
"""
Sync policies from the database to in-memory registry.
- Production versions are loaded into _policies (by policy name) for resolution.
- Config-loaded policies are preserved; on a name conflict the DB version wins.
- Draft and published versions are loaded into _policies_by_id so request-body
policy_<uuid> overrides can be resolved without DB access in the hot path.
"""
try:
self._policies = {}
production = await self.get_all_policies_from_db(prisma_client, version_status="production")
for policy_response in production:
policy = self._parse_policy(
db_policies = {
policy_response.policy_name: self._parse_policy(
policy_response.policy_name,
{
"inherit": policy_response.inherit,
@ -611,7 +649,16 @@ class PolicyRegistry:
"pipeline": policy_response.pipeline,
},
)
self.add_policy(policy_response.policy_name, policy)
for policy_response in production
}
for policy_name in set(db_policies) & set(self._config_policies):
verbose_proxy_logger.warning(
f"Policy '{policy_name}' is defined in both config.yaml and the DB; the DB version takes precedence"
)
config_sources: Mapping[str, Literal["db", "config"]] = {name: "config" for name in self._config_policies}
db_sources: Mapping[str, Literal["db", "config"]] = {name: "db" for name in db_policies}
self._policies = {**self._config_policies, **db_policies}
self._sources = {**config_sources, **db_sources}
self._policies_by_id = {}
non_production = await _policy_table(prisma_client).find_many(
@ -637,7 +684,8 @@ class PolicyRegistry:
self._initialized = True
verbose_proxy_logger.info(
f"Synced {len(production)} production policies and {len(non_production)} "
"draft/published (by ID) from DB to in-memory registry"
"draft/published (by ID) from DB to in-memory registry "
f"({len(self._config_policies)} config-defined policies preserved)"
)
except Exception as e:
verbose_proxy_logger.exception(f"Error syncing policies from DB: {e}")
@ -983,12 +1031,20 @@ class PolicyRegistry:
prisma_client: The Prisma client instance
Returns:
Dict with success message
Dict with "message" and optional "warning" if a config-defined policy took over.
"""
try:
await _policy_table(prisma_client).delete_many(where={"policy_name": policy_name})
self.remove_policy(policy_name)
return {"message": f"All versions of policy '{policy_name}' deleted successfully"}
message = f"All versions of policy '{policy_name}' deleted successfully"
if self.get_source(policy_name) == "config":
return {
"message": message,
"warning": (
"All DB versions were deleted. The config-defined policy with the same name is active again."
),
}
return {"message": message}
except Exception as e:
verbose_proxy_logger.exception(f"Error deleting all versions: {e}")
raise Exception(f"Error deleting all versions: {str(e)}")

View file

@ -7,7 +7,8 @@ import traceback
import uuid
from datetime import datetime
from functools import lru_cache
from typing import Any, Dict, List, Literal, Optional
from types import MappingProxyType
from typing import Any, Dict, List, Literal, Mapping, Optional
import httpx
from openai._streaming import SSEDecoder
@ -48,13 +49,32 @@ def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -
verbose_logger.error("%s failed: %s", task_name, exception)
_CLIENT_ERROR_CODES: frozenset[str] = frozenset(
(
"invalid_request_error",
"context_length_exceeded",
"content_policy_violation",
"model_not_found",
)
_ERROR_CODE_HTTP_STATUS: Mapping[str, int] = MappingProxyType(
{ # mutable-ok: immediately frozen by MappingProxyType
"server_error": 500,
"rate_limit_exceeded": 429,
"insufficient_quota": 429,
"vector_store_timeout": 504,
"invalid_prompt": 400,
"invalid_image": 400,
"invalid_image_format": 400,
"invalid_base64_image": 400,
"invalid_image_url": 400,
"image_too_large": 400,
"image_too_small": 400,
"image_parse_error": 400,
"image_content_policy_violation": 400,
"invalid_image_mode": 400,
"image_file_too_large": 400,
"unsupported_image_media_type": 400,
"empty_image_file": 400,
"failed_to_download_image": 400,
"image_file_not_found": 400,
"invalid_request_error": 400,
"context_length_exceeded": 400,
"content_policy_violation": 400,
"model_not_found": 400,
}
)
@ -78,12 +98,13 @@ def _error_event_fields(error_obj: object) -> tuple[str, Optional[str], Optional
def _status_code_for_error_fields(error_type: Optional[str], error_code: Optional[str]) -> int:
fields = tuple(field for field in (error_type, error_code) if field is not None)
fields = tuple(field for field in (error_code, error_type) if field is not None)
if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields):
return 429
if any(field in _CLIENT_ERROR_CODES for field in fields):
return 400
return 500
return next(
(_ERROR_CODE_HTTP_STATUS[field] for field in fields if field in _ERROR_CODE_HTTP_STATUS),
500,
)
class BaseResponsesAPIStreamingIterator:

View file

@ -6,7 +6,7 @@ the final guardrails list.
"""
from datetime import datetime
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, ConfigDict, Field
@ -220,6 +220,10 @@ class PolicyDBResponse(BaseModel):
updated_at: Optional[datetime] = Field(default=None, description="When the policy was last updated.")
created_by: Optional[str] = Field(default=None, description="Who created the policy.")
updated_by: Optional[str] = Field(default=None, description="Who last updated the policy.")
definition_location: Literal["db", "config"] = Field(
default="db",
description="Where this policy is defined: 'db' (database) or 'config' (config.yaml).",
)
class PolicyListDBResponse(BaseModel):
@ -317,6 +321,10 @@ class PolicyAttachmentDBResponse(BaseModel):
updated_at: Optional[datetime] = Field(default=None, description="When the attachment was last updated.")
created_by: Optional[str] = Field(default=None, description="Who created the attachment.")
updated_by: Optional[str] = Field(default=None, description="Who last updated the attachment.")
definition_location: Literal["db", "config"] = Field(
default="db",
description="Where this attachment is defined: 'db' (database) or 'config' (config.yaml).",
)
class PolicyAttachmentListResponse(BaseModel):

View file

@ -200,7 +200,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing
cache_creation_input_token_cost: Optional[float]
cache_creation_input_token_cost_above_200k_tokens: Optional[float]
cache_creation_input_token_cost_above_272k_tokens: Optional[float]
cache_creation_input_token_cost_above_272k_tokens_priority: Optional[float]
cache_creation_input_token_cost_above_272k_tokens_flex: Optional[float]
cache_creation_input_token_cost_above_1hr: Optional[float]
cache_creation_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing
cache_creation_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing
cache_read_input_token_cost: Optional[float]
cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing
cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing
@ -208,6 +213,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
cache_read_input_token_cost_above_200k_tokens_priority: Optional[float]
cache_read_input_token_cost_above_272k_tokens: Optional[float]
cache_read_input_token_cost_above_272k_tokens_priority: Optional[float]
cache_read_input_token_cost_above_272k_tokens_flex: Optional[float]
cache_read_input_token_cost_above_512k_tokens: Optional[float]
# Smallest prefix this model will actually cache, whatever caching mechanism its provider uses.
# Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT.
@ -219,6 +225,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
input_cost_per_token_above_200k_tokens_priority: Optional[float]
input_cost_per_token_above_272k_tokens: Optional[float] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input
input_cost_per_token_above_272k_tokens_priority: Optional[float]
input_cost_per_token_above_272k_tokens_flex: Optional[float]
input_cost_per_token_above_512k_tokens: Optional[float] # MiniMax-M3: prompts >512K priced at 2x input
input_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models
input_cost_per_query: Optional[float] # only for rerank models
@ -246,6 +253,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_token_above_200k_tokens_priority: Optional[float]
output_cost_per_token_above_272k_tokens: Optional[float] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output
output_cost_per_token_above_272k_tokens_priority: Optional[float]
output_cost_per_token_above_272k_tokens_flex: Optional[float]
output_cost_per_token_above_512k_tokens: Optional[float] # MiniMax-M3: prompts >512K priced at 2x output
output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models
output_cost_per_image: Optional[float]
@ -3182,6 +3190,11 @@ class CustomPricingLiteLLMParams(BaseModel):
cache_creation_input_token_cost: Optional[float] = None
cache_creation_input_token_cost_above_1hr: Optional[float] = None
cache_creation_input_token_cost_above_200k_tokens: Optional[float] = None
cache_creation_input_token_cost_above_272k_tokens: Optional[float] = None
cache_creation_input_token_cost_above_272k_tokens_priority: Optional[float] = None
cache_creation_input_token_cost_above_272k_tokens_flex: Optional[float] = None
cache_creation_input_token_cost_flex: Optional[float] = None
cache_creation_input_token_cost_priority: Optional[float] = None
cache_creation_input_audio_token_cost: Optional[float] = None
cache_read_input_token_cost: Optional[float] = None
cache_read_input_token_cost_flex: Optional[float] = None
@ -3189,6 +3202,7 @@ class CustomPricingLiteLLMParams(BaseModel):
cache_read_input_token_cost_above_200k_tokens: Optional[float] = None
cache_read_input_token_cost_above_200k_tokens_priority: Optional[float] = None
cache_read_input_token_cost_above_272k_tokens_priority: Optional[float] = None
cache_read_input_token_cost_above_272k_tokens_flex: Optional[float] = None
cache_read_input_audio_token_cost: Optional[float] = None
input_cost_per_character: Optional[float] = None
input_cost_per_character_above_128k_tokens: Optional[float] = None
@ -3198,6 +3212,7 @@ class CustomPricingLiteLLMParams(BaseModel):
input_cost_per_token_above_200k_tokens: Optional[float] = None
input_cost_per_token_above_200k_tokens_priority: Optional[float] = None
input_cost_per_token_above_272k_tokens_priority: Optional[float] = None
input_cost_per_token_above_272k_tokens_flex: Optional[float] = None
input_cost_per_query: Optional[float] = None
input_cost_per_image: Optional[float] = None
input_cost_per_image_above_128k_tokens: Optional[float] = None
@ -3217,6 +3232,7 @@ class CustomPricingLiteLLMParams(BaseModel):
output_cost_per_token_above_200k_tokens: Optional[float] = None
output_cost_per_token_above_200k_tokens_priority: Optional[float] = None
output_cost_per_token_above_272k_tokens_priority: Optional[float] = None
output_cost_per_token_above_272k_tokens_flex: Optional[float] = None
output_cost_per_character_above_128k_tokens: Optional[float] = None
output_cost_per_image: Optional[float] = None
output_cost_per_image_token: Optional[float] = None
@ -3851,6 +3867,7 @@ class ServiceTier(Enum):
AUTO = "auto"
FLEX = "flex"
PRIORITY = "priority"
FAST = "fast"
class DataResidency(Enum):

View file

@ -1048,7 +1048,7 @@ def function_setup(
if "metadata" in kwargs:
litellm_params["metadata"] = kwargs["metadata"]
if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict):
litellm_params["litellm_metadata"] = kwargs["litellm_metadata"].copy()
litellm_params["litellm_metadata"] = kwargs["litellm_metadata"]
# For endpoints like /v1/messages that use "litellm_metadata" instead
# of "metadata" (to avoid conflicting with provider API metadata fields),
# populate litellm_params["metadata"] so callbacks (e.g. Langfuse) that
@ -5410,6 +5410,19 @@ def _get_model_info_helper(
cache_creation_input_token_cost_above_200k_tokens=_model_info.get(
"cache_creation_input_token_cost_above_200k_tokens", None
),
cache_creation_input_token_cost_above_272k_tokens=_model_info.get(
"cache_creation_input_token_cost_above_272k_tokens", None
),
cache_creation_input_token_cost_above_272k_tokens_priority=_model_info.get(
"cache_creation_input_token_cost_above_272k_tokens_priority", None
),
cache_creation_input_token_cost_above_272k_tokens_flex=_model_info.get(
"cache_creation_input_token_cost_above_272k_tokens_flex", None
),
cache_creation_input_token_cost_flex=_model_info.get("cache_creation_input_token_cost_flex", None),
cache_creation_input_token_cost_priority=_model_info.get(
"cache_creation_input_token_cost_priority", None
),
cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None),
prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None),
cache_read_input_token_cost_above_200k_tokens=_model_info.get(
@ -5424,6 +5437,9 @@ def _get_model_info_helper(
cache_read_input_token_cost_above_272k_tokens_priority=_model_info.get(
"cache_read_input_token_cost_above_272k_tokens_priority", None
),
cache_read_input_token_cost_above_272k_tokens_flex=_model_info.get(
"cache_read_input_token_cost_above_272k_tokens_flex", None
),
cache_read_input_token_cost_above_512k_tokens=_model_info.get(
"cache_read_input_token_cost_above_512k_tokens", None
),
@ -5442,6 +5458,9 @@ def _get_model_info_helper(
input_cost_per_token_above_272k_tokens_priority=_model_info.get(
"input_cost_per_token_above_272k_tokens_priority", None
),
input_cost_per_token_above_272k_tokens_flex=_model_info.get(
"input_cost_per_token_above_272k_tokens_flex", None
),
input_cost_per_token_above_512k_tokens=_model_info.get("input_cost_per_token_above_512k_tokens", None),
input_cost_per_query=_model_info.get("input_cost_per_query", None),
input_cost_per_second=_model_info.get("input_cost_per_second", None),
@ -5483,6 +5502,9 @@ def _get_model_info_helper(
output_cost_per_token_above_272k_tokens_priority=_model_info.get(
"output_cost_per_token_above_272k_tokens_priority", None
),
output_cost_per_token_above_272k_tokens_flex=_model_info.get(
"output_cost_per_token_above_272k_tokens_flex", None
),
output_cost_per_token_above_512k_tokens=_model_info.get(
"output_cost_per_token_above_512k_tokens", None
),

View file

@ -23753,14 +23753,17 @@
"gpt-5.6": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06,
"cache_creation_input_token_cost_flex": 3.125e-06,
"cache_creation_input_token_cost_priority": 1.25e-05,
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 1e-06,
"input_cost_per_token": 5e-06,
"input_cost_per_token_above_272k_tokens": 1e-05,
"input_cost_per_token_above_272k_tokens_flex": 5e-06,
"input_cost_per_token_batches": 2.5e-06,
"input_cost_per_token_flex": 2.5e-06,
"input_cost_per_token_priority": 1e-05,
@ -23771,6 +23774,7 @@
"mode": "chat",
"output_cost_per_token": 3e-05,
"output_cost_per_token_above_272k_tokens": 4.5e-05,
"output_cost_per_token_above_272k_tokens_flex": 2.25e-05,
"output_cost_per_token_batches": 1.5e-05,
"output_cost_per_token_flex": 1.5e-05,
"output_cost_per_token_priority": 6e-05,
@ -23806,14 +23810,17 @@
"gpt-5.6-sol": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06,
"cache_creation_input_token_cost_flex": 3.125e-06,
"cache_creation_input_token_cost_priority": 1.25e-05,
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
"cache_read_input_token_cost_above_272k_tokens_flex": 5e-07,
"cache_read_input_token_cost_flex": 2.5e-07,
"cache_read_input_token_cost_priority": 1e-06,
"input_cost_per_token": 5e-06,
"input_cost_per_token_above_272k_tokens": 1e-05,
"input_cost_per_token_above_272k_tokens_flex": 5e-06,
"input_cost_per_token_batches": 2.5e-06,
"input_cost_per_token_flex": 2.5e-06,
"input_cost_per_token_priority": 1e-05,
@ -23824,6 +23831,7 @@
"mode": "chat",
"output_cost_per_token": 3e-05,
"output_cost_per_token_above_272k_tokens": 4.5e-05,
"output_cost_per_token_above_272k_tokens_flex": 2.25e-05,
"output_cost_per_token_batches": 1.5e-05,
"output_cost_per_token_flex": 1.5e-05,
"output_cost_per_token_priority": 6e-05,
@ -23857,29 +23865,33 @@
"supports_xhigh_reasoning_effort": true
},
"gpt-5.6-terra": {
"cache_creation_input_token_cost": 3.125e-06,
"cache_creation_input_token_cost_above_272k_tokens": 6.25e-06,
"cache_creation_input_token_cost_flex": 1.5625e-06,
"cache_creation_input_token_cost_priority": 6.25e-06,
"cache_read_input_token_cost": 2.5e-07,
"cache_read_input_token_cost_above_272k_tokens": 5e-07,
"cache_read_input_token_cost_flex": 1.25e-07,
"cache_read_input_token_cost_priority": 5e-07,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_above_272k_tokens": 5e-06,
"input_cost_per_token_batches": 1.25e-06,
"input_cost_per_token_flex": 1.25e-06,
"input_cost_per_token_priority": 5e-06,
"cache_creation_input_token_cost": 2.5e-06,
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
"cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06,
"cache_creation_input_token_cost_flex": 1.25e-06,
"cache_creation_input_token_cost_priority": 5e-06,
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_272k_tokens": 4e-07,
"cache_read_input_token_cost_above_272k_tokens_flex": 2e-07,
"cache_read_input_token_cost_flex": 1e-07,
"cache_read_input_token_cost_priority": 4e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"input_cost_per_token_above_272k_tokens_flex": 2e-06,
"input_cost_per_token_batches": 1e-06,
"input_cost_per_token_flex": 1e-06,
"input_cost_per_token_priority": 4e-06,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_above_272k_tokens": 2.25e-05,
"output_cost_per_token_batches": 7.5e-06,
"output_cost_per_token_flex": 7.5e-06,
"output_cost_per_token_priority": 3e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_272k_tokens": 1.8e-05,
"output_cost_per_token_above_272k_tokens_flex": 9e-06,
"output_cost_per_token_batches": 6e-06,
"output_cost_per_token_flex": 6e-06,
"output_cost_per_token_priority": 2.4e-05,
"regional_processing_uplift_multiplier_eu": 1.1,
"regional_processing_uplift_multiplier_us": 1.1,
"supported_endpoints": [
@ -23910,29 +23922,33 @@
"supports_xhigh_reasoning_effort": true
},
"gpt-5.6-luna": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_creation_input_token_cost_flex": 6.25e-07,
"cache_creation_input_token_cost_priority": 2.5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"cache_read_input_token_cost_flex": 5e-08,
"cache_read_input_token_cost_priority": 2e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"input_cost_per_token_batches": 5e-07,
"input_cost_per_token_flex": 5e-07,
"input_cost_per_token_priority": 2e-06,
"cache_creation_input_token_cost": 2.5e-07,
"cache_creation_input_token_cost_above_272k_tokens": 5e-07,
"cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07,
"cache_creation_input_token_cost_flex": 1.25e-07,
"cache_creation_input_token_cost_priority": 5e-07,
"cache_read_input_token_cost": 2e-08,
"cache_read_input_token_cost_above_272k_tokens": 4e-08,
"cache_read_input_token_cost_above_272k_tokens_flex": 2e-08,
"cache_read_input_token_cost_flex": 1e-08,
"cache_read_input_token_cost_priority": 4e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_272k_tokens": 4e-07,
"input_cost_per_token_above_272k_tokens_flex": 2e-07,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_token_flex": 1e-07,
"input_cost_per_token_priority": 4e-07,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"output_cost_per_token_above_272k_tokens": 9e-06,
"output_cost_per_token_batches": 3e-06,
"output_cost_per_token_flex": 3e-06,
"output_cost_per_token_priority": 1.2e-05,
"output_cost_per_token": 1.2e-06,
"output_cost_per_token_above_272k_tokens": 1.8e-06,
"output_cost_per_token_above_272k_tokens_flex": 9e-07,
"output_cost_per_token_batches": 6e-07,
"output_cost_per_token_flex": 6e-07,
"output_cost_per_token_priority": 2.4e-06,
"regional_processing_uplift_multiplier_eu": 1.1,
"regional_processing_uplift_multiplier_us": 1.1,
"supported_endpoints": [
@ -45241,10 +45257,10 @@
"supports_vision": true
},
"bedrock_mantle/openai.gpt-5.6-terra": {
"input_cost_per_token": 2.75e-06,
"cache_creation_input_token_cost": 3.4375e-06,
"cache_read_input_token_cost": 2.75e-07,
"output_cost_per_token": 1.65e-05,
"input_cost_per_token": 2.2e-06,
"cache_creation_input_token_cost": 2.75e-06,
"cache_read_input_token_cost": 2.2e-07,
"output_cost_per_token": 1.32e-05,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
@ -45269,10 +45285,10 @@
"supports_vision": true
},
"bedrock_mantle/openai.gpt-5.6-luna": {
"input_cost_per_token": 1.1e-06,
"cache_creation_input_token_cost": 1.375e-06,
"cache_read_input_token_cost": 1.1e-07,
"output_cost_per_token": 6.6e-06,
"input_cost_per_token": 2.2e-07,
"cache_creation_input_token_cost": 2.75e-07,
"cache_read_input_token_cost": 2.2e-08,
"output_cost_per_token": 1.32e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 272000,
"max_output_tokens": 128000,

View file

@ -99,6 +99,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
"description": "Flex service-tier rate for the same-named base field."
},
"cache_creation_input_token_cost_flex": {
"type": "number",
"minimum": 0,
@ -133,6 +138,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
"description": "Flex service-tier rate for the same-named base field."
},
"cache_read_input_token_cost_above_272k_tokens_priority": {
"type": "number",
"minimum": 0,
@ -262,6 +272,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
"description": "Flex service-tier rate for the same-named base field."
},
"input_cost_per_token_above_272k_tokens_priority": {
"type": "number",
"minimum": 0,
@ -434,6 +449,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
"description": "Flex service-tier rate for the same-named base field."
},
"output_cost_per_token_above_272k_tokens_priority": {
"type": "number",
"minimum": 0,

View file

@ -85,6 +85,8 @@ The metric is coverage: the share of registry rows that have a passing covering
Tests do not declare a dashboard module directly. They only declare the registry cell id with `@pytest.mark.covers("...")`; the registry row decides the module, tier, endpoint, and dashboard rollup. Run `python -m coverage_registry.collector --strict` when you want CI to reject unknown marker ids. Add `--fail-on-collection-errors` when the job should also fail on pytest collection errors.
Skipping a test gives its cell back to the gap list: the collector counts a cell as covered only when a test pytest would actually run declares it, and prints the cells left claimed only by skipped tests. So a `@pytest.mark.skip` on a red cell is honest bookkeeping, not a way to keep the number up.
### Naming grammar per module
LLMs - endpoint features (subject = the route), seeded from the Claude Code compat matrix. `chat_completions`, `messages`, and `responses` roll up to `Core LLMs`. Other LLM endpoints, including `batches` and `realtime`, roll up to `Non-Core LLMs`.

View file

@ -37,6 +37,16 @@ def test_openai_streaming_tool_calls(self) -> None:
It is static: a collect-only pass reads the markers, so it runs no test and needs no live
proxy. Whether a covered cell currently passes or fails is a separate, live concern.
A skipped test asserts nothing, so its markers do not count. A cell is covered only when
at least one test pytest would actually run declares it; a cell claimed by both a live
test and a skipped one stays covered. Skip state comes from pytest's own evaluator, so
`skip` and `skipif` resolve exactly as they do in the e2e run, which also means a
`skipif` on an absent credential makes that cell uncovered in the environments where the
test cannot run. Cells left uncovered this way are listed under the headline (and counted
by `litellm_e2e_coverage_skipped_markers`) so an unskipped-pending gap is visible rather
than inflating the number. The one skip the collector cannot see is `pytest.skip()`
called from inside a test body, since it does not exist until the test runs.
```
cd tests/e2e && PYTHONPATH=. python -m coverage_registry.collector
```

View file

@ -5,6 +5,13 @@ Coverage here is static: it reads the markers via a collect-only pass, so it run
no test and needs no live proxy. Whether a covered cell currently passes or fails
(covered_pass vs covered_fail) is a separate, live concern layered on top later.
A skipped test asserts nothing, so its markers do not count: a cell is covered
only when at least one test that pytest would actually run declares it. Skip
state is read with pytest's own evaluator, so `skip` and `skipif` are resolved
exactly as the e2e run resolves them in this environment. The one skip the
collector cannot see is `pytest.skip()` called from inside a test body, which
does not exist until the test runs.
cd tests/e2e && PYTHONPATH=. python -m coverage_registry.collector
"""
@ -20,6 +27,7 @@ from pathlib import Path
from typing import Literal
import pytest
from _pytest.skipping import evaluate_skip_marks
from pydantic import BaseModel
from .registry import load_registry
@ -28,22 +36,57 @@ from .schema import MODULE_ORDER, Cell, Tier, dashboard_module, loki_module_labe
E2E_DIR = Path(__file__).resolve().parent.parent
@dataclass(frozen=True, slots=True)
class CollectedMarkers:
"""What a collect-only pass saw: cell ids declared by tests that would run,
cell ids only ever declared by skipped tests, and nodes that failed to import."""
covered: frozenset[str]
skipped_only: frozenset[str]
collection_errors: tuple[str, ...]
def _is_skipped(item: pytest.Item) -> bool:
"""True when pytest would skip this test instead of running it.
A marker pytest cannot evaluate (for example a bare boolean `skipif` with no
reason) turns into a setup failure at run time, so the test asserts nothing
either way and is treated the same as a skip.
"""
try:
return evaluate_skip_marks(item) is not None
except (pytest.fail.Exception, TypeError):
return True
class _CoversSink:
"""Pytest plugin: after collection, capture every cell id declared via
@pytest.mark.covers(...), plus any nodes that failed to import."""
@pytest.mark.covers(...) split by whether its test would run, plus any nodes
that failed to import."""
def __init__(self) -> None:
self.covered_ids: frozenset[str] = frozenset()
self.skipped_only_ids: frozenset[str] = frozenset()
self.collection_errors: tuple[str, ...] = ()
def pytest_collection_finish(self, session: pytest.Session) -> None:
marker_args: tuple[tuple[object, ...], ...] = tuple(
marker.args
for item in session.items
marker_args: tuple[tuple[bool, tuple[object, ...]], ...] = tuple(
(skipped, marker.args)
for item, skipped in ((i, _is_skipped(i)) for i in session.items)
for marker in item.iter_markers(name="covers")
)
declared = tuple(
(skipped, arg)
for skipped, args in marker_args
for arg in args
if isinstance(arg, str)
)
self.covered_ids = frozenset(
arg for args in marker_args for arg in args if isinstance(arg, str)
cell_id for skipped, cell_id in declared if not skipped
)
self.skipped_only_ids = (
frozenset(cell_id for skipped, cell_id in declared if skipped)
- self.covered_ids
)
def pytest_collectreport(self, report: pytest.CollectReport) -> None:
@ -51,10 +94,8 @@ class _CoversSink:
self.collection_errors = (*self.collection_errors, report.nodeid)
def collect_covered_ids(
e2e_dir: Path = E2E_DIR,
) -> tuple[frozenset[str], tuple[str, ...]]:
"""Return (covered cell ids, nodeids that failed to import)."""
def collect_markers(e2e_dir: Path = E2E_DIR) -> CollectedMarkers:
"""Read every @pytest.mark.covers marker in `e2e_dir` via a collect-only pass."""
sink = _CoversSink()
with contextlib.redirect_stdout(io.StringIO()):
pytest.main(
@ -68,7 +109,11 @@ def collect_covered_ids(
],
plugins=[sink],
)
return sink.covered_ids, sink.collection_errors
return CollectedMarkers(
covered=sink.covered_ids,
skipped_only=sink.skipped_only_ids,
collection_errors=sink.collection_errors,
)
@dataclass(frozen=True, slots=True)
@ -93,6 +138,7 @@ class CoverageReport:
p0_covered: int
p0_gaps: tuple[str, ...]
orphan_markers: tuple[str, ...]
skipped_markers: tuple[str, ...]
collection_errors: tuple[str, ...]
@property
@ -122,6 +168,7 @@ def compute_coverage(
cells: tuple[Cell, ...],
covered: frozenset[str],
collection_errors: tuple[str, ...] = (),
skipped_only: frozenset[str] = frozenset(),
) -> CoverageReport:
p0_cells = tuple(c for c in cells if c.tier is Tier.P0)
registry_ids = frozenset(c.id for c in cells)
@ -132,7 +179,8 @@ def compute_coverage(
p0_total=len(p0_cells),
p0_covered=sum(1 for c in p0_cells if c.id in covered),
p0_gaps=tuple(sorted(c.id for c in p0_cells if c.id not in covered)),
orphan_markers=tuple(sorted(covered - registry_ids)),
orphan_markers=tuple(sorted((covered | skipped_only) - registry_ids)),
skipped_markers=tuple(sorted(skipped_only & registry_ids)),
collection_errors=collection_errors,
)
@ -161,6 +209,15 @@ def render(report: CoverageReport) -> str:
if report.orphan_markers
else ()
)
skipped = (
(
f"\n{len(report.skipped_markers)} cell(s) are claimed only by skipped tests, "
f"so they count as uncovered (unskip the test or drop the marker):\n "
+ "\n ".join(report.skipped_markers),
)
if report.skipped_markers
else ()
)
warning = (
(
f"\nWARNING: {len(report.collection_errors)} node(s) failed to import during "
@ -170,7 +227,7 @@ def render(report: CoverageReport) -> str:
if report.collection_errors
else ()
)
return "\n".join((*lines, *orphans, *warning))
return "\n".join((*lines, *orphans, *skipped, *warning))
def _report_dict(report: CoverageReport) -> dict[str, object]:
@ -190,6 +247,7 @@ def _report_dict(report: CoverageReport) -> dict[str, object]:
for m in report.modules
],
"orphan_markers": list(report.orphan_markers),
"skipped_markers": list(report.skipped_markers),
"collection_errors": list(report.collection_errors),
}
@ -234,6 +292,9 @@ def render_prometheus(report: CoverageReport) -> str:
"# HELP litellm_e2e_coverage_orphan_markers Coverage markers not found in the registry.",
"# TYPE litellm_e2e_coverage_orphan_markers gauge",
f"litellm_e2e_coverage_orphan_markers {len(report.orphan_markers)}",
"# HELP litellm_e2e_coverage_skipped_markers Registry cells claimed only by skipped tests.",
"# TYPE litellm_e2e_coverage_skipped_markers gauge",
f"litellm_e2e_coverage_skipped_markers {len(report.skipped_markers)}",
"# HELP litellm_e2e_coverage_collection_errors Pytest nodes that failed during collection.",
"# TYPE litellm_e2e_coverage_collection_errors gauge",
f"litellm_e2e_coverage_collection_errors {len(report.collection_errors)}",
@ -286,8 +347,13 @@ def main() -> int:
)
args = _CliArgs.model_validate(vars(parser.parse_args()))
cells = load_registry()
covered, errors = collect_covered_ids()
report = compute_coverage(cells, covered, errors)
markers = collect_markers()
report = compute_coverage(
cells,
markers.covered,
markers.collection_errors,
markers.skipped_only,
)
output = {
"text": render,
"json": render_json,

View file

@ -12,6 +12,7 @@ from pathlib import Path
import pytest
from coverage_registry.collector import (
collect_markers,
compute_coverage,
render,
render_json,
@ -61,6 +62,27 @@ def test_orphan_marker_is_reported_not_counted() -> None:
assert report.orphan_markers == ("llm.ghost",)
def test_cell_claimed_only_by_a_skipped_test_is_uncovered() -> None:
cells = (_llm("llm.a", Tier.P0), _llm("llm.b", Tier.P0))
report = compute_coverage(
cells, frozenset({"llm.a"}), skipped_only=frozenset({"llm.b"})
)
assert (report.covered, report.p0_covered) == (1, 1)
assert report.p0_gaps == ("llm.b",)
assert report.skipped_markers == ("llm.b",)
assert "only by skipped tests" in render(report)
assert '"skipped_markers": [\n "llm.b"\n ]' in render_json(report)
assert "litellm_e2e_coverage_skipped_markers 1" in render_prometheus(report)
def test_skipped_marker_outside_the_registry_is_still_an_orphan() -> None:
report = compute_coverage(
(_llm("llm.a", Tier.P0),), frozenset(), skipped_only=frozenset({"llm.ghost"})
)
assert report.orphan_markers == ("llm.ghost",)
assert report.skipped_markers == ()
def test_logging_and_guardrail_roll_up_into_one_module() -> None:
cells = (
LoggingCell(
@ -175,6 +197,81 @@ def test_loki_render_exposes_exact_stdout_lines_for_loki() -> None:
)
_MARKED_TESTS = '''
import pytest
@pytest.mark.covers("llm.runs")
def test_runs() -> None:
pass
@pytest.mark.skip(reason="stage red: product gap")
@pytest.mark.covers("llm.skipped")
def test_skipped() -> None:
pass
@pytest.mark.skipif(True, reason="credentials absent in this environment")
@pytest.mark.covers("llm.skipif_true")
def test_skipif_true() -> None:
pass
@pytest.mark.skipif(False, reason="credentials present in this environment")
@pytest.mark.covers("llm.skipif_false")
def test_skipif_false() -> None:
pass
@pytest.mark.skipif("True")
@pytest.mark.covers("llm.skipif_string")
def test_skipif_string_condition() -> None:
pass
@pytest.mark.covers("llm.shared")
def test_shared_cell_runs() -> None:
pass
@pytest.mark.skip(reason="stage red: product gap")
@pytest.mark.covers("llm.shared")
def test_shared_cell_skipped() -> None:
pass
'''
_MODULE_LEVEL_SKIP = '''
import pytest
pytestmark = pytest.mark.skipif(True, reason="whole module needs a session fixture")
@pytest.mark.covers("llm.module_skipped")
def test_module_level_skip() -> None:
pass
'''
def test_collection_counts_only_markers_on_tests_that_would_run(
tmp_path: Path,
) -> None:
"""The collect-only pass is the numerator, so a test pytest would skip must not
contribute its cell. A cell stays covered as long as one runnable test claims it."""
(tmp_path / "test_marked.py").write_text(_MARKED_TESTS)
(tmp_path / "test_module_skip.py").write_text(_MODULE_LEVEL_SKIP)
markers = collect_markers(tmp_path)
assert markers.covered == frozenset(
{"llm.runs", "llm.skipif_false", "llm.shared"}
)
assert markers.skipped_only == frozenset(
{"llm.skipped", "llm.skipif_true", "llm.skipif_string", "llm.module_skipped"}
)
assert markers.collection_errors == ()
def test_real_registry_loads_and_ids_are_unique() -> None:
cells = load_registry()
ids = [c.id for c in cells]

View file

@ -2010,6 +2010,9 @@ async def test_get_tools_for_single_server_applies_disallowed_tools_without_allo
mock_server.mcp_info = {"server_name": "zapier"}
mock_server.name = "zapier"
mock_server.server_id = "zapier"
mock_server.server_name = "zapier"
mock_server.alias = None
mock_server.short_prefix = None
mock_server.allowed_tools = None
mock_server.disallowed_tools = ["send_email"]
@ -2036,6 +2039,65 @@ async def test_get_tools_for_single_server_applies_disallowed_tools_without_allo
assert [tool.name for tool in result] == ["read_email"]
@pytest.mark.asyncio
async def test_rest_listing_hides_key_grants_dispatch_would_refuse():
"""REST listing must answer for exactly the key/team grants dispatch honors.
``mcp_tool_permissions`` and toolset rows name a tool on one server, so both
the MCP list path and ``tools/call`` compare them bare. A wire-form entry
therefore grants nothing, and REST listing that matched the prefixed
spelling would advertise a tool the very next call refuses.
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.rest_endpoints import (
_get_tools_for_single_server,
)
from litellm.proxy._types import UserAPIKeyAuth
from mcp.types import Tool as MCPTool
server_id = "3c6f6617-d23c-4f48-bfb0-f205e3b27bab"
mock_server = MagicMock()
mock_server.mcp_info = {"server_name": server_id}
mock_server.name = server_id
mock_server.server_id = server_id
mock_server.server_name = None
mock_server.alias = None
mock_server.short_prefix = None
mock_server.allowed_tools = None
mock_server.disallowed_tools = None
mock_server.tool_name_to_display_name = None
mock_tools = [
MCPTool(
name="read_wiki_contents",
description="Read a wiki",
inputSchema={"type": "object"},
),
]
with patch(
"litellm.proxy._experimental.mcp_server.rest_endpoints.global_mcp_server_manager"
) as mock_manager, patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
) as mock_server_manager, patch.object(
MCPRequestHandler,
"get_allowed_tools_for_server",
AsyncMock(return_value=[f"{server_id}-read_wiki_contents"]),
):
mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools)
mock_server_manager.get_mcp_server_by_id.return_value = mock_server
result = await _get_tools_for_single_server(
mock_server,
"Bearer test_token",
user_api_key_auth=UserAPIKeyAuth(api_key="sk-test"),
)
assert result == []
@pytest.mark.asyncio
async def test_list_tool_rest_api_with_server_specific_auth():
"""Test list_tool_rest_api with server-specific auth headers."""

View file

@ -0,0 +1,156 @@
from datetime import datetime
from unittest.mock import MagicMock, patch
import litellm
from litellm.integrations.s3 import S3Logger
TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id"
def _standard_logging_payload() -> dict:
return {
"id": "chatcmpl-test-id",
"metadata": {"user_api_key_team_alias": None},
}
def _log_event_kwargs() -> dict:
return {
"litellm_params": {"metadata": {}},
"standard_logging_object": _standard_logging_payload(),
}
def _run_log_event(callback_params: dict) -> MagicMock:
original = litellm.s3_callback_params
litellm.s3_callback_params = callback_params
try:
with patch("boto3.client") as mock_boto3_client:
mock_s3_client = MagicMock()
mock_boto3_client.return_value = mock_s3_client
logger = S3Logger()
logger.log_event(
kwargs=_log_event_kwargs(),
response_obj={},
start_time=datetime(2026, 7, 30, 12, 0, 0),
end_time=datetime(2026, 7, 30, 12, 0, 1),
print_verbose=lambda *args, **kwargs: None,
)
return mock_s3_client
finally:
litellm.s3_callback_params = original
def test_put_object_includes_sse_kms_params_when_configured():
"""
When s3_server_side_encryption and s3_sse_kms_key_id are set in
s3_callback_params, put_object must receive ServerSideEncryption and
SSEKMSKeyId so objects land encrypted with the customer-managed key.
"""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN
def test_put_object_supports_sse_s3_without_key_id():
"""SSE-S3 (AES256) needs only ServerSideEncryption, no key id."""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
"s3_server_side_encryption": "AES256",
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert put_object_kwargs["ServerSideEncryption"] == "AES256"
assert "SSEKMSKeyId" not in put_object_kwargs
def test_put_object_omits_sse_params_by_default():
"""Without SSE config, put_object kwargs must stay unchanged."""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert "ServerSideEncryption" not in put_object_kwargs
assert "SSEKMSKeyId" not in put_object_kwargs
def test_put_object_infers_aws_kms_when_only_key_id_set():
"""A key id without an algorithm must infer aws:kms instead of sending an invalid request."""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN
def test_put_object_drops_key_id_when_algorithm_is_not_kms():
"""AES256 plus a key id is invalid for S3; the key id must be dropped, not sent."""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
"s3_server_side_encryption": "AES256",
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert put_object_kwargs["ServerSideEncryption"] == "AES256"
assert "SSEKMSKeyId" not in put_object_kwargs
def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued():
"""
A YAML boolean in s3_server_side_encryption must not crash logger init and
must not discard the valid key id; aws:kms is inferred from the key id.
"""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
"s3_server_side_encryption": True,
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN
def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept():
"""A mistyped key id (unquoted YAML number) must not disable the valid algorithm."""
mock_s3_client = _run_log_event(
{
"s3_bucket_name": "test-bucket",
"s3_region_name": "us-east-1",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": 12345,
}
)
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
assert "SSEKMSKeyId" not in put_object_kwargs

View file

@ -1388,3 +1388,253 @@ def test_s3_server_side_encryption_read_from_callback_params():
assert logger.s3_server_side_encryption == "aws:kms"
finally:
litellm.s3_callback_params = original
@pytest.mark.asyncio
async def test_async_upload_sets_sse_kms_key_id_header_when_configured():
"""
When s3_sse_kms_key_id is set alongside aws:kms, the PUT must carry
x-amz-server-side-encryption-aws-kms-key-id so objects are encrypted
with the customer-managed KMS key instead of the bucket default.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_server_side_encryption="aws:kms",
s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id",
)
test_element = s3BatchLoggingElement(
s3_object_key="2025-09-14/test-sse-kms.json",
payload={"test": "sse-kms"},
s3_object_download_filename="test-sse-kms.json",
)
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put.return_value = response
await logger.async_upload_data_to_s3(test_element)
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
assert headers["x-amz-server-side-encryption"] == "aws:kms"
assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == (
"arn:aws:kms:us-east-1:111122223333:key/test-key-id"
)
def test_sync_upload_sets_sse_kms_key_id_header_when_configured():
"""The sync upload path must carry the same SSE-KMS headers."""
from unittest.mock import MagicMock
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_server_side_encryption="aws:kms",
s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id",
)
test_element = s3BatchLoggingElement(
s3_object_key="2025-09-14/test-sync-sse-kms.json",
payload={"test": "sync-sse-kms"},
s3_object_download_filename="test-sync-sse-kms.json",
)
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
mock_sync_client = MagicMock()
mock_sync_client.put.return_value = response
with patch(
"litellm.integrations.s3_v2._get_httpx_client",
return_value=mock_sync_client,
):
logger.upload_data_to_s3(test_element)
headers = mock_sync_client.put.call_args.kwargs["headers"]
assert headers["x-amz-server-side-encryption"] == "aws:kms"
assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == (
"arn:aws:kms:us-east-1:111122223333:key/test-key-id"
)
@pytest.mark.asyncio
async def test_async_upload_omits_kms_key_id_header_when_not_configured():
"""SSE without a key id must not emit the KMS key id header."""
from unittest.mock import AsyncMock, MagicMock
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_server_side_encryption="AES256",
)
test_element = s3BatchLoggingElement(
s3_object_key="2025-09-14/test-aes256.json",
payload={"test": "aes256"},
s3_object_download_filename="test-aes256.json",
)
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put.return_value = response
await logger.async_upload_data_to_s3(test_element)
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
assert headers["x-amz-server-side-encryption"] == "AES256"
assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers
def test_s3_sse_kms_key_id_read_from_callback_params():
"""s3_sse_kms_key_id can be configured via s3_callback_params."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
}
try:
logger = S3Logger()
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
finally:
litellm.s3_callback_params = original
@pytest.mark.asyncio
async def test_async_upload_infers_aws_kms_when_only_key_id_set():
"""
Setting only s3_sse_kms_key_id must not produce an invalid request
(S3 rejects a key id without an algorithm); aws:kms is inferred.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id",
)
test_element = s3BatchLoggingElement(
s3_object_key="2025-09-14/test-kms-only.json",
payload={"test": "kms-only"},
s3_object_download_filename="test-kms-only.json",
)
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put.return_value = response
await logger.async_upload_data_to_s3(test_element)
headers = logger.async_httpx_client.put.call_args.kwargs["headers"]
assert headers["x-amz-server-side-encryption"] == "aws:kms"
assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == (
"arn:aws:kms:us-east-1:111122223333:key/test-key-id"
)
def test_s3_sse_kms_key_id_read_from_audit_override_params():
"""The audit-log override path must honor s3_sse_kms_key_id too."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {"s3_bucket_name": "normal-logs-bucket"}
try:
logger = S3Logger(
s3_callback_params_override={
"s3_bucket_name": "audit-logs-bucket",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id",
}
)
assert logger.s3_bucket_name == "audit-logs-bucket"
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id")
finally:
litellm.s3_callback_params = original
def test_kms_key_id_dropped_when_algorithm_is_not_kms():
"""
AES256 plus a KMS key id is an invalid S3 combination; the key id must be
dropped at init so uploads keep working instead of silently 400ing.
"""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "AES256",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "AES256"
assert logger.s3_sse_kms_key_id is None
finally:
litellm.s3_callback_params = original
def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued():
"""
A YAML boolean in s3_server_side_encryption must not crash logger init and
must not discard the valid key id; aws:kms is inferred from the key id.
"""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": True,
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
finally:
litellm.s3_callback_params = original
def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept():
"""A mistyped key id (unquoted YAML number) must not disable the valid algorithm."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": 12345,
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
assert logger.s3_sse_kms_key_id is None
finally:
litellm.s3_callback_params = original

View file

@ -611,8 +611,8 @@ def test_generic_cost_per_token_gpt55_pro():
[
("gpt-5.6", 5e-6, 3e-5, 5e-7, 6.25e-6),
("gpt-5.6-sol", 5e-6, 3e-5, 5e-7, 6.25e-6),
("gpt-5.6-terra", 2.5e-6, 1.5e-5, 2.5e-7, 3.125e-6),
("gpt-5.6-luna", 1e-6, 6e-6, 1e-7, 1.25e-6),
("gpt-5.6-terra", 2e-6, 1.2e-5, 2e-7, 2.5e-6),
("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7),
],
)
def test_generic_cost_per_token_gpt56(
@ -661,6 +661,97 @@ def test_generic_cost_per_token_gpt56(
assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10)
@pytest.mark.parametrize(
"model,flex_long_input_cost,flex_long_output_cost",
[
("gpt-5.6", 5e-6, 2.25e-5),
("gpt-5.6-sol", 5e-6, 2.25e-5),
("gpt-5.6-terra", 2e-6, 9e-6),
("gpt-5.6-luna", 2e-7, 9e-7),
],
)
def test_generic_cost_per_token_gpt56_flex_above_272k(
model, flex_long_input_cost, flex_long_output_cost
):
"""A >272K flex request bills the flex long-context rate, not the standard one.
Flex long-context is half the standard long-context rate. Without the
``*_above_272k_tokens_flex`` keys these requests silently fell back to the
standard long-context price, billing 2x what OpenAI charges.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
prompt_tokens = 300000
completion_tokens = 1000
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider="openai",
service_tier="flex",
)
assert prompt_cost == pytest.approx(flex_long_input_cost * prompt_tokens)
assert completion_cost == pytest.approx(flex_long_output_cost * completion_tokens)
standard_long_prompt_cost, standard_long_completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider="openai",
service_tier=None,
)
assert prompt_cost == pytest.approx(standard_long_prompt_cost / 2)
assert completion_cost == pytest.approx(standard_long_completion_cost / 2)
@pytest.mark.parametrize(
"service_tier,prompt_tokens,input_rate,cache_write_rate,cache_read_rate",
[
(None, 100000, 2e-6, 2.5e-6, 2e-7),
("flex", 100000, 1e-6, 1.25e-6, 1e-7),
("priority", 100000, 4e-6, 5e-6, 4e-7),
(None, 300000, 4e-6, 5e-6, 4e-7),
("flex", 300000, 2e-6, 2.5e-6, 2e-7),
],
)
def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context(
service_tier, prompt_tokens, input_rate, cache_write_rate, cache_read_rate
):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
cached_tokens = 50000
cache_write_tokens = 40000
text_tokens = prompt_tokens - cached_tokens - cache_write_tokens
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=100,
total_tokens=prompt_tokens + 100,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens
),
)
prompt_cost, _ = generic_cost_per_token(
model="gpt-5.6-terra",
usage=usage,
custom_llm_provider="openai",
service_tier=service_tier,
)
expected_prompt_cost = (
text_tokens * input_rate
+ cached_tokens * cache_read_rate
+ cache_write_tokens * cache_write_rate
)
assert prompt_cost == pytest.approx(expected_prompt_cost)
@pytest.mark.parametrize(
"model,input_cost,output_cost,cache_read_cost",
[
@ -2399,3 +2490,66 @@ def test_generic_cost_per_token_gemini_35_flash_lite():
)
assert prompt_cost == pytest.approx(0.0003)
assert completion_cost == pytest.approx(0.00125)
def test_fast_service_tier_bills_at_the_priority_rate(_local_model_cost_map):
"""Regression: OpenAI's Fast mode replaced Priority Processing and costs 2x standard.
Before the fix "fast" fell through to standard pricing, so a Fast mode request
was billed at half of what it actually costs."""
from litellm.types.utils import Usage
usage = Usage(
prompt_tokens=1_000,
completion_tokens=500,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200),
)
standard = generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier=None
)
priority = generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="priority"
)
fast = generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast"
)
expected_prompt = 800 * 1e-05 + 200 * 1e-06
expected_completion = 500 * 6e-05
assert fast == priority
assert fast[0] == pytest.approx(expected_prompt, rel=1e-9)
assert fast[1] == pytest.approx(expected_completion, rel=1e-9)
assert fast[0] == pytest.approx(standard[0] * 2, rel=1e-9)
assert fast[1] == pytest.approx(standard[1] * 2, rel=1e-9)
def test_fast_service_tier_is_case_insensitive(_local_model_cost_map):
from litellm.types.utils import Usage
usage = Usage(prompt_tokens=1_000, completion_tokens=500)
assert generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="FAST"
) == generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast"
)
def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_model_cost_map):
"""The above-threshold branch resolves its own cost keys, so the alias has to hold there too."""
from litellm.types.utils import Usage
usage = Usage(prompt_tokens=300_000, completion_tokens=1_000)
fast = generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast"
)
priority = generic_cost_per_token(
model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="priority"
)
assert fast == priority
assert fast[0] == pytest.approx(300_000 * 1e-05, rel=1e-9)
assert fast[1] == pytest.approx(1_000 * 4.5e-05, rel=1e-9)

View file

@ -2992,6 +2992,58 @@ def test_function_setup_litellm_metadata_populates_metadata():
), "litellm_params['metadata'] should be a copy, not the same object"
def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
"""
Regression test for LIT-4512: guardrail writes into the request's
"litellm_metadata" bucket that happen AFTER function_setup (the proxy
initializes the logging object before pre-call guardrails run) must be
visible to the logging object and survive merge_litellm_metadata, so
/v1/messages spend logs carry guardrail_information and
applied_guardrails just like /v1/chat/completions.
"""
import litellm
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
kwargs = {
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "hello"}],
"litellm_call_id": "test-call-id-lit4512",
"litellm_metadata": {
"user_api_key_hash": "sk-hashed-lit4512",
"guardrails": ["pam-ethical-request"],
},
}
logging_obj, returned_kwargs = litellm.utils.function_setup(
original_function="anthropic_messages",
rules_obj=litellm.utils.Rules(),
start_time=time.time(),
**kwargs,
)
guardrail_entry = {
"guardrail_name": "pam-ethical-request",
"guardrail_mode": "pre_call",
"guardrail_status": "success",
}
_, metadata_bucket = get_or_create_metadata_bucket(returned_kwargs)
metadata_bucket["standard_logging_guardrail_information"] = [guardrail_entry]
metadata_bucket["applied_guardrails"] = ["pam-ethical-request"]
litellm_params = logging_obj.model_call_details.get("litellm_params", {})
litellm_metadata = litellm_params.get("litellm_metadata")
assert litellm_metadata is not None
assert litellm_metadata.get("standard_logging_guardrail_information") == [
guardrail_entry
], "guardrail writes after function_setup must be visible to the logging object"
assert litellm_metadata.get("applied_guardrails") == ["pam-ethical-request"]
merged = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params)
assert merged.get("standard_logging_guardrail_information") == [guardrail_entry]
assert merged.get("applied_guardrails") == ["pam-ethical-request"]
def test_function_setup_metadata_takes_precedence_over_litellm_metadata():
"""
Test that when BOTH metadata and litellm_metadata are present (e.g., user sets

View file

@ -1519,8 +1519,8 @@ class TestBedrockMantleResponsesPricing:
"model, input_cost, cache_creation_cost, cache_read_cost, output_cost",
[
("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05),
("openai.gpt-5.6-terra", 2.75e-06, 3.4375e-06, 2.75e-07, 1.65e-05),
("openai.gpt-5.6-luna", 1.1e-06, 1.375e-06, 1.1e-07, 6.6e-06),
("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05),
("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06),
],
)
def test_gpt_5_6_pricing_and_mode(
@ -1534,6 +1534,39 @@ class TestBedrockMantleResponsesPricing:
assert info["output_cost_per_token"] == pytest.approx(output_cost)
assert info["max_input_tokens"] == 272000
@pytest.mark.parametrize(
"model, input_cost, output_cost",
[
("openai.gpt-5.6-sol", 5.5e-06, 3.3e-05),
("openai.gpt-5.6-terra", 2.2e-06, 1.32e-05),
("openai.gpt-5.6-luna", 2.2e-07, 1.32e-06),
],
)
def test_gpt_5_6_responses_call_cost(self, local_cost_map, model, input_cost, output_cost):
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
input_tokens = 100000
output_tokens = 10000
response = ResponsesAPIResponse(
id="resp-1",
created_at=1700000000,
model=model,
output=[],
usage=ResponseAPIUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
),
)
cost = litellm.completion_cost(
completion_response=response,
model=f"bedrock_mantle/{model}",
custom_llm_provider="bedrock_mantle",
)
assert cost == pytest.approx(input_tokens * input_cost + output_tokens * output_cost)
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models

View file

@ -1029,6 +1029,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
working_server.server_name = "working_server"
working_server.auth_type = None
working_server.extra_headers = None
working_server.short_prefix = None
working_server.tool_name_to_display_name = None
working_server.tool_name_to_description = None
failing_server = MagicMock()
failing_server.name = "failing_server"
@ -1039,6 +1042,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
failing_server.server_name = "failing_server"
failing_server.auth_type = None
failing_server.extra_headers = None
failing_server.short_prefix = None
failing_server.tool_name_to_display_name = None
failing_server.tool_name_to_description = None
# Mock global_mcp_server_manager
mock_manager = MagicMock()
@ -4507,28 +4513,37 @@ def test_tool_name_matches_case_insensitive():
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="srv-per-store",
name="per_store",
server_name="per_store",
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
spec_path="/specs/petstore.yaml",
)
# Test case 1: Unprefixed tool name with camelCase in filter list
assert _tool_name_matches("addpet", ["addPet", "updatePet"]) is True
assert _tool_name_matches("updatepet", ["addPet", "updatePet"]) is True
assert _tool_name_matches("deletepet", ["addPet", "updatePet"]) is False
assert _tool_name_matches("addpet", ["addPet", "updatePet"], server) is True
assert _tool_name_matches("updatepet", ["addPet", "updatePet"], server) is True
assert _tool_name_matches("deletepet", ["addPet", "updatePet"], server) is False
# Test case 2: Prefixed tool name with camelCase in filter list
assert _tool_name_matches("per_store-addpet", ["addPet", "updatePet"]) is True
assert _tool_name_matches("per_store-updatepet", ["addPet", "updatePet"]) is True
assert _tool_name_matches("per_store-deletepet", ["addPet", "updatePet"]) is False
assert _tool_name_matches("per_store-addpet", ["addPet", "updatePet"], server) is True
assert _tool_name_matches("per_store-updatepet", ["addPet", "updatePet"], server) is True
assert _tool_name_matches("per_store-deletepet", ["addPet", "updatePet"], server) is False
# Test case 3: Mixed case variations
assert _tool_name_matches("findPetsByStatus", ["findpetsbystatus"]) is True
assert _tool_name_matches("findpetsbystatus", ["findPetsByStatus"]) is True
assert _tool_name_matches("FINDPETSBYSTATUS", ["findPetsByStatus"]) is True
assert _tool_name_matches("findPetsByStatus", ["findpetsbystatus"], server) is True
assert _tool_name_matches("findpetsbystatus", ["findPetsByStatus"], server) is True
assert _tool_name_matches("FINDPETSBYSTATUS", ["findPetsByStatus"], server) is True
# Test case 4: Full prefixed name in filter list (case-insensitive)
assert _tool_name_matches("server-addPet", ["server-addpet"]) is True
assert _tool_name_matches("server-addpet", ["server-addPet"]) is True
assert _tool_name_matches("server-addPet", ["server-addpet"], server) is True
assert _tool_name_matches("server-addpet", ["server-addPet"], server) is True
# Test case 5: Ensure non-matching names still don't match
assert _tool_name_matches("addpet", ["deletePet", "updatePet"]) is False
assert _tool_name_matches("server-addpet", ["deletePet", "updatePet"]) is False
assert _tool_name_matches("addpet", ["deletePet", "updatePet"], server) is False
assert _tool_name_matches("server-addpet", ["deletePet", "updatePet"], server) is False
def test_filter_tools_by_allowed_tools_case_insensitive():
@ -4581,7 +4596,9 @@ def test_filter_tools_by_allowed_tools_case_insensitive():
server = MCPServer(
server_id="test-server",
name="per_store",
server_name="per_store",
transport=MCPTransport.http,
spec_path="/specs/petstore.yaml",
allowed_tools=["addPet", "updatePet", "findPetsByStatus"],
)
@ -6121,6 +6138,70 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool
assert captured["name"] == "echo"
@pytest.mark.asyncio
async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator():
"""A server with no alias publishes its UUID server_id as the tool prefix.
Splitting that wire name at the first separator leaves a truncated UUID tail
glued to the tool name, which then travels to the upstream server as the tool
to call, into the spend log, and into the server-level allowed_tools check.
"""
from mcp.types import TextContent
from litellm.proxy._experimental.mcp_server import server as mcp_module
server_id = "117c814c-1a2b-4c4d-8e8f-0a1b2c3d4e5f"
alias_less_server = MCPServer(
server_id=server_id,
name=server_id,
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.api_key,
authentication_token="abc123",
)
captured: dict = {}
async def fake_handle_managed_mcp_tool(**kwargs):
captured.update(kwargs)
return mcp_module.CallToolResult(
content=[TextContent(type="text", text="ok")],
isError=False,
)
with (
patch.object(
mcp_module.global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=alias_less_server,
),
patch.object(
mcp_module,
"_handle_managed_mcp_tool",
new=fake_handle_managed_mcp_tool,
),
patch.object(
mcp_module.MCPRequestHandler,
"is_tool_allowed",
return_value=True,
),
patch.object(
mcp_module.global_mcp_tool_registry,
"get_tool",
return_value=None,
),
):
await mcp_module.execute_mcp_tool(
name=f"{server_id}-read_wiki_contents",
arguments={"repoName": "acme/wiki"},
allowed_mcp_servers=[alias_less_server],
start_time=datetime.now(),
)
assert captured["server_name"] == server_id
assert captured["name"] == "read_wiki_contents"
@pytest.mark.asyncio
async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credentials():
"""REST server_id must inject the requested server's auth, not a URL-collision peer's."""
@ -6426,6 +6507,8 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details():
fake_server.mcp_info = None
fake_server.server_id = "srv-1"
fake_server.server_name = "openapi-petstore"
fake_server.alias = None
fake_server.short_prefix = None
fake_tool = MagicMock()
fake_tool.name = "list_pets"
@ -6481,7 +6564,13 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details():
@pytest.mark.asyncio
async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requested_server():
"""A prefixed REST name that resolves to no tool must still dispatch to the server_id."""
"""A prefixed REST name that resolves to no tool must still dispatch to the server_id.
The prefix here belongs to a different server, so it is not a prefix boundary on the
routed server and the name travels upstream whole. Stripping it would invoke the routed
server's similarly named tool instead, which is what the tool_server_mismatch 403 exists
to prevent when the prefix does resolve.
"""
from mcp.types import TextContent
from litellm.proxy._experimental.mcp_server import server as mcp_module
@ -6553,7 +6642,7 @@ async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requeste
)
assert captured["server_name"] == "rest_target"
assert captured["name"] == "list_things"
assert captured["name"] == "known_prefix-list_things"
routed_server = {
requested_server.name: requested_server,
@ -7587,22 +7676,28 @@ async def test_aggregate_listing_reports_per_server_outcomes():
working_server = MagicMock()
working_server.name = "working_server"
working_server.alias = "working"
working_server.short_prefix = None
working_server.allowed_tools = None
working_server.disallowed_tools = None
working_server.server_id = "working_server"
working_server.server_name = "working_server"
working_server.auth_type = None
working_server.extra_headers = None
working_server.tool_name_to_display_name = None
working_server.tool_name_to_description = None
broken_server = MagicMock()
broken_server.name = "broken_server"
broken_server.alias = "broken"
broken_server.short_prefix = None
broken_server.allowed_tools = None
broken_server.disallowed_tools = None
broken_server.server_id = "broken_server"
broken_server.server_name = "broken_server"
broken_server.auth_type = None
broken_server.extra_headers = None
broken_server.tool_name_to_display_name = None
broken_server.tool_name_to_description = None
mock_manager = MagicMock()
mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["working_server", "broken_server"])
@ -7924,3 +8019,209 @@ async def test_post_mcp_call_guardrails_propagate_a_block():
user_api_key_auth=None,
request_data={},
)
class TestListFiltersHonorThePrefixBoundary:
"""The listing filters compare a published (prefixed) tool name against
configured entries, so they have to locate the boundary with the server's
registered prefixes. An alias-less server publishes its UUID server_id as
the prefix, and that prefix contains the separator, so cutting at the first
separator drops every tool on the server from the listing.
"""
SERVER_ID = "117c814c-1a2b-4c4d-8e8f-0a1b2c3d4e5f"
@staticmethod
def _alias_less_server(**overrides):
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer
return MCPServer(
server_id=TestListFiltersHonorThePrefixBoundary.SERVER_ID,
name=TestListFiltersHonorThePrefixBoundary.SERVER_ID,
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
**overrides,
)
def _published_tools(self, *bare_names: str):
from mcp.types import Tool as MCPTool
return [
MCPTool(name=f"{self.SERVER_ID}-{bare}", description=bare, inputSchema={"type": "object"})
for bare in bare_names
]
def test_bare_allowlist_entry_keeps_the_published_tool(self):
from litellm.proxy._experimental.mcp_server.server import (
filter_tools_by_allowed_tools,
)
server = self._alias_less_server(allowed_tools=["read_wiki_contents"])
tools = self._published_tools("read_wiki_contents", "read_wiki_structure")
kept = filter_tools_by_allowed_tools(tools, server)
assert [tool.name for tool in kept] == [f"{self.SERVER_ID}-read_wiki_contents"]
def test_bare_blocklist_entry_excludes_the_published_tool(self):
from litellm.proxy._experimental.mcp_server.server import (
filter_tools_by_allowed_tools,
)
server = self._alias_less_server(disallowed_tools=["read_wiki_structure"])
tools = self._published_tools("read_wiki_contents", "read_wiki_structure")
kept = filter_tools_by_allowed_tools(tools, server)
assert [tool.name for tool in kept] == [f"{self.SERVER_ID}-read_wiki_contents"]
def test_unrelated_entry_does_not_match(self):
from litellm.proxy._experimental.mcp_server.server import _tool_name_matches
server = self._alias_less_server()
assert not _tool_name_matches(f"{self.SERVER_ID}-read_wiki_contents", ["read_wiki_structure"], server)
def test_case_folding_applies_to_openapi_servers_and_not_to_native_ones(self):
# Registration rewrites operationIds through sanitize_openapi_tool_name, so
# folding recovers a spec-spelled entry on an OpenAPI server. A native server
# gets none of it: routing dispatches two names differing only in case as two
# tools, so one policy must not decide both.
from litellm.proxy._experimental.mcp_server.server import _tool_name_matches
native = self._alias_less_server()
openapi = self._alias_less_server(spec_path="/specs/petstore.yaml")
assert _tool_name_matches(f"{self.SERVER_ID}-findPetsByStatus", ["findpetsbystatus"], openapi)
assert not _tool_name_matches(f"{self.SERVER_ID}-findPetsByStatus", ["findpetsbystatus"], native)
assert _tool_name_matches(f"{self.SERVER_ID}-findpetsbystatus", ["findpetsbystatus"], native)
def test_alias_form_entry_matches_a_tool_published_under_the_short_prefix(self, monkeypatch):
# Routing accepts the alias form, so an entry stored before short
# prefixes were turned on still governs the tool. Matching only the
# published spelling left it advertised while dispatch refused it.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer
from litellm.proxy._experimental.mcp_server.server import _tool_name_matches
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = MCPServer(
server_id=self.SERVER_ID,
name="deepwiki_cfg",
alias="deepwiki_cfg",
short_prefix="eiG",
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
)
assert _tool_name_matches("eiG-read_wiki_contents", ["deepwiki_cfg-read_wiki_contents"], server)
def test_discovery_hides_exactly_what_dispatch_refuses(self, monkeypatch):
"""Both halves of one decision, driven through both production paths.
A spelling the blocklist enforces but the filter misses leaves a blocked
tool advertised; the reverse hides a tool that would have been callable.
Every spelling routing registers bans, and its upper-cased form bans nothing,
because a tool's identity is its exact name; asserting the verdict and not only
the agreement is what keeps this from passing on a matcher that answers wrongly
but consistently.
"""
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServer,
MCPServerManager,
)
from litellm.proxy._experimental.mcp_server.server import (
filter_tools_by_allowed_tools,
)
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
def _server(**overrides):
return MCPServer(
server_id=self.SERVER_ID,
name="deepwiki_prod",
alias="deepwiki",
server_name="deepwiki_prod",
short_prefix="eiG",
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
**overrides,
)
manager = MCPServerManager()
manager._create_prefixed_tools(
[MCPTool(name="read_wiki_contents", description="", inputSchema={"type": "object"})],
_server(),
)
registered = sorted(manager.tool_name_to_mcp_server_name_mapping)
assert len(registered) > 1
published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"})
for spelling in registered:
for entry, expected in ((spelling, True), (spelling.upper(), False)):
server = _server(disallowed_tools=[entry])
refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server)
hidden = filter_tools_by_allowed_tools([published], server) == []
assert refused == hidden, entry
assert refused is expected, entry
@pytest.mark.asyncio
@pytest.mark.parametrize(
"grants,expected",
[
(None, True),
([], False),
(["read_wiki_contents"], True),
(["read_wiki_structure"], False),
([f"{SERVER_ID}-read_wiki_contents"], False),
(["READ_WIKI_CONTENTS"], False),
],
)
async def test_key_team_listing_and_dispatch_agree(self, grants, expected):
"""The key/team grant question, driven through both production paths.
Listing and dispatch read one predicate, so a row where the tool is advertised
and then refused (or hidden while callable) cannot exist. Asserting the expected
verdict as well as the agreement matters: both paths reading one predicate makes
equality alone tautological, so a wrong predicate would keep them consistent.
Grants are stored bare, so the wire-form and case-variant rows deny; that is
deliberately unlike the server-level lists, which honor every spelling.
"""
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.server import (
filter_tools_by_key_team_permissions,
)
from litellm.proxy._types import UserAPIKeyAuth
server = MCPServer(
server_id=self.SERVER_ID,
name=self.SERVER_ID,
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
)
published = MCPTool(
name=f"{self.SERVER_ID}-read_wiki_contents", description="", inputSchema={"type": "object"}
)
auth = UserAPIKeyAuth(api_key="sk-test")
with patch.object(
MCPRequestHandler, "get_allowed_tools_for_server", AsyncMock(return_value=grants)
), patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
) as mock_manager:
mock_manager.get_mcp_server_by_id.return_value = server
listed = await filter_tools_by_key_team_permissions([published], self.SERVER_ID, auth) != []
callable_ = await MCPRequestHandler.is_tool_allowed_for_server(
tool_name="read_wiki_contents", server_id=self.SERVER_ID, user_api_key_auth=auth
)
assert listed == callable_, f"grants={grants!r} listed={listed} callable={callable_}"
assert listed is expected, f"grants={grants!r} expected={expected} got={listed}"

View file

@ -9189,3 +9189,502 @@ class TestDiscoveryFailureLogging:
assert "typo_row" in caplog.text
assert "authorization_url, token_url" in caplog.text
assert "unresolved" in caplog.text
def _unrestricted_auth() -> MagicMock:
"""A caller with no object_permission, so only server-level checks apply."""
user_api_key_auth = MagicMock()
user_api_key_auth.object_permission = None
user_api_key_auth.object_permission_id = None
return user_api_key_auth
def _permissive_proxy_logging() -> MagicMock:
proxy_logging_obj = MagicMock()
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
return proxy_logging_obj
ALIAS_LESS_SERVER_ID = "117c814c-1a2b-4c4d-8e8f-0a1b2c3d4e5f"
class TestServerToolListsHonorThePrefixBoundary:
"""The server-level allowed_tools / disallowed_tools / allowed_params checks
receive a BARE tool name. Every caller resolves the prefix boundary before
dispatch (``server.py``'s ``original_tool_name``, the Responses handler's
``sanitized_tool_name``), and ``call_tool`` hands that same value to the
upstream client verbatim, which only works because it carries no prefix.
Stored entries may also carry a prefix, and routing accepts *every* prefix
from ``iter_known_server_prefixes`` (short ID, alias, server_name,
server_id), so enforcement derives that whole set via
``iter_known_tool_name_spellings``. Rebuilding a single comparand as
``f"{server.name}-{tool_name}"`` used a field the prefix chain never reads
(``get_server_prefix`` is short_prefix, then alias, then server_name, then
server_id) and hardcoded the separator; deriving only ``get_server_prefix``
covers just the currently published spelling. Either way the check answers
for fewer spellings than are reachable, and on the blocklist arm that is a
fail-open.
"""
async def _run_check(self, server: MCPServer, name: str, arguments: dict[str, Any] | None = None) -> None:
await MCPServerManager().pre_call_tool_check(
name=name,
arguments=arguments if arguments is not None else {},
server_name=server.name,
user_api_key_auth=_unrestricted_auth(),
proxy_logging_obj=_permissive_proxy_logging(),
server=server,
)
@staticmethod
def _aliased_server(**overrides: Any) -> MCPServer:
return MCPServer(
server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21",
name="petstore_prod",
alias="petstore",
server_name="petstore_prod",
url="https://petstore.example.com/mcp",
transport=MCPTransport.http,
**overrides,
)
@staticmethod
def _alias_less_server(**overrides: Any) -> MCPServer:
# No alias and no server_name, so the published prefix is the UUID
# server_id, which itself contains the prefix separator.
return MCPServer(
server_id=ALIAS_LESS_SERVER_ID,
name=ALIAS_LESS_SERVER_ID,
url="https://wiki.example.com/mcp",
transport=MCPTransport.http,
**overrides,
)
@pytest.mark.asyncio
async def test_allowlist_entry_prefixed_with_the_alias_matches_a_bare_call(self):
# The dashboard shows tools under the published prefix, so admins store
# "petstore-getpetbyid"; the display name "petstore_prod" is not it.
server = self._aliased_server(allowed_tools=["petstore-getpetbyid"])
await self._run_check(server, "getpetbyid")
@pytest.mark.asyncio
async def test_bare_allowlist_entry_matches_on_an_alias_less_server(self):
server = self._alias_less_server(allowed_tools=["read_wiki_contents"])
await self._run_check(server, "read_wiki_contents")
@pytest.mark.asyncio
async def test_wire_form_allowlist_entry_matches_on_an_alias_less_server(self):
# The published prefix is the UUID server_id, so it contains the
# separator; the derived wire form has to reproduce it whole.
server = self._alias_less_server(allowed_tools=[f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents"])
await self._run_check(server, "read_wiki_contents")
@pytest.mark.asyncio
async def test_wire_form_entry_matches_a_native_name_that_opens_with_the_prefix(self):
# "petstore-getpetbyid" is a real upstream tool name here, so its wire
# form is "petstore-petstore-getpetbyid". Stripping the stored entry
# instead of deriving the wire form cut a boundary the caller had already
# consumed, leaving asymmetric operands that denied a permitted call.
server = self._aliased_server(allowed_tools=["petstore-petstore-getpetbyid"])
await self._run_check(server, "petstore-getpetbyid")
@pytest.mark.asyncio
async def test_wire_form_blocklist_entry_blocks_a_native_name_that_opens_with_the_prefix(self):
server = self._aliased_server(disallowed_tools=["petstore-petstore-getpetbyid"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "petstore-getpetbyid")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_wire_form_blocklist_entry_blocks_under_the_short_prefix_mode(self, monkeypatch):
# short_prefix wins in get_server_prefix but is never server.name, so the
# hand-built comparand could not match a stored wire-form entry and the
# blocklisted tool stayed callable.
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._aliased_server(short_prefix="F3X", disallowed_tools=["F3X-deletepet"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "deletepet")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_alias_form_blocklist_entry_still_blocks_under_the_short_prefix_mode(self, monkeypatch):
# Turning short prefixes on republishes every tool under the short ID,
# but routing still resolves the alias form, so an entry an admin stored
# before the flip stays reachable and has to stay enforced. Deriving only
# the published spelling silently stops honoring it: a fail-open on a
# config nobody edited.
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._aliased_server(short_prefix="F3X", disallowed_tools=["petstore-deletepet"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "deletepet")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_alias_form_allowlist_entry_still_matches_under_the_short_prefix_mode(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._aliased_server(short_prefix="F3X", allowed_tools=["petstore-getpetbyid"])
await self._run_check(server, "getpetbyid")
@pytest.mark.asyncio
async def test_server_name_form_blocklist_entry_still_blocks_under_the_short_prefix_mode(self, monkeypatch):
# server_name sits third in the prefix chain, so it is published only
# when alias and short_prefix are both absent, yet routing accepts it
# regardless.
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._aliased_server(short_prefix="F3X", disallowed_tools=["petstore_prod-deletepet"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "deletepet")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_raw_server_id_form_blocklist_entry_still_blocks_under_the_short_prefix_mode(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._alias_less_server(disallowed_tools=[f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "read_wiki_contents")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_allowed_params_keyed_by_the_alias_form_are_enforced_under_the_short_prefix_mode(self, monkeypatch):
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._aliased_server(short_prefix="F3X", allowed_params={"petstore-getpetbyid": ["petid"]})
with pytest.raises(HTTPException) as exc_info:
await self._run_check(
server,
"getpetbyid",
arguments={"petid": "7", "include_internal": "true"},
)
assert exc_info.value.status_code == 403
assert "include_internal" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_foreign_prefix_entry_does_not_match_under_the_short_prefix_mode(self, monkeypatch):
# Honoring every known prefix must not become "honor any prefix": the
# widened set is this server's spellings only.
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
server = self._aliased_server(short_prefix="F3X", allowed_tools=["other_server-getpetbyid"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "getpetbyid")
assert exc_info.value.status_code == 403
@pytest.mark.parametrize("short_prefix_mode", [False, True])
@pytest.mark.asyncio
async def test_every_spelling_routing_registers_is_also_enforced(self, monkeypatch, short_prefix_mode):
"""The invariant, driven through production code on both sides.
``_create_prefixed_tools`` decides which spellings reach dispatch, so
every key it registers has to be a spelling the blocklist can refuse.
Any key routing accepts but enforcement misses is a callable blocked
tool.
"""
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true" if short_prefix_mode else "false")
shape = self._aliased_server(short_prefix="F3X")
manager = MCPServerManager()
manager._create_prefixed_tools([MCPTool(name="deletepet", description="", inputSchema={})], shape)
registered = sorted(manager.tool_name_to_mcp_server_name_mapping)
assert len(registered) > 1
for spelling in registered:
server = self._aliased_server(short_prefix="F3X", disallowed_tools=[spelling])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "deletepet")
assert exc_info.value.status_code == 403, spelling
@pytest.mark.asyncio
async def test_wire_form_allowlist_entry_follows_a_non_default_separator(self):
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
server = self._aliased_server(allowed_tools=["petstore__getpetbyid"])
with patch.object(mcp_utils, "MCP_TOOL_PREFIX_SEPARATOR", "__"):
await self._run_check(server, "getpetbyid")
@pytest.mark.asyncio
async def test_tool_outside_the_allowlist_is_still_denied(self):
server = self._aliased_server(allowed_tools=["petstore-getpetbyid"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "deletepet")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_allowlist_entry_prefixed_for_another_server_does_not_match(self):
# Reducing both sides must not widen the allowlist across servers: a
# foreign prefix is not one of this server's known prefixes, so the
# entry keeps it and never collapses onto a bare name.
server = self._aliased_server(allowed_tools=["other_server-getpetbyid"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "getpetbyid")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_prefixed_disallowed_entry_blocks_a_bare_call(self):
# Fail-open regression: the blocklist arm answered "not banned" whenever
# the stored entry carried a prefix it failed to reconstruct.
server = self._aliased_server(disallowed_tools=["petstore-deletepet"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "deletepet")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_tool_outside_the_blocklist_is_still_allowed(self):
server = self._aliased_server(disallowed_tools=["petstore-deletepet"])
await self._run_check(server, "getpetbyid")
@pytest.mark.asyncio
async def test_allowed_params_are_enforced_for_a_bare_key(self):
server = self._alias_less_server(allowed_params={"read_wiki_contents": ["repo"]})
with pytest.raises(HTTPException) as exc_info:
await self._run_check(
server,
"read_wiki_contents",
arguments={"repo": "acme/wiki", "internal_only": "true"},
)
assert exc_info.value.status_code == 403
assert "internal_only" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_allowed_params_are_enforced_for_a_wire_form_key(self):
# A key stored under the published prefix matched nothing, so the lookup
# returned None and the check silently allowed every parameter instead of
# enforcing the configured list.
server = self._alias_less_server(allowed_params={f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents": ["repo"]})
with pytest.raises(HTTPException) as exc_info:
await self._run_check(
server,
"read_wiki_contents",
arguments={"repo": "acme/wiki", "internal_only": "true"},
)
assert exc_info.value.status_code == 403
assert "internal_only" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_allowed_params_still_accept_the_configured_parameters(self):
server = self._alias_less_server(allowed_params={f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents": ["repo"]})
await self._run_check(server, "read_wiki_contents", arguments={"repo": "acme/wiki"})
@pytest.mark.asyncio
async def test_allowed_params_are_enforced_for_a_native_name_that_opens_with_the_prefix(self):
server = self._aliased_server(allowed_params={"petstore-petstore-getpetbyid": ["petid"]})
with pytest.raises(HTTPException) as exc_info:
await self._run_check(
server,
"petstore-getpetbyid",
arguments={"petid": "7", "include_internal": "true"},
)
assert exc_info.value.status_code == 403
assert "include_internal" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_an_entry_does_not_decide_an_operation_id_registration_keeps_separate(self):
server = self._aliased_server(disallowed_tools=["foo/bar"], spec_path="/specs/petstore.yaml")
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "foo/bar")
assert exc_info.value.status_code == 403
await self._run_check(server, "foo.bar")
@pytest.mark.asyncio
async def test_a_blocklist_entry_does_not_reach_a_case_variant_sibling_tool(self):
server = self._aliased_server(disallowed_tools=["petstore-getPet"])
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "getPet")
assert exc_info.value.status_code == 403
await self._run_check(server, "getpet")
@pytest.mark.asyncio
async def test_an_allowlist_entry_does_not_grant_a_case_variant_sibling_tool(self):
server = self._aliased_server(allowed_tools=["petstore-getPet"])
await self._run_check(server, "getPet")
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "getpet")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_an_explicitly_empty_allowed_params_list_refuses_every_parameter(self):
server = self._alias_less_server(allowed_params={"read_wiki_contents": []})
with pytest.raises(HTTPException) as exc_info:
await self._run_check(server, "read_wiki_contents", arguments={"repo": "acme/wiki"})
assert exc_info.value.status_code == 403
assert "repo" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_an_explicitly_empty_allowed_params_list_still_permits_an_argument_free_call(self):
server = self._alias_less_server(allowed_params={"read_wiki_contents": []})
await self._run_check(server, "read_wiki_contents", arguments={})
class TestOpenAPIRegistryKeyMatchesRegistration:
"""OpenAPI tools are registered under ``add_server_prefix_to_name(base, get_server_prefix(server))``,
so the dispatch lookup has to build its key the same way from the bare name ``call_tool``
hands it. Rebuilding it as ``f"{server.name}-{bare_name}"`` used a field the prefix chain
never reads and hardcoded the separator, so every call on a server whose published prefix
differs from its display name failed with "not found in registry" instead of dispatching.
"""
@staticmethod
def _register(server: MCPServer, base_tool_name: str) -> str:
from litellm.proxy._experimental.mcp_server.utils import (
add_server_prefix_to_name,
get_server_prefix,
)
return add_server_prefix_to_name(base_tool_name, get_server_prefix(server))
async def _call(self, server: MCPServer, registered_key: str, bare_tool_name: str) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
async def handler(**kwargs: Any) -> str:
return "dispatched"
tool = MagicMock()
tool.handler = handler
with patch.dict(global_mcp_tool_registry.tools, {registered_key: tool}, clear=True):
return await MCPServerManager()._call_openapi_tool_handler(server, bare_tool_name, {})
@pytest.mark.asyncio
async def test_aliased_server_dispatches_when_name_differs_from_published_prefix(self):
server = MCPServer(
server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21",
name="petstore_prod",
alias="petstore",
server_name="petstore_prod",
url=None,
transport=MCPTransport.http,
spec_path="https://example.com/petstore.yaml",
)
registered_key = self._register(server, "list_pets")
assert registered_key == "petstore-list_pets"
result = await self._call(server, registered_key, "list_pets")
assert result.isError is False
assert result.content[0].text == "dispatched"
@pytest.mark.asyncio
async def test_alias_less_server_dispatches_when_the_prefix_contains_the_separator(self):
server = MCPServer(
server_id=ALIAS_LESS_SERVER_ID,
name=ALIAS_LESS_SERVER_ID,
url=None,
transport=MCPTransport.http,
spec_path="https://example.com/wiki.yaml",
)
registered_key = self._register(server, "read_wiki_contents")
assert registered_key == f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents"
result = await self._call(server, registered_key, "read_wiki_contents")
assert result.isError is False
assert result.content[0].text == "dispatched"
@pytest.mark.asyncio
async def test_dispatch_keeps_a_native_name_that_opens_with_the_prefix(self):
# Registration prefixes the upstream name whatever it looks like, so
# "petstore-list_pets" is registered as "petstore-petstore-list_pets".
# Stripping the bare name again before rebuilding the key cut that
# leading segment back off and the lookup missed.
server = MCPServer(
server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21",
name="petstore_prod",
alias="petstore",
server_name="petstore_prod",
url=None,
transport=MCPTransport.http,
spec_path="https://example.com/petstore.yaml",
)
registered_key = self._register(server, "petstore-list_pets")
assert registered_key == "petstore-petstore-list_pets"
result = await self._call(server, registered_key, "petstore-list_pets")
assert result.isError is False
assert result.content[0].text == "dispatched"
@pytest.mark.asyncio
async def test_dispatch_follows_a_non_default_prefix_separator(self):
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
server = MCPServer(
server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21",
name="petstore_prod",
alias="petstore",
server_name="petstore_prod",
url=None,
transport=MCPTransport.http,
spec_path="https://example.com/petstore.yaml",
)
with patch.object(mcp_utils, "MCP_TOOL_PREFIX_SEPARATOR", "__"):
registered_key = self._register(server, "list_pets")
assert registered_key == "petstore__list_pets"
result = await self._call(server, registered_key, "list_pets")
assert result.isError is False
assert result.content[0].text == "dispatched"
@pytest.mark.asyncio
async def test_unregistered_tool_is_still_reported_missing(self):
server = MCPServer(
server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21",
name="petstore_prod",
alias="petstore",
server_name="petstore_prod",
url=None,
transport=MCPTransport.http,
spec_path="https://example.com/petstore.yaml",
)
result = await self._call(server, "petstore-list_pets", "delete_pet")
assert result.isError is True
assert "not found in registry" in result.content[0].text

View file

@ -32,6 +32,8 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
fake_server.mcp_info = None
fake_server.server_id = "srv-1"
fake_server.server_name = "openapi-petstore"
fake_server.alias = None
fake_server.short_prefix = None
fake_tool = MagicMock()
fake_tool.name = "list_pets"
@ -111,6 +113,8 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises():
fake_server.mcp_info = None
fake_server.server_id = "srv-1"
fake_server.server_name = "openapi-petstore"
fake_server.alias = None
fake_server.short_prefix = None
fake_tool = MagicMock()
fake_tool.name = "delete_pet"

View file

@ -769,8 +769,179 @@ class TestListToolsRestAPI:
assert captured["server"] is stub_server
assert result["tools"] == ["tool-1"]
assert result["error"] is None
async def test_non_admin_ui_session_resolves_as_admitted_subject(self, monkeypatch):
"""LIT-4861: a non-admin dashboard session must act as the admitted subject on this
route, so server reachability AND tool ceilings bind to the user's grants exactly as
they do for a gateway session, never to the bare session key."""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
session_auth = UserAPIKeyAuth(
team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user"
)
admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org")
async def fake_reload(user_id):
assert user_id == "grant-user"
return admitted_auth
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
fake_reload,
)
seen_server_resolution_auths = []
async def fake_get_allowed_mcp_servers(user_api_key_auth=None, **kwargs):
seen_server_resolution_auths.append(user_api_key_auth)
return ["server-1"]
class StubServer:
alias = "server-1"
server_name = "server-1"
name = "stub"
allowed_tools = None
mcp_info = {"server_name": "stub"}
available_on_public_internet = True
stub_server = StubServer()
captured = {}
async def fake_get_tools(
server,
server_auth_header,
raw_headers=None,
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
):
captured["user_api_key_auth"] = user_api_key_auth
return ["tool-1"]
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_allowed_mcp_servers",
fake_get_allowed_mcp_servers,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_mcp_server_by_id",
lambda server_id: stub_server if server_id == "server-1" else None,
raising=False,
)
monkeypatch.setattr(
rest_endpoints,
"_get_tools_for_single_server",
fake_get_tools,
raising=False,
)
request = _build_request(path="/mcp-rest/tools/list", method="GET")
result = await rest_endpoints.list_tool_rest_api(
request,
server_id="server-1",
user_api_key_dict=session_auth,
)
resolved = [*seen_server_resolution_auths, captured["user_api_key_auth"]]
assert seen_server_resolution_auths
assert all(a.org_id == "admitted-org" and a.team_id is None for a in resolved)
assert result["tools"] == ["tool-1"]
assert result["message"] == "Successfully retrieved tools"
async def test_toolset_scoped_request_keeps_the_caller_credential(self, monkeypatch):
"""LIT-4861: the admitted subject resolves per grant source and a team source deliberately
carries none of the caller's own object_permission, so a toolset narrowing layered on top
would evaporate on every team-granted server. A toolset-scoped request therefore stays on
the caller's own credential, exactly as it did before the acting-as-user swap."""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
session_auth = UserAPIKeyAuth(
team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user"
)
scoped_auth = UserAPIKeyAuth(
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="toolset-scope",
mcp_servers=["toolset-server-1"],
)
)
reload_calls: list[str] = []
scope_inputs: list[UserAPIKeyAuth] = []
async def record_reload(user_id):
reload_calls.append(user_id)
return UserAPIKeyAuth(user_id=user_id)
class StubToolset:
toolset_id = "toolset-1"
class StubServer:
alias = "toolset-server-1"
server_name = "toolset-server-1"
name = "toolset-server-1"
allowed_tools = None
mcp_info = {"server_name": "toolset-server-1"}
available_on_public_internet = True
stub_server = StubServer()
async def fake_get_toolset_by_name_cached(prisma_client, toolset_name):
return StubToolset()
async def fake_apply_toolset_scope(user_api_key_auth, toolset_id):
scope_inputs.append(user_api_key_auth)
return scoped_auth
async def fake_get_allowed_mcp_servers(user_api_key_auth=None, **kwargs):
assert user_api_key_auth is scoped_auth
return ["toolset-server-1"]
async def fake_get_tools(server, server_auth_header, *args, **kwargs):
return ["toolset-tool-1"]
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
record_reload,
)
monkeypatch.setattr(
"litellm.proxy.utils.get_prisma_client_or_throw",
lambda *args, **kwargs: MagicMock(),
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_toolset_by_name_cached",
fake_get_toolset_by_name_cached,
raising=False,
)
monkeypatch.setattr(rest_endpoints, "_apply_toolset_scope", fake_apply_toolset_scope, raising=False)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_allowed_mcp_servers",
fake_get_allowed_mcp_servers,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_mcp_server_by_id",
lambda server_id: stub_server if server_id == "toolset-server-1" else None,
raising=False,
)
monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False)
request = _build_request(path="/mcp-rest/tools/list", method="GET")
result = await rest_endpoints.list_tool_rest_api(
request,
server_id=None,
toolset_name="research_tools",
user_api_key_dict=session_auth,
)
assert result["tools"] == ["toolset-tool-1"]
assert scope_inputs == [session_auth]
assert reload_calls == []
async def test_include_disabled_tools_is_admin_only(self, monkeypatch):
"""include_disabled_tools skips the allowlist filter only for PROXY_ADMIN;
a non-admin passing it stays filtered so the REST endpoint can't be used

View file

@ -20,7 +20,9 @@ from litellm.proxy._experimental.mcp_server.utils import (
compute_short_server_prefix,
get_server_prefix,
is_short_mcp_tool_prefix_enabled,
is_tool_name_prefixed,
iter_known_server_prefixes,
match_known_server_prefix,
strip_known_server_prefix,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -195,6 +197,70 @@ class TestStripKnownServerPrefix:
assert strip_known_server_prefix("svc-tool", None) == "tool"
# ---------------------------------------------------------------------------
# match_known_server_prefix — the shared boundary primitive
# ---------------------------------------------------------------------------
class TestMatchKnownServerPrefix:
"""Locates the boundary by matching registered prefixes instead of cutting
at the first separator, preferring the longest candidate so a prefix that
itself contains the separator still wins."""
def test_returns_matched_prefix_and_bare_name(self):
assert match_known_server_prefix("deepwiki-contents", ["deepwiki"]) == (
"deepwiki",
"contents",
)
def test_returns_none_when_no_candidate_matches(self):
assert match_known_server_prefix("contents", ["deepwiki"]) is None
def test_separator_must_follow_the_prefix(self):
assert match_known_server_prefix("deepwikicontents", ["deepwiki"]) is None
def test_uuid_prefix_survives_its_own_separators(self):
server_id = "117c814c-1a2b-3c4d-9e8f"
assert match_known_server_prefix(f"{server_id}-contents", [server_id]) == (
server_id,
"contents",
)
def test_longest_candidate_wins_over_leading_segment(self):
# "svc" is a registered prefix in its own right and also the leading
# segment of "svc-prod". A first-separator split hands the tool to
# "svc" with a bare name of "prod-run", attributing it to the wrong
# server; longest-match keeps it on "svc-prod".
assert match_known_server_prefix("svc-prod-run", ["svc", "svc-prod"]) == (
"svc-prod",
"run",
)
def test_candidates_are_normalised_before_matching(self):
assert match_known_server_prefix("my_server-run", ["my server"]) == (
"my_server",
"run",
)
def test_empty_candidate_never_matches_a_leading_separator(self):
assert match_known_server_prefix("-run", [""]) is None
class TestIsToolNamePrefixedBoundary:
"""The known-prefix gate decides which branch the call path takes, so it has
to agree with the prefix the list path actually emitted."""
def test_uuid_prefix_is_recognised(self):
server_id = "117c814c-1a2b-3c4d-9e8f"
assert is_tool_name_prefixed(f"{server_id}-contents", known_server_prefixes={server_id})
def test_unrelated_hyphenated_tool_is_still_not_prefixed(self):
# Negative control: an upstream tool whose own name contains the
# separator must not start reading as prefixed just because the gate
# got more permissive about where the boundary can fall.
assert not is_tool_name_prefixed("text-to-speech", known_server_prefixes={"deepwiki"})
# ---------------------------------------------------------------------------
# Manager-level behaviour: list + reverse-lookup
# ---------------------------------------------------------------------------

View file

@ -3,6 +3,7 @@ from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from fastapi import HTTPException
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import UserAPIKeyAuth
@ -120,3 +121,140 @@ async def test_build_effective_auth_contexts_handles_unpicklable_parent_span(
assert contexts[0].team_id == "team-span"
assert contexts[0].parent_otel_span is parent_span
@pytest.mark.asyncio
async def test_build_effective_auth_contexts_appends_admitted_user_context(monkeypatch):
"""LIT-4861: the dashboard session must resolve with the user's admitted identity so the
page list and every per-server action endpoint see user-level grants the same way the
gateway session does."""
user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-42")
admitted_auth = UserAPIKeyAuth(user_id="user-42")
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids",
AsyncMock(return_value=["team-one"]),
)
reload_mock = AsyncMock(return_value=admitted_auth)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
)
contexts = await build_effective_auth_contexts(user_auth)
assert contexts[-1].user_id == "user-42" and contexts[-1].team_id is None
assert [ctx.team_id for ctx in contexts[:-1]] == ["team-one"]
reload_mock.assert_awaited_once_with("user-42")
@pytest.mark.asyncio
async def test_build_effective_auth_contexts_never_widens_caller_passed_keys(monkeypatch):
normal_user = UserAPIKeyAuth(team_id="regular-team", user_id="user-1")
reload_mock = AsyncMock()
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
)
contexts = await build_effective_auth_contexts(normal_user)
assert contexts == [normal_user]
reload_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_build_effective_auth_contexts_survives_admitted_reload_failure(monkeypatch):
user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-9")
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids",
AsyncMock(return_value=["team-a"]),
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
AsyncMock(side_effect=HTTPException(status_code=503, detail="db down")),
)
contexts = await build_effective_auth_contexts(user_auth)
assert [ctx.team_id for ctx in contexts] == ["team-a"]
@pytest.mark.asyncio
async def test_acting_user_auth_returns_admitted_subject_for_non_admin_sessions(monkeypatch):
"""LIT-4861: acting-as-user MCP routes must resolve a non-admin dashboard session as the
admitted subject so tool ceilings, reachability, and limits bind exactly as on /mcp."""
from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth
user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-42", user_role="internal_user")
admitted_auth = UserAPIKeyAuth(user_id="user-42")
reload_mock = AsyncMock(return_value=admitted_auth)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
)
result = await acting_user_auth(user_auth)
assert result.user_id == "user-42" and result.team_id is None
reload_mock.assert_awaited_once_with("user-42")
@pytest.mark.asyncio
async def test_acting_user_auth_keeps_admin_sessions_and_passed_keys_unchanged(monkeypatch):
from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth
reload_mock = AsyncMock()
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
)
admin_session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="admin-1", user_role="proxy_admin")
assert await acting_user_auth(admin_session) is admin_session
passed_key = UserAPIKeyAuth(team_id="regular-team", user_id="user-1", user_role="internal_user")
assert await acting_user_auth(passed_key) is passed_key
reload_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_acting_user_auth_falls_back_to_session_auth_on_reload_failure(monkeypatch):
from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth
user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-9", user_role="internal_user")
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
AsyncMock(side_effect=HTTPException(status_code=503, detail="db down")),
)
assert await acting_user_auth(user_auth) is user_auth
@pytest.mark.asyncio
async def test_admitted_user_context_carries_the_request_span(monkeypatch):
"""Swapping the principal must not drop the request: the admitted subject is rebuilt from the
user row and carries no span of its own, so every consumer would otherwise lose trace linkage
for the resolution and logging it drives."""
from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth
class DummySpan:
def __init__(self) -> None:
self._lock = threading.RLock()
parent_span = DummySpan()
user_auth = UserAPIKeyAuth(
team_id=UI_SESSION_TOKEN_TEAM_ID,
user_id="user-42",
user_role="internal_user",
parent_otel_span=parent_span,
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
AsyncMock(return_value=UserAPIKeyAuth(user_id="user-42")),
)
assert (await acting_user_auth(user_auth)).parent_otel_span is parent_span
assert (await build_effective_auth_contexts(user_auth))[-1].parent_otel_span is parent_span

View file

@ -6138,3 +6138,255 @@ def test_bundled_openapi_registry_parses_and_entries_are_well_formed():
)
for tool in entry.get("key_tools", []):
assert tool.get("name") and tool.get("description"), f"{entry['name']}: malformed key_tool"
class TestConnectedAppViewAnnotation:
"""LIT-4861: GET /v1/mcp/server?connected_app_view=true must annotate each server with
whether the caller's gateway OAuth sessions (connected apps) are served it on /mcp.
The view is honored only for the dashboard's UI session credential; a caller-passed
virtual key must never be widened to its owning user's identity."""
def _ui_session_auth(self, user_role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> UserAPIKeyAuth:
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
return generate_mock_user_api_key_auth(user_role=user_role, team_id=UI_SESSION_TOKEN_TEAM_ID)
def _mock_manager(self, servers, reachable_ids):
mock_manager = MagicMock()
mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=servers)
mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=servers)
mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=reachable_ids)
return mock_manager
def _servers(self):
return [
generate_mock_mcp_server_db_record(server_id="server-1", alias="Granted"),
generate_mock_mcp_server_db_record(server_id="server-2", alias="Ungranted"),
]
@pytest.mark.asyncio
async def test_connected_app_view_annotates_reachability_via_admitted_resolver(self):
caller_auth = self._ui_session_auth()
admitted_auth = UserAPIKeyAuth(user_id="test_user_id")
admitted_auth.mcp_admitted_user_subject = True
mock_manager = self._mock_manager(self._servers(), ["server-1"])
reload_mock = AsyncMock(return_value=admitted_auth)
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[caller_auth]),
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
flags = {server.server_id: server.connected_app_reachable for server in result}
assert flags == {"server-1": True, "server-2": False}
reload_mock.assert_awaited_once_with("test_user_id")
mock_manager.get_allowed_mcp_servers.assert_awaited_once_with(admitted_auth)
@pytest.mark.asyncio
async def test_connected_app_view_stamps_view_all_list_and_survives_non_admin_sanitizer(self):
"""view_all preempts the manager's admin shortcut with a second whole-registry
shortcut; the annotation must still land, and must survive the non-admin sanitizer."""
caller_auth = self._ui_session_auth(user_role=LitellmUserRoles.INTERNAL_USER)
mock_manager = self._mock_manager(self._servers(), ["server-2"])
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode",
return_value="view_all",
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id")),
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
mock_manager.get_all_mcp_servers_unfiltered.assert_awaited_once()
flags = {server.server_id: server.connected_app_reachable for server in result}
assert flags == {"server-1": False, "server-2": True}
@pytest.mark.asyncio
async def test_connected_app_view_lists_user_granted_servers_via_admitted_context(self):
"""A server granted only through the user's own object permission must be listed and
flagged reachable: the REAL build_effective_auth_contexts appends the admitted-user
context, so the page and every action endpoint resolve it identically."""
caller_auth = self._ui_session_auth(user_role=LitellmUserRoles.INTERNAL_USER)
admitted_auth = UserAPIKeyAuth(user_id="test_user_id", org_id="admitted-org")
listed_row = generate_mock_mcp_server_db_record(server_id="server-1", alias="TeamGranted")
user_granted_row = generate_mock_mcp_server_db_record(server_id="server-2", alias="UserGranted")
async def per_context_servers(user_api_key_auth=None):
if user_api_key_auth is not None and user_api_key_auth.org_id == "admitted-org":
return [listed_row, user_granted_row]
return [listed_row]
mock_manager = MagicMock()
mock_manager.get_all_allowed_mcp_servers = AsyncMock(side_effect=per_context_servers)
mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["server-1", "server-2"])
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids",
AsyncMock(return_value=[]),
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
AsyncMock(return_value=admitted_auth),
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
flags = {server.server_id: server.connected_app_reachable for server in result}
assert flags == {"server-1": True, "server-2": True}
@pytest.mark.asyncio
async def test_connected_app_view_fails_closed_when_admitted_reload_fails(self):
caller_auth = self._ui_session_auth()
mock_manager = self._mock_manager(self._servers(), ["server-1"])
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[caller_auth]),
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
AsyncMock(side_effect=HTTPException(status_code=401, detail="expired")),
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
assert all(server.connected_app_reachable is False for server in result)
@pytest.mark.asyncio
async def test_connected_app_view_off_leaves_field_unset(self):
caller_auth = generate_mock_user_api_key_auth()
mock_manager = self._mock_manager(self._servers(), ["server-1"])
reload_mock = AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id"))
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[caller_auth]),
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth)
assert all(server.connected_app_reachable is None for server in result)
reload_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_connected_app_view_userless_ui_credential_leaves_field_unset(self):
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
caller_auth = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="test_api_key", team_id=UI_SESSION_TOKEN_TEAM_ID
)
caller_auth.user_id = None
mock_manager = self._mock_manager(self._servers(), ["server-1"])
reload_mock = AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id"))
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[caller_auth]),
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
assert all(server.connected_app_reachable is None for server in result)
reload_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_connected_app_view_ignored_for_caller_passed_virtual_keys(self):
"""A virtual key the user passes themselves is never widened to the owning user's
identity: the view param is a no-op and the admitted resolver is never consulted."""
caller_auth = generate_mock_user_api_key_auth(team_id="some-real-team")
mock_manager = self._mock_manager(self._servers(), ["server-1"])
reload_mock = AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id"))
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[caller_auth]),
),
patch(
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user",
reload_mock,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_all_mcp_servers,
)
result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True)
assert all(server.connected_app_reachable is None for server in result)
reload_mock.assert_not_awaited()

View file

@ -294,19 +294,22 @@ class TestUnifiedGuardrailCallTypeResolution:
response_body = {"candidates": [{"content": {"parts": [{"text": "hello"}]}}]}
with patch(
"litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail.load_guardrail_translation_mappings"
) as mock_load:
mock_handler_instance = AsyncMock()
mock_handler_instance.process_output_response = AsyncMock(
return_value=response_body
)
mock_handler_class = MagicMock(return_value=mock_handler_instance)
mock_handler_instance = AsyncMock()
mock_handler_instance.process_output_response = AsyncMock(
return_value=response_body
)
mock_handler_class = MagicMock(return_value=mock_handler_instance)
from litellm.types.utils import CallTypes
mock_load.return_value = {CallTypes.pass_through: mock_handler_class}
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
unified_guardrail as unified_guardrail_module,
)
from litellm.types.utils import CallTypes
with patch.object(
unified_guardrail_module,
"endpoint_guardrail_translation_mappings",
{CallTypes.pass_through: mock_handler_class},
):
result = await unified.async_post_call_success_hook(
data=data,
user_api_key_dict=user_api_key_dict,

View file

@ -4,6 +4,9 @@ Unit tests for AttachmentRegistry - tests policy attachment matching.
Tests the main entry point: get_attached_policies()
"""
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.policy_engine.attachment_registry import (
@ -389,3 +392,75 @@ class TestAttachmentRegistrySingleton:
registry1 = get_attachment_registry()
registry2 = get_attachment_registry()
assert registry1 is registry2
def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scope=None, teams=None):
row = MagicMock()
row.attachment_id = attachment_id
row.policy_name = policy_name
row.scope = scope
row.teams = teams or []
row.keys = []
row.models = []
row.tags = []
row.created_at = datetime.now(timezone.utc)
row.updated_at = datetime.now(timezone.utc)
row.created_by = None
row.updated_by = None
return row
def _prisma_with_attachment_rows(rows):
prisma = MagicMock()
prisma.db.litellm_policyattachmenttable.find_many = AsyncMock(return_value=rows)
return prisma
class TestConfigAttachmentsPreservedAcrossDbSync:
"""Config-defined attachments must survive sync_attachments_from_db (regression for issue #35255)."""
@pytest.mark.asyncio
async def test_sync_with_empty_db_preserves_config_attachments(self):
registry = AttachmentRegistry()
registry.load_attachments([{"policy": "config-policy", "scope": "*"}])
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([]))
context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="gpt-5.2")
assert registry.get_attached_policies(context) == ["config-policy"]
@pytest.mark.asyncio
async def test_sync_merges_db_attachments_with_config_attachments(self):
registry = AttachmentRegistry()
registry.load_attachments([{"policy": "config-policy", "scope": "*"}])
db_row = _make_db_attachment_row(policy_name="db-policy", teams=["db-team"])
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row]))
assert len(registry.get_all_attachments()) == 2
assert len(registry.get_config_attachments()) == 1
context = PolicyMatchContext(team_alias="db-team", key_alias="k", model="gpt-5.2")
attached = registry.get_attached_policies(context)
assert "config-policy" in attached
assert "db-policy" in attached
@pytest.mark.asyncio
async def test_repeated_syncs_do_not_duplicate_config_attachments(self):
registry = AttachmentRegistry()
registry.load_attachments([{"policy": "config-policy", "scope": "*"}])
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([]))
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([]))
assert len(registry.get_all_attachments()) == 1
@pytest.mark.asyncio
async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
registry = AttachmentRegistry()
registry.load_attachments([{"policy": "config-policy", "scope": "*"}])
registry.clear()
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([]))
assert registry.get_all_attachments() == []
assert registry.get_config_attachments() == ()

View file

@ -0,0 +1,248 @@
"""
Unit tests for policy_engine/policy_endpoints.py list endpoints.
Regression tests for issue #35255: config-defined policies and attachments must be
returned by the list endpoints (marked definition_location="config"), DB rows must keep
their exact shape, and the endpoints must not 500 when no database is connected.
"""
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm.proxy.policy_engine.policy_endpoints as policy_endpoints
from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry
from litellm.proxy.policy_engine.policy_registry import PolicyRegistry
def _make_policy_row(
policy_id="uuid-1",
policy_name="db-policy",
version_status="production",
guardrails_add=None,
):
row = MagicMock()
row.policy_id = policy_id
row.policy_name = policy_name
row.version_number = 1
row.version_status = version_status
row.parent_version_id = None
row.is_latest = True
row.published_at = None
row.production_at = None
row.inherit = None
row.description = "db description"
row.guardrails_add = guardrails_add or []
row.guardrails_remove = []
row.condition = None
row.pipeline = None
row.created_at = datetime.now(timezone.utc)
row.updated_at = datetime.now(timezone.utc)
row.created_by = "admin"
row.updated_by = "admin"
return row
def _make_attachment_row(attachment_id="att-1", policy_name="db-policy", scope="*"):
row = MagicMock()
row.attachment_id = attachment_id
row.policy_name = policy_name
row.scope = scope
row.teams = []
row.keys = []
row.models = []
row.tags = []
row.created_at = datetime.now(timezone.utc)
row.updated_at = datetime.now(timezone.utc)
row.created_by = "admin"
row.updated_by = "admin"
return row
@pytest.fixture
def policy_registry(monkeypatch):
registry = PolicyRegistry()
monkeypatch.setattr(policy_endpoints, "get_policy_registry", lambda: registry)
return registry
@pytest.fixture
def attachment_registry(monkeypatch):
registry = AttachmentRegistry()
monkeypatch.setattr(policy_endpoints, "get_attachment_registry", lambda: registry)
return registry
def _set_prisma(monkeypatch, prisma):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
class TestListPoliciesIncludesConfig:
@pytest.mark.asyncio
async def test_returns_config_policies_without_prisma(self, policy_registry, monkeypatch):
_set_prisma(monkeypatch, None)
policy_registry.load_policies(
{"config-policy": {"description": "from config", "guardrails": {"add": ["tooling"]}}}
)
response = await policy_endpoints.list_policies()
assert response.total_count == 1
entry = response.policies[0]
assert entry.policy_name == "config-policy"
assert entry.policy_id == "config-policy"
assert entry.definition_location == "config"
assert entry.version_status == "production"
assert entry.guardrails_add == ["tooling"]
assert entry.description == "from config"
assert entry.created_at is None
@pytest.mark.asyncio
async def test_merges_db_rows_with_config_and_keeps_db_row_shape(self, policy_registry, monkeypatch):
row = _make_policy_row(policy_id="uuid-1", policy_name="db-policy", guardrails_add=["db-guard"])
prisma = MagicMock()
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row])
_set_prisma(monkeypatch, prisma)
policy_registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
response = await policy_endpoints.list_policies()
assert response.total_count == 2
db_entry = next(p for p in response.policies if p.policy_name == "db-policy")
assert db_entry.definition_location == "db"
assert db_entry.policy_id == "uuid-1"
assert db_entry.guardrails_add == ["db-guard"]
assert db_entry.description == "db description"
assert db_entry.created_at == row.created_at
assert db_entry.created_by == "admin"
config_entry = next(p for p in response.policies if p.policy_name == "config-policy")
assert config_entry.definition_location == "config"
@pytest.mark.asyncio
async def test_db_policy_shadows_config_policy_with_same_name(self, policy_registry, monkeypatch):
row = _make_policy_row(policy_id="uuid-1", policy_name="shared-name", guardrails_add=["db-guard"])
prisma = MagicMock()
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row])
_set_prisma(monkeypatch, prisma)
policy_registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
response = await policy_endpoints.list_policies()
assert response.total_count == 1
assert response.policies[0].definition_location == "db"
assert response.policies[0].guardrails_add == ["db-guard"]
@pytest.mark.asyncio
async def test_draft_db_policy_does_not_hide_enforced_config_policy(self, policy_registry, monkeypatch):
"""
Runtime sync only lets production DB versions override a config policy,
so a draft or published DB version sharing the name must not suppress
the config entry: the config version is still the one being enforced,
and hiding it makes the list API disagree with actual enforcement.
"""
row = _make_policy_row(
policy_id="uuid-1", policy_name="shared-name", version_status="draft", guardrails_add=["db-guard"]
)
prisma = MagicMock()
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row])
_set_prisma(monkeypatch, prisma)
policy_registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
response = await policy_endpoints.list_policies()
assert response.total_count == 2
config_entry = next(p for p in response.policies if p.definition_location == "config")
assert config_entry.policy_name == "shared-name"
assert config_entry.version_status == "production"
assert config_entry.guardrails_add == ["config-guard"]
db_entry = next(p for p in response.policies if p.definition_location == "db")
assert db_entry.version_status == "draft"
@pytest.mark.asyncio
async def test_stale_registry_provenance_does_not_hide_config_policy(self, policy_registry, monkeypatch):
"""
Another proxy instance can delete or demote the production DB override
between registry syncs. The endpoint's fresh DB query is the source of
truth for conflicts; stale in-memory provenance from the last sync must
not suppress the config entry once no production override exists.
"""
policy_registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
production_row = _make_policy_row(policy_id="uuid-1", policy_name="shared-name", guardrails_add=["db-guard"])
sync_prisma = MagicMock()
sync_prisma.db.litellm_policytable.find_many = AsyncMock(side_effect=[[production_row], []])
await policy_registry.sync_policies_from_db(sync_prisma)
assert policy_registry.get_source("shared-name") == "db"
fresh_prisma = MagicMock()
fresh_prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[])
_set_prisma(monkeypatch, fresh_prisma)
response = await policy_endpoints.list_policies()
assert response.total_count == 1
entry = response.policies[0]
assert entry.policy_name == "shared-name"
assert entry.definition_location == "config"
assert entry.guardrails_add == ["config-guard"]
@pytest.mark.asyncio
async def test_version_status_filter_excludes_config_policies(self, policy_registry, monkeypatch):
row = _make_policy_row(policy_id="uuid-1", policy_name="db-policy", version_status="draft")
prisma = MagicMock()
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row])
_set_prisma(monkeypatch, prisma)
policy_registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
response = await policy_endpoints.list_policies(version_status="draft")
assert response.total_count == 1
assert response.policies[0].policy_name == "db-policy"
assert response.policies[0].definition_location == "db"
@pytest.mark.asyncio
async def test_production_filter_includes_config_policies(self, policy_registry, monkeypatch):
_set_prisma(monkeypatch, None)
policy_registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
response = await policy_endpoints.list_policies(version_status="production")
assert response.total_count == 1
assert response.policies[0].definition_location == "config"
class TestListAttachmentsIncludesConfig:
@pytest.mark.asyncio
async def test_returns_config_attachments_without_prisma(self, attachment_registry, monkeypatch):
_set_prisma(monkeypatch, None)
attachment_registry.load_attachments([{"policy": "config-policy", "scope": "*"}])
response = await policy_endpoints.list_policy_attachments()
assert response.total_count == 1
entry = response.attachments[0]
assert entry.attachment_id == "config-0"
assert entry.policy_name == "config-policy"
assert entry.scope == "*"
assert entry.definition_location == "config"
assert entry.created_at is None
@pytest.mark.asyncio
async def test_merges_db_attachments_with_config_and_keeps_db_row_shape(self, attachment_registry, monkeypatch):
row = _make_attachment_row(attachment_id="att-1", policy_name="db-policy")
prisma = MagicMock()
prisma.db.litellm_policyattachmenttable.find_many = AsyncMock(return_value=[row])
_set_prisma(monkeypatch, prisma)
attachment_registry.load_attachments([{"policy": "config-policy", "scope": "*"}])
response = await policy_endpoints.list_policy_attachments()
assert response.total_count == 2
db_entry = next(a for a in response.attachments if a.policy_name == "db-policy")
assert db_entry.attachment_id == "att-1"
assert db_entry.definition_location == "db"
assert db_entry.created_at == row.created_at
assert db_entry.created_by == "admin"
config_entry = next(a for a in response.attachments if a.policy_name == "config-policy")
assert config_entry.attachment_id == "config-0"
assert config_entry.definition_location == "config"

View file

@ -13,8 +13,10 @@ from litellm.proxy.policy_engine.policy_registry import (
get_policy_registry,
)
from litellm.types.proxy.policy_engine import (
Policy,
PolicyCreateRequest,
PolicyDBResponse,
PolicyGuardrails,
PolicyUpdateRequest,
)
@ -450,3 +452,182 @@ class TestGetPolicyRegistrySingleton:
a = get_policy_registry()
b = get_policy_registry()
assert a is b
def _prisma_with_policy_rows(production_rows, non_production_rows=None):
prisma = MagicMock()
prisma.db.litellm_policytable.find_many = AsyncMock(side_effect=[production_rows, non_production_rows or []])
return prisma
class TestConfigPoliciesPreservedAcrossDbSync:
"""Config-defined policies must survive sync_policies_from_db (regression for issue #35255)."""
@pytest.mark.asyncio
async def test_sync_with_empty_db_preserves_config_policies(self):
registry = PolicyRegistry()
registry.load_policies({"config-policy": {"description": "from config", "guardrails": {"add": ["tooling"]}}})
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
assert registry.has_policy("config-policy")
policy = registry.get_policy("config-policy")
assert policy is not None
assert policy.guardrails.add == ["tooling"]
assert registry.get_source("config-policy") == "config"
@pytest.mark.asyncio
async def test_sync_merges_db_policies_with_config_policies(self):
registry = PolicyRegistry()
registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
db_row = _make_row(policy_id="db-1", policy_name="db-policy", guardrails_add=["db-guard"])
await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
assert registry.has_policy("config-policy")
assert registry.has_policy("db-policy")
assert registry.get_source("config-policy") == "config"
assert registry.get_source("db-policy") == "db"
@pytest.mark.asyncio
async def test_db_wins_on_policy_name_conflict(self):
registry = PolicyRegistry()
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"])
await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
policy = registry.get_policy("shared-name")
assert policy is not None
assert policy.guardrails.add == ["db-guard"]
assert registry.get_source("shared-name") == "db"
@pytest.mark.asyncio
async def test_config_policy_restored_after_conflicting_db_row_deleted(self):
registry = PolicyRegistry()
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"])
await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
policy = registry.get_policy("shared-name")
assert policy is not None
assert policy.guardrails.add == ["config-guard"]
assert registry.get_source("shared-name") == "config"
@pytest.mark.asyncio
async def test_config_policy_resolves_guardrails_after_sync(self):
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
registry = PolicyRegistry()
registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
resolved = PolicyResolver.resolve_policy_guardrails(
policy_name="config-policy",
policies=registry.get_all_policies(),
context=None,
)
assert resolved.guardrails == ["tooling"]
@pytest.mark.asyncio
async def test_add_policy_with_config_source_survives_sync(self):
registry = PolicyRegistry()
registry.add_policy(
"late-config-policy",
Policy(guardrails=PolicyGuardrails(add=["tooling"])),
source="config",
)
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
assert registry.has_policy("late-config-policy")
assert registry.get_source("late-config-policy") == "config"
@pytest.mark.asyncio
async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
registry = PolicyRegistry()
registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
registry.clear()
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
assert not registry.has_policy("config-policy")
assert registry.get_source("config-policy") is None
class TestRemovePolicyRestoresConfigFallback:
"""Deleting a same-named DB override must re-activate the config policy immediately, not at the next sync."""
def test_remove_policy_restores_config_version_immediately(self):
registry = PolicyRegistry()
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
assert registry.remove_policy("shared-name") is True
policy = registry.get_policy("shared-name")
assert policy is not None
assert policy.guardrails.add == ["config-guard"]
assert registry.get_source("shared-name") == "config"
def test_remove_policy_without_config_fallback_removes_entirely(self):
registry = PolicyRegistry()
registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])))
assert registry.remove_policy("db-only") is True
assert not registry.has_policy("db-only")
assert registry.get_source("db-only") is None
def test_remove_missing_policy_returns_false(self):
registry = PolicyRegistry()
assert registry.remove_policy("missing") is False
@pytest.mark.asyncio
async def test_delete_production_override_reactivates_config_policy_and_says_so(self):
registry = PolicyRegistry()
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
prisma = MagicMock()
prod_row = _make_row(policy_id="prod-1", policy_name="shared-name", version_status="production")
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)
prisma.db.litellm_policytable.delete = AsyncMock()
result = await registry.delete_policy_from_db(policy_id="prod-1", prisma_client=prisma)
assert "config" in result["warning"]
policy = registry.get_policy("shared-name")
assert policy is not None
assert policy.guardrails.add == ["config-guard"]
assert registry.get_source("shared-name") == "config"
@pytest.mark.asyncio
async def test_delete_all_versions_reactivates_config_policy(self):
registry = PolicyRegistry()
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
prisma = MagicMock()
prisma.db.litellm_policytable.delete_many = AsyncMock()
result = await registry.delete_all_versions(policy_name="shared-name", prisma_client=prisma)
assert registry.get_source("shared-name") == "config"
policy = registry.get_policy("shared-name")
assert policy is not None
assert policy.guardrails.add == ["config-guard"]
assert "config" in result["warning"]
async def test_delete_all_versions_without_config_twin_has_no_warning(self):
registry = PolicyRegistry()
registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
prisma = MagicMock()
prisma.db.litellm_policytable.delete_many = AsyncMock()
result = await registry.delete_all_versions(policy_name="db-only", prisma_client=prisma)
assert registry.get_policy("db-only") is None
assert "warning" not in result

View file

@ -28,9 +28,11 @@ from litellm.exceptions import MidStreamFallbackError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import (
_ERROR_CODE_HTTP_STATUS,
BaseResponsesAPIStreamingIterator,
ResponsesAPIStreamingIterator,
SyncResponsesAPIStreamingIterator,
_status_code_for_error_fields,
)
from litellm.types.llms.openai import (
ErrorEvent,
@ -355,3 +357,70 @@ def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event():
pass
assert exc_info.value.status_code == 429
assert isinstance(exc_info.value.original_exception, litellm.APIError)
def test_every_openai_sdk_response_error_code_has_explicit_status_mapping():
from typing import get_args
from openai.types.responses.response_error import ResponseError
sdk_codes = set(get_args(ResponseError.model_fields["code"].annotation))
unmapped = sdk_codes - set(_ERROR_CODE_HTTP_STATUS)
assert unmapped == set(), (
f"OpenAI SDK ResponseError codes missing from _ERROR_CODE_HTTP_STATUS: {sorted(unmapped)}; "
"classify each new code with an explicit HTTP status instead of letting it default to 500"
)
@pytest.mark.parametrize(
"code,expected_status",
[
("server_error", 500),
("rate_limit_exceeded", 429),
("insufficient_quota", 429),
("vector_store_timeout", 504),
("invalid_prompt", 400),
("invalid_image", 400),
("invalid_image_format", 400),
("invalid_base64_image", 400),
("invalid_image_url", 400),
("image_too_large", 400),
("image_too_small", 400),
("image_parse_error", 400),
("image_content_policy_violation", 400),
("invalid_image_mode", 400),
("image_file_too_large", 400),
("unsupported_image_media_type", 400),
("empty_image_file", 400),
("failed_to_download_image", 400),
("image_file_not_found", 400),
("totally_unknown_future_code", 500),
],
)
def test_status_code_for_documented_response_error_codes(code: str, expected_status: int):
assert _status_code_for_error_fields(None, code) == expected_status
def test_specific_error_code_wins_over_generic_error_type():
assert _status_code_for_error_fields("server_error", "invalid_image") == 400
def test_maybe_raise_for_response_failed_event_maps_image_code_to_400():
iterator = _make_iterator()
mock_response_obj = Mock()
mock_response_obj.error = {"code": "image_content_policy_violation", "message": "image rejected"}
chunk = Mock()
chunk.type = "response.failed"
chunk.response = mock_response_obj
with pytest.raises(litellm.APIError) as exc_info:
iterator._maybe_raise_for_error_event(chunk)
assert exc_info.value.status_code == 400
assert not isinstance(exc_info.value, MidStreamFallbackError)
def test_maybe_raise_for_error_event_maps_vector_store_timeout_to_retriable_504():
iterator = _make_iterator()
chunk = _make_error_chunk("server_error", "vector_store_timeout", "vector store timed out")
with pytest.raises(MidStreamFallbackError) as exc_info:
iterator._maybe_raise_for_error_event(chunk)
assert exc_info.value.status_code == 504

View file

@ -1,159 +0,0 @@
import json
from pathlib import Path
import pytest
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
GPT_5_6_MODELS = ("gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna")
STANDARD_PRICING = {
"gpt-5.6": (5e-06, 3e-05, 5e-07, 6.25e-06),
"gpt-5.6-sol": (5e-06, 3e-05, 5e-07, 6.25e-06),
"gpt-5.6-terra": (2.5e-06, 1.5e-05, 2.5e-07, 3.125e-06),
"gpt-5.6-luna": (1e-06, 6e-06, 1e-07, 1.25e-06),
}
@pytest.mark.parametrize("model", GPT_5_6_MODELS)
def test_openai_gpt_5_6_model_info(model):
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
with open(json_path) as f:
model_cost = json.load(f)
info = model_cost.get(model)
assert info is not None, f"{model} not found in model_prices_and_context_window.json"
assert info["litellm_provider"] == "openai"
assert info["mode"] == "chat"
input_cost, output_cost, cache_read_cost, cache_write_cost = STANDARD_PRICING[model]
assert info["input_cost_per_token"] == input_cost
assert info["output_cost_per_token"] == output_cost
assert info["cache_read_input_token_cost"] == cache_read_cost
assert info["cache_creation_input_token_cost"] == cache_write_cost
assert info["cache_creation_input_token_cost"] == pytest.approx(input_cost * 1.25)
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5)
assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2)
assert info["max_input_tokens"] == 1050000
assert info["max_output_tokens"] == 128000
assert info["max_tokens"] == 128000
assert info["supports_function_calling"] is True
assert info["supports_prompt_caching"] is True
assert info["supports_reasoning"] is True
assert info["supports_response_schema"] is True
assert info["supports_tool_choice"] is True
assert info["supports_vision"] is True
assert info["supports_web_search"] is True
assert info["supports_none_reasoning_effort"] is True
assert info["supports_xhigh_reasoning_effort"] is True
assert info["supports_minimal_reasoning_effort"] is False
assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/batch", "/v1/responses"]
assert info["supported_modalities"] == ["text", "image"]
assert info["supported_output_modalities"] == ["text"]
routed_model, provider, _, _ = get_llm_provider(model=f"openai/{model}")
assert routed_model == model
assert provider == "openai"
AZURE_GLOBAL_MODELS = (
"azure/gpt-5.6",
"azure/gpt-5.6-sol",
"azure/gpt-5.6-terra",
"azure/gpt-5.6-luna",
)
AZURE_REGIONAL_MODELS = tuple(
f"azure/{region}/{tier}"
for region in ("us", "eu")
for tier in ("gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna")
)
def _tier_key(azure_model):
return azure_model.split("/")[-1]
def _load_main():
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
with open(json_path) as f:
return json.load(f)
@pytest.mark.parametrize("model", AZURE_GLOBAL_MODELS)
def test_azure_gpt_5_6_global_model_info(model):
model_cost = _load_main()
info = model_cost.get(model)
assert info is not None, f"{model} not found in model_prices_and_context_window.json"
assert info["litellm_provider"] == "azure"
assert info["mode"] == "chat"
input_cost, output_cost, cache_read_cost, _ = STANDARD_PRICING[_tier_key(model)]
assert info["input_cost_per_token"] == input_cost
assert info["output_cost_per_token"] == output_cost
assert info["cache_read_input_token_cost"] == cache_read_cost
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5)
assert info["input_cost_per_token_priority"] == pytest.approx(input_cost * 2)
assert info["output_cost_per_token_priority"] == pytest.approx(output_cost * 2)
assert info["input_cost_per_token_above_272k_tokens_priority"] == pytest.approx(input_cost * 4)
assert info["output_cost_per_token_above_272k_tokens_priority"] == pytest.approx(output_cost * 3)
assert info["max_input_tokens"] == 1050000
assert info["max_output_tokens"] == 128000
assert info["supports_reasoning"] is True
routed_model, provider, _, _ = get_llm_provider(model=model)
assert provider == "azure"
@pytest.mark.parametrize("model", AZURE_REGIONAL_MODELS)
def test_azure_gpt_5_6_regional_model_info(model):
model_cost = _load_main()
info = model_cost.get(model)
assert info is not None, f"{model} not found in model_prices_and_context_window.json"
assert info["litellm_provider"] == "azure"
assert info["mode"] == "chat"
input_cost, output_cost, cache_read_cost, _ = STANDARD_PRICING[_tier_key(model)]
assert info["input_cost_per_token"] == pytest.approx(input_cost * 1.1)
assert info["output_cost_per_token"] == pytest.approx(output_cost * 1.1)
assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost * 1.1)
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2.2)
assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.65)
assert info["input_cost_per_token_priority"] == pytest.approx(input_cost * 2.75)
assert info["output_cost_per_token_priority"] == pytest.approx(output_cost * 2.75)
assert info["max_input_tokens"] == 1050000
assert info["max_output_tokens"] == 128000
assert info["supports_reasoning"] is True
_, provider, _, _ = get_llm_provider(model=model)
assert provider == "azure"
def test_gpt_5_6_backup_matches_main():
"""Ensure the bundled model cost map stays in sync with the canonical file."""
repo_root = Path(__file__).parents[2]
main_path = repo_root / "model_prices_and_context_window.json"
backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json"
with open(main_path) as f:
main_cost = json.load(f)
with open(backup_path) as f:
backup_cost = json.load(f)
for model in GPT_5_6_MODELS + AZURE_GLOBAL_MODELS + AZURE_REGIONAL_MODELS:
assert backup_cost.get(model) == main_cost.get(model), (
f"{model} differs between main and backup model cost maps"
)

View file

@ -768,11 +768,17 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"cache_creation_input_token_cost_above_1hr": {"type": "number"},
"cache_creation_input_token_cost_above_200k_tokens": {"type": "number"},
"cache_creation_input_token_cost_above_272k_tokens": {"type": "number"},
"cache_creation_input_token_cost_above_272k_tokens_flex": {
"type": "number"
},
"cache_creation_input_token_cost_flex": {"type": "number"},
"cache_creation_input_token_cost_priority": {"type": "number"},
"cache_read_input_token_cost": {"type": "number"},
"cache_read_input_token_cost_above_200k_tokens": {"type": "number"},
"cache_read_input_token_cost_above_272k_tokens": {"type": "number"},
"cache_read_input_token_cost_above_272k_tokens_flex": {
"type": "number"
},
"cache_read_input_token_cost_above_512k_tokens": {"type": "number"},
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": {
"type": "number"
@ -806,11 +812,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"input_cost_per_token_priority": {"type": "number"},
"input_cost_per_token_above_200k_tokens_priority": {"type": "number"},
"input_cost_per_token_above_272k_tokens_priority": {"type": "number"},
"input_cost_per_token_above_272k_tokens_flex": {"type": "number"},
"input_cost_per_audio_token_priority": {"type": "number"},
"output_cost_per_token_flex": {"type": "number"},
"output_cost_per_token_priority": {"type": "number"},
"output_cost_per_token_above_200k_tokens_priority": {"type": "number"},
"output_cost_per_token_above_272k_tokens_priority": {"type": "number"},
"output_cost_per_token_above_272k_tokens_flex": {"type": "number"},
"regional_processing_uplift_multiplier_eu": {"type": "number"},
"regional_processing_uplift_multiplier_us": {"type": "number"},
"input_cost_per_pixel": {"type": "number"},

View file

@ -2729,7 +2729,7 @@
"count": 1
},
"no-nested-ternary": {
"count": 6
"count": 4
}
},
"src/components/chat/MCPConnectPicker.tsx": {

View file

@ -41,7 +41,12 @@ interface AttachmentRowActionsProps {
onDeleteClick: (attachmentId: string) => void;
}
const CONFIG_ATTACHMENT_HINT =
"Config attachments are defined in the config file and cannot be deleted from the dashboard.";
function AttachmentRowActions({ attachment, isAdmin, onDeleteClick }: AttachmentRowActionsProps) {
const isConfigAttachment = attachment.definition_location === "config";
return (
<DropdownMenu>
<DropdownMenuTrigger
@ -65,6 +70,8 @@ function AttachmentRowActions({ attachment, isAdmin, onDeleteClick }: Attachment
<DropdownMenuItem
variant="destructive"
data-testid="attachment-action-delete"
disabled={isConfigAttachment}
title={isConfigAttachment ? CONFIG_ATTACHMENT_HINT : undefined}
onClick={() => onDeleteClick(attachment.attachment_id)}
>
<Trash2 />

View file

@ -145,4 +145,34 @@ describe("PolicyTable", () => {
await user.click(screen.getByRole("button", { name: /grouped/ }));
expect(defaultProps.onViewClick).toHaveBeenCalledWith("prod-id");
});
const sameNamedDbDraft: Partial<Policy> = {
policy_name: "config-policy",
policy_id: "db-draft-id",
version_status: "draft",
version_number: 2,
};
const configTwin: Partial<Policy> = {
policy_name: "config-policy",
policy_id: "config-policy",
version_status: "production",
definition_location: "config",
};
it("should render a config policy and a same-named DB draft as separate rows", () => {
const policies = [makePolicy(sameNamedDbDraft), makePolicy(configTwin)];
renderWithProviders(<PolicyTable {...defaultProps} policies={policies} />);
expect(screen.getAllByText("config-policy")).toHaveLength(2);
expect(screen.getByText("Config")).toBeInTheDocument();
});
it("should keep a same-named DB draft reachable next to a config policy", async () => {
const user = userEvent.setup();
const policies = [makePolicy(sameNamedDbDraft), makePolicy(configTwin)];
renderWithProviders(<PolicyTable {...defaultProps} policies={policies} />);
await user.click(screen.getByRole("button", { name: "config-policy" }));
expect(defaultProps.onViewClick).toHaveBeenCalledWith("db-draft-id");
await user.click(screen.getByTestId("policy-actions-db-draft-id"));
expect(await screen.findByTestId("policy-action-edit")).not.toHaveAttribute("data-disabled");
});
});

View file

@ -9,16 +9,21 @@ import { Policy } from "@/components/policies/types";
import { getPolicyTableColumns, PolicyRow } from "./PolicyTableColumns";
/** One row per policy name; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */
/** One row per DB policy name plus one row per config policy, so a config policy never hides same-named DB versions; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */
function groupPoliciesByName(policies: Policy[]): PolicyRow[] {
const names = Array.from(new Set(policies.map((policy) => policy.policy_name || "(unnamed)")));
return names.map((policyName) => {
const versions = policies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName);
const dbPolicies = policies.filter((policy) => policy.definition_location !== "config");
const names = Array.from(new Set(dbPolicies.map((policy) => policy.policy_name || "(unnamed)")));
const dbRows = names.map((policyName) => {
const versions = dbPolicies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName);
const primary =
versions.find((version) => version.version_status === "production") ??
[...versions].sort((a, b) => (b.version_number ?? 0) - (a.version_number ?? 0))[0];
return { policy_name: policyName, primaryPolicy: primary, versionCount: versions.length };
});
const configRows = policies
.filter((policy) => policy.definition_location === "config")
.map((policy) => ({ policy_name: policy.policy_name || "(unnamed)", primaryPolicy: policy, versionCount: 1 }));
return [...dbRows, ...configRows];
}
interface PolicyTableProps {
@ -67,7 +72,7 @@ const PolicyTable: React.FC<PolicyTableProps> = ({
<DataTable
data={rows}
columns={columns}
getRowId={(row) => row.policy_name}
getRowId={(row) => `${row.primaryPolicy.definition_location ?? "db"}:${row.policy_name}`}
sortingMode="client"
sorting={sorting}
onSortingChange={setSorting}

View file

@ -22,6 +22,9 @@ export interface PolicyRow {
versionCount: number;
}
const CONFIG_POLICY_HINT =
"Config policies are defined in the config file and cannot be edited or deleted from the dashboard.";
function GuardrailChips({ guardrails, tone }: { guardrails: string[]; tone: "success" | "error" }) {
if (guardrails.length === 0) {
return <span className="text-muted-foreground">-</span>;
@ -45,6 +48,8 @@ interface PolicyRowActionsProps {
}
function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActionsProps) {
const isConfigPolicy = policy.definition_location === "config";
return (
<DropdownMenu>
<DropdownMenuTrigger
@ -55,7 +60,12 @@ function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActio
<MoreHorizontal className="size-4" />
</DropdownMenuTrigger>
<DropdownMenuContent align="end" className="w-52">
<DropdownMenuItem data-testid="policy-action-edit" onClick={() => onEditClick(policy)}>
<DropdownMenuItem
data-testid="policy-action-edit"
disabled={isConfigPolicy}
title={isConfigPolicy ? CONFIG_POLICY_HINT : undefined}
onClick={() => onEditClick(policy)}
>
<Pencil />
Edit policy
</DropdownMenuItem>
@ -63,6 +73,8 @@ function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActio
<DropdownMenuItem
variant="destructive"
data-testid="policy-action-delete"
disabled={isConfigPolicy}
title={isConfigPolicy ? CONFIG_POLICY_HINT : undefined}
onClick={() => onDeleteClick(policy.policy_id, policy.policy_name || "Unnamed Policy")}
>
<Trash2 />
@ -93,18 +105,23 @@ export const getPolicyTableColumns = ({
header: ({ column }) => <DataTableSortHeader column={column} title="Name" />,
size: 220,
enableSorting: true,
cell: ({ row }) => (
<IdentityCell
title={row.original.policy_name}
titleClassName="max-w-60"
badge={
row.original.versionCount > 1 ? (
<StatusBadge tone="neutral" label={`${row.original.versionCount} versions`} />
) : undefined
}
onClick={() => onViewClick(row.original.primaryPolicy.policy_id)}
/>
),
cell: ({ row }) => {
const isConfigPolicy = row.original.primaryPolicy.definition_location === "config";
const versionBadge =
row.original.versionCount > 1 ? (
<StatusBadge tone="neutral" label={`${row.original.versionCount} versions`} />
) : undefined;
return (
<IdentityCell
title={row.original.policy_name}
titleClassName="max-w-60"
badge={
isConfigPolicy ? <StatusBadge tone="neutral" label="Config" tooltip={CONFIG_POLICY_HINT} /> : versionBadge
}
onClick={isConfigPolicy ? undefined : () => onViewClick(row.original.primaryPolicy.policy_id)}
/>
);
},
},
{
id: "description",

View file

@ -1,7 +1,13 @@
import { InfoCircleOutlined } from "@ant-design/icons";
import { Select as AntdSelect, Card, InputNumber, Radio, Space, Tooltip, Typography } from "antd";
import React from "react";
import { ClassifierType, ComplexityRouterConfigValue, DEFAULT_CLASSIFIER_TIMEOUT_MS } from "./ComplexityRouterConfig";
import {
ClassifierType,
ComplexityRouterConfigValue,
DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
DEFAULT_CLASSIFIER_TIMEOUT_MS,
} from "./ComplexityRouterConfig";
const { Text } = Typography;
@ -26,14 +32,23 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model;
const handleClassifierTypeChange = (classifierType: ClassifierType) => {
onChange({
const nextValue: ComplexityRouterConfigValue = {
...value,
classifier_type: classifierType,
classifier_llm_config:
classifierType === "llm"
? value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS }
: undefined,
});
classifier_context_window_size:
classifierType === "llm"
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
: undefined,
classifier_context_per_turn_chars:
classifierType === "llm"
? value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS
: undefined,
};
onChange(nextValue);
};
const handleClassifierModelChange = (model: string) => {
@ -56,6 +71,20 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
});
};
const handleClassifierContextWindowSizeChange = (windowSize: number | null) => {
onChange({
...value,
classifier_context_window_size: windowSize ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
});
};
const handleClassifierContextPerTurnCharsChange = (perTurnChars: number | null) => {
onChange({
...value,
classifier_context_per_turn_chars: perTurnChars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
});
};
return (
<>
<Radio.Group
@ -111,6 +140,36 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
response.
</Text>
</div>
<div>
<Text strong style={{ display: "block", marginBottom: 4 }}>
Context Window Size
</Text>
<InputNumber
value={value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE}
onChange={handleClassifierContextWindowSizeChange}
min={0}
style={{ width: "100%" }}
/>
<Text type="secondary" style={{ fontSize: 12 }}>
Number of prior user turns (tool output and harness reminders excluded) sent to the classifier as context,
so a referring follow-up like &quot;now do the same for the streaming path&quot; is classified against
what it refers to. Set to 0 to send only the current message.
</Text>
</div>
<div>
<Text strong style={{ display: "block", marginBottom: 4 }}>
Context Per-Turn Character Limit
</Text>
<InputNumber
value={value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS}
onChange={handleClassifierContextPerTurnCharsChange}
min={1}
style={{ width: "100%" }}
/>
<Text type="secondary" style={{ fontSize: 12 }}>
Prior turns longer than this are truncated.
</Text>
</div>
</div>
)}

View file

@ -99,11 +99,14 @@ describe("ComplexityRouterConfig", () => {
fireEvent.click(screen.getByText("Advanced: Classification Method"));
fireEvent.click(screen.getByText("LLM Classifier"));
expect(onChange).toHaveBeenCalledWith({
const expectedValue: ComplexityRouterConfigValue = {
...defaultValue,
classifier_type: "llm",
classifier_llm_config: { model: "", timeout_ms: 3000 },
});
classifier_context_window_size: 3,
classifier_context_per_turn_chars: 200,
};
expect(onChange).toHaveBeenCalledWith(expectedValue);
});
it("should show classifier fields and use the configured values when classifier_type is llm", () => {
@ -111,6 +114,8 @@ describe("ComplexityRouterConfig", () => {
...defaultValue,
classifier_type: "llm",
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 750 },
classifier_context_window_size: 5,
classifier_context_per_turn_chars: 400,
};
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={vi.fn()} />);
@ -119,6 +124,74 @@ describe("ComplexityRouterConfig", () => {
expect(screen.getByText("Classifier Model")).toBeInTheDocument();
expect(screen.getByText("Timeout (ms)")).toBeInTheDocument();
expect(screen.getByDisplayValue("750")).toBeInTheDocument();
expect(screen.getByText("Context Window Size")).toBeInTheDocument();
expect(screen.getByDisplayValue("5")).toBeInTheDocument();
expect(screen.getByText("Context Per-Turn Character Limit")).toBeInTheDocument();
expect(screen.getByDisplayValue("400")).toBeInTheDocument();
});
it("should default classifier context fields to 3 and 200 when llm is selected without explicit values", () => {
const llmValue: ComplexityRouterConfigValue = {
...defaultValue,
classifier_type: "llm",
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
};
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={vi.fn()} />);
fireEvent.click(screen.getByText("Advanced: Classification Method"));
const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement;
expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument();
const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement;
expect(within(perTurnCharsSection).getByDisplayValue("200")).toBeInTheDocument();
});
it("should hide classifier context fields when classifier_type is heuristic", () => {
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={defaultValue} onChange={vi.fn()} />);
fireEvent.click(screen.getByText("Advanced: Classification Method"));
expect(screen.queryByText("Context Window Size")).not.toBeInTheDocument();
expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument();
});
it("should call onChange with the updated classifier_context_window_size when edited", () => {
const onChange = vi.fn();
const llmValue: ComplexityRouterConfigValue = {
...defaultValue,
classifier_type: "llm",
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
};
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Classification Method"));
const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement;
const input = within(windowSizeSection).getByRole("spinbutton");
fireEvent.change(input, { target: { value: "7" } });
expect(onChange).toHaveBeenCalledWith({
...llmValue,
classifier_context_window_size: 7,
});
});
it("should call onChange with the updated classifier_context_per_turn_chars when edited", () => {
const onChange = vi.fn();
const llmValue: ComplexityRouterConfigValue = {
...defaultValue,
classifier_type: "llm",
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
};
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Classification Method"));
const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement;
const input = within(perTurnCharsSection).getByRole("spinbutton");
fireEvent.change(input, { target: { value: "500" } });
expect(onChange).toHaveBeenCalledWith({
...llmValue,
classifier_context_per_turn_chars: 500,
});
});
it("should render the custom technical keywords field", () => {

View file

@ -12,6 +12,8 @@ const { Text } = Typography;
export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000;
export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5;
export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3;
export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200;
export interface ComplexityTiers {
SIMPLE: string[];
@ -40,6 +42,8 @@ export interface ComplexityRouterConfigValue {
tiers: ComplexityTiers;
classifier_type: ClassifierType;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: number;
classifier_context_per_turn_chars?: number;
adaptive?: boolean;
adaptive_weights?: AdaptiveRouterWeights;
tier_distance_penalty?: number;

View file

@ -99,6 +99,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
tiers,
classifier_type: classifierType,
classifier_llm_config: classifierLlmConfig,
classifier_context_window_size: classifierContextWindowSize,
classifier_context_per_turn_chars: classifierContextPerTurnChars,
adaptive = false,
adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS,
tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY,
@ -142,6 +144,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
tiers,
classifierType,
classifierLlmConfig,
classifierContextWindowSize,
classifierContextPerTurnChars,
customTechnicalKeywords,
keywordTierRules,
semanticMatchingEnabled,

View file

@ -16,6 +16,8 @@ const baseParams: BuildComplexityRouterConfigParams = {
tiers,
classifierType: "heuristic",
classifierLlmConfig: undefined,
classifierContextWindowSize: undefined,
classifierContextPerTurnChars: undefined,
customTechnicalKeywords: [],
keywordTierRules: [],
semanticMatchingEnabled: false,
@ -75,6 +77,52 @@ describe("buildComplexityRouterConfig", () => {
expect(config.classifier_llm_config).toBeUndefined();
});
it("includes classifier_context_window_size and classifier_context_per_turn_chars only when classifier_type is llm", () => {
const params: BuildComplexityRouterConfigParams = {
...baseParams,
classifierType: "llm",
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
classifierContextWindowSize: 5,
classifierContextPerTurnChars: 300,
};
const config = buildComplexityRouterConfig(params);
expect(config.classifier_context_window_size).toBe(5);
expect(config.classifier_context_per_turn_chars).toBe(300);
});
it("omits classifier_context_window_size and classifier_context_per_turn_chars when classifier_type is heuristic even if values linger in state", () => {
const params: BuildComplexityRouterConfigParams = {
...baseParams,
classifierType: "heuristic",
classifierContextWindowSize: 5,
classifierContextPerTurnChars: 300,
};
const config = buildComplexityRouterConfig(params);
expect(config.classifier_context_window_size).toBeUndefined();
expect(config.classifier_context_per_turn_chars).toBeUndefined();
});
it("omits classifier_context_window_size and classifier_context_per_turn_chars when classifier_type is llm but neither was set, leaving the backend default", () => {
const config = buildComplexityRouterConfig({
...baseParams,
classifierType: "llm",
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
});
expect(config.classifier_context_window_size).toBeUndefined();
expect(config.classifier_context_per_turn_chars).toBeUndefined();
});
it("allows classifier_context_window_size of 0, distinct from unset, to send no prior-turn context", () => {
const params: BuildComplexityRouterConfigParams = {
...baseParams,
classifierType: "llm",
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
classifierContextWindowSize: 0,
};
const config = buildComplexityRouterConfig(params);
expect(config.classifier_context_window_size).toBe(0);
});
it("sends keyword_tier_rules with their per-tier targeting preserved (not flattened)", () => {
const params: BuildComplexityRouterConfigParams = {
...baseParams,

View file

@ -12,6 +12,8 @@ export interface BuildComplexityRouterConfigParams {
tiers: ComplexityTiers;
classifierType: ClassifierType;
classifierLlmConfig: ClassifierLLMConfig | undefined;
classifierContextWindowSize: number | undefined;
classifierContextPerTurnChars: number | undefined;
customTechnicalKeywords: string[];
keywordTierRules: KeywordTierRule[];
semanticMatchingEnabled: boolean;
@ -29,6 +31,8 @@ export interface ComplexityRouterConfigPayload {
tiers: ComplexityTiers;
classifier_type: ClassifierType;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: number;
classifier_context_per_turn_chars?: number;
custom_technical_keywords?: string[];
keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[];
semantic_keyword_matching?: boolean;
@ -69,6 +73,8 @@ export const buildComplexityRouterConfig = ({
tiers,
classifierType,
classifierLlmConfig,
classifierContextWindowSize,
classifierContextPerTurnChars,
customTechnicalKeywords,
keywordTierRules,
semanticMatchingEnabled,
@ -89,6 +95,14 @@ export const buildComplexityRouterConfig = ({
tiers,
classifier_type: classifierType,
...(classifierType === "llm" && classifierLlmConfig && { classifier_llm_config: classifierLlmConfig }),
...(classifierType === "llm" &&
classifierContextWindowSize !== undefined && {
classifier_context_window_size: classifierContextWindowSize,
}),
...(classifierType === "llm" &&
classifierContextPerTurnChars !== undefined && {
classifier_context_per_turn_chars: classifierContextPerTurnChars,
}),
...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }),
...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }),
escalation_keywords: cleanedEscalationKeywords,

View file

@ -1,6 +1,6 @@
import React from "react";
import { render, screen, fireEvent } from "@testing-library/react";
import { describe, it, expect, vi, afterEach } from "vitest";
import { render, screen, fireEvent, waitFor, act } from "@testing-library/react";
import { describe, it, expect, vi, afterEach, beforeEach } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import MCPAppsPanel from "./MCPAppsPanel";
import { fetchMCPServers, listMCPTools } from "../networking";
@ -86,3 +86,208 @@ describe("MCPAppsPanel logos", () => {
expect(screen.getByAltText("local_logo logo").getAttribute("src")).toBe("/litellm/ui/assets/logos/github.svg");
});
});
const connectServers = [
{
server_id: "s-reach",
server_name: "reachable_srv",
auth_type: "none",
connected_app_reachable: true,
},
{
server_id: "s-unreach",
server_name: "unreachable_srv",
auth_type: "none",
connected_app_reachable: false,
},
] as MCPServer[];
const renderConnectPanel = (connectMode: boolean, selectedServers: string[] = []) =>
render(
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
<MCPAppsPanel accessToken="tok" selectedServers={selectedServers} onChange={vi.fn()} connectMode={connectMode} />
</QueryClientProvider>,
);
describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("requests the connected-app view and hides unreachable servers in connect mode", async () => {
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
renderConnectPanel(true, ["reachable_srv", "unreachable_srv"]);
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true);
expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument();
expect(screen.getByText("Connected (1)")).toBeInTheDocument();
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
expect(toolCountFetchedIds).toContain("s-reach");
expect(toolCountFetchedIds).not.toContain("s-unreach");
});
it("blocks connecting an unsupported server from the detail view in connect mode", async () => {
const detailServers = [
...connectServers,
{
server_id: "s-unsup",
server_name: "unsupported_srv",
auth_type: "oauth2_token_exchange",
connected_app_reachable: true,
},
] as MCPServer[];
vi.mocked(fetchMCPServers).mockResolvedValue(detailServers);
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
renderConnectPanel(true);
fireEvent.click(await screen.findByText("unsupported_srv"));
expect(await screen.findByRole("heading", { name: "unsupported_srv" })).toBeInTheDocument();
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
expect(screen.getByText("Not supported on this connection")).toBeInTheDocument();
});
it("keeps the detail-view Connect action outside connect mode", async () => {
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
renderConnectPanel(false);
fireEvent.click(await screen.findByText("unreachable_srv"));
expect(await screen.findByRole("heading", { name: "unreachable_srv" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Connect" })).toBeInTheDocument();
});
it("ignores the flag and skips no server outside connect mode", async () => {
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
renderConnectPanel(false, ["reachable_srv", "unreachable_srv"]);
expect(await screen.findByText("unreachable_srv")).toBeInTheDocument();
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, false);
expect(screen.queryByText("Not available to connected apps")).not.toBeInTheDocument();
expect(screen.getByText("Connected (2)")).toBeInTheDocument();
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
expect(toolCountFetchedIds).toContain("s-unreach");
});
const revocable = (reachable: boolean) =>
[
{ server_id: "s-reach", server_name: "reachable_srv", auth_type: "none", connected_app_reachable: true },
{ server_id: "s-drop", server_name: "revoked_srv", auth_type: "none", connected_app_reachable: reachable },
] as MCPServer[];
const ConnectPanel = ({
token,
onChange,
client,
}: {
token: string;
onChange: (servers: string[]) => void;
client: QueryClient;
}) => (
<QueryClientProvider client={client}>
<MCPAppsPanel accessToken={token} selectedServers={[]} onChange={onChange} connectMode />
</QueryClientProvider>
);
const newClient = () => new QueryClient({ defaultOptions: { queries: { retry: false } } });
it("drops an open detail view when a refetch removes that server from the reachable set", async () => {
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
const client = newClient();
const { rerender } = render(<ConnectPanel token="tok" onChange={vi.fn()} client={client} />);
fireEvent.click(await screen.findByText("revoked_srv"));
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
rerender(<ConnectPanel token="tok-refreshed" onChange={vi.fn()} client={client} />);
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
expect(screen.getByText("reachable_srv")).toBeInTheDocument();
});
it("does not select a server whose Connect finishes after a refetch removed it", async () => {
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
const onChange = vi.fn();
const client = newClient();
const { rerender } = render(<ConnectPanel token="tok" onChange={onChange} client={client} />);
fireEvent.click(await screen.findByText("revoked_srv"));
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
let finishConnect: (result: { tools: never[] }) => void = () => {};
vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve)));
fireEvent.click(screen.getByRole("button", { name: "Connect" }));
rerender(<ConnectPanel token="tok-refreshed" onChange={onChange} client={client} />);
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
await act(async () => {
finishConnect({ tools: [] });
});
expect(onChange).not.toHaveBeenCalled();
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
expect(screen.getByText("Connected", { exact: false }).textContent).toBe("Connected");
});
it("does not select a server when Connect resolves in the same tick the refetch drops it", async () => {
let finishRefetch: (servers: MCPServer[]) => void = () => {};
vi.mocked(fetchMCPServers)
.mockResolvedValueOnce(revocable(true))
.mockImplementationOnce(() => new Promise((resolve) => (finishRefetch = resolve)));
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
const onChange = vi.fn();
const client = newClient();
const { rerender } = render(<ConnectPanel token="tok" onChange={onChange} client={client} />);
fireEvent.click(await screen.findByText("revoked_srv"));
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
let finishConnect: (result: { tools: never[] }) => void = () => {};
vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve)));
fireEvent.click(screen.getByRole("button", { name: "Connect" }));
rerender(<ConnectPanel token="tok-refreshed" onChange={onChange} client={client} />);
await act(async () => {
finishRefetch(revocable(false));
finishConnect({ tools: [] });
});
expect(onChange).not.toHaveBeenCalled();
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
});
it("does not let a superseded list load overwrite the current reachable set", async () => {
let finishStaleLoad: (servers: MCPServer[]) => void = () => {};
vi.mocked(fetchMCPServers)
.mockImplementationOnce(() => new Promise((resolve) => (finishStaleLoad = resolve)))
.mockResolvedValueOnce(revocable(false));
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
const client = newClient();
const { rerender } = render(<ConnectPanel token="tok" onChange={vi.fn()} client={client} />);
rerender(<ConnectPanel token="tok-refreshed" onChange={vi.fn()} client={client} />);
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
await act(async () => {
finishStaleLoad(revocable(true));
});
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
});
});

View file

@ -103,16 +103,17 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
const [query, setQuery] = useState("");
const [activeTab, setActiveTab] = useState<TabKey>("all");
const [togglingOn, setTogglingOn] = useState<Set<string>>(new Set());
const [detailServer, setDetailServer] = useState<MCPServer | null>(null);
const [detailServerId, setDetailServerId] = useState<string | null>(null);
const [toolCounts, setToolCounts] = useState<Record<string, number>>({});
const [loadingCounts, setLoadingCounts] = useState(false);
const [oauthConnected, setOauthConnected] = useState<Set<string>>(new Set());
const [oauthChecking, setOauthChecking] = useState<Set<string>>(new Set());
const serversRef = useRef<MCPServer[]>([]);
useEffect(() => {
serversRef.current = servers;
}, [servers]);
const commitServers = useCallback((next: MCPServer[]) => {
serversRef.current = next;
setServers(next);
}, []);
const selectedServersRef = useRef<string[]>(selectedServers);
useEffect(() => {
selectedServersRef.current = selectedServers;
@ -124,13 +125,30 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
const nameOf = (s: MCPServer) => s.server_name ?? s.alias ?? s.server_id;
const fetchLoadCancelledRef = useRef(false);
const detailServer = servers.find((s) => s.server_id === detailServerId);
const connectUnavailabilityLabel = useCallback(
(s: MCPServer): string | null => {
if (!connectMode) return null;
if (isUnsupportedOnGatewayConnect(s.auth_type)) return "Not supported on this connection";
return null;
},
[connectMode],
);
const connectableNow = useCallback(
(serverId: string): MCPServer | undefined => {
const current = serversRef.current.find((s) => s.server_id === serverId);
return current !== undefined && connectUnavailabilityLabel(current) === null ? current : undefined;
},
[connectUnavailabilityLabel],
);
const fetchToolCount = useCallback(
async (server: MCPServer) => {
async (server: MCPServer, isCurrentLoad: () => boolean) => {
try {
const toolsData = await listMCPTools(accessToken, server.server_id);
if (fetchLoadCancelledRef.current) return;
if (!isCurrentLoad()) return;
const tools: MCPTool[] = Array.isArray(toolsData?.tools) ? toolsData.tools : [];
setToolCounts((prev) => ({ ...prev, [nameOf(server)]: tools.length }));
} catch {
@ -141,17 +159,17 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
);
const checkOauthCredential = useCallback(
async (server: MCPServer) => {
async (server: MCPServer, isCurrentLoad: () => boolean) => {
try {
const status = await getMCPOAuthUserCredentialStatus(accessToken, server.server_id);
if (fetchLoadCancelledRef.current) return;
if (!isCurrentLoad()) return;
if (status.has_credential && !status.is_expired) {
setOauthConnected((prev) => new Set(prev).add(server.server_id));
}
} catch {
// ignore
} finally {
if (!fetchLoadCancelledRef.current) {
if (isCurrentLoad()) {
setOauthChecking((prev) => {
const next = new Set(prev);
next.delete(server.server_id);
@ -164,70 +182,77 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
);
useEffect(() => {
fetchLoadCancelledRef.current = false;
let current = true;
const isCurrentLoad = () => current;
fetchMCPServers(accessToken)
fetchMCPServers(accessToken, undefined, connectMode)
.then(async (serverData) => {
if (fetchLoadCancelledRef.current) return;
if (!isCurrentLoad()) return;
const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? [];
const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
setServers(list);
const reachable = connectMode ? list.filter((s) => s.connected_app_reachable !== false) : list;
const oauthServers = reachable.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
commitServers(reachable);
setOauthChecking(new Set(oauthServers.map((s) => s.server_id)));
setLoading(false);
oauthServers.forEach((s) => checkOauthCredential(s));
oauthServers.forEach((s) => checkOauthCredential(s, isCurrentLoad));
setLoadingCounts(true);
const chunks = Array.from({ length: Math.ceil(list.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) =>
list.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY),
const chunks = Array.from({ length: Math.ceil(reachable.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) =>
reachable.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY),
);
for (const chunk of chunks) {
if (fetchLoadCancelledRef.current) return;
await Promise.allSettled(chunk.map((s) => fetchToolCount(s)));
if (!isCurrentLoad()) return;
await Promise.allSettled(chunk.map((s) => fetchToolCount(s, isCurrentLoad)));
}
if (!fetchLoadCancelledRef.current) setLoadingCounts(false);
if (isCurrentLoad()) setLoadingCounts(false);
})
.catch(() => {
if (!fetchLoadCancelledRef.current) {
setServers([]);
if (isCurrentLoad()) {
commitServers([]);
setLoading(false);
}
});
return () => {
fetchLoadCancelledRef.current = true;
current = false;
};
}, [accessToken, fetchToolCount, checkOauthCredential]);
}, [accessToken, connectMode, commitServers, fetchToolCount, checkOauthCredential]);
useEffect(() => {
if (oauthConnected.size === 0) return;
const namesToAdd = serversRef.current
.filter((s) => oauthConnected.has(s.server_id) && !selectedServersRef.current.includes(nameOf(s)))
.filter(
(s) =>
oauthConnected.has(s.server_id) &&
!selectedServersRef.current.includes(nameOf(s)) &&
connectUnavailabilityLabel(s) === null,
)
.map(nameOf);
if (namesToAdd.length > 0) {
onChangeRef.current([...selectedServersRef.current, ...namesToAdd]);
}
}, [oauthConnected]);
}, [oauthConnected, connectUnavailabilityLabel]);
const handleToggle = async (serverName: string, checked: boolean, serverId?: string) => {
const handleToggle = async (server: MCPServer, checked: boolean) => {
const serverName = nameOf(server);
if (!checked) {
onChange(selectedServers.filter((s) => s !== serverName));
if (serverId) {
setOauthConnected((prev) => {
const next = new Set(prev);
next.delete(serverId);
return next;
});
}
setOauthConnected((prev) => {
const next = new Set(prev);
next.delete(server.server_id);
return next;
});
return;
}
if (connectableNow(server.server_id) === undefined) return;
setTogglingOn((prev) => new Set(prev).add(serverName));
try {
const idToFetch = serverId ?? serverName;
const result = await listMCPTools(accessToken, idToFetch);
const result = await listMCPTools(accessToken, server.server_id);
if (result?.error) {
MessageManager.warning(`Could not load tools for ${serverName}`);
return;
}
if (connectableNow(server.server_id) === undefined) return;
if (!selectedServersRef.current.includes(serverName)) {
onChange([...selectedServersRef.current, serverName]);
}
@ -243,11 +268,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
};
const renderConnectionIndicator = (server: MCPServer) => {
if (connectMode && isUnsupportedOnGatewayConnect(server.auth_type)) {
const unavailabilityLabel = connectUnavailabilityLabel(server);
if (unavailabilityLabel !== null) {
return (
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">
Not supported on this connection
</span>
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">{unavailabilityLabel}</span>
);
}
if (server.auth_type === AUTH_TYPE.OAUTH2) {
@ -285,11 +309,23 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
!query.trim() ||
name.toLowerCase().includes(query.toLowerCase()) ||
(s.description ?? "").toLowerCase().includes(query.toLowerCase());
const matchesTab = activeTab === "all" || selectedServers.includes(name);
const matchesTab =
activeTab === "all" || (selectedServers.includes(name) && connectUnavailabilityLabel(s) === null);
return matchesQuery && matchesTab;
});
const connectedCount = servers.filter((s) => selectedServers.includes(nameOf(s))).length;
const connectedCount = servers.filter(
(s) => selectedServers.includes(nameOf(s)) && connectUnavailabilityLabel(s) === null,
).length;
const emptyStateText = () => {
if (servers.length === 0) {
return connectMode
? "No MCP servers are available to this connection yet. Ask an admin to grant your user or team access."
: "No MCP servers configured. Add servers in Tools -> MCP Servers.";
}
return activeTab === "connected" ? "No servers connected yet." : "No servers match your search.";
};
const totalTools = Object.values(toolCounts).reduce((sum, n) => sum + n, 0);
if (detailServer) {
@ -298,12 +334,65 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
const isTogglingOn = togglingOn.has(name);
const color = getAvatarColor(name);
const renderDetailAction = () => {
const unavailabilityLabel = connectUnavailabilityLabel(detailServer);
if (unavailabilityLabel !== null) {
return <span className="text-[13px] text-muted-foreground py-2.5 shrink-0">{unavailabilityLabel}</span>;
}
if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) {
return (
<Button
variant={isConnected ? "outline" : "default"}
disabled={isTogglingOn}
onClick={() => handleToggle(detailServer, !isConnected)}
className="font-semibold h-[38px] min-w-[110px]"
>
{isTogglingOn && <Loader2 className="h-4 w-4 animate-spin mr-1.5" />}
{isConnected ? "Disconnect" : "Connect"}
</Button>
);
}
if (oauthConnected.has(detailServer.server_id)) {
return (
<Button
variant="destructive"
onClick={async () => {
try {
await deleteMCPOAuthUserCredential(accessToken, detailServer.server_id);
} catch (_) {
// Ignore
}
setOauthConnected((prev) => {
const n = new Set(prev);
n.delete(detailServer.server_id);
return n;
});
onChangeRef.current(selectedServersRef.current.filter((s) => s !== name));
}}
className="font-semibold h-[38px] min-w-[110px]"
>
Disconnect
</Button>
);
}
return (
<OAuth2ConnectButton
server={detailServer}
accessToken={accessToken}
onConnect={(id) => {
setOauthConnected((prev) => new Set(prev).add(id));
}}
variant="button"
/>
);
};
return (
<div className="w-full">
<Button
variant="ghost"
size="sm"
onClick={() => setDetailServer(null)}
onClick={() => setDetailServerId(null)}
className="-ml-3 mb-5 gap-1.5 text-muted-foreground hover:text-foreground"
>
<ArrowLeft className="h-3 w-3" />
@ -329,48 +418,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
<h2 className="m-0 mb-1 text-[22px] font-bold text-foreground">{name}</h2>
<p className="m-0 text-sm text-muted-foreground">{detailServer.description ?? "MCP server"}</p>
</div>
{detailServer.auth_type === AUTH_TYPE.OAUTH2 ? (
oauthConnected.has(detailServer.server_id) ? (
<Button
variant="destructive"
onClick={async () => {
try {
await deleteMCPOAuthUserCredential(accessToken, detailServer.server_id);
} catch (_) {
// Ignore
}
setOauthConnected((prev) => {
const n = new Set(prev);
n.delete(detailServer.server_id);
return n;
});
onChangeRef.current(selectedServersRef.current.filter((s) => s !== name));
}}
className="font-semibold h-[38px] min-w-[110px]"
>
Disconnect
</Button>
) : (
<OAuth2ConnectButton
server={detailServer}
accessToken={accessToken}
onConnect={(id) => {
setOauthConnected((prev) => new Set(prev).add(id));
}}
variant="button"
/>
)
) : (
<Button
variant={isConnected ? "outline" : "default"}
disabled={isTogglingOn}
onClick={() => handleToggle(name, !isConnected, detailServer.server_id)}
className="font-semibold h-[38px] min-w-[110px]"
>
{isTogglingOn && <Loader2 className="h-4 w-4 animate-spin mr-1.5" />}
{isConnected ? "Disconnect" : "Connect"}
</Button>
)}
{renderDetailAction()}
</div>
<h3 className="m-0 mb-3 text-[15px] font-semibold text-foreground">Information</h3>
@ -494,13 +542,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
))}
</div>
) : filtered.length === 0 ? (
<div className="text-center text-muted-foreground text-[13px] py-12 px-3">
{servers.length === 0
? "No MCP servers configured. Add servers in Tools -> MCP Servers."
: activeTab === "connected"
? "No servers connected yet."
: "No servers match your search."}
</div>
<div className="text-center text-muted-foreground text-[13px] py-12 px-3">{emptyStateText()}</div>
) : (
<div className="grid grid-cols-2 border rounded-lg overflow-hidden">
{filtered.map((server, idx) => {
@ -508,16 +550,16 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
const color = getAvatarColor(name);
const isLeftCol = idx % 2 === 0;
const count = toolCounts[name];
const unsupported = !!connectMode && isUnsupportedOnGatewayConnect(server.auth_type);
const unavailable = connectUnavailabilityLabel(server) !== null;
return (
<div
key={server.server_id}
onClick={() => setDetailServer(server)}
onClick={() => setDetailServerId(server.server_id)}
className={`flex items-center gap-3 p-4 bg-card cursor-pointer transition-colors hover:bg-accent/30 min-w-0 ${
isLeftCol ? "border-r" : ""
} ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""} ${
unsupported ? "opacity-50" : ""
unavailable ? "opacity-50" : ""
}`}
>
{server.mcp_info?.logo_url ? (

View file

@ -86,3 +86,68 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => {
expect(result.match_threshold).toBe(0.72);
});
});
const STORED_LLM = {
tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] },
classifier_type: "llm",
classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000 },
classifier_context_window_size: 5,
classifier_context_per_turn_chars: 300,
};
describe("buildUpdatedComplexityRouterConfig classifier context window", () => {
it("round-trips an untouched edit without changing the classifier context values", () => {
const formValue = {
tiers: STORED_LLM.tiers,
classifier_type: "llm" as const,
classifier_llm_config: STORED_LLM.classifier_llm_config,
classifier_context_window_size: 5,
classifier_context_per_turn_chars: 300,
};
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBe(5);
expect(result.classifier_context_per_turn_chars).toBe(300);
});
it("persists an edited classifier context window size and per-turn char limit", () => {
const formValue = {
tiers: STORED_LLM.tiers,
classifier_type: "llm" as const,
classifier_llm_config: STORED_LLM.classifier_llm_config,
classifier_context_window_size: 10,
classifier_context_per_turn_chars: 500,
};
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBe(10);
expect(result.classifier_context_per_turn_chars).toBe(500);
});
it("omits classifier context fields when classifier_type is heuristic even if values linger in state", () => {
const formValue = {
tiers: STORED_LLM.tiers,
classifier_type: "heuristic" as const,
classifier_context_window_size: 5,
classifier_context_per_turn_chars: 300,
};
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBeUndefined();
expect(result.classifier_context_per_turn_chars).toBeUndefined();
});
it("does not resurrect a stale stored classifier_context_window_size once the form's own value is unset", () => {
// classifier_context_window_size is a MANAGED key: the form's value must win over whatever
// is still sitting in the stored config, never fall back to it through preservedConfig.
const formValue = {
tiers: STORED_LLM.tiers,
classifier_type: "llm" as const,
classifier_llm_config: STORED_LLM.classifier_llm_config,
};
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBeUndefined();
expect(result.classifier_context_per_turn_chars).toBeUndefined();
});
});

View file

@ -1,7 +1,7 @@
import userEvent from "@testing-library/user-event";
import { describe, expect, it, vi } from "vitest";
import { renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils";
import NotificationsManager from "@/components/molecules/notifications_manager";
import EditAutoRouterModal from "./edit_auto_router_modal";
@ -119,3 +119,68 @@ describe("EditAutoRouterModal keyword matching", () => {
expect(modelPatchUpdateCall).not.toHaveBeenCalled();
});
});
describe("EditAutoRouterModal classifier context window", () => {
beforeEach(() => {
modelPatchUpdateCall.mockClear();
});
const STORED_LLM_CONFIG = {
tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o-mini"], COMPLEX: ["gpt-4o-mini"], REASONING: ["gpt-4o-mini"] },
classifier_type: "llm",
classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000 },
classifier_context_window_size: 5,
classifier_context_per_turn_chars: 300,
};
const renderLlmModal = () =>
renderWithProviders(
<EditAutoRouterModal
isVisible
onCancel={vi.fn()}
onSuccess={vi.fn()}
modelData={{
...MODEL_DATA,
litellm_params: { ...MODEL_DATA.litellm_params, complexity_router_config: STORED_LLM_CONFIG },
}}
accessToken="token"
userRole="Admin"
/>,
);
// Hydration bugs are invisible to the payload-builder unit tests, which only exercise
// buildUpdatedComplexityRouterConfig with a form value the caller already assembled by hand.
// Only driving the real component through open, then save with nothing touched, catches a
// missing initializeForm hydration line.
it("shows the stored classifier context values and preserves them through an untouched open-and-save", async () => {
const user = userEvent.setup();
renderLlmModal();
await user.click(await screen.findByText("Advanced: Classification Method"));
await screen.findByText("Context Window Size");
expect(screen.getByDisplayValue("5")).toBeInTheDocument();
expect(screen.getByDisplayValue("300")).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
const config = savedConfig();
expect(config.classifier_context_window_size).toBe(5);
expect(config.classifier_context_per_turn_chars).toBe(300);
});
it("persists an edited classifier context window size", async () => {
const user = userEvent.setup();
renderLlmModal();
await user.click(await screen.findByText("Advanced: Classification Method"));
const windowSizeSection = (await screen.findByText("Context Window Size")).closest("div") as HTMLElement;
const input = within(windowSizeSection).getByRole("spinbutton");
fireEvent.change(input, { target: { value: "8" } });
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().classifier_context_window_size).toBe(8);
});
});

View file

@ -33,6 +33,8 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
"tiers",
"classifier_type",
"classifier_llm_config",
"classifier_context_window_size",
"classifier_context_per_turn_chars",
"adaptive",
"adaptive_weights",
"tier_distance_penalty",
@ -86,6 +88,14 @@ export const buildUpdatedComplexityRouterConfig = (
tiers: value.tiers,
classifier_type: value.classifier_type,
...(value.classifier_type === "llm" ? { classifier_llm_config: value.classifier_llm_config } : {}),
...(value.classifier_type === "llm" &&
value.classifier_context_window_size !== undefined && {
classifier_context_window_size: value.classifier_context_window_size,
}),
...(value.classifier_type === "llm" &&
value.classifier_context_per_turn_chars !== undefined && {
classifier_context_per_turn_chars: value.classifier_context_per_turn_chars,
}),
...(customTechnicalKeywords &&
customTechnicalKeywords.length > 0 && {
custom_technical_keywords: customTechnicalKeywords,
@ -182,7 +192,7 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
parsedConfig = JSON.parse(parsedConfig);
}
setComplexityRouterConfig({
const hydratedComplexityRouterConfig: ComplexityRouterConfigValue = {
tiers: {
SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE),
MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM),
@ -191,12 +201,21 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
},
classifier_type: parsedConfig.classifier_type || "heuristic",
classifier_llm_config: parsedConfig.classifier_llm_config,
classifier_context_window_size:
typeof parsedConfig.classifier_context_window_size === "number"
? parsedConfig.classifier_context_window_size
: undefined,
classifier_context_per_turn_chars:
typeof parsedConfig.classifier_context_per_turn_chars === "number"
? parsedConfig.classifier_context_per_turn_chars
: undefined,
adaptive: parsedConfig.adaptive || false,
adaptive_weights: parsedConfig.adaptive_weights,
tier_distance_penalty: parsedConfig.tier_distance_penalty,
adaptive_eligible: parsedConfig.adaptive_eligible || "all",
return_raw_model_name: parsedConfig.return_raw_model_name || false,
});
};
setComplexityRouterConfig(hydratedComplexityRouterConfig);
setCustomTechnicalKeywords(
Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [],
);

View file

@ -436,6 +436,7 @@ export interface MCPServer {
byok_description?: string[] | null;
byok_api_key_help_url?: string | null;
has_user_credential?: boolean | null;
connected_app_reachable?: boolean | null;
/** GitHub / source repository URL */
source_url?: string | null;

View file

@ -4765,9 +4765,12 @@ export const fetchDiscoverableMCPServers = async (accessToken: string) => {
}
};
export const fetchMCPServers = async (accessToken: string, teamId?: string | null) => {
export const fetchMCPServers = async (accessToken: string, teamId?: string | null, connectedAppView?: boolean) => {
try {
return await apiClient.get(`/v1/mcp/server`, { accessToken, query: { team_id: teamId || undefined } });
return await apiClient.get(`/v1/mcp/server`, {
accessToken,
query: { team_id: teamId || undefined, connected_app_view: connectedAppView || undefined },
});
} catch (error) {
console.error("Failed to fetch MCP servers:", error);
throw error;

View file

@ -14,6 +14,7 @@ export interface Policy {
updated_at?: string;
created_by?: string;
updated_by?: string;
definition_location?: "db" | "config";
}
export interface PolicyCondition {
@ -47,6 +48,7 @@ export interface PolicyAttachment {
updated_at?: string;
created_by?: string;
updated_by?: string;
definition_location?: "db" | "config";
}
export interface PolicyCreateRequest {

View file

@ -9379,7 +9379,10 @@ export interface paths {
};
/**
* List Policy Attachments
* @description List all policy attachments from the database.
* @description List all policy attachments from the database and config.yaml.
*
* Config-defined attachments are returned with definition_location "config" and a
* synthetic attachment_id ("config-<index>").
*
* Example Request:
* ```bash
@ -9487,7 +9490,10 @@ export interface paths {
};
/**
* List Policies
* @description List all policies from the database. Optionally filter by version_status.
* @description List all policies from the database and config.yaml. Optionally filter by version_status.
*
* Config-defined policies are returned with definition_location "config" and are treated
* as production versions. On a name conflict with a DB policy, only the DB policy is returned.
*
* Query params:
* - version_status: Optional. One of "draft", "published", "production".
@ -25845,6 +25851,16 @@ export interface components {
cache_creation_input_token_cost_above_1hr?: number | null;
/** Cache Creation Input Token Cost Above 200K Tokens */
cache_creation_input_token_cost_above_200k_tokens?: number | null;
/** Cache Creation Input Token Cost Above 272K Tokens */
cache_creation_input_token_cost_above_272k_tokens?: number | null;
/** Cache Creation Input Token Cost Above 272K Tokens Flex */
cache_creation_input_token_cost_above_272k_tokens_flex?: number | null;
/** Cache Creation Input Token Cost Above 272K Tokens Priority */
cache_creation_input_token_cost_above_272k_tokens_priority?: number | null;
/** Cache Creation Input Token Cost Flex */
cache_creation_input_token_cost_flex?: number | null;
/** Cache Creation Input Token Cost Priority */
cache_creation_input_token_cost_priority?: number | null;
/** Cache Read Input Audio Token Cost */
cache_read_input_audio_token_cost?: number | null;
/** Cache Read Input Token Cost */
@ -25855,6 +25871,8 @@ export interface components {
cache_read_input_token_cost_above_200k_tokens_priority?: number | null;
/** Cache Read Input Token Cost Above 272K Tokens */
cache_read_input_token_cost_above_272k_tokens?: number | null;
/** Cache Read Input Token Cost Above 272K Tokens Flex */
cache_read_input_token_cost_above_272k_tokens_flex?: number | null;
/** Cache Read Input Token Cost Above 272K Tokens Priority */
cache_read_input_token_cost_above_272k_tokens_priority?: number | null;
/** Cache Read Input Token Cost Above 512K Tokens */
@ -25913,6 +25931,8 @@ export interface components {
input_cost_per_token_above_200k_tokens_priority?: number | null;
/** Input Cost Per Token Above 272K Tokens */
input_cost_per_token_above_272k_tokens?: number | null;
/** Input Cost Per Token Above 272K Tokens Flex */
input_cost_per_token_above_272k_tokens_flex?: number | null;
/** Input Cost Per Token Above 272K Tokens Priority */
input_cost_per_token_above_272k_tokens_priority?: number | null;
/** Input Cost Per Token Above 512K Tokens */
@ -26004,6 +26024,8 @@ export interface components {
output_cost_per_token_above_200k_tokens_priority?: number | null;
/** Output Cost Per Token Above 272K Tokens */
output_cost_per_token_above_272k_tokens?: number | null;
/** Output Cost Per Token Above 272K Tokens Flex */
output_cost_per_token_above_272k_tokens_flex?: number | null;
/** Output Cost Per Token Above 272K Tokens Priority */
output_cost_per_token_above_272k_tokens_priority?: number | null;
/** Output Cost Per Token Above 512K Tokens */
@ -29376,6 +29398,13 @@ export interface components {
* @description Who created the attachment.
*/
created_by?: string | null;
/**
* Definition Location
* @description Where this attachment is defined: 'db' (database) or 'config' (config.yaml).
* @default db
* @enum {string}
*/
definition_location: "db" | "config";
/**
* Keys
* @description Key patterns.
@ -29507,6 +29536,13 @@ export interface components {
* @description Who created the policy.
*/
created_by?: string | null;
/**
* Definition Location
* @description Where this policy is defined: 'db' (database) or 'config' (config.yaml).
* @default db
* @enum {string}
*/
definition_location: "db" | "config";
/**
* Description
* @description Policy description.
@ -33135,6 +33171,17 @@ export interface components {
*/
model?: string | null;
};
/** UsageChartPoint */
UsageChartPoint: {
/** Blocked */
blocked: number;
/** Date */
date: string;
/** Passed */
passed: number;
/** Score */
score?: number | null;
};
/** UsageDetailResponse */
UsageDetailResponse: {
/** Avglatency */
@ -33957,6 +34004,16 @@ export interface components {
cache_creation_input_token_cost_above_1hr?: number | null;
/** Cache Creation Input Token Cost Above 200K Tokens */
cache_creation_input_token_cost_above_200k_tokens?: number | null;
/** Cache Creation Input Token Cost Above 272K Tokens */
cache_creation_input_token_cost_above_272k_tokens?: number | null;
/** Cache Creation Input Token Cost Above 272K Tokens Flex */
cache_creation_input_token_cost_above_272k_tokens_flex?: number | null;
/** Cache Creation Input Token Cost Above 272K Tokens Priority */
cache_creation_input_token_cost_above_272k_tokens_priority?: number | null;
/** Cache Creation Input Token Cost Flex */
cache_creation_input_token_cost_flex?: number | null;
/** Cache Creation Input Token Cost Priority */
cache_creation_input_token_cost_priority?: number | null;
/** Cache Read Input Audio Token Cost */
cache_read_input_audio_token_cost?: number | null;
/** Cache Read Input Token Cost */
@ -33967,6 +34024,8 @@ export interface components {
cache_read_input_token_cost_above_200k_tokens_priority?: number | null;
/** Cache Read Input Token Cost Above 272K Tokens */
cache_read_input_token_cost_above_272k_tokens?: number | null;
/** Cache Read Input Token Cost Above 272K Tokens Flex */
cache_read_input_token_cost_above_272k_tokens_flex?: number | null;
/** Cache Read Input Token Cost Above 272K Tokens Priority */
cache_read_input_token_cost_above_272k_tokens_priority?: number | null;
/** Cache Read Input Token Cost Above 512K Tokens */
@ -34025,6 +34084,8 @@ export interface components {
input_cost_per_token_above_200k_tokens_priority?: number | null;
/** Input Cost Per Token Above 272K Tokens */
input_cost_per_token_above_272k_tokens?: number | null;
/** Input Cost Per Token Above 272K Tokens Flex */
input_cost_per_token_above_272k_tokens_flex?: number | null;
/** Input Cost Per Token Above 272K Tokens Priority */
input_cost_per_token_above_272k_tokens_priority?: number | null;
/** Input Cost Per Token Above 512K Tokens */
@ -34116,6 +34177,8 @@ export interface components {
output_cost_per_token_above_200k_tokens_priority?: number | null;
/** Output Cost Per Token Above 272K Tokens */
output_cost_per_token_above_272k_tokens?: number | null;
/** Output Cost Per Token Above 272K Tokens Flex */
output_cost_per_token_above_272k_tokens_flex?: number | null;
/** Output Cost Per Token Above 272K Tokens Priority */
output_cost_per_token_above_272k_tokens_priority?: number | null;
/** Output Cost Per Token Above 512K Tokens */