mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_otel_v2_admin_owned_destinations
This commit is contained in:
commit
a4bf719054
66 changed files with 4590 additions and 613 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -20402,6 +20402,16 @@
|
|||
"description": "Who created the attachment.",
|
||||
"title": "Created By"
|
||||
},
|
||||
"definition_location": {
|
||||
"default": "db",
|
||||
"description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).",
|
||||
"enum": [
|
||||
"db",
|
||||
"config"
|
||||
],
|
||||
"title": "Definition Location",
|
||||
"type": "string"
|
||||
},
|
||||
"keys": {
|
||||
"description": "Key patterns.",
|
||||
"items": {
|
||||
|
|
@ -20658,6 +20668,16 @@
|
|||
"description": "Who created the policy.",
|
||||
"title": "Created By"
|
||||
},
|
||||
"definition_location": {
|
||||
"default": "db",
|
||||
"description": "Where this policy is defined: 'db' (database) or 'config' (config.yaml).",
|
||||
"enum": [
|
||||
"db",
|
||||
"config"
|
||||
],
|
||||
"title": "Definition Location",
|
||||
"type": "string"
|
||||
},
|
||||
"description": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -21129,12 +21149,45 @@
|
|||
"title": "PolicyVersionStatusUpdateRequest",
|
||||
"type": "object"
|
||||
},
|
||||
"UsageChartPoint": {
|
||||
"properties": {
|
||||
"blocked": {
|
||||
"title": "Blocked",
|
||||
"type": "integer"
|
||||
},
|
||||
"date": {
|
||||
"title": "Date",
|
||||
"type": "string"
|
||||
},
|
||||
"passed": {
|
||||
"title": "Passed",
|
||||
"type": "integer"
|
||||
},
|
||||
"score": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Score"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"date",
|
||||
"passed",
|
||||
"blocked"
|
||||
],
|
||||
"title": "UsageChartPoint",
|
||||
"type": "object"
|
||||
},
|
||||
"UsageOverviewResponse": {
|
||||
"properties": {
|
||||
"chart": {
|
||||
"items": {
|
||||
"additionalProperties": true,
|
||||
"type": "object"
|
||||
"$ref": "#/components/schemas/UsageChartPoint"
|
||||
},
|
||||
"title": "Chart",
|
||||
"type": "array"
|
||||
|
|
@ -21243,6 +21296,13 @@
|
|||
},
|
||||
"ValidationError": {
|
||||
"properties": {
|
||||
"ctx": {
|
||||
"title": "Context",
|
||||
"type": "object"
|
||||
},
|
||||
"input": {
|
||||
"title": "Input"
|
||||
},
|
||||
"loc": {
|
||||
"items": {
|
||||
"anyOf": [
|
||||
|
|
@ -21420,7 +21480,7 @@
|
|||
},
|
||||
"/policies/attachments/list": {
|
||||
"get": {
|
||||
"description": "List all policy attachments from the database.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"attachments\": [\n {\n \"attachment_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"scope\": \"*\",\n \"teams\": [],\n \"keys\": [],\n \"models\": [],\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
|
||||
"description": "List all policy attachments from the database and config.yaml.\n\nConfig-defined attachments are returned with definition_location \"config\" and a\nsynthetic attachment_id (\"config-<index>\").\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"attachments\": [\n {\n \"attachment_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"scope\": \"*\",\n \"teams\": [],\n \"keys\": [],\n \"models\": [],\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
|
||||
"operationId": "list_policy_attachments_policies_attachments_list_get",
|
||||
"responses": {
|
||||
"200": {
|
||||
|
|
@ -21596,7 +21656,7 @@
|
|||
},
|
||||
"/policies/list": {
|
||||
"get": {
|
||||
"description": "List all policies from the database. Optionally filter by version_status.\n\nQuery params:\n- version_status: Optional. One of \"draft\", \"published\", \"production\".\n If omitted, all versions are returned.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"policies\": [\n {\n \"policy_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"version_number\": 1,\n \"version_status\": \"production\",\n \"inherit\": null,\n \"description\": \"Base guardrails for all requests\",\n \"guardrails_add\": [\"pii_masking\"],\n \"guardrails_remove\": [],\n \"condition\": null,\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
|
||||
"description": "List all policies from the database and config.yaml. Optionally filter by version_status.\n\nConfig-defined policies are returned with definition_location \"config\" and are treated\nas production versions. On a name conflict with a DB policy, only the DB policy is returned.\n\nQuery params:\n- version_status: Optional. One of \"draft\", \"published\", \"production\".\n If omitted, all versions are returned.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/list\" \\\n -H \"Authorization: Bearer <your_api_key>\"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```\n\nExample Response:\n```json\n{\n \"policies\": [\n {\n \"policy_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"version_number\": 1,\n \"version_status\": \"production\",\n \"inherit\": null,\n \"description\": \"Base guardrails for all requests\",\n \"guardrails_add\": [\"pii_masking\"],\n \"guardrails_remove\": [],\n \"condition\": null,\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```",
|
||||
"operationId": "list_policies_policies_list_get",
|
||||
"parameters": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)}")
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
|||
from litellm.types.proxy.policy_engine import (
|
||||
GuardrailPipeline,
|
||||
PipelineTestRequest,
|
||||
Policy,
|
||||
PolicyAttachment,
|
||||
PolicyAttachmentCreateRequest,
|
||||
PolicyAttachmentDBResponse,
|
||||
PolicyAttachmentListResponse,
|
||||
|
|
@ -33,6 +35,35 @@ from litellm.types.proxy.policy_engine import (
|
|||
router = APIRouter()
|
||||
|
||||
|
||||
def _config_policy_to_db_response(policy_name: str, policy: Policy) -> PolicyDBResponse:
|
||||
return PolicyDBResponse(
|
||||
policy_id=policy_name,
|
||||
policy_name=policy_name,
|
||||
version_number=1,
|
||||
version_status="production",
|
||||
inherit=policy.inherit,
|
||||
description=policy.description,
|
||||
guardrails_add=policy.guardrails.get_add(),
|
||||
guardrails_remove=policy.guardrails.get_remove(),
|
||||
condition=policy.condition.model_dump() if policy.condition else None,
|
||||
pipeline=policy.pipeline.model_dump() if policy.pipeline else None,
|
||||
definition_location="config",
|
||||
)
|
||||
|
||||
|
||||
def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment) -> PolicyAttachmentDBResponse:
|
||||
return PolicyAttachmentDBResponse(
|
||||
attachment_id=f"config-{index}",
|
||||
policy_name=attachment.policy,
|
||||
scope=attachment.scope,
|
||||
teams=attachment.teams or [],
|
||||
keys=attachment.keys or [],
|
||||
models=attachment.models or [],
|
||||
tags=attachment.tags or [],
|
||||
definition_location="config",
|
||||
)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Policy CRUD Endpoints
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -46,7 +77,13 @@ router = APIRouter()
|
|||
)
|
||||
async def list_policies(version_status: Optional[str] = None):
|
||||
"""
|
||||
List all policies from the database. Optionally filter by version_status.
|
||||
List all policies from the database and config.yaml. Optionally filter by version_status.
|
||||
|
||||
Config-defined policies are returned with definition_location "config" and are treated
|
||||
as production versions. On a name conflict with a production DB policy, only the DB policy
|
||||
is returned, mirroring runtime resolution where only production DB versions override config.
|
||||
A draft or published DB version does not hide the config policy, since the config version
|
||||
is still the one being enforced.
|
||||
|
||||
Query params:
|
||||
- version_status: Optional. One of "draft", "published", "production".
|
||||
|
|
@ -84,11 +121,27 @@ async def list_policies(version_status: Optional[str] = None):
|
|||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database not connected")
|
||||
|
||||
try:
|
||||
policies = await get_policy_registry().get_all_policies_from_db(prisma_client, version_status=version_status)
|
||||
registry = get_policy_registry()
|
||||
db_policies = (
|
||||
await registry.get_all_policies_from_db(prisma_client, version_status=version_status)
|
||||
if prisma_client is not None
|
||||
else []
|
||||
)
|
||||
db_policy_names = {
|
||||
db_policy.policy_name for db_policy in db_policies if db_policy.version_status == "production"
|
||||
}
|
||||
include_config = version_status in (None, "production")
|
||||
config_policies = (
|
||||
[
|
||||
_config_policy_to_db_response(policy_name, policy)
|
||||
for policy_name, policy in registry.list_config_policies().items()
|
||||
if policy_name not in db_policy_names
|
||||
]
|
||||
if include_config
|
||||
else []
|
||||
)
|
||||
policies = db_policies + config_policies
|
||||
return PolicyListDBResponse(policies=policies, total_count=len(policies))
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error listing policies: {e}")
|
||||
|
|
@ -606,7 +659,10 @@ async def test_pipeline(
|
|||
)
|
||||
async def list_policy_attachments():
|
||||
"""
|
||||
List all policy attachments from the database.
|
||||
List all policy attachments from the database and config.yaml.
|
||||
|
||||
Config-defined attachments are returned with definition_location "config" and a
|
||||
synthetic attachment_id ("config-<index>").
|
||||
|
||||
Example Request:
|
||||
```bash
|
||||
|
|
@ -635,11 +691,14 @@ async def list_policy_attachments():
|
|||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database not connected")
|
||||
|
||||
try:
|
||||
attachments = await get_attachment_registry().get_all_attachments_from_db(prisma_client)
|
||||
registry = get_attachment_registry()
|
||||
db_attachments = await registry.get_all_attachments_from_db(prisma_client) if prisma_client is not None else []
|
||||
config_attachments = [
|
||||
_config_attachment_to_db_response(index, attachment)
|
||||
for index, attachment in enumerate(registry.get_config_attachments())
|
||||
]
|
||||
attachments = db_attachments + config_attachments
|
||||
return PolicyAttachmentListResponse(attachments=attachments, total_count=len(attachments))
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error listing policy attachments: {e}")
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from datetime import datetime, timezone
|
|||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypedDict,
|
||||
|
|
@ -162,6 +163,8 @@ class PolicyRegistry:
|
|||
|
||||
def __init__(self):
|
||||
self._policies: dict[str, Policy] = {}
|
||||
self._config_policies: Mapping[str, Policy] = {}
|
||||
self._sources: Mapping[str, Literal["db", "config"]] = {}
|
||||
self._policies_by_id: dict[str, tuple[str, Policy]] = {}
|
||||
self._initialized: bool = False
|
||||
|
||||
|
|
@ -174,6 +177,8 @@ class PolicyRegistry:
|
|||
This is the raw config from the YAML file.
|
||||
"""
|
||||
self._policies = {}
|
||||
self._config_policies = {}
|
||||
self._sources = {}
|
||||
self._policies_by_id = {}
|
||||
|
||||
for policy_name, policy_data in policies_config.items():
|
||||
|
|
@ -185,6 +190,8 @@ class PolicyRegistry:
|
|||
verbose_proxy_logger.error(f"Error loading policy '{policy_name}': {str(e)}")
|
||||
raise ValueError(f"Invalid policy '{policy_name}': {str(e)}") from e
|
||||
|
||||
self._config_policies = dict(self._policies)
|
||||
self._sources = {policy_name: "config" for policy_name in self._policies}
|
||||
self._initialized = True
|
||||
verbose_proxy_logger.info(f"Loaded {len(self._policies)} policies")
|
||||
|
||||
|
|
@ -299,23 +306,42 @@ class PolicyRegistry:
|
|||
Clear all policies from the registry.
|
||||
"""
|
||||
self._policies = {}
|
||||
self._config_policies = {}
|
||||
self._sources = {}
|
||||
self._initialized = False
|
||||
|
||||
def add_policy(self, policy_name: str, policy: Policy) -> None:
|
||||
def get_source(self, policy_name: str) -> Optional[Literal["db", "config"]]:
|
||||
"""
|
||||
Return the provenance of an in-memory policy, or None if unknown.
|
||||
"""
|
||||
return self._sources.get(policy_name)
|
||||
|
||||
def list_config_policies(self) -> Mapping[str, Policy]:
|
||||
"""
|
||||
Return the policies loaded from config.yaml, keyed by policy name.
|
||||
"""
|
||||
return dict(self._config_policies)
|
||||
|
||||
def add_policy(self, policy_name: str, policy: Policy, source: Literal["db", "config"] = "db") -> None:
|
||||
"""
|
||||
Add or update a single policy.
|
||||
|
||||
Args:
|
||||
policy_name: Name of the policy
|
||||
policy: Policy object to add
|
||||
source: Provenance of the policy ("db" or "config")
|
||||
"""
|
||||
self._policies[policy_name] = policy
|
||||
self._sources = {**self._sources, policy_name: source}
|
||||
if source == "config":
|
||||
self._config_policies = {**self._config_policies, policy_name: policy}
|
||||
self._initialized = True
|
||||
verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}")
|
||||
|
||||
def remove_policy(self, policy_name: str) -> bool:
|
||||
"""
|
||||
Remove a policy by name.
|
||||
Remove a policy by name. If a config-defined policy shares the name,
|
||||
it is restored immediately instead of waiting for the next DB sync.
|
||||
|
||||
Args:
|
||||
policy_name: Name of the policy to remove
|
||||
|
|
@ -323,11 +349,18 @@ class PolicyRegistry:
|
|||
Returns:
|
||||
True if policy was removed, False if it didn't exist
|
||||
"""
|
||||
if policy_name in self._policies:
|
||||
del self._policies[policy_name]
|
||||
verbose_proxy_logger.debug(f"Removed policy: {policy_name}")
|
||||
if policy_name not in self._policies:
|
||||
return False
|
||||
config_fallback = self._config_policies.get(policy_name)
|
||||
if config_fallback is not None:
|
||||
self._policies[policy_name] = config_fallback
|
||||
self._sources = {**self._sources, policy_name: "config"}
|
||||
verbose_proxy_logger.debug(f"Removed policy: {policy_name}; restored config-defined version")
|
||||
return True
|
||||
return False
|
||||
del self._policies[policy_name]
|
||||
self._sources = {name: source for name, source in self._sources.items() if name != policy_name}
|
||||
verbose_proxy_logger.debug(f"Removed policy: {policy_name}")
|
||||
return True
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────
|
||||
# Database CRUD Methods
|
||||
|
|
@ -501,10 +534,15 @@ class PolicyRegistry:
|
|||
# Remove from in-memory registry only if this was the production version
|
||||
if version_status == "production":
|
||||
self.remove_policy(policy_name)
|
||||
result["warning"] = (
|
||||
"Production version was deleted. No other version was promoted. "
|
||||
"Promote another version to production if this policy should remain active."
|
||||
)
|
||||
if self.get_source(policy_name) == "config":
|
||||
result["warning"] = (
|
||||
"Production version was deleted. The config-defined policy with the same name is active again."
|
||||
)
|
||||
else:
|
||||
result["warning"] = (
|
||||
"Production version was deleted. No other version was promoted. "
|
||||
"Promote another version to production if this policy should remain active."
|
||||
)
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
|
|
@ -591,14 +629,14 @@ class PolicyRegistry:
|
|||
"""
|
||||
Sync policies from the database to in-memory registry.
|
||||
- Production versions are loaded into _policies (by policy name) for resolution.
|
||||
- Config-loaded policies are preserved; on a name conflict the DB version wins.
|
||||
- Draft and published versions are loaded into _policies_by_id so request-body
|
||||
policy_<uuid> overrides can be resolved without DB access in the hot path.
|
||||
"""
|
||||
try:
|
||||
self._policies = {}
|
||||
production = await self.get_all_policies_from_db(prisma_client, version_status="production")
|
||||
for policy_response in production:
|
||||
policy = self._parse_policy(
|
||||
db_policies = {
|
||||
policy_response.policy_name: self._parse_policy(
|
||||
policy_response.policy_name,
|
||||
{
|
||||
"inherit": policy_response.inherit,
|
||||
|
|
@ -611,7 +649,16 @@ class PolicyRegistry:
|
|||
"pipeline": policy_response.pipeline,
|
||||
},
|
||||
)
|
||||
self.add_policy(policy_response.policy_name, policy)
|
||||
for policy_response in production
|
||||
}
|
||||
for policy_name in set(db_policies) & set(self._config_policies):
|
||||
verbose_proxy_logger.warning(
|
||||
f"Policy '{policy_name}' is defined in both config.yaml and the DB; the DB version takes precedence"
|
||||
)
|
||||
config_sources: Mapping[str, Literal["db", "config"]] = {name: "config" for name in self._config_policies}
|
||||
db_sources: Mapping[str, Literal["db", "config"]] = {name: "db" for name in db_policies}
|
||||
self._policies = {**self._config_policies, **db_policies}
|
||||
self._sources = {**config_sources, **db_sources}
|
||||
|
||||
self._policies_by_id = {}
|
||||
non_production = await _policy_table(prisma_client).find_many(
|
||||
|
|
@ -637,7 +684,8 @@ class PolicyRegistry:
|
|||
self._initialized = True
|
||||
verbose_proxy_logger.info(
|
||||
f"Synced {len(production)} production policies and {len(non_production)} "
|
||||
"draft/published (by ID) from DB to in-memory registry"
|
||||
"draft/published (by ID) from DB to in-memory registry "
|
||||
f"({len(self._config_policies)} config-defined policies preserved)"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error syncing policies from DB: {e}")
|
||||
|
|
@ -983,12 +1031,20 @@ class PolicyRegistry:
|
|||
prisma_client: The Prisma client instance
|
||||
|
||||
Returns:
|
||||
Dict with success message
|
||||
Dict with "message" and optional "warning" if a config-defined policy took over.
|
||||
"""
|
||||
try:
|
||||
await _policy_table(prisma_client).delete_many(where={"policy_name": policy_name})
|
||||
self.remove_policy(policy_name)
|
||||
return {"message": f"All versions of policy '{policy_name}' deleted successfully"}
|
||||
message = f"All versions of policy '{policy_name}' deleted successfully"
|
||||
if self.get_source(policy_name) == "config":
|
||||
return {
|
||||
"message": message,
|
||||
"warning": (
|
||||
"All DB versions were deleted. The config-defined policy with the same name is active again."
|
||||
),
|
||||
}
|
||||
return {"message": message}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error deleting all versions: {e}")
|
||||
raise Exception(f"Error deleting all versions: {str(e)}")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
156
tests/test_litellm/integrations/test_s3.py
Normal file
156
tests/test_litellm/integrations/test_s3.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.s3 import S3Logger
|
||||
|
||||
TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id"
|
||||
|
||||
|
||||
def _standard_logging_payload() -> dict:
|
||||
return {
|
||||
"id": "chatcmpl-test-id",
|
||||
"metadata": {"user_api_key_team_alias": None},
|
||||
}
|
||||
|
||||
|
||||
def _log_event_kwargs() -> dict:
|
||||
return {
|
||||
"litellm_params": {"metadata": {}},
|
||||
"standard_logging_object": _standard_logging_payload(),
|
||||
}
|
||||
|
||||
|
||||
def _run_log_event(callback_params: dict) -> MagicMock:
|
||||
original = litellm.s3_callback_params
|
||||
litellm.s3_callback_params = callback_params
|
||||
try:
|
||||
with patch("boto3.client") as mock_boto3_client:
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto3_client.return_value = mock_s3_client
|
||||
logger = S3Logger()
|
||||
logger.log_event(
|
||||
kwargs=_log_event_kwargs(),
|
||||
response_obj={},
|
||||
start_time=datetime(2026, 7, 30, 12, 0, 0),
|
||||
end_time=datetime(2026, 7, 30, 12, 0, 1),
|
||||
print_verbose=lambda *args, **kwargs: None,
|
||||
)
|
||||
return mock_s3_client
|
||||
finally:
|
||||
litellm.s3_callback_params = original
|
||||
|
||||
|
||||
def test_put_object_includes_sse_kms_params_when_configured():
|
||||
"""
|
||||
When s3_server_side_encryption and s3_sse_kms_key_id are set in
|
||||
s3_callback_params, put_object must receive ServerSideEncryption and
|
||||
SSEKMSKeyId so objects land encrypted with the customer-managed key.
|
||||
"""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_server_side_encryption": "aws:kms",
|
||||
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
|
||||
assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN
|
||||
|
||||
|
||||
def test_put_object_supports_sse_s3_without_key_id():
|
||||
"""SSE-S3 (AES256) needs only ServerSideEncryption, no key id."""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_server_side_encryption": "AES256",
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert put_object_kwargs["ServerSideEncryption"] == "AES256"
|
||||
assert "SSEKMSKeyId" not in put_object_kwargs
|
||||
|
||||
|
||||
def test_put_object_omits_sse_params_by_default():
|
||||
"""Without SSE config, put_object kwargs must stay unchanged."""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert "ServerSideEncryption" not in put_object_kwargs
|
||||
assert "SSEKMSKeyId" not in put_object_kwargs
|
||||
|
||||
|
||||
def test_put_object_infers_aws_kms_when_only_key_id_set():
|
||||
"""A key id without an algorithm must infer aws:kms instead of sending an invalid request."""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
|
||||
assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN
|
||||
|
||||
|
||||
def test_put_object_drops_key_id_when_algorithm_is_not_kms():
|
||||
"""AES256 plus a key id is invalid for S3; the key id must be dropped, not sent."""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_server_side_encryption": "AES256",
|
||||
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert put_object_kwargs["ServerSideEncryption"] == "AES256"
|
||||
assert "SSEKMSKeyId" not in put_object_kwargs
|
||||
|
||||
|
||||
def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued():
|
||||
"""
|
||||
A YAML boolean in s3_server_side_encryption must not crash logger init and
|
||||
must not discard the valid key id; aws:kms is inferred from the key id.
|
||||
"""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_server_side_encryption": True,
|
||||
"s3_sse_kms_key_id": TEST_KMS_KEY_ARN,
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
|
||||
assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN
|
||||
|
||||
|
||||
def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept():
|
||||
"""A mistyped key id (unquoted YAML number) must not disable the valid algorithm."""
|
||||
mock_s3_client = _run_log_event(
|
||||
{
|
||||
"s3_bucket_name": "test-bucket",
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_server_side_encryption": "aws:kms",
|
||||
"s3_sse_kms_key_id": 12345,
|
||||
}
|
||||
)
|
||||
|
||||
put_object_kwargs = mock_s3_client.put_object.call_args.kwargs
|
||||
assert put_object_kwargs["ServerSideEncryption"] == "aws:kms"
|
||||
assert "SSEKMSKeyId" not in put_object_kwargs
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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() == ()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -2729,7 +2729,7 @@
|
|||
"count": 1
|
||||
},
|
||||
"no-nested-ternary": {
|
||||
"count": 6
|
||||
"count": 4
|
||||
}
|
||||
},
|
||||
"src/components/chat/MCPConnectPicker.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 (
|
||||
<DropdownMenu>
|
||||
<DropdownMenuTrigger
|
||||
|
|
@ -65,6 +70,8 @@ function AttachmentRowActions({ attachment, isAdmin, onDeleteClick }: Attachment
|
|||
<DropdownMenuItem
|
||||
variant="destructive"
|
||||
data-testid="attachment-action-delete"
|
||||
disabled={isConfigAttachment}
|
||||
title={isConfigAttachment ? CONFIG_ATTACHMENT_HINT : undefined}
|
||||
onClick={() => onDeleteClick(attachment.attachment_id)}
|
||||
>
|
||||
<Trash2 />
|
||||
|
|
|
|||
|
|
@ -145,4 +145,34 @@ describe("PolicyTable", () => {
|
|||
await user.click(screen.getByRole("button", { name: /grouped/ }));
|
||||
expect(defaultProps.onViewClick).toHaveBeenCalledWith("prod-id");
|
||||
});
|
||||
|
||||
const sameNamedDbDraft: Partial<Policy> = {
|
||||
policy_name: "config-policy",
|
||||
policy_id: "db-draft-id",
|
||||
version_status: "draft",
|
||||
version_number: 2,
|
||||
};
|
||||
const configTwin: Partial<Policy> = {
|
||||
policy_name: "config-policy",
|
||||
policy_id: "config-policy",
|
||||
version_status: "production",
|
||||
definition_location: "config",
|
||||
};
|
||||
|
||||
it("should render a config policy and a same-named DB draft as separate rows", () => {
|
||||
const policies = [makePolicy(sameNamedDbDraft), makePolicy(configTwin)];
|
||||
renderWithProviders(<PolicyTable {...defaultProps} policies={policies} />);
|
||||
expect(screen.getAllByText("config-policy")).toHaveLength(2);
|
||||
expect(screen.getByText("Config")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should keep a same-named DB draft reachable next to a config policy", async () => {
|
||||
const user = userEvent.setup();
|
||||
const policies = [makePolicy(sameNamedDbDraft), makePolicy(configTwin)];
|
||||
renderWithProviders(<PolicyTable {...defaultProps} policies={policies} />);
|
||||
await user.click(screen.getByRole("button", { name: "config-policy" }));
|
||||
expect(defaultProps.onViewClick).toHaveBeenCalledWith("db-draft-id");
|
||||
await user.click(screen.getByTestId("policy-actions-db-draft-id"));
|
||||
expect(await screen.findByTestId("policy-action-edit")).not.toHaveAttribute("data-disabled");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -9,16 +9,21 @@ import { Policy } from "@/components/policies/types";
|
|||
|
||||
import { getPolicyTableColumns, PolicyRow } from "./PolicyTableColumns";
|
||||
|
||||
/** One row per policy name; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */
|
||||
/** One row per DB policy name plus one row per config policy, so a config policy never hides same-named DB versions; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */
|
||||
function groupPoliciesByName(policies: Policy[]): PolicyRow[] {
|
||||
const names = Array.from(new Set(policies.map((policy) => policy.policy_name || "(unnamed)")));
|
||||
return names.map((policyName) => {
|
||||
const versions = policies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName);
|
||||
const dbPolicies = policies.filter((policy) => policy.definition_location !== "config");
|
||||
const names = Array.from(new Set(dbPolicies.map((policy) => policy.policy_name || "(unnamed)")));
|
||||
const dbRows = names.map((policyName) => {
|
||||
const versions = dbPolicies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName);
|
||||
const primary =
|
||||
versions.find((version) => version.version_status === "production") ??
|
||||
[...versions].sort((a, b) => (b.version_number ?? 0) - (a.version_number ?? 0))[0];
|
||||
return { policy_name: policyName, primaryPolicy: primary, versionCount: versions.length };
|
||||
});
|
||||
const configRows = policies
|
||||
.filter((policy) => policy.definition_location === "config")
|
||||
.map((policy) => ({ policy_name: policy.policy_name || "(unnamed)", primaryPolicy: policy, versionCount: 1 }));
|
||||
return [...dbRows, ...configRows];
|
||||
}
|
||||
|
||||
interface PolicyTableProps {
|
||||
|
|
@ -67,7 +72,7 @@ const PolicyTable: React.FC<PolicyTableProps> = ({
|
|||
<DataTable
|
||||
data={rows}
|
||||
columns={columns}
|
||||
getRowId={(row) => row.policy_name}
|
||||
getRowId={(row) => `${row.primaryPolicy.definition_location ?? "db"}:${row.policy_name}`}
|
||||
sortingMode="client"
|
||||
sorting={sorting}
|
||||
onSortingChange={setSorting}
|
||||
|
|
|
|||
|
|
@ -22,6 +22,9 @@ export interface PolicyRow {
|
|||
versionCount: number;
|
||||
}
|
||||
|
||||
const CONFIG_POLICY_HINT =
|
||||
"Config policies are defined in the config file and cannot be edited or deleted from the dashboard.";
|
||||
|
||||
function GuardrailChips({ guardrails, tone }: { guardrails: string[]; tone: "success" | "error" }) {
|
||||
if (guardrails.length === 0) {
|
||||
return <span className="text-muted-foreground">-</span>;
|
||||
|
|
@ -45,6 +48,8 @@ interface PolicyRowActionsProps {
|
|||
}
|
||||
|
||||
function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActionsProps) {
|
||||
const isConfigPolicy = policy.definition_location === "config";
|
||||
|
||||
return (
|
||||
<DropdownMenu>
|
||||
<DropdownMenuTrigger
|
||||
|
|
@ -55,7 +60,12 @@ function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActio
|
|||
<MoreHorizontal className="size-4" />
|
||||
</DropdownMenuTrigger>
|
||||
<DropdownMenuContent align="end" className="w-52">
|
||||
<DropdownMenuItem data-testid="policy-action-edit" onClick={() => onEditClick(policy)}>
|
||||
<DropdownMenuItem
|
||||
data-testid="policy-action-edit"
|
||||
disabled={isConfigPolicy}
|
||||
title={isConfigPolicy ? CONFIG_POLICY_HINT : undefined}
|
||||
onClick={() => onEditClick(policy)}
|
||||
>
|
||||
<Pencil />
|
||||
Edit policy
|
||||
</DropdownMenuItem>
|
||||
|
|
@ -63,6 +73,8 @@ function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActio
|
|||
<DropdownMenuItem
|
||||
variant="destructive"
|
||||
data-testid="policy-action-delete"
|
||||
disabled={isConfigPolicy}
|
||||
title={isConfigPolicy ? CONFIG_POLICY_HINT : undefined}
|
||||
onClick={() => onDeleteClick(policy.policy_id, policy.policy_name || "Unnamed Policy")}
|
||||
>
|
||||
<Trash2 />
|
||||
|
|
@ -93,18 +105,23 @@ export const getPolicyTableColumns = ({
|
|||
header: ({ column }) => <DataTableSortHeader column={column} title="Name" />,
|
||||
size: 220,
|
||||
enableSorting: true,
|
||||
cell: ({ row }) => (
|
||||
<IdentityCell
|
||||
title={row.original.policy_name}
|
||||
titleClassName="max-w-60"
|
||||
badge={
|
||||
row.original.versionCount > 1 ? (
|
||||
<StatusBadge tone="neutral" label={`${row.original.versionCount} versions`} />
|
||||
) : undefined
|
||||
}
|
||||
onClick={() => onViewClick(row.original.primaryPolicy.policy_id)}
|
||||
/>
|
||||
),
|
||||
cell: ({ row }) => {
|
||||
const isConfigPolicy = row.original.primaryPolicy.definition_location === "config";
|
||||
const versionBadge =
|
||||
row.original.versionCount > 1 ? (
|
||||
<StatusBadge tone="neutral" label={`${row.original.versionCount} versions`} />
|
||||
) : undefined;
|
||||
return (
|
||||
<IdentityCell
|
||||
title={row.original.policy_name}
|
||||
titleClassName="max-w-60"
|
||||
badge={
|
||||
isConfigPolicy ? <StatusBadge tone="neutral" label="Config" tooltip={CONFIG_POLICY_HINT} /> : versionBadge
|
||||
}
|
||||
onClick={isConfigPolicy ? undefined : () => onViewClick(row.original.primaryPolicy.policy_id)}
|
||||
/>
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "description",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,13 @@
|
|||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { Select as AntdSelect, Card, InputNumber, Radio, Space, Tooltip, Typography } from "antd";
|
||||
import React from "react";
|
||||
import { ClassifierType, ComplexityRouterConfigValue, DEFAULT_CLASSIFIER_TIMEOUT_MS } from "./ComplexityRouterConfig";
|
||||
import {
|
||||
ClassifierType,
|
||||
ComplexityRouterConfigValue,
|
||||
DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
|
||||
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
|
||||
DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
} from "./ComplexityRouterConfig";
|
||||
|
||||
const { Text } = Typography;
|
||||
|
||||
|
|
@ -26,14 +32,23 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model;
|
||||
|
||||
const handleClassifierTypeChange = (classifierType: ClassifierType) => {
|
||||
onChange({
|
||||
const nextValue: ComplexityRouterConfigValue = {
|
||||
...value,
|
||||
classifier_type: classifierType,
|
||||
classifier_llm_config:
|
||||
classifierType === "llm"
|
||||
? value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS }
|
||||
: undefined,
|
||||
});
|
||||
classifier_context_window_size:
|
||||
classifierType === "llm"
|
||||
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
|
||||
: undefined,
|
||||
classifier_context_per_turn_chars:
|
||||
classifierType === "llm"
|
||||
? value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS
|
||||
: undefined,
|
||||
};
|
||||
onChange(nextValue);
|
||||
};
|
||||
|
||||
const handleClassifierModelChange = (model: string) => {
|
||||
|
|
@ -56,6 +71,20 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
});
|
||||
};
|
||||
|
||||
const handleClassifierContextWindowSizeChange = (windowSize: number | null) => {
|
||||
onChange({
|
||||
...value,
|
||||
classifier_context_window_size: windowSize ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
|
||||
});
|
||||
};
|
||||
|
||||
const handleClassifierContextPerTurnCharsChange = (perTurnChars: number | null) => {
|
||||
onChange({
|
||||
...value,
|
||||
classifier_context_per_turn_chars: perTurnChars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
|
||||
});
|
||||
};
|
||||
|
||||
return (
|
||||
<>
|
||||
<Radio.Group
|
||||
|
|
@ -111,6 +140,36 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
response.
|
||||
</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text strong style={{ display: "block", marginBottom: 4 }}>
|
||||
Context Window Size
|
||||
</Text>
|
||||
<InputNumber
|
||||
value={value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE}
|
||||
onChange={handleClassifierContextWindowSizeChange}
|
||||
min={0}
|
||||
style={{ width: "100%" }}
|
||||
/>
|
||||
<Text type="secondary" style={{ fontSize: 12 }}>
|
||||
Number of prior user turns (tool output and harness reminders excluded) sent to the classifier as context,
|
||||
so a referring follow-up like "now do the same for the streaming path" is classified against
|
||||
what it refers to. Set to 0 to send only the current message.
|
||||
</Text>
|
||||
</div>
|
||||
<div>
|
||||
<Text strong style={{ display: "block", marginBottom: 4 }}>
|
||||
Context Per-Turn Character Limit
|
||||
</Text>
|
||||
<InputNumber
|
||||
value={value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS}
|
||||
onChange={handleClassifierContextPerTurnCharsChange}
|
||||
min={1}
|
||||
style={{ width: "100%" }}
|
||||
/>
|
||||
<Text type="secondary" style={{ fontSize: 12 }}>
|
||||
Prior turns longer than this are truncated.
|
||||
</Text>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
|
|
|
|||
|
|
@ -99,11 +99,14 @@ describe("ComplexityRouterConfig", () => {
|
|||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
fireEvent.click(screen.getByText("LLM Classifier"));
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith({
|
||||
const expectedValue: ComplexityRouterConfigValue = {
|
||||
...defaultValue,
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "", timeout_ms: 3000 },
|
||||
});
|
||||
classifier_context_window_size: 3,
|
||||
classifier_context_per_turn_chars: 200,
|
||||
};
|
||||
expect(onChange).toHaveBeenCalledWith(expectedValue);
|
||||
});
|
||||
|
||||
it("should show classifier fields and use the configured values when classifier_type is llm", () => {
|
||||
|
|
@ -111,6 +114,8 @@ describe("ComplexityRouterConfig", () => {
|
|||
...defaultValue,
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 750 },
|
||||
classifier_context_window_size: 5,
|
||||
classifier_context_per_turn_chars: 400,
|
||||
};
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={vi.fn()} />);
|
||||
|
||||
|
|
@ -119,6 +124,74 @@ describe("ComplexityRouterConfig", () => {
|
|||
expect(screen.getByText("Classifier Model")).toBeInTheDocument();
|
||||
expect(screen.getByText("Timeout (ms)")).toBeInTheDocument();
|
||||
expect(screen.getByDisplayValue("750")).toBeInTheDocument();
|
||||
expect(screen.getByText("Context Window Size")).toBeInTheDocument();
|
||||
expect(screen.getByDisplayValue("5")).toBeInTheDocument();
|
||||
expect(screen.getByText("Context Per-Turn Character Limit")).toBeInTheDocument();
|
||||
expect(screen.getByDisplayValue("400")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should default classifier context fields to 3 and 200 when llm is selected without explicit values", () => {
|
||||
const llmValue: ComplexityRouterConfigValue = {
|
||||
...defaultValue,
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
|
||||
};
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={vi.fn()} />);
|
||||
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
|
||||
const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement;
|
||||
expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument();
|
||||
|
||||
const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement;
|
||||
expect(within(perTurnCharsSection).getByDisplayValue("200")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should hide classifier context fields when classifier_type is heuristic", () => {
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={defaultValue} onChange={vi.fn()} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
expect(screen.queryByText("Context Window Size")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should call onChange with the updated classifier_context_window_size when edited", () => {
|
||||
const onChange = vi.fn();
|
||||
const llmValue: ComplexityRouterConfigValue = {
|
||||
...defaultValue,
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
|
||||
};
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={onChange} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
|
||||
const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement;
|
||||
const input = within(windowSizeSection).getByRole("spinbutton");
|
||||
fireEvent.change(input, { target: { value: "7" } });
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith({
|
||||
...llmValue,
|
||||
classifier_context_window_size: 7,
|
||||
});
|
||||
});
|
||||
|
||||
it("should call onChange with the updated classifier_context_per_turn_chars when edited", () => {
|
||||
const onChange = vi.fn();
|
||||
const llmValue: ComplexityRouterConfigValue = {
|
||||
...defaultValue,
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
|
||||
};
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={llmValue} onChange={onChange} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
|
||||
const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement;
|
||||
const input = within(perTurnCharsSection).getByRole("spinbutton");
|
||||
fireEvent.change(input, { target: { value: "500" } });
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith({
|
||||
...llmValue,
|
||||
classifier_context_per_turn_chars: 500,
|
||||
});
|
||||
});
|
||||
|
||||
it("should render the custom technical keywords field", () => {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -99,6 +99,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
tiers,
|
||||
classifier_type: classifierType,
|
||||
classifier_llm_config: classifierLlmConfig,
|
||||
classifier_context_window_size: classifierContextWindowSize,
|
||||
classifier_context_per_turn_chars: classifierContextPerTurnChars,
|
||||
adaptive = false,
|
||||
adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS,
|
||||
tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY,
|
||||
|
|
@ -142,6 +144,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
tiers,
|
||||
classifierType,
|
||||
classifierLlmConfig,
|
||||
classifierContextWindowSize,
|
||||
classifierContextPerTurnChars,
|
||||
customTechnicalKeywords,
|
||||
keywordTierRules,
|
||||
semanticMatchingEnabled,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import React from "react";
|
||||
import { render, screen, fireEvent } from "@testing-library/react";
|
||||
import { describe, it, expect, vi, afterEach } from "vitest";
|
||||
import { render, screen, fireEvent, waitFor, act } from "@testing-library/react";
|
||||
import { describe, it, expect, vi, afterEach, beforeEach } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import MCPAppsPanel from "./MCPAppsPanel";
|
||||
import { fetchMCPServers, listMCPTools } from "../networking";
|
||||
|
|
@ -86,3 +86,208 @@ describe("MCPAppsPanel logos", () => {
|
|||
expect(screen.getByAltText("local_logo logo").getAttribute("src")).toBe("/litellm/ui/assets/logos/github.svg");
|
||||
});
|
||||
});
|
||||
|
||||
const connectServers = [
|
||||
{
|
||||
server_id: "s-reach",
|
||||
server_name: "reachable_srv",
|
||||
auth_type: "none",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
{
|
||||
server_id: "s-unreach",
|
||||
server_name: "unreachable_srv",
|
||||
auth_type: "none",
|
||||
connected_app_reachable: false,
|
||||
},
|
||||
] as MCPServer[];
|
||||
|
||||
const renderConnectPanel = (connectMode: boolean, selectedServers: string[] = []) =>
|
||||
render(
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel accessToken="tok" selectedServers={selectedServers} onChange={vi.fn()} connectMode={connectMode} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("requests the connected-app view and hides unreachable servers in connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(true, ["reachable_srv", "unreachable_srv"]);
|
||||
|
||||
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
|
||||
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true);
|
||||
expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected (1)")).toBeInTheDocument();
|
||||
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
|
||||
expect(toolCountFetchedIds).toContain("s-reach");
|
||||
expect(toolCountFetchedIds).not.toContain("s-unreach");
|
||||
});
|
||||
|
||||
it("blocks connecting an unsupported server from the detail view in connect mode", async () => {
|
||||
const detailServers = [
|
||||
...connectServers,
|
||||
{
|
||||
server_id: "s-unsup",
|
||||
server_name: "unsupported_srv",
|
||||
auth_type: "oauth2_token_exchange",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
] as MCPServer[];
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(detailServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(true);
|
||||
|
||||
fireEvent.click(await screen.findByText("unsupported_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "unsupported_srv" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Not supported on this connection")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the detail-view Connect action outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(false);
|
||||
|
||||
fireEvent.click(await screen.findByText("unreachable_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "unreachable_srv" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Connect" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("ignores the flag and skips no server outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(false, ["reachable_srv", "unreachable_srv"]);
|
||||
|
||||
expect(await screen.findByText("unreachable_srv")).toBeInTheDocument();
|
||||
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, false);
|
||||
expect(screen.queryByText("Not available to connected apps")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected (2)")).toBeInTheDocument();
|
||||
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
|
||||
expect(toolCountFetchedIds).toContain("s-unreach");
|
||||
});
|
||||
|
||||
const revocable = (reachable: boolean) =>
|
||||
[
|
||||
{ server_id: "s-reach", server_name: "reachable_srv", auth_type: "none", connected_app_reachable: true },
|
||||
{ server_id: "s-drop", server_name: "revoked_srv", auth_type: "none", connected_app_reachable: reachable },
|
||||
] as MCPServer[];
|
||||
|
||||
const ConnectPanel = ({
|
||||
token,
|
||||
onChange,
|
||||
client,
|
||||
}: {
|
||||
token: string;
|
||||
onChange: (servers: string[]) => void;
|
||||
client: QueryClient;
|
||||
}) => (
|
||||
<QueryClientProvider client={client}>
|
||||
<MCPAppsPanel accessToken={token} selectedServers={[]} onChange={onChange} connectMode />
|
||||
</QueryClientProvider>
|
||||
);
|
||||
|
||||
const newClient = () => new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
|
||||
it("drops an open detail view when a refetch removes that server from the reachable set", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={vi.fn()} client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={vi.fn()} client={client} />);
|
||||
|
||||
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
|
||||
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("reachable_srv")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not select a server whose Connect finishes after a refetch removed it", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const onChange = vi.fn();
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={onChange} client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
let finishConnect: (result: { tools: never[] }) => void = () => {};
|
||||
vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve)));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Connect" }));
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={onChange} client={client} />);
|
||||
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
|
||||
|
||||
await act(async () => {
|
||||
finishConnect({ tools: [] });
|
||||
});
|
||||
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected", { exact: false }).textContent).toBe("Connected");
|
||||
});
|
||||
|
||||
it("does not select a server when Connect resolves in the same tick the refetch drops it", async () => {
|
||||
let finishRefetch: (servers: MCPServer[]) => void = () => {};
|
||||
vi.mocked(fetchMCPServers)
|
||||
.mockResolvedValueOnce(revocable(true))
|
||||
.mockImplementationOnce(() => new Promise((resolve) => (finishRefetch = resolve)));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const onChange = vi.fn();
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={onChange} client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
let finishConnect: (result: { tools: never[] }) => void = () => {};
|
||||
vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve)));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Connect" }));
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={onChange} client={client} />);
|
||||
|
||||
await act(async () => {
|
||||
finishRefetch(revocable(false));
|
||||
finishConnect({ tools: [] });
|
||||
});
|
||||
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not let a superseded list load overwrite the current reachable set", async () => {
|
||||
let finishStaleLoad: (servers: MCPServer[]) => void = () => {};
|
||||
vi.mocked(fetchMCPServers)
|
||||
.mockImplementationOnce(() => new Promise((resolve) => (finishStaleLoad = resolve)))
|
||||
.mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={vi.fn()} client={client} />);
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={vi.fn()} client={client} />);
|
||||
|
||||
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
|
||||
|
||||
await act(async () => {
|
||||
finishStaleLoad(revocable(true));
|
||||
});
|
||||
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -103,16 +103,17 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
const [query, setQuery] = useState("");
|
||||
const [activeTab, setActiveTab] = useState<TabKey>("all");
|
||||
const [togglingOn, setTogglingOn] = useState<Set<string>>(new Set());
|
||||
const [detailServer, setDetailServer] = useState<MCPServer | null>(null);
|
||||
const [detailServerId, setDetailServerId] = useState<string | null>(null);
|
||||
const [toolCounts, setToolCounts] = useState<Record<string, number>>({});
|
||||
const [loadingCounts, setLoadingCounts] = useState(false);
|
||||
const [oauthConnected, setOauthConnected] = useState<Set<string>>(new Set());
|
||||
const [oauthChecking, setOauthChecking] = useState<Set<string>>(new Set());
|
||||
|
||||
const serversRef = useRef<MCPServer[]>([]);
|
||||
useEffect(() => {
|
||||
serversRef.current = servers;
|
||||
}, [servers]);
|
||||
const commitServers = useCallback((next: MCPServer[]) => {
|
||||
serversRef.current = next;
|
||||
setServers(next);
|
||||
}, []);
|
||||
const selectedServersRef = useRef<string[]>(selectedServers);
|
||||
useEffect(() => {
|
||||
selectedServersRef.current = selectedServers;
|
||||
|
|
@ -124,13 +125,30 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
|
||||
const nameOf = (s: MCPServer) => s.server_name ?? s.alias ?? s.server_id;
|
||||
|
||||
const fetchLoadCancelledRef = useRef(false);
|
||||
const detailServer = servers.find((s) => s.server_id === detailServerId);
|
||||
|
||||
const connectUnavailabilityLabel = useCallback(
|
||||
(s: MCPServer): string | null => {
|
||||
if (!connectMode) return null;
|
||||
if (isUnsupportedOnGatewayConnect(s.auth_type)) return "Not supported on this connection";
|
||||
return null;
|
||||
},
|
||||
[connectMode],
|
||||
);
|
||||
|
||||
const connectableNow = useCallback(
|
||||
(serverId: string): MCPServer | undefined => {
|
||||
const current = serversRef.current.find((s) => s.server_id === serverId);
|
||||
return current !== undefined && connectUnavailabilityLabel(current) === null ? current : undefined;
|
||||
},
|
||||
[connectUnavailabilityLabel],
|
||||
);
|
||||
|
||||
const fetchToolCount = useCallback(
|
||||
async (server: MCPServer) => {
|
||||
async (server: MCPServer, isCurrentLoad: () => boolean) => {
|
||||
try {
|
||||
const toolsData = await listMCPTools(accessToken, server.server_id);
|
||||
if (fetchLoadCancelledRef.current) return;
|
||||
if (!isCurrentLoad()) return;
|
||||
const tools: MCPTool[] = Array.isArray(toolsData?.tools) ? toolsData.tools : [];
|
||||
setToolCounts((prev) => ({ ...prev, [nameOf(server)]: tools.length }));
|
||||
} catch {
|
||||
|
|
@ -141,17 +159,17 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
);
|
||||
|
||||
const checkOauthCredential = useCallback(
|
||||
async (server: MCPServer) => {
|
||||
async (server: MCPServer, isCurrentLoad: () => boolean) => {
|
||||
try {
|
||||
const status = await getMCPOAuthUserCredentialStatus(accessToken, server.server_id);
|
||||
if (fetchLoadCancelledRef.current) return;
|
||||
if (!isCurrentLoad()) return;
|
||||
if (status.has_credential && !status.is_expired) {
|
||||
setOauthConnected((prev) => new Set(prev).add(server.server_id));
|
||||
}
|
||||
} catch {
|
||||
// ignore
|
||||
} finally {
|
||||
if (!fetchLoadCancelledRef.current) {
|
||||
if (isCurrentLoad()) {
|
||||
setOauthChecking((prev) => {
|
||||
const next = new Set(prev);
|
||||
next.delete(server.server_id);
|
||||
|
|
@ -164,70 +182,77 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
);
|
||||
|
||||
useEffect(() => {
|
||||
fetchLoadCancelledRef.current = false;
|
||||
let current = true;
|
||||
const isCurrentLoad = () => current;
|
||||
|
||||
fetchMCPServers(accessToken)
|
||||
fetchMCPServers(accessToken, undefined, connectMode)
|
||||
.then(async (serverData) => {
|
||||
if (fetchLoadCancelledRef.current) return;
|
||||
if (!isCurrentLoad()) return;
|
||||
const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? [];
|
||||
const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
|
||||
setServers(list);
|
||||
const reachable = connectMode ? list.filter((s) => s.connected_app_reachable !== false) : list;
|
||||
const oauthServers = reachable.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
|
||||
commitServers(reachable);
|
||||
setOauthChecking(new Set(oauthServers.map((s) => s.server_id)));
|
||||
setLoading(false);
|
||||
|
||||
oauthServers.forEach((s) => checkOauthCredential(s));
|
||||
oauthServers.forEach((s) => checkOauthCredential(s, isCurrentLoad));
|
||||
|
||||
setLoadingCounts(true);
|
||||
const chunks = Array.from({ length: Math.ceil(list.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) =>
|
||||
list.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY),
|
||||
const chunks = Array.from({ length: Math.ceil(reachable.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) =>
|
||||
reachable.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY),
|
||||
);
|
||||
for (const chunk of chunks) {
|
||||
if (fetchLoadCancelledRef.current) return;
|
||||
await Promise.allSettled(chunk.map((s) => fetchToolCount(s)));
|
||||
if (!isCurrentLoad()) return;
|
||||
await Promise.allSettled(chunk.map((s) => fetchToolCount(s, isCurrentLoad)));
|
||||
}
|
||||
if (!fetchLoadCancelledRef.current) setLoadingCounts(false);
|
||||
if (isCurrentLoad()) setLoadingCounts(false);
|
||||
})
|
||||
.catch(() => {
|
||||
if (!fetchLoadCancelledRef.current) {
|
||||
setServers([]);
|
||||
if (isCurrentLoad()) {
|
||||
commitServers([]);
|
||||
setLoading(false);
|
||||
}
|
||||
});
|
||||
return () => {
|
||||
fetchLoadCancelledRef.current = true;
|
||||
current = false;
|
||||
};
|
||||
}, [accessToken, fetchToolCount, checkOauthCredential]);
|
||||
}, [accessToken, connectMode, commitServers, fetchToolCount, checkOauthCredential]);
|
||||
|
||||
useEffect(() => {
|
||||
if (oauthConnected.size === 0) return;
|
||||
const namesToAdd = serversRef.current
|
||||
.filter((s) => oauthConnected.has(s.server_id) && !selectedServersRef.current.includes(nameOf(s)))
|
||||
.filter(
|
||||
(s) =>
|
||||
oauthConnected.has(s.server_id) &&
|
||||
!selectedServersRef.current.includes(nameOf(s)) &&
|
||||
connectUnavailabilityLabel(s) === null,
|
||||
)
|
||||
.map(nameOf);
|
||||
if (namesToAdd.length > 0) {
|
||||
onChangeRef.current([...selectedServersRef.current, ...namesToAdd]);
|
||||
}
|
||||
}, [oauthConnected]);
|
||||
}, [oauthConnected, connectUnavailabilityLabel]);
|
||||
|
||||
const handleToggle = async (serverName: string, checked: boolean, serverId?: string) => {
|
||||
const handleToggle = async (server: MCPServer, checked: boolean) => {
|
||||
const serverName = nameOf(server);
|
||||
if (!checked) {
|
||||
onChange(selectedServers.filter((s) => s !== serverName));
|
||||
if (serverId) {
|
||||
setOauthConnected((prev) => {
|
||||
const next = new Set(prev);
|
||||
next.delete(serverId);
|
||||
return next;
|
||||
});
|
||||
}
|
||||
setOauthConnected((prev) => {
|
||||
const next = new Set(prev);
|
||||
next.delete(server.server_id);
|
||||
return next;
|
||||
});
|
||||
return;
|
||||
}
|
||||
if (connectableNow(server.server_id) === undefined) return;
|
||||
setTogglingOn((prev) => new Set(prev).add(serverName));
|
||||
try {
|
||||
const idToFetch = serverId ?? serverName;
|
||||
const result = await listMCPTools(accessToken, idToFetch);
|
||||
const result = await listMCPTools(accessToken, server.server_id);
|
||||
if (result?.error) {
|
||||
MessageManager.warning(`Could not load tools for ${serverName}`);
|
||||
return;
|
||||
}
|
||||
if (connectableNow(server.server_id) === undefined) return;
|
||||
if (!selectedServersRef.current.includes(serverName)) {
|
||||
onChange([...selectedServersRef.current, serverName]);
|
||||
}
|
||||
|
|
@ -243,11 +268,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
};
|
||||
|
||||
const renderConnectionIndicator = (server: MCPServer) => {
|
||||
if (connectMode && isUnsupportedOnGatewayConnect(server.auth_type)) {
|
||||
const unavailabilityLabel = connectUnavailabilityLabel(server);
|
||||
if (unavailabilityLabel !== null) {
|
||||
return (
|
||||
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">
|
||||
Not supported on this connection
|
||||
</span>
|
||||
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">{unavailabilityLabel}</span>
|
||||
);
|
||||
}
|
||||
if (server.auth_type === AUTH_TYPE.OAUTH2) {
|
||||
|
|
@ -285,11 +309,23 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
!query.trim() ||
|
||||
name.toLowerCase().includes(query.toLowerCase()) ||
|
||||
(s.description ?? "").toLowerCase().includes(query.toLowerCase());
|
||||
const matchesTab = activeTab === "all" || selectedServers.includes(name);
|
||||
const matchesTab =
|
||||
activeTab === "all" || (selectedServers.includes(name) && connectUnavailabilityLabel(s) === null);
|
||||
return matchesQuery && matchesTab;
|
||||
});
|
||||
|
||||
const connectedCount = servers.filter((s) => selectedServers.includes(nameOf(s))).length;
|
||||
const connectedCount = servers.filter(
|
||||
(s) => selectedServers.includes(nameOf(s)) && connectUnavailabilityLabel(s) === null,
|
||||
).length;
|
||||
|
||||
const emptyStateText = () => {
|
||||
if (servers.length === 0) {
|
||||
return connectMode
|
||||
? "No MCP servers are available to this connection yet. Ask an admin to grant your user or team access."
|
||||
: "No MCP servers configured. Add servers in Tools -> MCP Servers.";
|
||||
}
|
||||
return activeTab === "connected" ? "No servers connected yet." : "No servers match your search.";
|
||||
};
|
||||
const totalTools = Object.values(toolCounts).reduce((sum, n) => sum + n, 0);
|
||||
|
||||
if (detailServer) {
|
||||
|
|
@ -298,12 +334,65 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
const isTogglingOn = togglingOn.has(name);
|
||||
const color = getAvatarColor(name);
|
||||
|
||||
const renderDetailAction = () => {
|
||||
const unavailabilityLabel = connectUnavailabilityLabel(detailServer);
|
||||
if (unavailabilityLabel !== null) {
|
||||
return <span className="text-[13px] text-muted-foreground py-2.5 shrink-0">{unavailabilityLabel}</span>;
|
||||
}
|
||||
if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) {
|
||||
return (
|
||||
<Button
|
||||
variant={isConnected ? "outline" : "default"}
|
||||
disabled={isTogglingOn}
|
||||
onClick={() => handleToggle(detailServer, !isConnected)}
|
||||
className="font-semibold h-[38px] min-w-[110px]"
|
||||
>
|
||||
{isTogglingOn && <Loader2 className="h-4 w-4 animate-spin mr-1.5" />}
|
||||
{isConnected ? "Disconnect" : "Connect"}
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
if (oauthConnected.has(detailServer.server_id)) {
|
||||
return (
|
||||
<Button
|
||||
variant="destructive"
|
||||
onClick={async () => {
|
||||
try {
|
||||
await deleteMCPOAuthUserCredential(accessToken, detailServer.server_id);
|
||||
} catch (_) {
|
||||
// Ignore
|
||||
}
|
||||
setOauthConnected((prev) => {
|
||||
const n = new Set(prev);
|
||||
n.delete(detailServer.server_id);
|
||||
return n;
|
||||
});
|
||||
onChangeRef.current(selectedServersRef.current.filter((s) => s !== name));
|
||||
}}
|
||||
className="font-semibold h-[38px] min-w-[110px]"
|
||||
>
|
||||
Disconnect
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
return (
|
||||
<OAuth2ConnectButton
|
||||
server={detailServer}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => {
|
||||
setOauthConnected((prev) => new Set(prev).add(id));
|
||||
}}
|
||||
variant="button"
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="w-full">
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
onClick={() => setDetailServer(null)}
|
||||
onClick={() => setDetailServerId(null)}
|
||||
className="-ml-3 mb-5 gap-1.5 text-muted-foreground hover:text-foreground"
|
||||
>
|
||||
<ArrowLeft className="h-3 w-3" />
|
||||
|
|
@ -329,48 +418,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
<h2 className="m-0 mb-1 text-[22px] font-bold text-foreground">{name}</h2>
|
||||
<p className="m-0 text-sm text-muted-foreground">{detailServer.description ?? "MCP server"}</p>
|
||||
</div>
|
||||
{detailServer.auth_type === AUTH_TYPE.OAUTH2 ? (
|
||||
oauthConnected.has(detailServer.server_id) ? (
|
||||
<Button
|
||||
variant="destructive"
|
||||
onClick={async () => {
|
||||
try {
|
||||
await deleteMCPOAuthUserCredential(accessToken, detailServer.server_id);
|
||||
} catch (_) {
|
||||
// Ignore
|
||||
}
|
||||
setOauthConnected((prev) => {
|
||||
const n = new Set(prev);
|
||||
n.delete(detailServer.server_id);
|
||||
return n;
|
||||
});
|
||||
onChangeRef.current(selectedServersRef.current.filter((s) => s !== name));
|
||||
}}
|
||||
className="font-semibold h-[38px] min-w-[110px]"
|
||||
>
|
||||
Disconnect
|
||||
</Button>
|
||||
) : (
|
||||
<OAuth2ConnectButton
|
||||
server={detailServer}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => {
|
||||
setOauthConnected((prev) => new Set(prev).add(id));
|
||||
}}
|
||||
variant="button"
|
||||
/>
|
||||
)
|
||||
) : (
|
||||
<Button
|
||||
variant={isConnected ? "outline" : "default"}
|
||||
disabled={isTogglingOn}
|
||||
onClick={() => handleToggle(name, !isConnected, detailServer.server_id)}
|
||||
className="font-semibold h-[38px] min-w-[110px]"
|
||||
>
|
||||
{isTogglingOn && <Loader2 className="h-4 w-4 animate-spin mr-1.5" />}
|
||||
{isConnected ? "Disconnect" : "Connect"}
|
||||
</Button>
|
||||
)}
|
||||
{renderDetailAction()}
|
||||
</div>
|
||||
|
||||
<h3 className="m-0 mb-3 text-[15px] font-semibold text-foreground">Information</h3>
|
||||
|
|
@ -494,13 +542,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
))}
|
||||
</div>
|
||||
) : filtered.length === 0 ? (
|
||||
<div className="text-center text-muted-foreground text-[13px] py-12 px-3">
|
||||
{servers.length === 0
|
||||
? "No MCP servers configured. Add servers in Tools -> MCP Servers."
|
||||
: activeTab === "connected"
|
||||
? "No servers connected yet."
|
||||
: "No servers match your search."}
|
||||
</div>
|
||||
<div className="text-center text-muted-foreground text-[13px] py-12 px-3">{emptyStateText()}</div>
|
||||
) : (
|
||||
<div className="grid grid-cols-2 border rounded-lg overflow-hidden">
|
||||
{filtered.map((server, idx) => {
|
||||
|
|
@ -508,16 +550,16 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
const color = getAvatarColor(name);
|
||||
const isLeftCol = idx % 2 === 0;
|
||||
const count = toolCounts[name];
|
||||
const unsupported = !!connectMode && isUnsupportedOnGatewayConnect(server.auth_type);
|
||||
const unavailable = connectUnavailabilityLabel(server) !== null;
|
||||
|
||||
return (
|
||||
<div
|
||||
key={server.server_id}
|
||||
onClick={() => setDetailServer(server)}
|
||||
onClick={() => setDetailServerId(server.server_id)}
|
||||
className={`flex items-center gap-3 p-4 bg-card cursor-pointer transition-colors hover:bg-accent/30 min-w-0 ${
|
||||
isLeftCol ? "border-r" : ""
|
||||
} ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""} ${
|
||||
unsupported ? "opacity-50" : ""
|
||||
unavailable ? "opacity-50" : ""
|
||||
}`}
|
||||
>
|
||||
{server.mcp_info?.logo_url ? (
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
|
||||
import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils";
|
||||
|
||||
import NotificationsManager from "@/components/molecules/notifications_manager";
|
||||
import EditAutoRouterModal from "./edit_auto_router_modal";
|
||||
|
|
@ -119,3 +119,68 @@ describe("EditAutoRouterModal keyword matching", () => {
|
|||
expect(modelPatchUpdateCall).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("EditAutoRouterModal classifier context window", () => {
|
||||
beforeEach(() => {
|
||||
modelPatchUpdateCall.mockClear();
|
||||
});
|
||||
|
||||
const STORED_LLM_CONFIG = {
|
||||
tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o-mini"], COMPLEX: ["gpt-4o-mini"], REASONING: ["gpt-4o-mini"] },
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000 },
|
||||
classifier_context_window_size: 5,
|
||||
classifier_context_per_turn_chars: 300,
|
||||
};
|
||||
|
||||
const renderLlmModal = () =>
|
||||
renderWithProviders(
|
||||
<EditAutoRouterModal
|
||||
isVisible
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
modelData={{
|
||||
...MODEL_DATA,
|
||||
litellm_params: { ...MODEL_DATA.litellm_params, complexity_router_config: STORED_LLM_CONFIG },
|
||||
}}
|
||||
accessToken="token"
|
||||
userRole="Admin"
|
||||
/>,
|
||||
);
|
||||
|
||||
// Hydration bugs are invisible to the payload-builder unit tests, which only exercise
|
||||
// buildUpdatedComplexityRouterConfig with a form value the caller already assembled by hand.
|
||||
// Only driving the real component through open, then save with nothing touched, catches a
|
||||
// missing initializeForm hydration line.
|
||||
it("shows the stored classifier context values and preserves them through an untouched open-and-save", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderLlmModal();
|
||||
|
||||
await user.click(await screen.findByText("Advanced: Classification Method"));
|
||||
await screen.findByText("Context Window Size");
|
||||
expect(screen.getByDisplayValue("5")).toBeInTheDocument();
|
||||
expect(screen.getByDisplayValue("300")).toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
const config = savedConfig();
|
||||
expect(config.classifier_context_window_size).toBe(5);
|
||||
expect(config.classifier_context_per_turn_chars).toBe(300);
|
||||
});
|
||||
|
||||
it("persists an edited classifier context window size", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderLlmModal();
|
||||
|
||||
await user.click(await screen.findByText("Advanced: Classification Method"));
|
||||
const windowSizeSection = (await screen.findByText("Context Window Size")).closest("div") as HTMLElement;
|
||||
const input = within(windowSizeSection).getByRole("spinbutton");
|
||||
fireEvent.change(input, { target: { value: "8" } });
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedConfig().classifier_context_window_size).toBe(8);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -33,6 +33,8 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
|
|||
"tiers",
|
||||
"classifier_type",
|
||||
"classifier_llm_config",
|
||||
"classifier_context_window_size",
|
||||
"classifier_context_per_turn_chars",
|
||||
"adaptive",
|
||||
"adaptive_weights",
|
||||
"tier_distance_penalty",
|
||||
|
|
@ -86,6 +88,14 @@ export const buildUpdatedComplexityRouterConfig = (
|
|||
tiers: value.tiers,
|
||||
classifier_type: value.classifier_type,
|
||||
...(value.classifier_type === "llm" ? { classifier_llm_config: value.classifier_llm_config } : {}),
|
||||
...(value.classifier_type === "llm" &&
|
||||
value.classifier_context_window_size !== undefined && {
|
||||
classifier_context_window_size: value.classifier_context_window_size,
|
||||
}),
|
||||
...(value.classifier_type === "llm" &&
|
||||
value.classifier_context_per_turn_chars !== undefined && {
|
||||
classifier_context_per_turn_chars: value.classifier_context_per_turn_chars,
|
||||
}),
|
||||
...(customTechnicalKeywords &&
|
||||
customTechnicalKeywords.length > 0 && {
|
||||
custom_technical_keywords: customTechnicalKeywords,
|
||||
|
|
@ -182,7 +192,7 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
parsedConfig = JSON.parse(parsedConfig);
|
||||
}
|
||||
|
||||
setComplexityRouterConfig({
|
||||
const hydratedComplexityRouterConfig: ComplexityRouterConfigValue = {
|
||||
tiers: {
|
||||
SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE),
|
||||
MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM),
|
||||
|
|
@ -191,12 +201,21 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
},
|
||||
classifier_type: parsedConfig.classifier_type || "heuristic",
|
||||
classifier_llm_config: parsedConfig.classifier_llm_config,
|
||||
classifier_context_window_size:
|
||||
typeof parsedConfig.classifier_context_window_size === "number"
|
||||
? parsedConfig.classifier_context_window_size
|
||||
: undefined,
|
||||
classifier_context_per_turn_chars:
|
||||
typeof parsedConfig.classifier_context_per_turn_chars === "number"
|
||||
? parsedConfig.classifier_context_per_turn_chars
|
||||
: undefined,
|
||||
adaptive: parsedConfig.adaptive || false,
|
||||
adaptive_weights: parsedConfig.adaptive_weights,
|
||||
tier_distance_penalty: parsedConfig.tier_distance_penalty,
|
||||
adaptive_eligible: parsedConfig.adaptive_eligible || "all",
|
||||
return_raw_model_name: parsedConfig.return_raw_model_name || false,
|
||||
});
|
||||
};
|
||||
setComplexityRouterConfig(hydratedComplexityRouterConfig);
|
||||
setCustomTechnicalKeywords(
|
||||
Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [],
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
67
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
67
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -9379,7 +9379,10 @@ export interface paths {
|
|||
};
|
||||
/**
|
||||
* List Policy Attachments
|
||||
* @description List all policy attachments from the database.
|
||||
* @description List all policy attachments from the database and config.yaml.
|
||||
*
|
||||
* Config-defined attachments are returned with definition_location "config" and a
|
||||
* synthetic attachment_id ("config-<index>").
|
||||
*
|
||||
* Example Request:
|
||||
* ```bash
|
||||
|
|
@ -9487,7 +9490,10 @@ export interface paths {
|
|||
};
|
||||
/**
|
||||
* List Policies
|
||||
* @description List all policies from the database. Optionally filter by version_status.
|
||||
* @description List all policies from the database and config.yaml. Optionally filter by version_status.
|
||||
*
|
||||
* Config-defined policies are returned with definition_location "config" and are treated
|
||||
* as production versions. On a name conflict with a DB policy, only the DB policy is returned.
|
||||
*
|
||||
* Query params:
|
||||
* - version_status: Optional. One of "draft", "published", "production".
|
||||
|
|
@ -25845,6 +25851,16 @@ export interface components {
|
|||
cache_creation_input_token_cost_above_1hr?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 200K Tokens */
|
||||
cache_creation_input_token_cost_above_200k_tokens?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 272K Tokens */
|
||||
cache_creation_input_token_cost_above_272k_tokens?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 272K Tokens Flex */
|
||||
cache_creation_input_token_cost_above_272k_tokens_flex?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 272K Tokens Priority */
|
||||
cache_creation_input_token_cost_above_272k_tokens_priority?: number | null;
|
||||
/** Cache Creation Input Token Cost Flex */
|
||||
cache_creation_input_token_cost_flex?: number | null;
|
||||
/** Cache Creation Input Token Cost Priority */
|
||||
cache_creation_input_token_cost_priority?: number | null;
|
||||
/** Cache Read Input Audio Token Cost */
|
||||
cache_read_input_audio_token_cost?: number | null;
|
||||
/** Cache Read Input Token Cost */
|
||||
|
|
@ -25855,6 +25871,8 @@ export interface components {
|
|||
cache_read_input_token_cost_above_200k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens */
|
||||
cache_read_input_token_cost_above_272k_tokens?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens Flex */
|
||||
cache_read_input_token_cost_above_272k_tokens_flex?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens Priority */
|
||||
cache_read_input_token_cost_above_272k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Above 512K Tokens */
|
||||
|
|
@ -25913,6 +25931,8 @@ export interface components {
|
|||
input_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens */
|
||||
input_cost_per_token_above_272k_tokens?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens Flex */
|
||||
input_cost_per_token_above_272k_tokens_flex?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens Priority */
|
||||
input_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Above 512K Tokens */
|
||||
|
|
@ -26004,6 +26024,8 @@ export interface components {
|
|||
output_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens */
|
||||
output_cost_per_token_above_272k_tokens?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens Flex */
|
||||
output_cost_per_token_above_272k_tokens_flex?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens Priority */
|
||||
output_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Above 512K Tokens */
|
||||
|
|
@ -29376,6 +29398,13 @@ export interface components {
|
|||
* @description Who created the attachment.
|
||||
*/
|
||||
created_by?: string | null;
|
||||
/**
|
||||
* Definition Location
|
||||
* @description Where this attachment is defined: 'db' (database) or 'config' (config.yaml).
|
||||
* @default db
|
||||
* @enum {string}
|
||||
*/
|
||||
definition_location: "db" | "config";
|
||||
/**
|
||||
* Keys
|
||||
* @description Key patterns.
|
||||
|
|
@ -29507,6 +29536,13 @@ export interface components {
|
|||
* @description Who created the policy.
|
||||
*/
|
||||
created_by?: string | null;
|
||||
/**
|
||||
* Definition Location
|
||||
* @description Where this policy is defined: 'db' (database) or 'config' (config.yaml).
|
||||
* @default db
|
||||
* @enum {string}
|
||||
*/
|
||||
definition_location: "db" | "config";
|
||||
/**
|
||||
* Description
|
||||
* @description Policy description.
|
||||
|
|
@ -33135,6 +33171,17 @@ export interface components {
|
|||
*/
|
||||
model?: string | null;
|
||||
};
|
||||
/** UsageChartPoint */
|
||||
UsageChartPoint: {
|
||||
/** Blocked */
|
||||
blocked: number;
|
||||
/** Date */
|
||||
date: string;
|
||||
/** Passed */
|
||||
passed: number;
|
||||
/** Score */
|
||||
score?: number | null;
|
||||
};
|
||||
/** UsageDetailResponse */
|
||||
UsageDetailResponse: {
|
||||
/** Avglatency */
|
||||
|
|
@ -33957,6 +34004,16 @@ export interface components {
|
|||
cache_creation_input_token_cost_above_1hr?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 200K Tokens */
|
||||
cache_creation_input_token_cost_above_200k_tokens?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 272K Tokens */
|
||||
cache_creation_input_token_cost_above_272k_tokens?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 272K Tokens Flex */
|
||||
cache_creation_input_token_cost_above_272k_tokens_flex?: number | null;
|
||||
/** Cache Creation Input Token Cost Above 272K Tokens Priority */
|
||||
cache_creation_input_token_cost_above_272k_tokens_priority?: number | null;
|
||||
/** Cache Creation Input Token Cost Flex */
|
||||
cache_creation_input_token_cost_flex?: number | null;
|
||||
/** Cache Creation Input Token Cost Priority */
|
||||
cache_creation_input_token_cost_priority?: number | null;
|
||||
/** Cache Read Input Audio Token Cost */
|
||||
cache_read_input_audio_token_cost?: number | null;
|
||||
/** Cache Read Input Token Cost */
|
||||
|
|
@ -33967,6 +34024,8 @@ export interface components {
|
|||
cache_read_input_token_cost_above_200k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens */
|
||||
cache_read_input_token_cost_above_272k_tokens?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens Flex */
|
||||
cache_read_input_token_cost_above_272k_tokens_flex?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens Priority */
|
||||
cache_read_input_token_cost_above_272k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Above 512K Tokens */
|
||||
|
|
@ -34025,6 +34084,8 @@ export interface components {
|
|||
input_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens */
|
||||
input_cost_per_token_above_272k_tokens?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens Flex */
|
||||
input_cost_per_token_above_272k_tokens_flex?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens Priority */
|
||||
input_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Above 512K Tokens */
|
||||
|
|
@ -34116,6 +34177,8 @@ export interface components {
|
|||
output_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens */
|
||||
output_cost_per_token_above_272k_tokens?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens Flex */
|
||||
output_cost_per_token_above_272k_tokens_flex?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens Priority */
|
||||
output_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Above 512K Tokens */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue