fix(mcp): preserve legacy behavior on SDK2 and streamline verification

This commit is contained in:
Joshua Valluru 2026-09-18 22:28:31 -07:00
parent 77cf6c2fbd
commit aea13ee03b
41 changed files with 1109 additions and 1199 deletions

View file

@ -7,14 +7,6 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "pyproject.toml"
- "uv.lock"
- "litellm/experimental_mcp_client/**"
- "litellm/proxy/_experimental/mcp_server/**"
- "litellm/types/mcp.py"
- "scripts/check_mcp_sdk_install.py"
- ".github/workflows/test-mcp-dependency-resolution.yml"
permissions:
contents: read
@ -63,28 +55,42 @@ jobs:
run: |
uv lock --check
- name: Install locked dependencies
- name: Check locked runtime installations
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --python ${{ matrix.python-version }} --extra mcp --extra proxy
for extra in core mcp proxy; do
args=()
if [ "$extra" != core ]; then args=(--extra "$extra"); fi
UV_PROJECT_ENVIRONMENT=".venv-$extra" .github/scripts/uv_sync_with_retries.sh --frozen --no-dev --no-editable --python ${{ matrix.python-version }} "${args[@]}"
uv pip check --python ".venv-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-$extra/bin/python" "${checker[@]}")
done
- name: Check locked MCP SDK installation
- name: Build the public wheel
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync python scripts/check_mcp_sdk_install.py
run: uv build --all-packages --wheel --out-dir dist/mcp-check
- name: Resolve lowest direct dependencies
- name: Check lowest direct runtime installations
if: steps.changes.outputs.decision != 'skip'
run: |
uv pip compile pyproject.toml --python-version ${{ matrix.python-version }} --extra mcp --extra proxy --resolution lowest-direct -o lowest-direct.txt
- name: Install lowest direct dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
uv venv --python ${{ matrix.python-version }} .venv-lowest
uv pip install --python .venv-lowest -r lowest-direct.txt -e .
- name: Check lowest-direct MCP SDK installation
if: steps.changes.outputs.decision != 'skip'
run: |
.venv-lowest/bin/python scripts/check_mcp_sdk_install.py
wheel=$(realpath dist/mcp-check/litellm-[0-9]*.whl)
for extra in core mcp proxy; do
args=()
if [ "$extra" != core ]; then args=(--extra "$extra"); fi
uv pip compile pyproject.toml --no-sources --find-links dist/mcp-check "${args[@]}" --python-version ${{ matrix.python-version }} --resolution lowest-direct -o "lowest-$extra.txt"
uv venv --python ${{ matrix.python-version }} ".venv-lowest-$extra"
uv pip sync --find-links dist/mcp-check --python ".venv-lowest-$extra" "lowest-$extra.txt"
uv pip install --python ".venv-lowest-$extra" --no-deps "$wheel"
uv pip check --python ".venv-lowest-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-lowest-$extra/bin/python" "${checker[@]}")
done

View file

@ -57,6 +57,13 @@ jobs:
uv lock --check
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
- name: Install the unchanged SDK1 peer
if: steps.changes.outputs.decision != 'skip'
run: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "MCP_TEST_PEER_PYTHON=$GITHUB_WORKSPACE/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
- name: Run MCP tests
if: steps.changes.outputs.decision != 'skip'
run: |

View file

@ -1,6 +1,15 @@
# LiteLLM MCP Client
LiteLLM MCP Client is a client that allows you to use MCP tools with LiteLLM.
LiteLLM MCP Client allows you to use MCP tools with LiteLLM
## MCP Python SDK compatibility
The `mcp` and `proxy` extras require MCP Python SDK 2.2 or newer within the 2.x release line. Installing core LiteLLM without these extras does not require MCP
Existing MCP SDK1 clients can continue connecting to the gateway over the supported legacy MCP protocols. The client and gateway can use different SDK versions in separate Python environments. Modern protocol advertisement remains disabled during the Phase 0 upgrade
Code sharing the gateway's Python environment must support SDK2. Its Python API has breaking changes, including renamed imports and snake_case model attributes such as `input_schema`, `is_error`, and `structured_content`. This also applies to callers consuming SDK objects returned by LiteLLM's experimental MCP client. MCP JSON fields retain their protocol spelling, such as `inputSchema` and `isError`
Upgrade SDK1-dependent libraries before installing them alongside `litellm[mcp]` or `litellm[proxy]`, or keep those clients in a separate environment and connect over the network. For example, `langchain-mcp-adapters==0.2.1` uses SDK1 Python APIs and is tested as a separate legacy client, not as a shared SDK2 dependency
See the official [SDK migration guide](https://py.sdk.modelcontextprotocol.io/migration/) for Python API changes

View file

@ -14,6 +14,8 @@ from types import MappingProxyType
from typing import Any, Final, TypeAlias, TypeVar
import httpx2
from httpx2._client import UseClientDefault
from httpx2._types import AuthTypes
from mcp import ClientSession, MCPError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
@ -147,6 +149,23 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
TSessionResult = TypeVar("TSessionResult")
class _MCPHTTPClient(httpx2.AsyncClient):
async def send(
self,
request: httpx2.Request,
*,
stream: bool = False,
auth: AuthTypes | UseClientDefault | None = httpx2.USE_CLIENT_DEFAULT,
follow_redirects: bool | UseClientDefault = httpx2.USE_CLIENT_DEFAULT,
) -> httpx2.Response:
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
# Check after the auth flow completes so a refreshable 401 can still be retried.
if request.method == "POST" and response.is_error:
await response.aclose()
response.raise_for_status()
return response
class MCPSigV4Auth(httpx2.Auth):
"""
httpx2 Auth class that signs each request with AWS SigV4.
@ -448,7 +467,7 @@ class MCPClient:
async def receive_message(
message: ServerNotification | Exception,
) -> None:
if not isinstance(message, (ValueError, httpx2.RequestError, OSError)):
if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)):
return
if not stream_error.done():
stream_error.set_result(message)
@ -592,7 +611,9 @@ class MCPClient:
headers.update(injected or {})
return _strip_header_whitespace(headers)
def _create_httpx_client_factory(self) -> Callable[..., httpx2.AsyncClient]:
def _create_httpx_client_factory(
self, *, transport: httpx2.AsyncBaseTransport | None = None
) -> Callable[..., httpx2.AsyncClient]:
"""
Create a custom httpx2 client factory that uses LiteLLM's SSL configuration.
This factory follows the same CA bundle path logic as http_handler.py:
@ -618,7 +639,8 @@ class MCPClient:
fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth
effective_auth: Final = auth if auth is not None else fallback_auth
guard: Final = credential_redirect_hook(self.server_url, self._credential_slot)
return httpx2.AsyncClient(
return _MCPHTTPClient(
transport=transport,
headers=headers,
timeout=timeout,
auth=effective_auth,

View file

@ -100,6 +100,8 @@ Usage with curl::
http://localhost:4000/mcp/atlassian_mcp
"""
from __future__ import annotations
import asyncio
import base64
import io
@ -109,7 +111,7 @@ from collections.abc import AsyncIterator, Callable, Mapping
from http.cookies import CookieError, SimpleCookie
from itertools import islice
from types import MappingProxyType
from typing import Final
from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl, quote, quote_plus, unquote_plus, urlencode
import httpx
@ -120,7 +122,9 @@ from starlette.types import Message, Send
from litellm.litellm_core_utils.secret_redaction import REDACTED, redact_string
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
# Header the client sends to opt into debug mode
MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug"
@ -151,6 +155,8 @@ class MCPAuthDiagnostics:
self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),)
def resolution(self) -> str:
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
match self._outcomes:
case ():
return AuthResolution.unresolved.value
@ -160,6 +166,8 @@ class MCPAuthDiagnostics:
return AuthResolution.multiple.value
def headers(self) -> Mapping[str, str]:
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
if len(self._outcomes) <= 1:
return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()})
return MappingProxyType(
@ -373,6 +381,8 @@ class MCPDebug:
server_url: str | None = None
server_auth_type: str | None = None
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
auth_resolution: Final = AuthResolution.unresolved.value
for server_name in mcp_servers or []:

View file

@ -166,7 +166,8 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
)
if exc.error.code == -32700 or exc.error.message.startswith("Failed to parse"):
return (
"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. "
f"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response "
f"(JSON-RPC code {exc.error.code}). "
"Check the MCP endpoint URL and the server's protocol implementation."
)
if exc.error.code == -32000 and exc.error.message == "Connection closed":
@ -1652,7 +1653,7 @@ if MCP_AVAILABLE:
"message": f"Timed out listing tools after {listing_deadline} seconds. "
"The MCP server may be responding slowly or paginating excessively.",
}
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
model_dumped_tools: Final[list[dict]] = [tool.model_dump(by_alias=True) for tool in list_tools_result]
return {
"tools": model_dumped_tools,
"error": None,

View file

@ -374,7 +374,7 @@ def _convert_single_content(
# ToolResultContent → proper OpenAI tool-role message.
# Marked so the message-level converter can emit it as a
# separate ``{"role": "tool", ...}`` message.
tool_result_use_id: Final = getattr(content, "toolUseId", "")
tool_result_use_id: Final = getattr(content, "tool_use_id", "")
nested_content: Final[Sequence[ContentBlock]] = getattr(content, "content", [])
if isinstance(nested_content, list):
text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"]
@ -537,7 +537,7 @@ def _extract_tool_results(
results: Final = []
for item in items:
if getattr(item, "type", None) == "tool_result":
tool_use_id = getattr(item, "toolUseId", "")
tool_use_id = getattr(item, "tool_use_id", "")
# Extract text from nested content
nested_content: Sequence[ContentBlock] = getattr(item, "content", [])
if isinstance(nested_content, list):

View file

@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol
import httpx
from fastapi import FastAPI, HTTPException
from pydantic import AnyUrl, ConfigDict, TypeAdapter, ValidationError
from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter, ValidationError
from starlette.requests import Request as StarletteRequest
from starlette.responses import JSONResponse
from starlette.types import Message, Receive, Scope, Send
@ -541,7 +541,7 @@ if MCP_AVAILABLE:
Object returned by the /tools/list REST API route.
"""
mcp_info: MCPInfo | None = None
mcp_info: MCPInfo | None = Field(default=None, alias="mcp_info")
model_config = ConfigDict(arbitrary_types_allowed=True)
def _gateway_create_initialization_options(
@ -910,7 +910,7 @@ if MCP_AVAILABLE:
if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta):
return None
host_token: Final = getattr(host_ctx.meta, "progress_token", None)
host_token: Final = host_ctx.meta.get("progress_token")
if host_token is None or not (hasattr(host_ctx, "session") and host_ctx.session):
return None
host_session: Final = host_ctx.session
@ -3790,7 +3790,7 @@ if MCP_AVAILABLE:
def _extract_initialize_client_info(body: bytes) -> Implementation | None:
try:
return InitializeRequest.model_validate_json(body).params.clientInfo
return InitializeRequest.model_validate_json(body, by_name=False).params.client_info
except ValidationError:
return None

View file

@ -7,7 +7,7 @@ while preserving the existing public import path.
from collections.abc import Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final, Optional, cast
from typing import TYPE_CHECKING, Final, Optional
from fastapi import HTTPException
@ -45,15 +45,6 @@ def _serialize_mcp_content_item(item: object) -> dict[str, object]:
return {"type": "text", "text": str(item)}
def _coerce_pair_list_source(source: object) -> object:
if not isinstance(source, list):
return source
try:
return dict(cast("Sequence[tuple[str, object]]", source)) # pyright: ignore[reportUnknownArgumentType] # response_obj arrives untyped; dict() rejects non-pair shapes
except (TypeError, ValueError):
return source
def _source_field(source: object, key: str, snake_key: str) -> object:
if isinstance(source, dict):
for candidate in (key, snake_key):
@ -526,10 +517,9 @@ class _CiscoAIDefenseMcpMixin:
content: Sequence[object],
source: object = None,
) -> dict[str, object]:
source_map: Final[object] = _coerce_pair_list_source(source)
result: Final[dict[str, object]] = {"content": [_serialize_mcp_content_item(item) for item in content]}
for key, snake_key in (("structuredContent", "structured_content"), ("isError", "is_error")):
value = _source_field(source_map, key, snake_key)
value = _source_field(source, key, snake_key)
if value is not None and (key != "isError" or isinstance(value, bool)):
result[key] = value
return result

View file

@ -1,3 +1,4 @@
import argparse
import importlib
import importlib.metadata
import sys
@ -20,7 +21,10 @@ def _version_tuple(distribution: str) -> tuple[int, ...]:
def main() -> int:
for module_name in IMPORTED_MODULES:
parser: Final = argparse.ArgumentParser()
parser.add_argument("--extra", choices=("mcp", "proxy"), default="proxy")
extra: Final = parser.parse_args().extra
for module_name in IMPORTED_MODULES if extra == "proxy" else IMPORTED_MODULES[:3]:
try:
importlib.import_module(module_name)
except Exception as exc:
@ -39,22 +43,23 @@ def main() -> int:
sys.stderr.write(f"HANDSHAKE_PROTOCOL_VERSIONS missing {required}\n")
return 1
scope: Final = {
"type": "http",
"method": "POST",
"path": "/mcp",
"headers": [(b"mcp-protocol-version", b"2026-07-28")],
}
mcp_server: Final = sys.modules["litellm.proxy._experimental.mcp_server.server"]
if mcp_server.unsupported_protocol_version(scope) != "2026-07-28":
sys.stderr.write("unsupported_protocol_version accepted a modern-only version\n")
return 1
if (
mcp_server.unsupported_protocol_version(dict(scope, headers=[(b"mcp-protocol-version", b"2025-06-18")]))
is not None
):
sys.stderr.write("unsupported_protocol_version rejected a handshake version\n")
return 1
if extra == "proxy":
scope: Final = {
"type": "http",
"method": "POST",
"path": "/mcp",
"headers": [(b"mcp-protocol-version", b"2026-07-28")],
}
mcp_server: Final = sys.modules["litellm.proxy._experimental.mcp_server.server"]
if mcp_server.unsupported_protocol_version(scope) != "2026-07-28":
sys.stderr.write("unsupported_protocol_version accepted a modern-only version\n")
return 1
if (
mcp_server.unsupported_protocol_version(dict(scope, headers=[(b"mcp-protocol-version", b"2025-06-18")]))
is not None
):
sys.stderr.write("unsupported_protocol_version rejected a handshake version\n")
return 1
sys.stdout.write(
"python {} mcp {} httpx2 {} pydantic {} litellm {}\n".format(

View file

@ -74,3 +74,14 @@ def pytest_collection_modifyitems(config, items):
# Reorder the items list
items[:] = custom_logger_tests + other_tests
@pytest.fixture
def config_only_mcp_manager_factory():
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
class ConfigOnlyManager(MCPServerManager):
def initialize_tool_name_to_mcp_server_name_mapping(self):
return None
return ConfigOnlyManager

View file

@ -51,6 +51,21 @@ def request_headers(ctx: Context) -> dict[str, str]:
}
@mcp.prompt()
def greeting(name: str) -> str:
return f"Hello, {name}"
@mcp.resource("memo://status")
def status() -> str:
return "ready"
@mcp.resource("memo://greeting/{name}")
def greeting_resource(name: str) -> str:
return f"Hello, {name}"
def main() -> None:
args = _parse_args()
transport = (args.transport or "stdio").lower()

View file

@ -1,6 +1,7 @@
import logging
import os
import pytest
from mcp.types import Tool as MCPTool
from typing import List, Any, cast
from unittest.mock import AsyncMock, patch
@ -371,48 +372,32 @@ async def test_mcp_allowed_tools_filtering():
# Mock MCP tools returned from the server (simulating all available tools)
mock_mcp_tools_from_server = [
# Mock MCP tool object with name attribute
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "search_tiktoken_documentation",
"description": "Search tiktoken documentation",
"inputSchema": {
"type": "object",
"properties": {"query": {"type": "string"}},
},
},
)(),
type(
"MCPTool",
(),
{
}, by_name=False),
MCPTool.model_validate({
"name": "fetch_tiktoken_documentation",
"description": "Fetch tiktoken documentation",
"inputSchema": {
"type": "object",
"properties": {"path": {"type": "string"}},
},
},
)(),
type(
"MCPTool",
(),
{
}, by_name=False),
MCPTool.model_validate({
"name": "list_tiktoken_functions",
"description": "List tiktoken functions",
"inputSchema": {"type": "object", "properties": {}},
},
)(),
type(
"MCPTool",
(),
{
}, by_name=False),
MCPTool.model_validate({
"name": "get_tiktoken_examples",
"description": "Get tiktoken examples",
"inputSchema": {"type": "object", "properties": {}},
},
)(),
}, by_name=False),
]
allowed_mcp_servers = ["gitmcp"]
@ -491,10 +476,7 @@ async def test_mcp_allowed_tools_filtering():
# Test Case 3: Test deduplication of duplicate tools
mock_mcp_tools_with_duplicates = [
# First instance of duplicate tool
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "GitMCP-fetch_litellm_documentation",
"description": "Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.",
"inputSchema": {
@ -502,13 +484,9 @@ async def test_mcp_allowed_tools_filtering():
"properties": {},
"additionalProperties": False,
},
},
)(),
}, by_name=False),
# Second instance of duplicate tool (should be filtered out)
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "GitMCP-fetch_litellm_documentation",
"description": "Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.",
"inputSchema": {
@ -516,13 +494,9 @@ async def test_mcp_allowed_tools_filtering():
"properties": {},
"additionalProperties": False,
},
},
)(),
}, by_name=False),
# Other unique tools
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "GitMCP-search_litellm_documentation",
"description": "Semantically search within the fetched documentation from GitHub repository: BerriAI/litellm. Useful for specific queries.",
"inputSchema": {
@ -531,8 +505,7 @@ async def test_mcp_allowed_tools_filtering():
"required": ["query"],
"additionalProperties": False,
},
},
)(),
}, by_name=False),
]
mcp_tool_config_with_duplicates = [
@ -680,10 +653,7 @@ async def test_streaming_mcp_events_validation():
# Mock MCP tools that would be returned from the manager
mock_mcp_tools = [
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "search_repo",
"description": "Search BerriAI/litellm repository for information",
"inputSchema": {
@ -693,12 +663,8 @@ async def test_streaming_mcp_events_validation():
},
"required": ["query"],
},
},
)(),
type(
"MCPTool",
(),
{
}, by_name=False),
MCPTool.model_validate({
"name": "get_repo_info",
"description": "Get repository information",
"inputSchema": {
@ -711,8 +677,7 @@ async def test_streaming_mcp_events_validation():
},
"required": ["repo_name"],
},
},
)(),
}, by_name=False),
]
# Build fake streaming chunks that the inner aresponses() call would yield
@ -920,10 +885,7 @@ async def test_streaming_responses_api_with_mcp_tools(
# Mock MCP tools that would be returned from the manager
mock_mcp_tools = [
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "search_repo",
"description": "Search BerriAI/litellm repository for information",
"inputSchema": {
@ -933,8 +895,7 @@ async def test_streaming_responses_api_with_mcp_tools(
},
"required": ["query"],
},
},
)()
}, by_name=False)
]
# Only mock the MCP-specific operations, let LLM responses be real
@ -1263,10 +1224,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e():
# Mock MCP tools that would be returned from the manager
mock_mcp_tools = [
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "search_docs",
"description": "Search documentation for information",
"inputSchema": {
@ -1276,12 +1234,8 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e():
},
"required": ["query"],
},
},
)(),
type(
"MCPTool",
(),
{
}, by_name=False),
MCPTool.model_validate({
"name": "get_file_content",
"description": "Get content of a specific file",
"inputSchema": {
@ -1291,8 +1245,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e():
},
"required": ["file_path"],
},
},
)(),
}, by_name=False),
]
# Track all calls to the underlying LLM to detect duplicates
@ -1499,10 +1452,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency(
from unittest.mock import AsyncMock, patch
mock_mcp_tools = [
type(
"MCPTool",
(),
{
MCPTool.model_validate({
"name": "get_weather",
"description": "Get weather for a city",
"inputSchema": {
@ -1512,8 +1462,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency(
},
"required": ["city"],
},
},
)()
}, by_name=False)
]
with caplog.at_level(logging.ERROR):

View file

@ -44,14 +44,14 @@ async def test_mcp_server_works_without_config_auth_value():
@pytest.mark.parametrize("token_key", ["authentication_token", "auth_value"])
async def test_mcp_server_config_auth_value_header_used(token_key):
async def test_mcp_server_config_auth_value_header_used(token_key, config_only_mcp_manager_factory):
"""Ensure the configured auth token is emitted as the upstream Authorization header.
The token is resolved through the v2 credential resolver and rides on the client's
httpx.Auth, so assert the header it writes onto the request rather than the (now
credential-free) _get_auth_headers() dict.
"""
import httpx
import httpx2
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import (
StaticHeaderAuth,
@ -66,13 +66,13 @@ async def test_mcp_server_config_auth_value_header_used(token_key):
}
}
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
await manager.load_servers_from_config(config)
server = next(iter(manager.config_mcp_servers.values()))
client = await manager._create_mcp_client(server)
assert isinstance(client._resolved_auth, StaticHeaderAuth)
emitted = next(client._resolved_auth.auth_flow(httpx.Request("POST", server.url)))
emitted = next(client._resolved_auth.auth_flow(httpx2.Request("POST", server.url)))
assert emitted.headers["Authorization"] == "Bearer example_token"
assert client.auth_type == MCPAuth.bearer_token

View file

@ -169,7 +169,7 @@ class TestMCPClientUnitTests:
MCPTool(
name="test_tool",
description="Test tool",
input_schema={
inputSchema={
"type": "object",
"properties": {"arg1": {"type": "string"}},
"required": ["arg1"],
@ -207,12 +207,12 @@ class TestMCPClientUnitTests:
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
first_page_tools = [
MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", input_schema={}) for idx in range(100)
MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", inputSchema={}) for idx in range(100)
]
second_page_tool = MCPTool(
name="tool_100",
description="Tool 100",
input_schema={},
inputSchema={},
)
mock_session_instance.list_tools.side_effect = [
ListToolsResult(tools=first_page_tools, nextCursor="page-2"),
@ -249,7 +249,7 @@ class TestMCPClientUnitTests:
mock_session_instance.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="Tool 0", input_schema={})],
tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})],
nextCursor="page-2",
),
RuntimeError("transient upstream failure"),

View file

@ -84,7 +84,7 @@ async def test_mcp_cost_tracking():
# Create a mock tool call result
litellm.logging_callback_manager._reset_all_callbacks()
mock_result = CallToolResult(
content=[TextContent(type="text", text="Test response")], is_error=False
content=[TextContent(type="text", text="Test response")], isError=False
)
# Create a mock MCPClient
@ -95,7 +95,7 @@ async def test_mcp_cost_tracking():
MCPTool(
name="add_tools",
description="Test tool",
input_schema={
inputSchema={
"type": "object",
"properties": {"test": {"type": "string"}},
},
@ -209,7 +209,7 @@ async def test_mcp_cost_tracking_per_tool():
# Create a mock tool call result
litellm.logging_callback_manager._reset_all_callbacks()
mock_result = CallToolResult(
content=[TextContent(type="text", text="Test response")], is_error=False
content=[TextContent(type="text", text="Test response")], isError=False
)
# Create a mock MCPClient
@ -220,7 +220,7 @@ async def test_mcp_cost_tracking_per_tool():
MCPTool(
name="expensive_tool",
description="Expensive tool",
input_schema={
inputSchema={
"type": "object",
"properties": {"data": {"type": "string"}},
},
@ -228,7 +228,7 @@ async def test_mcp_cost_tracking_per_tool():
MCPTool(
name="cheap_tool",
description="Cheap tool",
input_schema={
inputSchema={
"type": "object",
"properties": {"data": {"type": "string"}},
},
@ -390,7 +390,7 @@ async def test_mcp_tool_call_hook():
# Create a mock tool call result
litellm.logging_callback_manager._reset_all_callbacks()
mock_result = CallToolResult(
content=[TextContent(type="text", text="Test response")], is_error=False
content=[TextContent(type="text", text="Test response")], isError=False
)
# Create a mock MCPClient
@ -401,7 +401,7 @@ async def test_mcp_tool_call_hook():
MCPTool(
name="add_tools",
description="Test tool",
input_schema={
inputSchema={
"type": "object",
"properties": {"test": {"type": "string"}},
},

View file

@ -44,7 +44,7 @@ async def test_mcp_server_manager_https_server():
MCPTool(
name="gmail_send_email",
description="Send an email via Gmail",
input_schema={
inputSchema={
"type": "object",
"properties": {
"body": {"type": "string"},
@ -58,7 +58,7 @@ async def test_mcp_server_manager_https_server():
mock_result = CallToolResult(
content=[TextContent(type="text", text="Email sent successfully")],
is_error=False,
isError=False,
)
# Create a mock MCPClient
@ -143,7 +143,7 @@ async def test_mcp_http_transport_list_tools_mock():
MCPTool(
name="gmail_send_email",
description="Send an email via Gmail",
input_schema={
inputSchema={
"type": "object",
"properties": {
"to": {"type": "string"},
@ -156,7 +156,7 @@ async def test_mcp_http_transport_list_tools_mock():
MCPTool(
name="calendar_create_event",
description="Create a calendar event",
input_schema={
inputSchema={
"type": "object",
"properties": {
"title": {"type": "string"},
@ -242,7 +242,7 @@ async def test_mcp_http_transport_call_tool_mock():
content=[
TextContent(type="text", text="Email sent successfully to test@example.com")
],
is_error=False,
isError=False,
)
# Create a mock MCPClient that returns our test result
@ -308,7 +308,7 @@ async def test_mcp_http_transport_call_tool_error_mock():
# Mock tool call error result
mock_error_result = CallToolResult(
content=[TextContent(type="text", text="Error: Invalid email address")],
is_error=True,
isError=True,
)
# Create a mock MCPClient that returns our test error result
@ -361,11 +361,11 @@ async def test_mcp_http_transport_call_tool_error_mock():
@pytest.mark.asyncio
async def test_mcp_http_transport_tool_not_found():
async def test_mcp_http_transport_tool_not_found(config_only_mcp_manager_factory):
"""Test calling a tool that doesn't exist"""
# Create a fresh manager for testing
test_manager = MCPServerManager()
test_manager = config_only_mcp_manager_factory()
# Load server config
await test_manager.load_servers_from_config(
@ -892,8 +892,8 @@ async def test_get_tools_from_mcp_servers():
transport=MCPTransport.http,
access_groups=["group-a"],
)
mock_tool_1 = MCPTool(name="tool1", description="test tool 1", input_schema={})
mock_tool_2 = MCPTool(name="tool2", description="test tool 2", input_schema={})
mock_tool_1 = MCPTool(name="tool1", description="test tool 1", inputSchema={})
mock_tool_2 = MCPTool(name="tool2", description="test tool 2", inputSchema={})
# Test Case 1: With specific MCP servers
try:
@ -1058,14 +1058,14 @@ async def test_list_tools_only_returns_allowed_servers(monkeypatch):
MCPTool(
name="send_email",
description="Send an email via Server A",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
mock_tools_b = [
MCPTool(
name="create_event",
description="Create an event via Server B",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -1097,11 +1097,11 @@ async def test_list_tools_only_returns_allowed_servers(monkeypatch):
@pytest.mark.asyncio
async def test_mcp_server_manager_access_groups_from_config():
async def test_mcp_server_manager_access_groups_from_config(config_only_mcp_manager_factory):
"""
Test that access_groups are loaded from config and can be resolved.
"""
test_manager = MCPServerManager()
test_manager = config_only_mcp_manager_factory()
await test_manager.load_servers_from_config(
{
"config_server": {
@ -1168,7 +1168,7 @@ async def test_mcp_server_manager_access_groups_from_config():
@pytest.mark.asyncio
async def test_mcp_server_manager_config_integration_with_database():
async def test_mcp_server_manager_config_integration_with_database(config_only_mcp_manager_factory):
"""
Test that config-based servers properly integrate with database servers,
specifically testing access_groups and description fields.
@ -1176,7 +1176,7 @@ async def test_mcp_server_manager_config_integration_with_database():
import datetime
from litellm.proxy._types import LiteLLM_MCPServerTable
test_manager = MCPServerManager()
test_manager = config_only_mcp_manager_factory()
# Test 1: Load config with access_groups and description
await test_manager.load_servers_from_config(
@ -1365,7 +1365,7 @@ async def test_mcp_server_manager_alias_tool_prefixing():
MCPTool(
name="send_email",
description="Send an email",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -1425,7 +1425,7 @@ async def test_mcp_server_manager_server_name_tool_prefixing():
MCPTool(
name="send_email",
description="Send an email",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -1485,7 +1485,7 @@ async def test_mcp_server_manager_server_id_tool_prefixing():
MCPTool(
name="send_email",
description="Send an email",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -1904,12 +1904,12 @@ def test_create_tool_response_objects():
MCPTool(
name="send_email",
description="Send an email",
input_schema={"type": "object", "properties": {"to": {"type": "string"}}},
inputSchema={"type": "object", "properties": {"to": {"type": "string"}}},
),
MCPTool(
name="create_event",
description="Create a calendar event",
input_schema={"type": "object", "properties": {"title": {"type": "string"}}},
inputSchema={"type": "object", "properties": {"title": {"type": "string"}}},
),
]
@ -1962,7 +1962,7 @@ async def test_get_tools_for_single_server():
MCPTool(
name="send_email",
description="Send an email",
input_schema={"type": "object", "properties": {"to": {"type": "string"}}},
inputSchema={"type": "object", "properties": {"to": {"type": "string"}}},
)
]
@ -2016,12 +2016,12 @@ async def test_get_tools_for_single_server_applies_disallowed_tools_without_allo
MCPTool(
name="send_email",
description="Send an email",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="read_email",
description="Read an email",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
]
@ -2069,7 +2069,7 @@ async def test_rest_listing_hides_key_grants_dispatch_would_refuse():
MCPTool(
name="read_wiki_contents",
description="Read a wiki",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
]
@ -2430,22 +2430,22 @@ async def test_filter_tools_by_allowed_tools_integration():
MCPTool(
name="allowed_tool_1",
description="This tool should be allowed",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="allowed_tool_2",
description="This tool should also be allowed",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="blocked_tool_1",
description="This tool should be blocked",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="blocked_tool_2",
description="This tool should also be blocked",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
]
@ -2545,22 +2545,22 @@ async def test_filter_tools_by_disallowed_tools_integration():
MCPTool(
name="safe_tool_1",
description="This tool should be allowed",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="safe_tool_2",
description="This tool should also be allowed",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="dangerous_tool_1",
description="This tool should be blocked",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="dangerous_tool_2",
description="This tool should also be blocked",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
]
@ -2659,12 +2659,12 @@ async def test_filter_tools_no_restrictions_integration():
MCPTool(
name="tool_1",
description="Tool 1",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="tool_2",
description="Tool 2",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
]
@ -2811,7 +2811,7 @@ async def test_mcp_access_group_permission_intersection_integration():
@pytest.mark.asyncio
async def test_mcp_server_manager_with_access_groups_integration():
async def test_mcp_server_manager_with_access_groups_integration(config_only_mcp_manager_factory):
"""Integration test for MCPServerManager with access group filtering"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
@ -2820,7 +2820,7 @@ async def test_mcp_server_manager_with_access_groups_integration():
from litellm.proxy._types import UserAPIKeyAuth
# Create a test manager
test_manager = MCPServerManager()
test_manager = config_only_mcp_manager_factory()
# Load servers with access groups
await test_manager.load_servers_from_config(
@ -2863,13 +2863,13 @@ async def test_mcp_server_manager_with_access_groups_integration():
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_returns_registry_for_admin():
async def test_get_allowed_mcp_servers_returns_registry_for_admin(config_only_mcp_manager_factory):
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
test_manager = MCPServerManager()
test_manager = config_only_mcp_manager_factory()
await test_manager.load_servers_from_config(
{
"alpha_server": {
@ -2898,14 +2898,14 @@ async def test_get_allowed_mcp_servers_returns_registry_for_admin():
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permissions():
async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permissions(config_only_mcp_manager_factory):
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
MCPServerAccess,
)
test_manager = MCPServerManager()
test_manager = config_only_mcp_manager_factory()
await test_manager.load_servers_from_config(
{
"alpha_server": {

View file

@ -15,6 +15,7 @@ from datetime import datetime
from pathlib import Path
import httpx
import httpx2
import pytest
import uvicorn
import yaml
@ -36,6 +37,7 @@ from litellm.proxy.proxy_server import (
CONFIG_TEMPLATE_PATH = Path("tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml")
MCP_SERVER_SCRIPT = Path("tests/mcp_tests/mcp_server.py")
MCP_PEER_PYTHON = os.environ.get("MCP_TEST_PEER_PYTHON", sys.executable)
PROJECT_ROOT = Path(__file__).resolve().parents[2]
PROXY_START_TIMEOUT = 30
@ -125,7 +127,7 @@ def _math_http_server(offset: int) -> typing.Iterator[str]:
with tempfile.TemporaryFile() as server_log:
process = subprocess.Popen(
[sys.executable, str(MCP_SERVER_SCRIPT), "--transport", "http", "--host", host, "--port", str(port)],
[MCP_PEER_PYTHON, str(MCP_SERVER_SCRIPT), "--transport", "http", "--host", host, "--port", str(port)],
cwd=str(PROJECT_ROOT),
stdout=server_log,
stderr=subprocess.STDOUT,
@ -175,7 +177,7 @@ def _proxy_server(
config_dir = tmp_path_factory.mktemp("mcp_e2e")
config_path = config_dir / "config.yaml"
config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text())
config["mcp_servers"]["math_stdio"]["command"] = sys.executable
config["mcp_servers"]["math_stdio"]["command"] = MCP_PEER_PYTHON
config["mcp_servers"]["math_streamable_http"]["url"] = f"{math_streamable_http_server}/mcp"
config["mcp_servers"]["math_restricted"]["url"] = f"{math_restricted_server}/mcp"
config["general_settings"]["custom_auth"] = f"{__name__}.authorize_proxy_key"
@ -202,17 +204,90 @@ def proxy_server_url(_proxy_server: ProxyRig, setup_and_teardown: None) -> str:
return _proxy_server.url
@asynccontextmanager
async def _http_streams(url: str, headers: dict[str, str]):
async with httpx2.AsyncClient(headers=headers) as http_client:
async with streamable_http_client(url, http_client=http_client) as streams:
yield streams
@pytest.mark.asyncio
async def test_unchanged_sdk1_langchain_peer_can_list_and_call(proxy_server_url: str) -> None:
script = """
import asyncio, json, sys
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from langchain_mcp_adapters.tools import load_mcp_tools
async def main():
async with streamablehttp_client(sys.argv[1] + '/mcp', headers={'Authorization': 'Bearer sk-1234'}) as (read, write, _):
async with ClientSession(read, write) as session:
await session.initialize()
tools = await load_mcp_tools(session)
results = {}
for name in ('math_stdio-add', 'math_streamable_http-add'):
tool = next(tool for tool in tools if tool.name == name)
results[name] = await tool.ainvoke({'a': 3, 'b': 4})
print(json.dumps(results))
asyncio.run(main())
"""
completed = await asyncio.to_thread(
subprocess.run, [MCP_PEER_PYTHON, "-c", script, proxy_server_url],
capture_output=True, text=True, timeout=30, check=True,
)
results = json.loads(completed.stdout)
assert [(item["type"], item["text"]) for item in results["math_stdio-add"]] == [("text", "7")]
assert [(item["type"], item["text"]) for item in results["math_streamable_http-add"]] == [("text", "107")]
@pytest.mark.parametrize("requested", ["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25", "2026-07-28"])
def test_initialize_keeps_legacy_negotiation(proxy_server_url: str, requested: str) -> None:
response = httpx.post(
proxy_server_url + "/mcp",
headers={"Authorization": PROXY_AUTHORIZATION_HEADER, "Accept": "application/json, text/event-stream"},
json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {
"protocolVersion": requested, "capabilities": {}, "clientInfo": {"name": "legacy-test", "version": "1"},
}},
timeout=10,
)
assert response.status_code == 200
result = _rpc_result(response)
assert result["protocolVersion"] == ("2025-11-25" if requested == "2026-07-28" else requested)
@pytest.mark.asyncio
async def test_legacy_prompts_and_resources_round_trip(proxy_server_url: str) -> None:
async with _http_streams(
proxy_server_url + "/mcp",
{"Authorization": PROXY_AUTHORIZATION_HEADER, "x-mcp-servers": "math_streamable_http"},
) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
prompts = await session.list_prompts()
greeting = next(prompt for prompt in prompts.prompts if prompt.name.endswith("greeting"))
prompt = await session.get_prompt(greeting.name, {"name": "Ada"})
assert prompt.messages[0].content.text == "Hello, Ada"
resources = await session.list_resources()
status = next(resource for resource in resources.resources if resource.name.endswith("status"))
contents = await session.read_resource(status.uri)
assert contents.contents[0].text == "ready"
templates = await session.list_resource_templates()
greeting_template = next(template for template in templates.resource_templates if "greeting" in template.name)
contents = await session.read_resource(greeting_template.uri_template.replace("{name}", "Ada"))
assert contents.contents[0].text == "Hello, Ada"
class TestProxyMcpSimpleConnections:
@pytest.mark.asyncio
async def test_proxy_mcp_stdio_roundtrip(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with streamable_http_client(
async with _http_streams(
url=f"{proxy_server_url}/mcp",
headers={
"Authorization": PROXY_AUTHORIZATION_HEADER,
"x-mcp-servers": "math_stdio",
},
) as (read, write, _get_session_id):
) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
@ -227,13 +302,13 @@ class TestProxyMcpSimpleConnections:
@pytest.mark.asyncio
async def test_proxy_mcp_streamable_http_roundtrip(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with streamable_http_client(
async with _http_streams(
url=f"{proxy_server_url}/mcp",
headers={
"Authorization": PROXY_AUTHORIZATION_HEADER,
"x-mcp-servers": "math_streamable_http",
},
) as (read, write, _get_session_id):
) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
@ -248,10 +323,10 @@ class TestProxyMcpSimpleConnections:
@pytest.mark.asyncio
async def test_proxy_mcp_lists_all_servers_without_header(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with streamable_http_client(
async with _http_streams(
url=f"{proxy_server_url}/mcp",
headers={"Authorization": PROXY_AUTHORIZATION_HEADER},
) as (read, write, _get_session_id):
) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
@ -296,16 +371,16 @@ class TestProxyMcpStatelessBehavior:
"""Two independent clients connect and operate without sharing session state."""
async with asyncio.timeout(30):
# --- Client A: connect, initialize, call tool ---
async with streamable_http_client(
async with _http_streams(
url=f"{proxy_server_url}/mcp",
headers={
"Authorization": PROXY_AUTHORIZATION_HEADER,
"x-mcp-servers": "math_stdio",
},
) as (read_a, write_a, _get_sid_a):
) as (read_a, write_a):
async with ClientSession(read_a, write_a) as session_a:
await session_a.initialize()
result_a = await session_a.call_tool("add", arguments={"a": 10, "b": 20})
result_a = await session_a.call_tool("math_stdio-add", arguments={"a": 10, "b": 20})
assert result_a.content
text_a = getattr(result_a.content[0], "text", None)
assert text_a == "30"
@ -316,18 +391,18 @@ class TestProxyMcpStatelessBehavior:
await asyncio.sleep(0.5)
# --- Client B: completely independent connection ---
async with streamable_http_client(
async with _http_streams(
url=f"{proxy_server_url}/mcp",
headers={
"Authorization": PROXY_AUTHORIZATION_HEADER,
"x-mcp-servers": "math_stdio",
},
) as (read_b, write_b, _get_sid_b):
) as (read_b, write_b):
async with ClientSession(read_b, write_b) as session_b:
await session_b.initialize()
tools = await session_b.list_tools()
assert any(t.name.endswith("add") for t in tools.tools)
result_b = await session_b.call_tool("add", arguments={"a": 100, "b": 200})
result_b = await session_b.call_tool("math_stdio-add", arguments={"a": 100, "b": 200})
assert result_b.content
text_b = getattr(result_b.content[0], "text", None)
assert text_b == "300"
@ -342,7 +417,7 @@ def _payload(result: typing.Any) -> typing.Any:
def _proxy_session(proxy_server_url: str, **extra_headers: str):
return streamable_http_client(
return _http_streams(
url=f"{proxy_server_url}/mcp/proxy",
headers={"Authorization": PROXY_AUTHORIZATION_HEADER, **extra_headers},
)
@ -356,7 +431,7 @@ class TestProxyMcpSchemaDiscoveryMode:
@pytest.mark.asyncio
async def test_initialize_and_list_expose_only_discovery_tools(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with _proxy_session(proxy_server_url) as (read, write, _sid):
async with _proxy_session(proxy_server_url) as (read, write):
async with ClientSession(read, write) as session:
init = await session.initialize()
assert init.capabilities.tools is not None
@ -369,7 +444,7 @@ class TestProxyMcpSchemaDiscoveryMode:
@pytest.mark.asyncio
async def test_search_schema_and_call_round_trip_keeps_server_identity(self, proxy_server_url: str) -> None:
async with asyncio.timeout(30):
async with _proxy_session(proxy_server_url) as (read, write, _sid):
async with _proxy_session(proxy_server_url) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
@ -408,7 +483,6 @@ class TestProxyMcpSchemaDiscoveryMode:
async with _proxy_session(proxy_server_url, **{"x-mcp-servers": "math_streamable_http"}) as (
read,
write,
_sid,
):
async with ClientSession(read, write) as session:
await session.initialize()
@ -421,7 +495,7 @@ class TestProxyMcpSchemaDiscoveryMode:
from mcp.types import METHOD_NOT_FOUND
async with asyncio.timeout(30):
async with _proxy_session(proxy_server_url) as (read, write, _sid):
async with _proxy_session(proxy_server_url) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"}))
@ -494,7 +568,7 @@ proxy_call_recorder = ProxyCallRecorder()
@asynccontextmanager
async def _scoped_session(url: str, key: str = "sk-1234", **headers: str) -> typing.AsyncIterator[ClientSession]:
async with asyncio.timeout(30):
async with _proxy_session(url, Authorization=f"Bearer {key}", **headers) as (read, write, _sid):
async with _proxy_session(url, Authorization=f"Bearer {key}", **headers) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
yield session

View file

@ -58,46 +58,46 @@ async def test_e2e_semantic_filter():
MCPTool(
name="gmail_send",
description="Send an email via Gmail",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="calendar_create",
description="Create a calendar event",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="file_upload",
description="Upload a file",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="web_search",
description="Search the web",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="slack_send",
description="Send Slack message",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="doc_read", description="Read document", input_schema={"type": "object"}
name="doc_read", description="Read document", inputSchema={"type": "object"}
),
MCPTool(
name="db_query",
description="Query database",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="api_call", description="Make API call", input_schema={"type": "object"}
name="api_call", description="Make API call", inputSchema={"type": "object"}
),
MCPTool(
name="task_create",
description="Create task",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="note_add", description="Add note", input_schema={"type": "object"}
name="note_add", description="Add note", inputSchema={"type": "object"}
),
]

View file

@ -1,11 +1,18 @@
# Create server parameters for stdio connection
import asyncio
import os
from mcp import ClientSession
from mcp.client.sse import sse_client
async def main():
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
model = ChatOpenAI(model="gpt-4o", api_key="sk-12")
async with sse_client(url="http://localhost:4000/mcp/") as (read, write):
async with ClientSession(read, write) as session:
# Initialize the connection
@ -15,15 +22,13 @@ async def main():
# Get tools
print("Loading tools")
tools = await session.list_tools()
tools = await load_mcp_tools(session)
print("Tools loaded")
print(tools)
if tools.tools:
first = tools.tools[0]
print(f"Calling tool {first.name}")
result = await session.call_tool(first.name, {})
print(result)
# # Create and run the agent
# agent = create_react_agent(model, tools)
# agent_response = await agent.ainvoke({"messages": "what's (3 + 5) x 12?"})
# Run the async function

View file

@ -1326,6 +1326,7 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None:
("application/json", b"", MCPError),
("application/json", b'{"secret":"invalid-rpc"}', MCPError),
("application/json", b'{"jsonrpc":"2.0","id":0}', MCPError),
("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"bad-schema"}}', ValidationError),
],
)
async def test_invalid_http_response_surfaces_without_waiting_for_timeout(
@ -1334,6 +1335,8 @@ async def test_invalid_http_response_surfaces_without_waiting_for_timeout(
from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message
def respond(request: httpx2.Request) -> httpx2.Response:
if expected_type is ValidationError:
return httpx2.Response(200, json={**json.loads(body), "id": json.loads(request.content)["id"]})
return httpx2.Response(200, headers={"Content-Type": content_type}, content=body)
async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client:
@ -1354,7 +1357,7 @@ async def test_invalid_http_response_surfaces_without_waiting_for_timeout(
@pytest.mark.asyncio
@pytest.mark.parametrize("status_code", [200, 401, 503])
@pytest.mark.parametrize("status_code", [200, 401, 403, 429, 503])
async def test_http_response_handler_preserves_success_and_http_errors(status_code: int) -> None:
def respond(request: httpx2.Request) -> httpx2.Response:
if request.method == "DELETE":
@ -1373,8 +1376,8 @@ async def test_http_response_handler_preserves_success_and_http_errors(status_co
)
return httpx2.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result})
async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http_client:
client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30)
client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30)
async with client._create_httpx_client_factory(transport=httpx2.MockTransport(respond))() as http_client:
operation: Final = client._execute_session_operation(
streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools()
)
@ -1382,9 +1385,33 @@ async def test_http_response_handler_preserves_success_and_http_errors(status_co
result: Final = await asyncio.wait_for(operation, timeout=3)
assert result.tools == []
else:
with pytest.raises(MCPError) as caught:
with pytest.raises(httpx2.HTTPStatusError) as caught:
await asyncio.wait_for(operation, timeout=3)
assert caught.value.error.code == INTERNAL_ERROR
assert caught.value.response.status_code == status_code
@pytest.mark.asyncio
async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth
seen = []
async def refresh(failed):
assert failed == "stale"
return "fresh"
def respond(request):
seen.append(request.headers["authorization"])
return httpx2.Response(401 if len(seen) == 1 else 200, json={"ok": True})
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ClientCredentialsConfig
auth = ClientCredentialsBearerAuth("stale", refresh, ClientCredentialsConfig())
client = MCPClient(server_url="https://example.com/mcp", resolved_auth=auth)
async with client._create_httpx_client_factory(transport=httpx2.MockTransport(respond))() as http_client:
response = await http_client.post(client.server_url, json={"method": "tools/list"})
assert response.status_code == 200
assert seen == ["Bearer stale", "Bearer fresh"]
@pytest.mark.asyncio
@ -1619,7 +1646,6 @@ async def test_sse_read_failure_is_preserved() -> None:
@pytest.mark.parametrize("mode", ["ok", "closed", "silent"])
async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None:
from mcp import ClientSession
from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message
logging_callback: Final = AsyncMock()

View file

@ -32,7 +32,7 @@ def mock_mcp_tool():
return MCPTool(
name="test_tool",
description="A test tool",
input_schema={"type": "object", "properties": {"test": {"type": "string"}}},
inputSchema={"type": "object", "properties": {"test": {"type": "string"}}},
)
@ -51,7 +51,7 @@ def mock_list_tools_result():
MCPTool(
name="test_tool",
description="A test tool",
input_schema={
inputSchema={
"type": "object",
"properties": {"test": {"type": "string"}},
},
@ -113,12 +113,12 @@ async def test_load_mcp_tools_follows_pagination(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[
MCPTool(name="tool_a", description="a", input_schema={}),
MCPTool(name="tool_b", description="b", input_schema={}),
MCPTool(name="tool_a", description="a", inputSchema={}),
MCPTool(name="tool_b", description="b", inputSchema={}),
],
nextCursor="page-2",
),
ListToolsResult(tools=[MCPTool(name="tool_c", description="c", input_schema={})]),
ListToolsResult(tools=[MCPTool(name="tool_c", description="c", inputSchema={})]),
]
result = await load_mcp_tools(mock_session, format="mcp")
assert [tool.name for tool in result] == ["tool_a", "tool_b", "tool_c"]
@ -133,14 +133,14 @@ async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch):
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_MAX_PAGES", 2)
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="0", input_schema={})],
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
nextCursor="page-2",
),
ListToolsResult(
tools=[MCPTool(name="tool_1", description="1", input_schema={})],
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
nextCursor="page-3",
),
ListToolsResult(tools=[MCPTool(name="tool_2", description="2", input_schema={})]),
ListToolsResult(tools=[MCPTool(name="tool_2", description="2", inputSchema={})]),
]
result = await list_tools_with_pagination(mock_session)
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
@ -151,11 +151,11 @@ async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch):
async def test_pagination_walk_stops_on_repeated_cursor(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="0", input_schema={})],
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
nextCursor="same-cursor",
),
ListToolsResult(
tools=[MCPTool(name="tool_1", description="1", input_schema={})],
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
nextCursor="same-cursor",
),
]
@ -168,7 +168,7 @@ async def test_pagination_walk_stops_on_repeated_cursor(mock_session):
async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_0", description="0", input_schema={})],
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
nextCursor="",
),
]
@ -190,7 +190,7 @@ async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkey
await anyio.sleep(0.15)
idx = int(params.cursor) if params is not None else 0
return ListToolsResult(
tools=[MCPTool(name=f"tool_{idx}", description=str(idx), input_schema={})],
tools=[MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})],
nextCursor=str(idx + 1),
)
@ -212,7 +212,7 @@ async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_sessio
async def slow_page(params=None):
await anyio.sleep(0.15)
idx = int(params.cursor) if params is not None else 0
tools = [MCPTool(name=f"tool_{idx}", description=str(idx), input_schema={})]
tools = [MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})]
if idx == 0:
return ListToolsResult(tools=tools, nextCursor="1")
return ListToolsResult(tools=tools)
@ -227,10 +227,10 @@ async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_sessio
async def test_load_mcp_tools_openai_format_spans_pages(mock_session):
mock_session.list_tools.side_effect = [
ListToolsResult(
tools=[MCPTool(name="tool_a", description="a", input_schema={})],
tools=[MCPTool(name="tool_a", description="a", inputSchema={})],
nextCursor="page-2",
),
ListToolsResult(tools=[MCPTool(name="tool_b", description="b", input_schema={})]),
ListToolsResult(tools=[MCPTool(name="tool_b", description="b", inputSchema={})]),
]
result = await load_mcp_tools(mock_session, format="openai")
assert [t["function"]["name"] for t in result] == ["tool_a", "tool_b"]
@ -349,7 +349,7 @@ def test_transform_mcp_tool_to_openai_responses_api_tool():
minimal_tool = MCPTool(
name="GitMCP-fetch_litellm_documentation",
description="Fetch entire documentation file from GitHub repository",
input_schema={"type": "object"}, # This was causing the error
inputSchema={"type": "object"}, # This was causing the error
)
openai_tool = transform_mcp_tool_to_openai_responses_api_tool(minimal_tool)
@ -364,7 +364,7 @@ def test_transform_mcp_tool_to_openai_responses_api_tool():
complete_tool = MCPTool(
name="test_tool_complete",
description="A test tool with complete schema",
input_schema={
inputSchema={
"type": "object",
"properties": {"query": {"type": "string", "description": "Search query"}},
"required": ["query"],
@ -395,7 +395,7 @@ def test_transform_mcp_tool_to_anthropic_tool():
tool = MCPTool(
name="read_wiki_structure",
description="Get a list of documentation topics",
input_schema={
inputSchema={
"type": "object",
"properties": {"repoName": {"type": "string"}},
"required": ["repoName"],
@ -417,7 +417,7 @@ def test_transform_mcp_tool_to_anthropic_tool():
def test_transform_mcp_tool_to_anthropic_tool_normalizes_empty_schema():
"""A tool with no declared arguments must still present a valid object schema."""
anthropic_tool = transform_mcp_tool_to_anthropic_tool(
MCPTool(name="noargs", description=None, input_schema={})
MCPTool(name="noargs", description=None, inputSchema={})
)
assert anthropic_tool["name"] == "noargs"
@ -445,7 +445,7 @@ def test_transform_mcp_tool_to_anthropic_tool_strips_keys_anthropic_rejects():
tool = MCPTool(
name="rich",
description="tool with a dirty schema",
input_schema={
inputSchema={
"type": "object",
"properties": {"q": {"type": "string"}},
"required": ["q"],

View file

@ -70,7 +70,9 @@ def test_arize_set_attributes():
# Simulated LLM response object
response_obj = ModelResponse(
usage={"total_tokens": 100, "completion_tokens": 60, "prompt_tokens": 40},
choices=[Choices(message={"role": "assistant", "content": "Basic Response Content"})],
choices=[
Choices(message={"role": "assistant", "content": "Basic Response Content"})
],
model="gpt-4o",
id="chatcmpl-ID",
)
@ -87,7 +89,9 @@ def test_arize_set_attributes():
assert span.set_attribute.call_count == 26
# Metadata attached to the span
span.set_attribute.assert_any_call(SpanAttributes.METADATA, json.dumps({"key_1": "value_1", "key_2": None}))
span.set_attribute.assert_any_call(
SpanAttributes.METADATA, json.dumps({"key_1": "value_1", "key_2": None})
)
# Basic LLM information
span.set_attribute.assert_any_call(SpanAttributes.LLM_MODEL_NAME, "gpt-4o")
@ -110,12 +114,16 @@ def test_arize_set_attributes():
span.set_attribute.assert_any_call(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM")
# And TOOL must never be written for an LLM chat completion call.
span_kind_writes = [
c.args[1] for c in span.set_attribute.call_args_list if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
c.args[1]
for c in span.set_attribute.call_args_list
if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
]
assert "TOOL" not in span_kind_writes
# Request message content and metadata
span.set_attribute.assert_any_call(SpanAttributes.INPUT_VALUE, "Basic Request Content")
span.set_attribute.assert_any_call(
SpanAttributes.INPUT_VALUE, "Basic Request Content"
)
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}",
"user",
@ -126,7 +134,9 @@ def test_arize_set_attributes():
)
# Tool call definitions and function names
span.set_attribute.assert_any_call(f"{SpanAttributes.LLM_TOOLS}.0.name", "get_weather")
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_TOOLS}.0.name", "get_weather"
)
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_TOOLS}.0.description",
"Fetches weather details.",
@ -136,20 +146,26 @@ def test_arize_set_attributes():
json.dumps(
{
"type": "object",
"properties": {"location": {"type": "string", "description": "City name"}},
"properties": {
"location": {"type": "string", "description": "City name"}
},
"required": ["location"],
}
),
)
# Invocation parameters
span.set_attribute.assert_any_call(SpanAttributes.LLM_INVOCATION_PARAMETERS, '{"user": "test_user"}')
span.set_attribute.assert_any_call(
SpanAttributes.LLM_INVOCATION_PARAMETERS, '{"user": "test_user"}'
)
# User ID
span.set_attribute.assert_any_call(SpanAttributes.USER_ID, "test_user")
# Output message content
span.set_attribute.assert_any_call(SpanAttributes.OUTPUT_VALUE, "Basic Response Content")
span.set_attribute.assert_any_call(
SpanAttributes.OUTPUT_VALUE, "Basic Response Content"
)
span.set_attribute.assert_any_call(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}",
"assistant",
@ -212,7 +228,9 @@ def test_arize_set_attributes_responses_api():
ResponseReasoningItem(
id="reasoning-001",
type="reasoning",
summary=[Summary(text="First, I need to analyze...", type="summary_text")],
summary=[
Summary(text="First, I need to analyze...", type="summary_text")
],
),
ResponseOutputMessage(
id="msg-001",
@ -259,7 +277,9 @@ def test_arize_set_attributes_responses_api():
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 370)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 250)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 120)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180)
span.set_attribute.assert_any_call(
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180
)
def test_set_usage_outputs_pydantic_completion_usage():
@ -307,7 +327,9 @@ def test_set_usage_outputs_pydantic_completion_usage():
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 40)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 60)
# reasoning_tokens for chat completions live in completion_tokens_details
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 25)
span.set_attribute.assert_any_call(
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 25
)
def test_set_usage_outputs_pydantic_response_api_usage():
@ -340,7 +362,9 @@ def test_set_usage_outputs_pydantic_response_api_usage():
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 370)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 120)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 250)
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180)
span.set_attribute.assert_any_call(
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180
)
class TestArizeLogger(CustomLogger):
@ -351,12 +375,16 @@ class TestArizeLogger(CustomLogger):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = None
self.standard_callback_dynamic_params: Optional[
StandardCallbackDynamicParams
] = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
# Capture dynamic params and print them for verification
print("logged kwargs", json.dumps(kwargs, indent=4, default=str))
self.standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params")
self.standard_callback_dynamic_params = kwargs.get(
"standard_callback_dynamic_params"
)
@pytest.mark.asyncio
@ -382,8 +410,14 @@ async def test_arize_dynamic_params():
# Assert dynamic parameters were received in the callback
assert test_arize_logger.standard_callback_dynamic_params is not None
assert test_arize_logger.standard_callback_dynamic_params.get("arize_api_key") == "test_api_key_dynamic"
assert test_arize_logger.standard_callback_dynamic_params.get("arize_space_key") == "test_space_key_dynamic"
assert (
test_arize_logger.standard_callback_dynamic_params.get("arize_api_key")
== "test_api_key_dynamic"
)
assert (
test_arize_logger.standard_callback_dynamic_params.get("arize_space_key")
== "test_space_key_dynamic"
)
def test_construct_dynamic_arize_headers():
@ -394,7 +428,9 @@ def test_construct_dynamic_arize_headers():
from litellm.types.utils import StandardCallbackDynamicParams
# Test with all parameters present
dynamic_params_full = StandardCallbackDynamicParams(arize_api_key="test_api_key", arize_space_id="test_space_id")
dynamic_params_full = StandardCallbackDynamicParams(
arize_api_key="test_api_key", arize_space_id="test_space_id"
)
arize_logger = ArizeLogger()
headers = arize_logger.construct_dynamic_otel_headers(dynamic_params_full)
@ -402,7 +438,9 @@ def test_construct_dynamic_arize_headers():
assert headers == expected_headers
# Test with only space_id
dynamic_params_space_id_only = StandardCallbackDynamicParams(arize_space_id="test_space_id")
dynamic_params_space_id_only = StandardCallbackDynamicParams(
arize_space_id="test_space_id"
)
headers = arize_logger.construct_dynamic_otel_headers(dynamic_params_space_id_only)
expected_headers = {"arize-space-id": "test_space_id"}
@ -418,7 +456,9 @@ def test_construct_dynamic_arize_headers():
dynamic_params_space_key_and_api_key = StandardCallbackDynamicParams(
arize_space_key="test_space_key", arize_api_key="test_api_key"
)
headers = arize_logger.construct_dynamic_otel_headers(dynamic_params_space_key_and_api_key)
headers = arize_logger.construct_dynamic_otel_headers(
dynamic_params_space_key_and_api_key
)
expected_headers = {"arize-space-id": "test_space_key", "api_key": "test_api_key"}
@ -488,7 +528,9 @@ def test_arize_emits_no_cache_tokens_when_absent():
from litellm.integrations.arize._utils import _set_usage_outputs
span = MagicMock()
response_obj = {"usage": {"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6}}
response_obj = {
"usage": {"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6}
}
_set_usage_outputs(span, response_obj, SpanAttributes)
attrs = _collect_calls(span)
assert SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ not in attrs
@ -500,8 +542,14 @@ def test_passthrough_call_type_resolves_to_llm_span_kind():
from litellm.integrations._types.open_inference import OpenInferenceSpanKindValues
from litellm.integrations.arize._utils import _infer_open_inference_span_kind
assert _infer_open_inference_span_kind("allm_passthrough_route") == OpenInferenceSpanKindValues.LLM.value
assert _infer_open_inference_span_kind("llm_passthrough_route") == OpenInferenceSpanKindValues.LLM.value
assert (
_infer_open_inference_span_kind("allm_passthrough_route")
== OpenInferenceSpanKindValues.LLM.value
)
assert (
_infer_open_inference_span_kind("llm_passthrough_route")
== OpenInferenceSpanKindValues.LLM.value
)
def test_arize_chat_completion_with_tools_stays_llm_span_kind():
@ -557,7 +605,9 @@ def test_arize_chat_completion_with_tools_stays_llm_span_kind():
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
span_kind_writes = [
c.args[1] for c in span.set_attribute.call_args_list if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
c.args[1]
for c in span.set_attribute.call_args_list
if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
]
assert span_kind_writes, "span.kind must be written"
assert all(v == "LLM" for v in span_kind_writes)
@ -609,8 +659,13 @@ def test_arize_emits_assistant_tool_calls_on_output_message():
attrs = _collect_calls(span)
base = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_TOOL_CALLS}.0"
assert attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_ID}"] == "call_abc"
assert attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}"] == "get_weather"
assert attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}"] == '{"location": "SF"}'
assert (
attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}"] == "get_weather"
)
assert (
attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}"]
== '{"location": "SF"}'
)
def test_arize_output_value_falls_back_to_tool_calls_summary():
@ -763,7 +818,9 @@ def test_arize_emits_tool_call_id_and_name_on_input_tool_message():
assert attrs[f"{assistant_base}.{ToolCallAttributes.TOOL_CALL_ID}"] == "call_abc"
# Tool message at index 2
tool_prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.2"
assert attrs[f"{tool_prefix}.{MessageAttributes.MESSAGE_TOOL_CALL_ID}"] == "call_abc"
assert (
attrs[f"{tool_prefix}.{MessageAttributes.MESSAGE_TOOL_CALL_ID}"] == "call_abc"
)
assert attrs[f"{tool_prefix}.{MessageAttributes.MESSAGE_NAME}"] == "get_weather"
@ -809,7 +866,10 @@ def test_arize_emits_multimodal_input_contents():
assert attrs[f"{base}.0.message_content.type"] == "text"
assert attrs[f"{base}.0.message_content.text"] == "What is in this image?"
assert attrs[f"{base}.1.message_content.type"] == "image"
assert attrs[f"{base}.1.message_content.image.image.url"] == "https://example.com/cat.png"
assert (
attrs[f"{base}.1.message_content.image.image.url"]
== "https://example.com/cat.png"
)
def test_arize_emits_session_and_user_attrs_from_metadata():
@ -914,7 +974,11 @@ def test_arize_does_not_overwrite_user_id_from_optional_params():
id="r2",
)
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
user_id_writes = [c.args[1] for c in span.set_attribute.call_args_list if c.args[0] == SpanAttributes.USER_ID]
user_id_writes = [
c.args[1]
for c in span.set_attribute.call_args_list
if c.args[0] == SpanAttributes.USER_ID
]
assert "from_metadata" not in user_id_writes
@ -984,7 +1048,9 @@ def test_arize_passthrough_bedrock_anthropic_normalization():
"complete_input_dict": {
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 64,
"messages": [{"role": "user", "content": "What is the capital of France?"}],
"messages": [
{"role": "user", "content": "What is the capital of France?"}
],
}
},
"standard_logging_object": {
@ -1002,13 +1068,19 @@ def test_arize_passthrough_bedrock_anthropic_normalization():
assert attrs[SpanAttributes.INPUT_VALUE] == "What is the capital of France?"
msg0 = f"{SpanAttributes.LLM_INPUT_MESSAGES}.0"
assert attrs[f"{msg0}.{MessageAttributes.MESSAGE_ROLE}"] == "user"
assert attrs[f"{msg0}.{MessageAttributes.MESSAGE_CONTENT}"] == "What is the capital of France?"
assert (
attrs[f"{msg0}.{MessageAttributes.MESSAGE_CONTENT}"]
== "What is the capital of France?"
)
# Output rendering (Anthropic content[].text)
assert attrs[SpanAttributes.OUTPUT_VALUE] == "The capital of France is Paris."
out0 = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0"
assert attrs[f"{out0}.{MessageAttributes.MESSAGE_ROLE}"] == "assistant"
assert attrs[f"{out0}.{MessageAttributes.MESSAGE_CONTENT}"] == "The capital of France is Paris."
assert (
attrs[f"{out0}.{MessageAttributes.MESSAGE_CONTENT}"]
== "The capital of France is Paris."
)
# Token counts (Bedrock input_tokens/output_tokens) — extracted via
# coercion of the non-dict response.
@ -1017,7 +1089,9 @@ def test_arize_passthrough_bedrock_anthropic_normalization():
# Span kind defended even though the call_type is a passthrough variant.
span_kind_writes = [
c.args[1] for c in span.set_attribute.call_args_list if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
c.args[1]
for c in span.set_attribute.call_args_list
if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
]
assert span_kind_writes # at least one
assert all(v == "LLM" for v in span_kind_writes)
@ -1035,7 +1109,11 @@ def test_arize_passthrough_call_type_does_not_run_on_chat_completion():
span = MagicMock()
_maybe_normalize_passthrough(
span,
{"additional_args": {"complete_input_dict": {"messages": [{"role": "user", "content": "x"}]}}},
{
"additional_args": {
"complete_input_dict": {"messages": [{"role": "user", "content": "x"}]}
}
},
{"choices": [{"message": {"role": "assistant", "content": "y"}}]},
{"choices": [{"message": {"role": "assistant", "content": "y"}}]},
{"call_type": "completion"},
@ -1055,7 +1133,11 @@ def test_arize_passthrough_skipped_when_message_redaction_enabled():
span = MagicMock()
kwargs = {
"additional_args": {
"complete_input_dict": {"messages": [{"role": "user", "content": "Patient John Doe, SSN 123-45-6789"}]}
"complete_input_dict": {
"messages": [
{"role": "user", "content": "Patient John Doe, SSN 123-45-6789"}
]
}
},
# Enables redaction via the dynamic-param path inside
# should_redact_message_logging(), without touching globals.
@ -1129,7 +1211,9 @@ def test_arize_mcp_call_tool_result_does_not_break_attribute_setting():
"optional_params": {},
"litellm_params": {"custom_llm_provider": "mcp"},
}
response_obj = CallToolResult(content=[TextContent(type="text", text="sunny, 21C")], is_error=False)
response_obj = CallToolResult(
content=[TextContent(type="text", text="sunny, 21C")], isError=False
)
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
@ -1147,7 +1231,7 @@ def test_arize_coerce_response_obj_dumps_pydantic_without_get():
from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs
result = CallToolResult(content=[TextContent(type="text", text="hi")], is_error=False)
result = CallToolResult(content=[TextContent(type="text", text="hi")], isError=False)
coerced = _coerce_response_obj_for_attrs(result)
assert isinstance(coerced, dict)
@ -1211,7 +1295,9 @@ def test_arize_mcp_tool_span_renders_name_input_and_output():
from mcp.types import CallToolResult, TextContent
span = MagicMock()
response_obj = CallToolResult(content=[TextContent(type="text", text="sunny, 21C")], is_error=False)
response_obj = CallToolResult(
content=[TextContent(type="text", text="sunny, 21C")], isError=False
)
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
@ -1232,7 +1318,7 @@ def test_arize_mcp_tool_span_serializes_non_text_content():
span = MagicMock()
response_obj = CallToolResult(
content=[ImageContent(type="image", data="Zm9v", mimeType="image/png")],
is_error=False,
isError=False,
)
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
@ -1250,7 +1336,9 @@ def test_arize_mcp_tool_span_respects_message_redaction():
from mcp.types import CallToolResult, TextContent
span = MagicMock()
response_obj = CallToolResult(content=[TextContent(type="text", text="SSN 123-45-6789")], is_error=False)
response_obj = CallToolResult(
content=[TextContent(type="text", text="SSN 123-45-6789")], isError=False
)
ArizeLogger.set_arize_attributes(
span,
@ -1302,7 +1390,7 @@ def test_arize_mcp_tool_span_renders_empty_arguments():
span = MagicMock()
kwargs = _mcp_kwargs(mcp_tool_call_metadata={"name": "ping", "arguments": {}})
response_obj = CallToolResult(content=[TextContent(type="text", text="pong")], is_error=False)
response_obj = CallToolResult(content=[TextContent(type="text", text="pong")], isError=False)
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
@ -1317,7 +1405,7 @@ def test_arize_mcp_tool_span_renders_empty_content():
from mcp.types import CallToolResult
span = MagicMock()
response_obj = CallToolResult(content=[], is_error=False)
response_obj = CallToolResult(content=[], isError=False)
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
@ -1332,7 +1420,7 @@ def test_arize_mcp_tool_span_falls_back_to_structured_content():
from mcp.types import CallToolResult
span = MagicMock()
response_obj = CallToolResult(content=[], structured_content={"temp_c": 21}, is_error=False)
response_obj = CallToolResult(content=[], structuredContent={"temp_c": 21}, isError=False)
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
@ -1375,7 +1463,7 @@ def test_arize_mcp_tool_span_serializes_mixed_text_and_media():
TextContent(type="text", text="see image"),
ImageContent(type="image", data="Zm9v", mimeType="image/png"),
],
is_error=False,
isError=False,
)
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)

View file

@ -44,3 +44,37 @@ def _hermetic_server_root_path():
finally:
if saved is not None:
os.environ["SERVER_ROOT_PATH"] = saved
@pytest.fixture
def config_only_mcp_manager_factory():
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
class ConfigOnlyManager(MCPServerManager):
def initialize_tool_name_to_mcp_server_name_mapping(self):
return None
return ConfigOnlyManager
@pytest.fixture
def _mcp_request_ctx():
def _mcp_request_ctx(**overrides):
from types import SimpleNamespace
from mcp.server.context import ServerRequestContext
kwargs = {
"session": SimpleNamespace(),
"lifespan_context": {},
"protocol_version": "2025-06-18",
"method": "",
"params": None,
"request_id": 1,
"meta": None,
"request": None,
}
kwargs.update(overrides)
return ServerRequestContext(**kwargs)
return _mcp_request_ctx

View file

@ -533,7 +533,7 @@ async def test_process_output_response_masks_text_content():
TextContent(type="text", text="email jane@example.com"),
TextContent(type="text", text="call 415-555-0132"),
],
is_error=False,
isError=False,
)
returned = await handler.process_output_response(
@ -569,7 +569,7 @@ async def test_process_output_response_propagates_block():
guardrail = MaskingGuardrail(
raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="masking-mcp-guardrail")
)
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], is_error=False)
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
with pytest.raises(BlockedPiiEntityError):
await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
@ -582,7 +582,7 @@ async def test_process_output_response_skips_non_text_content():
guardrail = MaskingGuardrail(masked_texts=["should not be used"])
result = CallToolResult(
content=[ImageContent(type="image", data="aGk=", mimeType="image/png")],
is_error=False,
isError=False,
)
returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
@ -613,7 +613,7 @@ async def test_process_output_response_blocks_on_text_count_mismatch():
TextContent(type="text", text="jane@example.com"),
TextContent(type="text", text="415-555-0132"),
],
is_error=False,
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
@ -645,14 +645,14 @@ async def test_structured_content_is_masked_alongside_content():
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="email jane@example.com")],
structured_content={"contact": {"email": "jane@example.com"}, "balance": 42.0},
is_error=False,
structuredContent={"contact": {"email": "jane@example.com"}, "balance": 42.0},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert returned.content[0].text == "email <EMAIL_ADDRESS>"
assert returned.structured_content== {"contact": {"email": "<EMAIL_ADDRESS>"}, "balance": 42.0}
assert returned.structured_content == {"contact": {"email": "<EMAIL_ADDRESS>"}, "balance": 42.0}
@pytest.mark.asyncio
@ -666,14 +666,14 @@ async def test_value_present_only_in_structured_content_is_masked():
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structured_content={"records": [{"email": "jane@example.com"}]},
is_error=False,
structuredContent={"records": [{"email": "jane@example.com"}]},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert "jane@example.com" in guardrail.seen_texts
assert returned.structured_content== {"records": [{"email": "<EMAIL_ADDRESS>"}]}
assert returned.structured_content == {"records": [{"email": "<EMAIL_ADDRESS>"}]}
assert returned.content[0].text == "lookup complete"
@ -684,13 +684,13 @@ async def test_structured_content_without_a_match_is_untouched():
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structured_content={"record_id": "C-1001", "balance": 42.0, "active": True, "note": None},
is_error=False,
structuredContent={"record_id": "C-1001", "balance": 42.0, "active": True, "note": None},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert returned.structured_content== {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None}
assert returned.structured_content == {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None}
@pytest.mark.asyncio
@ -707,8 +707,8 @@ async def test_structured_content_nested_too_deeply_is_blocked():
nested = {"next": nested}
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structured_content=nested,
is_error=False,
structuredContent=nested,
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
@ -754,8 +754,8 @@ async def test_sensitive_structured_content_key_is_blocked():
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structured_content={"jane@example.com": {"balance": 42.0}},
is_error=False,
structuredContent={"jane@example.com": {"balance": 42.0}},
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
@ -774,8 +774,8 @@ async def test_sensitive_structured_content_numeric_value_is_blocked():
guardrail = SubstitutingGuardrail("4155550199", "<PHONE_NUMBER>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structured_content={"phone": 4155550199},
is_error=False,
structuredContent={"phone": 4155550199},
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
@ -791,11 +791,11 @@ async def test_clean_structured_content_keys_do_not_block():
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="email jane@example.com")],
structured_content={"record_id": "C-1001", "balance": 42.0, "count": 3},
is_error=False,
structuredContent={"record_id": "C-1001", "balance": 42.0, "count": 3},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert returned.content[0].text == "email <EMAIL_ADDRESS>"
assert returned.structured_content== {"record_id": "C-1001", "balance": 42.0, "count": 3}
assert returned.structured_content == {"record_id": "C-1001", "balance": 42.0, "count": 3}

View file

@ -18,9 +18,9 @@ from litellm.proxy._types import LiteLLM_MCPServerTable
class TestMCPCustomFields:
"""Test custom fields functionality in MCP server configuration."""
async def test_custom_fields_preserved_from_config(self):
async def test_custom_fields_preserved_from_config(self, config_only_mcp_manager_factory):
"""Test that custom fields in mcp_info are preserved when loading from config."""
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
# Mock config with custom fields
mock_config = {
@ -62,9 +62,9 @@ class TestMCPCustomFields:
assert mcp_info["priority"] == 10
assert mcp_info["tags"] == ["production", "api"]
async def test_custom_fields_preserved_from_database(self):
async def test_custom_fields_preserved_from_database(self, config_only_mcp_manager_factory):
"""Test that custom fields in mcp_info are preserved when adding from database."""
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
# Mock database record with custom fields
mock_server = LiteLLM_MCPServerTable(
@ -106,9 +106,9 @@ class TestMCPCustomFields:
assert mcp_info["metadata"] == {"source": "database"}
assert mcp_info["version"] == "1.0.0"
async def test_empty_mcp_info_handled_gracefully(self):
async def test_empty_mcp_info_handled_gracefully(self, config_only_mcp_manager_factory):
"""Test that empty or missing mcp_info is handled gracefully."""
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
# Config with empty mcp_info
mock_config = {
@ -130,9 +130,9 @@ class TestMCPCustomFields:
# Should have default server_name
assert mcp_info["server_name"] == "test_server"
async def test_missing_mcp_info_creates_defaults(self):
async def test_missing_mcp_info_creates_defaults(self, config_only_mcp_manager_factory):
"""Test that missing mcp_info creates appropriate defaults."""
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
# Config without mcp_info
mock_config = {
@ -155,9 +155,9 @@ class TestMCPCustomFields:
assert mcp_info["server_name"] == "test_server"
assert mcp_info["description"] == "Server description"
async def test_config_description_fallback(self):
async def test_config_description_fallback(self, config_only_mcp_manager_factory):
"""Test that description from config level is used as fallback."""
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
# Config with description at server level but not in mcp_info
mock_config = {
@ -179,9 +179,9 @@ class TestMCPCustomFields:
assert mcp_info["description"] == "Config level description"
assert mcp_info["custom_field"] == "custom_value"
async def test_mcp_info_description_takes_precedence(self):
async def test_mcp_info_description_takes_precedence(self, config_only_mcp_manager_factory):
"""Test that description in mcp_info takes precedence over config level."""
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
# Config with description at both levels
mock_config = {

View file

@ -262,23 +262,6 @@ class TestDescribeUpstreamHttpFailure:
assert describe_upstream_http_failure(ConnectionError("refused")) is None
def _mcp_request_ctx(**overrides):
from types import SimpleNamespace
from mcp.server.context import ServerRequestContext
kwargs = {
"session": SimpleNamespace(),
"lifespan_context": {},
"protocol_version": "2025-06-18",
"method": "",
"params": None,
"request_id": 1,
"meta": None,
"request": None,
}
kwargs.update(overrides)
return ServerRequestContext(**kwargs)
@pytest.mark.parametrize("body", [
b'{"password":"first second","token":"demo-secret"}',
@ -479,7 +462,7 @@ def test_diagnostics_keep_requests_separate_and_do_not_collapse_multiple_servers
@pytest.mark.asyncio
async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None:
async def test_concurrent_mcp_messages_record_on_their_own_http_scope(_mcp_request_ctx) -> None:
from unittest.mock import MagicMock
from starlette.requests import Request
@ -557,7 +540,6 @@ def test_oversized_request_omits_potentially_reflected_response_credentials():
@pytest.mark.asyncio
async def test_streamed_error_redacts_reflected_credentials_before_capture():
import json
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
secret = "generic-credential-123"

View file

@ -1714,7 +1714,7 @@ async def test_missing_user_env_vars_error_renders_in_mcp_call_tool():
result = CallToolResult(
content=[TextContent(text=str(err), type="text")],
is_error=True,
isError=True,
)
assert result.is_error is True
text = result.content[0].text # type: ignore[union-attr]

View file

@ -38,7 +38,7 @@ class TestMCPMetadataPreservation:
tool_with_metadata = MCPTool(
name="hello_widget",
description="Display a greeting widget",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
meta={
"openai/outputTemplate": "ui://widget/hello.html",
"openai/widgetDescription": "A greeting widget",

View file

@ -332,7 +332,7 @@ async def test_aggregate_list_tools_absorbs_one_unauthenticated_server():
"s1", "delegate_docs", auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True
)
working = _http_server("s2", "working_docs", auth_type=MCPAuth.none)
good_tool = MCPTool(name="working_docs-read", description="d", input_schema={"type": "object"})
good_tool = MCPTool(name="working_docs-read", description="d", inputSchema={"type": "object"})
async def fake_get_tools(server, **kwargs):
if server.server_id == delegate.server_id:

View file

@ -10,6 +10,8 @@ import json
from types import SimpleNamespace
from typing import Any, Dict
from mcp.types import TextContent, ToolResultContent
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_messages_to_openai,
_convert_single_content,
@ -21,8 +23,8 @@ from litellm.proxy._experimental.mcp_server.sampling_handler import (
# ---------------------------------------------------------------------------
def _text(text: str) -> SimpleNamespace:
return SimpleNamespace(type="text", text=text)
def _text(text: str) -> TextContent:
return TextContent(type="text", text=text)
def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleNamespace:
@ -31,11 +33,9 @@ def _tool_use(*, name: str, tool_id: str, input_data: Dict[str, Any]) -> SimpleN
def _tool_result(
*, tool_use_id: str, content: Any = None, is_error: bool = False
) -> SimpleNamespace:
if content is None:
content = []
return SimpleNamespace(
type="tool_result", toolUseId=tool_use_id, content=content, is_error=is_error
) -> ToolResultContent:
return ToolResultContent(
tool_use_id=tool_use_id, content=[] if content is None else content, is_error=is_error
)

View file

@ -81,23 +81,6 @@ def cleanup_mcp_global_state():
def _mcp_request_ctx(**overrides):
from types import SimpleNamespace
from mcp.server.context import ServerRequestContext
kwargs = {
"session": SimpleNamespace(),
"lifespan_context": {},
"protocol_version": "2025-06-18",
"method": "",
"params": None,
"request_id": 1,
"meta": None,
"request": None,
}
kwargs.update(overrides)
return ServerRequestContext(**kwargs)
def _call_tool_params(name, arguments=None):
@ -112,7 +95,7 @@ def _paged_params():
return PaginatedRequestParams()
@pytest.mark.asyncio
async def test_mcp_server_tool_call_body_contains_request_data():
async def test_mcp_server_tool_call_body_contains_request_data(_mcp_request_ctx):
"""Test that proxy_server_request body contains name and arguments"""
try:
from litellm.proxy._experimental.mcp_server.server import (
@ -173,7 +156,7 @@ async def test_mcp_server_tool_call_body_contains_request_data():
@pytest.mark.asyncio
async def test_mcp_server_tool_call_forwards_client_headers_to_logging():
async def test_mcp_server_tool_call_forwards_client_headers_to_logging(_mcp_request_ctx):
"""The MCP protocol path must hand the connection's client headers to the pre-call
pipeline, so logging callbacks and guardrails see them the way the REST path does."""
try:
@ -222,7 +205,7 @@ async def test_mcp_server_tool_call_forwards_client_headers_to_logging():
@pytest.mark.asyncio
async def test_mcp_server_tool_call_strips_custom_litellm_key_header():
async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_request_ctx):
"""The deployment can rename the proxy key header via general_settings.litellm_key_header_name.
The pre-call pipeline only knows that name if it is passed in, so without it the virtual key
reaches metadata.headers and proxy_server_request.headers in plaintext."""
@ -274,7 +257,7 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header():
@pytest.mark.asyncio
async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror():
async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx):
"""The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session
tool call cannot emit a raw 401 the way the REST path does. mcp_server_tool_call must turn an
upstream MCPUpstreamAuthError into an explicit isError result naming the status, not a masked
@ -1360,7 +1343,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
tool1 = MagicMock()
tool1.name = "working_tool_1"
tool1.description = "Working tool 1"
tool1.input_schema= {}
tool1.input_schema = {}
return [tool1]
else:
# Failing server raises an exception
@ -1736,7 +1719,7 @@ async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_na
@pytest.mark.asyncio
async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error():
async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(_mcp_request_ctx):
"""The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error
(MCPError, INVALID_REQUEST) carrying the denial message, instead of a raw 500."""
try:
@ -1768,7 +1751,7 @@ async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(
@pytest.mark.asyncio
async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict():
async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(_mcp_request_ctx):
try:
from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call
except ImportError:
@ -1794,7 +1777,7 @@ async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict():
@pytest.mark.asyncio
async def test_mcp_server_tool_call_body_with_none_arguments():
async def test_mcp_server_tool_call_body_with_none_arguments(_mcp_request_ctx):
"""Test that proxy_server_request body handles None arguments correctly"""
try:
from litellm.proxy._experimental.mcp_server.server import (
@ -2011,7 +1994,7 @@ async def test_streamable_http_session_manager_is_stateless():
("DELETE", b"", False),
),
)
async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(
async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(_mcp_request_ctx,
debug: bool, method: str, request_body: bytes, stateful: bool
) -> None:
from starlette.requests import Request
@ -4167,7 +4150,7 @@ async def test_list_tools_single_server_unprefixed_names():
tool = MagicMock()
tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
tool.description = "desc"
tool.input_schema= {}
tool.input_schema = {}
return [tool]
mock_manager._get_tools_from_server = mock_get_tools_from_server
@ -4246,7 +4229,7 @@ async def test_list_tools_multiple_servers_prefixed_names():
# When multiple servers, add_prefix should be True -> prefixed names
tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
tool.description = "desc"
tool.input_schema= {}
tool.input_schema = {}
return [tool]
mock_manager._get_tools_from_server = mock_get_tools_from_server
@ -4659,22 +4642,22 @@ async def test_list_tools_filters_by_key_team_permissions():
tool1 = MagicMock()
tool1.name = "tool1"
tool1.description = "Tool 1"
tool1.input_schema= {}
tool1.input_schema = {}
tool2 = MagicMock()
tool2.name = "tool2"
tool2.description = "Tool 2"
tool2.input_schema= {}
tool2.input_schema = {}
tool3 = MagicMock()
tool3.name = "tool3"
tool3.description = "Tool 3 - not allowed"
tool3.input_schema= {}
tool3.input_schema = {}
tool4 = MagicMock()
tool4.name = "tool4"
tool4.description = "Tool 4 - not allowed"
tool4.input_schema= {}
tool4.input_schema = {}
return [tool1, tool2, tool3, tool4]
@ -4770,22 +4753,22 @@ async def test_list_tools_with_team_tool_permissions_inheritance():
tool1 = MagicMock()
tool1.name = "tool1"
tool1.description = "Tool 1"
tool1.input_schema= {}
tool1.input_schema = {}
tool2 = MagicMock()
tool2.name = "tool2"
tool2.description = "Tool 2"
tool2.input_schema= {}
tool2.input_schema = {}
tool3 = MagicMock()
tool3.name = "tool3"
tool3.description = "Tool 3"
tool3.input_schema= {}
tool3.input_schema = {}
tool4 = MagicMock()
tool4.name = "tool4"
tool4.description = "Tool 4"
tool4.input_schema= {}
tool4.input_schema = {}
return [tool1, tool2, tool3, tool4]
@ -4867,17 +4850,17 @@ async def test_list_tools_with_no_tool_permissions_shows_all():
tool1 = MagicMock()
tool1.name = "tool1"
tool1.description = "Tool 1"
tool1.input_schema= {}
tool1.input_schema = {}
tool2 = MagicMock()
tool2.name = "tool2"
tool2.description = "Tool 2"
tool2.input_schema= {}
tool2.input_schema = {}
tool3 = MagicMock()
tool3.name = "tool3"
tool3.description = "Tool 3"
tool3.input_schema= {}
tool3.input_schema = {}
return [tool1, tool2, tool3]
@ -4968,22 +4951,22 @@ async def test_list_tools_strips_prefix_when_matching_permissions():
tool1 = MagicMock()
tool1.name = "GITMCP-fetch_litellm_documentation" # Prefixed
tool1.description = "Fetch docs"
tool1.input_schema= {}
tool1.input_schema = {}
tool2 = MagicMock()
tool2.name = "GITMCP-search_litellm_documentation" # Prefixed, not in allowed list
tool2.description = "Search docs"
tool2.input_schema= {}
tool2.input_schema = {}
tool3 = MagicMock()
tool3.name = "GITMCP-search_litellm_code" # Prefixed
tool3.description = "Search code"
tool3.input_schema= {}
tool3.input_schema = {}
tool4 = MagicMock()
tool4.name = "GITMCP-fetch_generic_url_content" # Prefixed, not in allowed list
tool4.description = "Fetch URL"
tool4.input_schema= {}
tool4.input_schema = {}
return [tool1, tool2, tool3, tool4]
@ -5033,7 +5016,7 @@ def test_filter_tools_by_allowed_tools():
name="my_api_mcp-getpetbyid",
title=None,
description="Find pet by ID",
input_schema={
inputSchema={
"type": "object",
"properties": {"petId": {"type": "integer", "description": ""}},
"required": ["petId"],
@ -5045,7 +5028,7 @@ def test_filter_tools_by_allowed_tools():
name="my_api_mcp-findpetsbystatus",
title=None,
description="Finds Pets by status",
input_schema={
inputSchema={
"type": "object",
"properties": {"status": {"type": "string", "description": ""}},
"required": ["status"],
@ -5057,7 +5040,7 @@ def test_filter_tools_by_allowed_tools():
name="my_api_mcp-addpet",
title=None,
description="Add a new pet to the store",
input_schema={
inputSchema={
"type": "object",
"properties": {
"body": {
@ -5103,7 +5086,7 @@ def test_apply_tool_overrides():
name="my_api_mcp-getpetbyid",
title=None,
description="Original description",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
outputSchema=None,
annotations=None,
),
@ -5111,7 +5094,7 @@ def test_apply_tool_overrides():
name="my_api_mcp-findpetsbystatus",
title=None,
description="Finds Pets by status",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
outputSchema=None,
annotations=None,
),
@ -5145,7 +5128,7 @@ def test_apply_tool_overrides_no_overrides():
name="my_api_mcp-getpetbyid",
title=None,
description="Original description",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
outputSchema=None,
annotations=None,
),
@ -5487,7 +5470,7 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab
tool_1 = MCPTool(
name="server_a-tool_1",
description="test tool",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
dummy_logging_obj = MagicMock()
@ -5793,7 +5776,7 @@ def test_filter_tools_enforced_empty_allowlist_blocks_all():
name="read_wiki_structure",
title=None,
description="",
input_schema={"type": "object"},
inputSchema={"type": "object"},
outputSchema=None,
annotations=None,
),
@ -5823,7 +5806,7 @@ def test_filter_tools_legacy_empty_allowlist_allows_all():
name="read_wiki_structure",
title=None,
description="",
input_schema={"type": "object"},
inputSchema={"type": "object"},
outputSchema=None,
annotations=None,
),
@ -8223,7 +8206,7 @@ class TestMCPMetaTraceCarrier:
@pytest.mark.asyncio
async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations() -> None:
async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations(_mcp_request_ctx) -> None:
from types import SimpleNamespace
from litellm.integrations.otel.model.destination import OtelDestination
@ -8371,7 +8354,7 @@ async def test_get_active_submitted_mcp_server_ids_for_user_empty_user_id_skips_
def _call_tool_result(is_error: bool, text: str) -> CallToolResult:
return CallToolResult(content=[TextContent(type="text", text=text)], is_error=is_error)
return CallToolResult(content=[TextContent(type="text", text=text)], isError=is_error)
def _mock_mcp_logging_obj() -> MagicMock:
@ -8399,7 +8382,7 @@ def test_extract_mcp_tool_result_error_message():
assert extract_mcp_tool_result_error_message(_call_tool_result(True, "boom")) == "boom"
assert extract_mcp_tool_result_error_message(_call_tool_result(False, "ok")) is None
assert (
extract_mcp_tool_result_error_message(CallToolResult(content=[], is_error=True))
extract_mcp_tool_result_error_message(CallToolResult(content=[], isError=True))
== "MCP tool call returned isError=true"
)
assert (
@ -8875,7 +8858,7 @@ async def test_aggregate_listing_reports_per_server_outcomes():
tool1 = MagicMock()
tool1.name = "working_tool_1"
tool1.description = "Working tool 1"
tool1.input_schema= {}
tool1.input_schema = {}
return [tool1]
raise MCPServerListError(ServerListFault(tag="upstream_error", status_code=500), server.name)
@ -8924,7 +8907,7 @@ async def test_outcome_keys_use_display_prefix_never_canonical_names():
@pytest.mark.asyncio
async def test_handle_list_tools_attaches_outcome_meta():
async def test_handle_list_tools_attaches_outcome_meta(_mcp_request_ctx):
"""The protocol handler returns a ListToolsResult whose _meta carries the per-server outcomes,
so MCP clients can tell a degraded listing from a genuinely empty one."""
try:
@ -8941,7 +8924,7 @@ async def test_handle_list_tools_attaches_outcome_meta():
ServerListOk,
)
tool = Tool(name="t1", input_schema={"type": "object"})
tool = Tool(name="t1", inputSchema={"type": "object"})
listing = AggregateToolListing(
tools=[tool],
outcomes={"healthy": ServerListOk(tool_count=1), "broken": ServerListFault(tag="unreachable")},
@ -9505,7 +9488,7 @@ class TestListFiltersHonorThePrefixBoundary:
from mcp.types import Tool as MCPTool
return [
MCPTool(name=f"{self.SERVER_ID}-{bare}", description=bare, input_schema={"type": "object"})
MCPTool(name=f"{self.SERVER_ID}-{bare}", description=bare, inputSchema={"type": "object"})
for bare in bare_names
]
@ -9609,13 +9592,13 @@ class TestListFiltersHonorThePrefixBoundary:
manager = MCPServerManager()
manager._create_prefixed_tools(
[MCPTool(name="read_wiki_contents", description="", input_schema={"type": "object"})],
[MCPTool(name="read_wiki_contents", description="", inputSchema={"type": "object"})],
_server(),
)
registered = sorted(manager.tool_name_to_mcp_server_name_mapping)
assert len(registered) > 1
published = MCPTool(name="eiG-read_wiki_contents", description="", input_schema={"type": "object"})
published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"})
for spelling in registered:
for entry, expected in ((spelling, True), (spelling.upper(), False)):
server = _server(disallowed_tools=[entry])
@ -9664,7 +9647,7 @@ class TestListFiltersHonorThePrefixBoundary:
url="http://127.0.0.1:5115/mcp",
transport=MCPTransport.http,
)
published = MCPTool(name=f"{self.SERVER_ID}-read_wiki_contents", description="", input_schema={"type": "object"})
published = MCPTool(name=f"{self.SERVER_ID}-read_wiki_contents", description="", inputSchema={"type": "object"})
auth = UserAPIKeyAuth(api_key="sk-test")
with (
@ -9721,7 +9704,7 @@ async def test_list_tools_injects_byok_credential_for_non_oauth2_auth_types(auth
tool = MagicMock()
tool.name = f"{server.alias}-toolA" if add_prefix else "toolA"
tool.description = "desc"
tool.input_schema= {}
tool.input_schema = {}
return [tool]
mock_manager = MagicMock()
@ -9751,28 +9734,8 @@ async def test_list_tools_injects_byok_credential_for_non_oauth2_auth_types(auth
assert [tool.name for tool in listing.tools] == ["byok-toolA"]
@pytest.mark.parametrize(
"method,handler_name",
[
("tools/list", "handle_list_tools"),
("tools/call", "mcp_server_tool_call"),
("prompts/list", "list_prompts"),
("prompts/get", "get_prompt"),
("resources/list", "list_resources"),
("resources/templates/list", "list_resource_templates"),
("resources/read", "read_resource"),
],
)
def test_mcp_server_registers_all_spec_handlers(method: str, handler_name: str) -> None:
from litellm.proxy._experimental.mcp_server import server as mcp_module
entry = mcp_module.server.get_request_handler(method)
assert entry is not None
assert getattr(mcp_module, handler_name) is entry.handler
@pytest.mark.asyncio
async def test_active_request_ctx_var_feeds_get_current_session() -> None:
async def test_active_request_ctx_var_feeds_get_current_session(_mcp_request_ctx) -> None:
from litellm.proxy._experimental.mcp_server.server import _get_current_session
session = SimpleNamespace()
@ -9786,7 +9749,7 @@ async def test_active_request_ctx_var_feeds_get_current_session() -> None:
@pytest.mark.asyncio
async def test_active_request_ctx_var_feeds_auth_resolution_recording() -> None:
async def test_active_request_ctx_var_feeds_auth_resolution_recording(_mcp_request_ctx) -> None:
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.mcp_debug import (
@ -9849,23 +9812,3 @@ async def test_streamable_http_rejects_modern_protocol_version(header_value: str
assert header_value in body["error"]["message"]
for version in body["error"]["message"].split("supported: ")[1].split(", "):
assert version in HANDSHAKE_PROTOCOL_VERSIONS
@pytest.mark.asyncio
async def test_initialize_never_negotiates_outside_handshake_versions() -> None:
from mcp.server.runner import ServerRunner
from litellm.proxy._experimental.mcp_server import server as mcp_module
negotiate = ServerRunner._negotiate_initialize
for requested in ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25", "9999-01-01"):
_, negotiated = negotiate({"protocolVersion": requested, "capabilities": {}, "clientInfo": {"name": "t", "version": "0"}})
assert negotiated in HANDSHAKE_PROTOCOL_VERSIONS
from mcp.server.connection import Connection
runner = ServerRunner(mcp_module.server, Connection.from_envelope(LATEST_HANDSHAKE_VERSION, None, None), None)
result = runner._handle_initialize(
{"protocolVersion": "9999-01-01", "capabilities": {}, "clientInfo": {"name": "t", "version": "0"}}
)
assert result.protocol_version in HANDSHAKE_PROTOCOL_VERSIONS

View file

@ -380,7 +380,7 @@ class TestMCPServerManagerSigV4:
"""Tests for MCPServerManager config loading with SigV4."""
@pytest.mark.asyncio
async def test_load_config_with_aws_sigv4(self):
async def test_load_config_with_aws_sigv4(self, config_only_mcp_manager_factory):
"""Config loading correctly parses aws_sigv4 auth type and AWS fields."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
@ -398,7 +398,7 @@ class TestMCPServerManagerSigV4:
}
}
manager = MCPServerManager()
manager = config_only_mcp_manager_factory()
await manager.load_servers_from_config(config)
server = next(iter(manager.config_mcp_servers.values()))

View file

@ -40,7 +40,7 @@ from litellm.types.mcp import MCPToolSearchSettings
def _make_tools(specs: list[tuple[str, str]]) -> tuple[Tool, ...]:
return tuple(
Tool(name=name, description=desc, input_schema={"type": "object", "properties": {}}) for name, desc in specs
Tool(name=name, description=desc, inputSchema={"type": "object", "properties": {}}) for name, desc in specs
)
@ -62,17 +62,17 @@ SAMPLE_TOOLS = _make_tools(
FX_TOOL = Tool(
name="treasury-get_rates",
description="Get foreign exchange rates for a currency pair",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
)
WEATHER_TOOL = Tool(
name="weather-forecast",
description="Get the weather forecast for a city",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
)
CALENDAR_TOOL = Tool(
name="calendar-create_event",
description="Create a calendar event",
input_schema={"type": "object", "properties": {}},
inputSchema={"type": "object", "properties": {}},
)
CATALOG = (FX_TOOL, WEATHER_TOOL, CALENDAR_TOOL)
@ -85,23 +85,6 @@ FAKE_VECTORS: dict[str, Vector] = {
}
def _mcp_request_ctx(**overrides):
from types import SimpleNamespace
from mcp.server.context import ServerRequestContext
kwargs = {
"session": SimpleNamespace(),
"lifespan_context": {},
"protocol_version": "2025-06-18",
"method": "",
"params": None,
"request_id": 1,
"meta": None,
"request": None,
}
kwargs.update(overrides)
return ServerRequestContext(**kwargs)
def _paged_params():
@ -586,7 +569,7 @@ class TestCallToolRestApiVirtualTools:
mock_tool = MagicMock()
mock_tool.name = "github-create_issue"
mock_tool.description = "Create a GitHub issue"
mock_tool.input_schema= {"type": "object", "properties": {}}
mock_tool.input_schema = {"type": "object", "properties": {}}
with patch(
"litellm.proxy._experimental.mcp_server.server._list_mcp_tools",
@ -628,7 +611,7 @@ class TestCallToolRestApiVirtualTools:
fake_result = CallToolResult(
content=[TextContent(type="text", text="Issue created")],
is_error=False,
isError=False,
)
with (
@ -678,7 +661,7 @@ class TestCallToolRestApiVirtualTools:
}
)
fake_result = CallToolResult(content=[TextContent(type="text", text="ok")], is_error=False)
fake_result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
with (
patch(
@ -782,7 +765,7 @@ class TestCallToolRestApiVirtualTools:
request = self._make_request(
{"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}}
)
fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], is_error=False)
fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False)
with patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam
"litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search",
new_callable=AsyncMock,
@ -1097,7 +1080,7 @@ class TestDispatchVirtualMcpTool:
)
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
fake = CallToolResult(content=[TextContent(type="text", text="ok")], is_error=False)
fake = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
with (
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
@ -1168,76 +1151,28 @@ class TestDispatchVirtualMcpTool:
class TestCaptureHostProgressCallback:
"""Covers the host progress-forwarding helper extracted from the tool call path."""
@pytest.mark.parametrize("meta", [None, {}, {"traceparent": "trace"}])
def test_returns_none_without_progress(self, _mcp_request_ctx, meta) -> None:
from litellm.proxy._experimental.mcp_server.server import _capture_host_progress_callback
def test_returns_none_when_no_meta(self) -> None:
from types import SimpleNamespace
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
assert _capture_host_progress_callback(SimpleNamespace(meta=None, session=object())) is None
def test_returns_none_when_no_progress_token(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
from types import SimpleNamespace
host = SimpleNamespace(meta=SimpleNamespace(progress_token=None), session=MagicMock())
assert _capture_host_progress_callback(host) is None
def test_returns_callable_when_token_present(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
from types import SimpleNamespace
host = SimpleNamespace(meta=SimpleNamespace(progress_token="tok12345"), session=MagicMock())
assert callable(_capture_host_progress_callback(host))
def test_returns_callable_when_token_is_integer(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
from types import SimpleNamespace
host = SimpleNamespace(meta=SimpleNamespace(progress_token=12345), session=MagicMock())
assert callable(_capture_host_progress_callback(host))
def test_returns_callable_when_token_is_zero(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
from types import SimpleNamespace
host = SimpleNamespace(meta=SimpleNamespace(progress_token=0), session=MagicMock())
assert callable(_capture_host_progress_callback(host))
assert _capture_host_progress_callback(_mcp_request_ctx(meta=meta)) is None
@pytest.mark.asyncio
async def test_forwarded_progress_token_preserves_integer_value(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
@pytest.mark.parametrize("token", ["tok12345", 12345, 0])
async def test_forwards_wire_progress_token(self, _mcp_request_ctx, token) -> None:
from mcp.types import CallToolRequestParams
from litellm.proxy._experimental.mcp_server.server import _capture_host_progress_callback
params = CallToolRequestParams.model_validate(
{"name": "tool", "_meta": {"progressToken": token}}, by_name=False
)
from types import SimpleNamespace
session = AsyncMock()
host = SimpleNamespace(meta=SimpleNamespace(progress_token=12345), session=session)
callback = _capture_host_progress_callback(host)
callback = _capture_host_progress_callback(_mcp_request_ctx(meta=params.meta, session=session))
assert callback is not None
await callback(0.5, 1.0)
session.send_progress_notification.assert_awaited_once_with(
progress_token=12345,
progress=0.5,
total=1.0,
progress_token=token, progress=0.5, total=1.0
)
@ -1245,7 +1180,7 @@ class TestHandleListToolsVirtual:
"""Covers the protocol list_tools early-return when the flag is enabled."""
@pytest.mark.asyncio
async def test_returns_virtual_tools_when_flag_enabled(self) -> None:
async def test_returns_virtual_tools_when_flag_enabled(self, _mcp_request_ctx) -> None:
from litellm.proxy._experimental.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
@ -1269,7 +1204,7 @@ class TestMcpServerToolCallErrorHandling:
isError CallToolResult instead of letting them raise out of the handler."""
@pytest.mark.asyncio
async def test_virtual_tool_error_returns_iserror_not_raised(self) -> None:
async def test_virtual_tool_error_returns_iserror_not_raised(self, _mcp_request_ctx) -> None:
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server import server as srv

View file

@ -285,7 +285,7 @@ class TestToolsetPrefixResolution:
live_tools = [
MCPTool(
name=add_server_prefix_to_name(name, prefix),
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for name in ("read_wiki_contents", "read_wiki_structure", "not_granted")
]
@ -414,7 +414,7 @@ class TestToolsetPrefixResolution:
live_tools = [
MCPTool(
name=add_server_prefix_to_name(name, prefix),
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for name in (granted, sibling)
]
@ -472,7 +472,7 @@ class TestToolsetPrefixResolution:
live_tools = [
MCPTool(
name=add_server_prefix_to_name(granted, prefix),
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]

View file

@ -963,7 +963,7 @@ class TestTestToolsList:
class QuickClient:
async def list_tools(self, raise_on_error=False):
return [MCPTool(name="quick_tool", description="q", input_schema={})]
return [MCPTool(name="quick_tool", description="q", inputSchema={})]
async def fake_execute(
request,
@ -1008,7 +1008,7 @@ class TestTestToolsList:
async def list_tools(self, raise_on_error=False):
await asyncio.sleep(0.2)
return [MCPTool(name="slow_tool", description="s", input_schema={})]
return [MCPTool(name="slow_tool", description="s", inputSchema={})]
async def fake_execute(
request,
@ -1512,7 +1512,7 @@ class TestListToolsRestAPI:
MCPTool(
name="first_page_tool",
description="First page tool",
input_schema={},
inputSchema={},
)
],
nextCursor="page-2",
@ -1522,7 +1522,7 @@ class TestListToolsRestAPI:
MCPTool(
name="second_page_tool",
description="Second page tool",
input_schema={},
inputSchema={},
)
]
),
@ -3198,7 +3198,7 @@ class TestGetToolsForSingleServer:
def __init__(self, name, description):
self.name = name
self.description = description
self.input_schema= {}
self.input_schema = {}
mock_tools = [
MockTool("tool1", "First tool"),
@ -3259,7 +3259,7 @@ class TestGetToolsForSingleServer:
def __init__(self, name, description):
self.name = name
self.description = description
self.input_schema= {}
self.input_schema = {}
mock_tools = [
MockTool("tool1", "First tool"),
@ -3307,7 +3307,7 @@ class TestGetToolsForSingleServer:
def __init__(self, name, description):
self.name = name
self.description = description
self.input_schema= {}
self.input_schema = {}
mock_tools = [
MockTool("tool1", "First tool"),
@ -3360,7 +3360,7 @@ class TestGetToolsForSingleServer:
def __init__(self, name, description):
self.name = name
self.description = description
self.input_schema= {}
self.input_schema = {}
mock_tools = [
MockTool("tool1", "First tool"),
@ -3413,7 +3413,7 @@ class TestGetToolsForSingleServer:
def __init__(self, name, description):
self.name = name
self.description = description
self.input_schema= {}
self.input_schema = {}
mock_tools = [
MockTool("tool1", "First tool"),
@ -3475,7 +3475,7 @@ class TestGetToolsForSingleServer:
def __init__(self, name):
self.name = name
self.description = name
self.input_schema= {}
self.input_schema = {}
mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")]
@ -4138,7 +4138,7 @@ class TestToolResponseMcpInfoEnrichment:
MCPTool(
name="get_issue",
description="Fetch a Jira issue",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -4150,6 +4150,12 @@ class TestToolResponseMcpInfoEnrichment:
"alias": "atlassian",
}
from fastapi.encoders import jsonable_encoder
wire = jsonable_encoder(result[0])
assert wire["inputSchema"] == {"type": "object"}
assert wire["mcp_info"] == result[0].mcp_info
def test_alias_none_is_explicit_in_mcp_info(self):
from mcp.types import Tool as MCPTool
@ -4168,7 +4174,7 @@ class TestToolResponseMcpInfoEnrichment:
MCPTool(
name="ping",
description="Ping",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -4210,8 +4216,8 @@ class TestRestListToolsetFiltering:
stub_server.mcp_info = {"server_name": "stubtools"}
upstream_tools = [
MCPTool(name="lookup_status", input_schema={"type": "object"}),
MCPTool(name="delete_everything", input_schema={"type": "object"}),
MCPTool(name="lookup_status", inputSchema={"type": "object"}),
MCPTool(name="delete_everything", inputSchema={"type": "object"}),
]
key_object_permission = MagicMock()

View file

@ -42,52 +42,52 @@ async def test_semantic_filter_basic_filtering():
MCPTool(
name="gmail_send",
description="Send an email via Gmail",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="outlook_send",
description="Send an email via Outlook",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="calendar_create",
description="Create a calendar event",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="calendar_update",
description="Update a calendar event",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="email_read",
description="Read emails from inbox",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="email_delete",
description="Delete an email",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="calendar_delete",
description="Delete a calendar event",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="email_search",
description="Search for emails",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="calendar_list",
description="List calendar events",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
MCPTool(
name="email_forward",
description="Forward an email to someone",
input_schema={"type": "object"},
inputSchema={"type": "object"},
),
]
@ -170,7 +170,7 @@ async def test_semantic_filter_top_k_limiting():
MCPTool(
name=f"tool_{i}",
description=f"Tool number {i} for testing",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(20)
]
@ -228,7 +228,7 @@ async def test_semantic_filter_disabled():
tools = [
MCPTool(
name=f"tool_{i}", description=f"Tool {i}", input_schema={"type": "object"}
name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}
)
for i in range(10)
]
@ -375,7 +375,7 @@ async def test_semantic_filter_hook_triggers_on_completion():
# Prepare data - completion request with tools
tools = [
MCPTool(
name=f"tool_{i}", description=f"Tool {i}", input_schema={"type": "object"}
name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}
)
for i in range(10)
]
@ -508,7 +508,7 @@ async def test_semantic_filter_hook_preserves_native_tools():
MCPTool(
name=f"mcp_tool_{i}",
description=f"MCP tool {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(5)
]
@ -624,7 +624,7 @@ async def test_semantic_filter_hook_all_native_tools():
MCPTool(
name="some_mcp_tool",
description="An MCP tool",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
@ -741,7 +741,7 @@ async def test_semantic_filter_hook_responses_api_name_collision():
MCPTool(
name="github-search",
description="Search GitHub repos",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
]
filter_instance._build_router(mcp_tools)
@ -836,7 +836,7 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools():
MCPTool(
name=f"srv-tool_{i}",
description=f"Registry tool {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(5)
]
@ -958,7 +958,7 @@ async def test_semantic_filter_hook_narrows_mcp_reference_for_chat_completions()
MCPTool(
name=f"srv-tool_{i}",
description=f"Registry tool {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(5)
]
@ -1065,7 +1065,7 @@ async def test_semantic_filter_hook_zero_matches_exposes_all_tools_on_both_paths
MCPTool(
name=f"srv-tool_{i}",
description=f"Registry tool {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(3)
]
@ -1182,7 +1182,7 @@ async def test_semantic_filter_hook_filters_expanded_tools_with_string_input():
MCPTool(
name=f"srv-tool_{i}",
description=f"Registry tool {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(5)
]
@ -1326,12 +1326,12 @@ async def test_semantic_filter_hook_preserves_tool_order():
mcp_tool_a = MCPTool(
name="github-search",
description="Search GitHub",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
mcp_tool_b = MCPTool(
name="github-issue",
description="Create GitHub issue",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
filter_instance._build_router([mcp_tool_a, mcp_tool_b])
@ -1683,7 +1683,7 @@ async def test_semantic_filter_fails_closed_on_query_time_context_window_error()
filter_instance = _make_context_window_filter(state)
tools = [
MCPTool(name=f"tool_{i}", description=f"Tool {i}", input_schema={"type": "object"})
MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"})
for i in range(5)
]
filter_instance._build_router(tools)
@ -1716,7 +1716,7 @@ async def test_semantic_filter_records_build_time_context_window_error():
filter_instance = _make_context_window_filter(state)
tools = [
MCPTool(name=f"tool_{i}", description=f"Tool {i}", input_schema={"type": "object"})
MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"})
for i in range(5)
]
filter_instance._build_router(tools)
@ -1750,7 +1750,7 @@ async def test_semantic_filter_hook_fails_closed_on_context_window_error():
filter_instance = _make_context_window_filter(state)
tools = [
MCPTool(name=f"tool_{i}", description=f"Tool {i}", input_schema={"type": "object"})
MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"})
for i in range(5)
]
filter_instance._build_router(tools)
@ -1798,7 +1798,7 @@ async def test_semantic_filter_hook_fails_closed_on_expanded_tools_context_windo
filter_instance = _make_context_window_filter(state)
registry_tools = [
MCPTool(name=f"srv-tool_{i}", description=f"Registry tool {i}", input_schema={"type": "object"})
MCPTool(name=f"srv-tool_{i}", description=f"Registry tool {i}", inputSchema={"type": "object"})
for i in range(5)
]
filter_instance._build_router(registry_tools)
@ -1862,7 +1862,7 @@ async def test_semantic_filter_hook_ignores_build_error_for_native_only_tools():
filter_instance = _make_context_window_filter(state)
mcp_tools = [
MCPTool(name=f"tool_{i}", description=f"Tool {i}", input_schema={"type": "object"})
MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"})
for i in range(3)
]
filter_instance._build_router(mcp_tools)
@ -2019,7 +2019,7 @@ def _linear_issue_tool():
return MCPTool(
name="linear_stub-get_issue",
description="Get a Linear issue (ticket) by its identifier such as LIT-1234",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
@ -2027,7 +2027,7 @@ def _linear_list_tool():
return MCPTool(
name="linear_stub-list_issues",
description="List Linear issues (tickets) in the workspace",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
@ -2035,7 +2035,7 @@ def _weather_tool():
return MCPTool(
name="weather_stub-get_weather",
description="Get the current weather conditions for a city",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
@ -2135,8 +2135,8 @@ async def test_request_time_context_window_error_is_request_scoped():
state = {"raise_context_error": True}
filter_instance = _make_context_window_filter(state)
tools = [
MCPTool(name="tool_a", description="Tool A", input_schema={"type": "object"}),
MCPTool(name="tool_b", description="Tool B", input_schema={"type": "object"}),
MCPTool(name="tool_a", description="Tool A", inputSchema={"type": "object"}),
MCPTool(name="tool_b", description="Tool B", inputSchema={"type": "object"}),
]
with pytest.raises(SemanticToolFilterContextWindowError):
@ -2171,7 +2171,7 @@ async def test_foreign_index_routes_cannot_displace_available_tools():
MCPTool(
name=f"other_user-linear_tool_{i}",
description=f"Get a Linear issue variant {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(6)
]
@ -2180,7 +2180,7 @@ async def test_foreign_index_routes_cannot_displace_available_tools():
my_kanban = MCPTool(
name="mine-kanban_board",
description="Manage kanban board cards",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
filtered = await filter_instance.filter_tools(
query="what is Linear ticket LIT-3794 about",
@ -2204,7 +2204,7 @@ async def test_top_k_above_router_default_is_respected():
MCPTool(
name=f"linear_stub-tool_{i}",
description=f"Work with Linear issues part {i}",
input_schema={"type": "object"},
inputSchema={"type": "object"},
)
for i in range(6)
]

View file

@ -268,8 +268,8 @@ class TestIsToolNamePrefixedBoundary:
def _stub_tools() -> List[MCPTool]:
return [
MCPTool(name="get_repo", description="", input_schema={"type": "object"}),
MCPTool(name="list_issues", description="", input_schema={"type": "object"}),
MCPTool(name="get_repo", description="", inputSchema={"type": "object"}),
MCPTool(name="list_issues", description="", inputSchema={"type": "object"}),
]

View file

@ -269,3 +269,17 @@ class TestBuildSyntheticMcpRequest:
)
assert request.headers.get("x-user-email") == "alice@corp.example"
@pytest.mark.parametrize("field", ["structuredContent", "structured_content"])
def test_structured_content_redaction_updates_shared_dictionary(field):
from litellm.proxy._experimental.mcp_server.utils import (
mcp_tool_result_structured_content,
set_mcp_tool_result_structured_content,
)
result = {field: {"secret": "sensitive"}, "content": []}
logging_reference = result
assert set_mcp_tool_result_structured_content(result, {"secret": "[REDACTED]"}) is True
assert mcp_tool_result_structured_content(logging_reference) == {"secret": "[REDACTED]"}
assert set(result) == {field, "content"}

View file

@ -51,7 +51,9 @@ class TestCiscoAIDefenseMCPMode:
@pytest.mark.asyncio
async def test_mcp_mode_inspects_mcp_request(self):
g = _make_guardrail(inspection_type="mcp", event_hook="pre_mcp_call")
data = _mcp_request(name="send_email", args={"to": "x@y.com"}, litellm_call_id="call-1")
data = _mcp_request(
name="send_email", args={"to": "x@y.com"}, litellm_call_id="call-1"
)
post_mock = AsyncMock(return_value=_safe_response(url=MCP_URL))
with _patch_inspection_post(g, post_mock):
result = await g.async_pre_call_hook(
@ -76,7 +78,9 @@ class TestCiscoAIDefenseMCPMode:
async def test_mcp_mode_blocks_violation(self):
g = _make_guardrail(inspection_type="mcp", event_hook="pre_mcp_call")
data = _mcp_request(name="leak_secrets", args={"target": "evil"})
with _patch_inspection_post(g, AsyncMock(return_value=_violation_response(url=MCP_URL))):
with _patch_inspection_post(
g, AsyncMock(return_value=_violation_response(url=MCP_URL))
):
with pytest.raises(HTTPException) as exc:
await g.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
@ -161,7 +165,9 @@ class TestCiscoAIDefenseMCPMode:
call_type="mcp_call",
)
forwarded = ProxyLogging(user_api_key_cache=UserApiKeyCache())._convert_mcp_hook_response_to_kwargs(
forwarded = ProxyLogging(
user_api_key_cache=UserApiKeyCache()
)._convert_mcp_hook_response_to_kwargs(
response_data=result, original_kwargs={"arguments": dict(original_args)}
)
assert forwarded["arguments"] == sanitized_args, (
@ -173,10 +179,14 @@ class TestCiscoAIDefenseMCPMode:
@pytest.mark.asyncio
async def test_mcp_response_hook_inspects_tool_output(self):
g = _make_guardrail(inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"])
g = _make_guardrail(
inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]
)
response_obj = _mcp_response(
SimpleNamespace(content=[{"type": "text", "text": "Here is the secret API key abc123"}])
SimpleNamespace(
content=[{"type": "text", "text": "Here is the secret API key abc123"}]
)
)
post_mock = AsyncMock(return_value=_safe_response(url=MCP_URL))
@ -205,7 +215,9 @@ class TestCiscoAIDefenseMCPMode:
"name": "lookup_secret",
"arguments": {"key": "production"},
}
assert sent_payload["result"]["content"][0]["text"] == ("Here is the secret API key abc123")
assert sent_payload["result"]["content"][0]["text"] == (
"Here is the secret API key abc123"
)
assert "request" not in sent_payload
assert "metadata" not in sent_payload
@ -213,8 +225,12 @@ class TestCiscoAIDefenseMCPMode:
async def test_mcp_response_hook_blocks_violation(self):
from litellm.types.mcp import MCPPostCallResponseObject
g = _make_guardrail(inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"])
response_obj = _mcp_response(SimpleNamespace(content=[{"type": "text", "text": "leaked"}]))
g = _make_guardrail(
inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]
)
response_obj = _mcp_response(
SimpleNamespace(content=[{"type": "text", "text": "leaked"}])
)
post_mock = AsyncMock(return_value=_violation_response(url=MCP_URL))
with _patch_inspection_post(g, post_mock):
@ -241,7 +257,9 @@ class TestCiscoAIDefenseMCPMode:
@pytest.mark.asyncio
async def test_mcp_response_hook_skipped_in_chat_mode(self):
g = _make_guardrail()
response_obj = _mcp_response(SimpleNamespace(content=[{"type": "text", "text": "hi"}]))
response_obj = _mcp_response(
SimpleNamespace(content=[{"type": "text", "text": "hi"}])
)
post_mock = AsyncMock()
with _patch_inspection_post(g, post_mock):
@ -273,7 +291,11 @@ class TestCiscoAIDefenseMCPMode:
@pytest.mark.asyncio
async def test_mcp_response_hook_runs_with_pre_mcp_call_only(self):
g = _make_guardrail(inspection_type="mcp", event_hook="pre_mcp_call")
response_obj = _mcp_response(SimpleNamespace(content=[{"type": "text", "text": "would have been scanned"}]))
response_obj = _mcp_response(
SimpleNamespace(
content=[{"type": "text", "text": "would have been scanned"}]
)
)
post_mock = AsyncMock(return_value=_safe_response(url=MCP_URL))
with _patch_inspection_post(g, post_mock):
@ -295,18 +317,26 @@ class TestCiscoAIDefenseMCPMode:
[("safe", False), ("violation", True)],
)
@pytest.mark.asyncio
async def test_mcp_response_hook_handles_raw_list_content(self, cisco_response_kind, expected_block):
async def test_mcp_response_hook_handles_raw_list_content(
self, cisco_response_kind, expected_block
):
from litellm.types.mcp import MCPPostCallResponseObject
g = _make_guardrail(inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"])
g = _make_guardrail(
inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]
)
text_content = (
"exfiltrated data: ..." if cisco_response_kind == "violation" else "Here is the secret API key abc123"
"exfiltrated data: ..."
if cisco_response_kind == "violation"
else "Here is the secret API key abc123"
)
response_obj = _mcp_response([{"type": "text", "text": text_content}])
cisco_resp = (
_violation_response(url=MCP_URL) if cisco_response_kind == "violation" else _safe_response(url=MCP_URL)
_violation_response(url=MCP_URL)
if cisco_response_kind == "violation"
else _safe_response(url=MCP_URL)
)
post_mock = AsyncMock(return_value=cisco_resp)
kwargs = {
@ -324,7 +354,8 @@ class TestCiscoAIDefenseMCPMode:
)
assert post_mock.called, (
"MCP response inspect was silently skipped for raw-list shape — _normalize_mcp_response failed."
"MCP response inspect was silently skipped for raw-list "
"shape — _normalize_mcp_response failed."
)
assert post_mock.call_args.kwargs["url"] == MCP_URL
@ -351,12 +382,14 @@ class TestCiscoAIDefenseMCPMode:
from litellm.types.mcp import MCPPostCallResponseObject
g = _make_guardrail(inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"])
g = _make_guardrail(
inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]
)
real_result = CallToolResult(
content=[TextContent(type="text", text="leak 9045629876")],
structured_content={"patient": {"ssn": "123-45-6789"}},
is_error=False,
structuredContent={"patient": {"ssn": "123-45-6789"}},
isError=False,
)
wrapped = MCPPostCallResponseObject(
mcp_tool_call_response=real_result,
@ -364,8 +397,12 @@ class TestCiscoAIDefenseMCPMode:
)
assert isinstance(wrapped.mcp_tool_call_response, list)
assert all(isinstance(item, tuple) and len(item) == 2 for item in wrapped.mcp_tool_call_response), (
"Pydantic coercion shape changed — update the normalizer to match the new wire format."
assert all(
isinstance(item, tuple) and len(item) == 2
for item in wrapped.mcp_tool_call_response
), (
"Pydantic coercion shape changed — update the normalizer to "
"match the new wire format."
)
post_mock = AsyncMock(return_value=_safe_response(url=MCP_URL))
@ -404,7 +441,9 @@ class TestCiscoAIDefenseMCPMode:
f"``content`` field."
)
assert content_items[0].get("type") == "text"
assert sent_payload["result"]["structuredContent"] == {"patient": {"ssn": "123-45-6789"}}
assert sent_payload["result"]["structuredContent"] == {
"patient": {"ssn": "123-45-6789"}
}
assert sent_payload["result"]["isError"] is False
assert sent_payload["id"] == "real-wire-call"
assert sent_payload["method"] == "tools/call"
@ -519,16 +558,20 @@ class TestCiscoAIDefenseRedactListShape:
original_response = CallToolResult(
content=[TextContent(type="text", text="SSN: 123-45-6789")],
structured_content={"patient": {"ssn": "123-45-6789"}},
is_error=False,
structuredContent={"patient": {"ssn": "123-45-6789"}},
isError=False,
)
wrapper = MCPPostCallResponseObject(
mcp_tool_call_response=original_response,
hidden_params=HiddenParams(),
)
g = _make_guardrail(inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"])
with _patch_inspection_post(g, AsyncMock(return_value=self._violation_with_redact_response())):
g = _make_guardrail(
inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]
)
with _patch_inspection_post(
g, AsyncMock(return_value=self._violation_with_redact_response())
):
await g.async_post_mcp_tool_call_hook(
kwargs={
"name": "leak",
@ -556,7 +599,9 @@ class TestCiscoAIDefenseMcpInputRedactionFallback:
@pytest.mark.asyncio
async def test_single_string_arg_is_rewritten(self):
g = _make_guardrail(inspection_type="mcp", event_hook="pre_mcp_call")
data = _mcp_request(name="search", args={"query": "my SSN is 123-45-6789", "limit": 10})
data = _mcp_request(
name="search", args={"query": "my SSN is 123-45-6789", "limit": 10}
)
cisco = _redact_response(sanitized_text="my SSN is [REDACTED]", url=MCP_URL)
with _patch_inspection_post(g, AsyncMock(return_value=cisco)):
result = await g.async_pre_call_hook(
@ -611,6 +656,7 @@ class TestCiscoAIDefenseMcpInputRedactionFallback:
class TestCiscoAIDefenseMCPBlockingContract:
@pytest.mark.asyncio
async def test_block_response_survives_dispatcher_contract(self):
from litellm.litellm_core_utils.litellm_logging import Logging
@ -624,8 +670,8 @@ class TestCiscoAIDefenseMCPBlockingContract:
)
raw_response = CallToolResult(
content=[TextContent(type="text", text="exfiltrated")],
structured_content={"result": "exfiltrated"},
is_error=False,
structuredContent={"result": "exfiltrated"},
isError=False,
)
response_obj = MCPPostCallResponseObject(
mcp_tool_call_response=raw_response,
@ -672,6 +718,7 @@ class TestCiscoAIDefenseMCPBlockingContract:
class TestCiscoAIDefenseJsonRpcSuccessEnvelope:
@staticmethod
def _cisco_mcp_envelope(*, is_safe: bool, action: str = "Block") -> Response:
return _mock_inspect_response(
@ -707,8 +754,12 @@ class TestCiscoAIDefenseJsonRpcSuccessEnvelope:
],
)
@pytest.mark.asyncio
async def test_mcp_jsonrpc_envelope_respects_verdict(self, is_safe, action, should_block):
g = _make_guardrail(name="cisco-mcp", inspection_type="mcp", event_hook="pre_mcp_call")
async def test_mcp_jsonrpc_envelope_respects_verdict(
self, is_safe, action, should_block
):
g = _make_guardrail(
name="cisco-mcp", inspection_type="mcp", event_hook="pre_mcp_call"
)
data = _mcp_request(
name="ask_question",
args={
@ -718,7 +769,9 @@ class TestCiscoAIDefenseJsonRpcSuccessEnvelope:
)
with _patch_inspection_post(
g,
AsyncMock(return_value=self._cisco_mcp_envelope(is_safe=is_safe, action=action)),
AsyncMock(
return_value=self._cisco_mcp_envelope(is_safe=is_safe, action=action)
),
):
if should_block:
with pytest.raises(HTTPException) as exc:
@ -730,7 +783,10 @@ class TestCiscoAIDefenseJsonRpcSuccessEnvelope:
)
assert exc.value.status_code == 400
assert exc.value.detail["surface"] == "mcp"
assert exc.value.detail["event_id"] == "645d9d22-b016-47e0-a12c-9d587fb11c57"
assert (
exc.value.detail["event_id"]
== "645d9d22-b016-47e0-a12c-9d587fb11c57"
)
else:
result = await g.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),