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:
yuneng-jiang 2026-06-15 20:08:03 -07:00 committed by GitHub
commit 39832310a7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
75 changed files with 5684 additions and 598 deletions

View file

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

View 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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

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

View file

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

View file

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

View file

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

View file

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

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