diff --git a/.github/workflows/test-mcp-dependency-resolution.yml b/.github/workflows/test-mcp-dependency-resolution.yml index a0c8057e28b..251dffccd4f 100644 --- a/.github/workflows/test-mcp-dependency-resolution.yml +++ b/.github/workflows/test-mcp-dependency-resolution.yml @@ -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 diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 93ffcbe0586..9d6b0194df9 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -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: | diff --git a/litellm/experimental_mcp_client/Readme.md b/litellm/experimental_mcp_client/Readme.md index 4fbd624369c..7807f6a7379 100644 --- a/litellm/experimental_mcp_client/Readme.md +++ b/litellm/experimental_mcp_client/Readme.md @@ -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 diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index fa4d76ecbed..a1f0e5c0830 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 32bbfc7d913..ff482b80b50 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -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 []: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d8890ccad56..29ca2d6a064 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index f57ad4bfad5..361d8d5ae31 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -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): diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ad8c721db7a..f67b50368e9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py index 7bbe785b4fa..15b4f713a50 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py @@ -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 diff --git a/scripts/check_mcp_sdk_install.py b/scripts/check_mcp_sdk_install.py index 9b5106118e7..f5ab51b2f55 100644 --- a/scripts/check_mcp_sdk_install.py +++ b/scripts/check_mcp_sdk_install.py @@ -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( diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py index eff32f27aec..ca3e25949ba 100644 --- a/tests/mcp_tests/conftest.py +++ b/tests/mcp_tests/conftest.py @@ -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 diff --git a/tests/mcp_tests/mcp_server.py b/tests/mcp_tests/mcp_server.py index eba7cae1bca..f38b6a02139 100644 --- a/tests/mcp_tests/mcp_server.py +++ b/tests/mcp_tests/mcp_server.py @@ -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() diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index 7a48c366003..eb6f78b57a1 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -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): diff --git a/tests/mcp_tests/test_mcp_auth_priority.py b/tests/mcp_tests/test_mcp_auth_priority.py index 7ae0f59afe5..21a89d7ffcc 100644 --- a/tests/mcp_tests/test_mcp_auth_priority.py +++ b/tests/mcp_tests/test_mcp_auth_priority.py @@ -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 diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index 8e5a0cd30b9..6438525706a 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -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"), diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 04218e6d0ce..ed8829945e5 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -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"}}, }, diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 45be1f72207..94cf35b675d 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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": { diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index 018a09b5e89..99c03b3438d 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -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 diff --git a/tests/mcp_tests/test_semantic_tool_filter_e2e.py b/tests/mcp_tests/test_semantic_tool_filter_e2e.py index d2ebdb3a4dd..aa25c98107e 100644 --- a/tests/mcp_tests/test_semantic_tool_filter_e2e.py +++ b/tests/mcp_tests/test_semantic_tool_filter_e2e.py @@ -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"} ), ] diff --git a/tests/pass_through_tests/test_mcp_routes.py b/tests/pass_through_tests/test_mcp_routes.py index 9a4d4f9e865..e9d18193e7c 100644 --- a/tests/pass_through_tests/test_mcp_routes.py +++ b/tests/pass_through_tests/test_mcp_routes.py @@ -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 diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index f1f459fbc5b..ad58ce5f00f 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -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() diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/test_litellm/experimental_mcp_client/test_tools.py index 55eccbb8fbf..6645b06664d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_tools.py +++ b/tests/test_litellm/experimental_mcp_client/test_tools.py @@ -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"], diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/test_litellm/integrations/arize/test_arize_utils.py index 165b7bc94d4..167b083e147 100644 --- a/tests/test_litellm/integrations/arize/test_arize_utils.py +++ b/tests/test_litellm/integrations/arize/test_arize_utils.py @@ -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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index 9a66f130d24..76e92efd31a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index 9dd88ff18bd..77e9b987e74 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -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", "") 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 " - assert returned.structured_content== {"contact": {"email": ""}, "balance": 42.0} + assert returned.structured_content == {"contact": {"email": ""}, "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", "") 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": ""}]} + assert returned.structured_content == {"records": [{"email": ""}]} 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", "") 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", "") 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", "") 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", "") 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 " - 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} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py index 333d4c98899..e3437bf16f6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py @@ -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 = { diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index f1ca0f46fd2..46ecd4df716 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index ca9f774e8f6..93b894f7645 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -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] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py index 6c6f996977a..86748d99063 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py @@ -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", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index b5260aaa4e9..3f5d4ad83ea 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py index 90ec1ab9061..167847afe1f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_tool_conversion.py @@ -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 ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 5594cee8ca5..41287c122a0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index fbecdd60a26..dc1eed9ed7f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -85,22 +85,6 @@ def _reload_mcp_manager_module(): return reloaded -def _mcp_request_ctx(**overrides): - from mcp.server.context import ServerRequestContext - from types import SimpleNamespace - - 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.fixture(autouse=True) @@ -438,10 +422,10 @@ class TestMCPServerManager: assert "gateway-client" in dump assert "https://org-idp.example/oauth2/token" in dump - async def test_load_servers_from_config_warns_on_invalid_alias(self, caplog): + async def test_load_servers_from_config_warns_on_invalid_alias(self, config_only_mcp_manager_factory, caplog): """Invalid aliases from config should emit warnings during load.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "validserver": { "alias": "bad/name", @@ -456,10 +440,10 @@ class TestMCPServerManager: assert any("invalid alias 'bad/name'" in message for message in caplog.messages) @pytest.mark.asyncio - async def test_load_servers_from_config_accepts_valid_alias(self, caplog): + async def test_load_servers_from_config_accepts_valid_alias(self, config_only_mcp_manager_factory, caplog): """Valid aliases should be accepted and populate the registry.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "validserver": { "alias": "friendly_alias", @@ -1229,8 +1213,8 @@ class TestMCPServerManager: assert server.scopes == ["read"] @pytest.mark.asyncio - async def test_load_servers_from_config_non_oauth2_needs_no_flow(self): - manager = MCPServerManager() + async def test_load_servers_from_config_non_oauth2_needs_no_flow(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() config = { "apiserver": { "url": "https://example.com/mcp", @@ -1276,10 +1260,10 @@ class TestMCPServerManager: assert not any("oauth2_id_jag" in message for message in caplog.messages) @pytest.mark.asyncio - async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, monkeypatch, caplog): + async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, config_only_mcp_manager_factory, monkeypatch, caplog): self._clear_sso_env(monkeypatch) monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "api_key_server": { "url": "https://example.com/mcp", @@ -1416,9 +1400,9 @@ class TestMCPServerManager: assert server.is_dcr_bridge is False @pytest.mark.asyncio - async def test_load_servers_from_config_coerces_cost_string_to_float(self): + async def test_load_servers_from_config_coerces_cost_string_to_float(self, config_only_mcp_manager_factory): """YAML 1.1 parses `7e-05` as a string; ingest must coerce it to float.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "google_maps": { "url": "https://example.com/mcp", @@ -1442,9 +1426,9 @@ class TestMCPServerManager: assert isinstance(cost_info["tool_name_to_cost_per_query"]["geocode"], float) @pytest.mark.asyncio - async def test_load_servers_from_config_sets_token_endpoint_auth_method(self): + async def test_load_servers_from_config_sets_token_endpoint_auth_method(self, config_only_mcp_manager_factory): """token_endpoint_auth_method from config is carried onto the MCPServer (LIT-4091).""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "basic_provider": { "url": "https://example.com/mcp", @@ -1686,7 +1670,7 @@ class TestMCPServerManager: ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -1890,7 +1874,7 @@ class TestMCPServerManager: never wrapped as MCPUpstreamAuthError or replaced by error_tool_result.""" server = self._passthrough_call_server(MCPAuth.true_passthrough, server_id=f"pt-ok-{is_error}") manager = MCPServerManager() - expected = CallToolResult(content=[], is_error=is_error) + expected = CallToolResult(content=[], isError=is_error) mock_client = AsyncMock() mock_client.call_tool = AsyncMock(return_value=expected) manager._create_mcp_client = AsyncMock(return_value=mock_client) @@ -1940,7 +1924,7 @@ class TestMCPServerManager: ) manager = MCPServerManager() mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) manager._create_mcp_client = AsyncMock(return_value=mock_client) result = await manager._call_regular_mcp_tool( @@ -3111,7 +3095,7 @@ class TestMCPServerManager: ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -3170,7 +3154,7 @@ class TestMCPServerManager: assert _should_strip_caller_authorization(mcp_server=server, raw_headers=None, user_api_key_auth=None) is True mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = "unset" async def capture_create_mcp_client( @@ -3238,7 +3222,7 @@ class TestMCPServerManager: ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -3295,7 +3279,7 @@ class TestMCPServerManager: ) mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured_extra_headers = None async def capture_create_mcp_client( @@ -3330,7 +3314,7 @@ class TestMCPServerManager: async def _capture_call_extra_headers(self, server, oauth2_headers, raw_headers, user_api_key_auth): manager = MCPServerManager() mock_client = AsyncMock() - mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) captured = {"extra_headers": "unset"} async def capture_create_mcp_client( @@ -4559,9 +4543,7 @@ class TestMCPServerManager: @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2]) @pytest.mark.parametrize("is_byok", [False, True]) @pytest.mark.parametrize("scheme", ["http", "https"]) - async def test_openapi_health_loads_spec_without_mcp_handshake( - self, respx_mock, monkeypatch, auth_type, is_byok, scheme - ): + async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( @@ -4611,28 +4593,14 @@ class TestMCPServerManager: @pytest.mark.parametrize( ("failure", "expected_status", "expected_error"), [ - ( - httpx.Response(401, text="secret response content"), - "unhealthy", - "OpenAPI specification request failed (HTTP 401)", - ), + (httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"), (httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"), (httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"), - ( - httpx.ConnectError("secret network details"), - "unhealthy", - "OpenAPI specification could not be loaded (ConnectError)", - ), - ( - httpx.Response(200, text="secret invalid JSON body"), - "unhealthy", - "OpenAPI specification could not be loaded (JSONDecodeError)", - ), + (httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"), + (httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"), ], ) - async def test_openapi_health_reports_safe_failures( - self, respx_mock, monkeypatch, failure, expected_status, expected_error - ): + async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( @@ -5167,15 +5135,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, - method, - operation, - base_url, - headers=None, - server_label=None, - relays_upstream_auth=False, - auth_type=None, - upstream_token_header=None, + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers captured["server_label"] = server_label @@ -5260,15 +5221,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, - method, - operation, - base_url, - headers=None, - server_label=None, - relays_upstream_auth=False, - auth_type=None, - upstream_token_header=None, + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers @@ -5540,7 +5494,7 @@ class TestMCPServerManager: upstream_tool = MCPTool( name="send_email", description="Send an email", - input_schema={}, + inputSchema={}, ) manager._fetch_tools_with_timeout = AsyncMock(return_value=[upstream_tool]) @@ -6072,12 +6026,12 @@ class TestMCPServerManager: t1 = MCPTool( name="create_issue", description="", - input_schema={}, + inputSchema={}, ) t2 = MCPTool( name="close_issue", description="", - input_schema={}, + inputSchema={}, ) # Do not add prefix in returned objects @@ -6111,7 +6065,7 @@ class TestMCPServerManager: base_tool = MCPTool( name="create_zap", description="", - input_schema={}, + inputSchema={}, ) _ = manager._create_prefixed_tools([base_tool], server, add_prefix=False) @@ -7939,9 +7893,9 @@ class TestMCPServerTimestamps: assert client.timeout == 0.0 @pytest.mark.asyncio - async def test_load_servers_from_config_preserves_timeout(self): + async def test_load_servers_from_config_preserves_timeout(self, config_only_mcp_manager_factory): """timeout from proxy config is loaded into MCPServer.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() config = { "my_server": { "url": "https://example.com/mcp", @@ -8354,9 +8308,9 @@ class TestMCPServerManagerUpstreamInstructionsCache: assert manager._upstream_initialize_instructions_by_server_id.get("srv") is None @pytest.mark.asyncio - async def test_load_servers_from_config_clears_cache(self): + async def test_load_servers_from_config_clears_cache(self, config_only_mcp_manager_factory): """Reloading config clears any previously cached upstream instructions.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() manager._upstream_initialize_instructions_by_server_id["old"] = "stale" await manager.load_servers_from_config( mcp_servers_config={ @@ -8369,9 +8323,9 @@ class TestMCPServerManagerUpstreamInstructionsCache: assert manager._upstream_initialize_instructions_by_server_id.get("old") is None @pytest.mark.asyncio - async def test_load_servers_reads_instructions_from_config(self): + async def test_load_servers_reads_instructions_from_config(self, config_only_mcp_manager_factory): """instructions field from YAML config is persisted on the MCPServer.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( mcp_servers_config={ "srv_a": { @@ -9806,7 +9760,7 @@ class TestMCPToolsListAuthSurfacing: manager.get_mcp_server_by_id = MagicMock( side_effect=lambda server_id: {"good": good, "bad": bad}.get(server_id) ) - good_tool = MCPTool(name="good-do_thing", description="do thing", input_schema={}) + good_tool = MCPTool(name="good-do_thing", description="do thing", inputSchema={}) async def fake_get_tools(server, **kwargs): if server.server_id == "bad": @@ -9921,7 +9875,7 @@ class TestOBOCallToolRetry: @pytest.mark.asyncio async def test_upstream_401_invalidates_and_retries_once(self): manager = self._manager() - success = CallToolResult(content=[], is_error=False) + success = CallToolResult(content=[], isError=False) first = _RetryFakeClient(raises=_UpstreamAuthError(401)) retry = _RetryFakeClient(result=success) manager._create_mcp_client = AsyncMock(return_value=retry) @@ -9952,7 +9906,7 @@ class TestOBOCallToolRetry: ) manager = self._manager() - success = CallToolResult(content=[], is_error=False) + success = CallToolResult(content=[], isError=False) first = _RetryFakeClient(raises=_UpstreamAuthError(401)) retry = _RetryFakeClient(result=success) manager._create_mcp_client = AsyncMock(return_value=retry) @@ -9991,7 +9945,7 @@ class TestOBOCallToolRetry: """An oauth2_id_jag tool call with a subject token must take the invalidate-and-retry branch of _call_regular_mcp_tool, not the plain single call, so an upstream 401 re-exchanges.""" manager = self._manager() - success = CallToolResult(content=[], is_error=False) + success = CallToolResult(content=[], isError=False) first = _RetryFakeClient(raises=_UpstreamAuthError(401)) retry = _RetryFakeClient(result=success) manager._create_mcp_client = AsyncMock(side_effect=[first, retry]) @@ -10106,7 +10060,7 @@ class TestOBOConcurrencyLimit: await release.wait() finally: inflight["current"] -= 1 - return CallToolResult(content=[], is_error=False) + return CallToolResult(content=[], isError=False) manager = MCPServerManager() manager._create_mcp_client = AsyncMock(return_value=_ConcurrencyRecordingClient()) @@ -10320,7 +10274,7 @@ async def test_aggregate_list_still_absorbs_step_up_challenged_server(): ca = MCPServer(server_id="ca", name="ca", transport=MCPTransport.http) manager.get_allowed_mcp_servers = AsyncMock(return_value=["good", "ca"]) manager.get_mcp_server_by_id = MagicMock(side_effect=lambda server_id: {"good": good, "ca": ca}.get(server_id)) - good_tool = MCPTool(name="good-do_thing", description="do thing", input_schema={}) + good_tool = MCPTool(name="good-do_thing", description="do thing", inputSchema={}) async def fake_get_tools(server, **kwargs): if server.server_id == "ca": @@ -11068,7 +11022,7 @@ class TestServerToolListsHonorThePrefixBoundary: shape = self._aliased_server(short_prefix="F3X") manager = MCPServerManager() - manager._create_prefixed_tools([MCPTool(name="deletepet", description="", input_schema={})], shape) + manager._create_prefixed_tools([MCPTool(name="deletepet", description="", inputSchema={})], shape) registered = sorted(manager.tool_name_to_mcp_server_name_mapping) assert len(registered) > 1 @@ -11393,7 +11347,7 @@ class TestToolAuthorizationIsNotConditionalOnLogging: @pytest.mark.asyncio async def test_unentitled_tool_refused_without_proxy_logging_obj(self): manager, user = self._manager_with_scoped_server() - upstream = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + upstream = AsyncMock(return_value=CallToolResult(content=[], isError=False)) with patch.object(manager, "_call_regular_mcp_tool", new=upstream): with pytest.raises(HTTPException) as exc: @@ -11413,7 +11367,7 @@ class TestToolAuthorizationIsNotConditionalOnLogging: """The gate must refuse only what the entitlement excludes; an allowed tool still reaches the upstream when there is no logging object.""" manager, user = self._manager_with_scoped_server() - upstream = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + upstream = AsyncMock(return_value=CallToolResult(content=[], isError=False)) with patch.object(manager, "_call_regular_mcp_tool", new=upstream): await manager.call_tool( @@ -11626,7 +11580,7 @@ class TestClientForwardedDiscoveryFailureIsNotFatal: server = await self._registered(manager, auth_type, None) manager._set_oauth_discovery_deferred(server.server_id, True) manager._fetch_tools_with_timeout = AsyncMock( - return_value=[MCPTool(name="list_reports", description="d", input_schema={"type": "object"})] + return_value=[MCPTool(name="list_reports", description="d", inputSchema={"type": "object"})] ) with patch.object(manager, "_discover_oauth_metadata_for_server", new=AsyncMock(return_value=None)): @@ -11866,9 +11820,9 @@ class TestConfigServerIdPinning: } @pytest.mark.asyncio - async def test_derived_id_churns_when_connection_fields_change(self): + async def test_derived_id_churns_when_connection_fields_change(self, config_only_mcp_manager_factory): """The behavior the pin exists to escape: editing the url mints a brand-new id.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config()) before = next(iter(manager.config_mcp_servers)) @@ -11880,8 +11834,8 @@ class TestConfigServerIdPinning: assert before != after @pytest.mark.asyncio - async def test_pinned_id_survives_url_transport_auth_and_alias_edits(self): - manager = MCPServerManager() + async def test_pinned_id_survives_url_transport_auth_and_alias_edits(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) assert list(manager.config_mcp_servers) == ["docs-prod-1"] @@ -11902,8 +11856,8 @@ class TestConfigServerIdPinning: assert manager.config_mcp_servers["docs-prod-1"].url == "https://prod.example.com/mcp" @pytest.mark.asyncio - async def test_absent_server_id_keeps_the_derived_hash(self): - manager = MCPServerManager() + async def test_absent_server_id_keeps_the_derived_hash(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config()) @@ -11918,15 +11872,15 @@ class TestConfigServerIdPinning: @pytest.mark.asyncio @pytest.mark.parametrize("bad_value", ["", " ", 123, True, ["docs-prod-1"]]) - async def test_blank_or_non_string_server_id_is_rejected(self, bad_value: Any): - manager = MCPServerManager() + async def test_blank_or_non_string_server_id_is_rejected(self, config_only_mcp_manager_factory, bad_value: Any): + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_id must be a non-empty string"): await manager.load_servers_from_config(self._config(server_id=bad_value)) @pytest.mark.asyncio - async def test_two_servers_pinning_the_same_id_are_rejected(self): - manager = MCPServerManager() + async def test_two_servers_pinning_the_same_id_are_rejected(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() config: Dict[str, Any] = { "docs_server": {"url": "https://a.example.com/mcp", "server_id": "shared-id"}, "wiki_server": {"url": "https://b.example.com/mcp", "server_id": "shared-id"}, @@ -11936,9 +11890,9 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(config) @pytest.mark.asyncio - async def test_pinned_id_colliding_with_a_derived_id_is_rejected(self): + async def test_pinned_id_colliding_with_a_derived_id_is_rejected(self, config_only_mcp_manager_factory): """A pin that lands on another entry's derived hash collides just as hard.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() derived = manager._generate_stable_server_id( server_name="docs_server", url="https://a.example.com/mcp", @@ -11955,14 +11909,14 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(config) @pytest.mark.asyncio - async def test_pinned_id_colliding_with_a_db_backed_server_is_rejected(self): + async def test_pinned_id_colliding_with_a_db_backed_server_is_rejected(self, config_only_mcp_manager_factory): """get_registry() is ``config | registry``, so the db row would hide the config server. The registry is seeded by hand because on a real startup the config loads before the database does, so this check only fires on a later reload. The startup ordering is covered by ``test_db_row_arriving_on_a_pinned_config_id_warns``; the warning there is not redundant. """ - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() manager.registry["db-uuid-1"] = MCPServer( server_id="db-uuid-1", name="db_server", @@ -11974,9 +11928,9 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(self._config(server_id="db-uuid-1")) @pytest.mark.asyncio - async def test_derived_id_matching_a_db_backed_server_is_not_rejected(self): + async def test_derived_id_matching_a_db_backed_server_is_not_rejected(self, config_only_mcp_manager_factory): """Only a pinned id is an authoring error; a hash collision must not fail startup.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() derived = manager._generate_stable_server_id( server_name="docs_server", url="https://example.com/mcp", @@ -11996,8 +11950,8 @@ class TestConfigServerIdPinning: assert derived in manager.config_mcp_servers @pytest.mark.asyncio - async def test_pinned_id_is_stripped_of_surrounding_whitespace(self): - manager = MCPServerManager() + async def test_pinned_id_is_stripped_of_surrounding_whitespace(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id=" docs-prod-1 ")) @@ -12037,9 +11991,9 @@ class TestConfigServerIdPinning: await manager.reload_servers_from_database() @pytest.mark.asyncio - async def test_db_row_arriving_on_a_pinned_config_id_warns(self, caplog): + async def test_db_row_arriving_on_a_pinned_config_id_warns(self, config_only_mcp_manager_factory, caplog): """The db row loads after config on startup, so the config server is hidden then, not at load.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12049,8 +12003,8 @@ class TestConfigServerIdPinning: assert manager.get_registry()["docs-prod-1"].url == "https://db.example.com/mcp" @pytest.mark.asyncio - async def test_db_row_with_a_distinct_id_does_not_warn(self, caplog): - manager = MCPServerManager() + async def test_db_row_with_a_distinct_id_does_not_warn(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12060,9 +12014,9 @@ class TestConfigServerIdPinning: assert set(manager.get_registry()) == {"docs-prod-1", "db-uuid-1"} @pytest.mark.asyncio - async def test_pinned_id_matching_another_entrys_server_name_is_rejected(self): + async def test_pinned_id_matching_another_entrys_server_name_is_rejected(self, config_only_mcp_manager_factory): """expand_permission_list resolves against registry keys first, so this steals the grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12077,8 +12031,8 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinned_id_matching_another_entrys_alias_is_rejected(self): - manager = MCPServerManager() + async def test_pinned_id_matching_another_entrys_alias_is_rejected(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12097,17 +12051,17 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinning_a_servers_own_name_is_allowed(self): + async def test_pinning_a_servers_own_name_is_allowed(self, config_only_mcp_manager_factory): """The most natural pin an operator writes; it resolves to the same server either way.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs_server")) assert list(manager.config_mcp_servers) == ["docs_server"] @pytest.mark.asyncio - async def test_pinning_a_servers_own_alias_is_allowed(self): - manager = MCPServerManager() + async def test_pinning_a_servers_own_alias_is_allowed(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(alias="docs", server_id="docs")) @@ -12115,9 +12069,9 @@ class TestConfigServerIdPinning: @pytest.mark.asyncio @pytest.mark.parametrize("aliasing_entry_first", [True, False]) - async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, aliasing_entry_first: bool): + async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, config_only_mcp_manager_factory, aliasing_entry_first: bool): """A grant naming 'docs_server' reaches both servers unpinned; the pin would narrow it to one.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() wiki = ( "wiki_server", {"alias": "docs_server", "url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, @@ -12131,8 +12085,8 @@ class TestConfigServerIdPinning: await manager.load_servers_from_config(dict((wiki, docs) if aliasing_entry_first else (docs, wiki))) @pytest.mark.asyncio - async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self): - manager = MCPServerManager() + async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12148,9 +12102,9 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinning_own_alias_shared_with_a_later_entry_is_rejected(self): + async def test_pinning_own_alias_shared_with_a_later_entry_is_rejected(self, config_only_mcp_manager_factory): """Nothing rejects duplicate aliases, so the first entry's pin would answer the second's grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'docs_server'"): await manager.load_servers_from_config( @@ -12170,9 +12124,9 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_own_name_pin_resolves_grants_like_the_unpinned_name(self): + async def test_own_name_pin_resolves_grants_like_the_unpinned_name(self, config_only_mcp_manager_factory): """The negative control: a sole-owner self-pin must keep loading and answer the same grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12190,9 +12144,9 @@ class TestConfigServerIdPinning: assert manager.expand_permission_list(["wiki"]) == [wiki_id] @pytest.mark.asyncio - async def test_derived_id_is_not_checked_against_names(self): + async def test_derived_id_is_not_checked_against_names(self, config_only_mcp_manager_factory): """Unpinned configs must keep loading; only a pinned id can be an authoring error.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12204,9 +12158,9 @@ class TestConfigServerIdPinning: assert len(manager.config_mcp_servers) == 2 @pytest.mark.asyncio - async def test_shadow_warning_is_not_repeated_on_every_reload(self, caplog): + async def test_shadow_warning_is_not_repeated_on_every_reload(self, config_only_mcp_manager_factory, caplog): """reload_servers_from_database runs on the config-reload timer; one warning, not one a tick.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12219,8 +12173,8 @@ class TestConfigServerIdPinning: assert second_round == first_round @pytest.mark.asyncio - async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, caplog): - manager = MCPServerManager() + async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12231,9 +12185,9 @@ class TestConfigServerIdPinning: assert len([m for m in caplog.messages if "database entry takes precedence" in m]) == 2 @pytest.mark.asyncio - async def test_pinned_id_matching_a_mapped_alias_is_rejected(self): + async def test_pinned_id_matching_a_mapped_alias_is_rejected(self, config_only_mcp_manager_factory): """An alias can also arrive from litellm_settings.mcp_aliases; it is reserved just the same.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): await manager.load_servers_from_config( @@ -12249,8 +12203,8 @@ class TestConfigServerIdPinning: ) @pytest.mark.asyncio - async def test_pinning_a_servers_own_mapped_alias_is_allowed(self): - manager = MCPServerManager() + async def test_pinning_a_servers_own_mapped_alias_is_allowed(self, config_only_mcp_manager_factory): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( self._config(server_id="docs"), @@ -12260,9 +12214,9 @@ class TestConfigServerIdPinning: assert list(manager.config_mcp_servers) == ["docs"] @pytest.mark.asyncio - async def test_mapped_alias_for_an_unknown_server_reserves_nothing(self): + async def test_mapped_alias_for_an_unknown_server_reserves_nothing(self, config_only_mcp_manager_factory): """A dangling mcp_aliases entry is never applied, so it must not fail an unrelated pin.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( self._config(server_id="wiki"), @@ -12272,9 +12226,9 @@ class TestConfigServerIdPinning: assert list(manager.config_mcp_servers) == ["wiki"] @pytest.mark.asyncio - async def test_config_id_that_is_a_db_server_name_warns(self, caplog): + async def test_config_id_that_is_a_db_server_name_warns(self, config_only_mcp_manager_factory, caplog): """The mirror of the shadow case: here the config entry captures the db server's grants.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="db_server")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12283,8 +12237,8 @@ class TestConfigServerIdPinning: assert any("db_server" in m and "name or alias of a database-backed" in m for m in caplog.messages) @pytest.mark.asyncio - async def test_capture_warning_is_not_repeated_on_every_reload(self, caplog): - manager = MCPServerManager() + async def test_capture_warning_is_not_repeated_on_every_reload(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="db_server")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12294,8 +12248,8 @@ class TestConfigServerIdPinning: assert len([m for m in caplog.messages if "name or alias of a database-backed" in m]) == 1 @pytest.mark.asyncio - async def test_config_id_unrelated_to_db_names_does_not_warn(self, caplog): - manager = MCPServerManager() + async def test_config_id_unrelated_to_db_names_does_not_warn(self, config_only_mcp_manager_factory, caplog): + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12304,9 +12258,9 @@ class TestConfigServerIdPinning: assert all("name or alias of a database-backed" not in m for m in caplog.messages) @pytest.mark.asyncio - async def test_mapped_alias_for_a_server_with_its_own_alias_reserves_nothing(self): + async def test_mapped_alias_for_a_server_with_its_own_alias_reserves_nothing(self, config_only_mcp_manager_factory): """load_servers_from_config ignores the mapping when the entry sets alias, so it is free.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12328,9 +12282,9 @@ class TestConfigServerIdPinning: assert len(manager.config_mcp_servers) == 2 @pytest.mark.asyncio - async def test_only_the_first_mapped_alias_for_a_server_is_reserved(self): + async def test_only_the_first_mapped_alias_for_a_server_is_reserved(self, config_only_mcp_manager_factory): """Only the first mapping is applied, so pinning the second one must still load.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12347,15 +12301,15 @@ class TestConfigServerIdPinning: assert "wiki_two" in manager.config_mcp_servers @pytest.mark.asyncio - async def test_invalid_name_is_reported_before_any_entry_body_is_read(self): + async def test_invalid_name_is_reported_before_any_entry_body_is_read(self, config_only_mcp_manager_factory): """The identifier index walks every entry up front, so a bad name must still fail on the name.""" with pytest.raises(Exception, match="Server name cannot contain"): - await MCPServerManager().load_servers_from_config({"my-server": None}) + await config_only_mcp_manager_factory().load_servers_from_config({"my-server": None}) @pytest.mark.asyncio - async def test_a_shadowing_db_server_reports_only_the_shadow_warning(self, caplog): + async def test_a_shadowing_db_server_reports_only_the_shadow_warning(self, config_only_mcp_manager_factory, caplog): """The db row wins the id outright, so the capture message would contradict the shadow one.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config(self._config(server_id="db_server")) with caplog.at_level(logging.WARNING, logger="LiteLLM"): @@ -12366,9 +12320,9 @@ class TestConfigServerIdPinning: assert manager.get_registry()["db_server"].url == "https://db.example.com/mcp" @pytest.mark.asyncio - async def test_an_explicitly_blank_alias_still_blocks_the_mapping(self): + async def test_an_explicitly_blank_alias_still_blocks_the_mapping(self, config_only_mcp_manager_factory): """The loader only consults mcp_aliases when the key is absent, so a blank alias frees it.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { @@ -12390,9 +12344,9 @@ class TestConfigServerIdPinning: assert manager.config_mcp_servers["wiki"].url == "https://example.com/mcp" @pytest.mark.asyncio - async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, caplog): + async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, config_only_mcp_manager_factory, caplog): """Skipping is per identifier, not per row, so the second collision is not lost.""" - manager = MCPServerManager() + manager = config_only_mcp_manager_factory() await manager.load_servers_from_config( { "docs_server": { @@ -12472,7 +12426,7 @@ class TestLitellmAdmissionKeyIsNeverTheSubjectToken: def _manager_with_recording_client() -> MCPServerManager: manager: Final = MCPServerManager() client: Final = AsyncMock() - client.call_tool = AsyncMock(return_value=CallToolResult(content=[], is_error=False)) + client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) client.list_prompts = AsyncMock(return_value=[]) client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[])) manager._create_mcp_client = AsyncMock(return_value=client) @@ -12762,7 +12716,7 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, ("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), ], ) -async def test_debug_resolution_matches_final_header_conflict_winner( +async def test_debug_resolution_matches_final_header_conflict_winner(_mcp_request_ctx, config: Literal["stored", "static", "none"], extra_headers: dict[str, str] | None, expected_source: str, @@ -12833,7 +12787,7 @@ async def test_debug_resolution_matches_final_header_conflict_winner( @pytest.mark.asyncio @pytest.mark.parametrize("transport", ["http", "stdio"]) -async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None: +async def test_debug_reports_legacy_signing_and_non_http_transport(_mcp_request_ctx, transport: Literal["http", "stdio"]) -> None: from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from starlette.requests import Request @@ -12877,16 +12831,12 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li async def test_temporary_server_discovery_reuses_resolved_metadata_without_publishing() -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="temporary-oauth-discovery", - name="temporary", - url="https://idp.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.true_passthrough, + server_id="temporary-oauth-discovery", name="temporary", url="https://idp.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, ) manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", ) with patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery: @@ -12906,18 +12856,13 @@ async def test_temporary_server_discovery_reuses_resolved_metadata_without_publi async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="repeated-stale", - name="stale", - url="https://idp.example.com/mcp", - transport=MCPTransport.http, - auth_type=auth_type, - oauth2_flow="authorization_code", + server_id="repeated-stale", name="stale", url="https://idp.example.com/mcp", + transport=MCPTransport.http, auth_type=auth_type, oauth2_flow="authorization_code", ) manager.registry[server.server_id] = server manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", ) with ( patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery, @@ -12937,20 +12882,13 @@ async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> async def test_stale_discovery_falls_back_to_resolved_registered_server() -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="resolved-replacement", - name="replacement", - url="https://old.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - oauth2_flow="authorization_code", - ) - replacement: Final = original.model_copy( - update={ - "url": "https://new.example.com/mcp", - "authorization_url": "https://new.example.com/authorize", - "token_url": "https://new.example.com/token", - } + server_id="resolved-replacement", name="replacement", url="https://old.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", ) + replacement: Final = original.model_copy(update={ + "url": "https://new.example.com/mcp", "authorization_url": "https://new.example.com/authorize", + "token_url": "https://new.example.com/token", + }) manager.registry[original.server_id] = replacement assert await manager._rejoin_oauth_metadata_discovery(original, retry_stale=False) is replacement @@ -12958,11 +12896,8 @@ async def test_stale_discovery_falls_back_to_resolved_registered_server() -> Non def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="stale-publication", - name="publication", - url="https://old.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, + server_id="stale-publication", name="publication", url="https://old.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2, ) manager._set_oauth_discovery_deferred(original.server_id, True) original_slot: Final = manager._oauth_discovery_slot(original.server_id) @@ -12978,13 +12913,9 @@ def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: async def test_temporary_oauth_discovery_expires_without_more_requests() -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="expiring-session", - name="temporary", - url="https://idp.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.true_passthrough, - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", + server_id="expiring-session", name="temporary", url="https://idp.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", ) manager._set_oauth_discovery_deferred(server.server_id, True) resolved: Final = await manager.ensure_oauth_metadata_discovered(server) @@ -13085,9 +13016,7 @@ async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(r result = await manager.health_check_server(server.server_id) cached = await manager.health_check_server(server.server_id) assert result.status == "unknown" - assert ( - result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" - ) + assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" assert cached.health_check_error == result.health_check_error assert cached.last_health_check == result.last_health_check assert route.call_count == 1 @@ -13099,11 +13028,8 @@ async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, mon monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( - server_id="cancelled-cache", - name="cancelled-cache", - transport=MCPTransport.http, - spec_path="https://93.184.216.34/cancelled-cache.json", - auth_type=MCPAuth.none, + server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http, + spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none, ) manager.registry = {server.server_id: server} started = asyncio.Event() @@ -13225,9 +13151,7 @@ class _DiscoveryUpstream: def _discovery_server() -> MCPServer: - return MCPServer( - server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http - ) + return MCPServer(server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http) @pytest.mark.asyncio @@ -13372,9 +13296,7 @@ async def test_discovery_cache_can_be_disabled(monkeypatch: pytest.MonkeyPatch) assert upstream.initializes == 2 -@pytest.mark.parametrize( - "value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5)) -) +@pytest.mark.parametrize("value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5))) def test_discovery_cache_ttl_validation(value: str, expected: float, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _mcp_discovery_cache_ttl @@ -13637,45 +13559,26 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( class TestProtectedCredentialPreparation: @pytest.mark.asyncio - @pytest.mark.parametrize( - "auth_type,credential", - [ - (MCPAuth.bearer_token, None), - (MCPAuth.bearer_token, "Bearer"), - (MCPAuth.api_key, None), - (MCPAuth.basic, "Basic"), - ], - ) + @pytest.mark.parametrize("auth_type,credential", [ + (MCPAuth.bearer_token, None), + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.api_key, None), + (MCPAuth.basic, "Basic"), + ]) @pytest.mark.parametrize("dispatch", ["managed", "local"]) async def test_openapi_dispatch_rejects_unusable_effective_credentials( - self, - tmp_path: Path, - respx_mock: MockRouter, - monkeypatch: pytest.MonkeyPatch, - auth_type: MCPAuthType, - credential: str | None, - dispatch: str, + self, tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, credential: str | None, dispatch: str, ) -> None: from litellm.proxy._experimental.mcp_server.server import _handle_local_mcp_tool from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix spec_path: Final = tmp_path / "openapi.json" - spec_path.write_text( - json.dumps( - { - "openapi": "3.0.0", - "info": {"title": "Auth", "version": "1"}, - "paths": {"/echo": {"get": {"operationId": "echo"}}}, - } - ) - ) + spec_path.write_text(json.dumps({"openapi": "3.0.0", "info": {"title": "Auth", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}})) server: Final = MCPServer( - server_id="dispatch-auth", - name="dispatch-auth", - url="https://upstream.example", - transport=MCPTransport.http, - auth_type=auth_type, - authentication_token=credential, + server_id="dispatch-auth", name="dispatch-auth", url="https://upstream.example", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=credential, ) manager: Final = MCPServerManager() await manager._register_openapi_tools(str(spec_path), server, server.url) @@ -13698,21 +13601,14 @@ class TestProtectedCredentialPreparation: self, transport: MCPTransport, client_secret: str | None, subject: str | None ) -> None: server = MCPServer( - server_id="incomplete-obo", - name="incomplete-obo", - url="https://upstream.example/mcp", - transport=transport, - auth_type=MCPAuth.oauth2_token_exchange, - client_id="gateway", - client_secret=client_secret, - token_exchange_endpoint="https://idp.example/token", - authentication_token="static-fallback", + server_id="incomplete-obo", name="incomplete-obo", url="https://upstream.example/mcp", + transport=transport, auth_type=MCPAuth.oauth2_token_exchange, + client_id="gateway", client_secret=client_secret, + token_exchange_endpoint="https://idp.example/token", authentication_token="static-fallback", ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client( - server, - mcp_auth_header="Bearer override", - subject_token=subject, + server, mcp_auth_header="Bearer override", subject_token=subject, ) assert exc.value.status_code == (401 if subject is None else 500) assert "static-fallback" not in str(exc.value.detail) @@ -13725,11 +13621,8 @@ class TestProtectedCredentialPreparation: self, auth_type: MCPAuthType, credential: str | dict[str, str] | None ) -> None: server = MCPServer( - server_id="empty-static", - name="empty-static", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, + server_id="empty-static", name="empty-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header=credential) @@ -13737,22 +13630,16 @@ class TestProtectedCredentialPreparation: assert "credential" in str(exc.value.detail).lower() @pytest.mark.asyncio - @pytest.mark.parametrize( - "auth_type,headers", - [ - (MCPAuth.api_key, {"X-API-Key": "key"}), - (MCPAuth.bearer_token, {"Authorization": "Bearer token"}), - ], - ) + @pytest.mark.parametrize("auth_type,headers", [ + (MCPAuth.api_key, {"X-API-Key": "key"}), + (MCPAuth.bearer_token, {"Authorization": "Bearer token"}), + ]) async def test_static_auth_accepts_actual_forwarded_credential( self, auth_type: MCPAuthType, headers: dict[str, str] ) -> None: server = MCPServer( - server_id="header-static", - name="header-static", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, + server_id="header-static", name="header-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, ) client = await MCPServerManager()._create_mcp_client(server, extra_headers=headers) assert client._get_auth_headers() == headers @@ -13761,48 +13648,29 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange]) async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: server = MCPServer( - server_id="openapi-empty", - name="openapi-empty", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, + server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, token_exchange_endpoint="https://idp.example/token", ) with pytest.raises(HTTPException) as exc: await MCPServerManager().resolve_openapi_upstream_auth( - mcp_server=server, - oauth2_headers=None, - raw_headers=None, - mcp_auth_header=None, - user_api_key_auth=None, - forwarded_headers=None, + mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, + user_api_key_auth=None, forwarded_headers=None, ) assert exc.value.status_code in (401, 500) @pytest.mark.asyncio - @pytest.mark.parametrize( - "auth_type,slot,value", - [ - (MCPAuth.api_key, "X-API-Key", "token"), - (MCPAuth.authorization, "Authorization", "opaque-secret-value"), - (MCPAuth.authorization, "Authorization", "Bearer abc"), - (MCPAuth.authorization, "Authorization", "Custom abc"), - ], - ) + @pytest.mark.parametrize("auth_type,slot,value", [ + (MCPAuth.api_key, "X-API-Key", "token"), + (MCPAuth.authorization, "Authorization", "opaque-secret-value"), + (MCPAuth.authorization, "Authorization", "Bearer abc"), + (MCPAuth.authorization, "Authorization", "Custom abc"), + ]) async def test_raw_static_credentials_are_forwarded_unchanged( - self, - auth_type: MCPAuthType, - slot: str, - value: str, + self, auth_type: MCPAuthType, slot: str, value: str, ) -> None: - server = MCPServer( - server_id="raw-key", - name="raw-key", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, - authentication_token=value, - ) + server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value) client = await MCPServerManager()._create_mcp_client(server) assert client._resolved_auth is not None request = httpx.Request("GET", server.url) @@ -13816,24 +13684,17 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("value", ["Bearer", "basic", "token", "ApiKey", " bEaReR ", "\tTOKEN\t"]) @pytest.mark.parametrize("source", ["configured", "caller", "forwarded"]) async def test_raw_authorization_rejects_bare_schemes_before_dispatch( - self, - respx_mock: MockRouter, - value: str, - source: str, + self, respx_mock: MockRouter, value: str, source: str, ) -> None: server: Final = MCPServer( - server_id="raw-empty", - name="raw-empty", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.authorization, + server_id="raw-empty", name="raw-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.authorization, authentication_token=value if source == "configured" else None, ) destination: Final = respx_mock.route().respond(200) with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: await MCPServerManager()._create_mcp_client( - server, - mcp_auth_header=value if source == "caller" else None, + server, mcp_auth_header=value if source == "caller" else None, extra_headers={"Authorization": value} if source == "forwarded" else None, ) assert exc.value.status_code == 500 @@ -13841,15 +13702,9 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None: - server = MCPServer( - server_id="obo-byok", - name="obo-byok", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - is_byok=True, - token_exchange_endpoint="https://idp.example/token", - ) + server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, is_byok=True, + token_exchange_endpoint="https://idp.example/token") with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header="Bearer override") assert exc.value.status_code == 401 @@ -13857,66 +13712,41 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio @pytest.mark.parametrize("configured,override", [(None, "Bearer usable"), ("shared", "Bearer usable")]) async def test_bearer_override_remains_usable(self, configured: str | None, override: str) -> None: - server = MCPServer( - server_id="override", - name="override", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.bearer_token, - authentication_token=configured, - ) + server = MCPServer(server_id="override", name="override", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=configured) client = await MCPServerManager()._create_mcp_client(server, mcp_auth_header=override) assert client._get_auth_headers()["Authorization"] == override @pytest.mark.asyncio @pytest.mark.parametrize("token", [None, "shared"]) async def test_empty_injected_header_cannot_satisfy_protected_auth(self, token: str | None) -> None: - server = MCPServer( - server_id="empty-header", - name="empty-header", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.bearer_token, - authentication_token=token, - ) + server = MCPServer(server_id="empty-header", name="empty-header", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=token) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"authorization": " "}) assert exc.value.status_code == 500 @pytest.mark.asyncio async def test_custom_slot_uses_its_actual_credential(self) -> None: - server = MCPServer( - server_id="custom", - name="custom", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - upstream_token_header="X-Custom", - authentication_token="key", - ) + server = MCPServer(server_id="custom", name="custom", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, + upstream_token_header="X-Custom", authentication_token="key") client = await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Trace": "trace"}) assert client._credential_slot == "X-Custom" assert await client.discovery_auth_fingerprint() @pytest.mark.asyncio - @pytest.mark.parametrize( - "static_headers,accepted", - [ - ({"apikey": "static-key"}, True), - ({"apikey": ""}, False), - ({"X-Tenant": "tenant"}, True), - ], - ) + @pytest.mark.parametrize("static_headers,accepted", [ + ({"apikey": "static-key"}, True), + ({"apikey": ""}, False), + ({"X-Tenant": "tenant"}, True), + ]) async def test_api_key_carried_by_static_header_passes_fail_closed_check( self, static_headers: dict[str, str], accepted: bool ) -> None: server: Final = MCPServer( - server_id="static-slot", - name="static-slot", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - static_headers=static_headers, + server_id="static-slot", name="static-slot", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static_headers, ) if not accepted: with pytest.raises(HTTPException) as exc: @@ -13928,36 +13758,21 @@ class TestProtectedCredentialPreparation: assert all(request.headers[name] == value for name, value in static_headers.items()) @pytest.mark.asyncio - @pytest.mark.parametrize( - "static,forwarded,caller", - [ - ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), - ({}, {"X-API-Key": "forwarded"}, None), - ({}, None, "ApiKey caller"), - ({"X-API-Key": "static"}, {"Authorization": ""}, None), - ], - ) + @pytest.mark.parametrize("static,forwarded,caller", [ + ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), + ({}, {"X-API-Key": "forwarded"}, None), + ({}, None, "ApiKey caller"), + ({"X-API-Key": "static"}, {"Authorization": ""}, None), + ]) async def test_openapi_static_credentials_remain_supported( - self, - respx_mock: MockRouter, - monkeypatch: pytest.MonkeyPatch, - static: dict[str, str], - forwarded: dict[str, str] | None, - caller: str | None, + self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None ) -> None: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - _request_auth_header, - _request_extra_headers, - create_tool_function, + _request_auth_header, _request_extra_headers, create_tool_function, ) - tool: Final = create_tool_function( - "/echo", - "get", - {}, - "https://upstream.example", - headers=static, - auth_type=MCPAuth.api_key, + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.api_key, ) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") @@ -13991,13 +13806,8 @@ class TestProtectedCredentialPreparation: self.closed = True auth = CancelledAuth() - server = MCPServer( - server_id="cancel", - name="cancel", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - ) + server = MCPServer(server_id="cancel", name="cancel", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key) client = MCPClient(server_url=server.url, auth_type=MCPAuth.api_key, resolved_auth=auth) with pytest.raises(asyncio.CancelledError): await prepare_mcp_client(server, client) @@ -14006,14 +13816,8 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio @pytest.mark.parametrize("auth_type", [MCPAuth.basic, MCPAuth.token, MCPAuth.authorization]) async def test_other_static_schemes_reject_whitespace_credentials(self, auth_type: MCPAuthType) -> None: - server = MCPServer( - server_id="blank-static", - name="blank-static", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, - authentication_token=" ", - ) + server = MCPServer(server_id="blank-static", name="blank-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=" ") with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server) assert exc.value.status_code == 500 @@ -14021,13 +13825,8 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM=", "Basic bm8tY29sb24="]) async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: - server = MCPServer( - server_id="bad-basic", - name="bad-basic", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.basic, - ) + server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header}) assert exc.value.status_code == 500 @@ -14036,48 +13835,34 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("value", ["Basic", "Basic ", "basic"]) @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_basic_scheme_alone_is_not_a_credential(self, value: str, source: str) -> None: - server = MCPServer( - server_id="basic-scheme", - name="basic-scheme", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.basic, - authentication_token=value if source == "configured" else None, - ) + server = MCPServer(server_id="basic-scheme", name="basic-scheme", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, + authentication_token=value if source == "configured" else None) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize( - "auth_type,value,default_slot", - [ - (MCPAuth.api_key, "fixture-key", "X-API-Key"), - (MCPAuth.bearer_token, "fixture-key", "Authorization"), - (MCPAuth.basic, "user:pass", "Authorization"), - (MCPAuth.token, "fixture-key", "Authorization"), - (MCPAuth.authorization, "fixture-key", "Authorization"), - ], - ) + @pytest.mark.parametrize("auth_type,value,default_slot", [ + (MCPAuth.api_key, "fixture-key", "X-API-Key"), + (MCPAuth.bearer_token, "fixture-key", "Authorization"), + (MCPAuth.basic, "user:pass", "Authorization"), + (MCPAuth.token, "fixture-key", "Authorization"), + (MCPAuth.authorization, "fixture-key", "Authorization"), + ]) @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_usable_credential_survives_an_empty_alternate_header( self, auth_type: MCPAuthType, value: str, default_slot: str, source: str ) -> None: server: Final = MCPServer( - server_id="alternate", - name="alternate", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, - upstream_token_header="X-Custom", + server_id="alternate", name="alternate", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, upstream_token_header="X-Custom", authentication_token=value if source == "configured" else None, ) empty_slot: Final = default_slot if source == "configured" else "X-Custom" selected_slot: Final = "X-Custom" if source == "configured" else default_slot client: Final = await MCPServerManager()._create_mcp_client( - server, - mcp_auth_header=value if source == "caller" else None, - extra_headers={empty_slot: ""}, + server, mcp_auth_header=value if source == "caller" else None, extra_headers={empty_slot: ""}, ) request: Final = await client.prepare_request_auth() assert request.headers[selected_slot] @@ -14086,12 +13871,8 @@ class TestProtectedCredentialPreparation: @pytest.mark.asyncio async def test_empty_custom_and_default_headers_do_not_satisfy_auth(self) -> None: server: Final = MCPServer( - server_id="both-empty", - name="both-empty", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - upstream_token_header="X-Custom", + server_id="both-empty", name="both-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""}) @@ -14104,17 +13885,12 @@ class TestProtectedCredentialPreparation: self, custom_slot: str | None, source: str ) -> None: server: Final = MCPServer( - server_id="caller-auth", - name="caller-auth", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - upstream_token_header=custom_slot, + server_id="caller-auth", name="caller-auth", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header=custom_slot, ) headers: Final = {"Authorization": "Bearer caller-credential", "X-API-Key": ""} client: Final = await MCPServerManager()._create_mcp_client( - server, - mcp_auth_header=headers if source == "caller" else None, + server, mcp_auth_header=headers if source == "caller" else None, extra_headers=headers if source == "forwarded" else None, ) request: Final = await client.prepare_request_auth() @@ -14123,29 +13899,14 @@ class TestProtectedCredentialPreparation: assert custom_slot is None or custom_slot not in request.headers @pytest.mark.asyncio - @pytest.mark.parametrize( - "value", - [ - "", - " ", - "Bearer", - "Basic", - "token", - "ApiKey", - "Bearer Bearer", - "ApiKey ApiKey", - "token token", - "bEaReR BEARER", - "aPiKeY\tAPIKEY", - ], - ) + @pytest.mark.parametrize("value", [ + "", " ", "Bearer", "Basic", "token", "ApiKey", + "Bearer Bearer", "ApiKey ApiKey", "token token", "bEaReR BEARER", "aPiKeY\tAPIKEY", + ]) async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None: server: Final = MCPServer( - server_id="caller-empty", - name="caller-empty", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, + server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, ) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value}) @@ -14156,11 +13917,8 @@ class TestProtectedCredentialPreparation: @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_basic_requires_a_username_password_separator(self, value: str, source: str) -> None: server: Final = MCPServer( - server_id="basic-pair", - name="basic-pair", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.basic, + server_id="basic-pair", name="basic-pair", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token=value if source == "configured" else None, ) with pytest.raises(HTTPException) as exc: @@ -14173,12 +13931,8 @@ class TestProtectedCredentialPreparation: import base64 server: Final = MCPServer( - server_id="basic-valid", - name="basic-valid", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.basic, - authentication_token=value, + server_id="basic-valid", name="basic-valid", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token=value, ) client: Final = await MCPServerManager()._create_mcp_client(server) request: Final = await client.prepare_request_auth() @@ -14187,27 +13941,17 @@ class TestProtectedCredentialPreparation: assert base64.b64decode(encoded) == value.encode() @pytest.mark.asyncio - @pytest.mark.parametrize( - "auth_type,value", - [ - (MCPAuth.bearer_token, "Bearer"), - (MCPAuth.bearer_token, "Bearer "), - (MCPAuth.bearer_token, "bearer"), - (MCPAuth.token, "token"), - (MCPAuth.token, "token "), - (MCPAuth.token, "TOKEN"), - ], - ) + @pytest.mark.parametrize("auth_type,value", [ + (MCPAuth.bearer_token, "Bearer"), (MCPAuth.bearer_token, "Bearer "), (MCPAuth.bearer_token, "bearer"), + (MCPAuth.token, "token"), (MCPAuth.token, "token "), (MCPAuth.token, "TOKEN"), + ]) @pytest.mark.parametrize("source", ["configured", "caller"]) async def test_static_scheme_only_input_cannot_hide_behind_rendered_prefix( self, auth_type: MCPAuthType, value: str, source: str ) -> None: server: Final = MCPServer( - server_id="empty-scheme", - name="empty-scheme", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, + server_id="empty-scheme", name="empty-scheme", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value if source == "configured" else None, ) with pytest.raises(HTTPException) as exc: @@ -14215,24 +13959,17 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize( - "auth_type,value,expected", - [ - (MCPAuth.bearer_token, "token", "Bearer token"), - (MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"), - (MCPAuth.token, "tokenish", "token tokenish"), - ], - ) + @pytest.mark.parametrize("auth_type,value,expected", [ + (MCPAuth.bearer_token, "token", "Bearer token"), + (MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"), + (MCPAuth.token, "tokenish", "token tokenish"), + ]) async def test_static_credentials_that_resemble_schemes_remain_usable( self, auth_type: MCPAuthType, value: str, expected: str ) -> None: server: Final = MCPServer( - server_id="real-token", - name="real-token", - url="https://upstream.example/mcp", - transport=MCPTransport.http, - auth_type=auth_type, - authentication_token=value, + server_id="real-token", name="real-token", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value, ) client: Final = await MCPServerManager()._create_mcp_client(server) request: Final = await client.prepare_request_auth() @@ -14271,31 +14008,16 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon registry.register_tool("observer-execute", "Execute", {"type": "object"}, upstream) monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) manager = MCPServerManager() - manager.registry = { - "observer": MCPServer( - server_id="observer", - name="observer", - server_name="observer", - transport="http", - url="https://observer.example/mcp", - spec_path="observer.json", - auth_type="none", - ) - } + manager.registry = {"observer": MCPServer( + server_id="observer", name="observer", server_name="observer", transport="http", + url="https://observer.example/mcp", spec_path="observer.json", auth_type="none", + )} manager.tool_name_to_mcp_server_name_mapping = {"observer-execute": "observer"} - result = await asyncio.wait_for( - manager.call_tool( - server_name="observer", - name="execute", - arguments={"text": "hello"}, - user_api_key_auth=UserAPIKeyAuth(), - proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), - guardrail_context=MCPRequestContext.resolve_guardrail_context( - {"metadata": {"guardrails": ["observe"] if selected else []}} - ), - ), - timeout=5, - ) + result = await asyncio.wait_for(manager.call_tool( + server_name="observer", name="execute", arguments={"text": "hello"}, + user_api_key_auth=UserAPIKeyAuth(), proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + guardrail_context=MCPRequestContext.resolve_guardrail_context({"metadata": {"guardrails": ["observe"] if selected else []}}), + ), timeout=5) assert tool_started.is_set() assert guardrail_started.is_set() is selected assert result.is_error is False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index 66d5f0e56f9..8cf3bc6fcc7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -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())) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index efb841a4e01..cb43d2c2592 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index c4e1f1e4a6e..519acc241c6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -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"}, ) ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index a0320661fa2..07468a682ac 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -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() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 64ec6d2e78e..f0b4e94f72f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -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) ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index 8528f20fe89..941e5deee93 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -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"}), ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py index 0252fb9843d..842859e5a1e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py @@ -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"} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py index 07436199a8d..bc784923eb5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py @@ -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(),