scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge pull request #42292 from BerriAI/litellm_extract_mcp_operation_context_7741

Browse source 
refactor(mcp): extract explicit operation context and dispatch
This commit is contained in:
joshua-berri 2026-09-21 22:38:07 +00:00 • committed by GitHub
commit 3353c57e76
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
29 changed files with 4686 additions and 3240 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.github/workflows/test-linting.yml vendored
View file

@ -130,6 +130,10 @@ jobs:
echo "File content around line 43:"
head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10
- name: Check MCP operation boundary
if: steps.changes.outputs.decision != 'skip'
run: uv run --no-sync python scripts/check_mcp_operation_boundary.py
- name: Run Ruff linting
if: steps.changes.outputs.decision != 'skip'
run: |

1
Makefile
View file

@ -164,6 +164,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
# Linting targets
lint-ruff: $(LINT_DEP_INSTALL)
$(UV_RUN) python scripts/check_mcp_operation_boundary.py
cd litellm && $(UV_RUN) ruff check . && cd ..
$(UV_RUN) ruff check --config ruff-tests.toml tests

95
litellm/proxy/_experimental/mcp_server/contracts.py Normal file
View file

@ -0,0 +1,95 @@
from collections.abc import Mapping
from copy import deepcopy
from dataclasses import dataclass, field
from datetime import datetime
from types import MappingProxyType
from typing import Final, Protocol
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def copy_caller(auth: UserAPIKeyAuth | None) -> UserAPIKeyAuth | None:
if auth is None:
return None
span: Final = auth.parent_otel_span
return deepcopy(auth, {id(span): span} if span is not None else None) # mutable-ok: deepcopy mutates its memo
@dataclass(frozen=True, slots=True)
class OperationContext:
_caller: UserAPIKeyAuth | None = field(repr=False)
mcp_auth_header: str | None = field(default=None, repr=False)
mcp_servers: tuple[str, ...] | None = None
mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None = field(default=None, repr=False)
oauth2_headers: Mapping[str, str] | None = field(default=None, repr=False)
raw_headers: Mapping[str, str] | None = field(default=None, repr=False)
client_ip: str | None = None
mcp_proxy_mode: bool = False
def __post_init__(self) -> None:
object.__setattr__(self, "_caller", copy_caller(self._caller))
object.__setattr__(self, "mcp_servers", tuple(self.mcp_servers) if self.mcp_servers is not None else None)
object.__setattr__(
self,
"oauth2_headers",
MappingProxyType(dict(self.oauth2_headers)) if self.oauth2_headers is not None else None,
)
object.__setattr__(
self, "raw_headers", MappingProxyType(dict(self.raw_headers)) if self.raw_headers is not None else None
)
object.__setattr__(
self,
"mcp_server_auth_headers",
MappingProxyType(
{key: MappingProxyType(dict(value)) for key, value in self.mcp_server_auth_headers.items()}
)
if self.mcp_server_auth_headers is not None
else None,
)
@property
def user_api_key_auth(self) -> UserAPIKeyAuth | None:
return copy_caller(self._caller)
def legacy_auth(
self,
) -> tuple[
UserAPIKeyAuth | None,
str | None,
list[str] | None, # mutable-ok: detached legacy server-list payload
dict[str, dict[str, str]] | None, # mutable-ok: legacy auth dispatch requires concrete dict headers
dict[str, str] | None, # mutable-ok: detached legacy header payload
dict[str, str] | None, # mutable-ok: detached legacy header payload
str | None,
]:
return (
self.user_api_key_auth,
self.mcp_auth_header,
list(self.mcp_servers) if self.mcp_servers is not None else None, # mutable-ok: legacy policy list input
{
key: dict(value) for key, value in self.mcp_server_auth_headers.items()
} # mutable-ok: legacy auth dispatch checks concrete dict headers
if self.mcp_server_auth_headers is not None
else None,
dict(self.oauth2_headers)
if self.oauth2_headers is not None
else None, # mutable-ok: legacy OAuth header input
dict(self.raw_headers) if self.raw_headers is not None else None, # mutable-ok: legacy request header input
self.client_ip,
)
class ProgressCallback(Protocol):
async def __call__(self, progress: float, total: float | None, /) -> None: ...
@dataclass(frozen=True, slots=True)
class AuthorizedToolCall:
name: str
arguments: Mapping[str, object]
allowed_mcp_servers: tuple[MCPServer, ...]
start_time: datetime
host_progress_callback: ProgressCallback | None
guardrail_context: Mapping[str, object] | None
logging_data: Mapping[str, object]

83
litellm/proxy/_experimental/mcp_server/legacy_callbacks.py Normal file
View file

@ -0,0 +1,83 @@
from collections.abc import Mapping
from typing import Final, Protocol
from mcp.client.session import ClientRequestContext
from mcp.types import (
CreateMessageRequestParams,
CreateMessageResult,
CreateMessageResultWithTools,
ElicitRequestParams,
ElicitResult,
ErrorData,
)
from litellm.proxy._experimental.mcp_server.contracts import OperationContext
from litellm.proxy._types import UserAPIKeyAuth
class SamplingCallback(Protocol):
async def __call__(
self, context: ClientRequestContext, params: CreateMessageRequestParams, /
) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData: ...
class ElicitationCallback(Protocol):
async def __call__(self, context: object, params: ElicitRequestParams, /) -> ElicitResult | ErrorData: ...
def create_sampling_callback(
user_api_key_auth: UserAPIKeyAuth | None = None,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
operation_context: OperationContext | None = None,
) -> SamplingCallback:
from litellm.proxy._experimental.mcp_server.server import get_active_auth_context
auth: Final = get_active_auth_context() if operation_context is None and user_api_key_auth is None else None
captured: Final = (
operation_context
if operation_context is not None
else OperationContext(
_caller=user_api_key_auth if user_api_key_auth is not None else (auth.user_api_key_auth if auth else None),
raw_headers=raw_headers if raw_headers is not None else (auth.raw_headers if auth else None),
client_ip=client_ip if client_ip is not None else (auth.client_ip if auth else None),
)
)
async def callback(
context: ClientRequestContext, params: CreateMessageRequestParams
) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData:
import litellm
from litellm.proxy._experimental.mcp_server.sampling_handler import handle_sampling_create_message
return await handle_sampling_create_message(
context=context,
params=params,
default_model=getattr(litellm, "default_mcp_sampling_model", None),
user_api_key_auth=captured.user_api_key_auth,
raw_headers=dict(captured.raw_headers)
if captured.raw_headers is not None
else None, # mutable-ok: handler consumes an owned request header dict
client_ip=captured.client_ip,
)
return callback
def create_elicitation_callback() -> ElicitationCallback:
from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session
downstream_session: Final = get_active_mcp_session()
downstream_capabilities: Final = getattr(downstream_session, "capabilities", None)
async def callback(context: object, params: ElicitRequestParams) -> ElicitResult | ErrorData:
from litellm.proxy._experimental.mcp_server.elicitation_handler import handle_elicitation_request
return await handle_elicitation_request(
context=context,
params=params,
downstream_session=downstream_session,
downstream_capabilities=downstream_capabilities,
)
return callback

