diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index e8252d87572..07bd957b5a3 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -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, diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 5b953035cfd..7fa78f39460 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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() diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 85ed0665ebf..fbc06b76c72 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 46220e1c29f..07f04136bd0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 23b26bd8e89..e428d20f99d 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 5d8ac8d678f..dc796837277 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ae1095da336..89c83459524 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 9b51513f4ac..24d37f81787 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -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") diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ec07d33f24d..48effdb0f6e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py index 1b37b884987..d1d28574988 100644 --- a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py +++ b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py @@ -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] diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index afd396adc4c..698df147247 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -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 diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 96f6ee89d56..12da0a26708 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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 \"\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-\").\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer \"\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 \"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer \"\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 \"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer \"\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": [ { diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 64cc13a5543..dcc4dc36d83 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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: diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 8ef509810ba..fe9ad3bef6d 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -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)}") diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index a879f6b6f7e..cff1378c676 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -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-"). 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}") diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index e1afbf2f5f2..01b88836387 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -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_ 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)}") diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 357a7ecefe6..dab666ff0d9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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: diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index 2c7e8d5afc9..b4096cd2044 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -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): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c77183a61b8..7db12e15be0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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): diff --git a/litellm/utils.py b/litellm/utils.py index 944bb61d5e7..ad661fd3d16 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 ), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 314a076dc20..346f613ea3e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 988f56655fc..882f514b199 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -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, diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 17aee22560c..180639b53e3 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -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`. diff --git a/tests/e2e/coverage_registry/README.md b/tests/e2e/coverage_registry/README.md index aef4c16c89a..5627c88dee4 100644 --- a/tests/e2e/coverage_registry/README.md +++ b/tests/e2e/coverage_registry/README.md @@ -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 ``` diff --git a/tests/e2e/coverage_registry/collector.py b/tests/e2e/coverage_registry/collector.py index 50ef23bcb1a..e20f7884f55 100644 --- a/tests/e2e/coverage_registry/collector.py +++ b/tests/e2e/coverage_registry/collector.py @@ -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, diff --git a/tests/e2e/coverage_registry/test_collector.py b/tests/e2e/coverage_registry/test_collector.py index 079ee215866..a85190cc3ba 100644 --- a/tests/e2e/coverage_registry/test_collector.py +++ b/tests/e2e/coverage_registry/test_collector.py @@ -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] diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index f1e6539439a..434a9bc3809 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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.""" diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py new file mode 100644 index 00000000000..7e997870852 --- /dev/null +++ b/tests/test_litellm/integrations/test_s3.py @@ -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 diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index f0a33f2ebfc..3977daae92f 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -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 diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index d282e656ce8..fbdf9b64bc0 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -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) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 641c552d144..e8488a72d00 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 4dd1c663de0..bea979aec64 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index c0affdf46b3..3a84428add1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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}" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 8a8dea0ba28..db5b64ef131 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 1e4349c3143..c4b3c7f5f67 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index cec62f79e33..329b6d5c45d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index 662ef585c6b..6e3ac014840 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py index 52120207f76..cd4cba51908 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 53dda8f6648..bf119c4fb2f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -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() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index a2f7476abd1..470179a0429 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -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, diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 1ae0b4d3d48..cc231e383a3 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -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() == () diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py new file mode 100644 index 00000000000..1ca830dc1e6 --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py @@ -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" diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index dd20021d0e1..ebebfde5cd3 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -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 diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 1a2dcd0fcb7..3b87246ebdb 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -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 diff --git a/tests/test_litellm/test_gpt_5_6_model_metadata.py b/tests/test_litellm/test_gpt_5_6_model_metadata.py deleted file mode 100644 index 5a7b621d521..00000000000 --- a/tests/test_litellm/test_gpt_5_6_model_metadata.py +++ /dev/null @@ -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" - ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3b2bc7e647d..b23d3333ea7 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -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"}, diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 53a10605c53..7ef729818c3 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2729,7 +2729,7 @@ "count": 1 }, "no-nested-ternary": { - "count": 6 + "count": 4 } }, "src/components/chat/MCPConnectPicker.tsx": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx index 9e0f8d6715d..ded9e3a1e6d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx @@ -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 ( onDeleteClick(attachment.attachment_id)} > diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx index 06c939aa151..9be1bd60ec8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx @@ -145,4 +145,34 @@ describe("PolicyTable", () => { await user.click(screen.getByRole("button", { name: /grouped/ })); expect(defaultProps.onViewClick).toHaveBeenCalledWith("prod-id"); }); + + const sameNamedDbDraft: Partial = { + policy_name: "config-policy", + policy_id: "db-draft-id", + version_status: "draft", + version_number: 2, + }; + const configTwin: Partial = { + 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(); + 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(); + 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"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx index d6e841c2119..3405ac6b6bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx @@ -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 = ({ row.policy_name} + getRowId={(row) => `${row.primaryPolicy.definition_location ?? "db"}:${row.policy_name}`} sortingMode="client" sorting={sorting} onSortingChange={setSorting} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx index dd036d83283..488de728bad 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx @@ -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 -; @@ -45,6 +48,8 @@ interface PolicyRowActionsProps { } function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActionsProps) { + const isConfigPolicy = policy.definition_location === "config"; + return ( - onEditClick(policy)}> + onEditClick(policy)} + > Edit policy @@ -63,6 +73,8 @@ function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActio onDeleteClick(policy.policy_id, policy.policy_name || "Unnamed Policy")} > @@ -93,18 +105,23 @@ export const getPolicyTableColumns = ({ header: ({ column }) => , size: 220, enableSorting: true, - cell: ({ row }) => ( - 1 ? ( - - ) : undefined - } - onClick={() => onViewClick(row.original.primaryPolicy.policy_id)} - /> - ), + cell: ({ row }) => { + const isConfigPolicy = row.original.primaryPolicy.definition_location === "config"; + const versionBadge = + row.original.versionCount > 1 ? ( + + ) : undefined; + return ( + : versionBadge + } + onClick={isConfigPolicy ? undefined : () => onViewClick(row.original.primaryPolicy.policy_id)} + /> + ); + }, }, { id: "description", diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 92df8029edc..75acf87a7d8 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -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 = ({ 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 = ({ }); }; + 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 ( <> = ({ response. +
+ + Context Window Size + + + + Number of prior user turns (tool output and harness reminders excluded) sent to the classifier as context, + so a referring follow-up like "now do the same for the streaming path" is classified against + what it refers to. Set to 0 to send only the current message. + +
+
+ + Context Per-Turn Character Limit + + + + Prior turns longer than this are truncated. + +
)} diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 6b3c1961468..a1e992bd46b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -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(); @@ -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(); + + 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(); + 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(); + 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(); + 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", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 1f2edf697a9..de32d15d5a7 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index cd3294b347f..fb7eabd7110 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -99,6 +99,8 @@ const AddAutoRouterTab: React.FC = ({ 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 = ({ tiers, classifierType, classifierLlmConfig, + classifierContextWindowSize, + classifierContextPerTurnChars, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index b5973bf7101..e269a3c9028 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 02be9b280aa..04a9b0f9bd4 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx index 656ef157363..5fc3e75195b 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx @@ -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( + + + , + ); + +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; + }) => ( + + + + ); + + 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(); + + fireEvent.click(await screen.findByText("revoked_srv")); + expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument(); + + rerender(); + + 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(); + + 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(); + 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(); + + 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(); + + 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(); + rerender(); + + expect(await screen.findByText("reachable_srv")).toBeInTheDocument(); + + await act(async () => { + finishStaleLoad(revocable(true)); + }); + + expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 329e4ec99fa..7dbfc058a77 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -103,16 +103,17 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, const [query, setQuery] = useState(""); const [activeTab, setActiveTab] = useState("all"); const [togglingOn, setTogglingOn] = useState>(new Set()); - const [detailServer, setDetailServer] = useState(null); + const [detailServerId, setDetailServerId] = useState(null); const [toolCounts, setToolCounts] = useState>({}); const [loadingCounts, setLoadingCounts] = useState(false); const [oauthConnected, setOauthConnected] = useState>(new Set()); const [oauthChecking, setOauthChecking] = useState>(new Set()); const serversRef = useRef([]); - useEffect(() => { - serversRef.current = servers; - }, [servers]); + const commitServers = useCallback((next: MCPServer[]) => { + serversRef.current = next; + setServers(next); + }, []); const selectedServersRef = useRef(selectedServers); useEffect(() => { selectedServersRef.current = selectedServers; @@ -124,13 +125,30 @@ const MCPAppsPanel: React.FC = ({ 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 = ({ 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 = ({ 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 = ({ accessToken, selectedServers, onChange, }; const renderConnectionIndicator = (server: MCPServer) => { - if (connectMode && isUnsupportedOnGatewayConnect(server.auth_type)) { + const unavailabilityLabel = connectUnavailabilityLabel(server); + if (unavailabilityLabel !== null) { return ( - - Not supported on this connection - + {unavailabilityLabel} ); } if (server.auth_type === AUTH_TYPE.OAUTH2) { @@ -285,11 +309,23 @@ const MCPAppsPanel: React.FC = ({ 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 = ({ accessToken, selectedServers, onChange, const isTogglingOn = togglingOn.has(name); const color = getAvatarColor(name); + const renderDetailAction = () => { + const unavailabilityLabel = connectUnavailabilityLabel(detailServer); + if (unavailabilityLabel !== null) { + return {unavailabilityLabel}; + } + if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) { + return ( + + ); + } + if (oauthConnected.has(detailServer.server_id)) { + return ( + + ); + } + return ( + { + setOauthConnected((prev) => new Set(prev).add(id)); + }} + variant="button" + /> + ); + }; + return (
- {detailServer.auth_type === AUTH_TYPE.OAUTH2 ? ( - oauthConnected.has(detailServer.server_id) ? ( - - ) : ( - { - setOauthConnected((prev) => new Set(prev).add(id)); - }} - variant="button" - /> - ) - ) : ( - - )} + {renderDetailAction()}

Information

@@ -494,13 +542,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, ))} ) : filtered.length === 0 ? ( -
- {servers.length === 0 - ? "No MCP servers configured. Add servers in Tools -> MCP Servers." - : activeTab === "connected" - ? "No servers connected yet." - : "No servers match your search."} -
+
{emptyStateText()}
) : (
{filtered.map((server, idx) => { @@ -508,16 +550,16 @@ const MCPAppsPanel: React.FC = ({ 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 (
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 ? ( diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index e39a7b6e444..eef5e1e4d06 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -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(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index e47bbecf8bd..504234e8977 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -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( + , + ); + + // 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); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 8fbca822165..2686f99307e 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -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 = ({ 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 = ({ }, 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 : [], ); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index de497d91afe..44ac25ea955 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 79b88a21299..10f5cccb512 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/policies/types.ts b/ui/litellm-dashboard/src/components/policies/types.ts index 887781ff943..6ac110e3c0a 100644 --- a/ui/litellm-dashboard/src/components/policies/types.ts +++ b/ui/litellm-dashboard/src/components/policies/types.ts @@ -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 { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a64fdae735a..c4c7819040d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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-"). * * 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 */