mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(mcp): preserve legacy behavior on SDK2 and streamline verification
This commit is contained in:
parent
77cf6c2fbd
commit
aea13ee03b
41 changed files with 1109 additions and 1199 deletions
|
|
@ -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
|
||||
|
|
|
|||
7
.github/workflows/test-mcp.yml
vendored
7
.github/workflows/test-mcp.yml
vendored
|
|
@ -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: |
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 []:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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"}},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
),
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -533,7 +533,7 @@ async def test_process_output_response_masks_text_content():
|
|||
TextContent(type="text", text="email jane@example.com"),
|
||||
TextContent(type="text", text="call 415-555-0132"),
|
||||
],
|
||||
is_error=False,
|
||||
isError=False,
|
||||
)
|
||||
|
||||
returned = await handler.process_output_response(
|
||||
|
|
@ -569,7 +569,7 @@ async def test_process_output_response_propagates_block():
|
|||
guardrail = MaskingGuardrail(
|
||||
raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="masking-mcp-guardrail")
|
||||
)
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], is_error=False)
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
with pytest.raises(BlockedPiiEntityError):
|
||||
await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
|
||||
|
|
@ -582,7 +582,7 @@ async def test_process_output_response_skips_non_text_content():
|
|||
guardrail = MaskingGuardrail(masked_texts=["should not be used"])
|
||||
result = CallToolResult(
|
||||
content=[ImageContent(type="image", data="aGk=", mimeType="image/png")],
|
||||
is_error=False,
|
||||
isError=False,
|
||||
)
|
||||
|
||||
returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
|
||||
|
|
@ -613,7 +613,7 @@ async def test_process_output_response_blocks_on_text_count_mismatch():
|
|||
TextContent(type="text", text="jane@example.com"),
|
||||
TextContent(type="text", text="415-555-0132"),
|
||||
],
|
||||
is_error=False,
|
||||
isError=False,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -645,14 +645,14 @@ async def test_structured_content_is_masked_alongside_content():
|
|||
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="email jane@example.com")],
|
||||
structured_content={"contact": {"email": "jane@example.com"}, "balance": 42.0},
|
||||
is_error=False,
|
||||
structuredContent={"contact": {"email": "jane@example.com"}, "balance": 42.0},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
|
||||
|
||||
assert returned.content[0].text == "email <EMAIL_ADDRESS>"
|
||||
assert returned.structured_content== {"contact": {"email": "<EMAIL_ADDRESS>"}, "balance": 42.0}
|
||||
assert returned.structured_content == {"contact": {"email": "<EMAIL_ADDRESS>"}, "balance": 42.0}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -666,14 +666,14 @@ async def test_value_present_only_in_structured_content_is_masked():
|
|||
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="lookup complete")],
|
||||
structured_content={"records": [{"email": "jane@example.com"}]},
|
||||
is_error=False,
|
||||
structuredContent={"records": [{"email": "jane@example.com"}]},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
|
||||
|
||||
assert "jane@example.com" in guardrail.seen_texts
|
||||
assert returned.structured_content== {"records": [{"email": "<EMAIL_ADDRESS>"}]}
|
||||
assert returned.structured_content == {"records": [{"email": "<EMAIL_ADDRESS>"}]}
|
||||
assert returned.content[0].text == "lookup complete"
|
||||
|
||||
|
||||
|
|
@ -684,13 +684,13 @@ async def test_structured_content_without_a_match_is_untouched():
|
|||
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="lookup complete")],
|
||||
structured_content={"record_id": "C-1001", "balance": 42.0, "active": True, "note": None},
|
||||
is_error=False,
|
||||
structuredContent={"record_id": "C-1001", "balance": 42.0, "active": True, "note": None},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
|
||||
|
||||
assert returned.structured_content== {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None}
|
||||
assert returned.structured_content == {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -707,8 +707,8 @@ async def test_structured_content_nested_too_deeply_is_blocked():
|
|||
nested = {"next": nested}
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="lookup complete")],
|
||||
structured_content=nested,
|
||||
is_error=False,
|
||||
structuredContent=nested,
|
||||
isError=False,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -754,8 +754,8 @@ async def test_sensitive_structured_content_key_is_blocked():
|
|||
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="lookup complete")],
|
||||
structured_content={"jane@example.com": {"balance": 42.0}},
|
||||
is_error=False,
|
||||
structuredContent={"jane@example.com": {"balance": 42.0}},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -774,8 +774,8 @@ async def test_sensitive_structured_content_numeric_value_is_blocked():
|
|||
guardrail = SubstitutingGuardrail("4155550199", "<PHONE_NUMBER>")
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="lookup complete")],
|
||||
structured_content={"phone": 4155550199},
|
||||
is_error=False,
|
||||
structuredContent={"phone": 4155550199},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -791,11 +791,11 @@ async def test_clean_structured_content_keys_do_not_block():
|
|||
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
|
||||
response = CallToolResult(
|
||||
content=[TextContent(type="text", text="email jane@example.com")],
|
||||
structured_content={"record_id": "C-1001", "balance": 42.0, "count": 3},
|
||||
is_error=False,
|
||||
structuredContent={"record_id": "C-1001", "balance": 42.0, "count": 3},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
|
||||
|
||||
assert returned.content[0].text == "email <EMAIL_ADDRESS>"
|
||||
assert returned.structured_content== {"record_id": "C-1001", "balance": 42.0, "count": 3}
|
||||
assert returned.structured_content == {"record_id": "C-1001", "balance": 42.0, "count": 3}
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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()))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
)
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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"}),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue