Add schema validation to lowlevel server (#1005)

This commit is contained in:
bhosmer-ant 2025-06-25 09:23:49 -04:00 • committed by GitHub
commit c8bbfc034d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 1155 additions and 8 deletions

View file

@ -0,0 +1,310 @@
"""Test input schema validation for lowlevel server."""
import logging
from collections.abc import Awaitable, Callable
from typing import Any
import anyio
import pytest
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import CallToolResult, ClientResult, ServerNotification, ServerRequest, TextContent, Tool
async def run_tool_test(
tools: list[Tool],
call_tool_handler: Callable[[str, dict[str, Any]], Awaitable[list[TextContent]]],
test_callback: Callable[[ClientSession], Awaitable[CallToolResult]],
) -> CallToolResult:
"""Helper to run a tool test with minimal boilerplate.
Args:
tools: List of tools to register
call_tool_handler: Handler function for tool calls
test_callback: Async function that performs the test using the client session
Returns:
The result of the tool call
"""
server = Server("test")
@server.list_tools()
async def list_tools():
return tools
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
return await call_tool_handler(name, arguments)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
# Message handler for client
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
if isinstance(message, Exception):
raise message
# Server task
async def run_server():
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
),
) as server_session:
async with anyio.create_task_group() as tg:
async def handle_messages():
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, {}, False)
tg.start_soon(handle_messages)
await anyio.sleep_forever()
# Run the test
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
# Initialize the session
await client_session.initialize()
# Run the test callback
result = await test_callback(client_session)
# Cancel the server task
tg.cancel_scope.cancel()
return result
def create_add_tool() -> Tool:
"""Create a standard 'add' tool for testing."""
return Tool(
name="add",
description="Add two numbers",
inputSchema={
"type": "object",
"properties": {
"a": {"type": "number"},
"b": {"type": "number"},
},
"required": ["a", "b"],
"additionalProperties": False,
},
)
@pytest.mark.anyio
async def test_valid_tool_call():
"""Test that valid arguments pass validation."""
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
if name == "add":
result = arguments["a"] + arguments["b"]
return [TextContent(type="text", text=f"Result: {result}")]
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("add", {"a": 5, "b": 3})
result = await run_tool_test([create_add_tool()], call_tool_handler, test_callback)
# Verify results
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "Result: 8"
@pytest.mark.anyio
async def test_invalid_tool_call_missing_required():
"""Test that missing required arguments fail validation."""
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
# This should not be reached due to validation
raise RuntimeError("Should not reach here")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("add", {"a": 5}) # missing 'b'
result = await run_tool_test([create_add_tool()], call_tool_handler, test_callback)
# Verify results
assert result is not None
assert result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert "Input validation error" in result.content[0].text
assert "'b' is a required property" in result.content[0].text
@pytest.mark.anyio
async def test_invalid_tool_call_wrong_type():
"""Test that wrong argument types fail validation."""
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
# This should not be reached due to validation
raise RuntimeError("Should not reach here")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("add", {"a": "five", "b": 3}) # 'a' should be number
result = await run_tool_test([create_add_tool()], call_tool_handler, test_callback)
# Verify results
assert result is not None
assert result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert "Input validation error" in result.content[0].text
assert "'five' is not of type 'number'" in result.content[0].text
@pytest.mark.anyio
async def test_cache_refresh_on_missing_tool():
"""Test that tool cache is refreshed when tool is not found."""
tools = [
Tool(
name="multiply",
description="Multiply two numbers",
inputSchema={
"type": "object",
"properties": {
"x": {"type": "number"},
"y": {"type": "number"},
},
"required": ["x", "y"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
if name == "multiply":
result = arguments["x"] * arguments["y"]
return [TextContent(type="text", text=f"Result: {result}")]
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
# Call tool without first listing tools (cache should be empty)
# The cache should be refreshed automatically
return await client_session.call_tool("multiply", {"x": 10, "y": 20})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results - should work because cache will be refreshed
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "Result: 200"
@pytest.mark.anyio
async def test_enum_constraint_validation():
"""Test that enum constraints are validated."""
tools = [
Tool(
name="greet",
description="Greet someone",
inputSchema={
"type": "object",
"properties": {
"name": {"type": "string"},
"title": {"type": "string", "enum": ["Mr", "Ms", "Dr"]},
},
"required": ["name"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
# This should not be reached due to validation failure
raise RuntimeError("Should not reach here")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("greet", {"name": "Smith", "title": "Prof"}) # Invalid title
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results
assert result is not None
assert result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert "Input validation error" in result.content[0].text
assert "'Prof' is not one of" in result.content[0].text
@pytest.mark.anyio
async def test_tool_not_in_list_logs_warning(caplog):
"""Test that calling a tool not in list_tools logs a warning and skips validation."""
tools = [
Tool(
name="add",
description="Add two numbers",
inputSchema={
"type": "object",
"properties": {
"a": {"type": "number"},
"b": {"type": "number"},
},
"required": ["a", "b"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
# This should be reached since validation is skipped for unknown tools
if name == "unknown_tool":
# Even with invalid arguments, this should execute since validation is skipped
return [TextContent(type="text", text="Unknown tool executed without validation")]
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
# Call a tool that's not in the list with invalid arguments
# This should trigger the warning about validation not being performed
return await client_session.call_tool("unknown_tool", {"invalid": "args"})
with caplog.at_level(logging.WARNING):
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results - should succeed because validation is skipped for unknown tools
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "Unknown tool executed without validation"
# Verify warning was logged
assert any(
"Tool 'unknown_tool' not listed, no validation will be performed" in record.message for record in caplog.records
)

View file

@ -0,0 +1,433 @@
"""Test output schema validation for lowlevel server."""
import json
from collections.abc import Awaitable, Callable
from typing import Any
import anyio
import pytest
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import CallToolResult, ClientResult, ServerNotification, ServerRequest, TextContent, Tool
async def run_tool_test(
tools: list[Tool],
call_tool_handler: Callable[[str, dict[str, Any]], Awaitable[Any]],
test_callback: Callable[[ClientSession], Awaitable[CallToolResult]],
) -> CallToolResult:
"""Helper to run a tool test with minimal boilerplate.
Args:
tools: List of tools to register
call_tool_handler: Handler function for tool calls
test_callback: Async function that performs the test using the client session
Returns:
The result of the tool call
"""
server = Server("test")
@server.list_tools()
async def list_tools():
return tools
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
return await call_tool_handler(name, arguments)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
# Message handler for client
async def message_handler(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
if isinstance(message, Exception):
raise message
# Server task
async def run_server():
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="test-server",
server_version="1.0.0",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
),
) as server_session:
async with anyio.create_task_group() as tg:
async def handle_messages():
async for message in server_session.incoming_messages:
await server._handle_message(message, server_session, {}, False)
tg.start_soon(handle_messages)
await anyio.sleep_forever()
# Run the test
async with anyio.create_task_group() as tg:
tg.start_soon(run_server)
async with ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as client_session:
# Initialize the session
await client_session.initialize()
# Run the test callback
result = await test_callback(client_session)
# Cancel the server task
tg.cancel_scope.cancel()
return result
@pytest.mark.anyio
async def test_content_only_without_output_schema():
"""Test returning content only when no outputSchema is defined."""
tools = [
Tool(
name="echo",
description="Echo a message",
inputSchema={
"type": "object",
"properties": {
"message": {"type": "string"},
},
"required": ["message"],
},
# No outputSchema defined
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
if name == "echo":
return [TextContent(type="text", text=f"Echo: {arguments['message']}")]
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("echo", {"message": "Hello"})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "Echo: Hello"
assert result.structuredContent is None
@pytest.mark.anyio
async def test_dict_only_without_output_schema():
"""Test returning dict only when no outputSchema is defined."""
tools = [
Tool(
name="get_info",
description="Get structured information",
inputSchema={
"type": "object",
"properties": {},
},
# No outputSchema defined
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
if name == "get_info":
return {"status": "ok", "data": {"value": 42}}
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("get_info", {})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
# Check that the content is the JSON serialization
assert json.loads(result.content[0].text) == {"status": "ok", "data": {"value": 42}}
assert result.structuredContent == {"status": "ok", "data": {"value": 42}}
@pytest.mark.anyio
async def test_both_content_and_dict_without_output_schema():
"""Test returning both content and dict when no outputSchema is defined."""
tools = [
Tool(
name="process",
description="Process data",
inputSchema={
"type": "object",
"properties": {},
},
# No outputSchema defined
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> tuple[list[TextContent], dict[str, Any]]:
if name == "process":
content = [TextContent(type="text", text="Processing complete")]
data = {"result": "success", "count": 10}
return (content, data)
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("process", {})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "Processing complete"
assert result.structuredContent == {"result": "success", "count": 10}
@pytest.mark.anyio
async def test_content_only_with_output_schema_error():
"""Test error when outputSchema is defined but only content is returned."""
tools = [
Tool(
name="structured_tool",
description="Tool expecting structured output",
inputSchema={
"type": "object",
"properties": {},
},
outputSchema={
"type": "object",
"properties": {
"result": {"type": "string"},
},
"required": ["result"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> list[TextContent]:
# This returns only content, but outputSchema expects structured data
return [TextContent(type="text", text="This is not structured")]
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("structured_tool", {})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify error
assert result is not None
assert result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert "Output validation error: outputSchema defined but no structured output returned" in result.content[0].text
@pytest.mark.anyio
async def test_valid_dict_with_output_schema():
"""Test valid dict output matching outputSchema."""
tools = [
Tool(
name="calc",
description="Calculate result",
inputSchema={
"type": "object",
"properties": {
"x": {"type": "number"},
"y": {"type": "number"},
},
"required": ["x", "y"],
},
outputSchema={
"type": "object",
"properties": {
"sum": {"type": "number"},
"product": {"type": "number"},
},
"required": ["sum", "product"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
if name == "calc":
x = arguments["x"]
y = arguments["y"]
return {"sum": x + y, "product": x * y}
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("calc", {"x": 3, "y": 4})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
# Check JSON serialization
assert json.loads(result.content[0].text) == {"sum": 7, "product": 12}
assert result.structuredContent == {"sum": 7, "product": 12}
@pytest.mark.anyio
async def test_invalid_dict_with_output_schema():
"""Test dict output that doesn't match outputSchema."""
tools = [
Tool(
name="user_info",
description="Get user information",
inputSchema={
"type": "object",
"properties": {},
},
outputSchema={
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"},
},
"required": ["name", "age"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
if name == "user_info":
# Missing required 'age' field
return {"name": "Alice"}
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("user_info", {})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify error
assert result is not None
assert result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert isinstance(result.content[0], TextContent)
assert "Output validation error:" in result.content[0].text
assert "'age' is a required property" in result.content[0].text
@pytest.mark.anyio
async def test_both_content_and_valid_dict_with_output_schema():
"""Test returning both content and valid dict with outputSchema."""
tools = [
Tool(
name="analyze",
description="Analyze data",
inputSchema={
"type": "object",
"properties": {
"text": {"type": "string"},
},
"required": ["text"],
},
outputSchema={
"type": "object",
"properties": {
"sentiment": {"type": "string", "enum": ["positive", "negative", "neutral"]},
"confidence": {"type": "number", "minimum": 0, "maximum": 1},
},
"required": ["sentiment", "confidence"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> tuple[list[TextContent], dict[str, Any]]:
if name == "analyze":
content = [TextContent(type="text", text=f"Analysis of: {arguments['text']}")]
data = {"sentiment": "positive", "confidence": 0.95}
return (content, data)
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("analyze", {"text": "Great job!"})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify results
assert result is not None
assert not result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert result.content[0].text == "Analysis of: Great job!"
assert result.structuredContent == {"sentiment": "positive", "confidence": 0.95}
@pytest.mark.anyio
async def test_output_schema_type_validation():
"""Test outputSchema validates types correctly."""
tools = [
Tool(
name="stats",
description="Get statistics",
inputSchema={
"type": "object",
"properties": {},
},
outputSchema={
"type": "object",
"properties": {
"count": {"type": "integer"},
"average": {"type": "number"},
"items": {"type": "array", "items": {"type": "string"}},
},
"required": ["count", "average", "items"],
},
)
]
async def call_tool_handler(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
if name == "stats":
# Wrong type for 'count' - should be integer
return {"count": "five", "average": 2.5, "items": ["a", "b"]}
else:
raise ValueError(f"Unknown tool: {name}")
async def test_callback(client_session: ClientSession) -> CallToolResult:
return await client_session.call_tool("stats", {})
result = await run_tool_test(tools, call_tool_handler, test_callback)
# Verify error
assert result is not None
assert result.isError
assert len(result.content) == 1
assert result.content[0].type == "text"
assert "Output validation error:" in result.content[0].text
assert "'five' is not of type 'integer'" in result.content[0].text