Marcelo Trylesinski
260 lines
10 KiB
Python
260 lines
10 KiB
Python
import inspect
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import jsonschema
|
|
import pytest
|
|
|
|
import mcp.types as types
|
|
from mcp import Client
|
|
from mcp.server.lowlevel import Server
|
|
from mcp.server.session import ServerSession
|
|
from mcp.shared.context import RequestContext
|
|
from mcp.types import Tool
|
|
|
|
|
|
@contextmanager
|
|
def bypass_server_output_validation():
|
|
"""Context manager that bypasses server-side output validation.
|
|
This simulates a malicious or non-compliant server that doesn't validate
|
|
its outputs, allowing us to test client-side validation.
|
|
"""
|
|
# Save the original validate function
|
|
original_validate = jsonschema.validate
|
|
|
|
# Create a mock that tracks which module is calling it
|
|
def selective_mock(instance: Any = None, schema: Any = None, *args: Any, **kwargs: Any) -> None:
|
|
# Check the call stack to see where this is being called from
|
|
for frame_info in inspect.stack():
|
|
# If called from the server module, skip validation
|
|
# TODO: fix this as it's a rather gross workaround and will eventually break
|
|
# Normalize path separators for cross-platform compatibility
|
|
normalized_path = frame_info.filename.replace("\\", "/")
|
|
if "mcp/server/lowlevel/server.py" in normalized_path:
|
|
return None
|
|
# Otherwise, use the real validation (for client-side)
|
|
return original_validate(instance=instance, schema=schema, *args, **kwargs)
|
|
|
|
with patch("jsonschema.validate", selective_mock):
|
|
yield
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tool_structured_output_client_side_validation_basemodel():
|
|
"""Test that client validates structured content against schema for BaseModel outputs"""
|
|
# Define the expected schema for our tool
|
|
output_schema = {
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string", "title": "Name"}, "age": {"type": "integer", "title": "Age"}},
|
|
"required": ["name", "age"],
|
|
"title": "UserOutput",
|
|
}
|
|
|
|
async def on_list_tools(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.PaginatedRequestParams | None,
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(
|
|
tools=[
|
|
Tool(
|
|
name="get_user",
|
|
description="Get user data",
|
|
input_schema={"type": "object"},
|
|
output_schema=output_schema,
|
|
)
|
|
]
|
|
)
|
|
|
|
async def on_call_tool(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CallToolRequestParams,
|
|
) -> types.CallToolResult:
|
|
# Return invalid structured content - age is string instead of integer
|
|
return types.CallToolResult(
|
|
content=[],
|
|
structured_content={"name": "John", "age": "invalid"}, # Invalid: age should be int
|
|
)
|
|
|
|
# Create a malicious low-level server that returns invalid structured content
|
|
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
|
|
|
|
# Test that client validates the structured content
|
|
with bypass_server_output_validation():
|
|
async with Client(server) as client:
|
|
# The client validates structured content and should raise an error
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await client.call_tool("get_user", {})
|
|
# Verify it's a validation error
|
|
assert "Invalid structured content returned by tool get_user" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tool_structured_output_client_side_validation_primitive():
|
|
"""Test that client validates structured content for primitive outputs"""
|
|
# Primitive types are wrapped in {"result": value}
|
|
output_schema = {
|
|
"type": "object",
|
|
"properties": {"result": {"type": "integer", "title": "Result"}},
|
|
"required": ["result"],
|
|
"title": "calculate_Output",
|
|
}
|
|
|
|
async def on_list_tools(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.PaginatedRequestParams | None,
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(
|
|
tools=[
|
|
Tool(
|
|
name="calculate",
|
|
description="Calculate something",
|
|
input_schema={"type": "object"},
|
|
output_schema=output_schema,
|
|
)
|
|
]
|
|
)
|
|
|
|
async def on_call_tool(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CallToolRequestParams,
|
|
) -> types.CallToolResult:
|
|
# Return invalid structured content - result is string instead of integer
|
|
return types.CallToolResult(
|
|
content=[],
|
|
structured_content={"result": "not_a_number"}, # Invalid: should be int
|
|
)
|
|
|
|
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
|
|
|
|
with bypass_server_output_validation():
|
|
async with Client(server) as client:
|
|
# The client validates structured content and should raise an error
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await client.call_tool("calculate", {})
|
|
assert "Invalid structured content returned by tool calculate" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tool_structured_output_client_side_validation_dict_typed():
|
|
"""Test that client validates dict[str, T] structured content"""
|
|
# dict[str, int] schema
|
|
output_schema = {"type": "object", "additionalProperties": {"type": "integer"}, "title": "get_scores_Output"}
|
|
|
|
async def on_list_tools(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.PaginatedRequestParams | None,
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(
|
|
tools=[
|
|
Tool(
|
|
name="get_scores",
|
|
description="Get scores",
|
|
input_schema={"type": "object"},
|
|
output_schema=output_schema,
|
|
)
|
|
]
|
|
)
|
|
|
|
async def on_call_tool(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CallToolRequestParams,
|
|
) -> types.CallToolResult:
|
|
# Return invalid structured content - values should be integers
|
|
return types.CallToolResult(
|
|
content=[],
|
|
structured_content={"alice": "100", "bob": "85"}, # Invalid: values should be int
|
|
)
|
|
|
|
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
|
|
|
|
with bypass_server_output_validation():
|
|
async with Client(server) as client:
|
|
# The client validates structured content and should raise an error
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await client.call_tool("get_scores", {})
|
|
assert "Invalid structured content returned by tool get_scores" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tool_structured_output_client_side_validation_missing_required():
|
|
"""Test that client validates missing required fields"""
|
|
output_schema = {
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}, "email": {"type": "string"}},
|
|
"required": ["name", "age", "email"], # All fields required
|
|
"title": "PersonOutput",
|
|
}
|
|
|
|
async def on_list_tools(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.PaginatedRequestParams | None,
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(
|
|
tools=[
|
|
Tool(
|
|
name="get_person",
|
|
description="Get person data",
|
|
input_schema={"type": "object"},
|
|
output_schema=output_schema,
|
|
)
|
|
]
|
|
)
|
|
|
|
async def on_call_tool(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CallToolRequestParams,
|
|
) -> types.CallToolResult:
|
|
# Return structured content missing required field 'email'
|
|
return types.CallToolResult(
|
|
content=[],
|
|
structured_content={"name": "John", "age": 30}, # Missing required 'email'
|
|
)
|
|
|
|
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
|
|
|
|
with bypass_server_output_validation():
|
|
async with Client(server) as client:
|
|
# The client validates structured content and should raise an error
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await client.call_tool("get_person", {})
|
|
assert "Invalid structured content returned by tool get_person" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_tool_not_listed_warning(caplog: pytest.LogCaptureFixture):
|
|
"""Test that client logs warning when tool is not in list_tools but has output_schema"""
|
|
|
|
async def on_list_tools(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.PaginatedRequestParams | None,
|
|
) -> types.ListToolsResult:
|
|
# Return empty list - tool is not listed
|
|
return types.ListToolsResult(tools=[])
|
|
|
|
async def on_call_tool(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CallToolRequestParams,
|
|
) -> types.CallToolResult:
|
|
# Server still responds to the tool call with structured content
|
|
return types.CallToolResult(
|
|
content=[],
|
|
structured_content={"result": 42},
|
|
)
|
|
|
|
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
|
|
|
|
# Set logging level to capture warnings
|
|
caplog.set_level(logging.WARNING)
|
|
|
|
with bypass_server_output_validation():
|
|
async with Client(server) as client:
|
|
# Call a tool that wasn't listed
|
|
result = await client.call_tool("mystery_tool", {})
|
|
assert result.structured_content == {"result": 42}
|
|
assert result.is_error is False
|
|
|
|
# Check that warning was logged
|
|
assert "Tool mystery_tool not listed" in caplog.text
|