129
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -73,6 +73,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPServerAccess,
_is_mcp_admitted_user_subject,
)
from litellm.proxy._experimental.mcp_server.contracts import OperationContext
from litellm.proxy._experimental.mcp_server.elicitation_handler import (
MCP_ELICITATION_AVAILABLE,
)
@ -195,9 +196,6 @@ from litellm.types.mcp_server.mcp_server_manager import (
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from mcp.client.session import ClientRequestContext
from mcp.types import CreateMessageRequestParams
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.mcp_server.mcp_toolset import MCPToolset
@ -1218,7 +1216,7 @@ async def _resolve_byok_mcp_auth_header(
if not mcp_server.is_byok:
return mcp_auth_header
from litellm.proxy._experimental.mcp_server.server import (
from litellm.proxy._experimental.mcp_server.operations import (
_check_byok_credential,
_get_byok_credential,
)
@ -1577,77 +1575,25 @@ def _normalize_mcp_server_cost_info(mcp_info: MCPInfo) -> None:
mcp_info["mcp_server_cost_info"] = normalized
def _create_sampling_callback(user_api_key_auth: UserAPIKeyAuth | None = None):
"""
Create a sampling callback for MCP ClientSession.
Returns a callable that handles sampling/createMessage requests from
upstream MCP servers by routing them through litellm.acompletion().
"""
def _create_sampling_callback(
user_api_key_auth: UserAPIKeyAuth | None = None,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
operation_context: OperationContext | None = None,
):
if not MCP_SAMPLING_AVAILABLE:
return None
from litellm.proxy._experimental.mcp_server.legacy_callbacks import create_sampling_callback
async def _sampling_callback(
context: "ClientRequestContext",
params: "CreateMessageRequestParams",
):
import litellm
from litellm.proxy._experimental.mcp_server.sampling_handler import (
handle_sampling_create_message,
)
from litellm.proxy._experimental.mcp_server.server import (
get_active_auth_context,
)
auth_context: Final = get_active_auth_context()
resolved_auth: Final = user_api_key_auth or (auth_context.user_api_key_auth if auth_context else None)
# Forward original HTTP headers and client IP so that
# header-dependent guardrails, tag-based routing, trace
# correlation, and forward_llm_provider_auth_headers work
# correctly for sampling sub-calls.
_raw_headers: Final = getattr(auth_context, "raw_headers", None)
_client_ip: Final = getattr(auth_context, "client_ip", None)
return await handle_sampling_create_message(
context=context,
params=params,
default_model=getattr(litellm, "default_mcp_sampling_model", None),
user_api_key_auth=resolved_auth,
raw_headers=_raw_headers,
client_ip=_client_ip,
)
return _sampling_callback
return create_sampling_callback(user_api_key_auth, raw_headers, client_ip, operation_context)
def _create_elicitation_callback():
"""
Create an elicitation callback for MCP ClientSession.
Returns a callable that handles elicitation/create requests from
upstream MCP servers. In gateway mode, this relays to the downstream
client; in tool bridge mode, it returns a decline response.
"""
if not MCP_ELICITATION_AVAILABLE:
return None
from litellm.proxy._experimental.mcp_server.legacy_callbacks import create_elicitation_callback
async def _elicitation_callback(context, params):
from litellm.proxy._experimental.mcp_server.elicitation_handler import (
handle_elicitation_request,
)
from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session
# In Gateway mode, we relay the elicitation request to the downstream client
# that triggered the current operation.
downstream_session: Final = get_active_mcp_session()
downstream_capabilities = getattr(downstream_session, "capabilities", None) if downstream_session else None
return await handle_elicitation_request(
context=context,
params=params,
downstream_session=downstream_session,
downstream_capabilities=downstream_capabilities,
)
return _elicitation_callback
return create_elicitation_callback()
def _record_mcp_guardrail_evaluations(
@ -3386,17 +3332,13 @@ class MCPServerManager:
listable but uninvokable.
Empty inside a toolset scope: toolset_mcp_route / dynamic_mcp_route set
``_mcp_active_toolset_id`` before calling the handler, pinning the request to the toolset's
the caller's server-only ``mcp_toolset_id`` before calling the handler, pinning the request to the toolset's
own servers (checking op.mcp_toolsets==[] instead would false-positive on DB-default rows
where Postgres initialises the column to ARRAY[]::TEXT[]).
``allow_all_server_ids`` / ``submitted_server_ids`` are injectable so the server union,
which precomputes both for its fallback path, does not compute them twice."""
from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415
_mcp_active_toolset_id,
)
if _mcp_active_toolset_id.get() is not None:
if user_api_key_auth is not None and user_api_key_auth.mcp_toolset_id is not None:
return set()
if allow_all_server_ids is None:
allow_all_server_ids = self.get_allow_all_keys_server_ids()
@ -4164,6 +4106,8 @@ class MCPServerManager:
subject_token: str | None = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
cred_provider: UpstreamCredentialProvider | None = None,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
@ -4212,7 +4156,13 @@ class MCPServerManager:
# Create sampling and elicitation callbacks for this client
sampling_cb = (
_create_sampling_callback(user_api_key_auth=user_api_key_auth) if resolved_server.allow_sampling else None
_create_sampling_callback(
operation_context=OperationContext(
_caller=user_api_key_auth, raw_headers=raw_headers, client_ip=client_ip
)
)
if resolved_server.allow_sampling
else None
)
elicitation_cb: Final = _create_elicitation_callback() if resolved_server.allow_elicitation else None
@ -4357,6 +4307,7 @@ class MCPServerManager:
raw_headers: dict[str, str] | None = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
oauth2_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[MCPTool]:
"""
Helper method to get tools from a single MCP server with prefixed names.
@ -4446,6 +4397,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
## HANDLE OPENAPI TOOLS
@ -4556,6 +4509,7 @@ class MCPServerManager:
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[Prompt]:
try:
headers: Final = (
@ -4576,6 +4530,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
credential_fingerprint: Final = await client.discovery_auth_fingerprint()
key: Final = self._discovery_key(
@ -4599,6 +4555,7 @@ class MCPServerManager:
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[Resource]:
try:
headers: Final = (
@ -4619,6 +4576,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
credential_fingerprint: Final = await client.discovery_auth_fingerprint()
key: Final = self._discovery_key(
@ -4642,6 +4601,7 @@ class MCPServerManager:
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[ResourceTemplate]:
try:
headers: Final = (
@ -4662,6 +4622,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
credential_fingerprint: Final = await client.discovery_auth_fingerprint()
key: Final = self._discovery_key(
@ -4685,6 +4647,7 @@ class MCPServerManager:
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> ReadResourceResult:
"""Read resource contents from a specific MCP server."""
@ -4705,6 +4668,9 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
raw_headers=raw_headers,
client_ip=client_ip,
user_api_key_auth=user_api_key_auth,
)
return await client.read_resource(url)
@ -4718,6 +4684,7 @@ class MCPServerManager:
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> GetPromptResult:
"""Fetch a specific prompt definition from a single MCP server."""
@ -4738,6 +4705,9 @@ class MCPServerManager:
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
raw_headers=raw_headers,
client_ip=client_ip,
user_api_key_auth=user_api_key_auth,
)
get_prompt_request_params: Final = GetPromptRequestParams(
@ -5818,6 +5788,8 @@ class MCPServerManager:
stdio_env: dict[str, str] | None,
subject_token: str | None,
user_api_key_auth: UserAPIKeyAuth | None,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> CallToolResult:
"""Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry.
@ -5843,6 +5815,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
return await retry_client.call_tool(call_tool_params, host_progress_callback=host_progress_callback)
@ -5860,6 +5834,7 @@ class MCPServerManager:
host_progress_callback: Callable | None = None,
hook_extra_headers: dict[str, str] | None = None,
user_api_key_auth: UserAPIKeyAuth | None = None,
client_ip: str | None = None,
) -> CallToolResult:
"""
Call a regular MCP tool using the MCP client.
@ -6004,6 +5979,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
call_tool_params: Final = MCPCallToolRequestParams(
@ -6027,6 +6004,8 @@ class MCPServerManager:
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
tool_call_coro = _obo_call_tool_limited()
@ -6202,7 +6181,7 @@ class MCPServerManager:
return oauth2_headers
try:
from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415
from litellm.proxy._experimental.mcp_server.operations import ( # noqa: PLC0415
_get_user_oauth_extra_headers_from_db,
)
@ -6308,6 +6287,7 @@ class MCPServerManager:
host_progress_callback: Callable | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
guardrail_context: Mapping[str, object] | None = None,
client_ip: str | None = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments
@ -6434,6 +6414,7 @@ class MCPServerManager:
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=client_ip,
proxy_logging_obj=proxy_logging_obj,
host_progress_callback=host_progress_callback,
hook_extra_headers=hook_result.get("extra_headers"),

3102
litellm/proxy/_experimental/mcp_server/operations.py Normal file
View file

File diff suppressed because it is too large Load diff

13
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -203,17 +203,19 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.oauth_utils import (
get_request_base_url,
)
from litellm.proxy._experimental.mcp_server.server import (
from litellm.proxy._experimental.mcp_server.operations import (
ListMCPToolsRestAPIResponseObject,
MCPInfo,
MCPServer,
_aggregate_server_key, # pyright: ignore[reportPrivateUsage] # same per-server key as the tools/list _meta outcomes
_apply_toolset_scope,
_aggregate_server_key,
_fire_mcp_tool_call_logging,
execute_mcp_tool,
filter_tools_by_allowed_tools,
filter_tools_by_key_team_permissions,
fire_mcp_tool_call_failure_logging,
)
from litellm.proxy._experimental.mcp_server.server import (
_apply_toolset_scope,
reject_disallowed_mcp_client,
)
@ -670,6 +672,7 @@ if MCP_AVAILABLE:
user_api_key_auth: UserAPIKeyAuth | None = None,
extra_headers: dict[str, str] | None = None,
apply_tool_filters: bool = True,
client_ip: str | None = None,
):
"""Helper function to get tools for a single server.
@ -684,6 +687,7 @@ if MCP_AVAILABLE:
extra_headers=extra_headers,
add_prefix=False,
raw_headers=raw_headers,
client_ip=client_ip,
user_api_key_auth=user_api_key_auth,
)
@ -797,6 +801,7 @@ if MCP_AVAILABLE:
user_api_key_dict,
extra_headers=user_oauth_extra_headers,
apply_tool_filters=apply_tool_filters,
client_ip=rest_client_ip,
)
except MCPUpstreamAuthError:
# Surface the upstream 401/403 to the caller so it can emit the
@ -1016,6 +1021,7 @@ if MCP_AVAILABLE:
user_api_key_dict,
extra_headers=user_oauth_extra_headers,
apply_tool_filters=apply_tool_filters,
client_ip=_rest_client_ip,
)
except Exception as e:
verbose_logger.warning(
@ -1193,6 +1199,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers=data.get("mcp_server_auth_headers"),
oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"),
raw_headers=data.get("raw_headers"),
client_ip=IPAddressUtils.get_mcp_client_ip(request),
litellm_logging_obj=data.get("litellm_logging_obj"),
guardrail_context=MCPRequestContext.resolve_guardrail_context(data),
requested_server_id=canonical_server_id,

3062
litellm/proxy/_experimental/mcp_server/server.py
View file

File diff suppressed because it is too large Load diff

11
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -463,8 +463,8 @@ async def handle_mcp_tool_search(
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.server import (
_list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner
from litellm.proxy._experimental.mcp_server.operations import (
_list_mcp_tools,
)
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
@ -519,8 +519,8 @@ async def handle_mcp_proxy_tool(
from jsonschema import validate
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.server import ( # pyright: ignore[reportPrivateUsage] # shared catalog owner
_list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner
from litellm.proxy._experimental.mcp_server.operations import (
_list_mcp_tools,
)
listing: Final = await _list_mcp_tools(
@ -607,7 +607,7 @@ async def handle_mcp_tool_call(
requested_server_id: str | None = None,
guardrail_context: Mapping[str, object] | None = None,
) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.server import (
from litellm.proxy._experimental.mcp_server.operations import (
_get_allowed_mcp_servers,
execute_mcp_tool,
raise_denied_scoped_mcp_access,
@ -643,6 +643,7 @@ async def handle_mcp_tool_call(
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=client_ip,
litellm_logging_obj=litellm_logging_obj,
requested_server_id=requested_server_id,
guardrail_context=guardrail_context,

2
litellm/proxy/_types.py
View file

@ -3239,6 +3239,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
# above; a forged value could at most narrow, but the stripping keeps the field's provenance
# single-owner so its meaning stays trustworthy.
mcp_session_resource_server_id: str | None = Field(default=None, exclude=True)
mcp_toolset_id: str | None = Field(default=None, exclude=True)
via_virtual_key: bool = Field(
default=False,
exclude=True,
@ -3280,6 +3281,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
values.pop("mcp_admitted_user_subject", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("mcp_session_resource_server_id", None)
values.pop("mcp_toolset_id", None)
values.pop("via_virtual_key", None)
if values.get("api_key") is not None:
values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))})

65
scripts/check_mcp_operation_boundary.py Normal file
View file

@ -0,0 +1,65 @@
import ast
import sys
from pathlib import Path
from typing import Final
PACKAGE: Final = Path("litellm/proxy/_experimental/mcp_server")
LEGACY_ADAPTERS: Final = frozenset({"server.py", "legacy_callbacks.py", "mcp_context.py", "mcp_debug.py"})
CONFINED_NAMES: Final = frozenset(
{
"auth_context_var",
"active_mcp_session_var",
"active_mcp_request_ctx_var",
"get_active_auth_context",
"get_active_mcp_session",
"get_active_mcp_request_ctx",
"get_or_extract_auth_context",
"_session_obj_auth_storage",
"WeakKeyDictionary",
"_mcp_active_toolset_id",
"_mcp_gateway_initialize_instructions",
"_mcp_gateway_server_name",
"_mcp_proxy_mode",
}
)
def is_confined(name: str) -> bool:
return name in CONFINED_NAMES or name.startswith("_stateful_session_")
def violations(path: Path, source: str) -> tuple[str, ...]:
if path.name in LEGACY_ADAPTERS:
return ()
tree: Final = ast.parse(source, filename=str(path))
return tuple(
f"{path}:{node.lineno}: MCP request/session state belongs in a legacy adapter"
for node in ast.walk(tree)
if (
isinstance(node, ast.ImportFrom)
and (
(node.module or "").endswith(".mcp_context")
or any(is_confined(alias.name) for alias in node.names)
or (path.name in {"operations.py", "contracts.py"} and (node.module or "").endswith(".server"))
)
or isinstance(node, ast.Name)
and is_confined(node.id)
or isinstance(node, ast.Attribute)
and is_confined(node.attr)
)
)
def main() -> int:
findings: Final = tuple(
finding for path in sorted(PACKAGE.rglob("*.py")) for finding in violations(path, path.read_text())
)
if findings:
print("\n".join(findings), file=sys.stderr)
return 1
print("MCP operation boundary: passed")
return 0
if __name__ == "__main__":
sys.exit(main())

3
scripts/pre_commit_lint.sh
View file

@ -102,6 +102,9 @@ ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|s
ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
litellm_py_files=$(scope_match "$litellm_py_pattern")
if [ -n "$(scope_match '^(litellm/proxy/_experimental/mcp_server/|scripts/check_mcp_operation_boundary\.py)')" ]; then
uv run --no-sync python scripts/check_mcp_operation_boundary.py || exit 1
fi
e2e_py_files=$(scope_match "$e2e_py_pattern")
test_tree_files=$(scope_match "$test_tree_pattern")
# ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.

6
tests/mcp_tests/test_mcp_logging.py
View file

@ -142,7 +142,7 @@ async def test_mcp_cost_tracking():
local_mcp_server_manager,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
local_mcp_server_manager,
),
):
@ -293,7 +293,7 @@ async def test_mcp_cost_tracking_per_tool():
local_mcp_server_manager,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
local_mcp_server_manager,
),
):
@ -451,7 +451,7 @@ async def test_mcp_tool_call_hook():
local_mcp_server_manager,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
local_mcp_server_manager,
),
):

24
tests/mcp_tests/test_mcp_server.py
View file

@ -922,7 +922,7 @@ async def test_get_tools_from_mcp_servers():
)
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
mock_manager,
):
# Test with specific servers
@ -950,6 +950,7 @@ async def test_get_tools_from_mcp_servers():
extra_headers=None,
add_prefix=False,
raw_headers=None,
client_ip=None,
user_api_key_auth=None,
oauth2_headers=None,
):
@ -966,7 +967,7 @@ async def test_get_tools_from_mcp_servers():
)
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
mock_manager_2,
):
result = await _get_tools_from_mcp_servers(
@ -998,7 +999,7 @@ async def test_get_tools_from_mcp_servers():
)
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
mock_manager,
):
with patch(
@ -1981,6 +1982,7 @@ async def test_get_tools_for_single_server():
extra_headers=None,
add_prefix=False,
raw_headers=None,
client_ip=None,
user_api_key_auth=None,
)
@ -2076,7 +2078,7 @@ async def test_rest_listing_hides_key_grants_dispatch_would_refuse():
with patch(
"litellm.proxy._experimental.mcp_server.rest_endpoints.global_mcp_server_manager"
) as mock_manager, patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager"
) as mock_server_manager, patch.object(
MCPRequestHandler,
"get_allowed_tools_for_server",
@ -2473,7 +2475,7 @@ async def test_filter_tools_by_allowed_tools_integration():
# Mock the global MCP server manager
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager"
) as mock_manager:
# Mock manager methods
mock_manager.get_allowed_mcp_servers = AsyncMock(
@ -2588,7 +2590,7 @@ async def test_filter_tools_by_disallowed_tools_integration():
# Mock the global MCP server manager
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager"
) as mock_manager:
# Mock manager methods
mock_manager.get_allowed_mcp_servers = AsyncMock(
@ -2689,7 +2691,7 @@ async def test_filter_tools_no_restrictions_integration():
# Mock the global MCP server manager
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager"
) as mock_manager:
# Mock manager methods
mock_manager.get_allowed_mcp_servers = AsyncMock(
@ -2970,10 +2972,10 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
return_value=mock_server,
) as mock_get_server,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_tool_registry"
"litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry"
) as mock_tool_registry,
patch(
"litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_managed_mcp_tool",
new_callable=AsyncMock,
) as mock_handle_managed,
patch(
@ -3046,10 +3048,10 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
return_value=mock_server,
) as mock_get_server,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_tool_registry"
"litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry"
) as mock_tool_registry,
patch(
"litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_managed_mcp_tool",
new_callable=AsyncMock,
) as mock_handle_managed,
patch(

19
tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py
View file

@ -1,3 +1,4 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
"""
Unit tests for the BYOK OAuth 2.1 authorization server endpoints.
@ -592,7 +593,7 @@ async def test_check_byok_credential_missing_credential(monkeypatch):
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
server_module.byok_credential_cache.flush_cache()
mcp_operations.byok_credential_cache.flush_cache()
mock_prisma = MagicMock()
with (
@ -628,13 +629,13 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy")
mcp_module.byok_credential_cache.flush_cache()
mcp_operations.byok_credential_cache.flush_cache()
server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True)
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
with pytest.raises(HTTPException) as exc_info:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="list_regions",
arguments={},
allowed_mcp_servers=[server],
@ -687,7 +688,7 @@ async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same
server = MCPServer(server_id="byok-revoke", name="byok-server", transport=MCPTransport.http, is_byok=True)
user_auth = UserAPIKeyAuth(user_id="mallory", api_key="sk-test")
server_module.byok_credential_cache.flush_cache()
mcp_operations.byok_credential_cache.flush_cache()
db_lookup = AsyncMock(side_effect=["sk-before-revoke", None])
publish = AsyncMock()
@ -699,13 +700,13 @@ async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch.object( # test-quality-ok: the redis publisher is module-level; asserting the broadcast without a redis
server_module, "publish_auth_cache_invalidation", new=publish
mcp_operations, "publish_auth_cache_invalidation", new=publish
),
):
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
await server_module._invalidate_byok_cred_cache("mallory", "byok-revoke")
assert await server_module._get_byok_credential(server, user_auth) is None
assert await mcp_operations._get_byok_credential(server, user_auth) == "sk-before-revoke"
assert await mcp_operations._get_byok_credential(server, user_auth) == "sk-before-revoke"
await mcp_operations._invalidate_byok_cred_cache("mallory", "byok-revoke")
assert await mcp_operations._get_byok_credential(server, user_auth) is None
assert db_lookup.await_count == 2
publish.assert_awaited_once_with(cache_key=byok_credential_cache_key("mallory", "byok-revoke"))

60
tests/test_litellm/proxy/_experimental/mcp_server/test_contracts.py Normal file
View file

@ -0,0 +1,60 @@
from dataclasses import FrozenInstanceError
import pytest
from litellm.proxy._experimental.mcp_server.operations import prepare_context
from litellm.proxy._types import UserAPIKeyAuth
def test_operation_context_isolates_nested_headers_and_caller_permissions():
caller = UserAPIKeyAuth(user_id="alpha", models=["allowed"])
caller.mcp_admitted_user_subject = True
caller.mcp_session_resource_server_id = "alpha-server"
caller.mcp_toolset_id = "toolset-alpha"
caller.mcp_source_team_rpm_limits = {"team": {"alpha-server": 2}}
headers = {"x-caller": "alpha"}
server_headers = {"alpha-server": {"authorization": "alpha-token"}}
context = prepare_context(caller, raw_headers=headers, mcp_server_auth_headers=server_headers)
caller.models.append("forbidden")
caller.mcp_source_team_rpm_limits["team"]["alpha-server"] = 999
headers["x-caller"] = "bravo"
server_headers["alpha-server"]["authorization"] = "bravo-token"
captured = context.user_api_key_auth
assert captured is not None
assert captured.models == ["allowed"]
assert captured.mcp_admitted_user_subject is True
assert captured.mcp_session_resource_server_id == "alpha-server"
assert captured.mcp_toolset_id == "toolset-alpha"
assert captured.mcp_source_team_rpm_limits == {"team": {"alpha-server": 2}}
captured.models.append("also-forbidden")
assert context.user_api_key_auth.models == ["allowed"]
assert context.raw_headers == {"x-caller": "alpha"}
assert context.mcp_server_auth_headers == {"alpha-server": {"authorization": "alpha-token"}}
with pytest.raises(TypeError):
context.raw_headers["x-caller"] = "changed"
with pytest.raises(TypeError):
context.mcp_server_auth_headers["alpha-server"]["authorization"] = "changed"
with pytest.raises(FrozenInstanceError):
context.client_ip = "untrusted"
def test_operation_context_preserves_missing_and_empty_inputs():
missing = prepare_context()
empty = prepare_context(mcp_servers=[], raw_headers={}, oauth2_headers={}, mcp_server_auth_headers={})
assert missing.user_api_key_auth is None
assert missing.mcp_servers is None
assert missing.raw_headers is None
assert missing.oauth2_headers is None
assert missing.mcp_server_auth_headers is None
assert empty.mcp_servers == ()
assert empty.raw_headers == {}
assert empty.oauth2_headers == {}
assert empty.mcp_server_auth_headers == {}
def test_toolset_request_marker_cannot_be_supplied_by_caller_or_serialized():
auth = UserAPIKeyAuth.model_validate({"user_id": "alpha", "mcp_toolset_id": "forged"})
assert auth.mcp_toolset_id is None
auth.mcp_toolset_id = "server-resolved"
assert "mcp_toolset_id" not in auth.model_dump()

5
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py
View file

@ -1,5 +1,6 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
"""Tests for guardrail-block recording in
``litellm.proxy._experimental.mcp_server.server.call_mcp_tool``.
``litellm.proxy._experimental.mcp_server.operations.call_mcp_tool``.
A pre-call MCP guardrail block *raises* into ``call_mcp_tool``'s
``except Exception``. The failure spend-log row that the Guardrails Monitor's
@ -70,7 +71,7 @@ async def _call_block(logging_obj, order: list, *, user_api_key_auth=mock.sentin
with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}):
with contextlib.suppress(HTTPException):
await server.call_mcp_tool.__wrapped__(
await mcp_operations.call_mcp_tool.__wrapped__(
name="t",
arguments=None,
user_api_key_auth=user_api_key_auth,

6
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
View file

@ -1229,7 +1229,7 @@ class TestResolveByokMcpAuthHeader:
user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard")
with patch(
"litellm.proxy._experimental.mcp_server.server._get_byok_credential",
"litellm.proxy._experimental.mcp_server.operations._get_byok_credential",
new=AsyncMock(return_value="stored-cred"),
):
result = await _resolve_byok_mcp_auth_header(server, user_auth, None)
@ -1249,7 +1249,7 @@ class TestResolveByokMcpAuthHeader:
user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard")
with patch(
"litellm.proxy._experimental.mcp_server.server._get_byok_credential",
"litellm.proxy._experimental.mcp_server.operations._get_byok_credential",
new=AsyncMock(return_value=None),
):
with pytest.raises(HTTPException) as exc_info:
@ -1272,7 +1272,7 @@ class TestResolveByokMcpAuthHeader:
check_mock = AsyncMock(return_value=None)
with patch(
"litellm.proxy._experimental.mcp_server.server._check_byok_credential",
"litellm.proxy._experimental.mcp_server.operations._check_byok_credential",
new=check_mock,
):
result = await _resolve_byok_mcp_auth_header(server, user_auth, "caller-header")

61
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
View file

@ -1,3 +1,4 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
"""Unit tests for MCP OAuth passthrough tool-fetch behavior."""
import logging
@ -339,16 +340,16 @@ async def test_aggregate_list_tools_absorbs_one_unauthenticated_server():
raise MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name=server.name)
return [good_tool]
with patch.object(mcp_server, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate, working])), patch.object(
mcp_server, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={})
), patch.object(mcp_server, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object(
mcp_server, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None)
with patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate, working])), patch.object(
mcp_operations, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={})
), patch.object(mcp_operations, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object(
mcp_operations, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None)
), patch.object(
mcp_server, "filter_tools_by_key_team_permissions", AsyncMock(side_effect=lambda tools, **k: tools)
mcp_operations, "filter_tools_by_key_team_permissions", AsyncMock(side_effect=lambda tools, **k: tools)
), patch.object(
mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools)
mcp_operations.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools)
):
listing = await mcp_server._get_tools_from_mcp_servers(
listing = await mcp_operations._get_tools_from_mcp_servers(
user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"),
mcp_auth_header=None,
mcp_servers=None,
@ -382,14 +383,14 @@ async def test_single_server_route_also_absorbs_upstream_auth_error():
# /<server>/mcp sets the path-derived single-server scope; absorption must hold even then.
token = _mcp_gateway_server_name.set("delegate_docs")
try:
with patch.object(mcp_server, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object(
mcp_server, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={})
), patch.object(mcp_server, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object(
mcp_server, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None)
with patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object(
mcp_operations, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={})
), patch.object(mcp_operations, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object(
mcp_operations, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None)
), patch.object(
mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools)
mcp_operations.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools)
):
listing = await mcp_server._get_tools_from_mcp_servers(
listing = await mcp_operations._get_tools_from_mcp_servers(
user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"),
mcp_auth_header=None,
mcp_servers=["delegate_docs"],
@ -419,15 +420,15 @@ async def test_aggregate_with_single_accessible_server_still_absorbs():
async def fake_get_tools(server, **kwargs):
raise MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name=server.name)
with patch.object(mcp_server, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object(
mcp_server, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={})
), patch.object(mcp_server, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object(
mcp_server, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None)
with patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object(
mcp_operations, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={})
), patch.object(mcp_operations, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object(
mcp_operations, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None)
), patch.object(
mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools)
mcp_operations.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools)
):
# Aggregate route: no explicit server filter, even though only one server is accessible.
listing = await mcp_server._get_tools_from_mcp_servers(
listing = await mcp_operations._get_tools_from_mcp_servers(
user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"),
mcp_auth_header=None,
mcp_servers=None,
@ -475,3 +476,25 @@ async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, capl
await manager._get_tools_from_server(server)
assert "POST https://upstream/ -> HTTP 500" in caplog.text
assert "missing_scope" in caplog.text and "query-secret" not in caplog.text
@pytest.mark.parametrize(
"oauth_headers,server_headers,authorized",
[
({"Authorization": "Bearer upstream"}, None, True),
({"AUTHORIZATION": "Bearer upstream"}, None, True),
({"x-unrelated": "present"}, None, False),
(None, {"catalog": {"Authorization": "Bearer scoped"}}, True),
(None, {"other-server": {"Authorization": "Bearer unrelated"}}, False),
(None, {"catalog": {"x-unrelated": "present"}}, False),
(None, {"catalog": "Bearer legacy"}, True),
(None, {"catalog": " "}, False),
],
)
def test_passthrough_admission_recognizes_only_matching_authorization(oauth_headers, server_headers, authorized):
from litellm.proxy._experimental.mcp_server.operations import _client_has_passthrough_authorization
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(server_id="catalog", name="catalog", alias="catalog", transport=MCPTransport.http)
assert _client_has_passthrough_authorization(server, oauth_headers, server_headers) is authorized

8
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py
View file

@ -1,3 +1,4 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
import json
from datetime import datetime
@ -27,8 +28,8 @@ def proxy_mode():
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_mode")
async def test_proxy_call_rejects_non_proxy_tool_names() -> None:
result = await server._dispatch_virtual_mcp_tool(
name="math_stdio-add", arguments={"a": 1, "b": 2}, user_api_key_auth=AUTH, client_ip=None
result = await mcp_operations._dispatch_virtual_mcp_tool(
name="math_stdio-add", arguments={"a": 1, "b": 2}, user_api_key_auth=AUTH, client_ip=None, mcp_proxy_mode=True
)
assert result is not None
@ -105,12 +106,13 @@ async def test_proxy_scope_exception_emits_failure_log(monkeypatch: pytest.Monke
arguments = {"tool_id": "denied-scope", "arguments": {}}
with pytest.raises(HTTPException) as denied:
await server._dispatch_virtual_mcp_tool(
await mcp_operations._dispatch_virtual_mcp_tool(
name="call_tool",
arguments=arguments,
user_api_key_auth=auth,
client_ip=None,
mcp_servers=["ungranted"],
mcp_proxy_mode=True,
raw_headers={"authorization": "Bearer raw-scope-secret", "x-litellm-call-id": "scope-denial"},
)

385
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

File diff suppressed because it is too large Load diff

179
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -73,6 +73,135 @@ from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
@pytest.mark.asyncio
async def test_manager_sampling_preserves_explicit_headers_without_ambient_context():
from litellm.proxy._experimental.mcp_server import server as legacy_server
caller = UserAPIKeyAuth(user_id="sampling-caller")
upstream = MCPServer(
server_id="sampling-context",
name="sampling_context",
url="https://example.invalid/mcp",
transport=MCPTransport.http,
allow_sampling=True,
)
sampling = AsyncMock()
client = MagicMock()
client.call_tool = AsyncMock(return_value=CallToolResult(content=[]))
assert legacy_server.get_active_auth_context() is None
with (
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", return_value=client) as factory,
patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling),
):
await MCPServerManager()._call_regular_mcp_tool(
mcp_server=upstream,
original_tool_name="probe",
arguments={},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
raw_headers={"x-test-caller": "sampling-caller"},
proxy_logging_obj=None,
user_api_key_auth=caller,
)
callback = factory.call_args.kwargs["sampling_callback"]
await callback(None, None)
assert sampling.await_args.kwargs["user_api_key_auth"].user_id == "sampling-caller"
assert sampling.await_args.kwargs["raw_headers"] == {"x-test-caller": "sampling-caller"}
@pytest.mark.asyncio
async def test_sampling_callback_keeps_creation_context_after_caller_switch():
from mcp.server.auth.middleware.auth_context import auth_context_var
from litellm.proxy._experimental.mcp_server import server as legacy_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback
token = auth_context_var.set(None)
recorder = AsyncMock()
try:
original = UserAPIKeyAuth(user_id="alpha", models=["alpha-model"])
original.mcp_admitted_user_subject = True
headers = {"x-caller": "alpha"}
legacy_server.set_auth_context(original, raw_headers=headers, client_ip="192.0.2.1")
callback = _create_sampling_callback()
original.models.append("bravo-model")
headers["x-caller"] = "bravo"
legacy_server.set_auth_context(UserAPIKeyAuth(user_id="bravo"), raw_headers={"x-caller": "bravo"})
with patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", recorder):
await callback(None, None)
observed = recorder.await_args.kwargs
assert observed["user_api_key_auth"].user_id == "alpha"
assert observed["user_api_key_auth"].models == ["alpha-model"]
assert observed["user_api_key_auth"].mcp_admitted_user_subject is True
assert observed["raw_headers"] == {"x-caller": "alpha"}
assert observed["client_ip"] == "192.0.2.1"
finally:
auth_context_var.reset(token)
@pytest.mark.asyncio
async def test_elicitation_callback_keeps_initiating_session():
from litellm.proxy._experimental.mcp_server import server as legacy_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_elicitation_callback
initiating = MagicMock()
replacement = MagicMock()
recorder = AsyncMock()
token = legacy_server.active_mcp_session_var.set(initiating)
try:
callback = _create_elicitation_callback()
legacy_server.active_mcp_session_var.set(replacement)
with patch("litellm.proxy._experimental.mcp_server.elicitation_handler.handle_elicitation_request", recorder):
await callback(None, None)
assert recorder.await_args.kwargs["downstream_session"] is initiating
assert recorder.await_args.kwargs["downstream_capabilities"] is initiating.capabilities
finally:
legacy_server.active_mcp_session_var.reset(token)
@pytest.mark.asyncio
async def test_sampling_callbacks_isolate_callers_and_cancellation():
from mcp.types import ErrorData
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback
started = asyncio.Event()
cancelled = asyncio.Event()
observed = {}
async def record_sampling(*, user_api_key_auth, raw_headers, **kwargs):
label = user_api_key_auth.user_id
if label == "cancelled":
started.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
cancelled.set()
raise
await asyncio.sleep(0)
observed[label] = raw_headers["x-caller"]
return ErrorData(code=-1, message=label)
callbacks = tuple(
_create_sampling_callback(UserAPIKeyAuth(user_id=label), raw_headers={"x-caller": label})
for label in ("alpha", "bravo", "cancelled")
)
with patch(
"litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", record_sampling
):
tasks = tuple(asyncio.create_task(callback(None, None)) for callback in callbacks)
await asyncio.wait_for(started.wait(), timeout=2)
tasks[2].cancel()
results = await asyncio.gather(*tasks, return_exceptions=True)
assert observed == {"alpha": "alpha", "bravo": "bravo"}
assert [result.message for result in results[:2]] == ["alpha", "bravo"]
assert isinstance(results[2], asyncio.CancelledError)
assert cancelled.is_set()
def _reload_mcp_manager_module():
utils_module = sys.modules["litellm.proxy._experimental.mcp_server.utils"]
manager_module = sys.modules["litellm.proxy._experimental.mcp_server.mcp_server_manager"]
@ -84,6 +213,9 @@ def _reload_mcp_manager_module():
server_module = sys.modules.get("litellm.proxy._experimental.mcp_server.server")
if server_module is not None and hasattr(server_module, "global_mcp_server_manager"):
server_module.global_mcp_server_manager = reloaded.global_mcp_server_manager
operations_module = sys.modules.get("litellm.proxy._experimental.mcp_server.operations")
if operations_module is not None:
operations_module.global_mcp_server_manager = reloaded.global_mcp_server_manager
return reloaded
@ -3923,6 +4055,7 @@ class TestMCPServerManager:
result = await manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=None,
raw_headers=None,
mcp_auth_header="auth",
extra_headers=None,
add_prefix=False,
@ -3935,6 +4068,8 @@ class TestMCPServerManager:
stdio_env=None,
subject_token=None,
user_api_key_auth=None,
raw_headers=None,
client_ip=None,
)
mock_client.list_resource_templates.assert_awaited_once()
assert result == expected_templates
@ -5849,7 +5984,7 @@ class TestMCPServerManager:
stored = {"Authorization": "Bearer stored-user-token"}
with patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db",
new=AsyncMock(return_value=stored),
) as mock_lookup:
result = await manager._resolve_oauth2_headers_for_tool_call(
@ -5876,7 +6011,7 @@ class TestMCPServerManager:
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
with patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db",
new=AsyncMock(return_value={"Authorization": "Bearer should-not-be-used"}),
) as mock_lookup:
result = await manager._resolve_oauth2_headers_for_tool_call(
@ -5902,7 +6037,7 @@ class TestMCPServerManager:
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
with patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db",
new=AsyncMock(side_effect=RuntimeError("redis down")),
):
result = await manager._resolve_oauth2_headers_for_tool_call(
@ -6058,7 +6193,7 @@ class TestMCPServerManager:
user_auth = UserAPIKeyAuth(api_key="sk-test")
with patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db",
new=AsyncMock(return_value={"Authorization": "Bearer x"}),
) as mock_lookup:
result = await manager._resolve_oauth2_headers_for_tool_call(
@ -6862,7 +6997,8 @@ class TestMCPServerManager:
}
user_api_key_auth = UserAPIKeyAuth(api_key="sk-test", user_id="user-123")
token = _mcp_active_toolset_id.set("toolset-abc")
user_api_key_auth.mcp_toolset_id = "toolset-abc"
token = _mcp_active_toolset_id.set("unrelated-ambient-toolset")
try:
with (
patch.object(proxy_server_module, "user_api_key_cache", cache),
@ -14332,3 +14468,36 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon
assert guardrail_started.is_set() is selected
assert result.is_error is False
assert result.content[0].text == "executed"
@pytest.mark.asyncio
@pytest.mark.parametrize("with_caller,legacy_factory", [(True, False), (False, False), (True, True)])
async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller, legacy_factory):
from mcp.server.auth.middleware.auth_context import auth_context_var
from litellm.proxy._experimental.mcp_server import server as legacy_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback
upstream = MCPServer(server_id="explicit-empty", name="explicit_empty", url="https://example.invalid/mcp", transport=MCPTransport.http, allow_sampling=True)
token = auth_context_var.set(None)
sampling = AsyncMock()
try:
legacy_server.set_auth_context(UserAPIKeyAuth(user_id="unrelated"), raw_headers={"authorization": "unrelated-credential"}, client_ip="192.0.2.99")
with (
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as factory,
patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling),
):
if legacy_factory:
callback = _create_sampling_callback(user_api_key_auth=UserAPIKeyAuth(user_id="explicit"))
else:
await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None)
callback = factory.call_args.kwargs["sampling_callback"]
await callback(None, None)
captured = sampling.await_args.kwargs
if with_caller:
assert captured["user_api_key_auth"].user_id == "explicit"
else:
assert captured["user_api_key_auth"] is None
assert captured["raw_headers"] is None
assert captured["client_ip"] is None
finally:
auth_context_var.reset(token)

40
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py
View file

@ -639,12 +639,12 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
return_value=False,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
),
patch.object(
@ -727,12 +727,12 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
return_value=False,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
),
patch.object(
@ -833,11 +833,11 @@ async def test_client_credentials_server_is_not_preemptively_challenged(m2m_fiel
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=m2m_server,
),
patch.object(session_manager_stateless, "handle_request", new_callable=AsyncMock) as mock_handle_request,
@ -929,16 +929,16 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=delegated_server,
),
patch( # test-quality-ok: registry is empty in unit tests; key owns the delegated server
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
"litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[delegated_server],
),
@ -1022,12 +1022,12 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
return_value=True,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
),
patch.object(
@ -1126,11 +1126,11 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=delegated_server,
),
patch.object(
@ -1218,7 +1218,7 @@ async def test_handle_streamable_http_mcp_token_exchange_without_subject_returns
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=obo_server,
),
patch.object(
@ -1317,7 +1317,7 @@ async def test_handle_streamable_http_mcp_oauth_delegate_without_token_returns_g
True,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=od_server,
),
patch.object(
@ -1391,7 +1391,7 @@ async def test_handle_streamable_http_mcp_oauth_delegate_with_forwarded_token_sk
new_callable=AsyncMock,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=od_server,
),
patch.object(
@ -1453,7 +1453,7 @@ async def _run_passthrough_connect(
new_callable=AsyncMock,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=server,
),
patch.object(session_manager_stateless, "handle_request", new_callable=AsyncMock) as mock_handle_request,
@ -1574,7 +1574,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_without_token_surface
return_value=probe_client,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=tp_server,
),
patch.object(
@ -1642,7 +1642,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_dcr_bridge_challenges
return_value=probe_client,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=bridge_server,
),
patch.object(
@ -1720,7 +1720,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_with_token_skips_prob
return_value=probe_client,
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name",
return_value=tp_server,
),
patch.object(

40
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
View file

@ -1,3 +1,4 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
"""
Tests for MCP tool search feature.
@ -572,7 +573,7 @@ class TestCallToolRestApiVirtualTools:
mock_tool.input_schema = {"type": "object", "properties": {}}
with patch(
"litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
"litellm.proxy._experimental.mcp_server.operations._list_mcp_tools",
new_callable=AsyncMock,
return_value=AggregateToolListing(tools=[mock_tool], outcomes={}),
):
@ -616,12 +617,12 @@ class TestCallToolRestApiVirtualTools:
with (
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
"litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[MagicMock()],
),
patch(
"litellm.proxy._experimental.mcp_server.server.execute_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool",
new_callable=AsyncMock,
return_value=fake_result,
) as mock_execute,
@ -669,12 +670,12 @@ class TestCallToolRestApiVirtualTools:
return_value="203.0.113.7",
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
"litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[MagicMock()],
) as mock_allowed,
patch(
"litellm.proxy._experimental.mcp_server.server.execute_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool",
new_callable=AsyncMock,
return_value=fake_result,
),
@ -699,7 +700,7 @@ class TestCallToolRestApiVirtualTools:
return_value="203.0.113.7",
),
patch(
"litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
"litellm.proxy._experimental.mcp_server.operations._list_mcp_tools",
new_callable=AsyncMock,
return_value=AggregateToolListing(tools=[], outcomes={}),
) as mock_list,
@ -832,7 +833,7 @@ class TestCallToolRestApiVirtualTools:
"litellm.proxy.proxy_server.proxy_logging_obj", key_limits
),
patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real
"litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
"litellm.proxy._experimental.mcp_server.operations._list_mcp_tools",
new_callable=AsyncMock,
return_value=AggregateToolListing(tools=list(CATALOG), outcomes={}),
) as mock_list,
@ -939,7 +940,7 @@ class TestDispatchVirtualMcpTool:
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
) as mock_search:
result = await srv._dispatch_virtual_mcp_tool(
result = await mcp_operations._dispatch_virtual_mcp_tool(
name=MCP_TOOL_SEARCH_TOOL_NAME,
arguments={"query": "q", "top_k": 3},
user_api_key_auth=uak,
@ -961,7 +962,7 @@ class TestDispatchVirtualMcpTool:
new_callable=AsyncMock,
return_value="AGENT_RESULT",
) as mock_agent_search:
result = await srv._dispatch_virtual_mcp_tool(
result = await mcp_operations._dispatch_virtual_mcp_tool(
name=AGENT_SEARCH_TOOL_NAME,
arguments={"query": "translate a document", "top_k": "2"},
user_api_key_auth=uak,
@ -996,7 +997,7 @@ class TestDispatchVirtualMcpTool:
new_callable=AsyncMock,
return_value="CALL_RESULT",
) as mock_call:
result = await srv._dispatch_virtual_mcp_tool(
result = await mcp_operations._dispatch_virtual_mcp_tool(
name=MCP_TOOL_CALL_TOOL_NAME,
arguments={"tool_name": "math-add", "arguments": {"a": 1, "b": 2}},
user_api_key_auth=uak,
@ -1027,8 +1028,7 @@ class TestDispatchVirtualMcpTool:
sentinel_logging_obj = object()
with (
patch.object(
srv,
"_build_virtual_call_logging_obj",
mcp_operations, "_build_virtual_call_logging_obj",
new_callable=AsyncMock,
return_value=sentinel_logging_obj,
) as mock_build,
@ -1038,7 +1038,7 @@ class TestDispatchVirtualMcpTool:
return_value="CALL_RESULT",
) as mock_call,
):
await srv._dispatch_virtual_mcp_tool(
await mcp_operations._dispatch_virtual_mcp_tool(
name=MCP_TOOL_CALL_TOOL_NAME,
arguments={"tool_name": "math-add", "arguments": {"a": 1}},
user_api_key_auth=uak,
@ -1060,7 +1060,7 @@ class TestDispatchVirtualMcpTool:
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
) as mock_search:
await srv._dispatch_virtual_mcp_tool(
await mcp_operations._dispatch_virtual_mcp_tool(
name=MCP_TOOL_SEARCH_TOOL_NAME,
arguments={"query": "issue", "top_k": "not-a-number"},
user_api_key_auth=uak,
@ -1083,12 +1083,12 @@ class TestDispatchVirtualMcpTool:
fake = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
with (
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
"litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[MagicMock()],
) as mock_allowed,
patch(
"litellm.proxy._experimental.mcp_server.server.execute_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool",
new_callable=AsyncMock,
return_value=fake,
) as mock_exec,
@ -1130,12 +1130,12 @@ class TestDispatchVirtualMcpTool:
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
with (
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
"litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[],
),
patch(
"litellm.proxy._experimental.mcp_server.server.execute_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool",
new_callable=AsyncMock,
) as mock_exec,
):
@ -1217,7 +1217,7 @@ class TestMcpServerToolCallErrorHandling:
return_value=(uak, None, None, None, None, None, None),
),
patch(
"litellm.proxy._experimental.mcp_server.server._dispatch_virtual_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._dispatch_virtual_mcp_tool",
new_callable=AsyncMock,
side_effect=HTTPException(status_code=403, detail="User not allowed to call this tool"),
),
@ -1254,7 +1254,7 @@ async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> N
]
with patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
"litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(side_effect=resolve),
):
with pytest.raises(HTTPException) as exc_info:

18
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py
View file

@ -58,6 +58,22 @@ class TestApplyToolsetScope:
assert set(op.mcp_servers or []) == {"server-a", "server-b"}
assert op.mcp_tool_permissions == toolset_perms
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
from litellm.proxy._experimental.mcp_server.operations import prepare_context
manager = MCPServerManager()
unscoped_open = await manager.operator_open_server_ids(
auth, allow_all_server_ids=["operator-open-outside-toolset"], submitted_server_ids=[]
)
scoped_open = await manager.operator_open_server_ids(
prepare_context(result).user_api_key_auth,
allow_all_server_ids=["operator-open-outside-toolset"],
submitted_server_ids=[],
)
assert unscoped_open == {"operator-open-outside-toolset"}
assert scoped_open == set()
assert auth.mcp_toolset_id is None
@pytest.mark.asyncio
async def test_admin_creates_object_permission_when_none(self):
"""Admin key with object_permission=None can access any toolset."""
@ -564,7 +580,7 @@ class TestMCPActiveToolsetContextVar:
MagicMock(get_mcp_client_ip=MagicMock(return_value="127.0.0.1")),
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
"litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager",
MagicMock(get_mcp_server_by_name=MagicMock(return_value=None)),
),
patch(

75
tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py
View file

@ -1,3 +1,4 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
"""
VERIA-7 regression: OpenAPI-backed (local-registry) MCP tools must run
through `pre_call_tool_check` before dispatch, the same as managed
@ -49,22 +50,22 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
with (
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=fake_server,
),
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"pre_call_tool_check",
new=pre_call,
),
patch.object(
mcp_module.global_mcp_tool_registry,
mcp_operations.global_mcp_tool_registry,
"get_tool",
return_value=fake_tool,
),
patch(
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool",
new=handle_local,
),
patch(
@ -72,7 +73,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
return_value=True,
),
):
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="list_pets",
arguments={"limit": 10},
allowed_mcp_servers=[fake_server],
@ -92,7 +93,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
assert pre_call_kwargs["guardrail_context"] == {"metadata": {"guardrails": ("block-all",)}}
assert pre_call_kwargs["name"] == "list_pets"
assert pre_call_kwargs["server"] is fake_server
assert pre_call_kwargs["user_api_key_auth"] is user
assert pre_call_kwargs["user_api_key_auth"] == user
# `proxy_logging_obj` must be sourced from the canonical proxy_server
# module (same as the managed path) — passing None would crash the
# downstream `_create_mcp_request_object_from_kwargs` call with
@ -134,22 +135,22 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises():
with (
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=fake_server,
),
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"pre_call_tool_check",
new=pre_call,
),
patch.object(
mcp_module.global_mcp_tool_registry,
mcp_operations.global_mcp_tool_registry,
"get_tool",
return_value=fake_tool,
),
patch(
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool",
new=handle_local,
),
patch(
@ -158,7 +159,7 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises():
),
):
with pytest.raises(HTTPException) as exc:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="delete_pet",
arguments={},
allowed_mcp_servers=[fake_server],
@ -195,24 +196,24 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable():
# `_get_mcp_server_from_tool_name` returns None — no server context.
with (
patch.object(mcp_module, "_resolve_openapi_tool_auth", new=resolve_auth),
patch.object(mcp_operations, "_resolve_openapi_tool_auth", new=resolve_auth),
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=None,
),
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"pre_call_tool_check",
new=pre_call,
),
patch.object(
mcp_module.global_mcp_tool_registry,
mcp_operations.global_mcp_tool_registry,
"get_tool",
return_value=fake_tool,
),
patch(
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool",
new=handle_local,
),
patch(
@ -221,7 +222,7 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable():
),
):
with pytest.raises(HTTPException) as exc:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="list_pets",
arguments={},
allowed_mcp_servers=[],
@ -280,27 +281,27 @@ async def test_openapi_local_tool_injects_resolved_oauth_token():
with (
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=oauth_server,
),
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"pre_call_tool_check",
new=AsyncMock(return_value={}),
),
patch.object(
mcp_module.global_mcp_tool_registry,
mcp_operations.global_mcp_tool_registry,
"get_tool",
return_value=fake_tool,
),
patch.object(
mcp_module.global_mcp_server_manager._cred_provider,
mcp_operations.global_mcp_server_manager._cred_provider,
"resolve_credentials",
new=AsyncMock(return_value=Ok(StaticHeaderAuth("Bearer stored-user-token"))),
),
patch(
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool",
new=handle_local,
),
patch(
@ -308,7 +309,7 @@ async def test_openapi_local_tool_injects_resolved_oauth_token():
return_value=True,
),
):
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="get_values",
arguments={},
allowed_mcp_servers=[oauth_server],
@ -417,7 +418,7 @@ async def test_legacy_local_tool_fallback_refuses_unentitled_caller(legacy_local
)
with pytest.raises(HTTPException) as exc:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}",
arguments={},
allowed_mcp_servers=[server],
@ -451,7 +452,7 @@ async def test_legacy_local_tool_fallback_still_dispatches_entitled_caller(
server, executed = legacy_local_tool
user = _caller_entitled_to([LEGACY_TOOL])
result = await mcp_module.execute_mcp_tool(
result = await mcp_operations.execute_mcp_tool(
name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}",
arguments={},
allowed_mcp_servers=[server],
@ -481,7 +482,7 @@ async def test_legacy_local_tool_fallback_fails_closed_on_empty_prefix(
_server, executed = legacy_local_tool
with pytest.raises(HTTPException) as exc:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name=f"-{LEGACY_TOOL}",
arguments={},
allowed_mcp_servers=[],
@ -523,7 +524,7 @@ async def test_legacy_local_tool_fallback_fails_closed_when_prefix_names_no_serv
return_value=True,
):
with pytest.raises(HTTPException) as exc:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}",
arguments={},
allowed_mcp_servers=[other_server],
@ -546,7 +547,7 @@ async def test_unknown_tool_name_still_reports_not_found():
from litellm.proxy._experimental.mcp_server import server as mcp_module
with pytest.raises(HTTPException) as exc:
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="tool_no_registry_knows",
arguments={},
allowed_mcp_servers=[],
@ -610,7 +611,7 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc
captured["injected"] = _request_auth_header.get()
return []
manager = mcp_module.global_mcp_server_manager
manager = mcp_operations.global_mcp_server_manager
with (
patch.object(manager, "resolve_openapi_upstream_auth", new=fake_resolver),
patch.object(manager, "pre_call_tool_check", new=AsyncMock(return_value={})),
@ -620,9 +621,9 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc
fake_tool.name = "list_reports"
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server),
patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
patch.object(mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
patch(
"litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
"litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool",
new=capture_local,
),
patch(
@ -630,7 +631,7 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc
return_value=True,
),
):
await mcp_module.execute_mcp_tool(
await mcp_operations.execute_mcp_tool(
name="list_reports",
arguments={},
allowed_mcp_servers=[server],
@ -702,11 +703,11 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st
user = UserAPIKeyAuth(api_key="sk-user", user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value)
with (
patch.object(mcp_module.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=server),
patch.object(mcp_module.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={})),
patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
patch.object(mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=server),
patch.object(mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={})),
patch.object(mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
patch.object(
mcp_module.global_mcp_server_manager,
mcp_operations.global_mcp_server_manager,
"resolve_openapi_upstream_auth",
new=AsyncMock(return_value=(None, None)),
),
@ -715,7 +716,7 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st
return_value=True,
),
):
call = mcp_module.execute_mcp_tool(
call = mcp_operations.execute_mcp_tool(
name="list_reports",
arguments={},
allowed_mcp_servers=[server],

365
tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py Normal file
View file

@ -0,0 +1,365 @@
import asyncio
from unittest.mock import AsyncMock, patch
import pytest
from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult
from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, prepare_context
from litellm.proxy._types import UserAPIKeyAuth
@pytest.mark.asyncio
async def test_oauth_prefetch_failure_does_not_log_caller_or_exception_text(caplog):
from litellm.proxy._experimental.mcp_server.operations import _prefetch_oauth_creds_for_user
user_id = "caller\nFORGED-USER-LINE"
fetch = AsyncMock(side_effect=RuntimeError("database\nFORGED-ERROR-LINE"))
database = object()
with (
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=database),
patch("litellm.proxy._experimental.mcp_server.db.list_user_oauth_credentials", fetch),
caplog.at_level("WARNING", logger="LiteLLM"),
):
result = await _prefetch_oauth_creds_for_user(UserAPIKeyAuth(user_id=user_id))
assert result == {}
fetch.assert_awaited_once_with(database, user_id)
warnings = [record.getMessage() for record in caplog.records if "prefetch" in record.getMessage()]
assert len(warnings) == 1
assert "failed" in warnings[0]
assert "\n" not in warnings[0]
assert "FORGED" not in warnings[0]
@pytest.mark.asyncio
async def test_dispatch_uses_explicit_context_when_ambient_caller_differs():
from mcp.server.auth.middleware.auth_context import auth_context_var
from litellm.proxy._experimental.mcp_server.server import set_auth_context
context = prepare_context(
UserAPIKeyAuth(user_id="alpha"),
raw_headers={"x-caller": "alpha"},
mcp_servers=["alpha-server"],
client_ip="192.0.2.1",
)
token = auth_context_var.set(None)
handler = AsyncMock(return_value=GetPromptResult(messages=[]))
try:
set_auth_context(UserAPIKeyAuth(user_id="bravo"), raw_headers={"x-caller": "bravo"})
with patch("litellm.proxy._experimental.mcp_server.operations.mcp_get_prompt", handler):
result = await GatewayOperations().execute(
GetPromptRequest(params=GetPromptRequestParams(name="alpha-prompt")), context
)
assert result.messages == []
assert handler.await_args.kwargs["name"] == "alpha-prompt"
assert handler.await_args.kwargs["user_api_key_auth"].user_id == "alpha"
assert handler.await_args.kwargs["raw_headers"] == {"x-caller": "alpha"}
assert handler.await_args.kwargs["mcp_servers"] == ["alpha-server"]
assert handler.await_args.kwargs["client_ip"] == "192.0.2.1"
finally:
auth_context_var.reset(token)
@pytest.mark.asyncio
async def test_legacy_adapter_cleans_context_after_cancelled_operation():
from types import SimpleNamespace
from litellm.proxy._experimental.mcp_server import server
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
previous_session = server.active_mcp_session_var.get()
previous_request = active_mcp_request_ctx_var.get()
request = SimpleNamespace(session=object())
auth = (None, None, None, None, None, None, None)
async def cancelled_operation():
async with server._legacy_operation_context(request, trace=False):
assert server.active_mcp_session_var.get() is request.session
assert active_mcp_request_ctx_var.get() is request
raise asyncio.CancelledError
with patch(
"litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", AsyncMock(return_value=auth)
):
with pytest.raises(asyncio.CancelledError):
await cancelled_operation()
assert server.active_mcp_session_var.get() is previous_session
assert active_mcp_request_ctx_var.get() is previous_request
@pytest.mark.asyncio
async def test_legacy_adapter_cleans_context_when_trace_setup_fails():
from types import SimpleNamespace
from litellm.proxy._experimental.mcp_server import server
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
previous_session = server.active_mcp_session_var.get()
previous_request = active_mcp_request_ctx_var.get()
request = SimpleNamespace(session=object())
async def enter_operation():
async with server._legacy_operation_context(request, trace=True):
pytest.fail("Trace setup failure must prevent dispatch")
with patch.object(server, "_otel_set_mcp_transport_span", side_effect=RuntimeError("trace failure")):
with pytest.raises(RuntimeError, match="trace failure"):
await enter_operation()
assert server.active_mcp_session_var.get() is previous_session
assert active_mcp_request_ctx_var.get() is previous_request
@pytest.mark.asyncio
async def test_prompt_sampling_receives_explicit_operation_caller_headers_and_ip():
from unittest.mock import MagicMock
from litellm.proxy._experimental.mcp_server import operations
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
upstream = MCPServer(
server_id="explicit-prompt",
name="explicit_prompt",
url="https://example.invalid/mcp",
transport=MCPTransport.http,
allow_sampling=True,
)
context = prepare_context(
UserAPIKeyAuth(user_id="prompt-caller"),
raw_headers={"x-caller": "prompt-caller"},
client_ip="192.0.2.41",
)
client = MagicMock()
client.get_prompt = AsyncMock(return_value=GetPromptResult(messages=[]))
sampling = AsyncMock()
with (
patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[upstream])),
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", return_value=client) as factory,
patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling),
):
result = await GatewayOperations().execute(
GetPromptRequest(params=GetPromptRequestParams(name="explicit_prompt-prompt")), context
)
assert result.messages == []
await factory.call_args.kwargs["sampling_callback"](None, None)
captured = sampling.await_args.kwargs
assert captured["user_api_key_auth"] is not None
assert captured["user_api_key_auth"].user_id == "prompt-caller"
assert captured["raw_headers"] == {"x-caller": "prompt-caller"}
assert captured["client_ip"] == "192.0.2.41"
def _catalog_case(method):
from mcp import types
cases = {
"prompts/list": (
types.ListPromptsRequest(),
"list_prompts",
"get_prompts_from_server",
[types.Prompt(name="catalog-prompt")],
"prompts",
),
"prompts/get": (
types.GetPromptRequest(
params=types.GetPromptRequestParams(name="catalog-prompt", arguments={"topic": "test"})
),
"get_prompt",
"get_prompt_from_server",
types.GetPromptResult(messages=[]),
None,
),
"resources/list": (
types.ListResourcesRequest(),
"list_resources",
"get_resources_from_server",
[types.Resource(name="document", uri="https://example.com/document")],
"resources",
),
"resources/templates/list": (
types.ListResourceTemplatesRequest(),
"list_resource_templates",
"get_resource_templates_from_server",
[types.ResourceTemplate(name="document", uri_template="https://example.com/{name}")],
"resource_templates",
),
"resources/read": (
types.ReadResourceRequest(params=types.ReadResourceRequestParams(uri="https://example.com/document")),
"read_resource",
"read_resource_from_server",
types.ReadResourceResult(
contents=[types.TextResourceContents(uri="https://example.com/document", text="document body")]
),
None,
),
}
return cases[method]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"method", ["prompts/list", "prompts/get", "resources/list", "resources/templates/list", "resources/read"]
)
@pytest.mark.parametrize("state", ["success", "denied", "upstream_failure", "scope_failure"])
async def test_native_catalog_operations_preserve_context_results_and_failure_policy(method, state):
from types import SimpleNamespace
from fastapi import HTTPException
from mcp.server.context import ServerRequestContext
from mcp.types import PaginatedRequestParams
from litellm.proxy._experimental.mcp_server import operations, server
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
operation, handler_name, manager_method, payload, collection = _catalog_case(method)
caller = UserAPIKeyAuth(user_id="catalog-caller")
headers = {"x-caller": "catalog-caller"}
upstream_server = MCPServer(server_id="catalog", name="catalog", transport=MCPTransport.http)
allowed = AsyncMock(
return_value=[] if state == "denied" else [upstream_server],
side_effect=HTTPException(status_code=403, detail="scope denied") if state == "scope_failure" else None,
)
upstream = AsyncMock(
return_value=payload, side_effect=RuntimeError("upstream unavailable") if state == "upstream_failure" else None
)
ctx = ServerRequestContext(
session=SimpleNamespace(), lifespan_context={}, protocol_version="2025-06-18", method=method
)
auth = (caller, None, ["catalog"], None, None, headers, "192.0.2.41")
with (
patch.object(server, "get_or_extract_auth_context", AsyncMock(return_value=auth)),
patch.object(operations, "_get_allowed_mcp_servers", allowed),
patch.object(operations.global_mcp_server_manager, manager_method, upstream),
):
if collection is None and state != "success":
expected_error = RuntimeError if state == "upstream_failure" else HTTPException
with pytest.raises(expected_error):
await getattr(server, handler_name)(ctx, operation.params)
else:
result = await getattr(server, handler_name)(ctx, operation.params or PaginatedRequestParams())
if collection:
assert getattr(result, collection) == (payload if state == "success" else [])
else:
assert result == payload
assert allowed.await_args.kwargs == {
"user_api_key_auth": caller,
"mcp_servers": ["catalog"],
"client_ip": "192.0.2.41",
}
if state in ("denied", "scope_failure"):
upstream.assert_not_awaited()
else:
upstream.assert_awaited_once()
forwarded = upstream.await_args.kwargs
assert forwarded["user_api_key_auth"] == caller
assert forwarded["raw_headers"] == headers
assert forwarded["client_ip"] == "192.0.2.41"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"method", ["prompts/list", "prompts/get", "resources/list", "resources/templates/list", "resources/read"]
)
async def test_explicit_proxy_context_rejects_catalog_operations_before_upstream_access(method):
from mcp.shared.exceptions import MCPError
from mcp.types import METHOD_NOT_FOUND
from litellm.proxy._experimental.mcp_server import operations
operation, _, manager_method, _, _ = _catalog_case(method)
upstream = AsyncMock()
with patch.object(operations.global_mcp_server_manager, manager_method, upstream):
with pytest.raises(MCPError) as rejected:
await GatewayOperations().execute(operation, prepare_context(mcp_proxy_mode=True))
assert rejected.value.error.code == METHOD_NOT_FOUND
assert rejected.value.error.message == "Operation unavailable on /mcp/proxy"
upstream.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ["missing_env", "pii", "guardrail", "unexpected"])
async def test_tool_operation_preserves_failure_messages_and_request_trace(failure):
from mcp.types import CallToolRequest, CallToolRequestParams
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.proxy._experimental.mcp_server import operations
from litellm.proxy._experimental.mcp_server.utils import MCPMissingUserEnvVarsError
failures = {
"missing_env": (
MCPMissingUserEnvVarsError(
server_id="server", server_name="server", missing=["TOKEN"], setup_url="https://example.com/setup"
),
"https://example.com/setup",
),
"pii": (
BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="test"),
"Blocked PII entity detected",
),
"guardrail": (GuardrailRaisedException(message="request denied"), "Guardrail violation"),
"unexpected": (RuntimeError("upstream unavailable"), "Error: upstream unavailable"),
}
error, expected = failures[failure]
dispatch = AsyncMock(side_effect=error)
context = prepare_context(
raw_headers={"x-litellm-trace-id": "operation-trace", "authorization": "private-test-header"}
)
with patch.object(operations, "call_mcp_tool", dispatch):
result = await GatewayOperations().execute(
CallToolRequest(params=CallToolRequestParams(name="catalog-tool", arguments={})), context
)
assert result.is_error is True
assert expected in result.content[0].text
assert "private-test-header" not in result.content[0].text
dispatch.assert_awaited_once()
assert dispatch.await_args.kwargs["litellm_trace_id"] == "operation-trace"
assert dispatch.await_args.kwargs["litellm_session_id"] == "operation-trace"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"method,helper",
[
("prompts/list", "_list_mcp_prompts"),
("resources/list", "_list_mcp_resources"),
("resources/templates/list", "_list_mcp_resource_templates"),
],
)
async def test_catalog_operation_preserves_empty_result_for_malformed_upstream_items(method, helper):
from litellm.proxy._experimental.mcp_server import operations
operation, _, _, _, collection = _catalog_case(method)
with patch.object(operations, helper, AsyncMock(return_value=[{"unexpected": "item"}])):
result = await GatewayOperations().execute(operation, prepare_context())
assert getattr(result, collection) == []
@pytest.mark.asyncio
@pytest.mark.parametrize("catalog_unavailable", [False, True])
async def test_tool_listing_returns_empty_result_without_dispatch_for_unavailable_catalog(catalog_unavailable):
from mcp.types import ListToolsRequest
from litellm.proxy._experimental.mcp_server import operations
allowed = AsyncMock(
return_value=[], side_effect=RuntimeError("catalog unavailable") if catalog_unavailable else None
)
upstream = AsyncMock()
with (
patch.object(operations, "_get_allowed_mcp_servers", allowed),
patch.object(operations.global_mcp_server_manager, "_get_tools_from_server", upstream),
):
result = await GatewayOperations().execute(ListToolsRequest(), prepare_context())
assert result.tools == []
allowed.assert_awaited_once()
upstream.assert_not_awaited()
@pytest.mark.asyncio
async def test_explicit_proxy_context_lists_builtin_tools_and_blocks_direct_tool_dispatch():
from mcp.types import CallToolRequest, CallToolRequestParams, ListToolsRequest
from litellm.proxy._experimental.mcp_server import operations
context = prepare_context(mcp_proxy_mode=True)
allowed = AsyncMock()
with patch.object(operations, "_get_allowed_mcp_servers", allowed):
listing = await GatewayOperations().execute(ListToolsRequest(), context)
denied = await GatewayOperations().execute(
CallToolRequest(params=CallToolRequestParams(name="catalog-tool", arguments={})), context
)
assert {tool.name for tool in listing.tools} == {"search_tools", "get_tool_schema", "call_tool"}
assert denied.is_error is True
assert "unavailable on /mcp/proxy" in denied.content[0].text
allowed.assert_not_awaited()

13
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -1,3 +1,4 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
import asyncio
import inspect
import json
@ -1253,6 +1254,7 @@ class TestListToolsRestAPI:
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
client_ip=None,
):
captured["called"] = True
captured["server"] = server
@ -1338,6 +1340,7 @@ class TestListToolsRestAPI:
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
client_ip=None,
):
captured["user_api_key_auth"] = user_api_key_auth
return ["tool-1"]
@ -1891,6 +1894,7 @@ class TestListToolsRestAPI:
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
client_ip=None,
):
captured["called"] = True
captured["server_arg"] = server
@ -2027,6 +2031,7 @@ class TestListToolsRestAPI:
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
client_ip=None,
):
captured["called"] = True
captured["server_arg"] = server
@ -2112,6 +2117,7 @@ class TestListToolsRestAPI:
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
client_ip=None,
):
return ["scoped-tool"]
@ -2319,6 +2325,7 @@ class TestListToolsRestAPI:
user_api_key_auth=None,
extra_headers=None,
apply_tool_filters=True,
client_ip=None,
):
captured["server"] = server
captured["auth_header"] = server_auth_header
@ -3145,10 +3152,10 @@ async def test_request_selected_tool_specific_guardrail_applies_to_virtual_execu
monkeypatch.setattr(litellm, "callbacks", [guardrail])
monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry)
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(server, "global_mcp_tool_registry", registry)
monkeypatch.setattr(server, "global_mcp_server_manager", manager)
monkeypatch.setattr(mcp_operations, "global_mcp_tool_registry", registry)
monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager)
monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager)
monkeypatch.setattr(server, "_get_allowed_mcp_servers", AsyncMock(return_value=[managed_server]))
monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[managed_server]))
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()))
monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", passthrough_request_data)
monkeypatch.setattr(proxy_server, "proxy_config", {})

52
tests/test_litellm/test_check_mcp_operation_boundary.py Normal file
View file

@ -0,0 +1,52 @@
from pathlib import Path
import pytest
from scripts.check_mcp_operation_boundary import main, violations
@pytest.mark.parametrize(
"source",
(
"from mcp.server.auth.middleware.auth_context import auth_context_var as hidden",
"from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode as mode",
"caller = legacy.get_active_auth_context()",
"owners = transport._stateful_session_owners",
"from weakref import WeakKeyDictionary",
"from litellm.proxy._experimental.mcp_server.server import get_auth_context",
),
)
def test_shared_operation_boundary_rejects_ambient_state(source):
assert violations(Path("operations.py"), source)
def test_legacy_adapter_may_resolve_context_but_policy_must_receive_it():
source = "from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode"
assert violations(Path("server.py"), source) == ()
assert violations(Path("legacy_callbacks.py"), source) == ()
assert violations(Path("operations.py"), "def execute(context):\n return context.client_ip") == ()
assert violations(Path("mcp_server_manager.py"), "def _mcp_registry_key(server):\n return server.name") == ()
def test_boundary_command_rejects_shared_state_and_accepts_explicit_context(tmp_path, monkeypatch, capsys):
import subprocess
import sys
package = tmp_path / "litellm/proxy/_experimental/mcp_server"
package.mkdir(parents=True)
module = package / "operations.py"
module.write_text("from mcp.server.auth.middleware.auth_context import auth_context_var as hidden\n")
command = [sys.executable, str(Path(__file__).resolve().parents[2] / "scripts/check_mcp_operation_boundary.py")]
monkeypatch.chdir(tmp_path)
assert main() == 1
assert "operations.py:1:" in capsys.readouterr().err
rejected = subprocess.run(command, cwd=tmp_path, capture_output=True, text=True, check=False)
assert rejected.returncode == 1
assert "operations.py:1: MCP request/session state belongs in a legacy adapter" in rejected.stderr
module.write_text("def execute(context):\n return context.client_ip\n")
assert main() == 0
assert "MCP operation boundary: passed" in capsys.readouterr().out
accepted = subprocess.run(command, cwd=tmp_path, capture_output=True, text=True, check=False)
assert accepted.returncode == 0
assert "MCP operation boundary: passed" in accepted.stdout