refactor: replace lowlevel Server decorators with on_* constructor kwargs (#1985)

This commit is contained in:
Max Isbey 2026-02-12 15:55:54 +00:00 • committed by GitHub
commit 0a22a9dc33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 3563 additions and 5481 deletions

View file

@ -11,7 +11,7 @@ from inline_snapshot import snapshot
from mcp import types
from mcp.client._memory import InMemoryTransport
from mcp.client.client import Client
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import (
CallToolResult,
@ -41,33 +41,36 @@ pytestmark = pytest.mark.anyio
@pytest.fixture
def simple_server() -> Server:
"""Create a simple MCP server for testing."""
server = Server(name="test_server")
@server.list_resources()
async def handle_list_resources():
return [Resource(uri="memory://test", name="Test Resource", description="A test resource")]
async def handle_list_resources(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> ListResourcesResult:
return ListResourcesResult(
resources=[Resource(uri="memory://test", name="Test Resource", description="A test resource")]
)
@server.subscribe_resource()
async def handle_subscribe_resource(uri: str):
pass
async def handle_subscribe_resource(ctx: ServerRequestContext, params: types.SubscribeRequestParams) -> EmptyResult:
return EmptyResult()
@server.unsubscribe_resource()
async def handle_unsubscribe_resource(uri: str):
pass
async def handle_unsubscribe_resource(
ctx: ServerRequestContext, params: types.UnsubscribeRequestParams
) -> EmptyResult:
return EmptyResult()
@server.set_logging_level()
async def handle_set_logging_level(level: str):
pass
async def handle_set_logging_level(ctx: ServerRequestContext, params: types.SetLevelRequestParams) -> EmptyResult:
return EmptyResult()
@server.completion()
async def handle_completion(
ref: types.PromptReference | types.ResourceTemplateReference,
argument: types.CompletionArgument,
context: types.CompletionContext | None,
) -> types.Completion | None:
return types.Completion(values=[])
async def handle_completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> types.CompleteResult:
return types.CompleteResult(completion=types.Completion(values=[]))
return server
return Server(
name="test_server",
on_list_resources=handle_list_resources,
on_subscribe_resource=handle_subscribe_resource,
on_unsubscribe_resource=handle_unsubscribe_resource,
on_set_logging_level=handle_set_logging_level,
on_completion=handle_completion,
)
@pytest.fixture
@ -202,19 +205,14 @@ async def test_client_send_progress_notification():
"""Test sending progress notification."""
received_from_client = None
event = anyio.Event()
server = Server(name="test_server")
@server.progress_notification()
async def handle_progress_notification(
progress_token: str | int,
progress: float = 0.0,
total: float | None = None,
message: str | None = None,
) -> None:
async def handle_progress(ctx: ServerRequestContext, params: types.ProgressNotificationParams) -> None:
nonlocal received_from_client
received_from_client = {"progress_token": progress_token, "progress": progress}
received_from_client = {"progress_token": params.progress_token, "progress": params.progress}
event.set()
server = Server(name="test_server", on_progress=handle_progress)
async with Client(server) as client:
await client.send_progress_notification(progress_token="token123", progress=50.0)
await event.wait()

View file

@ -8,7 +8,6 @@ import multiprocessing
import socket
from collections.abc import AsyncGenerator, Generator
from contextlib import asynccontextmanager
from typing import Any
import pytest
from starlette.applications import Starlette
@ -17,7 +16,7 @@ from starlette.routing import Mount
from mcp import types
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.types import TextContent, Tool
from tests.test_helpers import wait_for_server
@ -47,54 +46,56 @@ def run_unicode_server(port: int) -> None: # pragma: no cover
import uvicorn
# Need to recreate the server setup in this process
server = Server(name="unicode_test_server")
@server.list_tools()
async def list_tools() -> list[Tool]:
"""List tools with Unicode descriptions."""
return [
Tool(
name="echo_unicode",
description="🔤 Echo Unicode text - Hello 👋 World 🌍 - Testing 🧪 Unicode ✨",
input_schema={
"type": "object",
"properties": {
"text": {"type": "string", "description": "Text to echo back"},
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
return types.ListToolsResult(
tools=[
Tool(
name="echo_unicode",
description="🔤 Echo Unicode text - Hello 👋 World 🌍 - Testing 🧪 Unicode ✨",
input_schema={
"type": "object",
"properties": {
"text": {"type": "string", "description": "Text to echo back"},
},
"required": ["text"],
},
"required": ["text"],
},
),
]
),
]
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]:
"""Handle tool calls with Unicode content."""
if name == "echo_unicode":
text = arguments.get("text", "") if arguments else ""
return [
TextContent(
type="text",
text=f"Echo: {text}",
async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
if params.name == "echo_unicode":
text = params.arguments.get("text", "") if params.arguments else ""
return types.CallToolResult(
content=[
TextContent(
type="text",
text=f"Echo: {text}",
)
]
)
else:
raise ValueError(f"Unknown tool: {params.name}")
async def handle_list_prompts(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListPromptsResult:
return types.ListPromptsResult(
prompts=[
types.Prompt(
name="unicode_prompt",
description="Unicode prompt - Слой хранилища, где располагаются",
arguments=[],
)
]
else:
raise ValueError(f"Unknown tool: {name}")
)
@server.list_prompts()
async def list_prompts() -> list[types.Prompt]:
"""List prompts with Unicode names and descriptions."""
return [
types.Prompt(
name="unicode_prompt",
description="Unicode prompt - Слой хранилища, где располагаются",
arguments=[],
)
]
@server.get_prompt()
async def get_prompt(name: str, arguments: dict[str, Any] | None) -> types.GetPromptResult:
"""Get a prompt with Unicode content."""
if name == "unicode_prompt":
async def handle_get_prompt(
ctx: ServerRequestContext, params: types.GetPromptRequestParams
) -> types.GetPromptResult:
if params.name == "unicode_prompt":
return types.GetPromptResult(
messages=[
types.PromptMessage(
@ -106,7 +107,15 @@ def run_unicode_server(port: int) -> None: # pragma: no cover
)
]
)
raise ValueError(f"Unknown prompt: {name}")
raise ValueError(f"Unknown prompt: {params.name}")
server = Server(
name="unicode_test_server",
on_list_tools=handle_list_tools,
on_call_tool=handle_call_tool,
on_list_prompts=handle_list_prompts,
on_get_prompt=handle_get_prompt,
)
# Create the session manager
session_manager = StreamableHTTPSessionManager(

View file

@ -3,9 +3,9 @@ from collections.abc import Callable
import pytest
from mcp import Client, types
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import ListToolsRequest, ListToolsResult
from mcp.types import ListToolsResult
from .conftest import StreamSpyCollection
@ -105,14 +105,16 @@ async def test_list_tools_with_strict_server_validation(
async def test_list_tools_with_lowlevel_server():
"""Test that list_tools works with a lowlevel Server using params."""
server = Server("test-lowlevel")
@server.list_tools()
async def handle_list_tools(request: ListToolsRequest) -> ListToolsResult:
async def handle_list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> ListToolsResult:
# Echo back what cursor we received in the tool description
cursor = request.params.cursor if request.params else None
cursor = params.cursor if params else None
return ListToolsResult(tools=[types.Tool(name="test_tool", description=f"cursor={cursor}", input_schema={})])
server = Server("test-lowlevel", on_list_tools=handle_list_tools)
async with Client(server) as client:
result = await client.list_tools()
assert result.tools[0].description == "cursor=None"

View file

@ -1,50 +1,41 @@
import inspect
import logging
from contextlib import contextmanager
from typing import Any
from unittest.mock import patch
import jsonschema
import pytest
from mcp import Client
from mcp.server.lowlevel import Server
from mcp.types import Tool
from mcp.server import Server, ServerRequestContext
from mcp.types import (
CallToolRequestParams,
CallToolResult,
ListToolsResult,
PaginatedRequestParams,
TextContent,
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
def _make_server(
tools: list[Tool],
structured_content: dict[str, Any],
) -> Server:
"""Create a low-level server that returns the given structured_content for any tool call."""
# 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)
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(tools=tools)
with patch("jsonschema.validate", selective_mock):
yield
async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
return CallToolResult(
content=[TextContent(type="text", text="result")],
structured_content=structured_content,
)
return Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
@pytest.mark.anyio
async def test_tool_structured_output_client_side_validation_basemodel():
"""Test that client validates structured content against schema for BaseModel outputs"""
# Create a malicious low-level server that returns invalid structured content
server = Server("test-server")
# Define the expected schema for our tool
output_schema = {
"type": "object",
"properties": {"name": {"type": "string", "title": "Name"}, "age": {"type": "integer", "title": "Age"}},
@ -52,39 +43,27 @@ async def test_tool_structured_output_client_side_validation_basemodel():
"title": "UserOutput",
}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="get_user",
description="Get user data",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"name": "John", "age": "invalid"}, # Invalid: age should be int
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return invalid structured content - age is string instead of integer
# The low-level server will wrap this in CallToolResult
return {"name": "John", "age": "invalid"} # Invalid: age should be int
# 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)
async with Client(server) as client:
with pytest.raises(RuntimeError) as exc_info:
await client.call_tool("get_user", {})
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"""
server = Server("test-server")
# Primitive types are wrapped in {"result": value}
output_schema = {
"type": "object",
"properties": {"result": {"type": "integer", "title": "Result"}},
@ -92,122 +71,95 @@ async def test_tool_structured_output_client_side_validation_primitive():
"title": "calculate_Output",
}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="calculate",
description="Calculate something",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"result": "not_a_number"}, # Invalid: should be int
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return invalid structured content - result is string instead of integer
return {"result": "not_a_number"} # Invalid: should be int
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)
async with Client(server) as client:
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"""
server = Server("test-server")
# dict[str, int] schema
output_schema = {"type": "object", "additionalProperties": {"type": "integer"}, "title": "get_scores_Output"}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="get_scores",
description="Get scores",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"alice": "100", "bob": "85"}, # Invalid: values should be int
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return invalid structured content - values should be integers
return {"alice": "100", "bob": "85"} # Invalid: values should be int
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)
async with Client(server) as client:
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"""
server = Server("test-server")
output_schema = {
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}, "email": {"type": "string"}},
"required": ["name", "age", "email"], # All fields required
"required": ["name", "age", "email"],
"title": "PersonOutput",
}
@server.list_tools()
async def list_tools():
return [
server = _make_server(
tools=[
Tool(
name="get_person",
description="Get person data",
input_schema={"type": "object"},
output_schema=output_schema,
)
]
],
structured_content={"name": "John", "age": 30}, # Missing required 'email'
)
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]):
# Return structured content missing required field 'email'
return {"name": "John", "age": 30} # Missing required 'email'
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)
async with Client(server) as client:
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"""
server = Server("test-server")
@server.list_tools()
async def list_tools() -> list[Tool]:
# Return empty list - tool is not listed
return []
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
return ListToolsResult(tools=[])
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
# Server still responds to the tool call with structured content
return {"result": 42}
async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
return CallToolResult(
content=[TextContent(type="text", text="result")],
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
async with Client(server) as client:
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
assert "Tool mystery_tool not listed" in caplog.text

View file

@ -2,31 +2,31 @@
import pytest
from mcp import Client
from mcp import Client, types
from mcp.client._memory import InMemoryTransport
from mcp.server import Server
from mcp.server import Server, ServerRequestContext
from mcp.server.mcpserver import MCPServer
from mcp.types import Resource
from mcp.types import ListResourcesResult, Resource
@pytest.fixture
def simple_server() -> Server:
"""Create a simple MCP server for testing."""
server = Server(name="test_server")
# pragma: no cover - handler exists only to register a resource capability.
# Transport tests verify stream creation, not handler invocation.
@server.list_resources()
async def handle_list_resources(): # pragma: no cover
return [
Resource(
uri="memory://test",
name="Test Resource",
description="A test resource",
)
]
async def handle_list_resources(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> ListResourcesResult: # pragma: no cover
return ListResourcesResult(
resources=[
Resource(
uri="memory://test",
name="Test Resource",
description="A test resource",
)
]
)
return server
return Server(name="test_server", on_list_resources=handle_list_resources)
@pytest.fixture