mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge pull request #30502 from BerriAI/litellm_backport_1_89_x_bp_1_89x_multi
chore(release): backport 1.84.8 patch set + MCP/model-info/DB fixes to stable/1.89.x and cut 1.89.1
This commit is contained in:
commit
39832310a7
75 changed files with 5684 additions and 598 deletions
|
|
@ -92,12 +92,22 @@ class DataDogLogger(
|
|||
# Class variables or attributes
|
||||
def __init__(
|
||||
self,
|
||||
dd_api_key: Optional[str] = None,
|
||||
dd_site: Optional[str] = None,
|
||||
dd_agent_host: Optional[str] = None,
|
||||
dd_agent_port: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initializes the datadog logger, checks if the correct env variables are set
|
||||
|
||||
Required environment variables (Direct API):
|
||||
Args:
|
||||
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var.
|
||||
dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var.
|
||||
dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var.
|
||||
dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var.
|
||||
|
||||
Required environment variables (Direct API) when kwargs not provided:
|
||||
`DD_API_KEY` - your datadog api key
|
||||
`DD_SITE` - your datadog site, example = `"us5.datadoghq.com"`
|
||||
|
||||
|
|
@ -130,12 +140,19 @@ class DataDogLogger(
|
|||
)
|
||||
|
||||
# Configure DataDog endpoint (Agent or Direct API)
|
||||
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
|
||||
dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")
|
||||
if dd_agent_host:
|
||||
self._configure_dd_agent(dd_agent_host=dd_agent_host)
|
||||
# Prefer explicit kwargs, then fall back to env vars
|
||||
resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST")
|
||||
if resolved_agent_host:
|
||||
self._configure_dd_agent(
|
||||
dd_agent_host=resolved_agent_host,
|
||||
dd_agent_port=dd_agent_port,
|
||||
dd_api_key=dd_api_key,
|
||||
)
|
||||
else:
|
||||
self._configure_dd_direct_api()
|
||||
self._configure_dd_direct_api(
|
||||
dd_api_key=dd_api_key,
|
||||
dd_site=dd_site,
|
||||
)
|
||||
|
||||
# Optional override for testing
|
||||
dd_base_url = get_datadog_base_url_from_env()
|
||||
|
|
@ -172,34 +189,54 @@ class DataDogLogger(
|
|||
).model_dump()
|
||||
return dict_datadog_params
|
||||
|
||||
def _configure_dd_agent(self, dd_agent_host: str) -> None:
|
||||
def _configure_dd_agent(
|
||||
self,
|
||||
dd_agent_host: str,
|
||||
dd_agent_port: Optional[str] = None,
|
||||
dd_api_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Configure DataDog Agent for log forwarding
|
||||
|
||||
Args:
|
||||
dd_agent_host: Hostname or IP of DataDog agent
|
||||
dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518).
|
||||
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. Optional when using agent.
|
||||
"""
|
||||
dd_agent_port = os.getenv(
|
||||
resolved_port = dd_agent_port or os.getenv(
|
||||
"LITELLM_DD_AGENT_PORT", "10518"
|
||||
) # default port for logs
|
||||
self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs"
|
||||
self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent
|
||||
self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs"
|
||||
self.DD_API_KEY = dd_api_key or os.getenv(
|
||||
"DD_API_KEY"
|
||||
) # Optional when using agent
|
||||
verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}")
|
||||
|
||||
def _configure_dd_direct_api(self) -> None:
|
||||
def _configure_dd_direct_api(
|
||||
self,
|
||||
dd_api_key: Optional[str] = None,
|
||||
dd_site: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Configure direct DataDog API connection
|
||||
|
||||
Args:
|
||||
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var.
|
||||
dd_site: Datadog site. Falls back to DD_SITE env var.
|
||||
|
||||
Raises:
|
||||
Exception: If required environment variables are not set
|
||||
Exception: If required credentials are not provided via args or env vars
|
||||
"""
|
||||
if os.getenv("DD_API_KEY", None) is None:
|
||||
resolved_api_key = dd_api_key or os.getenv("DD_API_KEY")
|
||||
resolved_site = dd_site or os.getenv("DD_SITE")
|
||||
|
||||
if resolved_api_key is None:
|
||||
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>")
|
||||
if os.getenv("DD_SITE", None) is None:
|
||||
if resolved_site is None:
|
||||
raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>")
|
||||
|
||||
self.DD_API_KEY = os.getenv("DD_API_KEY")
|
||||
self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs"
|
||||
self.DD_API_KEY = resolved_api_key
|
||||
self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs"
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
|
|
|
|||
117
litellm/integrations/datadog/datadog_team_handler.py
Normal file
117
litellm/integrations/datadog/datadog_team_handler.py
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
"""
|
||||
DataDog Team Handler
|
||||
|
||||
Used to get the DataDogLogger for a given request.
|
||||
Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
|
||||
|
||||
from .datadog import DataDogLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
|
||||
else:
|
||||
DynamicLoggingCache = Any
|
||||
|
||||
|
||||
class DatadogLoggingConfig(TypedDict):
|
||||
dd_api_key: Optional[str]
|
||||
dd_site: Optional[str]
|
||||
dd_agent_host: Optional[str]
|
||||
dd_agent_port: Optional[str]
|
||||
|
||||
|
||||
class DataDogHandler:
|
||||
@staticmethod
|
||||
def get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
in_memory_dynamic_logger_cache: DynamicLoggingCache,
|
||||
) -> DataDogLogger:
|
||||
"""
|
||||
Get a team-scoped DataDogLogger for a given request.
|
||||
|
||||
Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache,
|
||||
keyed by the team's DD credentials. Each unique set of credentials gets its own
|
||||
logger instance with its own batch/flush loop.
|
||||
|
||||
Note: This handler is only called when team-scoped DD credentials are present.
|
||||
The global (env-var based) DataDogLogger is managed separately by
|
||||
_init_custom_logger_compatible_class via _in_memory_loggers.
|
||||
"""
|
||||
_credentials = DataDogHandler.get_dynamic_datadog_logging_config(
|
||||
standard_callback_dynamic_params=standard_callback_dynamic_params,
|
||||
)
|
||||
credentials_dict = dict(_credentials)
|
||||
|
||||
# check if datadog logger is already cached
|
||||
temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache(
|
||||
credentials=credentials_dict, service_name="datadog"
|
||||
)
|
||||
|
||||
# if not cached, create a new datadog logger and cache it
|
||||
if temp_datadog_logger is None:
|
||||
temp_datadog_logger = (
|
||||
DataDogHandler._create_datadog_logger_from_credentials(
|
||||
credentials=credentials_dict,
|
||||
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||
)
|
||||
)
|
||||
|
||||
return temp_datadog_logger
|
||||
|
||||
@staticmethod
|
||||
def _create_datadog_logger_from_credentials(
|
||||
credentials: Dict,
|
||||
in_memory_dynamic_logger_cache: DynamicLoggingCache,
|
||||
) -> DataDogLogger:
|
||||
"""
|
||||
Create a DataDogLogger from the credentials and cache it.
|
||||
"""
|
||||
datadog_logger = DataDogLogger(
|
||||
dd_api_key=credentials.get("dd_api_key"),
|
||||
dd_site=credentials.get("dd_site"),
|
||||
dd_agent_host=credentials.get("dd_agent_host"),
|
||||
dd_agent_port=credentials.get("dd_agent_port"),
|
||||
)
|
||||
in_memory_dynamic_logger_cache.set_cache(
|
||||
credentials=credentials,
|
||||
service_name="datadog",
|
||||
logging_obj=datadog_logger,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Datadog: Created and cached new DataDogLogger for team-scoped credentials"
|
||||
)
|
||||
return datadog_logger
|
||||
|
||||
@staticmethod
|
||||
def get_dynamic_datadog_logging_config(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> DatadogLoggingConfig:
|
||||
"""
|
||||
Get the Datadog logging config for a given request from dynamic params.
|
||||
"""
|
||||
return DatadogLoggingConfig(
|
||||
dd_api_key=standard_callback_dynamic_params.get("dd_api_key"),
|
||||
dd_site=standard_callback_dynamic_params.get("dd_site"),
|
||||
dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"),
|
||||
dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _dynamic_datadog_credentials_are_passed(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params.
|
||||
"""
|
||||
if (
|
||||
standard_callback_dynamic_params.get("dd_api_key") is not None
|
||||
or standard_callback_dynamic_params.get("dd_site") is not None
|
||||
or standard_callback_dynamic_params.get("dd_agent_host") is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
@ -53,11 +53,19 @@ _supported_callback_params = [
|
|||
"braintrust_host",
|
||||
"slack_webhook_url",
|
||||
"lunary_public_key",
|
||||
"dd_api_key",
|
||||
"dd_site",
|
||||
"dd_agent_host",
|
||||
"dd_agent_port",
|
||||
]
|
||||
|
||||
_request_blocked_callback_params = {
|
||||
"gcs_bucket_name",
|
||||
"gcs_path_service_account",
|
||||
"dd_api_key",
|
||||
"dd_site",
|
||||
"dd_agent_host",
|
||||
"dd_agent_port",
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -376,13 +376,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
List[Union[str, Callable, CustomLogger]]
|
||||
] = dynamic_async_failure_callbacks
|
||||
|
||||
# Process dynamic callbacks
|
||||
self.process_dynamic_callbacks()
|
||||
|
||||
## DYNAMIC LANGFUSE / GCS / logging callback KEYS ##
|
||||
self.standard_callback_dynamic_params: StandardCallbackDynamicParams = (
|
||||
self.initialize_standard_callback_dynamic_params(kwargs)
|
||||
)
|
||||
|
||||
# Process dynamic callbacks (after standard_callback_dynamic_params is initialized,
|
||||
# so team-scoped credentials are available for callback initialization)
|
||||
self.process_dynamic_callbacks()
|
||||
self.standard_built_in_tools_params: StandardBuiltInToolsParams = (
|
||||
self.initialize_standard_built_in_tools_params(kwargs)
|
||||
)
|
||||
|
|
@ -477,8 +478,21 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
isinstance(callback, str)
|
||||
and callback in litellm._known_custom_logger_compatible_callbacks
|
||||
):
|
||||
# For callbacks that support team-scoped credentials (e.g. datadog),
|
||||
# pass only the relevant dynamic params as custom_logger_init_args.
|
||||
_custom_logger_init_args: Optional[dict] = None
|
||||
if callback == "datadog":
|
||||
_custom_logger_init_args = {
|
||||
k: v
|
||||
for k, v in self.standard_callback_dynamic_params.items()
|
||||
if k.startswith("dd_")
|
||||
}
|
||||
|
||||
callback_class = _init_custom_logger_compatible_class(
|
||||
callback, internal_usage_cache=None, llm_router=None # type: ignore
|
||||
callback, # type: ignore[arg-type]
|
||||
internal_usage_cache=None,
|
||||
llm_router=None, # type: ignore
|
||||
custom_logger_init_args=_custom_logger_init_args,
|
||||
)
|
||||
if callback_class is not None:
|
||||
processed_list.append(callback_class)
|
||||
|
|
@ -3890,6 +3904,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
_in_memory_loggers.append(_prometheus_logger)
|
||||
return _prometheus_logger # type: ignore
|
||||
elif logging_integration == "datadog":
|
||||
# Check if team-scoped credentials are provided
|
||||
_dd_api_key = custom_logger_init_args.get("dd_api_key")
|
||||
_dd_site = custom_logger_init_args.get("dd_site")
|
||||
_dd_agent_host = custom_logger_init_args.get("dd_agent_host")
|
||||
_dd_agent_port = custom_logger_init_args.get("dd_agent_port")
|
||||
|
||||
if _dd_api_key or _dd_site or _dd_agent_host:
|
||||
# Team-scoped credentials: use DynamicLoggingCache for per-credential isolation
|
||||
from litellm.integrations.datadog.datadog_team_handler import (
|
||||
DataDogHandler,
|
||||
)
|
||||
|
||||
return DataDogHandler.get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params=custom_logger_init_args, # type: ignore
|
||||
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||
)
|
||||
|
||||
# Global (env-var based): reuse cached instance
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, DataDogLogger):
|
||||
return callback # type: ignore
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple
|
|||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
|
||||
from litellm.types.llms.openai import (
|
||||
FileSearchTool,
|
||||
ResponsesAPIResponse,
|
||||
|
|
@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking:
|
|||
# and _handle_web_search_cost() is never called.
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and usage.server_tool_use is not None
|
||||
and usage.server_tool_use.web_search_requests is not None
|
||||
and _get_web_search_requests(usage.server_tool_use) is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking:
|
|||
elif usage is not None:
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and usage.server_tool_use is not None
|
||||
and usage.server_tool_use.web_search_requests is not None
|
||||
and _get_web_search_requests(usage.server_tool_use) is not None
|
||||
):
|
||||
return True
|
||||
elif (
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# What is this?
|
||||
## Helper utilities for cost_per_token()
|
||||
|
||||
from typing import Literal, Optional, Tuple, TypedDict, cast
|
||||
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]:
|
|||
return value if isinstance(value, int) else None
|
||||
|
||||
|
||||
def _get_web_search_requests(server_tool_use: Any) -> Optional[int]:
|
||||
"""
|
||||
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
|
||||
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
|
||||
or any other object supporting attribute access.
|
||||
|
||||
Returns ``None`` when the value cannot be resolved — callers can
|
||||
distinguish "absent" from "zero" using ``is None``.
|
||||
|
||||
See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder``
|
||||
historically left this as a plain ``dict``, which broke direct attribute
|
||||
access in cost calculation.
|
||||
"""
|
||||
if server_tool_use is None:
|
||||
return None
|
||||
if isinstance(server_tool_use, dict):
|
||||
return server_tool_use.get("web_search_requests")
|
||||
return getattr(server_tool_use, "web_search_requests", None)
|
||||
|
||||
|
||||
def _is_above_128k(tokens: float) -> bool:
|
||||
if tokens > 128000:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -588,7 +588,18 @@ class ChunkProcessor:
|
|||
hasattr(usage_chunk, "server_tool_use")
|
||||
and usage_chunk.server_tool_use is not None
|
||||
):
|
||||
server_tool_use = usage_chunk.server_tool_use
|
||||
# Coerce dict to ServerToolUse so downstream cost-calc code
|
||||
# (which accesses .web_search_requests as an attribute)
|
||||
# doesn't raise AttributeError. Some providers / streaming
|
||||
# paths leave server_tool_use as a plain dict on the chunk.
|
||||
if isinstance(usage_chunk.server_tool_use, dict):
|
||||
server_tool_use = ServerToolUse(**usage_chunk.server_tool_use)
|
||||
elif isinstance(usage_chunk.server_tool_use, ServerToolUse):
|
||||
server_tool_use = usage_chunk.server_tool_use
|
||||
else:
|
||||
server_tool_use = ServerToolUse.model_validate(
|
||||
usage_chunk.server_tool_use
|
||||
)
|
||||
if (
|
||||
usage_chunk_dict["prompt_tokens_details"] is not None
|
||||
and getattr(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple
|
|||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
_get_token_base_cost,
|
||||
_get_web_search_requests,
|
||||
_parse_prompt_tokens_details,
|
||||
calculate_cache_writing_cost,
|
||||
generic_cost_per_token,
|
||||
|
|
@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search(
|
|||
if model_info is None:
|
||||
return 0.0
|
||||
|
||||
if (
|
||||
usage is None
|
||||
or usage.server_tool_use is None
|
||||
or usage.server_tool_use.web_search_requests is None
|
||||
):
|
||||
if usage is None:
|
||||
return 0.0
|
||||
web_search_requests = _get_web_search_requests(
|
||||
getattr(usage, "server_tool_use", None)
|
||||
)
|
||||
if web_search_requests is None:
|
||||
return 0.0
|
||||
|
||||
## Get the cost per web search request
|
||||
|
|
@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search(
|
|||
return 0.0
|
||||
|
||||
## Calculate the total cost
|
||||
total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests
|
||||
total_cost = cost_per_web_search_request * web_search_requests
|
||||
return total_cost
|
||||
|
|
|
|||
|
|
@ -62,6 +62,26 @@ class BaseResponsesAPIConfig(ABC):
|
|||
"""
|
||||
return False
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
optional_params: dict,
|
||||
request_data: dict,
|
||||
api_base: str,
|
||||
api_key: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
stream: Optional[bool] = None,
|
||||
fake_stream: Optional[bool] = None,
|
||||
) -> Tuple[dict, Optional[bytes]]:
|
||||
"""Sign the request after the body is finalized.
|
||||
|
||||
Default is a no-op (returns headers unchanged, no signed body). Providers
|
||||
whose endpoint requires request signing (e.g. Bedrock Mantle SigV4)
|
||||
override this and return the signed body bytes so the handler sends those
|
||||
exact bytes.
|
||||
"""
|
||||
return headers, None
|
||||
|
||||
@abstractmethod
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -4,14 +4,26 @@ Amazon Bedrock Mantle - Responses API backend.
|
|||
gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses`
|
||||
path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI
|
||||
Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
|
||||
only the endpoint URL and Bearer authentication.
|
||||
only the endpoint URL and authentication.
|
||||
|
||||
Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
|
||||
standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
|
||||
Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard
|
||||
AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise
|
||||
AWS SigV4 (service name "bedrock") using the standard credential chain (IAM
|
||||
role / access key / profile / web identity), signed via the shared
|
||||
BaseAWSLLM._sign_request after the request body is finalized.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
import re
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from botocore.exceptions import (
|
||||
CredentialRetrievalError,
|
||||
NoCredentialsError,
|
||||
PartialCredentialsError,
|
||||
ProfileNotFound,
|
||||
)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
|
@ -29,22 +41,44 @@ _BASE_SUFFIXES_TO_STRIP = (
|
|||
"/v1",
|
||||
)
|
||||
|
||||
# Standard Mantle host: https://bedrock-mantle.<region>.api.aws (group 1 = region).
|
||||
_MANTLE_HOST_RE = re.compile(
|
||||
r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE
|
||||
)
|
||||
|
||||
|
||||
class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
||||
def __init__(self, aws_signer: Optional[BaseAWSLLM] = None):
|
||||
super().__init__()
|
||||
self._aws_signer = aws_signer or BaseAWSLLM()
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.BEDROCK_MANTLE
|
||||
|
||||
@staticmethod
|
||||
def _resolve_region(params: dict) -> str:
|
||||
region = params.get("aws_region_name")
|
||||
if region:
|
||||
return region
|
||||
base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
if base:
|
||||
match = _MANTLE_HOST_RE.match(base.rstrip("/"))
|
||||
if match:
|
||||
return match.group(1)
|
||||
return (
|
||||
get_secret_str("BEDROCK_MANTLE_REGION")
|
||||
or get_secret_str("AWS_REGION_NAME")
|
||||
or get_secret_str("AWS_REGION")
|
||||
or BEDROCK_MANTLE_DEFAULT_REGION
|
||||
)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
litellm_params: dict,
|
||||
) -> str:
|
||||
region = (
|
||||
get_secret_str("BEDROCK_MANTLE_REGION")
|
||||
or get_secret_str("AWS_REGION")
|
||||
or BEDROCK_MANTLE_DEFAULT_REGION
|
||||
)
|
||||
region = self._resolve_region({**litellm_params, "api_base": api_base})
|
||||
base = (
|
||||
api_base
|
||||
or get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
|
|
@ -55,6 +89,11 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
if base.endswith(suffix):
|
||||
base = base[: -len(suffix)]
|
||||
break
|
||||
# For the standard Mantle host (including the default-region base that
|
||||
# responses/main.py auto-injects into litellm_params.api_base), pin to the
|
||||
# single resolved region so aws_region_name wins; preserve custom proxy hosts.
|
||||
if _MANTLE_HOST_RE.match(base):
|
||||
base = f"https://bedrock-mantle.{region}.api.aws"
|
||||
return f"{base}/openai/v1/responses"
|
||||
|
||||
def validate_environment(
|
||||
|
|
@ -66,12 +105,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
|
||||
"(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
def supports_native_file_search(self) -> bool:
|
||||
|
|
@ -79,3 +114,58 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
|
||||
def supports_native_websocket(self) -> bool:
|
||||
return False
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
optional_params: dict,
|
||||
request_data: dict,
|
||||
api_base: str,
|
||||
api_key: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
stream: Optional[bool] = None,
|
||||
fake_stream: Optional[bool] = None,
|
||||
) -> Tuple[dict, Optional[bytes]]:
|
||||
bearer = (
|
||||
api_key
|
||||
or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
)
|
||||
if not bearer:
|
||||
# SigV4 path. Pin the credential-scope region to the region of the actual
|
||||
# signing URL (api_base, already region-resolved by get_complete_url) so the
|
||||
# SigV4 scope and the URL host can never disagree. Resolve from api_base first,
|
||||
# then fall back to the regular precedence. Also drop any caller Authorization
|
||||
# so _sign_request's restore-original-Authorization step cannot override the
|
||||
# SigV4 header.
|
||||
optional_params = {
|
||||
**optional_params,
|
||||
"aws_region_name": self._resolve_region(
|
||||
{**optional_params, "api_base": api_base}
|
||||
),
|
||||
}
|
||||
headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
|
||||
try:
|
||||
return self._aws_signer._sign_request(
|
||||
service_name="bedrock",
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
request_data=request_data,
|
||||
api_base=api_base,
|
||||
api_key=bearer,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
except (
|
||||
NoCredentialsError,
|
||||
PartialCredentialsError,
|
||||
ProfileNotFound,
|
||||
CredentialRetrievalError,
|
||||
) as e:
|
||||
raise ValueError(
|
||||
"Bedrock Mantle auth failed: no Bearer token and no usable AWS "
|
||||
"credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) "
|
||||
"or pass api_key for Bearer auth, or provide AWS credentials "
|
||||
"(IAM role / access key / profile / web identity) for SigV4."
|
||||
) from e
|
||||
|
|
|
|||
|
|
@ -2318,6 +2318,31 @@ class BaseLLMHTTPHandler:
|
|||
# but never included in the outbound provider payload.
|
||||
request_context["litellm_params"] = dict(litellm_params)
|
||||
|
||||
is_stream_request = bool(stream)
|
||||
if is_stream_request and fake_stream is True:
|
||||
stream, data = self._prepare_fake_stream_request(
|
||||
stream=stream,
|
||||
data=data,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
# Sign after the body is final (post-transform/normalize/extra_body and post
|
||||
# fake-stream prep) so signed bytes match what we send. No-op for providers
|
||||
# that inherit the default sign_request.
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
body_kwargs: Dict[str, Any] = (
|
||||
{"data": signed_body} if signed_body is not None else {"json": data}
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
|
|
@ -2330,22 +2355,14 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
try:
|
||||
if stream:
|
||||
# For streaming, use stream=True in the request
|
||||
if fake_stream is True:
|
||||
stream, data = self._prepare_fake_stream_request(
|
||||
stream=stream,
|
||||
data=data,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
if is_stream_request:
|
||||
response = sync_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
stream=stream,
|
||||
**body_kwargs,
|
||||
)
|
||||
if fake_stream is True:
|
||||
return MockResponsesAPIStreamingIterator(
|
||||
|
|
@ -2370,13 +2387,12 @@ class BaseLLMHTTPHandler:
|
|||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
else:
|
||||
# For non-streaming requests
|
||||
response = sync_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
**body_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
raise self._handle_error(
|
||||
|
|
@ -2464,6 +2480,28 @@ class BaseLLMHTTPHandler:
|
|||
# but never included in the outbound provider payload.
|
||||
request_context["litellm_params"] = dict(litellm_params)
|
||||
|
||||
is_stream_request = bool(stream)
|
||||
if is_stream_request and fake_stream is True:
|
||||
stream, data = self._prepare_fake_stream_request(
|
||||
stream=stream,
|
||||
data=data,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
body_kwargs: Dict[str, Any] = (
|
||||
{"data": signed_body} if signed_body is not None else {"json": data}
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
|
|
@ -2476,22 +2514,14 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
try:
|
||||
if stream:
|
||||
# For streaming, we need to use stream=True in the request
|
||||
if fake_stream is True:
|
||||
stream, data = self._prepare_fake_stream_request(
|
||||
stream=stream,
|
||||
data=data,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
if is_stream_request:
|
||||
response = await async_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
stream=stream,
|
||||
**body_kwargs,
|
||||
)
|
||||
|
||||
if fake_stream is True:
|
||||
|
|
@ -2518,13 +2548,12 @@ class BaseLLMHTTPHandler:
|
|||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
else:
|
||||
# For non-streaming, proceed as before
|
||||
response = await async_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
**body_kwargs,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -4005,6 +4034,18 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=url,
|
||||
api_key=litellm_params.api_key,
|
||||
model=model,
|
||||
)
|
||||
body_kwargs: Dict[str, Any] = (
|
||||
{"data": signed_body} if signed_body is not None else {"json": data}
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
|
|
@ -4018,7 +4059,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
try:
|
||||
response = sync_httpx_client.post(
|
||||
url=url, headers=headers, json=data, timeout=timeout
|
||||
url=url, headers=headers, timeout=timeout, **body_kwargs
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -4088,6 +4129,18 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=url,
|
||||
api_key=litellm_params.api_key,
|
||||
model=model,
|
||||
)
|
||||
body_kwargs: Dict[str, Any] = (
|
||||
{"data": signed_body} if signed_body is not None else {"json": data}
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
|
|
@ -4101,7 +4154,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
try:
|
||||
response = await async_httpx_client.post(
|
||||
url=url, headers=headers, json=data, timeout=timeout
|
||||
url=url, headers=headers, timeout=timeout, **body_kwargs
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -63,9 +63,10 @@ def _is_mcp_passthrough_cold_start(
|
|||
spec-compliant WWW-Authenticate challenge instead of surfacing a generic
|
||||
admission error.
|
||||
|
||||
Uses "all" semantics (mirrors :meth:`MCPRequestHandler._target_servers_use_oauth2`):
|
||||
one non-passthrough target in a co-targeted set must not flip the bypass
|
||||
open for the others. Fails closed when any target cannot be resolved."""
|
||||
Uses "all" semantics (mirrors
|
||||
:meth:`MCPRequestHandler._target_servers_delegate_auth_to_upstream`): one
|
||||
non-passthrough target in a co-targeted set must not flip the bypass open
|
||||
for the others. Fails closed when any target cannot be resolved."""
|
||||
if not mcp_servers:
|
||||
return False
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
|
|
@ -210,101 +211,64 @@ class MCPRequestHandler:
|
|||
# Only OAuth metadata routes registered under /.well-known/ are public.
|
||||
if request_route.startswith("/.well-known/"):
|
||||
validated_user_api_key_auth = UserAPIKeyAuth()
|
||||
elif (
|
||||
not litellm_api_key
|
||||
and MCPRequestHandler._target_servers_delegate_auth_to_upstream( # noqa: E501
|
||||
path=request_route,
|
||||
mcp_servers=mcp_servers,
|
||||
client_ip=IPAddressUtils.get_mcp_client_ip(request),
|
||||
)
|
||||
):
|
||||
# Operator opted this oauth2 server into upstream-delegated auth
|
||||
# (PKCE passthrough): skip LiteLLM API-key/SSO entirely so the
|
||||
# client authenticates directly with the upstream MCP server.
|
||||
# Fires ONLY when neither x-litellm-api-key nor Authorization is
|
||||
# present. If any LiteLLM key is supplied (primary or secondary
|
||||
# header), we fall through so user_id is resolved, spend/rate
|
||||
# limiting apply, and any stored OAuth token can be retrieved
|
||||
# and forwarded upstream. Gated by
|
||||
# _target_servers_delegate_auth_to_upstream, which only returns
|
||||
# True when EVERY target is auth_type=oauth2 AND has the
|
||||
# delegate_auth_to_upstream flag set — fails closed otherwise.
|
||||
validated_user_api_key_auth = UserAPIKeyAuth()
|
||||
elif has_explicit_litellm_key:
|
||||
# Explicit x-litellm-api-key provided - always validate normally
|
||||
# An explicit x-litellm-api-key is always a LiteLLM credential, even
|
||||
# for a delegated server, so validate it: identity / spend / rate
|
||||
# limits resolve and any stored upstream token can be forwarded.
|
||||
validated_user_api_key_auth = await user_api_key_auth(
|
||||
api_key=litellm_api_key, request=request
|
||||
)
|
||||
elif MCPRequestHandler._target_servers_delegate_auth_to_upstream(
|
||||
path=request_route,
|
||||
mcp_servers=mcp_servers,
|
||||
client_ip=IPAddressUtils.get_mcp_client_ip(request),
|
||||
):
|
||||
# Operator opted this oauth2 server into upstream-delegated auth: the
|
||||
# client authenticates directly with the upstream MCP server, so any
|
||||
# Authorization bearer is an upstream token, never a LiteLLM key. Skip
|
||||
# LiteLLM validation entirely — covering both the no-credential
|
||||
# discovery request and the authenticated call carrying the upstream
|
||||
# bearer — so a tool call that succeeds never carries a phantom 401
|
||||
# auth span; the bearer is forwarded upstream unchanged. Gated by
|
||||
# _target_servers_delegate_auth_to_upstream, which returns True only
|
||||
# when EVERY target is auth_type=oauth2 with delegate_auth_to_upstream
|
||||
# set; fails closed otherwise.
|
||||
validated_user_api_key_auth = UserAPIKeyAuth()
|
||||
elif oauth2_headers:
|
||||
# No x-litellm-api-key, but Authorization header present.
|
||||
# Could be a LiteLLM key (backward compat) OR an opaque OAuth2 token
|
||||
# the operator wants forwarded to an upstream OAuth2-mode MCP server.
|
||||
# Try LiteLLM auth first; on auth failure, only fall back to anonymous
|
||||
# passthrough when the request actually targets a server whose operator
|
||||
# configured ``auth_type=oauth2``. For any other server (api_key,
|
||||
# bearer_token, basic, etc.), a failed LiteLLM auth is a real failure
|
||||
# and must propagate — otherwise an attacker can exchange any garbage
|
||||
# bearer for an anonymous session.
|
||||
# Authorization on a non-delegated server: the bearer must be a real
|
||||
# LiteLLM credential, so a failed validation is a genuine 401/403 and
|
||||
# propagates. The sole anonymous fallback is the auth_type=none
|
||||
# pass-through cold-start (RFC 9728 discovery return), gated on a 401
|
||||
# so a recognized-but-forbidden key still fails closed.
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
try:
|
||||
validated_user_api_key_auth = await user_api_key_auth(
|
||||
api_key=litellm_api_key, request=request
|
||||
)
|
||||
except (HTTPException, ProxyException) as e:
|
||||
# HTTPException.status_code is int; ProxyException.code is
|
||||
# normalized to str in its __init__ but can be ``"None"`` or any
|
||||
# non-numeric string when the caller didn't supply a numeric
|
||||
# code, so we compare against both int and str forms rather
|
||||
# than coercing (``int("None")`` would raise ValueError and
|
||||
# rewrite the auth error as a 500).
|
||||
# ProxyException.code is normalized to str (possibly "None"), so
|
||||
# compare both int and str forms rather than coercing.
|
||||
status = e.status_code if isinstance(e, HTTPException) else e.code
|
||||
is_auth_error = status in (401, 403, "401", "403")
|
||||
is_unauthenticated = status in (401, "401")
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
if is_auth_error and MCPRequestHandler._target_servers_use_oauth2(
|
||||
path=request_route,
|
||||
mcp_servers=mcp_servers,
|
||||
client_ip=client_ip,
|
||||
mcp_servers_from_path = _parse_mcp_server_names_from_path(
|
||||
request_route, mcp_servers
|
||||
)
|
||||
if (
|
||||
is_unauthenticated
|
||||
and mcp_servers_from_path is not None
|
||||
and not _has_client_supplied_mcp_auth(
|
||||
mcp_auth_header,
|
||||
mcp_server_auth_headers,
|
||||
)
|
||||
and _is_mcp_passthrough_cold_start(
|
||||
mcp_servers_from_path, client_ip=client_ip
|
||||
)
|
||||
):
|
||||
verbose_logger.debug(
|
||||
"MCP OAuth2: target server is OAuth2-mode, treating "
|
||||
"Authorization as upstream OAuth2 token passthrough"
|
||||
"MCP pass-through return: forwarding Authorization as "
|
||||
"upstream OAuth token for delegated auth"
|
||||
)
|
||||
validated_user_api_key_auth = UserAPIKeyAuth()
|
||||
elif is_unauthenticated:
|
||||
# Pass-through cold-start return: per RFC 9728 / MCP
|
||||
# Authorization spec the client completes upstream OAuth
|
||||
# discovery and returns with ``Authorization: Bearer
|
||||
# <upstream-token>``. For ``auth_type=none`` passthrough
|
||||
# servers that bearer is not a LiteLLM key (auth above
|
||||
# failed) but is meant to be forwarded upstream
|
||||
# unchanged. Fall back to anonymous admission so the
|
||||
# caller is not rejected for following the discovery
|
||||
# flow without also setting ``x-litellm-api-key``.
|
||||
# Only trigger on 401 (token unrecognized); a 403 means
|
||||
# the key WAS recognized but is forbidden (e.g. over
|
||||
# budget / rate limited) and must propagate so those
|
||||
# controls are not bypassed via anonymous admission.
|
||||
mcp_servers_from_path = _parse_mcp_server_names_from_path(
|
||||
request_route, mcp_servers
|
||||
)
|
||||
if (
|
||||
mcp_servers_from_path is not None
|
||||
and not _has_client_supplied_mcp_auth(
|
||||
mcp_auth_header,
|
||||
mcp_server_auth_headers,
|
||||
)
|
||||
and _is_mcp_passthrough_cold_start(
|
||||
mcp_servers_from_path, client_ip=client_ip
|
||||
)
|
||||
):
|
||||
verbose_logger.debug(
|
||||
"MCP pass-through return: target server is "
|
||||
"passthrough, treating Authorization as "
|
||||
"upstream OAuth token for delegated auth"
|
||||
)
|
||||
validated_user_api_key_auth = UserAPIKeyAuth()
|
||||
else:
|
||||
raise
|
||||
else:
|
||||
raise
|
||||
else:
|
||||
|
|
@ -408,45 +372,6 @@ class MCPRequestHandler:
|
|||
return [single_server_match.group(1)]
|
||||
return [servers_and_path]
|
||||
|
||||
@staticmethod
|
||||
def _target_servers_use_oauth2(
|
||||
path: str, mcp_servers: Optional[List[str]], client_ip: Optional[str]
|
||||
) -> bool:
|
||||
"""
|
||||
True only when EVERY MCP server the request targets is configured for
|
||||
``auth_type == oauth2``. If any target is non-OAuth2 — or if the target
|
||||
cannot be resolved at all — return False so the caller fails closed.
|
||||
|
||||
Used to gate the "treat Authorization as opaque OAuth2 token" fallback
|
||||
in :meth:`process_mcp_request` so a failed LiteLLM-auth cannot be
|
||||
exchanged for an anonymous session against a non-OAuth2 server.
|
||||
"""
|
||||
# Inline imports avoid a circular dependency: mcp_server_manager imports
|
||||
# from this module.
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
# Resolve the same target list downstream routing will use. For
|
||||
# ``/mcp/...`` routes, ``extract_mcp_auth_context`` overrides the
|
||||
# ``x-mcp-servers`` header with path-derived names, so we must mirror
|
||||
# that here — otherwise a caller could set the header to a permissive
|
||||
# server while the path targets a stricter one (header/path TOCTOU).
|
||||
target_names = MCPRequestHandler._resolve_target_server_names(
|
||||
path=path, mcp_servers_header=mcp_servers
|
||||
)
|
||||
if not target_names:
|
||||
return False
|
||||
|
||||
for name in target_names:
|
||||
server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
name, client_ip=client_ip
|
||||
)
|
||||
if server is None or server.auth_type != MCPAuth.oauth2:
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _target_servers_delegate_auth_to_upstream(
|
||||
path: str, mcp_servers: Optional[List[str]], client_ip: Optional[str]
|
||||
|
|
@ -468,8 +393,8 @@ class MCPRequestHandler:
|
|||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
# See _target_servers_use_oauth2: must mirror the downstream
|
||||
# header-vs-path override or an attacker could set
|
||||
# Must mirror the downstream header-vs-path override
|
||||
# (``extract_mcp_auth_context``) or an attacker could set
|
||||
# ``x-mcp-servers`` to a delegate-enabled server while the URL path
|
||||
# targets a non-delegate server, skipping LiteLLM auth for it.
|
||||
target_names = MCPRequestHandler._resolve_target_server_names(
|
||||
|
|
|
|||
|
|
@ -557,10 +557,12 @@ async def delete_mcp_server(
|
|||
"""
|
||||
Delete the mcp server from the db by server_id
|
||||
|
||||
The server-row delete is the commit point. Per-user env var rows have no FK
|
||||
cascade, so they are cleaned up afterwards on a best-effort basis: a transient
|
||||
failure there leaves only orphaned rows pointing at a now-missing server and
|
||||
must not turn a successful delete into a caller-visible error.
|
||||
The server-row delete is the commit point. Per-user credential and env var
|
||||
rows have no FK cascade, so they are cleaned up afterwards on a best-effort
|
||||
basis: a transient failure there leaves only orphaned rows pointing at a
|
||||
now-missing server and must not turn a successful delete into a
|
||||
caller-visible error. Each table is cleaned independently so a failure on one
|
||||
still attempts the other.
|
||||
|
||||
Returns the deleted mcp server record if it exists, otherwise None
|
||||
"""
|
||||
|
|
@ -570,17 +572,20 @@ async def delete_mcp_server(
|
|||
},
|
||||
)
|
||||
if deleted_server is not None:
|
||||
try:
|
||||
await prisma_client.db.litellm_mcpuserenvvars.delete_many(
|
||||
where={"server_id": server_id}
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"MCP server %s deleted but per-user env var cleanup failed; "
|
||||
"orphaned rows can be removed on a later delete: %s",
|
||||
server_id,
|
||||
e,
|
||||
)
|
||||
for model, label in (
|
||||
(prisma_client.db.litellm_mcpusercredentials, "credential"),
|
||||
(prisma_client.db.litellm_mcpuserenvvars, "env var"),
|
||||
):
|
||||
try:
|
||||
await model.delete_many(where={"server_id": server_id})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"MCP server %s deleted but per-user %s cleanup failed; "
|
||||
"orphaned rows can be removed on a later delete: %s",
|
||||
server_id,
|
||||
label,
|
||||
e,
|
||||
)
|
||||
return deleted_server
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -512,12 +512,13 @@ async def exchange_token_with_server(
|
|||
result = {
|
||||
"access_token": access_token,
|
||||
"token_type": token_response.get("token_type", "Bearer"),
|
||||
"expires_in": token_response.get("expires_in", 3600),
|
||||
}
|
||||
|
||||
if "refresh_token" in token_response and token_response["refresh_token"]:
|
||||
if token_response.get("expires_in") is not None:
|
||||
result["expires_in"] = token_response["expires_in"]
|
||||
if token_response.get("refresh_token"):
|
||||
result["refresh_token"] = token_response["refresh_token"]
|
||||
if "scope" in token_response and token_response["scope"]:
|
||||
if token_response.get("scope"):
|
||||
result["scope"] = token_response["scope"]
|
||||
|
||||
# RFC 6749 §5.1: token responses must not be cached.
|
||||
|
|
|
|||
|
|
@ -386,8 +386,15 @@ if MCP_AVAILABLE:
|
|||
raw_headers: Optional[Dict[str, str]] = None,
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
apply_tool_filters: bool = True,
|
||||
):
|
||||
"""Helper function to get tools for a single server."""
|
||||
"""Helper function to get tools for a single server.
|
||||
|
||||
When ``apply_tool_filters`` is False the raw server catalog is returned
|
||||
without the allowed_tools/disallowed_tools gate or the per-key tool
|
||||
permissions. This is the admin-only configuration view; every runtime
|
||||
path keeps the default True so callable tools stay filtered.
|
||||
"""
|
||||
tools = await global_mcp_server_manager._get_tools_from_server(
|
||||
server=server,
|
||||
mcp_auth_header=server_auth_header,
|
||||
|
|
@ -397,6 +404,9 @@ if MCP_AVAILABLE:
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
|
||||
if not apply_tool_filters:
|
||||
return _create_tool_response_objects(tools, server.mcp_info)
|
||||
|
||||
# Always apply allowed_tools/disallowed_tools so the blacklist is
|
||||
# enforced even when no allowlist is set (matches the SSE/HTTP path).
|
||||
tools = filter_tools_by_allowed_tools(tools, server)
|
||||
|
|
@ -463,6 +473,7 @@ if MCP_AVAILABLE:
|
|||
mcp_auth_header: Optional[str],
|
||||
raw_headers_from_request: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
apply_tool_filters: bool = True,
|
||||
) -> dict:
|
||||
"""Handle tool listing for a single server_id request."""
|
||||
# Resolve a server name to its UUID if needed
|
||||
|
|
@ -527,6 +538,7 @@ if MCP_AVAILABLE:
|
|||
raw_headers_from_request,
|
||||
user_api_key_dict,
|
||||
extra_headers=user_oauth_extra_headers,
|
||||
apply_tool_filters=apply_tool_filters,
|
||||
)
|
||||
except MCPUpstreamAuthError:
|
||||
# Surface the upstream 401/403 to the caller so it can emit the
|
||||
|
|
@ -552,6 +564,14 @@ if MCP_AVAILABLE:
|
|||
server_id: Optional[str] = Query(
|
||||
None, description="The server id to list tools for"
|
||||
),
|
||||
include_disabled_tools: bool = Query(
|
||||
False,
|
||||
description=(
|
||||
"Admin only. Return the full server tool catalog without the "
|
||||
"allowed_tools filter or per-key tool permissions, so the MCP "
|
||||
"settings UI can configure the allowlist. Ignored for non-admins."
|
||||
),
|
||||
),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
) -> dict:
|
||||
"""
|
||||
|
|
@ -579,6 +599,13 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
|
||||
try:
|
||||
# The full catalog (allowlist filter skipped) is admin-only so the
|
||||
# REST endpoint can't be used to enumerate deliberately-disabled tools.
|
||||
apply_tool_filters = not (
|
||||
include_disabled_tools
|
||||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
# Extract auth headers from request
|
||||
headers = request.headers
|
||||
raw_headers_from_request = dict(headers)
|
||||
|
|
@ -620,6 +647,7 @@ if MCP_AVAILABLE:
|
|||
mcp_auth_header=mcp_auth_header,
|
||||
raw_headers_from_request=raw_headers_from_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
apply_tool_filters=apply_tool_filters,
|
||||
)
|
||||
else:
|
||||
if not allowed_server_ids:
|
||||
|
|
@ -677,6 +705,7 @@ if MCP_AVAILABLE:
|
|||
raw_headers_from_request,
|
||||
user_api_key_dict,
|
||||
extra_headers=user_oauth_extra_headers,
|
||||
apply_tool_filters=apply_tool_filters,
|
||||
)
|
||||
list_tools_result.extend(tools_result)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -2473,6 +2473,17 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"`statement_cache_size`). Keys here override any default LiteLLM sets."
|
||||
),
|
||||
)
|
||||
database_disable_prepared_statements: Optional[bool] = Field(
|
||||
None,
|
||||
description=(
|
||||
"Disable server-side prepared statements by setting Prisma's "
|
||||
"`pgbouncer=true` URL param. Use this for pgbouncer transaction-pooling "
|
||||
"deployments, or to prevent the 'cached plan must not change result "
|
||||
"type' error that pooled connections hit during rolling schema "
|
||||
"migrations. An explicit `pgbouncer` in `database_extra_connection_params` "
|
||||
"takes precedence."
|
||||
),
|
||||
)
|
||||
database_type: Optional[Literal["dynamo_db"]] = Field(
|
||||
None, description="to use dynamodb instead of postgres db"
|
||||
)
|
||||
|
|
@ -2609,6 +2620,24 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
None,
|
||||
description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).",
|
||||
)
|
||||
disable_budget_reservation: Optional[bool] = Field(
|
||||
None,
|
||||
description=(
|
||||
"If True, disables the optimistic per-request budget reservation "
|
||||
"introduced in v1.84.0. "
|
||||
"WARNING: This weakens hard budget enforcement. Without the reservation, "
|
||||
"a burst of concurrent requests from a single key can each pass the "
|
||||
"read-time spend check before any of them is charged, allowing a "
|
||||
"configured budget to be exceeded under high concurrency. "
|
||||
"Budgets are still evaluated on every request at read time, so "
|
||||
"an already-exhausted budget is still rejected. "
|
||||
"Enable only if your deployment is experiencing phantom "
|
||||
"BudgetExceededError responses caused by leaked reservations "
|
||||
"(see GitHub issue #27639). "
|
||||
"A proxy-level WARNING is logged on every request while this flag "
|
||||
"is active as a reminder that hard enforcement is relaxed."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConfigYAML(LiteLLMPydanticObjectBase):
|
||||
|
|
|
|||
|
|
@ -154,6 +154,16 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
)
|
||||
elif isinstance(e, ProxyException):
|
||||
raise e
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
"Service Unavailable, the authentication database is "
|
||||
"temporarily unreachable. Please retry shortly."
|
||||
),
|
||||
type=ProxyErrorTypes.no_db_connection,
|
||||
param="None",
|
||||
code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
)
|
||||
raise ProxyException(
|
||||
message="Authentication Error, " + str(e),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
|
|
|
|||
|
|
@ -2422,6 +2422,7 @@ async def _run_centralized_common_checks( # noqa: PLR0915
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
skip_budget_checks=skip_budget_checks,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2442,12 +2443,23 @@ async def _reserve_budget_after_common_checks(
|
|||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
skip_budget_checks: bool,
|
||||
general_settings: dict,
|
||||
end_user_id: Optional[str] = None,
|
||||
end_user_object: Optional[LiteLLM_EndUserTable] = None,
|
||||
) -> None:
|
||||
user_api_key_auth_obj.budget_reservation = None
|
||||
if skip_budget_checks:
|
||||
return
|
||||
if general_settings.get("disable_budget_reservation") is True:
|
||||
verbose_proxy_logger.warning(
|
||||
"disable_budget_reservation is enabled: skipping optimistic budget "
|
||||
"reservation. Budget enforcement is read-time only — concurrent "
|
||||
"requests can each pass the spend check before their cost is recorded, "
|
||||
"so a configured budget may be briefly exceeded under high concurrency. "
|
||||
"Set disable_budget_reservation to False or remove it to restore "
|
||||
"hard per-request budget enforcement."
|
||||
)
|
||||
return
|
||||
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
reserve_budget_for_request,
|
||||
|
|
|
|||
|
|
@ -109,6 +109,92 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def is_prisma_engine_internal_error(e: Exception) -> bool:
|
||||
"""True iff ``e`` is a non-``PrismaError`` exception raised from inside
|
||||
prisma-client-py's query-engine layer.
|
||||
|
||||
During the instant a DB connection is torn down, the query engine can
|
||||
return a malformed error payload (``user_facing_error.meta`` is
|
||||
``null``). prisma-client-py's ``handle_response_errors`` then crashes
|
||||
with ``AttributeError: 'NoneType' object has no attribute 'get'``
|
||||
before it can raise the proper P1001 "can't reach database server"
|
||||
error. That AttributeError carries no connection keyword, so it can't
|
||||
be matched by message; identify it by its ``prisma.engine`` origin
|
||||
instead.
|
||||
|
||||
Recognized ``PrismaError`` subclasses are excluded: connectivity ones
|
||||
are already classified by type/keyword above, and data-layer ones
|
||||
(the DB IS reachable) must stay 401.
|
||||
"""
|
||||
import prisma
|
||||
|
||||
if isinstance(e, prisma.errors.PrismaError):
|
||||
return False
|
||||
tb = getattr(e, "__traceback__", None)
|
||||
while tb is not None:
|
||||
if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"):
|
||||
return True
|
||||
tb = tb.tb_next
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def is_database_service_unavailable_error(e: Exception) -> bool:
|
||||
"""True iff the exception means the database could not answer at the
|
||||
infrastructure level (connection refused, socket/interface failure,
|
||||
timeout) rather than a genuine auth failure (key not found) or a
|
||||
data-layer error (the DB IS reachable and rejected the data).
|
||||
|
||||
Auth must answer 401 only for a key the DB confirms is invalid. When
|
||||
the DB itself is unreachable, the request has to surface as 503 so
|
||||
callers retry instead of treating valid keys as invalid during an
|
||||
outage.
|
||||
|
||||
Note: prisma-client-py mislabels the P1001 "can't reach database
|
||||
server" connectivity failure as a ``DataError`` (a data-layer type),
|
||||
so a type-only check misses real outages. ``is_database_transport_error``
|
||||
keyword-matches the connection message and catches that masquerade,
|
||||
while genuine data errors (no connection keyword) correctly stay 401.
|
||||
|
||||
The Postgres "cached plan must not change result type" error is matched
|
||||
here, not in ``is_database_transport_error``: it is a transient stale-DB-
|
||||
state condition (not an invalid key), but the connection is healthy so it
|
||||
must not trigger a reconnect.
|
||||
|
||||
A non-``PrismaError`` raised from inside the prisma query engine (e.g.
|
||||
the ``AttributeError`` from ``handle_response_errors`` when the engine
|
||||
returns a malformed error payload mid-tear-down) is also treated as
|
||||
unavailable; see ``is_prisma_engine_internal_error``.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
if PrismaDBExceptionHandler.is_database_connection_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_transport_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e):
|
||||
return True
|
||||
if "cached plan must not change result type" in str(e).lower():
|
||||
return True
|
||||
|
||||
# OSError already covers ConnectionError and (Py3.3+) TimeoutError.
|
||||
# asyncio.TimeoutError is a distinct class before Py3.11.
|
||||
if isinstance(e, (OSError, asyncio.TimeoutError)):
|
||||
return True
|
||||
|
||||
try:
|
||||
import asyncpg
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
return isinstance(
|
||||
e,
|
||||
(
|
||||
asyncpg.exceptions.PostgresConnectionError,
|
||||
asyncpg.exceptions.InterfaceError,
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def handle_db_exception(e: Exception):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -105,6 +105,16 @@ def _extract_text_from_content(content: object) -> str:
|
|||
return ""
|
||||
|
||||
|
||||
def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Optional[dict[str, Any]]:
|
||||
merged: dict[str, Any] = {}
|
||||
present = False
|
||||
for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")):
|
||||
if isinstance(bag, Mapping):
|
||||
present = True
|
||||
merged.update(bag)
|
||||
return merged if present else None
|
||||
|
||||
|
||||
class CrowdStrikeAIDRHandler(CustomGuardrail):
|
||||
"""
|
||||
CrowdStrike AIDR AI Guardrail handler to interact with the CrowdStrike AIDR
|
||||
|
|
@ -317,6 +327,22 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
|
|||
"event_type": event_type,
|
||||
}
|
||||
|
||||
model = inputs.get("model")
|
||||
if model:
|
||||
ai_guard_payload["model"] = model
|
||||
|
||||
metadata = _merge_metadata_bags(request_data)
|
||||
if metadata is not None:
|
||||
user_id = metadata.get("user_api_key_user_id")
|
||||
if user_id:
|
||||
ai_guard_payload["user_id"] = user_id
|
||||
|
||||
extra_info: dict[str, str] = {}
|
||||
user_email = metadata.get("user_api_key_user_email")
|
||||
if user_email:
|
||||
extra_info["user_name"] = user_email
|
||||
ai_guard_payload["extra_info"] = extra_info
|
||||
|
||||
ai_guard_response = await self._call_crowdstrike_aidr_guard(
|
||||
ai_guard_payload, hook_name
|
||||
)
|
||||
|
|
|
|||
|
|
@ -518,11 +518,17 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
# Check if this is a managed file (base64 encoded unified file ID)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
get_models_from_unified_file_id,
|
||||
)
|
||||
|
||||
# Managed files require bypassing the HTTP endpoint (which runs access-check hooks)
|
||||
# and calling the managed files hook directly with the user's credentials.
|
||||
is_managed_file = _is_base64_encoded_unified_file_id(file_id)
|
||||
target_model_names = (
|
||||
get_models_from_unified_file_id(is_managed_file)
|
||||
if is_managed_file
|
||||
else []
|
||||
)
|
||||
if is_managed_file and user_api_key_dict is not None:
|
||||
file_content = await self._fetch_managed_file_content(
|
||||
file_id=file_id,
|
||||
|
|
@ -560,6 +566,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
await self._enforce_batch_file_model_access(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
file_content_as_dict=file_content_as_dict,
|
||||
target_model_names=target_model_names or None,
|
||||
)
|
||||
|
||||
input_file_usage = _get_batch_job_input_file_usage(
|
||||
|
|
@ -595,9 +602,13 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
file_content_as_dict: List[dict],
|
||||
target_model_names: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Reject the batch if the caller is not authorized for every
|
||||
``body.model`` named inside the JSONL.
|
||||
"""Reject the batch if the caller is not authorized for the upload target.
|
||||
|
||||
For managed files, ``target_model_names`` (from the unified file id) is
|
||||
the proxy alias the file was uploaded for and is used directly for auth.
|
||||
For legacy/non-managed files, falls back to ``body.model`` values in the JSONL.
|
||||
|
||||
Reuses standard auth helpers so the same model access rules the proxy
|
||||
enforces on `/chat/completions` apply here.
|
||||
|
|
@ -614,9 +625,12 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
models = _get_models_from_batch_input_file_content(file_content_as_dict)
|
||||
if not models:
|
||||
return
|
||||
if target_model_names:
|
||||
models = target_model_names
|
||||
else:
|
||||
models = _get_models_from_batch_input_file_content(file_content_as_dict)
|
||||
if not models:
|
||||
return
|
||||
|
||||
team_object = None
|
||||
if (
|
||||
|
|
@ -647,12 +661,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
|
||||
llm_model_list = llm_router.model_list if llm_router is not None else None
|
||||
for model in models:
|
||||
# body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth.
|
||||
model_to_check = model
|
||||
if llm_router is not None:
|
||||
proxy_model_name = llm_router.resolve_model_name_from_model_id(model)
|
||||
if proxy_model_name is not None:
|
||||
model_to_check = proxy_model_name
|
||||
try:
|
||||
if team_object is not None:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ import json
|
|||
import os
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Iterable, List, Literal, Optional
|
||||
from typing import Any, Dict, Iterable, List, Literal, Optional, Set
|
||||
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
|
|
@ -1714,11 +1714,13 @@ if MCP_AVAILABLE:
|
|||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={"error": f"Access denied to MCP server {server_id}"},
|
||||
)
|
||||
allowed_server_ids = (
|
||||
await global_mcp_server_manager.get_allowed_mcp_servers(
|
||||
user_api_key_dict
|
||||
allowed_server_ids: Set[str] = set()
|
||||
for auth_context in await build_effective_auth_contexts(user_api_key_dict):
|
||||
allowed_server_ids.update(
|
||||
await global_mcp_server_manager.get_allowed_mcp_servers(
|
||||
auth_context
|
||||
)
|
||||
)
|
||||
)
|
||||
if server.server_id not in allowed_server_ids:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
|
|
|
|||
|
|
@ -4684,15 +4684,36 @@ async def team_model_add(
|
|||
detail={"error": "Only proxy admin or team admin can modify team models"},
|
||||
)
|
||||
|
||||
updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models)
|
||||
# Update team. `include` mirrors the relations the auth path consumes
|
||||
# off the cached team object so that `_refresh_cached_team` doesn't
|
||||
# null them out — see object_permission_utils.validate_key_search_tools_against_team
|
||||
# and the MCP/agent authz paths, which treat a missing object_permission
|
||||
# as "no team-level restriction".
|
||||
# Atomic array append with dedup at the database level so concurrent
|
||||
# BYOK model creates don't overwrite each other's team.models entries.
|
||||
# When the team currently has models=[] (unrestricted access), the
|
||||
# CASE expression inserts the 'all-proxy-models' sentinel first.
|
||||
models_to_add = list(data.models)
|
||||
await prisma_client.db.execute_raw(
|
||||
'UPDATE "LiteLLM_TeamTable" '
|
||||
"SET models = ("
|
||||
" SELECT ARRAY(SELECT DISTINCT unnest("
|
||||
" CASE WHEN cardinality(COALESCE(models, ARRAY[]::text[])) = 0 "
|
||||
" THEN ARRAY['all-proxy-models']::text[] "
|
||||
" ELSE models "
|
||||
" END || $1::text[]"
|
||||
" ))"
|
||||
") "
|
||||
"WHERE team_id = $2",
|
||||
models_to_add,
|
||||
data.team_id,
|
||||
)
|
||||
# Re-fetch via update (write-routed) instead of find_unique (read-routed)
|
||||
# to avoid returning stale data from a read replica. The models column was
|
||||
# already set by execute_raw above; this bumps updated_at. `include` mirrors
|
||||
# the relations the auth path consumes off the cached team object so that
|
||||
# `_refresh_cached_team` doesn't null them out — see
|
||||
# object_permission_utils.validate_key_search_tools_against_team and the
|
||||
# MCP/agent authz paths, which treat a missing object_permission as
|
||||
# "no team-level restriction".
|
||||
updated_team = await prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": data.team_id},
|
||||
data={"models": updated_models},
|
||||
data={"updated_at": datetime.now(timezone.utc)},
|
||||
include={"object_permission": True}, # type: ignore
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -100,6 +100,42 @@ class AnthropicPassthroughLoggingHandler:
|
|||
return get_end_user_id_from_request_body(request_body)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _resolve_costing_model(model: str, logging_obj: LiteLLMLoggingObj) -> str:
|
||||
if model and model != "unknown":
|
||||
return model
|
||||
litellm_params = (getattr(logging_obj, "model_call_details", {}) or {}).get(
|
||||
"litellm_params", {}
|
||||
) or {}
|
||||
deployment_model = litellm_params.get("model")
|
||||
if deployment_model and deployment_model != "unknown":
|
||||
return deployment_model
|
||||
model_group = (litellm_params.get("metadata", {}) or {}).get("model_group")
|
||||
if model_group:
|
||||
return model_group.removeprefix("passthrough/")
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _extract_model_from_anthropic_chunks(
|
||||
all_chunks: Sequence[Union[str, bytes]],
|
||||
) -> Optional[str]:
|
||||
for raw in all_chunks:
|
||||
text = raw.decode("utf-8") if isinstance(raw, bytes) else raw
|
||||
for line in text.splitlines():
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line[len("data:") :].strip())
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
if data.get("type") == "message_start":
|
||||
model = (data.get("message") or {}).get("model")
|
||||
if model:
|
||||
return model
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _create_anthropic_response_logging_payload(
|
||||
litellm_model_response: Union[ModelResponse, TextCompletionResponse],
|
||||
|
|
@ -127,6 +163,10 @@ class AnthropicPassthroughLoggingHandler:
|
|||
"custom_llm_provider"
|
||||
)
|
||||
|
||||
model = AnthropicPassthroughLoggingHandler._resolve_costing_model(
|
||||
model, logging_obj
|
||||
)
|
||||
|
||||
# Prepend custom_llm_provider to model if not already present
|
||||
model_for_cost = model
|
||||
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
|
||||
|
|
@ -213,6 +253,15 @@ class AnthropicPassthroughLoggingHandler:
|
|||
):
|
||||
model = cast(str, litellm_logging_obj.model_call_details.get("model"))
|
||||
|
||||
if not model or model == "unknown":
|
||||
chunk_model = (
|
||||
AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(
|
||||
all_chunks
|
||||
)
|
||||
)
|
||||
if chunk_model:
|
||||
model = chunk_model
|
||||
|
||||
complete_streaming_response = (
|
||||
AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=all_chunks,
|
||||
|
|
@ -468,6 +517,13 @@ class AnthropicPassthroughLoggingHandler:
|
|||
# Process each individual event
|
||||
for event_str in individual_events:
|
||||
try:
|
||||
# Skip OpenAI-style [DONE] sentinels some Anthropic-compatible
|
||||
# providers emit. Match the whole SSE line so a valid chunk whose
|
||||
# text payload happens to contain "[DONE]" is not dropped.
|
||||
if any(
|
||||
line.strip() == "data: [DONE]" for line in event_str.split("\n")
|
||||
):
|
||||
continue
|
||||
transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk(
|
||||
chunk=event_str
|
||||
)
|
||||
|
|
@ -476,6 +532,14 @@ class AnthropicPassthroughLoggingHandler:
|
|||
|
||||
except (StopIteration, StopAsyncIteration):
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
# Some upstreams emit non-JSON SSE lines; skip them so the
|
||||
# logging pipeline is not broken by a single bad frame.
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping non-JSON SSE event: %s",
|
||||
event_str[:200],
|
||||
)
|
||||
continue
|
||||
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=all_openai_chunks,
|
||||
|
|
|
|||
|
|
@ -44,15 +44,19 @@ def _build_db_connection_url_params(
|
|||
pool_timeout: Optional[Union[int, float]],
|
||||
connect_timeout: Optional[Union[int, float]] = None,
|
||||
socket_timeout: Optional[Union[int, float]] = None,
|
||||
disable_prepared_statements: bool = False,
|
||||
extra_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""Build the Prisma DATABASE_URL query params controlling connection pool behavior.
|
||||
|
||||
`connect_timeout` / `socket_timeout` map to the Prisma URL params of the same
|
||||
name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are
|
||||
omitted when None so Prisma's defaults apply. `extra_params` is an
|
||||
untyped passthrough — keys it provides win over the named arguments above,
|
||||
so it can be used to override any default we set here.
|
||||
omitted when None so Prisma's defaults apply. `disable_prepared_statements`
|
||||
sets `pgbouncer=true`, which makes Prisma stop using server-side prepared
|
||||
statements (pgbouncer transaction-pool compatible; also sidesteps the
|
||||
"cached plan must not change result type" error during rolling migrations).
|
||||
`extra_params` is an untyped passthrough — keys it provides win over the
|
||||
named arguments above, so it can be used to override any default we set here.
|
||||
"""
|
||||
params: dict = {
|
||||
"connection_limit": connection_limit,
|
||||
|
|
@ -63,6 +67,8 @@ def _build_db_connection_url_params(
|
|||
params["connect_timeout"] = connect_timeout
|
||||
if socket_timeout is not None:
|
||||
params["socket_timeout"] = socket_timeout
|
||||
if disable_prepared_statements:
|
||||
params["pgbouncer"] = "true"
|
||||
if extra_params:
|
||||
params.update(extra_params)
|
||||
return params
|
||||
|
|
@ -947,6 +953,7 @@ def run_server( # noqa: PLR0915
|
|||
db_connection_timeout: Optional[Union[int, float]] = 60
|
||||
db_connect_timeout: Optional[Union[int, float]] = None
|
||||
db_socket_timeout: Optional[Union[int, float]] = None
|
||||
db_disable_prepared_statements: bool = False
|
||||
db_extra_connection_params: Optional[dict] = None
|
||||
general_settings = {}
|
||||
### GET DB TOKEN FOR IAM AUTH ###
|
||||
|
|
@ -1067,6 +1074,17 @@ def run_server( # noqa: PLR0915
|
|||
)
|
||||
db_connect_timeout = general_settings.get("database_connect_timeout")
|
||||
db_socket_timeout = general_settings.get("database_socket_timeout")
|
||||
_disable_prepared_statements = general_settings.get(
|
||||
"database_disable_prepared_statements", False
|
||||
)
|
||||
if isinstance(_disable_prepared_statements, str):
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
db_disable_prepared_statements = (
|
||||
str_to_bool(_disable_prepared_statements) is True
|
||||
)
|
||||
else:
|
||||
db_disable_prepared_statements = bool(_disable_prepared_statements)
|
||||
db_extra_connection_params = general_settings.get(
|
||||
"database_extra_connection_params"
|
||||
)
|
||||
|
|
@ -1114,6 +1132,7 @@ def run_server( # noqa: PLR0915
|
|||
pool_timeout=db_connection_timeout,
|
||||
connect_timeout=db_connect_timeout,
|
||||
socket_timeout=db_socket_timeout,
|
||||
disable_prepared_statements=db_disable_prepared_statements,
|
||||
extra_params=db_extra_connection_params,
|
||||
)
|
||||
if os.getenv("DATABASE_URL", None) is not None:
|
||||
|
|
|
|||
|
|
@ -10920,16 +10920,26 @@ def get_direct_access_models(
|
|||
return direct_access_models
|
||||
|
||||
|
||||
async def get_all_team_and_direct_access_models(
|
||||
def _filter_models_to_user_accessible(all_models: List[Dict]) -> List[Dict]:
|
||||
"""Keep only deployments the caller can use via direct access or team membership."""
|
||||
return [
|
||||
_model
|
||||
for _model in all_models
|
||||
if _model.get("model_info", {}).get("direct_access", False)
|
||||
or _model.get("model_info", {}).get("access_via_team_ids", [])
|
||||
]
|
||||
|
||||
|
||||
async def _populate_team_access_on_models(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
llm_router: Router,
|
||||
all_models: List[Dict],
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Get all models across all teams user is in.
|
||||
Populate `model_info.access_via_team_ids` and `model_info.direct_access`
|
||||
without filtering the model list.
|
||||
"""
|
||||
|
||||
user_teams: Optional[Union[List[str], Literal["*"]]] = None
|
||||
direct_access_models: List[str] = []
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
|
|
@ -10948,7 +10958,6 @@ async def get_all_team_and_direct_access_models(
|
|||
user_db_object=user_object,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
## ADD ACCESS_VIA_TEAM_IDS TO ALL MODELS
|
||||
if user_teams is not None:
|
||||
team_models = await get_all_team_models(
|
||||
user_teams=user_teams,
|
||||
|
|
@ -10971,23 +10980,33 @@ async def get_all_team_and_direct_access_models(
|
|||
model_id, []
|
||||
)
|
||||
|
||||
## ADD DIRECT_ACCESS TO RELEVANT MODELS
|
||||
|
||||
direct_access_model_ids = set(direct_access_models)
|
||||
for _model in all_models:
|
||||
model_id = _model.get("model_info", {}).get("id", None)
|
||||
if model_id is not None and model_id in direct_access_models:
|
||||
_model["model_info"]["direct_access"] = True
|
||||
if model_id is not None:
|
||||
_model["model_info"]["direct_access"] = model_id in direct_access_model_ids
|
||||
|
||||
## FILTER OUT MODELS THAT ARE NOT IN DIRECT_ACCESS_MODELS OR ACCESS_VIA_TEAM_IDS - only show user models they can call
|
||||
all_models = [
|
||||
_model
|
||||
for _model in all_models
|
||||
if _model.get("model_info", {}).get("direct_access", False)
|
||||
or _model.get("model_info", {}).get("access_via_team_ids", [])
|
||||
]
|
||||
return all_models
|
||||
|
||||
|
||||
async def get_all_team_and_direct_access_models(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
llm_router: Router,
|
||||
all_models: List[Dict],
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Get all models across all teams user is in.
|
||||
"""
|
||||
all_models = await _populate_team_access_on_models(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
all_models=all_models,
|
||||
)
|
||||
return _filter_models_to_user_accessible(all_models)
|
||||
|
||||
|
||||
def _enrich_model_info_with_litellm_data(
|
||||
model: Dict[str, Any], debug: bool = False, llm_router: Optional[Router] = None
|
||||
) -> Dict[str, Any]:
|
||||
|
|
@ -11096,6 +11115,22 @@ async def _get_caller_byok_team_scope(
|
|||
return set(user_row.teams or [])
|
||||
|
||||
|
||||
def _byok_row_outside_caller_teams(
|
||||
model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]]
|
||||
) -> bool:
|
||||
"""Whether a team BYOK row belongs to a team the caller is not a member of.
|
||||
|
||||
`team_id` is only set on team BYOK rows; non-team rows fall through
|
||||
unaffected. `allowed_team_ids is None` means no scoping (e.g. admins).
|
||||
"""
|
||||
if allowed_team_ids is None:
|
||||
return False
|
||||
team_id = model_info_dict.get("team_id")
|
||||
if team_id is None:
|
||||
return False
|
||||
return team_id not in allowed_team_ids
|
||||
|
||||
|
||||
# Hard cap on rows the DB-side BYOK search may pull when results need to be
|
||||
# sorted across the full match set. Without this, an authenticated caller
|
||||
# can hit `/v2/model/info?search=<broad>&sortBy=<field>` and force the
|
||||
|
|
@ -11217,15 +11252,7 @@ async def _apply_search_filter_to_models(
|
|||
)
|
||||
|
||||
def _is_byok_outside_caller_teams(model_info_dict: Dict[str, Any]) -> bool:
|
||||
# `team_id` is only set on team BYOK rows. Non-team rows fall
|
||||
# through unaffected — they are gated by other paths (router
|
||||
# membership, direct_access, include_team_models).
|
||||
if allowed_team_ids is None:
|
||||
return False
|
||||
team_id = model_info_dict.get("team_id")
|
||||
if team_id is None:
|
||||
return False
|
||||
return team_id not in allowed_team_ids
|
||||
return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids)
|
||||
|
||||
def _model_matches_search(m: Dict[str, Any]) -> bool:
|
||||
# Team BYOK models persist an internal `model_name`
|
||||
|
|
@ -11723,10 +11750,8 @@ async def _find_model_by_id(
|
|||
|
||||
@router.get(
|
||||
"/v2/model/info",
|
||||
description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true",
|
||||
tags=["model management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
include_in_schema=False,
|
||||
)
|
||||
async def model_info_v2(
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
|
|
@ -11762,7 +11787,49 @@ async def model_info_v2(
|
|||
),
|
||||
):
|
||||
"""
|
||||
BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now.
|
||||
Paginated model metadata for proxy deployments (pricing, provider, team access).
|
||||
|
||||
Returns configured router deployments with enriched `model_info` (costs, provider,
|
||||
context window, etc.). Sensitive fields such as API keys and api_base are omitted.
|
||||
|
||||
Query parameters:
|
||||
model: Filter to a single public `model_name`.
|
||||
user_models_only: When true, only return models created by the calling user.
|
||||
include_team_models: When true, populate `access_via_team_ids` and `direct_access`
|
||||
on each model and filter to deployments the caller can use.
|
||||
page / size: Pagination controls (defaults: page=1, size=50).
|
||||
search: Case-insensitive partial match on model name or team public name.
|
||||
modelId: Return a single deployment by LiteLLM model id.
|
||||
teamId: Filter to models with direct access or team membership for this team id.
|
||||
sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status.
|
||||
|
||||
Example request:
|
||||
```
|
||||
curl -X GET 'http://localhost:4000/v2/model/info?include_team_models=true&page=1&size=50' \\
|
||||
--header 'Authorization: Bearer sk-1234'
|
||||
```
|
||||
|
||||
Example response:
|
||||
```json
|
||||
{
|
||||
"data": [
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "openai/gpt-4.1"},
|
||||
"model_info": {
|
||||
"id": "abc123",
|
||||
"litellm_provider": "openai",
|
||||
"access_via_team_ids": ["team-1"],
|
||||
"direct_access": true
|
||||
}
|
||||
}
|
||||
],
|
||||
"total_count": 1,
|
||||
"current_page": 1,
|
||||
"total_pages": 1,
|
||||
"size": 50
|
||||
}
|
||||
```
|
||||
"""
|
||||
global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router
|
||||
|
||||
|
|
@ -12325,6 +12392,72 @@ async def model_metrics_exceptions(
|
|||
return {"data": response, "exception_types": list(exception_types)}
|
||||
|
||||
|
||||
def _deployment_matches_allowed_model_names(
|
||||
model: Dict[str, Any], allowed_model_names: Set[str]
|
||||
) -> bool:
|
||||
"""Match a router deployment against allowed public model names.
|
||||
|
||||
Team-scoped rows store an internal routing key in ``model_name``; callers
|
||||
with key/team restrictions still refer to the public name in
|
||||
``model_info.team_public_model_name``.
|
||||
"""
|
||||
if model.get("model_name") in allowed_model_names:
|
||||
return True
|
||||
model_info = model.get("model_info")
|
||||
if not isinstance(model_info, dict):
|
||||
return False
|
||||
team_public_model_name = model_info.get("team_public_model_name")
|
||||
return (
|
||||
isinstance(team_public_model_name, str)
|
||||
and team_public_model_name in allowed_model_names
|
||||
)
|
||||
|
||||
|
||||
def _get_v1_model_info_allowed_model_names(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
llm_router: Router,
|
||||
) -> Optional[Set[str]]:
|
||||
"""Return key/team allowlisted public model names, or None if unrestricted."""
|
||||
model_access_groups = llm_router.get_model_access_groups()
|
||||
proxy_model_list = llm_router.get_model_names()
|
||||
key_models = get_key_models(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
team_models = get_team_models(
|
||||
team_models=user_api_key_dict.team_models,
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
if not key_models and not team_models:
|
||||
return None
|
||||
return set(
|
||||
get_complete_model_list(
|
||||
key_models=key_models,
|
||||
team_models=team_models,
|
||||
proxy_model_list=proxy_model_list,
|
||||
user_model=user_model,
|
||||
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
|
||||
llm_router=llm_router,
|
||||
return_wildcard_routes=False,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _filter_v1_model_info_deployments(
|
||||
all_models: List[dict],
|
||||
allowed_model_names: Optional[Set[str]],
|
||||
) -> List[dict]:
|
||||
if allowed_model_names is None:
|
||||
return all_models
|
||||
return [
|
||||
model
|
||||
for model in all_models
|
||||
if _deployment_matches_allowed_model_names(model, allowed_model_names)
|
||||
]
|
||||
|
||||
|
||||
def _translate_model_name_for_response(model: dict) -> dict:
|
||||
"""For team-scoped DB rows, replace `model_name` with the public name
|
||||
in `model_info.team_public_model_name` before returning. The DB column
|
||||
|
|
@ -12408,6 +12541,14 @@ def _get_proxy_model_info(model: dict) -> dict:
|
|||
async def model_info_v1( # noqa: PLR0915
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
litellm_model_id: Optional[str] = None,
|
||||
include_team_models: Optional[bool] = fastapi.Query(
|
||||
False,
|
||||
description="When true, filter to deployments the caller can use via direct access or team membership.",
|
||||
),
|
||||
teamId: Optional[str] = fastapi.Query(
|
||||
None,
|
||||
description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids",
|
||||
),
|
||||
):
|
||||
"""
|
||||
Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)
|
||||
|
|
@ -12417,6 +12558,11 @@ async def model_info_v1( # noqa: PLR0915
|
|||
|
||||
- When litellm_model_id is passed, it will return the info for that specific model
|
||||
- When litellm_model_id is not passed, it will return the info for all models
|
||||
- include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info).
|
||||
- teamId: Filter to models accessible by the given team.
|
||||
|
||||
Each model in the list response includes `model_info.access_via_team_ids` and
|
||||
`model_info.direct_access` when the proxy database is connected.
|
||||
|
||||
Returns:
|
||||
Returns a dictionary containing information about each model.
|
||||
|
|
@ -12443,6 +12589,12 @@ async def model_info_v1( # noqa: PLR0915
|
|||
"""
|
||||
global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router, user_model
|
||||
|
||||
# Unit tests call this handler directly; FastAPI normally resolves Query defaults.
|
||||
if not isinstance(include_team_models, bool):
|
||||
include_team_models = False
|
||||
if not isinstance(teamId, str):
|
||||
teamId = None
|
||||
|
||||
if user_model is not None:
|
||||
# user is trying to get specific model from litellm router
|
||||
try:
|
||||
|
|
@ -12479,6 +12631,14 @@ async def model_info_v1( # noqa: PLR0915
|
|||
},
|
||||
)
|
||||
|
||||
if prisma_client is None and (
|
||||
include_team_models or (teamId is not None and teamId.strip())
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||||
)
|
||||
|
||||
if litellm_model_id is not None:
|
||||
# user is trying to get specific model from litellm router
|
||||
deployment_info = llm_router.get_deployment(model_id=litellm_model_id)
|
||||
|
|
@ -12492,51 +12652,82 @@ async def model_info_v1( # noqa: PLR0915
|
|||
_deployment_info_dict = _get_proxy_model_info(
|
||||
model=deployment_info.model_dump(exclude_none=True)
|
||||
)
|
||||
return {"data": [_deployment_info_dict]}
|
||||
single_model_list: List[dict] = [_deployment_info_dict]
|
||||
if prisma_client is not None:
|
||||
single_model_list = await _populate_team_access_on_models(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
all_models=single_model_list,
|
||||
)
|
||||
if include_team_models:
|
||||
single_model_list = _filter_models_to_user_accessible(single_model_list)
|
||||
if teamId is not None and teamId.strip():
|
||||
single_model_list = await _filter_models_by_team_id(
|
||||
all_models=single_model_list,
|
||||
team_id=teamId.strip(),
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
return {"data": single_model_list}
|
||||
|
||||
all_models: List[dict] = []
|
||||
model_access_groups: Dict[str, List[str]] = defaultdict(list)
|
||||
## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ##
|
||||
if llm_router is None:
|
||||
proxy_model_list = []
|
||||
else:
|
||||
proxy_model_list = llm_router.get_model_names()
|
||||
model_access_groups = llm_router.get_model_access_groups()
|
||||
key_models = get_key_models(
|
||||
# Return router deployments (same source as /v2/model/info), not wildcard-
|
||||
# expanded model names from get_complete_model_list(). Team-scoped rows
|
||||
# use internal routing keys (model_name_{team_id}_{uuid}) and were omitted
|
||||
# when v1 resolved models only via public model_name strings.
|
||||
all_models: List[dict] = copy.deepcopy(llm_router.model_list)
|
||||
allowed_model_names = _get_v1_model_info_allowed_model_names(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
team_models = get_team_models(
|
||||
team_models=user_api_key_dict.team_models,
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
all_models_str = get_complete_model_list(
|
||||
key_models=key_models,
|
||||
team_models=team_models,
|
||||
proxy_model_list=proxy_model_list,
|
||||
user_model=user_model,
|
||||
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
if len(all_models_str) > 0:
|
||||
_relevant_models = []
|
||||
for model in all_models_str:
|
||||
router_models = llm_router.get_model_list(model_name=model)
|
||||
if router_models is not None:
|
||||
_relevant_models.extend(router_models)
|
||||
if llm_model_list is not None:
|
||||
all_models = copy.deepcopy(_relevant_models) # type: ignore
|
||||
else:
|
||||
all_models = []
|
||||
all_models = _filter_v1_model_info_deployments(
|
||||
all_models=all_models,
|
||||
allowed_model_names=allowed_model_names,
|
||||
)
|
||||
|
||||
# Reassign each entry: _get_proxy_model_info returns a (possibly new)
|
||||
# dict via _translate_model_name_for_response, which does NOT mutate in
|
||||
# place. Binding only the loop variable would drop the public-name swap
|
||||
# for team-scoped rows and leak the internal routing key (#28382).
|
||||
all_models = [_get_proxy_model_info(model=model) for model in all_models]
|
||||
# Team BYOK deployments carry an internal routing key and other teams'
|
||||
# public name/team_id/api_base; drop the ones the caller cannot access so
|
||||
# listing the full router model_list does not leak cross-team metadata.
|
||||
allowed_team_ids = await _get_caller_byok_team_scope(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
all_models = [
|
||||
model
|
||||
for model in all_models
|
||||
if not _byok_row_outside_caller_teams(
|
||||
model.get("model_info") or {}, allowed_team_ids
|
||||
)
|
||||
]
|
||||
|
||||
if prisma_client is not None:
|
||||
all_models = await _populate_team_access_on_models(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
all_models=all_models,
|
||||
)
|
||||
|
||||
if include_team_models:
|
||||
all_models = _filter_models_to_user_accessible(all_models)
|
||||
|
||||
all_models = [
|
||||
_translate_model_name_for_response(
|
||||
_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router)
|
||||
)
|
||||
for model in all_models
|
||||
]
|
||||
|
||||
if teamId is not None and teamId.strip():
|
||||
all_models = await _filter_models_by_team_id(
|
||||
all_models=all_models,
|
||||
team_id=teamId.strip(),
|
||||
prisma_client=cast(PrismaClient, prisma_client),
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("all_models: %s", all_models)
|
||||
return {"data": all_models}
|
||||
|
|
|
|||
|
|
@ -3239,40 +3239,49 @@ class PrismaClient:
|
|||
self, sql_query: str, *args
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
Execute a query with automatic fallback for PostgreSQL cached plan errors.
|
||||
Execute a query, recovering once from PostgreSQL's "cached plan must not
|
||||
change result type" error.
|
||||
|
||||
This handles the "cached plan must not change result type" error that occurs
|
||||
during rolling deployments when schema changes are applied while old pods
|
||||
still have cached query plans expecting the old schema.
|
||||
That error surfaces during rolling deployments when a schema change
|
||||
invalidates the prepared-statement plans that pooled connections still
|
||||
hold. Clearing only the server-side plans with DEALLOCATE ALL makes
|
||||
things worse: Prisma's query engine keeps a per-connection client-side
|
||||
cache of prepared-statement names, so once the server drops a plan the
|
||||
engine re-sends a name PostgreSQL no longer recognizes and the
|
||||
connection breaks with `prepared statement "sN" does not exist`. With a
|
||||
small pool that connection stays poisoned and every auth lookup fails.
|
||||
|
||||
Args:
|
||||
sql_query: SQL query string to execute
|
||||
Recreating the Prisma client kills the engine subprocess and drops the
|
||||
server-side plans and the engine's client-side name cache together, so
|
||||
the retried query is prepared fresh. We reconnect through
|
||||
`attempt_db_reconnect`, which is singleflight: when a schema change
|
||||
poisons every pooled connection at once, the first cached-plan error
|
||||
recreates the client and the concurrent waiters reuse that single
|
||||
recreate instead of racing to kill each other's fresh engine. We then
|
||||
retry the identical query exactly once.
|
||||
|
||||
Returns:
|
||||
Query result or None
|
||||
The retry reuses the original query byte-for-byte. Mutating the SQL
|
||||
(e.g. injecting a unique comment) would defeat PostgreSQL's plan cache,
|
||||
forcing a fresh plan on every request and pegging the database CPU.
|
||||
|
||||
Raises:
|
||||
Original exception if not a cached plan error
|
||||
If the reconnect is skipped because a recent reconnect is still within
|
||||
its cooldown, the retry runs against the same connection and may fail
|
||||
again; the get_data backoff decorator re-runs the lookup and a later
|
||||
attempt reconnects once the cooldown elapses.
|
||||
"""
|
||||
try:
|
||||
return await self.db.query_first(sql_query, *args)
|
||||
except Exception as e:
|
||||
error_str = str(e)
|
||||
if "cached plan must not change result type" in error_str:
|
||||
# Force PostgreSQL to re-plan by invalidating the cache
|
||||
# Add a unique comment to make the query different
|
||||
sql_query_retry = sql_query.replace(
|
||||
"SELECT",
|
||||
f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */",
|
||||
)
|
||||
verbose_proxy_logger.warning(
|
||||
"PostgreSQL cached plan error detected for token lookup, "
|
||||
"retrying with fresh plan. This may occur during rolling deployments "
|
||||
"when schema changes are applied."
|
||||
)
|
||||
return await self.db.query_first(sql_query_retry, *args)
|
||||
else:
|
||||
if "cached plan must not change result type" not in str(e):
|
||||
raise
|
||||
verbose_proxy_logger.warning(
|
||||
"PostgreSQL cached plan error detected for token lookup; "
|
||||
"recreating the database connection and retrying with the same "
|
||||
"query. This may occur during rolling deployments when schema "
|
||||
"changes are applied."
|
||||
)
|
||||
await self.attempt_db_reconnect(reason="postgres_cached_plan_error")
|
||||
return await self.db.query_first(sql_query, *args)
|
||||
|
||||
@backoff.on_exception(
|
||||
backoff.expo,
|
||||
|
|
@ -3628,7 +3637,10 @@ class PrismaClient:
|
|||
db=self.db, hashed_token=hashed_token
|
||||
)
|
||||
if active_token_id:
|
||||
response = await self.get_data(
|
||||
# The recursive call returns a finished
|
||||
# LiteLLM_VerificationTokenView; the dict
|
||||
# normalization below would crash subscripting it.
|
||||
deprecated_response = await self.get_data(
|
||||
token=active_token_id,
|
||||
table_name="combined_view",
|
||||
query_type="find_unique",
|
||||
|
|
@ -3636,10 +3648,11 @@ class PrismaClient:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
check_deprecated=False,
|
||||
)
|
||||
if response is not None:
|
||||
if deprecated_response is not None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Deprecated key used during grace period"
|
||||
)
|
||||
return deprecated_response
|
||||
|
||||
if response is not None:
|
||||
if response["team_models"] is None:
|
||||
|
|
|
|||
|
|
@ -1672,6 +1672,9 @@ class Usage(SafeAttributeModel, CompletionUsage):
|
|||
prompt_tokens_details=_prompt_tokens_details or None,
|
||||
)
|
||||
|
||||
if isinstance(server_tool_use, dict):
|
||||
server_tool_use = ServerToolUse(**server_tool_use)
|
||||
|
||||
if server_tool_use is not None:
|
||||
self.server_tool_use = server_tool_use
|
||||
else: # maintain openai compatibility in usage object if possible
|
||||
|
|
@ -3026,6 +3029,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
|
|||
wandb_api_key: Optional[str]
|
||||
weave_project_id: Optional[str]
|
||||
|
||||
# Datadog dynamic params
|
||||
dd_api_key: Optional[str]
|
||||
dd_site: Optional[str]
|
||||
dd_agent_host: Optional[str]
|
||||
dd_agent_port: Optional[str]
|
||||
|
||||
# Logging settings
|
||||
turn_off_message_logging: Optional[bool] # when true will not log messages
|
||||
litellm_disabled_callbacks: Optional[List[str]]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.89.0"
|
||||
version = "1.89.1"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.14"
|
||||
|
|
@ -125,7 +125,7 @@ proxy-runtime = [
|
|||
"mangum>=0.17.0,<1.0",
|
||||
"azure-ai-contentsafety>=1.0.0,<2.0",
|
||||
"azure-storage-file-datalake>=12.20.0,<13.0",
|
||||
"pypdf>=6.10.2,<7.0; python_version < '3.14'",
|
||||
"pypdf>=6.12.0,<7.0; python_version < '3.14'",
|
||||
"llm-sandbox>=0.3.39,<1.0",
|
||||
"detect-secrets>=1.5.0,<2.0",
|
||||
]
|
||||
|
|
@ -231,6 +231,10 @@ requires = ["uv_build==0.11.8"]
|
|||
build-backend = "uv_build"
|
||||
|
||||
[tool.uv]
|
||||
constraint-dependencies = [
|
||||
"tornado>=6.5.6",
|
||||
"aiohttp>=3.13.5,<3.14",
|
||||
]
|
||||
default-groups = ["dev"]
|
||||
required-version = ">=0.10.9"
|
||||
exclude-newer = "3 days"
|
||||
|
|
@ -260,7 +264,7 @@ source-exclude = [
|
|||
profile = "black"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.89.0"
|
||||
version = "1.89.1"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,194 @@
|
|||
"""
|
||||
Tests for team-scoped Datadog callback support.
|
||||
|
||||
Verifies that DataDogLogger can be instantiated with per-team credentials
|
||||
(dd_api_key, dd_site) instead of relying solely on environment variables,
|
||||
and that the DataDogHandler correctly resolves and caches per-team loggers.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.datadog.datadog import DataDogLogger
|
||||
from litellm.integrations.datadog.datadog_team_handler import (
|
||||
DataDogHandler,
|
||||
DatadogLoggingConfig,
|
||||
)
|
||||
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
|
||||
DynamicLoggingCache,
|
||||
)
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def datadog_env(monkeypatch):
|
||||
"""Set global DD env vars for the default/global logger."""
|
||||
monkeypatch.setenv("DD_API_KEY", "global_api_key")
|
||||
monkeypatch.setenv("DD_SITE", "us1.datadoghq.com")
|
||||
|
||||
|
||||
class TestDataDogLoggerCredentialKwargs:
|
||||
"""Test that DataDogLogger accepts credentials as kwargs."""
|
||||
|
||||
def test_init_with_explicit_credentials(self):
|
||||
"""Logger should use explicit kwargs instead of env vars."""
|
||||
with patch("asyncio.create_task"):
|
||||
logger = DataDogLogger(
|
||||
dd_api_key="team_api_key",
|
||||
dd_site="eu1.datadoghq.com",
|
||||
)
|
||||
|
||||
assert logger.DD_API_KEY == "team_api_key"
|
||||
assert "eu1.datadoghq.com" in logger.intake_url
|
||||
|
||||
def test_init_falls_back_to_env_vars(self, datadog_env):
|
||||
"""Logger should fall back to env vars when no kwargs provided."""
|
||||
with patch("asyncio.create_task"):
|
||||
logger = DataDogLogger()
|
||||
|
||||
assert logger.DD_API_KEY == "global_api_key"
|
||||
assert "us1.datadoghq.com" in logger.intake_url
|
||||
|
||||
def test_init_kwargs_override_env_vars(self, datadog_env):
|
||||
"""Explicit kwargs should take precedence over env vars."""
|
||||
with patch("asyncio.create_task"):
|
||||
logger = DataDogLogger(
|
||||
dd_api_key="override_key",
|
||||
dd_site="ap1.datadoghq.com",
|
||||
)
|
||||
|
||||
assert logger.DD_API_KEY == "override_key"
|
||||
assert "ap1.datadoghq.com" in logger.intake_url
|
||||
|
||||
def test_init_with_agent_credentials(self):
|
||||
"""Logger should use agent mode when dd_agent_host is provided."""
|
||||
with patch("asyncio.create_task"):
|
||||
logger = DataDogLogger(
|
||||
dd_agent_host="dd-agent.local",
|
||||
dd_agent_port="8125",
|
||||
dd_api_key="agent_api_key",
|
||||
)
|
||||
|
||||
assert "dd-agent.local:8125" in logger.intake_url
|
||||
assert logger.DD_API_KEY == "agent_api_key"
|
||||
|
||||
def test_init_raises_without_credentials(self, monkeypatch):
|
||||
"""Logger should raise if no credentials are available."""
|
||||
monkeypatch.delenv("DD_API_KEY", raising=False)
|
||||
monkeypatch.delenv("DD_SITE", raising=False)
|
||||
monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False)
|
||||
|
||||
with pytest.raises(Exception, match="DD_API_KEY"):
|
||||
with patch("asyncio.create_task"):
|
||||
DataDogLogger()
|
||||
|
||||
|
||||
class TestDataDogHandler:
|
||||
"""Test that DataDogHandler resolves the correct logger per team."""
|
||||
|
||||
def test_creates_team_logger_with_dynamic_credentials(self, datadog_env):
|
||||
"""Should create a new logger when team credentials are provided."""
|
||||
cache = DynamicLoggingCache()
|
||||
params = StandardCallbackDynamicParams(
|
||||
dd_api_key="team_a_key",
|
||||
dd_site="eu1.datadoghq.com",
|
||||
)
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result = DataDogHandler.get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params=params,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result.DD_API_KEY == "team_a_key"
|
||||
assert "eu1.datadoghq.com" in result.intake_url
|
||||
|
||||
def test_caches_team_logger(self, datadog_env):
|
||||
"""Same team credentials should return the same cached logger instance."""
|
||||
cache = DynamicLoggingCache()
|
||||
params = StandardCallbackDynamicParams(
|
||||
dd_api_key="team_b_key",
|
||||
dd_site="us5.datadoghq.com",
|
||||
)
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result1 = DataDogHandler.get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params=params,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
result2 = DataDogHandler.get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params=params,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result1 is result2
|
||||
|
||||
def test_different_teams_get_different_loggers(self, datadog_env):
|
||||
"""Different team credentials should create separate logger instances."""
|
||||
cache = DynamicLoggingCache()
|
||||
|
||||
params_a = StandardCallbackDynamicParams(
|
||||
dd_api_key="team_a_key",
|
||||
dd_site="us1.datadoghq.com",
|
||||
)
|
||||
params_b = StandardCallbackDynamicParams(
|
||||
dd_api_key="team_b_key",
|
||||
dd_site="eu1.datadoghq.com",
|
||||
)
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result_a = DataDogHandler.get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params=params_a,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
result_b = DataDogHandler.get_datadog_logger_for_request(
|
||||
standard_callback_dynamic_params=params_b,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result_a is not result_b
|
||||
assert result_a.DD_API_KEY == "team_a_key"
|
||||
assert result_b.DD_API_KEY == "team_b_key"
|
||||
|
||||
def test_request_blocked_callback_params_includes_dd(self):
|
||||
"""DD params should be blocked from request-level metadata (security)."""
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
_request_blocked_callback_params,
|
||||
)
|
||||
|
||||
assert "dd_api_key" in _request_blocked_callback_params
|
||||
assert "dd_site" in _request_blocked_callback_params
|
||||
assert "dd_agent_host" in _request_blocked_callback_params
|
||||
assert "dd_agent_port" in _request_blocked_callback_params
|
||||
|
||||
|
||||
class TestDynamicCredentialDetection:
|
||||
"""Test that _dynamic_datadog_credentials_are_passed works correctly."""
|
||||
|
||||
def test_no_credentials(self):
|
||||
params = StandardCallbackDynamicParams()
|
||||
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False
|
||||
|
||||
def test_dd_api_key_only(self):
|
||||
params = StandardCallbackDynamicParams(dd_api_key="key")
|
||||
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
|
||||
|
||||
def test_dd_site_only(self):
|
||||
params = StandardCallbackDynamicParams(dd_site="site")
|
||||
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
|
||||
|
||||
def test_dd_agent_host_only(self):
|
||||
params = StandardCallbackDynamicParams(dd_agent_host="host")
|
||||
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
|
||||
|
||||
|
||||
class TestStandardCallbackDynamicParamsIncludesDatadog:
|
||||
"""Verify that Datadog params are in the allow-list."""
|
||||
|
||||
def test_dd_params_in_annotations(self):
|
||||
annotations = StandardCallbackDynamicParams.__annotations__
|
||||
assert "dd_api_key" in annotations
|
||||
assert "dd_site" in annotations
|
||||
assert "dd_agent_host" in annotations
|
||||
assert "dd_agent_port" in annotations
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
"""
|
||||
Tests that the cost-tracking call sites tolerate ``server_tool_use`` being
|
||||
either a ``dict`` or a ``ServerToolUse`` pydantic instance.
|
||||
|
||||
See https://github.com/BerriAI/litellm/issues/26153.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
|
||||
StandardBuiltInToolCostTracking,
|
||||
_get_web_search_requests,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse, ServerToolUse, Usage
|
||||
|
||||
|
||||
class _UsageWithDictServerToolUse:
|
||||
"""
|
||||
Tiny stand-in that mimics the broken streaming-rebuild shape:
|
||||
``server_tool_use`` is a plain dict.
|
||||
"""
|
||||
|
||||
def __init__(self, server_tool_use):
|
||||
self.server_tool_use = server_tool_use
|
||||
self.prompt_tokens_details = None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_none():
|
||||
assert _get_web_search_requests(None) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict():
|
||||
assert _get_web_search_requests({"web_search_requests": 5}) == 5
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict_missing_key():
|
||||
assert _get_web_search_requests({}) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_pydantic():
|
||||
stu = ServerToolUse(web_search_requests=7)
|
||||
assert _get_web_search_requests(stu) == 7
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_pydantic_with_none_value():
|
||||
stu = ServerToolUse()
|
||||
assert _get_web_search_requests(stu) is None
|
||||
|
||||
|
||||
def test_response_object_includes_web_search_call_with_dict_server_tool_use():
|
||||
"""
|
||||
The exact bug: ``usage.server_tool_use`` is a dict and the check in
|
||||
``response_object_includes_web_search_call`` used to crash with
|
||||
``AttributeError``.
|
||||
"""
|
||||
response = ModelResponse()
|
||||
usage = _UsageWithDictServerToolUse({"web_search_requests": 2})
|
||||
|
||||
# Must not raise — and must correctly detect the web search call.
|
||||
result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
|
||||
response_object=response, usage=usage # type: ignore[arg-type]
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
def test_response_object_includes_web_search_call_with_pydantic_server_tool_use():
|
||||
response = ModelResponse()
|
||||
usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2))
|
||||
|
||||
result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
|
||||
response_object=response, usage=usage # type: ignore[arg-type]
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
def test_response_object_includes_web_search_call_with_none_server_tool_use():
|
||||
response = ModelResponse()
|
||||
usage = _UsageWithDictServerToolUse(None)
|
||||
|
||||
result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
|
||||
response_object=response, usage=usage # type: ignore[arg-type]
|
||||
)
|
||||
assert result is False
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
"""
|
||||
Regression tests for https://github.com/BerriAI/litellm/issues/26153
|
||||
|
||||
``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain
|
||||
``dict`` when reconstructing a streaming response. Downstream cost-calculation
|
||||
code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call``
|
||||
and ``get_cost_for_anthropic_web_search``) accesses
|
||||
``usage.server_tool_use.web_search_requests`` as an attribute, which raised
|
||||
``AttributeError: 'dict' object has no attribute 'web_search_requests'``.
|
||||
|
||||
These tests reconstruct streaming chunks for an Anthropic-style web_search
|
||||
response and assert:
|
||||
|
||||
1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for
|
||||
``usage.server_tool_use``.
|
||||
2. ``completion_cost`` runs end-to-end on the rebuilt response without
|
||||
raising ``AttributeError``.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm import completion_cost, stream_chunk_builder
|
||||
from litellm.types.utils import (
|
||||
Delta,
|
||||
ModelResponseStream,
|
||||
ServerToolUse,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
def _make_text_chunk(text: str) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-test-26153",
|
||||
created=1700000000,
|
||||
model="claude-3-haiku-20240307",
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(role="assistant", content=text),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream:
|
||||
"""Final chunk where server_tool_use is a *dict* — reproduces the bug shape."""
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-test-26153",
|
||||
created=1700000000,
|
||||
model="claude-3-haiku-20240307",
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
)
|
||||
],
|
||||
usage=Usage(
|
||||
prompt_tokens=42,
|
||||
completion_tokens=11,
|
||||
total_tokens=53,
|
||||
# NOTE: passed as a dict on purpose — this is the shape that
|
||||
# historically slipped through stream_chunk_builder unchanged.
|
||||
server_tool_use={"web_search_requests": 3},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic():
|
||||
"""
|
||||
Regression: stream_chunk_builder must produce ServerToolUse, not dict.
|
||||
"""
|
||||
chunks = [
|
||||
_make_text_chunk("Otters "),
|
||||
_make_text_chunk("are great."),
|
||||
_make_finish_chunk_with_usage_dict_server_tool_use(),
|
||||
]
|
||||
|
||||
rebuilt = stream_chunk_builder(chunks)
|
||||
|
||||
assert rebuilt is not None
|
||||
assert rebuilt.usage is not None # type: ignore[attr-defined]
|
||||
server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined]
|
||||
|
||||
assert (
|
||||
server_tool_use is not None
|
||||
), "server_tool_use should be carried through from the final chunk"
|
||||
assert isinstance(server_tool_use, ServerToolUse), (
|
||||
f"expected ServerToolUse, got {type(server_tool_use).__name__}: "
|
||||
f"{server_tool_use!r}"
|
||||
)
|
||||
# Attribute access must not raise (this is exactly what was broken).
|
||||
assert server_tool_use.web_search_requests == 3
|
||||
|
||||
|
||||
def test_completion_cost_does_not_raise_on_streaming_web_search_response():
|
||||
"""
|
||||
Regression: completion_cost(...) must not raise AttributeError when the
|
||||
response was reconstructed by stream_chunk_builder from a streaming
|
||||
Anthropic web_search call.
|
||||
"""
|
||||
chunks = [
|
||||
_make_text_chunk("hello"),
|
||||
_make_finish_chunk_with_usage_dict_server_tool_use(),
|
||||
]
|
||||
|
||||
rebuilt = stream_chunk_builder(chunks)
|
||||
assert rebuilt is not None
|
||||
|
||||
# The exact dollar amount depends on the model-pricing table; what matters
|
||||
# for this regression is that it does NOT raise AttributeError on
|
||||
# `dict has no attribute 'web_search_requests'`.
|
||||
try:
|
||||
cost = completion_cost(completion_response=rebuilt)
|
||||
except AttributeError as e: # pragma: no cover - regression guard
|
||||
pytest.fail(
|
||||
"completion_cost raised AttributeError after stream_chunk_builder "
|
||||
f"(issue #26153 regression): {e}"
|
||||
)
|
||||
|
||||
assert isinstance(cost, (int, float))
|
||||
|
|
@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search():
|
|||
assert usage.prompt_tokens == 50
|
||||
assert usage.completion_tokens == 27
|
||||
assert usage.total_tokens == 77
|
||||
assert usage.server_tool_use["web_search_requests"] == 2
|
||||
# server_tool_use must be a ServerToolUse pydantic so downstream cost-calc
|
||||
# (which uses attribute access) works. See issue #26153.
|
||||
assert isinstance(usage.server_tool_use, ServerToolUse)
|
||||
assert usage.server_tool_use.web_search_requests == 2
|
||||
|
||||
|
||||
def test_sort_chunks_handles_dict_hidden_params_created_at():
|
||||
|
|
|
|||
|
|
@ -0,0 +1,94 @@
|
|||
"""
|
||||
Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use``
|
||||
being either a ``dict`` or a ``ServerToolUse`` pydantic instance.
|
||||
|
||||
See https://github.com/BerriAI/litellm/issues/26153.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.llms.anthropic.cost_calculation import (
|
||||
_get_web_search_requests,
|
||||
get_cost_for_anthropic_web_search,
|
||||
)
|
||||
from litellm.types.utils import ModelInfo, ServerToolUse
|
||||
|
||||
|
||||
class _UsageWithServerToolUse:
|
||||
def __init__(self, server_tool_use):
|
||||
self.server_tool_use = server_tool_use
|
||||
|
||||
|
||||
def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo:
|
||||
info: ModelInfo = { # type: ignore[typeddict-item]
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_low": cost_per_query,
|
||||
"search_context_size_medium": cost_per_query,
|
||||
"search_context_size_high": cost_per_query,
|
||||
}
|
||||
}
|
||||
return info
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_none():
|
||||
assert _get_web_search_requests(None) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict():
|
||||
assert _get_web_search_requests({"web_search_requests": 4}) == 4
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict_missing_key():
|
||||
assert _get_web_search_requests({}) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_pydantic():
|
||||
assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2
|
||||
|
||||
|
||||
def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use():
|
||||
"""
|
||||
Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and
|
||||
direct attribute access on it raised ``AttributeError``.
|
||||
"""
|
||||
usage = _UsageWithServerToolUse({"web_search_requests": 3})
|
||||
info = _make_model_info(cost_per_query=0.01)
|
||||
|
||||
cost = get_cost_for_anthropic_web_search(
|
||||
model_info=info, usage=usage # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.03)
|
||||
|
||||
|
||||
def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use():
|
||||
usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3))
|
||||
info = _make_model_info(cost_per_query=0.01)
|
||||
|
||||
cost = get_cost_for_anthropic_web_search(
|
||||
model_info=info, usage=usage # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.03)
|
||||
|
||||
|
||||
def test_get_cost_for_anthropic_web_search_with_none_server_tool_use():
|
||||
usage = _UsageWithServerToolUse(None)
|
||||
info = _make_model_info(cost_per_query=0.01)
|
||||
|
||||
cost = get_cost_for_anthropic_web_search(
|
||||
model_info=info, usage=usage # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
def test_get_cost_for_anthropic_web_search_with_no_usage():
|
||||
info = _make_model_info(cost_per_query=0.01)
|
||||
cost = get_cost_for_anthropic_web_search(model_info=info, usage=None)
|
||||
assert cost == 0.0
|
||||
|
|
@ -12,6 +12,11 @@ import sys
|
|||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
import pytest
|
||||
from botocore.exceptions import (
|
||||
ConnectTimeoutError,
|
||||
PartialCredentialsError,
|
||||
ProfileNotFound,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock_mantle.responses.transformation import (
|
||||
|
|
@ -114,16 +119,15 @@ class TestBedrockMantleResponsesAuth:
|
|||
)
|
||||
assert headers["Authorization"] == "Bearer bearer-key"
|
||||
|
||||
def test_missing_key_raises(self, monkeypatch):
|
||||
def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch):
|
||||
# SigV4 may still apply, so validate_environment must defer instead of raising.
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
with pytest.raises(ValueError, match="Bedrock Mantle API key"):
|
||||
cfg.validate_environment(
|
||||
headers={},
|
||||
model="openai.gpt-5.5",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
)
|
||||
headers = cfg.validate_environment(
|
||||
headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams()
|
||||
)
|
||||
assert "Authorization" not in headers
|
||||
|
||||
def test_custom_llm_provider(self):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
|
|
@ -261,6 +265,386 @@ def local_cost_map(monkeypatch):
|
|||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesSigV4:
|
||||
def test_bearer_short_circuits_without_credentials(self, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(
|
||||
side_effect=AssertionError("get_credentials must not run for bearer auth")
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
headers, signed_body = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key="bearer-from-config",
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer bearer-from-config"
|
||||
assert signed_body == b'{"input": "hi"}'
|
||||
signer.get_credentials.assert_not_called()
|
||||
|
||||
def test_bearer_resolved_from_mantle_env_key(self, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(
|
||||
side_effect=AssertionError("get_credentials must not run for bearer auth")
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer env-bearer"
|
||||
|
||||
def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch):
|
||||
# The passed api_key (e.g. litellm_params.api_key) must win over the env
|
||||
# bearer; a reordered precedence chain would silently use the wrong token.
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(
|
||||
side_effect=AssertionError("get_credentials must not run for bearer auth")
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key="arg-bearer",
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer arg-bearer"
|
||||
signer.get_credentials.assert_not_called()
|
||||
|
||||
def test_access_key_produces_sigv4_headers(self, monkeypatch):
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
|
||||
headers, signed_body = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
"aws_session_token": "session-token-test",
|
||||
"aws_region_name": "us-east-2",
|
||||
},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert "Credential=AKIAEXAMPLE/" in headers["Authorization"]
|
||||
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
|
||||
assert "X-Amz-Date" in headers
|
||||
assert headers["X-Amz-Security-Token"] == "session-token-test"
|
||||
assert signed_body == b'{"input": "hi"}'
|
||||
|
||||
def test_assume_role_path_produces_sigv4_headers(self, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
from botocore.credentials import Credentials
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(
|
||||
return_value=Credentials(
|
||||
access_key="ASIAEXAMPLE",
|
||||
secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk",
|
||||
token="assumed-session-token",
|
||||
)
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={
|
||||
"aws_role_name": "arn:aws:iam::000000000000:role/test-role",
|
||||
"aws_session_name": "litellm-test",
|
||||
"aws_region_name": "us-east-2",
|
||||
},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
signer.get_credentials.assert_called_once()
|
||||
call = signer.get_credentials.call_args.kwargs
|
||||
assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role"
|
||||
assert call["aws_session_name"] == "litellm-test"
|
||||
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
|
||||
|
||||
def test_signed_body_matches_final_data_after_normalize(self, monkeypatch):
|
||||
"""Core regression: the signed bytes must equal the bytes actually sent.
|
||||
|
||||
Sign the *final* data dict and assert the returned signed_body decodes to
|
||||
exactly that dict, so a later change to the data would break the SigV4 hash.
|
||||
"""
|
||||
import json
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16}
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
|
||||
_, signed_body = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
"aws_region_name": "us-east-2",
|
||||
},
|
||||
request_data=final_data,
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
assert signed_body is not None
|
||||
assert json.loads(signed_body) == final_data
|
||||
|
||||
def test_region_comes_from_optional_params(self, monkeypatch):
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
|
||||
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
"aws_region_name": "eu-west-1",
|
||||
},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
|
||||
|
||||
def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch):
|
||||
"""Adversarial-review regression: a caller-supplied aws_region_name (no region
|
||||
env set) must shape BOTH the URL host and the SigV4 credential scope, or the
|
||||
request is signed for one region and sent to another -> 401.
|
||||
"""
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
params = {
|
||||
"aws_region_name": "ap-southeast-2",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
}
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
|
||||
url = cfg.get_complete_url(api_base=None, litellm_params=params)
|
||||
assert (
|
||||
url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses"
|
||||
)
|
||||
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params=params,
|
||||
request_data={"input": "hi"},
|
||||
api_base=url,
|
||||
api_key=None,
|
||||
)
|
||||
assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"]
|
||||
|
||||
def test_injected_default_region_base_does_not_override_aws_region_name(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""2nd-round adversarial regression: responses/main.py auto-injects
|
||||
litellm_params.api_base = https://bedrock-mantle.<DEFAULT>.api.aws/v1 (default
|
||||
region, ignoring aws_region_name). The config must still pin BOTH the URL host
|
||||
and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive
|
||||
'resolve region only when api_base is None' fix would fail this test.
|
||||
"""
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region
|
||||
params = {
|
||||
"aws_region_name": "us-east-2", # what the caller actually wants
|
||||
"api_base": injected_base,
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
}
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
|
||||
url = cfg.get_complete_url(api_base=injected_base, litellm_params=params)
|
||||
assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params=params,
|
||||
request_data={"input": "hi"},
|
||||
api_base=url,
|
||||
api_key=None,
|
||||
)
|
||||
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
|
||||
assert "us-east-1" not in headers["Authorization"]
|
||||
|
||||
def test_custom_proxy_host_is_preserved(self, monkeypatch):
|
||||
"""A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten
|
||||
to a bedrock-mantle host. Only standard Mantle hosts are region-pinned.
|
||||
"""
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
url = cfg.get_complete_url(
|
||||
api_base="https://mantle-proxy.internal.example/openai/v1",
|
||||
litellm_params={"aws_region_name": "us-east-2"},
|
||||
)
|
||||
assert url == "https://mantle-proxy.internal.example/openai/v1/responses"
|
||||
|
||||
def test_caller_authorization_does_not_override_sigv4(self, monkeypatch):
|
||||
"""Adversarial-review regression: a caller-supplied Authorization header (e.g.
|
||||
from extra_headers, surviving the relaxed validate_environment) must not clobber
|
||||
the SigV4 Authorization that _sign_request would otherwise restore.
|
||||
"""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM())
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={"Authorization": "Bearer stale-caller-token"},
|
||||
optional_params={
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
|
||||
"aws_region_name": "us-east-2",
|
||||
},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert "Bearer stale-caller-token" not in headers["Authorization"]
|
||||
|
||||
def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
from botocore.exceptions import NoCredentialsError
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(side_effect=NoCredentialsError())
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
with pytest.raises(ValueError) as exc:
|
||||
cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={"aws_region_name": "us-east-2"},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
msg = str(exc.value)
|
||||
assert "Bearer" in msg
|
||||
assert "SigV4" in msg or "IAM" in msg
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cred_error",
|
||||
[
|
||||
PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"),
|
||||
ProfileNotFound(profile="missing-profile"),
|
||||
],
|
||||
)
|
||||
def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error):
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(side_effect=cred_error)
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
with pytest.raises(ValueError) as exc:
|
||||
cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={"aws_region_name": "us-east-2"},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
msg = str(exc.value)
|
||||
assert "Bearer" in msg
|
||||
assert "SigV4" in msg or "IAM" in msg
|
||||
|
||||
def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch):
|
||||
# An AssumeRole / web-identity flow hits STS over the network, so a transient
|
||||
# connection error must surface as itself, not be rewritten into the
|
||||
# "no usable AWS credentials" message that would send the user to fix the
|
||||
# wrong thing.
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
|
||||
signer = BaseAWSLLM()
|
||||
signer.get_credentials = MagicMock(
|
||||
side_effect=ConnectTimeoutError(
|
||||
endpoint_url="https://sts.us-east-2.amazonaws.com"
|
||||
)
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer)
|
||||
|
||||
with pytest.raises(ConnectTimeoutError):
|
||||
cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={
|
||||
"aws_role_name": "arn:aws:iam::000000000000:role/test-role",
|
||||
"aws_region_name": "us-east-2",
|
||||
},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesPricing:
|
||||
def test_gpt_5_5_pricing_and_mode(self, local_cost_map):
|
||||
info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5")
|
||||
|
|
|
|||
|
|
@ -742,3 +742,241 @@ async def test_anthropic_post_retry_reserializes_mutated_body():
|
|||
assert first_sent == prebuilt # attempt 0 used prebuilt
|
||||
assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized
|
||||
assert "MUTATED" in second_sent # ... the mutated body
|
||||
|
||||
|
||||
def test_base_responses_config_sign_request_is_noop_by_default():
|
||||
"""Default responses sign_request must be a no-op: unchanged headers, no signed body.
|
||||
|
||||
Guards the 15 existing responses providers from accidental signing when the
|
||||
handler starts calling sign_request.
|
||||
"""
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
|
||||
cfg = OpenAIResponsesAPIConfig()
|
||||
headers = {"Authorization": "Bearer sk-existing"}
|
||||
out_headers, signed_body = cfg.sign_request(
|
||||
headers=headers,
|
||||
optional_params={},
|
||||
request_data={"input": "hi"},
|
||||
api_base="https://api.openai.com/v1/responses",
|
||||
)
|
||||
assert out_headers == {"Authorization": "Bearer sk-existing"}
|
||||
assert signed_body is None
|
||||
|
||||
|
||||
def _make_responses_handler_call(signed_body):
|
||||
"""Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider
|
||||
config + sync client, returning the kwargs the client.post was called with.
|
||||
|
||||
signed_body=None simulates a no-op (non-signing) provider; bytes simulates a
|
||||
signing provider (e.g. Bedrock Mantle).
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.validate_environment.return_value = {}
|
||||
provider_config.get_complete_url.return_value = (
|
||||
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
)
|
||||
provider_config.transform_responses_api_request.return_value = {"input": "hi"}
|
||||
provider_config.should_fake_stream.return_value = False
|
||||
provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body)
|
||||
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = MagicMock()
|
||||
|
||||
handler = BaseLLMHTTPHandler()
|
||||
handler.response_api_handler(
|
||||
model="openai.gpt-5.5",
|
||||
input="hi",
|
||||
responses_api_provider_config=provider_config,
|
||||
response_api_optional_request_params={},
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
|
||||
logging_obj=MagicMock(),
|
||||
client=mock_client,
|
||||
_is_async=False,
|
||||
)
|
||||
return mock_client.post.call_args.kwargs
|
||||
|
||||
|
||||
def test_responses_handler_sends_json_when_not_signed():
|
||||
"""No-op provider (signed_body is None) -> handler posts json=data, no data= bytes."""
|
||||
kwargs = _make_responses_handler_call(signed_body=None)
|
||||
assert kwargs.get("json") == {"input": "hi"}
|
||||
assert "data" not in kwargs
|
||||
|
||||
|
||||
def test_responses_handler_sends_signed_bytes_when_signed():
|
||||
"""Signing provider -> handler posts the exact signed bytes via data=, not json=."""
|
||||
kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}')
|
||||
assert kwargs.get("data") == b'{"input": "hi"}'
|
||||
assert "json" not in kwargs
|
||||
assert kwargs["headers"] == {"X-Signed": "1"}
|
||||
|
||||
|
||||
def test_responses_handler_signs_after_fake_stream_prep_strips_stream():
|
||||
"""Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT.
|
||||
|
||||
In the streaming + fake-stream path the handler first runs
|
||||
_prepare_fake_stream_request, which pops "stream" out of the body, and only
|
||||
then calls sign_request. If signing ran before that pop, the signed body
|
||||
would still carry "stream" while the body sent over the wire would not,
|
||||
producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment.
|
||||
We snapshot request_data at sign time and assert "stream" is already gone.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.validate_environment.return_value = {}
|
||||
provider_config.get_complete_url.return_value = (
|
||||
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
)
|
||||
provider_config.transform_responses_api_request.return_value = {
|
||||
"input": "hi",
|
||||
"stream": True,
|
||||
}
|
||||
provider_config.should_fake_stream.return_value = True
|
||||
provider_config.transform_response_api_response.return_value = ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=0,
|
||||
output=[],
|
||||
status="completed",
|
||||
model="openai.gpt-5.5",
|
||||
)
|
||||
|
||||
captured = {}
|
||||
|
||||
def _capture_sign(**kwargs):
|
||||
captured["request_data"] = dict(kwargs["request_data"])
|
||||
return ({"X-Signed": "1"}, b'{"input": "hi"}')
|
||||
|
||||
provider_config.sign_request.side_effect = _capture_sign
|
||||
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = MagicMock()
|
||||
|
||||
handler = BaseLLMHTTPHandler()
|
||||
handler.response_api_handler(
|
||||
model="openai.gpt-5.5",
|
||||
input="hi",
|
||||
responses_api_provider_config=provider_config,
|
||||
response_api_optional_request_params={"stream": True},
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
|
||||
logging_obj=MagicMock(),
|
||||
client=mock_client,
|
||||
_is_async=False,
|
||||
fake_stream=True,
|
||||
)
|
||||
|
||||
assert "stream" not in captured["request_data"]
|
||||
assert "input" in captured["request_data"]
|
||||
|
||||
post_kwargs = mock_client.post.call_args.kwargs
|
||||
assert post_kwargs.get("data") == b'{"input": "hi"}'
|
||||
assert "json" not in post_kwargs
|
||||
assert "stream" in post_kwargs
|
||||
|
||||
|
||||
def _make_compact_handler_call(signed_body, is_async):
|
||||
"""Drive (async_)compact_response_api_handler with a fully mocked provider config
|
||||
+ client, returning the kwargs the client.post was called with.
|
||||
|
||||
signed_body=None simulates a no-op (non-signing) provider; bytes simulates a
|
||||
signing provider (e.g. Bedrock Mantle SigV4 / bearer).
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact"
|
||||
provider_config = MagicMock()
|
||||
provider_config.validate_environment.return_value = {}
|
||||
provider_config.get_complete_url.return_value = (
|
||||
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
)
|
||||
provider_config.transform_compact_response_api_request.return_value = (
|
||||
compact_url,
|
||||
{"model": "openai.gpt-5.5", "input": "hi"},
|
||||
)
|
||||
provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body)
|
||||
provider_config.transform_compact_response_api_response.return_value = "ok"
|
||||
|
||||
spec = AsyncHTTPHandler if is_async else HTTPHandler
|
||||
mock_client = MagicMock(spec=spec)
|
||||
if is_async:
|
||||
mock_client.post = AsyncMock(return_value=MagicMock())
|
||||
else:
|
||||
mock_client.post.return_value = MagicMock()
|
||||
|
||||
handler = BaseLLMHTTPHandler()
|
||||
result = handler.compact_response_api_handler(
|
||||
model="openai.gpt-5.5",
|
||||
input="hi",
|
||||
responses_api_provider_config=provider_config,
|
||||
response_api_optional_request_params={},
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"),
|
||||
logging_obj=MagicMock(),
|
||||
client=mock_client,
|
||||
_is_async=is_async,
|
||||
)
|
||||
if is_async:
|
||||
asyncio.run(result)
|
||||
return provider_config, mock_client.post.call_args.kwargs
|
||||
|
||||
|
||||
def test_compact_handler_sends_json_when_not_signed():
|
||||
"""No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes."""
|
||||
provider_config, kwargs = _make_compact_handler_call(
|
||||
signed_body=None, is_async=False
|
||||
)
|
||||
provider_config.sign_request.assert_called_once()
|
||||
assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"}
|
||||
assert "data" not in kwargs
|
||||
|
||||
|
||||
def test_compact_handler_sends_signed_bytes_when_signed():
|
||||
"""Signing provider on compact -> posts the signed bytes via data=, not json=.
|
||||
|
||||
Regression for the adversarial-review finding that /responses/compact bypassed
|
||||
the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies.
|
||||
"""
|
||||
provider_config, kwargs = _make_compact_handler_call(
|
||||
signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False
|
||||
)
|
||||
assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}'
|
||||
assert "json" not in kwargs
|
||||
assert kwargs["headers"] == {"X-Signed": "1"}
|
||||
# signing must use the compact endpoint as api_base, not the create URL
|
||||
assert provider_config.sign_request.call_args.kwargs["api_base"].endswith(
|
||||
"/openai/v1/responses/compact"
|
||||
)
|
||||
|
||||
|
||||
def test_async_compact_handler_sends_signed_bytes_when_signed():
|
||||
"""Async compact must sign identically to sync (same omission in the async twin)."""
|
||||
provider_config, kwargs = _make_compact_handler_call(
|
||||
signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True
|
||||
)
|
||||
assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}'
|
||||
assert "json" not in kwargs
|
||||
assert kwargs["headers"] == {"X-Signed": "1"}
|
||||
|
||||
|
||||
def test_async_compact_handler_sends_json_when_not_signed():
|
||||
"""Async no-op provider on compact -> posts json=data, no data= bytes."""
|
||||
_provider_config, kwargs = _make_compact_handler_call(
|
||||
signed_body=None, is_async=True
|
||||
)
|
||||
assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"}
|
||||
assert "data" not in kwargs
|
||||
|
|
|
|||
|
|
@ -658,12 +658,11 @@ class TestMCPOAuth2AuthFlow:
|
|||
|
||||
async def test_oauth2_token_in_authorization_header_fallback(self):
|
||||
"""
|
||||
When only Authorization header is present with a non-LiteLLM OAuth2 token
|
||||
AND the target server is operator-configured for ``auth_type=oauth2``,
|
||||
auth should fall back to permissive mode (OAuth2 passthrough).
|
||||
When only the Authorization header is present with a non-LiteLLM OAuth2
|
||||
token AND the target server delegates auth to upstream, LiteLLM skips its
|
||||
own validation entirely (so the upstream token is never mistaken for a
|
||||
virtual key) and forwards the bearer upstream.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
scope = {
|
||||
|
|
@ -675,17 +674,16 @@ class TestMCPOAuth2AuthFlow:
|
|||
],
|
||||
}
|
||||
|
||||
async def mock_user_api_key_auth_fails(api_key, request):
|
||||
raise HTTPException(status_code=401, detail="Invalid API key")
|
||||
|
||||
oauth2_server = MagicMock()
|
||||
oauth2_server.auth_type = MCPAuth.oauth2
|
||||
oauth2_server.delegate_auth_to_upstream = True
|
||||
oauth2_server.has_client_credentials = False
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
|
||||
side_effect=mock_user_api_key_auth_fails,
|
||||
),
|
||||
new_callable=AsyncMock,
|
||||
) as mock_auth,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
|
||||
) as mock_mgr,
|
||||
|
|
@ -700,10 +698,10 @@ class TestMCPOAuth2AuthFlow:
|
|||
raw_headers,
|
||||
) = await MCPRequestHandler.process_mcp_request(scope)
|
||||
|
||||
# Should succeed with default UserAPIKeyAuth (OAuth2 fallback)
|
||||
assert auth_result is not None
|
||||
assert isinstance(auth_result, UserAPIKeyAuth)
|
||||
# OAuth2 headers should contain the token for upstream forwarding
|
||||
# The upstream token is never validated as a LiteLLM key ...
|
||||
mock_auth.assert_not_called()
|
||||
# ... and is preserved for upstream forwarding.
|
||||
assert (
|
||||
oauth2_headers.get("Authorization")
|
||||
== "Bearer atlassian-oauth2-access-token-xyz"
|
||||
|
|
@ -813,11 +811,12 @@ class TestMCPOAuth2AuthFlow:
|
|||
await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert exc_info.value.status_code == 500
|
||||
|
||||
async def test_proxy_exception_oauth2_fallback(self):
|
||||
async def test_proxy_exception_non_delegate_oauth2_propagates(self):
|
||||
"""
|
||||
user_api_key_auth raises ProxyException (not HTTPException) in production.
|
||||
The OAuth2 fallback must catch ProxyException with code 401/403 too,
|
||||
but only when the target server is operator-configured for ``auth_type=oauth2``.
|
||||
Production raises ProxyException (not HTTPException) on auth failure. For
|
||||
a non-delegate oauth2 server the bearer is treated as a LiteLLM credential
|
||||
and a 401 must propagate as a real auth error, not be exchanged for an
|
||||
anonymous upstream-passthrough session.
|
||||
"""
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
|
@ -841,6 +840,8 @@ class TestMCPOAuth2AuthFlow:
|
|||
|
||||
oauth2_server = MagicMock()
|
||||
oauth2_server.auth_type = MCPAuth.oauth2
|
||||
oauth2_server.delegate_auth_to_upstream = False
|
||||
oauth2_server.is_oauth_passthrough = False
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -852,22 +853,9 @@ class TestMCPOAuth2AuthFlow:
|
|||
) as mock_mgr,
|
||||
):
|
||||
mock_mgr.get_mcp_server_by_name.return_value = oauth2_server
|
||||
(
|
||||
auth_result,
|
||||
mcp_auth_header,
|
||||
mcp_servers,
|
||||
mcp_server_auth_headers,
|
||||
oauth2_headers,
|
||||
raw_headers,
|
||||
) = await MCPRequestHandler.process_mcp_request(scope)
|
||||
|
||||
# Should succeed with default UserAPIKeyAuth (OAuth2 fallback)
|
||||
assert auth_result is not None
|
||||
assert isinstance(auth_result, UserAPIKeyAuth)
|
||||
assert (
|
||||
oauth2_headers.get("Authorization")
|
||||
== "Bearer atlassian-oauth2-access-token-xyz"
|
||||
)
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert str(exc_info.value.code) == "401"
|
||||
|
||||
async def test_proxy_exception_non_auth_still_raises(self):
|
||||
"""
|
||||
|
|
@ -1355,11 +1343,15 @@ class TestMCPOAuth2FallbackTargetGating:
|
|||
await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
async def test_fallback_allowed_when_target_is_oauth2_mode(self):
|
||||
async def test_non_delegate_oauth2_does_not_fall_back_to_anonymous(self):
|
||||
"""
|
||||
Operator-configured OAuth2 passthrough still works: target server has
|
||||
``auth_type=oauth2`` → failed LiteLLM auth falls back to anonymous so
|
||||
the bearer can be forwarded to upstream.
|
||||
An ``auth_type=oauth2`` server that has NOT opted into
|
||||
``delegate_auth_to_upstream`` must not exchange a failed LiteLLM auth for
|
||||
an anonymous session: forwarding an arbitrary bearer upstream is only
|
||||
allowed once the operator explicitly delegates auth. A failed validation
|
||||
here is a genuine 401 and propagates (which is also what keeps the
|
||||
success-path trace free of a phantom 401, since no doomed validation runs
|
||||
for a delegated server).
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -1389,8 +1381,9 @@ class TestMCPOAuth2FallbackTargetGating:
|
|||
mock_mgr.get_mcp_server_by_name.return_value = (
|
||||
TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.oauth2)
|
||||
)
|
||||
auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert isinstance(auth_result, UserAPIKeyAuth)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
async def test_fallback_allowed_when_target_is_passthrough(self):
|
||||
"""
|
||||
|
|
@ -1668,19 +1661,16 @@ class TestMCPDelegateAuthToUpstream:
|
|||
assert isinstance(auth_result, UserAPIKeyAuth)
|
||||
mock_auth.assert_not_called()
|
||||
|
||||
async def test_delegate_with_upstream_token_in_authorization_falls_back_to_anonymous(
|
||||
async def test_delegate_with_upstream_token_in_authorization_skips_litellm_auth(
|
||||
self,
|
||||
):
|
||||
"""
|
||||
oauth2 + delegate_auth_to_upstream=True with an upstream OAuth token in
|
||||
``Authorization`` (not a LiteLLM key): LiteLLM auth is attempted first
|
||||
(and fails), then the existing oauth2 fallback returns anonymous so the
|
||||
bearer is forwarded upstream untouched. The delegate branch itself does
|
||||
not fire when Authorization is present — that is what protects spend
|
||||
tracking for callers using Authorization-style LiteLLM keys.
|
||||
``Authorization``: the delegate gate fires before any LiteLLM validation,
|
||||
so ``user_api_key_auth`` is never called and the bearer is forwarded
|
||||
upstream untouched. Skipping the doomed validation is what keeps a tool
|
||||
call that actually succeeds from carrying a phantom 401 auth span.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
scope = {
|
||||
|
|
@ -1690,14 +1680,11 @@ class TestMCPDelegateAuthToUpstream:
|
|||
"headers": [(b"authorization", b"Bearer upstream-pkce-token")],
|
||||
}
|
||||
|
||||
async def mock_user_api_key_auth_fails(api_key, request):
|
||||
raise HTTPException(status_code=401, detail="Invalid API key")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
|
||||
side_effect=mock_user_api_key_auth_fails,
|
||||
),
|
||||
new_callable=AsyncMock,
|
||||
) as mock_auth,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
|
||||
) as mock_mgr,
|
||||
|
|
@ -1718,6 +1705,7 @@ class TestMCPDelegateAuthToUpstream:
|
|||
) = await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert isinstance(auth_result, UserAPIKeyAuth)
|
||||
assert oauth2_headers.get("Authorization") == "Bearer upstream-pkce-token"
|
||||
mock_auth.assert_not_called()
|
||||
|
||||
async def test_delegate_off_still_requires_litellm_auth(self):
|
||||
"""
|
||||
|
|
@ -1912,12 +1900,15 @@ class TestMCPDelegateAuthToUpstream:
|
|||
assert auth_result.user_id == "real-user"
|
||||
mock_auth.assert_called_once()
|
||||
|
||||
async def test_litellm_key_via_authorization_header_not_bypassed(self):
|
||||
async def test_authorization_bearer_on_delegate_server_treated_as_upstream(self):
|
||||
"""
|
||||
Regression: a LiteLLM key sent via the secondary ``Authorization`` header
|
||||
(e.g. ``Authorization: Bearer sk-...``) must still trigger normal auth
|
||||
and not be silently swallowed by the delegate bypass — otherwise spend
|
||||
tracking and rate limiting are skipped for those callers.
|
||||
On a delegate server the ``Authorization`` header is, by contract, an
|
||||
upstream token rather than a LiteLLM key — even when it is sk-shaped. It
|
||||
is forwarded upstream without LiteLLM validation, so ``user_api_key_auth``
|
||||
is not called and no LiteLLM identity is resolved. Callers who need
|
||||
LiteLLM identity / spend tracking on a delegate server must supply
|
||||
``x-litellm-api-key`` (see
|
||||
test_explicit_litellm_key_takes_precedence_over_delegate).
|
||||
"""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
|
|
@ -1944,10 +1935,18 @@ class TestMCPDelegateAuthToUpstream:
|
|||
delegate_auth_to_upstream=True,
|
||||
)
|
||||
)
|
||||
auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope)
|
||||
(
|
||||
auth_result,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
oauth2_headers,
|
||||
_,
|
||||
) = await MCPRequestHandler.process_mcp_request(scope)
|
||||
assert isinstance(auth_result, UserAPIKeyAuth)
|
||||
assert auth_result.user_id == "real-user"
|
||||
mock_auth.assert_called_once()
|
||||
assert auth_result.user_id is None
|
||||
assert oauth2_headers.get("Authorization") == "Bearer sk-1234"
|
||||
mock_auth.assert_not_called()
|
||||
|
||||
async def test_delegate_ignored_for_client_credentials_server(self):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Tests for MCP OAuth discoverable endpoints"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -2661,3 +2662,74 @@ async def test_token_endpoint_sets_no_store_cache_control():
|
|||
|
||||
assert response.headers["cache-control"] == "no-store"
|
||||
assert response.headers["pragma"] == "no-cache"
|
||||
|
||||
|
||||
async def _exchange_with_upstream_token_response(upstream_body):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
exchange_token_with_server,
|
||||
)
|
||||
from litellm.proxy._types import MCPTransport
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="t",
|
||||
name="t",
|
||||
server_name="t",
|
||||
alias="t",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id="cid",
|
||||
client_secret="cs",
|
||||
authorization_url="https://provider.com/oauth/authorize",
|
||||
token_url="https://provider.com/oauth/token",
|
||||
)
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://litellm.example.com/"
|
||||
mock_request.headers = {}
|
||||
|
||||
fake_http_response = MagicMock()
|
||||
fake_http_response.json.return_value = upstream_body
|
||||
fake_http_response.raise_for_status = MagicMock()
|
||||
fake_http_client = MagicMock()
|
||||
fake_http_client.post = AsyncMock(return_value=fake_http_response)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=fake_http_client,
|
||||
):
|
||||
response = await exchange_token_with_server(
|
||||
request=mock_request,
|
||||
mcp_server=server,
|
||||
grant_type="authorization_code",
|
||||
code="c",
|
||||
redirect_uri="http://127.0.0.1:3000/cb",
|
||||
client_id="cid",
|
||||
client_secret=None,
|
||||
code_verifier=None,
|
||||
)
|
||||
return json.loads(response.body)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_exchange_omits_expires_in_when_upstream_omits_it():
|
||||
"""A provider that issues a non-expiring token (e.g. Slack without token
|
||||
rotation) returns no ``expires_in``. The exchange must mirror that and omit
|
||||
``expires_in`` rather than fabricate a 1-hour TTL, so the stored credential
|
||||
is treated as non-expiring instead of dying after an hour."""
|
||||
body = await _exchange_with_upstream_token_response(
|
||||
{"access_token": "tok", "token_type": "Bearer"}
|
||||
)
|
||||
assert "expires_in" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_exchange_passes_through_upstream_expires_in():
|
||||
"""When the provider does send ``expires_in`` (e.g. Slack with token
|
||||
rotation), the exchange forwards the real value unchanged."""
|
||||
body = await _exchange_with_upstream_token_response(
|
||||
{"access_token": "tok", "token_type": "Bearer", "expires_in": 43200}
|
||||
)
|
||||
assert body["expires_in"] == 43200
|
||||
|
|
|
|||
|
|
@ -872,6 +872,7 @@ def _mock_env_vars_prisma(row=None):
|
|||
prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[])
|
||||
prisma.db.litellm_mcpuserenvvars.upsert = AsyncMock()
|
||||
prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock()
|
||||
prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock()
|
||||
return prisma
|
||||
|
||||
|
||||
|
|
@ -1252,6 +1253,64 @@ async def test_delete_mcp_server_succeeds_when_orphan_cleanup_fails():
|
|||
prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_mcp_server_removes_orphaned_user_credentials():
|
||||
"""Deleting a server must also drop every user's stored BYOK/OAuth credential
|
||||
rows for it; there is no FK cascade, so skipping this leaves encrypted secrets
|
||||
pointing at a now-missing server."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import delete_mcp_server
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=object())
|
||||
|
||||
await delete_mcp_server(prisma, "srv-1")
|
||||
|
||||
prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once()
|
||||
call = prisma.db.litellm_mcpusercredentials.delete_many.call_args
|
||||
assert call.kwargs["where"] == {"server_id": "srv-1"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_mcp_server_skips_credential_cleanup_when_server_missing():
|
||||
"""A no-op delete (server not found) must not touch the credential table."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import delete_mcp_server
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=None)
|
||||
|
||||
result = await delete_mcp_server(prisma, "srv-1")
|
||||
|
||||
assert result is None
|
||||
prisma.db.litellm_mcpusercredentials.delete_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_mcp_server_credential_cleanup_failure_still_cleans_env_vars():
|
||||
"""Each per-user table is cleaned independently: a failure dropping credential
|
||||
rows must not skip the env var cleanup (or vice versa), and the delete must
|
||||
still succeed for the caller."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import delete_mcp_server
|
||||
|
||||
deleted = object()
|
||||
prisma = _mock_env_vars_prisma()
|
||||
prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=deleted)
|
||||
prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(
|
||||
side_effect=Exception("connection pool exhausted")
|
||||
)
|
||||
|
||||
result = await delete_mcp_server(prisma, "srv-1")
|
||||
|
||||
assert result is deleted
|
||||
prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once()
|
||||
prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once()
|
||||
|
||||
|
||||
# ── DB helpers: global env vars encrypted at rest ─────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -501,6 +501,7 @@ class TestListToolsRestAPI:
|
|||
raw_headers=None,
|
||||
user_api_key_auth=None,
|
||||
extra_headers=None,
|
||||
apply_tool_filters=True,
|
||||
):
|
||||
captured["called"] = True
|
||||
captured["server"] = server
|
||||
|
|
@ -545,6 +546,78 @@ class TestListToolsRestAPI:
|
|||
assert result["error"] is None
|
||||
assert result["message"] == "Successfully retrieved tools"
|
||||
|
||||
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
|
||||
to enumerate deliberately-disabled tools."""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
async def fake_contexts(user_api_key_auth):
|
||||
return [user_api_key_auth]
|
||||
|
||||
async def fake_get_allowed_mcp_servers(*args, **kwargs):
|
||||
return ["server-1"]
|
||||
|
||||
class StubServer:
|
||||
alias = "server-1"
|
||||
server_name = "server-1"
|
||||
name = "stub"
|
||||
allowed_tools = ["tool1"]
|
||||
mcp_info = {"server_name": "stub"}
|
||||
available_on_public_internet = True
|
||||
|
||||
stub_server = StubServer()
|
||||
captured = {}
|
||||
|
||||
async def fake_get_tools(
|
||||
server, server_auth_header, *args, apply_tool_filters=True, **kwargs
|
||||
):
|
||||
captured["apply_tool_filters"] = apply_tool_filters
|
||||
return ["tool-1"]
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
fake_contexts,
|
||||
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 == "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")
|
||||
|
||||
await rest_endpoints.list_tool_rest_api(
|
||||
request,
|
||||
server_id="server-1",
|
||||
include_disabled_tools=True,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
assert captured["apply_tool_filters"] is False
|
||||
|
||||
await rest_endpoints.list_tool_rest_api(
|
||||
request,
|
||||
server_id="server-1",
|
||||
include_disabled_tools=True,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
assert captured["apply_tool_filters"] is True
|
||||
|
||||
@pytest.mark.parametrize("upstream_status", [401, 403])
|
||||
async def test_upstream_auth_failure_surfaces_status_and_challenge(
|
||||
self, monkeypatch, upstream_status
|
||||
|
|
@ -649,6 +722,7 @@ class TestListToolsRestAPI:
|
|||
raw_headers=None,
|
||||
user_api_key_auth=None,
|
||||
extra_headers=None,
|
||||
apply_tool_filters=True,
|
||||
):
|
||||
captured["called"] = True
|
||||
captured["server_arg"] = server
|
||||
|
|
@ -792,6 +866,7 @@ class TestListToolsRestAPI:
|
|||
raw_headers=None,
|
||||
user_api_key_auth=None,
|
||||
extra_headers=None,
|
||||
apply_tool_filters=True,
|
||||
):
|
||||
captured["server"] = server
|
||||
captured["auth_header"] = server_auth_header
|
||||
|
|
@ -1284,6 +1359,56 @@ class TestGetToolsForSingleServer:
|
|||
assert "tool1" not in tool_names
|
||||
assert "tool4" not in tool_names
|
||||
|
||||
async def test_apply_tool_filters_false_returns_full_catalog(self, monkeypatch):
|
||||
"""apply_tool_filters=False returns the raw catalog without the server
|
||||
allowed_tools gate, so the config UI can render disabled tools as off."""
|
||||
from litellm.proxy._experimental.mcp_server.server import MCPServer
|
||||
from litellm.types.mcp import MCPTransport
|
||||
|
||||
class MockTool:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
self.description = name
|
||||
self.inputSchema = {}
|
||||
|
||||
mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")]
|
||||
|
||||
async def fake_get_tools_from_server(**kwargs):
|
||||
return mock_tools
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_get_tools_from_server",
|
||||
fake_get_tools_from_server,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
# Server enforces an allowlist of just tool1.
|
||||
server = MCPServer(
|
||||
server_id="test-server-id",
|
||||
name="test-server",
|
||||
transport=MCPTransport.sse,
|
||||
allowed_tools=["tool1"],
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key", object_permission=None)
|
||||
|
||||
# Runtime default: only the allowed tool comes back.
|
||||
filtered = await rest_endpoints._get_tools_for_single_server(
|
||||
server=server,
|
||||
server_auth_header=None,
|
||||
user_api_key_auth=user_api_key_dict,
|
||||
)
|
||||
assert [t.name for t in filtered] == ["tool1"]
|
||||
|
||||
# Config view: full catalog, including the disabled tools.
|
||||
full = await rest_endpoints._get_tools_for_single_server(
|
||||
server=server,
|
||||
server_auth_header=None,
|
||||
user_api_key_auth=user_api_key_dict,
|
||||
apply_tool_filters=False,
|
||||
)
|
||||
assert {t.name for t in full} == {"tool1", "tool2", "tool3"}
|
||||
|
||||
|
||||
class TestStdioCommandAllowlist:
|
||||
"""Tests for MCP stdio command allowlist validation."""
|
||||
|
|
|
|||
|
|
@ -112,6 +112,166 @@ async def test_handle_authentication_error_data_layer_errors_do_not_fall_back(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"db_error",
|
||||
[
|
||||
ConnectionError("connection refused"),
|
||||
TimeoutError("timed out"),
|
||||
asyncio.TimeoutError(),
|
||||
OSError("network is unreachable"),
|
||||
HTTPClientClosedError(),
|
||||
PrismaError("can't reach database server"),
|
||||
RawQueryError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "cached plan must not change result type",
|
||||
"meta": {"table": "t"},
|
||||
}
|
||||
}
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_handle_authentication_error_db_infra_error_returns_503(db_error):
|
||||
"""Regression for the outage where valid keys got 401 for 4 hours: an
|
||||
infrastructure-level DB failure during auth must surface as 503 (the DB
|
||||
could not confirm the key), never as 401 ("Invalid API key")."""
|
||||
handler = UserAPIKeyAuthExceptionHandler()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_exception_handler.seed_request_identity",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"allow_requests_on_db_unavailable": False},
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await handler._handle_authentication_error(
|
||||
db_error,
|
||||
MagicMock(),
|
||||
{},
|
||||
"/v1/chat/completions",
|
||||
None,
|
||||
"sk-valid-but-db-down",
|
||||
)
|
||||
|
||||
assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE
|
||||
assert exc_info.value.type == ProxyErrorTypes.no_db_connection
|
||||
assert "Invalid API key" not in str(exc_info.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_authentication_error_prisma_engine_teardown_returns_503():
|
||||
"""Regression for the first-request-of-an-outage edge case: at the instant
|
||||
the DB socket drops, the prisma query engine returns a malformed error
|
||||
payload and prisma-client-py crashes with a bare
|
||||
``AttributeError: 'NoneType' object has no attribute 'get'`` before it can
|
||||
raise P1001. That AttributeError reached auth and fell through to 401. It
|
||||
must surface as 503 like every other infra failure during the outage."""
|
||||
from prisma.engine import utils as prisma_engine_utils
|
||||
|
||||
malformed_payload = [
|
||||
{
|
||||
"error": "Can't reach database server",
|
||||
"user_facing_error": {
|
||||
"error_code": "P1001",
|
||||
"message": "Can't reach database server at `localhost`:`5503`",
|
||||
"meta": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
try:
|
||||
prisma_engine_utils.handle_response_errors(None, malformed_payload)
|
||||
raise AssertionError("expected prisma to raise AttributeError")
|
||||
except AttributeError as e:
|
||||
teardown_error = e
|
||||
|
||||
handler = UserAPIKeyAuthExceptionHandler()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_exception_handler.seed_request_identity",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"allow_requests_on_db_unavailable": False},
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await handler._handle_authentication_error(
|
||||
teardown_error,
|
||||
MagicMock(),
|
||||
{},
|
||||
"/v1/chat/completions",
|
||||
None,
|
||||
"sk-valid-but-db-down",
|
||||
)
|
||||
|
||||
assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE
|
||||
assert exc_info.value.type == ProxyErrorTypes.no_db_connection
|
||||
assert "Invalid API key" not in str(exc_info.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"auth_error",
|
||||
[
|
||||
# DB returned no row -> get_key_object raises this exact 401.
|
||||
ProxyException(
|
||||
message="Authentication Error, Invalid proxy server token passed.",
|
||||
type=ProxyErrorTypes.token_not_found_in_db,
|
||||
param="key",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
),
|
||||
# A bare auth failure raised as a plain Exception (e.g. master-key-only
|
||||
# route) must keep returning 401, not get reclassified as 503.
|
||||
Exception("Invalid proxy server token passed"),
|
||||
],
|
||||
)
|
||||
async def test_handle_authentication_error_genuine_auth_failure_stays_401(auth_error):
|
||||
"""Guard against the 503 conversion being too broad: a genuine auth
|
||||
failure (missing key / wrong key) must still be 401."""
|
||||
handler = UserAPIKeyAuthExceptionHandler()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_exception_handler.seed_request_identity",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"allow_requests_on_db_unavailable": False},
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await handler._handle_authentication_error(
|
||||
auth_error,
|
||||
MagicMock(),
|
||||
{},
|
||||
"/v1/chat/completions",
|
||||
None,
|
||||
"sk-bad-key",
|
||||
)
|
||||
|
||||
assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_authentication_error_budget_exceeded():
|
||||
handler = UserAPIKeyAuthExceptionHandler()
|
||||
|
|
|
|||
|
|
@ -112,11 +112,71 @@ async def test_should_clear_stale_budget_reservation_when_budget_checks_skip():
|
|||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
skip_budget_checks=True,
|
||||
general_settings={},
|
||||
)
|
||||
|
||||
assert user_api_key_auth_obj.budget_reservation is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disable_budget_reservation_skips_reservation():
|
||||
"""#27639: general_settings.disable_budget_reservation turns off the optimistic Redis
|
||||
reservation so operators hit by phantom BudgetExceededError can opt out of it."""
|
||||
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
|
||||
new=AsyncMock(return_value={"reserved_cost": 0.5, "entries": []}),
|
||||
) as mock_reserve:
|
||||
await _reserve_budget_after_common_checks(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
skip_budget_checks=False,
|
||||
general_settings={"disable_budget_reservation": True},
|
||||
)
|
||||
|
||||
mock_reserve.assert_not_called()
|
||||
assert user_api_key_auth_obj.budget_reservation is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_reservation_runs_when_not_disabled():
|
||||
"""Control for #27639: with the flag absent, the reservation still runs and is stored."""
|
||||
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
|
||||
reservation = {
|
||||
"reserved_cost": 0.5,
|
||||
"entries": [{"counter_key": "spend:key:test_token"}],
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
|
||||
new=AsyncMock(return_value=reservation),
|
||||
) as mock_reserve:
|
||||
await _reserve_budget_after_common_checks(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
skip_budget_checks=False,
|
||||
general_settings={},
|
||||
)
|
||||
|
||||
mock_reserve.assert_awaited_once()
|
||||
assert user_api_key_auth_obj.budget_reservation == reservation
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_reuse_cached_key_object_for_request_state():
|
||||
key_cache = DualCache()
|
||||
|
|
@ -1636,7 +1696,9 @@ class TestJWTOAuth2Coexistence:
|
|||
assert mock_auto_register.call_args.kwargs["team_id"] == "validated-team"
|
||||
assert mock_auto_register.call_args.kwargs["user_id"] == "validated-user"
|
||||
assert mock_auto_register.call_args.kwargs["org_id"] == "validated-org"
|
||||
assert mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user"
|
||||
assert (
|
||||
mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user"
|
||||
)
|
||||
assert result.org_id == "validated-org"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3548,3 +3610,118 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder()
|
|||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
def _proxy_attrs_for_db_lookup():
|
||||
"""Minimal proxy_server attributes for driving the real
|
||||
``_user_api_key_auth_builder`` down to the DB key lookup."""
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
return {
|
||||
"prisma_client": MagicMock(),
|
||||
"user_api_key_cache": DualCache(),
|
||||
"proxy_logging_obj": proxy_logging_obj,
|
||||
"master_key": "sk-test-master",
|
||||
"general_settings": {"allow_requests_on_db_unavailable": False},
|
||||
"llm_model_list": [],
|
||||
"llm_router": None,
|
||||
"open_telemetry_logger": None,
|
||||
"model_max_budget_limiter": MagicMock(),
|
||||
"user_custom_auth": None,
|
||||
"jwt_handler": None,
|
||||
"litellm_proxy_admin_name": "admin",
|
||||
}
|
||||
|
||||
|
||||
async def _run_builder_with_key_lookup(get_key_object_mock):
|
||||
"""Drive the real auth builder with ``get_key_object`` replaced by the
|
||||
given mock. Returns the builder result. Patches ``seed_request_identity``
|
||||
so the failure path doesn't touch OTEL."""
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
|
||||
|
||||
attrs = _proxy_attrs_for_db_lookup()
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_key_object",
|
||||
get_key_object_mock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_exception_handler.seed_request_identity",
|
||||
),
|
||||
):
|
||||
return await _user_api_key_auth_builder(
|
||||
request=request,
|
||||
api_key="Bearer sk-db-lookup-test",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_builder_returns_503_when_db_lookup_raises_infra_error():
|
||||
"""End-to-end: a DB infrastructure failure during the key lookup must
|
||||
propagate past the ``except ProxyException`` guard and surface as 503,
|
||||
not the 401 that masked the 4-hour outage. Killing the new 503 branch
|
||||
flips this to 401 and fails the test."""
|
||||
get_key_object = AsyncMock(side_effect=ConnectionError("connection refused"))
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _run_builder_with_key_lookup(get_key_object)
|
||||
|
||||
assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE
|
||||
assert exc_info.value.type == ProxyErrorTypes.no_db_connection
|
||||
assert "Invalid API key" not in str(exc_info.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_builder_returns_401_when_db_lookup_reports_missing_key():
|
||||
"""Regression guard: a genuinely missing key (DB returned no row, which
|
||||
``get_key_object`` raises as a 401 ProxyException) must still be 401."""
|
||||
missing_key_error = ProxyException(
|
||||
message="Authentication Error, Invalid proxy server token passed. key=..., not found in db.",
|
||||
type=ProxyErrorTypes.token_not_found_in_db,
|
||||
param="key",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
get_key_object = AsyncMock(side_effect=missing_key_error)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _run_builder_with_key_lookup(get_key_object)
|
||||
|
||||
assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_builder_succeeds_when_db_lookup_returns_valid_token():
|
||||
"""Regression guard: a valid key still authenticates. Proves the 503
|
||||
conversion only fires on the failure path and never intercepts success."""
|
||||
valid_token = UserAPIKeyAuth(api_key="sk-db-lookup-test", token="hashed-valid")
|
||||
get_key_object = AsyncMock(return_value=valid_token)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj",
|
||||
new_callable=AsyncMock,
|
||||
return_value=valid_token,
|
||||
) as mock_return:
|
||||
result = await _run_builder_with_key_lookup(get_key_object)
|
||||
|
||||
assert isinstance(result, UserAPIKeyAuth)
|
||||
# Reaching the success-assembly return (never the exception handler)
|
||||
# proves a valid key is unaffected by the 503 conversion.
|
||||
mock_return.assert_awaited_once()
|
||||
|
|
|
|||
|
|
@ -107,6 +107,201 @@ def test_is_database_connection_generic_errors():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[
|
||||
ConnectionError("connection refused"),
|
||||
TimeoutError("timed out"),
|
||||
OSError("network is unreachable"),
|
||||
asyncio.TimeoutError(),
|
||||
HTTPClientClosedError(),
|
||||
ClientNotConnectedError(),
|
||||
PrismaError("can't reach database server"),
|
||||
PrismaError(),
|
||||
],
|
||||
)
|
||||
def test_is_database_service_unavailable_error_infra_failures(error):
|
||||
"""Infrastructure-level failures (socket/connection/timeout, prisma
|
||||
transport, unknown PrismaError) mean the DB could not answer, so auth
|
||||
must surface 503 instead of treating a valid key as invalid."""
|
||||
assert PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is True
|
||||
|
||||
|
||||
def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_dataerror():
|
||||
"""Real-world regression: prisma-client-py raises the P1001 "can't reach
|
||||
database server" connectivity failure as a DataError (a data-layer type).
|
||||
A type-only check would miss it and return 401 during a genuine outage;
|
||||
the message keyword must still classify it as service-unavailable -> 503."""
|
||||
p1001_as_dataerror = DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "Can't reach database server at `127.0.0.1`:`5499`",
|
||||
"meta": {"table": "t"},
|
||||
}
|
||||
}
|
||||
)
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(
|
||||
p1001_as_dataerror
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_is_database_service_unavailable_error_cached_plan_escapes_as_503():
|
||||
"""Composes with the cached-plan retry: when that recovery fails and the
|
||||
Postgres "cached plan must not change result type" error escapes (raised by
|
||||
prisma as a data-layer RawQueryError), it is a transient stale-DB-state
|
||||
condition, not an invalid key, so it must classify as service-unavailable
|
||||
-> 503 rather than fall through to 401."""
|
||||
cached_plan_error = RawQueryError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "cached plan must not change result type",
|
||||
"meta": {"table": "t"},
|
||||
}
|
||||
}
|
||||
)
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(
|
||||
cached_plan_error
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_is_database_service_unavailable_error_prisma_engine_malformed_payload():
|
||||
"""Real-world regression: at the instant the DB socket drops, the prisma
|
||||
query engine returns a malformed error payload (``user_facing_error.meta``
|
||||
is ``null``). prisma-client-py's ``handle_response_errors`` then crashes
|
||||
with ``AttributeError: 'NoneType' object has no attribute 'get'`` before it
|
||||
can raise the proper P1001 error. That bare AttributeError has no
|
||||
connection keyword, so without the prisma-engine-origin check it falls
|
||||
through to 401 on the first request of an outage. Reproduce the exact
|
||||
prisma crash and assert it classifies as service-unavailable -> 503."""
|
||||
from prisma.engine import utils as prisma_engine_utils
|
||||
|
||||
malformed_payload = [
|
||||
{
|
||||
"error": "Can't reach database server",
|
||||
"user_facing_error": {
|
||||
"error_code": "P1001",
|
||||
"message": "Can't reach database server at `localhost`:`5503`",
|
||||
"meta": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
with pytest.raises(AttributeError) as exc_info:
|
||||
prisma_engine_utils.handle_response_errors(None, malformed_payload)
|
||||
|
||||
assert "no attribute 'get'" in str(exc_info.value)
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_is_prisma_engine_internal_error_excludes_application_attributeerror():
|
||||
"""The prisma-engine-origin check must stay narrow: a genuine AttributeError
|
||||
raised by application code (a real bug) must NOT be classified as
|
||||
service-unavailable, otherwise real bugs would silently become 503s."""
|
||||
|
||||
def application_bug():
|
||||
none_value = None
|
||||
return none_value.get("oops")
|
||||
|
||||
with pytest.raises(AttributeError) as exc_info:
|
||||
application_bug()
|
||||
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc_info.value)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_is_prisma_engine_internal_error_excludes_data_layer_prisma_error():
|
||||
"""A data-layer ``PrismaError`` (the DB IS reachable and rejected the data)
|
||||
must stay 401. These are always raised from prisma internals, so the check
|
||||
excludes any ``PrismaError`` by type before inspecting the traceback."""
|
||||
data_layer_error = UniqueViolationError(
|
||||
data={"user_facing_error": {"meta": {"table": "t"}}}
|
||||
)
|
||||
try:
|
||||
raise data_layer_error
|
||||
except UniqueViolationError as e:
|
||||
assert PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[
|
||||
DataError(data={"user_facing_error": {"meta": {"table": "t"}}}),
|
||||
UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}),
|
||||
RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}),
|
||||
Exception("some unrelated error"),
|
||||
ValueError("bad value"),
|
||||
],
|
||||
)
|
||||
def test_is_database_service_unavailable_error_excludes_non_infra(error):
|
||||
"""Data-layer errors (the DB IS reachable and answered) and generic
|
||||
non-DB errors must NOT be classified as service-unavailable, otherwise a
|
||||
genuine 401 would be masked as a transient 503."""
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is False
|
||||
)
|
||||
|
||||
|
||||
def test_is_database_service_unavailable_error_asyncpg(monkeypatch):
|
||||
"""asyncpg connection/interface errors map to service-unavailable. asyncpg
|
||||
is not a hard dependency, so inject a stand-in module to exercise the
|
||||
branch deterministically regardless of the install environment."""
|
||||
import sys
|
||||
import types
|
||||
|
||||
fake_asyncpg = types.ModuleType("asyncpg")
|
||||
fake_exceptions = types.ModuleType("asyncpg.exceptions")
|
||||
|
||||
class PostgresConnectionError(Exception):
|
||||
pass
|
||||
|
||||
class InterfaceError(Exception):
|
||||
pass
|
||||
|
||||
class UniqueViolationError(Exception): # data-layer, must stay False
|
||||
pass
|
||||
|
||||
fake_exceptions.PostgresConnectionError = PostgresConnectionError
|
||||
fake_exceptions.InterfaceError = InterfaceError
|
||||
fake_exceptions.UniqueViolationError = UniqueViolationError
|
||||
fake_asyncpg.exceptions = fake_exceptions
|
||||
|
||||
monkeypatch.setitem(sys.modules, "asyncpg", fake_asyncpg)
|
||||
monkeypatch.setitem(sys.modules, "asyncpg.exceptions", fake_exceptions)
|
||||
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(
|
||||
PostgresConnectionError("connection reset")
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(
|
||||
InterfaceError("connection was closed")
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
PrismaDBExceptionHandler.is_database_service_unavailable_error(
|
||||
UniqueViolationError("duplicate key")
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
# Test should_allow_request_on_db_unavailable method
|
||||
@patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
|
|
|
|||
|
|
@ -412,6 +412,171 @@ async def test_apply_guardrail_response_ok(
|
|||
assert result["texts"] == inputs["texts"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_sends_user_id_model_and_extra_info(
|
||||
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
|
||||
) -> None:
|
||||
inputs: GenericGuardrailAPIInputs = {
|
||||
"texts": ["Hello"],
|
||||
"structured_messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-4o",
|
||||
}
|
||||
request_data = {
|
||||
"messages": inputs["structured_messages"],
|
||||
"model": "gpt-4o",
|
||||
"litellm_metadata": {
|
||||
"user_api_key_user_id": "uid-abc",
|
||||
"user_api_key_user_email": "alice@example.com",
|
||||
},
|
||||
}
|
||||
guardrail_endpoint = (
|
||||
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
json={"result": {"blocked": False, "transformed": False}},
|
||||
request=httpx.Request(method="POST", url=guardrail_endpoint),
|
||||
),
|
||||
) as mock_method:
|
||||
await crowdstrike_aidr_guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
payload = mock_method.call_args.kwargs["json"]
|
||||
assert payload["user_id"] == "uid-abc"
|
||||
assert payload["model"] == "gpt-4o"
|
||||
assert payload["extra_info"] == {"user_name": "alice@example.com"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_empty_extra_info_when_no_email(
|
||||
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
|
||||
) -> None:
|
||||
inputs: GenericGuardrailAPIInputs = {
|
||||
"texts": ["Hello"],
|
||||
"structured_messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gemini-flash",
|
||||
}
|
||||
request_data = {
|
||||
"messages": inputs["structured_messages"],
|
||||
"model": "gemini-flash",
|
||||
"litellm_metadata": {
|
||||
"user_api_key_user_id": "uid-no-email",
|
||||
"user_api_key_user_email": None,
|
||||
},
|
||||
}
|
||||
guardrail_endpoint = (
|
||||
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
json={"result": {"blocked": False, "transformed": False}},
|
||||
request=httpx.Request(method="POST", url=guardrail_endpoint),
|
||||
),
|
||||
) as mock_method:
|
||||
await crowdstrike_aidr_guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
payload = mock_method.call_args.kwargs["json"]
|
||||
assert payload["user_id"] == "uid-no-email"
|
||||
assert payload["model"] == "gemini-flash"
|
||||
assert payload["extra_info"] == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_no_metadata_skips_user_fields(
|
||||
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
|
||||
) -> None:
|
||||
inputs: GenericGuardrailAPIInputs = {
|
||||
"texts": ["Hello"],
|
||||
"structured_messages": [{"role": "user", "content": "Hello"}],
|
||||
}
|
||||
request_data = {"messages": inputs["structured_messages"]}
|
||||
guardrail_endpoint = (
|
||||
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
json={"result": {"blocked": False, "transformed": False}},
|
||||
request=httpx.Request(method="POST", url=guardrail_endpoint),
|
||||
),
|
||||
) as mock_method:
|
||||
await crowdstrike_aidr_guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
payload = mock_method.call_args.kwargs["json"]
|
||||
assert "user_id" not in payload
|
||||
assert "model" not in payload
|
||||
assert "extra_info" not in payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_metadata, metadata",
|
||||
[
|
||||
(None, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}),
|
||||
({"trace_id": "t1"}, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}),
|
||||
(["unexpected"], {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}),
|
||||
({"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}, {"trace_id": "t1"}),
|
||||
],
|
||||
ids=["identity_in_metadata_llm_none", "identity_in_metadata_llm_user_dict", "identity_in_metadata_llm_non_mapping", "identity_in_litellm_metadata"],
|
||||
)
|
||||
async def test_apply_guardrail_reads_identity_from_either_metadata_bag(
|
||||
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
|
||||
litellm_metadata,
|
||||
metadata,
|
||||
) -> None:
|
||||
inputs: GenericGuardrailAPIInputs = {
|
||||
"texts": ["Hello"],
|
||||
"structured_messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-4o",
|
||||
}
|
||||
request_data = {
|
||||
"messages": inputs["structured_messages"],
|
||||
"model": "gpt-4o",
|
||||
"litellm_metadata": litellm_metadata,
|
||||
"metadata": metadata,
|
||||
}
|
||||
guardrail_endpoint = (
|
||||
f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions"
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
json={"result": {"blocked": False, "transformed": False}},
|
||||
request=httpx.Request(method="POST", url=guardrail_endpoint),
|
||||
),
|
||||
) as mock_method:
|
||||
await crowdstrike_aidr_guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
payload = mock_method.call_args.kwargs["json"]
|
||||
assert payload["user_id"] == "uid-abc"
|
||||
assert payload["extra_info"] == {"user_name": "alice@example.com"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_skipped_messages_stay_aligned(
|
||||
crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler,
|
||||
|
|
|
|||
|
|
@ -713,7 +713,8 @@ async def test_count_input_file_usage_decodes_model_embedded_file_id():
|
|||
@pytest.mark.asyncio
|
||||
async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias():
|
||||
"""After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5).
|
||||
Auth must check the proxy model_name the key was granted, not the stripped id."""
|
||||
Auth must check target_model_names from the unified file id, not reverse-map
|
||||
the stripped id."""
|
||||
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
||||
|
||||
rate_limiter = _PROXY_BatchRateLimiter(
|
||||
|
|
@ -732,7 +733,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(
|
|||
)
|
||||
mock_router = MagicMock()
|
||||
mock_router.model_list = []
|
||||
mock_router.resolve_model_name_from_model_id.return_value = proxy_alias
|
||||
can_key_call_model = AsyncMock(return_value=True)
|
||||
|
||||
with (
|
||||
|
|
@ -745,10 +745,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(
|
|||
await rate_limiter._enforce_batch_file_model_access(
|
||||
user_api_key_dict=user,
|
||||
file_content_as_dict=file_dict,
|
||||
target_model_names=[proxy_alias],
|
||||
)
|
||||
|
||||
can_key_call_model.assert_awaited_once()
|
||||
assert can_key_call_model.await_args.kwargs["model"] == proxy_alias
|
||||
mock_router.resolve_model_name_from_model_id.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_list_order",
|
||||
[
|
||||
[
|
||||
"openai/openai/gpt-5.5",
|
||||
"openai/openai/gpt-5.5-batch",
|
||||
"us/azure/openai/gpt-5.5",
|
||||
],
|
||||
[
|
||||
"us/azure/openai/gpt-5.5",
|
||||
"openai/openai/gpt-5.5",
|
||||
"openai/openai/gpt-5.5-batch",
|
||||
],
|
||||
[
|
||||
"openai/openai/gpt-5.5-batch",
|
||||
"us/azure/openai/gpt-5.5",
|
||||
"openai/openai/gpt-5.5",
|
||||
],
|
||||
],
|
||||
)
|
||||
async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup(
|
||||
model_list_order,
|
||||
):
|
||||
"""LIT-3593: three deployments strip to gpt-5.5; auth must use the upload
|
||||
target alias from target_model_names, not first-match reverse lookup."""
|
||||
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
||||
|
||||
rate_limiter = _PROXY_BatchRateLimiter(
|
||||
internal_usage_cache=MagicMock(),
|
||||
parallel_request_limiter=MagicMock(),
|
||||
)
|
||||
batch_alias = "openai/openai/gpt-5.5-batch"
|
||||
deployment_templates = {
|
||||
"openai/openai/gpt-5.5": {
|
||||
"model_name": "openai/openai/gpt-5.5",
|
||||
"litellm_params": {"model": "openai/gpt-5.5"},
|
||||
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
|
||||
},
|
||||
"openai/openai/gpt-5.5-batch": {
|
||||
"model_name": "openai/openai/gpt-5.5-batch",
|
||||
"litellm_params": {"model": "openai/gpt-5.5"},
|
||||
"model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"},
|
||||
},
|
||||
"us/azure/openai/gpt-5.5": {
|
||||
"model_name": "us/azure/openai/gpt-5.5",
|
||||
"litellm_params": {"model": "azure/gpt-5.5"},
|
||||
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
|
||||
},
|
||||
}
|
||||
mock_router = MagicMock()
|
||||
mock_router.model_list = [deployment_templates[name] for name in model_list_order]
|
||||
|
||||
def _resolve(model_id):
|
||||
for deployment in mock_router.model_list:
|
||||
actual_model = deployment.get("litellm_params", {}).get("model")
|
||||
if actual_model == model_id or (
|
||||
actual_model and actual_model.endswith(f"/{model_id}")
|
||||
):
|
||||
return deployment.get("model_name")
|
||||
return None
|
||||
|
||||
mock_router.resolve_model_name_from_model_id.side_effect = _resolve
|
||||
|
||||
file_dict = [
|
||||
{"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
|
||||
]
|
||||
user = UserAPIKeyAuth(
|
||||
api_key="sk-ok",
|
||||
user_id="alice",
|
||||
models=[batch_alias],
|
||||
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
)
|
||||
can_key_call_model = AsyncMock(return_value=True)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_checks.can_key_call_model",
|
||||
new=can_key_call_model,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
||||
):
|
||||
await rate_limiter._enforce_batch_file_model_access(
|
||||
user_api_key_dict=user,
|
||||
file_content_as_dict=file_dict,
|
||||
target_model_names=[batch_alias],
|
||||
)
|
||||
|
||||
can_key_call_model.assert_awaited_once()
|
||||
assert can_key_call_model.await_args.kwargs["model"] == batch_alias
|
||||
mock_router.resolve_model_name_from_model_id.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1482,6 +1482,10 @@ class TestTemporaryMCPSessionEndpoints:
|
|||
"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=[non_admin]),
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _get_cached_temporary_mcp_server_or_404("server-x", non_admin)
|
||||
|
|
@ -1514,6 +1518,10 @@ class TestTemporaryMCPSessionEndpoints:
|
|||
"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=[non_admin]),
|
||||
),
|
||||
):
|
||||
result = await _get_cached_temporary_mcp_server_or_404(
|
||||
"server-x", non_admin
|
||||
|
|
@ -1521,6 +1529,58 @@ class TestTemporaryMCPSessionEndpoints:
|
|||
|
||||
assert result is registry_server
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_temporary_mcp_server_non_admin_allowed_via_team_access_group(
|
||||
self,
|
||||
):
|
||||
"""Internal user whose only grant to the server flows through a team
|
||||
access-group must pass the authorize/token access check. The check has to
|
||||
expand the UI session into per-team contexts (build_effective_auth_contexts),
|
||||
the same way the server-list grid does; checking only the bare session
|
||||
context leaves the team grant invisible and 403s the user."""
|
||||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_get_cached_temporary_mcp_server_or_404,
|
||||
)
|
||||
|
||||
registry_server = generate_mock_mcp_server_config_record(server_id="server-x")
|
||||
ui_session_auth = generate_mock_user_api_key_auth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
team_context = ui_session_auth.model_copy()
|
||||
team_context.team_id = "team-with-mcp-grant"
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.get_mcp_server_by_id.return_value = registry_server
|
||||
mock_manager.get_mcp_server_by_name.return_value = None
|
||||
|
||||
def allowed_for(auth):
|
||||
return ["server-x"] if auth.team_id == "team-with-mcp-grant" else []
|
||||
|
||||
mock_manager.get_allowed_mcp_servers = AsyncMock(side_effect=allowed_for)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_cached_temporary_mcp_server",
|
||||
return_value=None,
|
||||
),
|
||||
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=[ui_session_auth, team_context]),
|
||||
),
|
||||
):
|
||||
result = await _get_cached_temporary_mcp_server_or_404(
|
||||
"server-x", ui_session_auth
|
||||
)
|
||||
|
||||
assert result is registry_server
|
||||
assert mock_manager.get_allowed_mcp_servers.await_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_temporary_mcp_server_temp_cache_non_admin_denied(self):
|
||||
"""Servers resolved from the admin-only temp cache reject non-admins."""
|
||||
|
|
|
|||
|
|
@ -1615,6 +1615,7 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name):
|
|||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=existing_team
|
||||
)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=None)
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(
|
||||
return_value=updated_team
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,77 @@
|
|||
"""
|
||||
Tests for atomic team model operations during BYOK model creation.
|
||||
|
||||
Regression tests for https://github.com/BerriAI/litellm/issues/22594
|
||||
Concurrent BYOK model creates must not overwrite each other's entries
|
||||
in team.models.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LitellmUserRoles,
|
||||
TeamModelAddRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
|
||||
|
||||
class TestTeamModelAddAtomicAppend:
|
||||
"""Verify team_model_add uses atomic SQL for the models array append."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_atomic_array_append_with_dedup(self):
|
||||
"""team_model_add must call execute_raw with DISTINCT unnest SQL."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_model_add
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user"
|
||||
)
|
||||
|
||||
existing_team = MagicMock()
|
||||
existing_team.model_dump.return_value = {
|
||||
"team_id": "team-1",
|
||||
"models": ["existing-model"],
|
||||
}
|
||||
|
||||
updated_team = MagicMock()
|
||||
updated_team.team_id = "team-1"
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=existing_team
|
||||
)
|
||||
mock_prisma.db.execute_raw = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_teamtable.update = AsyncMock(
|
||||
return_value=updated_team
|
||||
)
|
||||
|
||||
await team_model_add(
|
||||
data=TeamModelAddRequest(team_id="team-1", models=["new-model"]),
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
)
|
||||
|
||||
mock_prisma.db.execute_raw.assert_called_once()
|
||||
sql = mock_prisma.db.execute_raw.call_args[0][0]
|
||||
assert "DISTINCT unnest" in sql
|
||||
assert "all-proxy-models" in sql
|
||||
assert mock_prisma.db.execute_raw.call_args[0][1] == ["new-model"]
|
||||
assert mock_prisma.db.execute_raw.call_args[0][2] == "team-1"
|
||||
|
||||
# Should use write-routed update to re-fetch, not find_unique
|
||||
mock_prisma.db.litellm_teamtable.update.assert_called_once()
|
||||
|
|
@ -321,6 +321,307 @@ class TestAzureAnthropicCostCalculation:
|
|||
assert call_kwargs["model"] == "azure_ai/claude-sonnet-4-5_gb_20250929"
|
||||
assert call_kwargs["custom_llm_provider"] == "azure_ai"
|
||||
|
||||
@patch("litellm.completion_cost")
|
||||
def test_cost_calculation_resolves_unknown_model_from_litellm_params(
|
||||
self, mock_completion_cost
|
||||
):
|
||||
"""When the body model is the "unknown" sentinel, the deployment model
|
||||
from litellm_params must be used for costing, not "unknown" (which makes
|
||||
completion_cost raise and the cost silently fall back to $0)."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
mock_completion_cost.return_value = 0.001
|
||||
|
||||
logging_obj = self._create_mock_logging_obj(model="unknown")
|
||||
logging_obj.model_call_details["litellm_params"] = {
|
||||
"model": "anthropic/claude-3-5-haiku-20241022",
|
||||
"metadata": {
|
||||
"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"
|
||||
},
|
||||
}
|
||||
logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"]
|
||||
|
||||
mock_response = MagicMock(spec=ModelResponse)
|
||||
mock_response.id = "test-id"
|
||||
mock_response.model = "unknown"
|
||||
|
||||
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
|
||||
litellm_model_response=mock_response,
|
||||
model="unknown",
|
||||
kwargs={},
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
mock_completion_cost.assert_called_once()
|
||||
assert (
|
||||
mock_completion_cost.call_args[1]["model"]
|
||||
== "anthropic/claude-3-5-haiku-20241022"
|
||||
)
|
||||
assert kwargs["response_cost"] == 0.001
|
||||
assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022"
|
||||
|
||||
@patch("litellm.completion_cost")
|
||||
def test_cost_calculation_resolves_unknown_model_from_model_group(
|
||||
self, mock_completion_cost
|
||||
):
|
||||
"""With only model_group available (no deployment litellm_params.model),
|
||||
the leading passthrough/ prefix must be stripped so the cost map can
|
||||
resolve the model."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
mock_completion_cost.return_value = 0.002
|
||||
|
||||
logging_obj = self._create_mock_logging_obj(model="unknown")
|
||||
logging_obj.model_call_details["litellm_params"] = {
|
||||
"metadata": {
|
||||
"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"
|
||||
}
|
||||
}
|
||||
logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"]
|
||||
|
||||
mock_response = MagicMock(spec=ModelResponse)
|
||||
mock_response.id = "test-id"
|
||||
mock_response.model = "unknown"
|
||||
|
||||
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
|
||||
litellm_model_response=mock_response,
|
||||
model="unknown",
|
||||
kwargs={},
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
mock_completion_cost.assert_called_once()
|
||||
assert (
|
||||
mock_completion_cost.call_args[1]["model"]
|
||||
== "anthropic/claude-3-5-haiku-20241022"
|
||||
)
|
||||
assert kwargs["response_cost"] == 0.002
|
||||
|
||||
@patch("litellm.completion_cost")
|
||||
def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group(
|
||||
self, mock_completion_cost
|
||||
):
|
||||
"""When litellm_params.model is itself the "unknown" sentinel, the
|
||||
deployment-model branch must not short-circuit; resolution falls through
|
||||
to model_group so costing still prices the real model instead of "unknown"."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
mock_completion_cost.return_value = 0.003
|
||||
|
||||
logging_obj = self._create_mock_logging_obj(model="unknown")
|
||||
logging_obj.model_call_details["litellm_params"] = {
|
||||
"model": "unknown",
|
||||
"metadata": {
|
||||
"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"
|
||||
},
|
||||
}
|
||||
logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"]
|
||||
|
||||
mock_response = MagicMock(spec=ModelResponse)
|
||||
mock_response.id = "test-id"
|
||||
mock_response.model = "unknown"
|
||||
|
||||
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
|
||||
litellm_model_response=mock_response,
|
||||
model="unknown",
|
||||
kwargs={},
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
mock_completion_cost.assert_called_once()
|
||||
assert (
|
||||
mock_completion_cost.call_args[1]["model"]
|
||||
== "anthropic/claude-3-5-haiku-20241022"
|
||||
)
|
||||
assert kwargs["response_cost"] == 0.003
|
||||
assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022"
|
||||
|
||||
@patch("litellm.completion_cost")
|
||||
def test_streaming_cost_calculation_resolves_model_from_message_start_chunk(
|
||||
self, mock_completion_cost
|
||||
):
|
||||
"""On the bare /anthropic passthrough path litellm_params carries no model
|
||||
or model_group and the body model is the "unknown" sentinel; the model
|
||||
must be recovered from the message_start SSE event so completion_cost
|
||||
prices the real model instead of failing on "unknown" and logging $0."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
Logging as RealLoggingObj,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.streaming_handler import (
|
||||
PassThroughStreamingHandler,
|
||||
)
|
||||
|
||||
mock_completion_cost.return_value = 0.001
|
||||
|
||||
def _sse(event, data):
|
||||
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
|
||||
|
||||
frames = [
|
||||
_sse(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-3-5-haiku-20241022",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
_sse(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
),
|
||||
_sse(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "text_delta", "text": "hi"},
|
||||
},
|
||||
),
|
||||
_sse("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
_sse(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 1},
|
||||
},
|
||||
),
|
||||
_sse("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
all_chunks = list(
|
||||
PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames)
|
||||
)
|
||||
|
||||
logging_obj = RealLoggingObj(
|
||||
model="unknown",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="pass_through_endpoint",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-call-id",
|
||||
function_id="1",
|
||||
)
|
||||
logging_obj.model_call_details["model"] = "unknown"
|
||||
logging_obj.model_call_details["stream"] = True
|
||||
logging_obj.model_call_details["litellm_params"] = {}
|
||||
logging_obj.litellm_params = {}
|
||||
|
||||
result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
|
||||
litellm_logging_obj=logging_obj,
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/anthropic/v1/messages",
|
||||
request_body={"stream": True},
|
||||
endpoint_type="messages",
|
||||
start_time=datetime.now(),
|
||||
all_chunks=all_chunks,
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert result["result"] is not None
|
||||
mock_completion_cost.assert_called_once()
|
||||
assert mock_completion_cost.call_args[1]["model"] == "claude-3-5-haiku-20241022"
|
||||
assert result["kwargs"]["response_cost"] == 0.001
|
||||
assert result["kwargs"]["model"] == "claude-3-5-haiku-20241022"
|
||||
|
||||
def test_extract_model_skips_non_dict_data_payload(self):
|
||||
"""A scalar data: payload (e.g. `data: null`) must be skipped, not crash
|
||||
the streaming log handler with AttributeError, which would propagate out
|
||||
and break spend logging for the whole request."""
|
||||
chunks = [
|
||||
"event: ping\ndata: null\n\n",
|
||||
'event: message_start\ndata: {"type": "message_start", "message": '
|
||||
'{"model": "claude-3-5-haiku-20241022"}}\n\n',
|
||||
]
|
||||
|
||||
assert (
|
||||
AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(
|
||||
chunks
|
||||
)
|
||||
== "claude-3-5-haiku-20241022"
|
||||
)
|
||||
|
||||
def test_extract_model_parses_per_line_not_first_data_substring(self):
|
||||
"""A raw multi-line SSE event whose non-data line contains the substring
|
||||
"data:" must not derail parsing: matching only lines that start with
|
||||
"data:" recovers the message_start model, whereas a first-substring slice
|
||||
would consume the wrong offset, fail to parse JSON, and return None."""
|
||||
raw_event = (
|
||||
"event: ping data: not-json\n"
|
||||
'data: {"type": "message_start", "message": '
|
||||
'{"model": "claude-3-5-haiku-20241022"}}\n\n'
|
||||
)
|
||||
|
||||
assert (
|
||||
AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(
|
||||
[raw_event]
|
||||
)
|
||||
== "claude-3-5-haiku-20241022"
|
||||
)
|
||||
|
||||
def test_passthrough_logging_sets_response_cost_with_server_tool_use_dict(self):
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
logging_obj = self._create_mock_logging_obj(model="claude-3-7-sonnet-20250219")
|
||||
logging_obj.get_router_model_id.return_value = None
|
||||
logging_obj.litellm_params = {}
|
||||
|
||||
response = ModelResponse(
|
||||
id="test-id",
|
||||
choices=[
|
||||
Choices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
message=Message(content="test", role="assistant"),
|
||||
)
|
||||
],
|
||||
created=1234567890,
|
||||
model="claude-3-7-sonnet-20250219",
|
||||
usage={
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
"server_tool_use": {"web_search_requests": 1},
|
||||
},
|
||||
)
|
||||
|
||||
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
|
||||
litellm_model_response=response,
|
||||
model="claude-3-7-sonnet-20250219",
|
||||
kwargs={},
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert "response_cost" in kwargs
|
||||
assert kwargs["response_cost"] > 0
|
||||
|
||||
|
||||
class TestAnthropicBatchPassthroughCostTracking:
|
||||
"""Test cases for Anthropic batch passthrough cost tracking functionality"""
|
||||
|
|
@ -1045,6 +1346,74 @@ class TestPureTextFastPathParity:
|
|||
)
|
||||
|
||||
|
||||
class TestBuildCompleteStreamingResponseRobustness:
|
||||
"""_build_complete_streaming_response must tolerate non-standard SSE frames."""
|
||||
|
||||
def _build(self, chunks: List[str]):
|
||||
return AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=chunks,
|
||||
litellm_logging_obj=MagicMock(),
|
||||
model="claude-3-sonnet-20240229",
|
||||
)
|
||||
|
||||
def test_done_frame_is_skipped(self):
|
||||
"""A bare 'data: [DONE]' control frame must not break reconstruction."""
|
||||
chunks = [
|
||||
'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}',
|
||||
'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
|
||||
'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}',
|
||||
'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}',
|
||||
'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}',
|
||||
'event: message_stop\ndata: {"type":"message_stop"}',
|
||||
"data: [DONE]",
|
||||
]
|
||||
result = self._build(chunks)
|
||||
assert result is not None
|
||||
assert result.choices[0].message.content == "Hi"
|
||||
|
||||
def test_non_json_sse_line_is_skipped(self):
|
||||
"""Non-JSON SSE lines (comments, keep-alive pings) must be skipped."""
|
||||
chunks = [
|
||||
": ping",
|
||||
'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}',
|
||||
"this is not json at all",
|
||||
]
|
||||
# Must not raise; a malformed stream simply yields no usable response.
|
||||
result = self._build(chunks)
|
||||
assert result is None or hasattr(result, "choices")
|
||||
|
||||
def test_mixed_valid_and_invalid_frames(self):
|
||||
"""Valid events are still collected when interleaved with invalid ones."""
|
||||
chunks = [
|
||||
'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}',
|
||||
"data: [DONE]",
|
||||
": keep-alive",
|
||||
"not-json",
|
||||
'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
|
||||
'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}',
|
||||
'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}',
|
||||
'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}',
|
||||
'event: message_stop\ndata: {"type":"message_stop"}',
|
||||
]
|
||||
result = self._build(chunks)
|
||||
assert result is not None
|
||||
assert result.choices[0].message.content == "Hello"
|
||||
|
||||
def test_done_in_text_payload_is_not_dropped(self):
|
||||
"""A valid event whose text content contains '[DONE]' must NOT be skipped."""
|
||||
chunks = [
|
||||
'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}',
|
||||
'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
|
||||
'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"The stream ends with [DONE]"}}',
|
||||
'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}',
|
||||
'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":8}}',
|
||||
'event: message_stop\ndata: {"type":"message_stop"}',
|
||||
]
|
||||
result = self._build(chunks)
|
||||
assert result is not None
|
||||
assert result.choices[0].message.content == "The stream ends with [DONE]"
|
||||
|
||||
|
||||
class TestStreamFalseDeduplication:
|
||||
"""
|
||||
Regression tests for the duplicate-callback bug where a streaming pass-through
|
||||
|
|
|
|||
|
|
@ -60,6 +60,15 @@ def test_v2_model_info_invalid_page_returns_422(client, auth_as, empty_router):
|
|||
assert "detail" in response.json()
|
||||
|
||||
|
||||
def test_v2_model_info_in_openapi_schema():
|
||||
"""``GET /v2/model/info`` is published in the proxy OpenAPI/Swagger spec."""
|
||||
from litellm.proxy.proxy_server import get_openapi_schema
|
||||
|
||||
schema = get_openapi_schema()
|
||||
assert "/v2/model/info" in schema["paths"]
|
||||
assert "get" in schema["paths"]["/v2/model/info"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET /v1/model/info, GET /model/info
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -151,21 +151,24 @@ async def test_model_info_v2_translates_team_model_name(monkeypatch):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch):
|
||||
"""/v1/model/info list path (no litellm_model_id) must surface the public
|
||||
name. Covers the list comprehension that assigns _get_proxy_model_info's
|
||||
return back into all_models (#28382 review)."""
|
||||
"""/v1/model/info list path (no litellm_model_id) must include team-scoped
|
||||
deployments from the router model list and surface the public name (#28382)."""
|
||||
team_row = _team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4o",
|
||||
"litellm_params": {"model": "gpt-4o"},
|
||||
"model_info": {"id": "normal-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.get_model_names.return_value = ["team-claude-sonnet"]
|
||||
router.model_list = [team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4o"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
router.get_model_list.return_value = [_team_row()]
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", [_team_row()])
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "get_key_models", lambda **kw: [])
|
||||
monkeypatch.setattr(ps, "get_team_models", lambda **kw: [])
|
||||
monkeypatch.setattr(
|
||||
ps, "get_complete_model_list", lambda **kw: ["team-claude-sonnet"]
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
|
|
@ -176,3 +179,417 @@ async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch):
|
|||
names = [m["model_name"] for m in resp["data"]]
|
||||
assert "team-claude-sonnet" in names
|
||||
assert "model_name_team-abc-123_4a6b8" not in names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_unrestricted_key_returns_all_deployments(monkeypatch):
|
||||
"""Unrestricted keys must see all router deployments (legacy v1 access logic)."""
|
||||
deployment = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [deployment]
|
||||
router.get_model_names.return_value = ["gpt-4"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id="user-1",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
models=[],
|
||||
team_models=[],
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
assert [m["model_name"] for m in resp["data"]] == ["gpt-4"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_restricted_key_filters_deployments(monkeypatch):
|
||||
"""Key-level model allowlists must filter router deployments."""
|
||||
team_row = _team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4", "team-claude-sonnet"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id="user-1",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
models=["gpt-4"],
|
||||
team_models=[],
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
assert [m["model_name"] for m in resp["data"]] == ["gpt-4"]
|
||||
|
||||
|
||||
def _other_team_row() -> dict:
|
||||
return {
|
||||
"model_name": "model_name_team-other_9f2c1",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-5.2-low-rpm-testing",
|
||||
"api_base": "https://team-other-private.example.com",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "byok-id-other",
|
||||
"team_id": "team-other",
|
||||
"team_public_model_name": "team-claude-sonnet",
|
||||
"db_model": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch):
|
||||
"""Unrestricted non-admin keys must not enumerate other teams' BYOK
|
||||
deployments, but must still see global models and their own team's."""
|
||||
team_row = _team_row()
|
||||
other_team_row = _other_team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [team_row, other_team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
|
||||
prisma_client = MagicMock()
|
||||
caller_user_row = MagicMock()
|
||||
caller_user_row.teams = ["team-abc-123"]
|
||||
caller_user_row.model_dump.return_value = {
|
||||
"user_id": "user-1",
|
||||
"teams": ["team-abc-123"],
|
||||
"models": [],
|
||||
}
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=caller_user_row
|
||||
)
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={}))
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id="user-1",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
models=[],
|
||||
team_models=[],
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
returned_ids = {m["model_info"]["id"] for m in resp["data"]}
|
||||
assert returned_ids == {"global-id-1", "byok-id-1"}
|
||||
assert "byok-id-other" not in returned_ids
|
||||
names = [m["model_name"] for m in resp["data"]]
|
||||
assert "team-claude-sonnet" in names
|
||||
assert "gpt-4" in names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch):
|
||||
"""A key without a resolvable user (e.g. CI/service token) sees only
|
||||
global deployments, never any team-scoped BYOK rows."""
|
||||
team_row = _team_row()
|
||||
other_team_row = _other_team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [team_row, other_team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
|
||||
prisma_client = MagicMock()
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id=None,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id="team-abc-123",
|
||||
models=[],
|
||||
team_models=[],
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_populates_access_via_team_ids(monkeypatch):
|
||||
"""`/v1/model/info` must populate access_via_team_ids when the DB is connected."""
|
||||
team_id = "team-abc-123"
|
||||
team_row = _team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4o",
|
||||
"litellm_params": {"model": "gpt-4o"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4o", "team-claude-sonnet"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
router.get_model_ids.return_value = ["global-id-1"]
|
||||
|
||||
prisma_client = MagicMock()
|
||||
|
||||
async def _fake_populate(**kwargs):
|
||||
for model in kwargs["all_models"]:
|
||||
model_id = model["model_info"]["id"]
|
||||
if model_id == "byok-id-1":
|
||||
model["model_info"]["access_via_team_ids"] = [team_id]
|
||||
model["model_info"]["direct_access"] = False
|
||||
elif model_id == "global-id-1":
|
||||
model["model_info"]["direct_access"] = True
|
||||
return kwargs["all_models"]
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None)
|
||||
|
||||
by_id = {m["model_info"]["id"]: m for m in resp["data"]}
|
||||
assert by_id["byok-id-1"]["model_info"]["access_via_team_ids"] == [team_id]
|
||||
assert by_id["byok-id-1"]["model_info"]["direct_access"] is False
|
||||
assert by_id["global-id-1"]["model_info"]["direct_access"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_populate_team_access_sets_direct_access_false_by_default(monkeypatch):
|
||||
"""Team-accessible models without direct access must return direct_access=false."""
|
||||
team_row = _team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4o",
|
||||
"litellm_params": {"model": "gpt-4o"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.get_model_ids.return_value = ["global-id-1"]
|
||||
monkeypatch.setattr(
|
||||
ps,
|
||||
"get_all_team_models",
|
||||
AsyncMock(return_value={"byok-id-1": ["team-abc-123"]}),
|
||||
)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
result = await ps._populate_team_access_on_models(
|
||||
user_api_key_dict=admin,
|
||||
prisma_client=MagicMock(),
|
||||
llm_router=router,
|
||||
all_models=[team_row, global_row],
|
||||
)
|
||||
|
||||
by_id = {m["model_info"]["id"]: m for m in result}
|
||||
assert by_id["byok-id-1"]["model_info"]["direct_access"] is False
|
||||
assert by_id["global-id-1"]["model_info"]["direct_access"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_team_id_without_db_fails_fast(monkeypatch):
|
||||
"""`teamId` without a connected DB raises 500 before any enrichment work runs."""
|
||||
router = MagicMock()
|
||||
router.model_list = [_team_row()]
|
||||
|
||||
enrich_spy = MagicMock(side_effect=lambda model, **kw: model)
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
monkeypatch.setattr(ps, "_enrich_model_info_with_litellm_data", enrich_spy)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
|
||||
with pytest.raises(ps.HTTPException) as exc_info:
|
||||
await ps.model_info_v1(
|
||||
user_api_key_dict=admin, litellm_model_id=None, teamId="team-abc-123"
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "DB not connected" in exc_info.value.detail["error"]
|
||||
enrich_spy.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_include_team_models_without_db_fails_fast(monkeypatch):
|
||||
"""`include_team_models` without a connected DB raises 500 instead of silently
|
||||
returning an empty list (the access fields can only be populated from the DB)."""
|
||||
router = MagicMock()
|
||||
router.model_list = [_team_row()]
|
||||
|
||||
enrich_spy = MagicMock(side_effect=lambda model, **kw: model)
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
monkeypatch.setattr(ps, "_enrich_model_info_with_litellm_data", enrich_spy)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
|
||||
with pytest.raises(ps.HTTPException) as exc_info:
|
||||
await ps.model_info_v1(
|
||||
user_api_key_dict=admin, litellm_model_id=None, include_team_models=True
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "DB not connected" in exc_info.value.detail["error"]
|
||||
enrich_spy.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_litellm_model_id_team_id_without_db_fails_fast(
|
||||
monkeypatch,
|
||||
):
|
||||
"""`litellm_model_id` + `teamId` without a connected DB must raise 500 too, not
|
||||
return 200 with a model dict missing direct_access/access_via_team_ids."""
|
||||
router = MagicMock()
|
||||
router.model_list = [_team_row()]
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
|
||||
with pytest.raises(ps.HTTPException) as exc_info:
|
||||
await ps.model_info_v1(
|
||||
user_api_key_dict=admin,
|
||||
litellm_model_id="byok-id-1",
|
||||
teamId="team-abc-123",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "DB not connected" in exc_info.value.detail["error"]
|
||||
router.get_deployment.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_litellm_model_id_include_team_models_filters_inaccessible(
|
||||
monkeypatch,
|
||||
):
|
||||
"""`litellm_model_id` + `include_team_models` must drop a model the caller cannot
|
||||
use instead of returning it unconditionally from the single-model lookup."""
|
||||
team_row = _team_row()
|
||||
|
||||
router = MagicMock()
|
||||
deployment = MagicMock()
|
||||
deployment.model_dump.return_value = team_row
|
||||
router.get_deployment.return_value = deployment
|
||||
|
||||
async def _fake_populate(**kwargs):
|
||||
for model in kwargs["all_models"]:
|
||||
model["model_info"]["direct_access"] = False
|
||||
model["model_info"]["access_via_team_ids"] = []
|
||||
return kwargs["all_models"]
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", [team_row])
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row)
|
||||
monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[]
|
||||
)
|
||||
resp = await ps.model_info_v1(
|
||||
user_api_key_dict=caller,
|
||||
litellm_model_id="byok-id-1",
|
||||
include_team_models=True,
|
||||
)
|
||||
|
||||
assert resp["data"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_litellm_model_id_team_id_applies_team_filter(monkeypatch):
|
||||
"""`litellm_model_id` + `teamId` must run the teamId filter on the single model
|
||||
rather than returning it regardless of the team's access."""
|
||||
team_row = _team_row()
|
||||
|
||||
router = MagicMock()
|
||||
deployment = MagicMock()
|
||||
deployment.model_dump.return_value = team_row
|
||||
router.get_deployment.return_value = deployment
|
||||
|
||||
async def _fake_populate(**kwargs):
|
||||
return kwargs["all_models"]
|
||||
|
||||
team_filter = AsyncMock(return_value=[])
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", [team_row])
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row)
|
||||
monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate)
|
||||
monkeypatch.setattr(ps, "_filter_models_by_team_id", team_filter)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
resp = await ps.model_info_v1(
|
||||
user_api_key_dict=admin,
|
||||
litellm_model_id="byok-id-1",
|
||||
teamId="other-team",
|
||||
)
|
||||
|
||||
assert resp["data"] == []
|
||||
team_filter.assert_awaited_once()
|
||||
assert team_filter.await_args.kwargs["team_id"] == "other-team"
|
||||
assert team_filter.await_args.kwargs["all_models"] == [team_row]
|
||||
|
|
|
|||
|
|
@ -146,9 +146,9 @@ class TestModelInfoEndpointWithRouter:
|
|||
deployment_dict = deployment.model_dump(exclude_none=True)
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.model_list = [deployment_dict]
|
||||
mock_router.get_model_names.return_value = ["model1"]
|
||||
mock_router.get_model_access_groups.return_value = {}
|
||||
mock_router.get_model_list.return_value = [deployment_dict]
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
|
||||
|
||||
|
|
@ -156,6 +156,7 @@ class TestModelInfoEndpointWithRouter:
|
|||
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
||||
patch("litellm.proxy.proxy_server.llm_model_list", [deployment_dict]),
|
||||
patch("litellm.proxy.proxy_server.user_model", None),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None),
|
||||
patch("litellm.proxy.proxy_server.get_key_models", return_value=["model1"]),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.get_team_models", return_value=["model1"]
|
||||
|
|
|
|||
|
|
@ -795,6 +795,127 @@ class TestProxyInitializationHelpers:
|
|||
assert appended_params["pgbouncer"] == "true"
|
||||
assert appended_params["statement_cache_size"] == 0
|
||||
|
||||
def test_build_db_connection_url_params_disable_prepared_statements(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
||||
params = _build_db_connection_url_params(
|
||||
connection_limit=10,
|
||||
pool_timeout=60,
|
||||
disable_prepared_statements=True,
|
||||
)
|
||||
assert params["pgbouncer"] == "true"
|
||||
|
||||
def test_build_db_connection_url_params_no_pgbouncer_by_default(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
||||
params = _build_db_connection_url_params(
|
||||
connection_limit=10,
|
||||
pool_timeout=60,
|
||||
)
|
||||
assert "pgbouncer" not in params
|
||||
|
||||
def test_build_db_connection_url_params_extra_pgbouncer_overrides_flag(self):
|
||||
from litellm.proxy.proxy_cli import _build_db_connection_url_params
|
||||
|
||||
params = _build_db_connection_url_params(
|
||||
connection_limit=10,
|
||||
pool_timeout=60,
|
||||
disable_prepared_statements=True,
|
||||
extra_params={"pgbouncer": "false"},
|
||||
)
|
||||
assert params["pgbouncer"] == "false"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_value, expect_pgbouncer",
|
||||
[
|
||||
(True, True),
|
||||
(False, False),
|
||||
("true", True),
|
||||
("false", False),
|
||||
("not-a-bool", False),
|
||||
],
|
||||
)
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch(
|
||||
"litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False
|
||||
)
|
||||
def test_disable_prepared_statements_forwarded_to_url(
|
||||
self,
|
||||
mock_should_update,
|
||||
mock_setup_db,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
config_value,
|
||||
expect_pgbouncer,
|
||||
):
|
||||
from click.testing import CliRunner
|
||||
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
runner = CliRunner()
|
||||
mock_subprocess_run.return_value = MagicMock(returncode=0)
|
||||
|
||||
mock_proxy_module = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
|
||||
return_value={
|
||||
"general_settings": {
|
||||
"database_url": "postgresql://test:test@localhost:5432/test",
|
||||
"database_disable_prepared_statements": config_value,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL")
|
||||
}
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"proxy_server": mock_proxy_module,
|
||||
"litellm.proxy.proxy_server": mock_proxy_module,
|
||||
},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
) as mock_get_args,
|
||||
patch(
|
||||
"litellm.proxy.proxy_cli.append_query_params",
|
||||
side_effect=lambda url, params: str(url),
|
||||
) as mock_append_query_params,
|
||||
):
|
||||
mock_get_args.return_value = {
|
||||
"app": "litellm.proxy.proxy_server:app",
|
||||
"host": "localhost",
|
||||
"port": 8000,
|
||||
}
|
||||
|
||||
result = runner.invoke(
|
||||
run_server,
|
||||
["--local", "--config", "test-config.yaml", "--skip_server_startup"],
|
||||
)
|
||||
|
||||
assert (
|
||||
result.exit_code == 0
|
||||
), f"exit_code={result.exit_code}, output={result.output}"
|
||||
mock_append_query_params.assert_called()
|
||||
appended_params = mock_append_query_params.call_args.args[1]
|
||||
if expect_pgbouncer:
|
||||
assert appended_params["pgbouncer"] == "true"
|
||||
else:
|
||||
assert "pgbouncer" not in appended_params
|
||||
|
||||
@patch("uvicorn.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
|
|
|
|||
|
|
@ -3810,14 +3810,15 @@ async def test_model_info_v1_oci_secrets_not_leaked():
|
|||
|
||||
# Mock the llm_router to return our test data
|
||||
mock_router = MagicMock()
|
||||
mock_router.model_list = [mock_model_data]
|
||||
mock_router.get_model_names.return_value = ["oci-grok-test"]
|
||||
mock_router.get_model_access_groups.return_value = {}
|
||||
mock_router.get_model_list.return_value = [mock_model_data]
|
||||
|
||||
# Mock global variables
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
||||
patch("litellm.proxy.proxy_server.llm_model_list", [mock_model_data]),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"infer_model_from_keys": False},
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from __future__ import annotations
|
|||
|
||||
import hashlib
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -22,6 +23,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationTokenView
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
|
|
@ -193,6 +195,7 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row(
|
|||
) -> None:
|
||||
expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0}
|
||||
prisma_client.db.query_first = AsyncMock(return_value=expected)
|
||||
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
result = await prisma_client._query_first_with_cached_plan_fallback(
|
||||
"SELECT * FROM x WHERE token = $1", "abc"
|
||||
)
|
||||
|
|
@ -208,35 +211,110 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row(
|
|||
"args": ("SELECT * FROM x WHERE token = $1", "abc"),
|
||||
"matches": True,
|
||||
}
|
||||
prisma_client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_first_with_cached_plan_fallback_retries_on_cached_plan_error(
|
||||
async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_identical_query(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
original_query = 'SELECT * FROM "LiteLLM_VerificationToken" WHERE v.token = $1'
|
||||
expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0}
|
||||
manager = MagicMock()
|
||||
query_first = AsyncMock(
|
||||
side_effect=[
|
||||
RuntimeError("cached plan must not change result type"),
|
||||
expected,
|
||||
]
|
||||
)
|
||||
reconnect = AsyncMock(return_value=True)
|
||||
manager.attach_mock(query_first, "query_first")
|
||||
manager.attach_mock(reconnect, "attempt_db_reconnect")
|
||||
prisma_client.db.query_first = query_first
|
||||
prisma_client.attempt_db_reconnect = reconnect
|
||||
|
||||
result = await prisma_client._query_first_with_cached_plan_fallback(
|
||||
original_query, "abc"
|
||||
)
|
||||
|
||||
assert result == expected
|
||||
assert query_first.await_count == 2
|
||||
first_call, retry_call = query_first.await_args_list
|
||||
assert retry_call.args == first_call.args == (original_query, "abc")
|
||||
reconnect.assert_awaited_once()
|
||||
assert reconnect.await_args.kwargs.get("force", False) is False
|
||||
assert [name for name, *_ in manager.mock_calls] == [
|
||||
"query_first",
|
||||
"attempt_db_reconnect",
|
||||
"query_first",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_first_with_cached_plan_fallback_never_deallocates(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
expected = {"token": "abc"}
|
||||
prisma_client.db.query_first = AsyncMock(
|
||||
side_effect=[
|
||||
RuntimeError("cached plan must not change result type"),
|
||||
expected,
|
||||
]
|
||||
)
|
||||
result = await prisma_client._query_first_with_cached_plan_fallback(
|
||||
"SELECT * FROM x WHERE token = $1", "abc"
|
||||
prisma_client.db.execute_raw = AsyncMock(return_value=0)
|
||||
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
|
||||
|
||||
prisma_client.db.execute_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_first_with_cached_plan_fallback_propagates_when_retry_also_fails(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
plan_error = RuntimeError("cached plan must not change result type")
|
||||
prisma_client.db.query_first = AsyncMock(side_effect=[plan_error, plan_error])
|
||||
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
with pytest.raises(RuntimeError, match="cached plan must not change result type"):
|
||||
await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
|
||||
|
||||
assert prisma_client.db.query_first.await_count == 2
|
||||
prisma_client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_first_with_cached_plan_fallback_retries_when_reconnect_returns_false(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
expected = {"token": "abc"}
|
||||
prisma_client.db.query_first = AsyncMock(
|
||||
side_effect=[
|
||||
RuntimeError("cached plan must not change result type"),
|
||||
expected,
|
||||
]
|
||||
)
|
||||
prisma_client.attempt_db_reconnect = AsyncMock(return_value=False)
|
||||
|
||||
result = await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
|
||||
|
||||
assert result == expected
|
||||
assert prisma_client.db.query_first.await_count == 2
|
||||
second_call_sql = prisma_client.db.query_first.await_args_list[1].args[0]
|
||||
assert "cache_invalidated_" in second_call_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
prisma_client.db.query_first = AsyncMock(side_effect=RuntimeError("totally unrelated"))
|
||||
prisma_client.db.query_first = AsyncMock(
|
||||
side_effect=RuntimeError("totally unrelated")
|
||||
)
|
||||
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
with pytest.raises(RuntimeError, match="totally unrelated"):
|
||||
await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
|
||||
assert prisma_client.db.query_first.await_count == 1
|
||||
prisma_client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -351,7 +429,9 @@ async def test_get_data_token_find_unique_returns_record(
|
|||
async def test_get_data_token_find_unique_missing_token_raises_401(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
await prisma_client.get_data(token="sk-missing", table_name="key")
|
||||
err = excinfo.value
|
||||
|
|
@ -398,3 +478,39 @@ async def test_get_data_logs_and_raises_on_db_error(
|
|||
)
|
||||
with pytest.raises(RuntimeError, match="network split"):
|
||||
await prisma_client.get_data(token="sk-broken", table_name="key")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_data_combined_view_returns_view_for_deprecated_key(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
"""Grace-period rotation, full get_data flow: the old hash misses the
|
||||
combined view, the deprecated-key table resolves it to the active token,
|
||||
and get_data must return the recursive lookup's finished view instead of
|
||||
re-running dict normalization on it (which raised TypeError and turned
|
||||
every grace-period request into a 401)."""
|
||||
old_hash = "hashed-old-token-grace-e2e"
|
||||
active_hash = "hashed-active-token-grace-e2e"
|
||||
active_row = {
|
||||
"token": active_hash,
|
||||
"team_models": None,
|
||||
"team_blocked": None,
|
||||
"team_members_with_roles": None,
|
||||
"user_id": None,
|
||||
"expires": None,
|
||||
}
|
||||
prisma_client.db.query_first = AsyncMock(side_effect=[None, active_row])
|
||||
prisma_client.db.litellm_deprecatedverificationtoken = MagicMock()
|
||||
prisma_client.db.litellm_deprecatedverificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
active_token_id=active_hash,
|
||||
revoke_at=datetime.now(timezone.utc) + timedelta(hours=1),
|
||||
)
|
||||
)
|
||||
|
||||
response = await prisma_client.get_data(
|
||||
token=old_hash, table_name="combined_view", query_type="find_unique"
|
||||
)
|
||||
|
||||
assert isinstance(response, LiteLLM_VerificationTokenView)
|
||||
assert response.token == active_hash
|
||||
|
|
|
|||
|
|
@ -1351,11 +1351,6 @@
|
|||
"count": 1
|
||||
}
|
||||
},
|
||||
"src/components/mcp_tools/mcp_server_edit.test.tsx": {
|
||||
"unused-imports/no-unused-imports": {
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"src/components/mcp_tools/mcp_server_edit.tsx": {
|
||||
"no-restricted-imports": {
|
||||
"count": 1
|
||||
|
|
@ -1517,11 +1512,6 @@
|
|||
"count": 1
|
||||
}
|
||||
},
|
||||
"src/components/organisms/RegenerateKeyModal.tsx": {
|
||||
"react-hooks/set-state-in-effect": {
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"src/components/organisms/create_key_button.test.tsx": {
|
||||
"@typescript-eslint/no-require-imports": {
|
||||
"count": 2
|
||||
|
|
|
|||
116
ui/litellm-dashboard/package-lock.json
generated
116
ui/litellm-dashboard/package-lock.json
generated
|
|
@ -52,8 +52,8 @@
|
|||
"@types/react-dom": "18.3.7",
|
||||
"@types/react-syntax-highlighter": "15.5.13",
|
||||
"@types/uuid": "10.0.0",
|
||||
"@vitest/coverage-v8": "3.2.4",
|
||||
"@vitest/ui": "3.2.4",
|
||||
"@vitest/coverage-v8": "3.2.6",
|
||||
"@vitest/ui": "3.2.6",
|
||||
"autoprefixer": "10.4.24",
|
||||
"dotenv": "17.2.3",
|
||||
"eslint": "9.39.2",
|
||||
|
|
@ -69,7 +69,7 @@
|
|||
"typescript": "5.9.3",
|
||||
"typescript-eslint": "8.60.1",
|
||||
"vite": "7.3.2",
|
||||
"vitest": "3.2.4"
|
||||
"vitest": "3.2.6"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.9.0",
|
||||
|
|
@ -4276,9 +4276,9 @@
|
|||
]
|
||||
},
|
||||
"node_modules/@vitest/coverage-v8": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.4.tgz",
|
||||
"integrity": "sha512-EyF9SXU6kS5Ku/U82E259WSnvg6c8KTjppUncuNdm5QHpe17mwREHnjDzozC8x9MZ0xfBUFSaLkRv4TMA75ALQ==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.6.tgz",
|
||||
"integrity": "sha512-LsAdmUapA0qSN306d8+zOyawM0hFm2m2Hg9IwVNIKBm+qJV8cijiq2c+gxKZcB1HCfIWAy+0qEZDCUQA58A1cw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
|
|
@ -4300,8 +4300,8 @@
|
|||
"url": "https://opencollective.com/vitest"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@vitest/browser": "3.2.4",
|
||||
"vitest": "3.2.4"
|
||||
"@vitest/browser": "3.2.6",
|
||||
"vitest": "3.2.6"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
"@vitest/browser": {
|
||||
|
|
@ -4310,15 +4310,15 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/expect": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.4.tgz",
|
||||
"integrity": "sha512-Io0yyORnB6sikFlt8QW5K7slY4OjqNX9jmJQ02QDda8lyM6B5oNgVWoSoKPac8/kgnCUzuHQKrSLtu/uOqqrig==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.6.tgz",
|
||||
"integrity": "sha512-1+7q9BtaKzEmO+fmNT3kYvoNn5Y71XWAx2Q5HRim4tTVRQVRv4uJFAQ5FbK0OPUeNP/WmVCpxYxoJdvuHVjzBQ==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@types/chai": "^5.2.2",
|
||||
"@vitest/spy": "3.2.4",
|
||||
"@vitest/utils": "3.2.4",
|
||||
"@vitest/spy": "3.2.6",
|
||||
"@vitest/utils": "3.2.6",
|
||||
"chai": "^5.2.0",
|
||||
"tinyrainbow": "^2.0.0"
|
||||
},
|
||||
|
|
@ -4327,13 +4327,13 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/mocker": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.4.tgz",
|
||||
"integrity": "sha512-46ryTE9RZO/rfDd7pEqFl7etuyzekzEhUbTW3BvmeO/BcCMEgq59BKhek3dXDWgAj4oMK6OZi+vRr1wPW6qjEQ==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.6.tgz",
|
||||
"integrity": "sha512-EZOrpDbkKotFAP7wPAQV1UIyoGOk4oX7ynWhBhLB7v+meMHbQhU16oPpIYGTTe4oFlhpryGpgpcZP/sin3hYuw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@vitest/spy": "3.2.4",
|
||||
"@vitest/spy": "3.2.6",
|
||||
"estree-walker": "^3.0.3",
|
||||
"magic-string": "^0.30.17"
|
||||
},
|
||||
|
|
@ -4354,9 +4354,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/pretty-format": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.4.tgz",
|
||||
"integrity": "sha512-IVNZik8IVRJRTr9fxlitMKeJeXFFFN0JaB9PHPGQ8NKQbGpfjlTx9zO4RefN8gp7eqjNy8nyK3NZmBzOPeIxtA==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.6.tgz",
|
||||
"integrity": "sha512-lb7XXXzmm2h2ASzFnRvQpDo6onT1NmMJA3tkGTWiBFtRJ9lxGY3d3mm/Apt36gej2bkkOVLL/yTOtufDaFa/jA==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
|
|
@ -4367,13 +4367,13 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/runner": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.4.tgz",
|
||||
"integrity": "sha512-oukfKT9Mk41LreEW09vt45f8wx7DordoWUZMYdY/cyAk7w5TWkTRCNZYF7sX7n2wB7jyGAl74OxgwhPgKaqDMQ==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.6.tgz",
|
||||
"integrity": "sha512-HYcoSj1w5tcgUnzoF0HcyaAQjpA1gj9ftUJ7iSJSuipc02jW9gKkigwZbjFldAfYHA1fa8UZVRftdMY5msWM9Q==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@vitest/utils": "3.2.4",
|
||||
"@vitest/utils": "3.2.6",
|
||||
"pathe": "^2.0.3",
|
||||
"strip-literal": "^3.0.0"
|
||||
},
|
||||
|
|
@ -4382,13 +4382,13 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/snapshot": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.4.tgz",
|
||||
"integrity": "sha512-dEYtS7qQP2CjU27QBC5oUOxLE/v5eLkGqPE0ZKEIDGMs4vKWe7IjgLOeauHsR0D5YuuycGRO5oSRXnwnmA78fQ==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.6.tgz",
|
||||
"integrity": "sha512-H+ZjNTWGpObenh0YnlBctAPnJSI20P81PL8BPzWpx54YXLLTm8hEsWawtcYLMrwvpK48hGxLLbCS+1KRXhsKhw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@vitest/pretty-format": "3.2.4",
|
||||
"@vitest/pretty-format": "3.2.6",
|
||||
"magic-string": "^0.30.17",
|
||||
"pathe": "^2.0.3"
|
||||
},
|
||||
|
|
@ -4397,9 +4397,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/spy": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.4.tgz",
|
||||
"integrity": "sha512-vAfasCOe6AIK70iP5UD11Ac4siNUNJ9i/9PZ3NKx07sG6sUxeag1LWdNrMWeKKYBLlzuK+Gn65Yd5nyL6ds+nw==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.6.tgz",
|
||||
"integrity": "sha512-oq6BbH68WzcWmwtBrU9nqLeaXTR4XwJF7FSLkKEZo4i6eoXcrxjcwSuTvWBIRUTC6VC72nXYunzqgZA+IKdtxg==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
|
|
@ -4410,13 +4410,13 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@vitest/ui": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.4.tgz",
|
||||
"integrity": "sha512-hGISOaP18plkzbWEcP/QvtRW1xDXF2+96HbEX6byqQhAUbiS5oH6/9JwW+QsQCIYON2bI6QZBF+2PvOmrRZ9wA==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.6.tgz",
|
||||
"integrity": "sha512-mATfG3zVdhobE9U1rIpvtYD3DGuSSxqZ3Aj/8ityGqKXy8YDJ9BoAjZmAz6dZ1IZ1xI5V+MerkCczvVa+3QK9Q==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@vitest/utils": "3.2.4",
|
||||
"@vitest/utils": "3.2.6",
|
||||
"fflate": "^0.8.2",
|
||||
"flatted": "^3.3.3",
|
||||
"pathe": "^2.0.3",
|
||||
|
|
@ -4428,17 +4428,17 @@
|
|||
"url": "https://opencollective.com/vitest"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"vitest": "3.2.4"
|
||||
"vitest": "3.2.6"
|
||||
}
|
||||
},
|
||||
"node_modules/@vitest/utils": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.4.tgz",
|
||||
"integrity": "sha512-fB2V0JFrQSMsCo9HiSq3Ezpdv4iYaXRG1Sx8edX3MwxfyNn83mKiGzOcH+Fkxt4MHxr3y42fQi1oeAInqgX2QA==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.6.tgz",
|
||||
"integrity": "sha512-lI23nIs4bnT3T8NIoh+vFaz5s2/DdP0Jgt2jxwgWljvwn82cLJtyi/If+fjFyoLMGIOz0U/fKvWE0d4jsNQEfg==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@vitest/pretty-format": "3.2.4",
|
||||
"@vitest/pretty-format": "3.2.6",
|
||||
"loupe": "^3.1.4",
|
||||
"tinyrainbow": "^2.0.0"
|
||||
},
|
||||
|
|
@ -5027,9 +5027,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/brace-expansion": {
|
||||
"version": "5.0.5",
|
||||
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.5.tgz",
|
||||
"integrity": "sha512-VZznLgtwhn+Mact9tfiwx64fA9erHH/MCXEUfB/0bX/6Fz6ny5EGTXYltMocqg4xFAQZtnO3DHWWXi8RiuN7cQ==",
|
||||
"version": "5.0.6",
|
||||
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz",
|
||||
"integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
|
|
@ -6840,9 +6840,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/fflate": {
|
||||
"version": "0.8.2",
|
||||
"resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.2.tgz",
|
||||
"integrity": "sha512-cPJU47OaAoCbg0pBvzsgpTPhmhqI5eJjh/JIu8tPj5q+T7iLvW/JAYUqmE7KOB4R1ZyEhzBaIQpQpardBF5z8A==",
|
||||
"version": "0.8.3",
|
||||
"resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.3.tgz",
|
||||
"integrity": "sha512-tbZNuJrLwGUp3zshBtdy4W+ORxZuIh8a5ilyIEQDC5rY1f3U20JMry0Ll3WBzU58EZKsEuJFXhb5gwv8CsPvgA==",
|
||||
"dev": true,
|
||||
"license": "MIT"
|
||||
},
|
||||
|
|
@ -13499,20 +13499,20 @@
|
|||
}
|
||||
},
|
||||
"node_modules/vitest": {
|
||||
"version": "3.2.4",
|
||||
"resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.4.tgz",
|
||||
"integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==",
|
||||
"version": "3.2.6",
|
||||
"resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.6.tgz",
|
||||
"integrity": "sha512-xejya+bT/j/+R/AGa1XOfRxLmNUlLtlwjRsFUILF+xHfzElmGcmFydy2gqqIrd62ptIEfwVMofd19uNWD9L7Nw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@types/chai": "^5.2.2",
|
||||
"@vitest/expect": "3.2.4",
|
||||
"@vitest/mocker": "3.2.4",
|
||||
"@vitest/pretty-format": "^3.2.4",
|
||||
"@vitest/runner": "3.2.4",
|
||||
"@vitest/snapshot": "3.2.4",
|
||||
"@vitest/spy": "3.2.4",
|
||||
"@vitest/utils": "3.2.4",
|
||||
"@vitest/expect": "3.2.6",
|
||||
"@vitest/mocker": "3.2.6",
|
||||
"@vitest/pretty-format": "^3.2.6",
|
||||
"@vitest/runner": "3.2.6",
|
||||
"@vitest/snapshot": "3.2.6",
|
||||
"@vitest/spy": "3.2.6",
|
||||
"@vitest/utils": "3.2.6",
|
||||
"chai": "^5.2.0",
|
||||
"debug": "^4.4.1",
|
||||
"expect-type": "^1.2.1",
|
||||
|
|
@ -13542,8 +13542,8 @@
|
|||
"@edge-runtime/vm": "*",
|
||||
"@types/debug": "^4.1.12",
|
||||
"@types/node": "^18.0.0 || ^20.0.0 || >=22.0.0",
|
||||
"@vitest/browser": "3.2.4",
|
||||
"@vitest/ui": "3.2.4",
|
||||
"@vitest/browser": "3.2.6",
|
||||
"@vitest/ui": "3.2.6",
|
||||
"happy-dom": "*",
|
||||
"jsdom": "*"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -65,8 +65,8 @@
|
|||
"@types/react-dom": "18.3.7",
|
||||
"@types/react-syntax-highlighter": "15.5.13",
|
||||
"@types/uuid": "10.0.0",
|
||||
"@vitest/coverage-v8": "3.2.4",
|
||||
"@vitest/ui": "3.2.4",
|
||||
"@vitest/coverage-v8": "3.2.6",
|
||||
"@vitest/ui": "3.2.6",
|
||||
"autoprefixer": "10.4.24",
|
||||
"dotenv": "17.2.3",
|
||||
"eslint": "9.39.2",
|
||||
|
|
@ -82,7 +82,7 @@
|
|||
"typescript": "5.9.3",
|
||||
"typescript-eslint": "8.60.1",
|
||||
"vite": "7.3.2",
|
||||
"vitest": "3.2.4"
|
||||
"vitest": "3.2.6"
|
||||
},
|
||||
"overrides": {
|
||||
"prismjs": "1.30.0",
|
||||
|
|
@ -92,6 +92,7 @@
|
|||
"lodash": "4.18.1",
|
||||
"ws": "8.19.0",
|
||||
"braces": "3.0.3",
|
||||
"brace-expansion": "5.0.6",
|
||||
"axios": "1.13.6",
|
||||
"postcss": "8.5.13"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -17,15 +17,23 @@ vi.mock("@/utils/mcpTokenStore", () => ({
|
|||
}));
|
||||
|
||||
// Mutable holder so individual tests can simulate "Authorize & Fetch" having
|
||||
// produced a token before submit.
|
||||
const oauthHook = vi.hoisted(() => ({ tokenResponse: null as Record<string, unknown> | null }));
|
||||
// produced a token before submit, and inspect the reset wiring.
|
||||
const oauthHook = vi.hoisted(() => ({
|
||||
tokenResponse: null as Record<string, unknown> | null,
|
||||
reset: vi.fn(),
|
||||
onTokenReceived: null as ((token: Record<string, unknown> | null) => void) | null,
|
||||
}));
|
||||
vi.mock("@/hooks/useMcpOAuthFlow", () => ({
|
||||
useMcpOAuthFlow: () => ({
|
||||
startOAuthFlow: vi.fn(),
|
||||
status: "idle",
|
||||
error: null,
|
||||
tokenResponse: oauthHook.tokenResponse,
|
||||
}),
|
||||
useMcpOAuthFlow: (opts: { onTokenReceived: (token: Record<string, unknown> | null) => void }) => {
|
||||
oauthHook.onTokenReceived = opts.onTokenReceived;
|
||||
return {
|
||||
startOAuthFlow: vi.fn(),
|
||||
status: "idle",
|
||||
error: null,
|
||||
tokenResponse: oauthHook.tokenResponse,
|
||||
reset: oauthHook.reset,
|
||||
};
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("./mcp_server_cost_config", () => ({
|
||||
|
|
@ -59,7 +67,9 @@ vi.mock("./mcp_tool_configuration", () => ({
|
|||
}));
|
||||
|
||||
vi.mock("./mcp_connection_status", () => ({
|
||||
default: () => <div data-testid="mcp-connection-status" />,
|
||||
default: ({ tools }: { tools?: any[] }) => (
|
||||
<div data-testid="mcp-connection-status" data-tool-count={tools?.length ?? 0} />
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("./StdioConfiguration", () => ({
|
||||
|
|
@ -121,6 +131,7 @@ describe("CreateMCPServer", () => {
|
|||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
oauthHook.tokenResponse = null;
|
||||
oauthHook.onTokenReceived = null;
|
||||
});
|
||||
|
||||
it("should render the modal with title when visible", () => {
|
||||
|
|
@ -614,6 +625,100 @@ describe("CreateMCPServer", () => {
|
|||
|
||||
expect(defaultProps.setModalVisible).toHaveBeenCalledWith(false);
|
||||
});
|
||||
|
||||
it("does not leak a previous server's OAuth token into the next add-server session", async () => {
|
||||
const usedToken = (token: string) =>
|
||||
vi.mocked(networking.testMCPToolsListRequest).mock.calls.some((call) => call[2] === token);
|
||||
|
||||
const { rerender } = render(<CreateMCPServer {...defaultProps} />);
|
||||
|
||||
await selectAntOption("Transport Type", "Streamable HTTP");
|
||||
await waitFor(() => {
|
||||
expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument();
|
||||
});
|
||||
await selectAntOption("Authentication", "OAuth");
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await act(async () => {
|
||||
fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } });
|
||||
});
|
||||
|
||||
// Simulate "Authorize & Fetch Token" completing for server A.
|
||||
await act(async () => {
|
||||
oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 });
|
||||
});
|
||||
|
||||
// Precondition: the freshly fetched token drives the tool preview for server A.
|
||||
await waitFor(() => {
|
||||
expect(usedToken("stale-token-A")).toBe(true);
|
||||
});
|
||||
|
||||
// Parent hides the modal (Cancel / successful create both flip this prop).
|
||||
rerender(<CreateMCPServer {...defaultProps} isModalVisible={false} />);
|
||||
|
||||
// The OAuth flow state (source of the "Token fetched" badge) is reset on close.
|
||||
expect(oauthHook.reset).toHaveBeenCalled();
|
||||
|
||||
vi.mocked(networking.testMCPToolsListRequest).mockClear();
|
||||
|
||||
// Reopen for a brand-new server and enter a different URL without re-authorizing.
|
||||
rerender(<CreateMCPServer {...defaultProps} isModalVisible={true} />);
|
||||
const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await act(async () => {
|
||||
fireEvent.change(reopenedUrlInput, { target: { value: "https://server-b.example.com/mcp" } });
|
||||
});
|
||||
|
||||
// The previous server's token must never be replayed for the new session.
|
||||
expect(usedToken("stale-token-A")).toBe(false);
|
||||
});
|
||||
|
||||
it("clears the tool list and form fields when a parent dismisses the modal", async () => {
|
||||
vi.mocked(networking.testMCPToolsListRequest).mockResolvedValue({
|
||||
tools: [{ name: "tool_a" }],
|
||||
error: null,
|
||||
});
|
||||
const toolCount = () => screen.getByTestId("mcp-connection-status").getAttribute("data-tool-count");
|
||||
|
||||
const { rerender } = render(<CreateMCPServer {...defaultProps} />);
|
||||
|
||||
await selectAntOption("Transport Type", "Streamable HTTP");
|
||||
await waitFor(() => {
|
||||
expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument();
|
||||
});
|
||||
await selectAntOption("Authentication", "OAuth");
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await act(async () => {
|
||||
fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } });
|
||||
});
|
||||
await act(async () => {
|
||||
oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 });
|
||||
});
|
||||
|
||||
// Precondition: a tool list is shown for server A.
|
||||
await waitFor(() => {
|
||||
expect(toolCount()).toBe("1");
|
||||
});
|
||||
|
||||
// Parent dismisses the modal without routing through Cancel or create.
|
||||
rerender(<CreateMCPServer {...defaultProps} isModalVisible={false} />);
|
||||
|
||||
// Stale tools are cleared even though neither handler ran.
|
||||
await waitFor(() => {
|
||||
expect(toolCount()).toBe("0");
|
||||
});
|
||||
|
||||
// Reopening starts clean: the URL the prior server left in the Ant form store is gone.
|
||||
rerender(<CreateMCPServer {...defaultProps} isModalVisible={true} />);
|
||||
const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com") as HTMLInputElement;
|
||||
expect(reopenedUrlInput.value).toBe("");
|
||||
});
|
||||
});
|
||||
|
||||
describe("when stdio transport is selected", () => {
|
||||
|
|
|
|||
|
|
@ -134,6 +134,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
|||
status: oauthStatus,
|
||||
error: oauthError,
|
||||
tokenResponse: oauthTokenResponse,
|
||||
reset: resetOAuthFlow,
|
||||
} = useMcpOAuthFlow({
|
||||
accessToken,
|
||||
getCredentials: () => form.getFieldValue("credentials"),
|
||||
|
|
@ -188,6 +189,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
|||
}
|
||||
},
|
||||
onBeforeRedirect: persistCreateUiState,
|
||||
flowSource: "create",
|
||||
});
|
||||
|
||||
React.useEffect(() => {
|
||||
|
|
@ -553,12 +555,19 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
|||
}
|
||||
}, [formValues.server_name]);
|
||||
|
||||
// Clear formValues when modal closes to reset child components
|
||||
// Clear form, tools, and OAuth state when the modal closes so a previous server's
|
||||
// authorization, credentials, or tool list never bleed into the next "Add New MCP
|
||||
// Server" session, including when a parent dismisses the modal without routing
|
||||
// through handleCancel or handleCreate.
|
||||
React.useEffect(() => {
|
||||
if (!isModalVisible) {
|
||||
form.resetFields();
|
||||
setFormValues({});
|
||||
setOauthAccessToken(null);
|
||||
clearTools();
|
||||
resetOAuthFlow();
|
||||
}
|
||||
}, [isModalVisible]);
|
||||
}, [isModalVisible, form, clearTools, resetOAuthFlow]);
|
||||
|
||||
const isAdmin = isAdminRole(userRole);
|
||||
|
||||
|
|
@ -1088,7 +1097,6 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
|||
<div className="mt-6">
|
||||
<MCPToolConfiguration
|
||||
accessToken={accessToken}
|
||||
oauthAccessToken={oauthAccessToken}
|
||||
formValues={formValues}
|
||||
allowedTools={allowedTools}
|
||||
existingAllowedTools={null}
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import NotificationsManager from "../molecules/notifications_manager";
|
|||
vi.mock("../networking", () => ({
|
||||
updateMCPServer: vi.fn(),
|
||||
listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }),
|
||||
storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}),
|
||||
}));
|
||||
|
||||
vi.mock("../molecules/notifications_manager", () => ({
|
||||
|
|
@ -17,12 +18,13 @@ vi.mock("../molecules/notifications_manager", () => ({
|
|||
},
|
||||
}));
|
||||
|
||||
const mockOauth: { tokenResponse: any } = { tokenResponse: null };
|
||||
vi.mock("@/hooks/useMcpOAuthFlow", () => ({
|
||||
useMcpOAuthFlow: () => ({
|
||||
startOAuthFlow: vi.fn(),
|
||||
status: "idle",
|
||||
error: null,
|
||||
tokenResponse: null,
|
||||
tokenResponse: mockOauth.tokenResponse,
|
||||
}),
|
||||
}));
|
||||
|
||||
|
|
@ -37,12 +39,19 @@ vi.mock("./MCPPermissionManagement", () => ({
|
|||
vi.mock("./mcp_tool_configuration", () => ({
|
||||
default: ({
|
||||
existingAllowedTools,
|
||||
externalTools,
|
||||
externalError,
|
||||
onAllowedToolsChange,
|
||||
onToolAllowlistInteraction,
|
||||
onToolNameToDisplayNameChange,
|
||||
onToolNameToDescriptionChange,
|
||||
}: any) => (
|
||||
<div data-testid="mcp-tool-config" data-existing-allowed-tools={JSON.stringify(existingAllowedTools)}>
|
||||
<div
|
||||
data-testid="mcp-tool-config"
|
||||
data-existing-allowed-tools={JSON.stringify(existingAllowedTools)}
|
||||
data-external-tools={JSON.stringify(externalTools)}
|
||||
data-external-error={externalError ?? ""}
|
||||
>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
|
|
@ -65,6 +74,15 @@ vi.mock("./mcp_tool_configuration", () => ({
|
|||
),
|
||||
}));
|
||||
|
||||
const mockGetToken = vi.fn();
|
||||
const mockIsTokenValid = vi.fn();
|
||||
const mockSetToken = vi.fn();
|
||||
vi.mock("@/utils/mcpTokenStore", () => ({
|
||||
getToken: (...args: any[]) => mockGetToken(...args),
|
||||
isTokenValid: (...args: any[]) => mockIsTokenValid(...args),
|
||||
setToken: (...args: any[]) => mockSetToken(...args),
|
||||
}));
|
||||
|
||||
// ── fixtures ──────────────────────────────────────────────────────────────────
|
||||
|
||||
const interactiveOAuthServer = {
|
||||
|
|
@ -626,3 +644,299 @@ describe("MCPServerEdit (interactive OAuth)", () => {
|
|||
expect(payload.token_storage_ttl_seconds).toBe(7200);
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPServerEdit (tool list fetch)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: [], error: null });
|
||||
mockOauth.tokenResponse = null;
|
||||
});
|
||||
|
||||
it("loads an OBO server's tools via GET listMCPTools with no passthrough headers", async () => {
|
||||
vi.mocked(networking.listMCPTools).mockResolvedValue({
|
||||
tools: [{ name: "read_user" }],
|
||||
error: null,
|
||||
});
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
// includeDisabledTools=true so the config screen gets the full catalog.
|
||||
expect(networking.listMCPTools).toHaveBeenCalledWith("access-token", "oauth_server_1", undefined, true);
|
||||
});
|
||||
// OBO uses the backend-stored token; the browser passthrough store is never consulted.
|
||||
expect(mockIsTokenValid).not.toHaveBeenCalled();
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("mcp-tool-config")).toHaveAttribute(
|
||||
"data-external-tools",
|
||||
JSON.stringify([{ name: "read_user" }]),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it("forwards the sessionStorage token as the x-mcp passthrough header for a passthrough server", async () => {
|
||||
mockIsTokenValid.mockReturnValue(true);
|
||||
mockGetToken.mockReturnValue({ access_token: "browser-token" });
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={{ ...interactiveOAuthServer, delegate_auth_to_upstream: true }}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.listMCPTools).toHaveBeenCalledWith(
|
||||
"access-token",
|
||||
"oauth_server_1",
|
||||
{ "x-mcp-oauth_server-authorization": "Bearer browser-token" },
|
||||
true,
|
||||
);
|
||||
});
|
||||
expect(mockGetToken).toHaveBeenCalledWith("oauth_server_1", "user-1");
|
||||
});
|
||||
|
||||
it("uses the staged OAuth token to load passthrough tools after authorize", async () => {
|
||||
const passthroughServer = { ...interactiveOAuthServer, delegate_auth_to_upstream: true };
|
||||
mockIsTokenValid.mockReturnValue(false);
|
||||
vi.mocked(networking.listMCPTools).mockResolvedValue({
|
||||
tools: [{ name: "read_user" }],
|
||||
error: null,
|
||||
});
|
||||
|
||||
const props = {
|
||||
mcpServer: passthroughServer,
|
||||
accessToken: "access-token",
|
||||
userID: "user-1",
|
||||
onCancel: vi.fn(),
|
||||
onSuccess: vi.fn(),
|
||||
availableAccessGroups: [],
|
||||
};
|
||||
|
||||
const { rerender } = render(<MCPServerEdit {...props} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("mcp-tool-config").getAttribute("data-external-error")).toContain(
|
||||
"Authenticate with this server in the Tools tab",
|
||||
);
|
||||
});
|
||||
expect(networking.listMCPTools).not.toHaveBeenCalled();
|
||||
|
||||
mockOauth.tokenResponse = { access_token: "staged-token", expires_in: 1800 };
|
||||
rerender(<MCPServerEdit {...props} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.listMCPTools).toHaveBeenCalledWith(
|
||||
"access-token",
|
||||
"oauth_server_1",
|
||||
{ "x-mcp-oauth_server-authorization": "Bearer staged-token" },
|
||||
true,
|
||||
);
|
||||
});
|
||||
expect(mockGetToken).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("prompts to authenticate and does not fetch when a passthrough server has no session token", async () => {
|
||||
mockIsTokenValid.mockReturnValue(false);
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={{ ...interactiveOAuthServer, delegate_auth_to_upstream: true }}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("mcp-tool-config").getAttribute("data-external-error")).toContain(
|
||||
"Authenticate with this server in the Tools tab",
|
||||
);
|
||||
});
|
||||
expect(networking.listMCPTools).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPServerEdit (form resync)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: [], error: null });
|
||||
});
|
||||
|
||||
it("repopulates the form when the server data arrives after mount", async () => {
|
||||
const props = {
|
||||
accessToken: "access-token",
|
||||
onCancel: vi.fn(),
|
||||
onSuccess: vi.fn(),
|
||||
availableAccessGroups: [],
|
||||
};
|
||||
|
||||
// Mount before the server is loaded (mirrors landing on the page mid OAuth return).
|
||||
const { rerender } = render(<MCPServerEdit mcpServer={{ server_id: "" } as any} {...props} />);
|
||||
expect(screen.queryByDisplayValue("https://example.com/mcp")).not.toBeInTheDocument();
|
||||
|
||||
// Server data arrives; the form must repopulate rather than staying blank.
|
||||
rerender(<MCPServerEdit mcpServer={interactiveOAuthServer} {...props} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByDisplayValue("https://example.com/mcp")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPServerEdit (OAuth token persistence on save)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: [], error: null });
|
||||
mockOauth.tokenResponse = null;
|
||||
});
|
||||
|
||||
it("persists the OBO token to the DB on save after authorize", async () => {
|
||||
mockOauth.tokenResponse = {
|
||||
access_token: "obo-tok",
|
||||
refresh_token: "obo-refresh",
|
||||
expires_in: 3600,
|
||||
scope: "read write",
|
||||
};
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer });
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await act(async () => {
|
||||
fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledWith(
|
||||
"access-token",
|
||||
"oauth_server_1",
|
||||
expect.objectContaining({
|
||||
access_token: "obo-tok",
|
||||
refresh_token: "obo-refresh",
|
||||
expires_in: 3600,
|
||||
scopes: ["read", "write"],
|
||||
}),
|
||||
);
|
||||
});
|
||||
expect(mockSetToken).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not show success when OBO token persistence fails after update", async () => {
|
||||
mockOauth.tokenResponse = {
|
||||
access_token: "obo-tok",
|
||||
refresh_token: "obo-refresh",
|
||||
expires_in: 3600,
|
||||
scope: "read write",
|
||||
};
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer });
|
||||
vi.mocked(networking.storeMCPOAuthUserCredential).mockRejectedValueOnce(new Error("write failed"));
|
||||
const onSuccess = vi.fn();
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={onSuccess}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await act(async () => {
|
||||
fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalled();
|
||||
});
|
||||
expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(
|
||||
"MCP Server updated, but failed to persist OAuth token: write failed",
|
||||
);
|
||||
expect(NotificationsManager.success).not.toHaveBeenCalledWith("MCP Server updated successfully");
|
||||
expect(onSuccess).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("persists the passthrough token to sessionStorage on save after authorize", async () => {
|
||||
mockOauth.tokenResponse = { access_token: "pt-tok", expires_in: 1800, token_type: "bearer" };
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({
|
||||
...interactiveOAuthServer,
|
||||
delegate_auth_to_upstream: true,
|
||||
});
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={{ ...interactiveOAuthServer, delegate_auth_to_upstream: true }}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await act(async () => {
|
||||
fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockSetToken).toHaveBeenCalledWith(
|
||||
"oauth_server_1",
|
||||
expect.objectContaining({ access_token: "pt-tok", expires_in: 1800, token_type: "bearer" }),
|
||||
"user-1",
|
||||
);
|
||||
});
|
||||
expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("persists nothing on save when no token was fetched", async () => {
|
||||
mockOauth.tokenResponse = null;
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer });
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
userID="user-1"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await act(async () => {
|
||||
fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPServer).toHaveBeenCalled();
|
||||
});
|
||||
expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled();
|
||||
expect(mockSetToken).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -2,8 +2,18 @@ import React, { useState, useEffect } from "react";
|
|||
import { Form, Select, Button as AntdButton, Tooltip, Input, InputNumber } from "antd";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
|
||||
import { AUTH_TYPE, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types";
|
||||
import { updateMCPServer, listMCPTools } from "../networking";
|
||||
import {
|
||||
AUTH_TYPE,
|
||||
OAUTH_FLOW,
|
||||
MCP_OAUTH2_FLOW_M2M,
|
||||
MCPServer,
|
||||
MCPServerCostInfo,
|
||||
TRANSPORT,
|
||||
getMcpOAuthMode,
|
||||
} from "./types";
|
||||
import { updateMCPServer, listMCPTools, storeMCPOAuthUserCredential } from "../networking";
|
||||
import { getToken, isTokenValid, setToken } from "@/utils/mcpTokenStore";
|
||||
import { buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils";
|
||||
import MCPServerCostConfig from "./mcp_server_cost_config";
|
||||
import MCPPermissionManagement from "./MCPPermissionManagement";
|
||||
import MCPToolConfiguration from "./mcp_tool_configuration";
|
||||
|
|
@ -18,6 +28,7 @@ import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
|||
interface MCPServerEditProps {
|
||||
mcpServer: MCPServer;
|
||||
accessToken: string | null;
|
||||
userID?: string | null;
|
||||
onCancel: () => void;
|
||||
onSuccess: (server: MCPServer) => void;
|
||||
availableAccessGroups: string[];
|
||||
|
|
@ -25,11 +36,12 @@ interface MCPServerEditProps {
|
|||
|
||||
const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC];
|
||||
const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2, AUTH_TYPE.AWS_SIGV4];
|
||||
const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
|
||||
export const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
|
||||
|
||||
const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
mcpServer,
|
||||
accessToken,
|
||||
userID,
|
||||
onCancel,
|
||||
onSuccess,
|
||||
availableAccessGroups,
|
||||
|
|
@ -58,8 +70,6 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
const oauthFlowTypeValue = Form.useWatch("oauth_flow_type", form) as string | undefined;
|
||||
const isM2MFlow = isOAuthAuthType && oauthFlowTypeValue === OAUTH_FLOW.M2M;
|
||||
|
||||
const [oauthAccessToken, setOauthAccessToken] = useState<string | null>(null);
|
||||
|
||||
// Watch form fields that affect tool fetching
|
||||
const currentUrl = Form.useWatch("url", form);
|
||||
const currentSpecPath = Form.useWatch("spec_path", form);
|
||||
|
|
@ -140,8 +150,6 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
};
|
||||
},
|
||||
onTokenReceived: (token) => {
|
||||
setOauthAccessToken(token?.access_token ?? null);
|
||||
|
||||
if (token?.access_token) {
|
||||
const credentials = {
|
||||
access_token: token.access_token,
|
||||
|
|
@ -158,6 +166,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
}
|
||||
},
|
||||
onBeforeRedirect: persistEditUiState,
|
||||
flowSource: "edit",
|
||||
});
|
||||
|
||||
const initialStaticHeaders = React.useMemo(() => {
|
||||
|
|
@ -217,6 +226,20 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
[mcpServer, effectiveTransport, initialStaticHeaders, initialEnvVars, initialEnvJson],
|
||||
);
|
||||
|
||||
// antd applies `initialValues` only at first mount. When the server loads after
|
||||
// mount (e.g. returning from the OAuth redirect lands on Overview and the form
|
||||
// mounts before the server data is ready), the form would stay blank. Re-sync it
|
||||
// from the loaded server once per server_id so it always reflects the saved config;
|
||||
// the OAuth-restore effect below then overlays any in-progress edits on top.
|
||||
const syncedServerIdRef = React.useRef<string | null>(null);
|
||||
useEffect(() => {
|
||||
if (!mcpServer.server_id || syncedServerIdRef.current === mcpServer.server_id) {
|
||||
return;
|
||||
}
|
||||
syncedServerIdRef.current = mcpServer.server_id;
|
||||
form.setFieldsValue(initialValues);
|
||||
}, [mcpServer.server_id, initialValues, form]);
|
||||
|
||||
// Initialize cost config from existing server data
|
||||
useEffect(() => {
|
||||
if (mcpServer.mcp_info?.mcp_server_cost_info) {
|
||||
|
|
@ -280,6 +303,9 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
if (!pendingRestoredValues) {
|
||||
return;
|
||||
}
|
||||
// Set transport first so transport-dependent fields render, then apply the rest
|
||||
// on the re-run triggered by the transportType watch (without it the effect's
|
||||
// deps never change and the second pass never runs, leaving fields blank).
|
||||
const transport = pendingRestoredValues.transport || mcpServer.transport;
|
||||
if (transport && transport !== form.getFieldValue("transport")) {
|
||||
form.setFieldsValue({ transport });
|
||||
|
|
@ -287,7 +313,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
}
|
||||
form.setFieldsValue(pendingRestoredValues);
|
||||
setPendingRestoredValues(null);
|
||||
}, [pendingRestoredValues, form, mcpServer.transport]);
|
||||
}, [pendingRestoredValues, form, mcpServer.transport, transportType]);
|
||||
|
||||
// Transform string array to object array for initial form values
|
||||
useEffect(() => {
|
||||
|
|
@ -304,28 +330,52 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
return;
|
||||
}
|
||||
fetchTools();
|
||||
}, [mcpServer, accessToken]);
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token]);
|
||||
|
||||
const fetchTools = async () => {
|
||||
if (!accessToken || !mcpServer.server_id) return;
|
||||
|
||||
// OBO/M2M/static auth is attached server-side from the stored credential, so
|
||||
// a plain GET /tools/list?server_id suffices. PKCE passthrough holds the token
|
||||
// in the browser, so forward it from sessionStorage as the x-mcp header the
|
||||
// same way the Tools playground does.
|
||||
let customHeaders: Record<string, string> | undefined;
|
||||
const isPassthrough =
|
||||
getMcpOAuthMode({
|
||||
auth_type: mcpServer.auth_type,
|
||||
oauth2_flow: mcpServer.oauth2_flow,
|
||||
delegate_auth_to_upstream: mcpServer.delegate_auth_to_upstream,
|
||||
}) === "passthrough";
|
||||
if (isPassthrough) {
|
||||
const token =
|
||||
oauthTokenResponse?.access_token ??
|
||||
(isTokenValid(mcpServer.server_id, userID)
|
||||
? getToken(mcpServer.server_id, userID)?.access_token ?? null
|
||||
: null);
|
||||
if (!token) {
|
||||
setTools([]);
|
||||
setToolsError("Authenticate with this server in the Tools tab to load and configure its tools.");
|
||||
return;
|
||||
}
|
||||
customHeaders = buildMcpPassthroughAuthHeader(mcpServer.alias, token);
|
||||
}
|
||||
|
||||
setIsLoadingTools(true);
|
||||
setToolsError(null);
|
||||
|
||||
try {
|
||||
// Use the GET endpoint which looks up stored credentials by server_id,
|
||||
// rather than POST /test/tools/list which requires inline credentials.
|
||||
const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id);
|
||||
// include_disabled_tools: configuring the allowlist needs the full server
|
||||
// catalog, so tools toggled off still render (as unchecked) instead of vanishing.
|
||||
const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id, customHeaders, true);
|
||||
|
||||
if (toolsResponse.tools && !toolsResponse.error) {
|
||||
setTools(toolsResponse.tools);
|
||||
} else {
|
||||
console.error("Failed to fetch tools:", toolsResponse.message);
|
||||
setTools([]);
|
||||
setToolsError(toolsResponse.message || "Failed to load tools");
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Tools fetch error:", error);
|
||||
setTools([]);
|
||||
setToolsError(error instanceof Error ? error.message : "Failed to load tools");
|
||||
} finally {
|
||||
|
|
@ -627,6 +677,46 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
}
|
||||
|
||||
const updated = await updateMCPServer(accessToken, payload);
|
||||
|
||||
// Persist the token staged via "Authorize & Fetch" (mirrors the create flow's
|
||||
// commit-on-submit): OBO writes the per-user token to the DB, passthrough keeps
|
||||
// it in sessionStorage. M2M/static auth resolve server-side and need neither.
|
||||
if (oauthTokenResponse?.access_token) {
|
||||
const oauthMode = getMcpOAuthMode({
|
||||
auth_type: restValues.auth_type,
|
||||
oauth2_flow: isM2MFlow ? MCP_OAUTH2_FLOW_M2M : null,
|
||||
delegate_auth_to_upstream: Boolean(delegateAuthToUpstreamRaw ?? mcpServer.delegate_auth_to_upstream),
|
||||
});
|
||||
try {
|
||||
if (oauthMode === "obo") {
|
||||
const scope = oauthTokenResponse.scope;
|
||||
await storeMCPOAuthUserCredential(accessToken, mcpServer.server_id, {
|
||||
access_token: oauthTokenResponse.access_token,
|
||||
refresh_token: oauthTokenResponse.refresh_token,
|
||||
expires_in: oauthTokenResponse.expires_in,
|
||||
scopes: typeof scope === "string" && scope ? scope.split(" ") : undefined,
|
||||
});
|
||||
} else if (oauthMode === "passthrough") {
|
||||
setToken(
|
||||
mcpServer.server_id,
|
||||
{
|
||||
access_token: oauthTokenResponse.access_token,
|
||||
expires_in: oauthTokenResponse.expires_in,
|
||||
refresh_token: oauthTokenResponse.refresh_token,
|
||||
token_type: oauthTokenResponse.token_type,
|
||||
},
|
||||
userID,
|
||||
);
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
const message = error instanceof Error ? error.message : "";
|
||||
NotificationsManager.fromBackend(
|
||||
"MCP Server updated, but failed to persist OAuth token" + (message ? `: ${message}` : ""),
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
NotificationsManager.success("MCP Server updated successfully");
|
||||
onSuccess(updated);
|
||||
} catch (error: any) {
|
||||
|
|
@ -1146,7 +1236,6 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
<div className="mt-6">
|
||||
<MCPToolConfiguration
|
||||
accessToken={accessToken}
|
||||
oauthAccessToken={oauthAccessToken}
|
||||
formValues={{
|
||||
server_id: mcpServer.server_id,
|
||||
server_name: currentServerName ?? mcpServer.server_name,
|
||||
|
|
@ -1172,6 +1261,10 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
|||
toolNameToDescription={toolNameToDescription}
|
||||
onToolNameToDisplayNameChange={setToolNameToDisplayName}
|
||||
onToolNameToDescriptionChange={setToolNameToDescription}
|
||||
externalTools={tools}
|
||||
externalIsLoading={isLoadingTools}
|
||||
externalError={toolsError}
|
||||
externalCanFetch={true}
|
||||
/>
|
||||
</div>
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,8 @@ import { Title, Card, Button, Text, Grid, TabGroup, TabList, TabPanel, TabPanels
|
|||
import { MCPServer, handleTransport, handleAuth } from "./types";
|
||||
// TODO: Move Tools viewer from index file
|
||||
import { MCPToolsViewer } from ".";
|
||||
import MCPServerEdit from "./mcp_server_edit";
|
||||
import MCPServerEdit, { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit";
|
||||
import { getSecureItem } from "@/utils/secureStorage";
|
||||
import MCPServerCostDisplay from "./mcp_server_cost_display";
|
||||
import { getMaskedAndFullUrl } from "./utils";
|
||||
import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
|
||||
|
|
@ -21,6 +22,25 @@ interface MCPServerViewProps {
|
|||
userRole: string | null;
|
||||
userID: string | null;
|
||||
availableAccessGroups: string[];
|
||||
initialTabIndex?: number;
|
||||
}
|
||||
|
||||
// True when this render is the return from the edit-settings OAuth redirect for this
|
||||
// server: the edit form wrote its UI-state snapshot before redirecting. Used to open
|
||||
// the editing Settings tab on first render instead of defaulting to Overview.
|
||||
function isReturningFromEditOAuth(isProxyAdmin: boolean, serverId: string): boolean {
|
||||
if (typeof window === "undefined" || !isProxyAdmin) {
|
||||
return false;
|
||||
}
|
||||
const stored = getSecureItem(EDIT_OAUTH_UI_STATE_KEY);
|
||||
if (!stored) {
|
||||
return false;
|
||||
}
|
||||
try {
|
||||
return JSON.parse(stored)?.serverId === serverId;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
export const MCPServerView: React.FC<MCPServerViewProps> = ({
|
||||
|
|
@ -32,11 +52,15 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
|
|||
userRole,
|
||||
userID,
|
||||
availableAccessGroups,
|
||||
initialTabIndex = 0,
|
||||
}) => {
|
||||
const [editing, setEditing] = useState(isEditing);
|
||||
// Open the editing Settings tab on first render when returning from the edit OAuth
|
||||
// redirect, so the "token fetched" feedback shows where the user left off (Settings=2).
|
||||
const returningFromEditOAuth = isReturningFromEditOAuth(isProxyAdmin, mcpServer.server_id);
|
||||
const [editing, setEditing] = useState(isEditing || returningFromEditOAuth);
|
||||
const [showFullUrl, setShowFullUrl] = useState(false);
|
||||
const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
|
||||
const [selectedTabIndex, setSelectedTabIndex] = useState(0);
|
||||
const [selectedTabIndex, setSelectedTabIndex] = useState(returningFromEditOAuth ? 2 : initialTabIndex);
|
||||
|
||||
const handleSuccess = (updated: MCPServer) => {
|
||||
setEditing(false);
|
||||
|
|
@ -210,6 +234,7 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
|
|||
<MCPServerEdit
|
||||
mcpServer={mcpServer}
|
||||
accessToken={accessToken}
|
||||
userID={userID}
|
||||
onCancel={() => setEditing(false)}
|
||||
onSuccess={handleSuccess}
|
||||
availableAccessGroups={availableAccessGroups}
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ import MCPNetworkSettings from "./MCPNetworkSettings";
|
|||
import MCPDiscovery from "./mcp_discovery";
|
||||
import { ByokCredentialModal } from "./ByokCredentialModal";
|
||||
import { getSecureItem } from "@/utils/secureStorage";
|
||||
import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils";
|
||||
import UserEnvVarsModal from "./UserEnvVarsModal";
|
||||
import { listMCPUserEnvVarStatus } from "../networking";
|
||||
|
||||
|
|
@ -71,6 +72,23 @@ const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
|
|||
const { Text: AntdText, Title: AntdTitle } = Typography;
|
||||
const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
|
||||
|
||||
// Server id stashed by the Tools tab before an OBO OAuth redirect, read once at
|
||||
// mount so the redirect returns straight to that server's Tools tab.
|
||||
const readToolsOAuthServerId = (): string | null => {
|
||||
if (typeof window === "undefined") {
|
||||
return null;
|
||||
}
|
||||
try {
|
||||
const stored = getSecureItem(TOOLS_OAUTH_UI_STATE_KEY);
|
||||
if (!stored) {
|
||||
return null;
|
||||
}
|
||||
return JSON.parse(stored)?.serverId ?? null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
};
|
||||
|
||||
const { Option } = Select;
|
||||
|
||||
const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID }) => {
|
||||
|
|
@ -103,7 +121,12 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
|||
// state
|
||||
const [serverIdToDelete, setServerToDelete] = useState<string | null>(null);
|
||||
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
|
||||
const [selectedServerId, setSelectedServerId] = useState<string | null>(null);
|
||||
// Server whose Tools tab should be reopened after an OBO OAuth redirect; read
|
||||
// once from sessionStorage so the restored server selection is correct on the
|
||||
// first render. Cleared when the user navigates back to the list (handleBack)
|
||||
// so a later visit to the same server defaults to Overview, not the Tools tab.
|
||||
const [toolsTabServerId, setToolsTabServerId] = useState<string | null>(readToolsOAuthServerId);
|
||||
const [selectedServerId, setSelectedServerId] = useState<string | null>(toolsTabServerId);
|
||||
const [editServer, setEditServer] = useState(false);
|
||||
const [selectedTeam, setSelectedTeam] = useState<string>("all");
|
||||
const [selectedMcpAccessGroup, setSelectedMcpAccessGroup] = useState<string>("all");
|
||||
|
|
@ -178,6 +201,19 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
|||
}
|
||||
}, []);
|
||||
|
||||
// The restored server id was consumed by the initializer above; remove the
|
||||
// one-shot sessionStorage key so a full page reload doesn't reopen the Tools
|
||||
// tab (removeItem only, no setState).
|
||||
useEffect(() => {
|
||||
if (typeof window !== "undefined") {
|
||||
try {
|
||||
window.sessionStorage.removeItem(TOOLS_OAUTH_UI_STATE_KEY);
|
||||
} catch {
|
||||
// ignore storage errors
|
||||
}
|
||||
}
|
||||
}, []);
|
||||
|
||||
// Get unique teams from all servers
|
||||
const uniqueTeams = React.useMemo(() => {
|
||||
if (!serversWithHealth) return [];
|
||||
|
|
@ -338,6 +374,8 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
|||
const handleBack = React.useCallback(() => {
|
||||
setEditServer(false);
|
||||
setSelectedServerId(null);
|
||||
// Drop the post-redirect one-shot so re-selecting that server opens Overview.
|
||||
setToolsTabServerId(null);
|
||||
refetch();
|
||||
}, [refetch]);
|
||||
|
||||
|
|
@ -483,6 +521,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
|||
userID={userID}
|
||||
userRole={userRole}
|
||||
availableAccessGroups={uniqueMcpAccessGroups}
|
||||
initialTabIndex={selectedServerId === toolsTabServerId ? 1 : 0}
|
||||
/>
|
||||
) : (
|
||||
<div className="w-full h-full">
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ import React, { useEffect, useMemo, useRef, useState } from "react";
|
|||
import { Card, Title, Text } from "@tremor/react";
|
||||
import { ToolOutlined, CheckCircleOutlined, SearchOutlined, EditOutlined } from "@ant-design/icons";
|
||||
import { Badge, Spin, Checkbox, Input, Radio } from "antd";
|
||||
import { useTestMCPConnection } from "../../hooks/useTestMCPConnection";
|
||||
import McpCrudPermissionPanel from "./McpCrudPermissionPanel";
|
||||
|
||||
interface KeyTool {
|
||||
|
|
@ -12,7 +11,6 @@ interface KeyTool {
|
|||
|
||||
interface MCPToolConfigurationProps {
|
||||
accessToken: string | null;
|
||||
oauthAccessToken?: string | null;
|
||||
formValues: Record<string, any>;
|
||||
allowedTools: string[];
|
||||
existingAllowedTools: string[] | null;
|
||||
|
|
@ -144,7 +142,6 @@ const ToolRow: React.FC<ToolRowProps> = ({
|
|||
|
||||
const MCPToolConfiguration: React.FC<MCPToolConfigurationProps> = ({
|
||||
accessToken,
|
||||
oauthAccessToken,
|
||||
formValues,
|
||||
allowedTools,
|
||||
existingAllowedTools,
|
||||
|
|
@ -169,19 +166,13 @@ const MCPToolConfiguration: React.FC<MCPToolConfigurationProps> = ({
|
|||
const previousSuggestedToolNamesRef = useRef<string>("");
|
||||
const [expandedTools, setExpandedTools] = useState<Set<string>>(new Set());
|
||||
|
||||
// Use external tool state when provided (avoids duplicate fetch with MCPConnectionStatus).
|
||||
// Fall back to internal hook when used standalone (e.g., edit flow).
|
||||
const hasExternalState = externalTools !== undefined;
|
||||
const internalHook = useTestMCPConnection({
|
||||
accessToken,
|
||||
oauthAccessToken,
|
||||
formValues,
|
||||
enabled: !hasExternalState,
|
||||
});
|
||||
const tools: ToolEntry[] = hasExternalState ? externalTools : internalHook.tools;
|
||||
const isLoadingTools = hasExternalState ? externalIsLoading ?? false : internalHook.isLoadingTools;
|
||||
const toolsError = hasExternalState ? externalError ?? null : internalHook.toolsError;
|
||||
const canFetchTools = hasExternalState ? externalCanFetch ?? false : internalHook.canFetchTools;
|
||||
// Tool list is fetched by the parent (create/edit flow) and passed in. This
|
||||
// component renders that state; it never fetches on its own, so there is a
|
||||
// single source of truth and no risk of falling back to a different endpoint.
|
||||
const tools: ToolEntry[] = externalTools ?? [];
|
||||
const isLoadingTools = externalIsLoading ?? false;
|
||||
const toolsError = externalError ?? null;
|
||||
const canFetchTools = externalCanFetch ?? false;
|
||||
|
||||
// Fuzzy-match curated key tool names against actual loaded tool names
|
||||
const suggestedTools = useMemo(() => {
|
||||
|
|
|
|||
|
|
@ -2,12 +2,13 @@ import { render, screen, waitFor } from "@testing-library/react";
|
|||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { describe, expect, it, vi, beforeEach } from "vitest";
|
||||
import MCPToolsViewer from "./mcp_tools";
|
||||
import { listMCPTools } from "../networking";
|
||||
import { listMCPTools, getMCPOAuthUserCredentialStatus } from "../networking";
|
||||
import { isTokenValid, getToken } from "@/utils/mcpTokenStore";
|
||||
|
||||
vi.mock("../networking", () => ({
|
||||
listMCPTools: vi.fn(),
|
||||
callMCPTool: vi.fn(),
|
||||
getMCPOAuthUserCredentialStatus: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/utils/mcpTokenStore", () => ({
|
||||
|
|
@ -20,6 +21,10 @@ vi.mock("@/hooks/useToolsOAuthFlow", () => ({
|
|||
useToolsOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null }),
|
||||
}));
|
||||
|
||||
vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
|
||||
useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null }),
|
||||
}));
|
||||
|
||||
const GATE_TEXT = "Authentication required";
|
||||
// Realistic interactive servers carry a token endpoint; the old heuristic
|
||||
// (`oauth2 && !tokenUrl`) mislabeled exactly these as M2M. Setting it here is
|
||||
|
|
@ -42,6 +47,13 @@ const renderViewer = (props: Record<string, unknown>) =>
|
|||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
const credStatus = (overrides: Record<string, unknown> = {}) => ({
|
||||
server_id: "srv-1",
|
||||
has_credential: true,
|
||||
is_expired: false,
|
||||
...overrides,
|
||||
});
|
||||
|
||||
describe("MCPToolsViewer auth gate routing", () => {
|
||||
beforeEach(() => {
|
||||
vi.mocked(listMCPTools).mockReset().mockResolvedValue({ tools: [], error: null });
|
||||
|
|
@ -49,6 +61,8 @@ describe("MCPToolsViewer auth gate routing", () => {
|
|||
vi.mocked(getToken)
|
||||
.mockReset()
|
||||
.mockReturnValue(undefined as any);
|
||||
// Default: the OBO credential exists and is valid, so OBO servers list tools.
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockReset().mockResolvedValue(credStatus());
|
||||
});
|
||||
|
||||
it("shows the Authorize gate for a passthrough server with a token endpoint and does not list tools", async () => {
|
||||
|
|
@ -57,6 +71,8 @@ describe("MCPToolsViewer auth gate routing", () => {
|
|||
expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument();
|
||||
expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled();
|
||||
// Passthrough must not consult the per-user DB credential.
|
||||
expect(vi.mocked(getMCPOAuthUserCredentialStatus)).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("forwards the session token via the x-mcp header for a passthrough server that has one", async () => {
|
||||
|
|
@ -75,7 +91,40 @@ describe("MCPToolsViewer auth gate routing", () => {
|
|||
expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not gate an OBO server with a token endpoint; lists with the LiteLLM key and no x-mcp header", async () => {
|
||||
it("lists tools for an OBO server when the user has a DB credential, with no x-mcp header", async () => {
|
||||
renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false });
|
||||
|
||||
await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined));
|
||||
expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the Authorize gate for an OBO server when the user has no DB credential and does not list tools", async () => {
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue(credStatus({ has_credential: false }));
|
||||
|
||||
renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false });
|
||||
|
||||
expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument();
|
||||
expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("shows the Authorize gate for an OBO server when the credential-status check fails", async () => {
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockRejectedValue(new Error("network down"));
|
||||
|
||||
renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false });
|
||||
|
||||
expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument();
|
||||
expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not gate an OBO server whose stored token is expired; the list call refreshes it server-side", async () => {
|
||||
// has_credential=true with is_expired=true must NOT gate: resolve_valid_user_oauth_token
|
||||
// refreshes from the stored refresh_token on the list call, so the user never reauthorizes.
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue(
|
||||
credStatus({ has_credential: true, is_expired: true }),
|
||||
);
|
||||
|
||||
renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false });
|
||||
|
||||
await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined));
|
||||
|
|
@ -87,5 +136,7 @@ describe("MCPToolsViewer auth gate routing", () => {
|
|||
|
||||
await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined));
|
||||
expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument();
|
||||
// M2M uses the backend service token, not a per-user DB credential.
|
||||
expect(vi.mocked(getMCPOAuthUserCredentialStatus)).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
import React, { useEffect, useState } from "react";
|
||||
import React, { useCallback, useEffect, useState } from "react";
|
||||
import { useQuery, useMutation } from "@tanstack/react-query";
|
||||
import { ToolTestPanel } from "./ToolTestPanel";
|
||||
import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse, getMcpOAuthMode } from "./types";
|
||||
import { listMCPTools, callMCPTool } from "../networking";
|
||||
import { listMCPTools, callMCPTool, getMCPOAuthUserCredentialStatus } from "../networking";
|
||||
import { isTokenValid, getToken, removeToken } from "@/utils/mcpTokenStore";
|
||||
import { sanitizeMcpAliasForHeader } from "@/utils/mcpHeaderUtils";
|
||||
import { sanitizeMcpAliasForHeader, buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils";
|
||||
import { useToolsOAuthFlow } from "@/hooks/useToolsOAuthFlow";
|
||||
import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
|
||||
import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils";
|
||||
import { setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
import { Card, Title, Text } from "@tremor/react";
|
||||
import { RobotOutlined, ToolOutlined, SearchOutlined, KeyOutlined, LockOutlined } from "@ant-design/icons";
|
||||
|
|
@ -31,11 +34,14 @@ const MCPToolsViewer = ({
|
|||
const [passthroughHeaders, setPassthroughHeaders] = useState<Record<string, string>>({});
|
||||
const [showHeaderInput, setShowHeaderInput] = useState(false);
|
||||
|
||||
// Only PKCE passthrough uses a browser-held session token (sessionStorage,
|
||||
// cleared on tab/browser close) and a user-facing auth gate. OBO uses the
|
||||
// backend-stored per-user token and M2M uses the backend's own service token,
|
||||
// so neither needs a gate — they list tools with just the LiteLLM key.
|
||||
const isPassthrough = getMcpOAuthMode({ auth_type, oauth2_flow, delegate_auth_to_upstream }) === "passthrough";
|
||||
// PKCE passthrough holds a browser-side session token (sessionStorage) and
|
||||
// gates tool listing behind it. OBO uses a backend-stored per-user token that
|
||||
// the user must establish once via an interactive login; we gate on whether
|
||||
// that DB credential exists. M2M uses the backend's own service token and
|
||||
// needs no gate.
|
||||
const oauthMode = getMcpOAuthMode({ auth_type, oauth2_flow, delegate_auth_to_upstream });
|
||||
const isPassthrough = oauthMode === "passthrough";
|
||||
const isObo = oauthMode === "obo";
|
||||
const [oauthToken, setOauthToken] = useState<string | null>(() =>
|
||||
isPassthrough && isTokenValid(serverId, userID) ? getToken(serverId, userID)?.access_token ?? null : null,
|
||||
);
|
||||
|
|
@ -61,6 +67,31 @@ const MCPToolsViewer = ({
|
|||
onSuccess: setOauthToken,
|
||||
});
|
||||
|
||||
// OBO servers list tools using a per-user token the backend stores in the DB;
|
||||
// check whether the current user has a valid one so we can prompt them to
|
||||
// authorize when they don't (otherwise the backend silently returns no tools).
|
||||
const {
|
||||
data: oboCredStatus,
|
||||
isLoading: isLoadingOboCred,
|
||||
isError: isOboCredError,
|
||||
refetch: refetchOboCred,
|
||||
} = useQuery({
|
||||
queryKey: ["mcpOauthUserCredStatus", serverId, userID],
|
||||
queryFn: () => getMCPOAuthUserCredentialStatus(accessToken ?? "", serverId),
|
||||
enabled: !!accessToken && isObo,
|
||||
staleTime: 30000,
|
||||
});
|
||||
|
||||
// A stored credential is sufficient: the backend proactively refreshes an
|
||||
// expired or near-expiry token from the stored refresh_token on the next list
|
||||
// call, so the user only needs to authorize when no credential row exists. If
|
||||
// the status check itself fails we can't confirm a credential, so surface the
|
||||
// Authorize gate rather than a silent empty tool list; re-authorizing only
|
||||
// overwrites the user's own row, so it is safe when a credential did exist.
|
||||
const hasOboCred = !!oboCredStatus?.has_credential;
|
||||
const oboNeedsAuth = isObo && !isLoadingOboCred && (isOboCredError || (!!oboCredStatus && !hasOboCred));
|
||||
const oboStatusLoading = isObo && isLoadingOboCred;
|
||||
|
||||
// Check if this server has extra headers configured
|
||||
const hasExtraHeaders = extraHeaders && extraHeaders.length > 0;
|
||||
|
||||
|
|
@ -75,16 +106,7 @@ const MCPToolsViewer = ({
|
|||
// When no alias is available, fall back to x-mcp-auth (legacy but still supported).
|
||||
// Passthrough only: OBO/M2M tokens are attached server-side, not from the browser.
|
||||
if (isPassthrough && oauthToken) {
|
||||
if (serverAlias) {
|
||||
const safeAlias = sanitizeMcpAliasForHeader(serverAlias);
|
||||
if (safeAlias) {
|
||||
customHeaders[`x-mcp-${safeAlias}-authorization`] = `Bearer ${oauthToken}`;
|
||||
} else {
|
||||
customHeaders["x-mcp-auth"] = `Bearer ${oauthToken}`;
|
||||
}
|
||||
} else {
|
||||
customHeaders["x-mcp-auth"] = `Bearer ${oauthToken}`;
|
||||
}
|
||||
Object.assign(customHeaders, buildMcpPassthroughAuthHeader(serverAlias, oauthToken));
|
||||
}
|
||||
|
||||
// Add passthrough headers with server-specific prefix
|
||||
|
|
@ -135,8 +157,9 @@ const MCPToolsViewer = ({
|
|||
}
|
||||
return result;
|
||||
},
|
||||
// For OAuth servers, block the query until a session token is available
|
||||
enabled: !!accessToken && (!isPassthrough || oauthToken !== null),
|
||||
// Passthrough blocks until a browser session token exists; OBO blocks until
|
||||
// the user has a valid DB credential (else the backend returns no tools).
|
||||
enabled: !!accessToken && (isPassthrough ? oauthToken !== null : isObo ? hasOboCred : true),
|
||||
staleTime: 30000, // Consider data fresh for 30 seconds
|
||||
retry: (failureCount, error: any) => {
|
||||
// Don't retry on 401 — token is invalid, user must re-authenticate
|
||||
|
|
@ -145,6 +168,33 @@ const MCPToolsViewer = ({
|
|||
},
|
||||
});
|
||||
|
||||
// OBO authorize: same redirect+exchange flow as the admin "Authorize & Fetch"
|
||||
// and the chat "Connect" button, but persists the token to the per-user DB.
|
||||
const onOboAuthSuccess = useCallback(() => {
|
||||
refetchOboCred();
|
||||
refetchTools();
|
||||
}, [refetchOboCred, refetchTools]);
|
||||
|
||||
const {
|
||||
startOAuthFlow: startDbOAuthFlow,
|
||||
status: dbOAuthStatus,
|
||||
error: dbOAuthError,
|
||||
} = useUserMcpOAuthFlow({
|
||||
accessToken: accessToken ?? "",
|
||||
serverId,
|
||||
serverAlias,
|
||||
onSuccess: onOboAuthSuccess,
|
||||
});
|
||||
|
||||
// Stash which server started the redirect so the MCP Servers page can reopen
|
||||
// this Tools tab on return and let the flow resume to persist the credential.
|
||||
const startOboAuthorize = useCallback(() => {
|
||||
try {
|
||||
setSecureItem(TOOLS_OAUTH_UI_STATE_KEY, JSON.stringify({ serverId }));
|
||||
} catch (_) {}
|
||||
startDbOAuthFlow();
|
||||
}, [serverId, startDbOAuthFlow]);
|
||||
|
||||
// If the tools query fails with 401, the cached OAuth token is invalid —
|
||||
// clear it so the auth gate is shown again and the user can re-authenticate.
|
||||
useEffect(() => {
|
||||
|
|
@ -187,6 +237,13 @@ const MCPToolsViewer = ({
|
|||
|
||||
const toolsData = mcpToolsResponse?.tools || [];
|
||||
|
||||
// An auth gate replaces the tool list when the user must authenticate first:
|
||||
// passthrough needs a browser token, OBO needs a stored DB credential.
|
||||
const authGateActive = (isPassthrough && !oauthToken) || oboNeedsAuth;
|
||||
// Treat OBO credential-status loading as "tools loading" so the empty state
|
||||
// doesn't flash before we know whether the user needs to authorize.
|
||||
const toolsAreaLoading = isLoadingTools || oboStatusLoading;
|
||||
|
||||
// Filter tools based on search term
|
||||
const filteredTools = toolsData.filter((tool: MCPTool) => {
|
||||
const searchLower = toolSearchTerm.toLowerCase();
|
||||
|
|
@ -287,7 +344,7 @@ const MCPToolsViewer = ({
|
|||
)}
|
||||
</Text>
|
||||
|
||||
{/* OAuth Auth Gate — shown when token is absent for OAuth servers */}
|
||||
{/* Passthrough auth gate — browser session token absent */}
|
||||
{isPassthrough && !oauthToken && (
|
||||
<div className="p-4 text-center bg-white border border-gray-200 rounded-lg">
|
||||
<LockOutlined className="text-2xl text-gray-400 mb-2" />
|
||||
|
|
@ -306,8 +363,31 @@ const MCPToolsViewer = ({
|
|||
</div>
|
||||
)}
|
||||
|
||||
{/* OBO auth gate — only when no credential row exists for this user.
|
||||
An existing-but-expired token is refreshed server-side on the
|
||||
list call, so the gate never appears for a stored credential. */}
|
||||
{oboNeedsAuth && (
|
||||
<div className="p-4 text-center bg-white border border-gray-200 rounded-lg">
|
||||
<LockOutlined className="text-2xl text-gray-400 mb-2" />
|
||||
<p className="text-xs font-medium text-gray-700 mb-1">Authentication required</p>
|
||||
<p className="text-xs text-gray-500 mb-3">
|
||||
Authenticate with the upstream provider to view available tools
|
||||
</p>
|
||||
<AntdButton
|
||||
size="small"
|
||||
type="primary"
|
||||
loading={dbOAuthStatus === "authorizing" || dbOAuthStatus === "exchanging"}
|
||||
onClick={startOboAuthorize}
|
||||
disabled={!accessToken}
|
||||
>
|
||||
Authorize
|
||||
</AntdButton>
|
||||
{dbOAuthError && <p className="text-xs text-red-500 mt-2">{dbOAuthError}</p>}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Search Bar — only shown when tools are loaded */}
|
||||
{!isPassthrough || oauthToken ? (
|
||||
{!authGateActive ? (
|
||||
<>
|
||||
{toolsData.length > 0 && (
|
||||
<div className="mb-3">
|
||||
|
|
@ -324,7 +404,7 @@ const MCPToolsViewer = ({
|
|||
)}
|
||||
|
||||
{/* Loading State */}
|
||||
{isLoadingTools && (
|
||||
{toolsAreaLoading && (
|
||||
<div className="flex flex-col items-center justify-center py-8 bg-white border border-gray-200 rounded-lg">
|
||||
<div className="relative mb-3">
|
||||
<div className="animate-spin rounded-full h-6 w-6 border-2 border-gray-200"></div>
|
||||
|
|
@ -335,7 +415,7 @@ const MCPToolsViewer = ({
|
|||
)}
|
||||
|
||||
{/* Error State */}
|
||||
{(mcpToolsResponse?.error || mcpToolsError) && !isLoadingTools && !toolsData.length && (
|
||||
{(mcpToolsResponse?.error || mcpToolsError) && !toolsAreaLoading && !toolsData.length && (
|
||||
<div className="p-3 text-xs text-red-800 rounded-lg bg-red-50 border border-red-200">
|
||||
<p className="font-medium">
|
||||
Error: {mcpToolsResponse?.message || (mcpToolsError as Error)?.message}
|
||||
|
|
@ -344,7 +424,7 @@ const MCPToolsViewer = ({
|
|||
)}
|
||||
|
||||
{/* No Tools State */}
|
||||
{!isLoadingTools &&
|
||||
{!toolsAreaLoading &&
|
||||
!mcpToolsResponse?.error &&
|
||||
!mcpToolsError &&
|
||||
(!toolsData || toolsData.length === 0) && (
|
||||
|
|
@ -370,7 +450,7 @@ const MCPToolsViewer = ({
|
|||
)}
|
||||
|
||||
{/* Tools List */}
|
||||
{!isLoadingTools && !mcpToolsResponse?.error && toolsData.length > 0 && (
|
||||
{!toolsAreaLoading && !mcpToolsResponse?.error && toolsData.length > 0 && (
|
||||
<>
|
||||
{filteredTools.length === 0 ? (
|
||||
<div className="p-4 text-center bg-white border border-gray-200 rounded-lg">
|
||||
|
|
|
|||
|
|
@ -5225,11 +5225,16 @@ export const testSearchToolConnection = async (accessToken: string, litellmParam
|
|||
}
|
||||
};
|
||||
|
||||
export const listMCPTools = async (accessToken: string, serverId: string, customHeaders?: Record<string, string>) => {
|
||||
// Construct base URL
|
||||
let url = proxyBaseUrl
|
||||
? `${proxyBaseUrl}/mcp-rest/tools/list?server_id=${serverId}`
|
||||
: `/mcp-rest/tools/list?server_id=${serverId}`;
|
||||
export const listMCPTools = async (
|
||||
accessToken: string,
|
||||
serverId: string,
|
||||
customHeaders?: Record<string, string>,
|
||||
includeDisabledTools?: boolean,
|
||||
) => {
|
||||
// Construct base URL. include_disabled_tools returns the full server catalog
|
||||
// (admin-only, backend-enforced) so the settings UI can configure the allowlist.
|
||||
const query = `server_id=${serverId}${includeDisabledTools ? "&include_disabled_tools=true" : ""}`;
|
||||
let url = proxyBaseUrl ? `${proxyBaseUrl}/mcp-rest/tools/list?${query}` : `/mcp-rest/tools/list?${query}`;
|
||||
|
||||
console.log("Fetching MCP tools from:", url);
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,15 @@
|
|||
|
||||
import { getProxyBaseUrl, serverRootPath } from "@/components/networking";
|
||||
|
||||
/**
|
||||
* sessionStorage key used to restore the MCP server detail view on the Tools
|
||||
* tab after a full-page OAuth redirect. The OBO authorize flow redirects to the
|
||||
* IdP and back to the MCP Servers page; without this the user lands on the
|
||||
* server list and useUserMcpOAuthFlow never re-mounts to persist the credential.
|
||||
* Mirrors the admin edit flow's EDIT_OAUTH_UI_STATE_KEY.
|
||||
*/
|
||||
export const TOOLS_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-tools-state";
|
||||
|
||||
/**
|
||||
* Build the OAuth callback URL for the current UI deployment.
|
||||
*
|
||||
|
|
|
|||
117
ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx
Normal file
117
ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.test.tsx
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
import { act, renderHook, waitFor } from "@testing-library/react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import * as networking from "@/components/networking";
|
||||
import { setSecureItem } from "@/utils/secureStorage";
|
||||
import { useMcpOAuthFlow } from "./useMcpOAuthFlow";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
exchangeMcpOAuthToken: vi.fn(),
|
||||
cacheTemporaryMcpServer: vi.fn(),
|
||||
registerMcpOAuthClient: vi.fn(),
|
||||
buildMcpOAuthAuthorizeUrl: vi.fn(),
|
||||
getProxyBaseUrl: vi.fn(() => ""),
|
||||
serverRootPath: "",
|
||||
}));
|
||||
|
||||
vi.mock("@/components/molecules/notifications_manager", () => ({
|
||||
default: { success: vi.fn(), error: vi.fn() },
|
||||
}));
|
||||
|
||||
const FLOW_STATE_KEY = "litellm-mcp-oauth-flow-state";
|
||||
const RESULT_KEY = "litellm-mcp-oauth-result";
|
||||
|
||||
/** Seed the redirect result (the code returned by the IdP callback). */
|
||||
function seedResult(code: string) {
|
||||
setSecureItem(RESULT_KEY, JSON.stringify({ state: "state-1", code }));
|
||||
}
|
||||
|
||||
/** Seed the flow state stored before the redirect. */
|
||||
function seedFlowState() {
|
||||
setSecureItem(
|
||||
FLOW_STATE_KEY,
|
||||
JSON.stringify({
|
||||
state: "state-1",
|
||||
codeVerifier: "verifier-1",
|
||||
serverId: "server-1",
|
||||
clientId: "client-1",
|
||||
redirectUri: "https://app.example.com/ui/mcp/oauth/callback",
|
||||
flowSource: "create",
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/** Seed storage so the hook's on-mount resume flow exchanges a code for a token. */
|
||||
function seedCompletedRedirect() {
|
||||
seedResult("code-1");
|
||||
seedFlowState();
|
||||
}
|
||||
|
||||
function renderFlow(onTokenReceived = vi.fn()) {
|
||||
return renderHook(
|
||||
({ onTokenReceived: cb }: { onTokenReceived: (t: any) => void }) =>
|
||||
useMcpOAuthFlow({
|
||||
accessToken: "admin-token",
|
||||
getCredentials: () => ({}),
|
||||
getTemporaryPayload: () => ({ url: "https://server-1.example.com/mcp", transport: "http" }),
|
||||
onTokenReceived: cb,
|
||||
flowSource: "create",
|
||||
}),
|
||||
{ initialProps: { onTokenReceived } },
|
||||
);
|
||||
}
|
||||
|
||||
describe("useMcpOAuthFlow reset", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
window.sessionStorage.clear();
|
||||
window.localStorage.clear();
|
||||
});
|
||||
|
||||
it("clears a successfully fetched token so it cannot leak into the next session", async () => {
|
||||
const token = { access_token: "tok-123", expires_in: 3600 };
|
||||
vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue(token);
|
||||
seedCompletedRedirect();
|
||||
|
||||
const onTokenReceived = vi.fn();
|
||||
const { result } = renderFlow(onTokenReceived);
|
||||
|
||||
await waitFor(() => expect(result.current.status).toBe("success"));
|
||||
expect(result.current.tokenResponse).toEqual(token);
|
||||
expect(onTokenReceived).toHaveBeenCalledWith(token);
|
||||
|
||||
act(() => {
|
||||
result.current.reset();
|
||||
});
|
||||
|
||||
expect(result.current.status).toBe("idle");
|
||||
expect(result.current.tokenResponse).toBeNull();
|
||||
expect(result.current.error).toBeNull();
|
||||
});
|
||||
|
||||
it("clears the in-flight guard so a callback after a mid-exchange close is not swallowed", async () => {
|
||||
// First exchange hangs, mimicking the modal being closed while the token
|
||||
// endpoint is still in flight. processingRef is left true at that point.
|
||||
vi.mocked(networking.exchangeMcpOAuthToken).mockReturnValueOnce(new Promise<any>(() => {}));
|
||||
seedFlowState();
|
||||
seedResult("code-1");
|
||||
|
||||
const onTokenReceived1 = vi.fn();
|
||||
const { result, rerender } = renderFlow(onTokenReceived1);
|
||||
|
||||
await waitFor(() => expect(result.current.status).toBe("exchanging"));
|
||||
|
||||
act(() => {
|
||||
result.current.reset();
|
||||
});
|
||||
|
||||
// The reopened modal receives a fresh callback; it must be processed, not
|
||||
// dropped by a stale in-flight guard.
|
||||
const token = { access_token: "tok-2" };
|
||||
vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValueOnce(token);
|
||||
seedResult("code-2");
|
||||
const onTokenReceived2 = vi.fn();
|
||||
rerender({ onTokenReceived: onTokenReceived2 });
|
||||
|
||||
await waitFor(() => expect(onTokenReceived2).toHaveBeenCalledWith(token));
|
||||
});
|
||||
});
|
||||
|
|
@ -28,6 +28,11 @@ interface UseMcpOAuthFlowOptions {
|
|||
getTemporaryPayload: () => Record<string, any> | null;
|
||||
onTokenReceived: (tokenResponse: Record<string, any>) => void;
|
||||
onBeforeRedirect?: () => void;
|
||||
// Distinguishes which form started the flow (e.g. "create" vs "edit"). Both forms
|
||||
// mount this hook with shared storage keys, so the return handler only processes a
|
||||
// callback whose stored flowSource matches, preventing one form from grabbing the
|
||||
// other's OAuth result.
|
||||
flowSource: string;
|
||||
}
|
||||
|
||||
interface UseMcpOAuthFlowResult {
|
||||
|
|
@ -35,6 +40,7 @@ interface UseMcpOAuthFlowResult {
|
|||
status: McpOAuthStatus;
|
||||
error: string | null;
|
||||
tokenResponse: Record<string, any> | null;
|
||||
reset: () => void;
|
||||
}
|
||||
|
||||
export const useMcpOAuthFlow = ({
|
||||
|
|
@ -43,6 +49,7 @@ export const useMcpOAuthFlow = ({
|
|||
getTemporaryPayload,
|
||||
onTokenReceived,
|
||||
onBeforeRedirect,
|
||||
flowSource,
|
||||
}: UseMcpOAuthFlowOptions): UseMcpOAuthFlowResult => {
|
||||
const [status, setStatus] = useState<McpOAuthStatus>("idle");
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
|
@ -60,6 +67,7 @@ export const useMcpOAuthFlow = ({
|
|||
clientSecret?: string;
|
||||
serverId: string;
|
||||
redirectUri: string;
|
||||
flowSource?: string;
|
||||
};
|
||||
|
||||
const setStorageItem = (key: string, value: string) => {
|
||||
|
|
@ -179,6 +187,7 @@ export const useMcpOAuthFlow = ({
|
|||
clientSecret: registeredClient.clientSecret || credentials.client_secret,
|
||||
serverId,
|
||||
redirectUri: callbackUrl(),
|
||||
flowSource,
|
||||
};
|
||||
|
||||
if (typeof window === "undefined") {
|
||||
|
|
@ -257,6 +266,16 @@ export const useMcpOAuthFlow = ({
|
|||
return;
|
||||
}
|
||||
|
||||
// Only the form that started this redirect should consume the result. The create
|
||||
// form and the edit form both mount this hook with shared storage keys, so without
|
||||
// this another instance (e.g. the always-mounted create form) would grab and handle
|
||||
// an edit-page authorization. Bail out without clearing RESULT_KEY so the matching
|
||||
// instance can still process it.
|
||||
if (flowState?.flowSource !== flowSource) {
|
||||
processingRef.current = false;
|
||||
return;
|
||||
}
|
||||
|
||||
// Clear the result key after reading it
|
||||
if (typeof window !== "undefined") {
|
||||
try {
|
||||
|
|
@ -318,10 +337,18 @@ export const useMcpOAuthFlow = ({
|
|||
resumeOAuthFlow();
|
||||
}, [resumeOAuthFlow]);
|
||||
|
||||
const reset = useCallback(() => {
|
||||
setStatus("idle");
|
||||
setError(null);
|
||||
setTokenResponse(null);
|
||||
processingRef.current = false;
|
||||
}, []);
|
||||
|
||||
return {
|
||||
startOAuthFlow,
|
||||
status,
|
||||
error,
|
||||
tokenResponse,
|
||||
reset,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
import { useState, useEffect } from "react";
|
||||
import { useState, useEffect, useCallback } from "react";
|
||||
import { testMCPToolsListRequest } from "../components/networking";
|
||||
import { AUTH_TYPE, OAUTH_FLOW, TRANSPORT } from "@/components/mcp_tools/types";
|
||||
|
||||
|
|
@ -177,12 +177,12 @@ export const useTestMCPConnection = ({
|
|||
}
|
||||
};
|
||||
|
||||
const clearTools = () => {
|
||||
const clearTools = useCallback(() => {
|
||||
setTools([]);
|
||||
setToolsError(null);
|
||||
setToolsErrorStackTrace(null);
|
||||
setHasShownSuccessMessage(false);
|
||||
};
|
||||
}, []);
|
||||
|
||||
// Auto-fetch tools when form values change and required fields are available
|
||||
useEffect(() => {
|
||||
|
|
|
|||
18
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
18
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -5922,6 +5922,11 @@ export interface paths {
|
|||
*
|
||||
* - When litellm_model_id is passed, it will return the info for that specific model
|
||||
* - When litellm_model_id is not passed, it will return the info for all models
|
||||
* - include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info).
|
||||
* - teamId: Filter to models accessible by the given team.
|
||||
*
|
||||
* Each model in the list response includes `model_info.access_via_team_ids` and
|
||||
* `model_info.direct_access` when the proxy database is connected.
|
||||
*
|
||||
* Returns:
|
||||
* Returns a dictionary containing information about each model.
|
||||
|
|
@ -14383,6 +14388,11 @@ export interface paths {
|
|||
*
|
||||
* - When litellm_model_id is passed, it will return the info for that specific model
|
||||
* - When litellm_model_id is not passed, it will return the info for all models
|
||||
* - include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info).
|
||||
* - teamId: Filter to models accessible by the given team.
|
||||
*
|
||||
* Each model in the list response includes `model_info.access_via_team_ids` and
|
||||
* `model_info.direct_access` when the proxy database is connected.
|
||||
*
|
||||
* Returns:
|
||||
* Returns a dictionary containing information about each model.
|
||||
|
|
@ -37867,6 +37877,10 @@ export interface operations {
|
|||
parameters: {
|
||||
query?: {
|
||||
litellm_model_id?: string | null;
|
||||
/** @description When true, filter to deployments the caller can use via direct access or team membership. */
|
||||
include_team_models?: boolean | null;
|
||||
/** @description Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids */
|
||||
teamId?: string | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
@ -48127,6 +48141,10 @@ export interface operations {
|
|||
parameters: {
|
||||
query?: {
|
||||
litellm_model_id?: string | null;
|
||||
/** @description When true, filter to deployments the caller can use via direct access or team membership. */
|
||||
include_team_models?: boolean | null;
|
||||
/** @description Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids */
|
||||
teamId?: string | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
|
|||
|
|
@ -12,3 +12,18 @@ export function sanitizeMcpAliasForHeader(alias: string): string {
|
|||
.replace(/_+/g, "_")
|
||||
.replace(/^_|_$/g, "");
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the passthrough auth header that forwards a browser-held PKCE token to
|
||||
* the upstream MCP server. The backend picks up the x-mcp-{alias}-authorization
|
||||
* pattern and forwards it; without a usable alias it falls back to the legacy
|
||||
* x-mcp-auth header.
|
||||
*/
|
||||
export function buildMcpPassthroughAuthHeader(
|
||||
serverAlias: string | null | undefined,
|
||||
token: string,
|
||||
): Record<string, string> {
|
||||
const safeAlias = serverAlias ? sanitizeMcpAliasForHeader(serverAlias) : "";
|
||||
const headerName = safeAlias ? `x-mcp-${safeAlias}-authorization` : "x-mcp-auth";
|
||||
return { [headerName]: `Bearer ${token}` };
|
||||
}
|
||||
|
|
|
|||
38
uv.lock
generated
38
uv.lock
generated
|
|
@ -9,7 +9,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-06-03T21:40:52.018333Z"
|
||||
exclude-newer = "2026-06-13T01:42:46.429412Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -18,6 +18,10 @@ members = [
|
|||
"litellm-enterprise",
|
||||
"litellm-proxy-extras",
|
||||
]
|
||||
constraints = [
|
||||
{ name = "aiohttp", specifier = ">=3.13.5,<3.14" },
|
||||
{ name = "tornado", specifier = ">=6.5.6" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "a2a-sdk"
|
||||
|
|
@ -3276,7 +3280,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.89.0"
|
||||
version = "1.89.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
|
@ -3524,7 +3528,7 @@ requires-dist = [
|
|||
{ name = "pydantic-settings", marker = "extra == 'proxy'", specifier = ">=2.14.1,<3.0" },
|
||||
{ name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.12.0,<3.0" },
|
||||
{ name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" },
|
||||
{ name = "pypdf", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=6.10.2,<7.0" },
|
||||
{ name = "pypdf", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=6.12.0,<7.0" },
|
||||
{ name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" },
|
||||
{ name = "python-dotenv", specifier = ">=1.0.0,<2.0" },
|
||||
{ name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" },
|
||||
|
|
@ -6048,14 +6052,14 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "pypdf"
|
||||
version = "6.10.2"
|
||||
version = "6.13.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.11'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7b/3f/9f2167401c2e94833ca3b69535bad89e533b5de75fefe4197a2c224baec2/pypdf-6.10.2.tar.gz", hash = "sha256:7d09ce108eff6bf67465d461b6ef352dcb8d84f7a91befc02f904455c6eea11d", size = 5315679, upload-time = "2026-04-15T16:37:36.978Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/99/0a/48fe05c6bb3aa4bb4d2a4079a383d33c0dfec1edf613a642f07d8b8b5c2e/pypdf-6.13.2.tar.gz", hash = "sha256:5a96a17dbdfbf9c2ab24c0a13fa0aba182be22ba6f283098712c16fc242f509f", size = 6479250, upload-time = "2026-06-10T16:42:34.5Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/d6/1d5c60cc17bbdf37c1552d9c03862fc6d32c5836732a0415b2d637edc2d0/pypdf-6.10.2-py3-none-any.whl", hash = "sha256:aa53be9826655b51c96741e5d7983ca224d898ac0a77896e64636810517624aa", size = 336308, upload-time = "2026-04-15T16:37:34.851Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/17/378943705992f74e451a06de3401ce68e3213763c81e44d0614559c45599/pypdf-6.13.2-py3-none-any.whl", hash = "sha256:6eeb9e57693f29d41bd01255d02660cbbb41fd7fc818a982677389a35e4f2083", size = 346555, upload-time = "2026-06-10T16:42:32.37Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -7571,19 +7575,19 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "tornado"
|
||||
version = "6.5.5"
|
||||
version = "6.5.7"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f8/f1/3173dfa4a18db4a9b03e5d55325559dab51ee653763bb8745a75af491286/tornado-6.5.5.tar.gz", hash = "sha256:192b8f3ea91bd7f1f50c06955416ed76c6b72f96779b962f07f911b91e8d30e9", size = 516006, upload-time = "2026-03-10T21:31:02.067Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/59/8c/77f5097695f4dd8255ecbd08b2a1ed8ba8b953d337804dd7080f199e12bf/tornado-6.5.5-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:487dc9cc380e29f58c7ab88f9e27cdeef04b2140862e5076a66fb6bb68bb1bfa", size = 445983, upload-time = "2026-03-10T21:30:44.28Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ab/5e/7625b76cd10f98f1516c36ce0346de62061156352353ef2da44e5c21523c/tornado-6.5.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:65a7f1d46d4bb41df1ac99f5fcb685fb25c7e61613742d5108b010975a9a6521", size = 444246, upload-time = "2026-03-10T21:30:46.571Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b2/04/7b5705d5b3c0fab088f434f9c83edac1573830ca49ccf29fb83bf7178eec/tornado-6.5.5-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:e74c92e8e65086b338fd56333fb9a68b9f6f2fe7ad532645a290a464bcf46be5", size = 447229, upload-time = "2026-03-10T21:30:48.273Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/34/01/74e034a30ef59afb4097ef8659515e96a39d910b712a89af76f5e4e1f93c/tornado-6.5.5-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:435319e9e340276428bbdb4e7fa732c2d399386d1de5686cb331ec8eee754f07", size = 448192, upload-time = "2026-03-10T21:30:51.22Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/be/00/fe9e02c5a96429fce1a1d15a517f5d8444f9c412e0bb9eadfbe3b0fc55bf/tornado-6.5.5-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:3f54aa540bdbfee7b9eb268ead60e7d199de5021facd276819c193c0fb28ea4e", size = 448039, upload-time = "2026-03-10T21:30:53.52Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/9e/656ee4cec0398b1d18d0f1eb6372c41c6b889722641d84948351ae19556d/tornado-6.5.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:36abed1754faeb80fbd6e64db2758091e1320f6bba74a4cf8c09cd18ccce8aca", size = 447445, upload-time = "2026-03-10T21:30:55.541Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/76/4921c00511f88af86a33de770d64141170f1cfd9c00311aea689949e274e/tornado-6.5.5-cp39-abi3-win32.whl", hash = "sha256:dd3eafaaeec1c7f2f8fdcd5f964e8907ad788fe8a5a32c4426fbbdda621223b7", size = 448582, upload-time = "2026-03-10T21:30:57.142Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2c/23/f6c6112a04d28eed765e374435fb1a9198f73e1ec4b4024184f21faeb1ad/tornado-6.5.5-cp39-abi3-win_amd64.whl", hash = "sha256:6443a794ba961a9f619b1ae926a2e900ac20c34483eea67be4ed8f1e58d3ef7b", size = 448990, upload-time = "2026-03-10T21:30:58.857Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/c8/876602cbc96469911f0939f703453c1157b0c826ecb05bdd32e023397d4e/tornado-6.5.5-cp39-abi3-win_arm64.whl", hash = "sha256:2c9a876e094109333f888539ddb2de4361743e5d21eece20688e3e351e4990a6", size = 448016, upload-time = "2026-03-10T21:31:00.43Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue