Files
wehub-resource-sync 60e0ffc959
Upgrade checks / Notify on failure (push) Has been cancelled
Upgrade checks / Close issue on success (push) Has been cancelled
Schema Crash Test / Real-world schema crash test (232K schemas) (push) Has been cancelled
Run static analysis / static_analysis (push) Has been cancelled
Tests / Tests: Python 3.10 on ubuntu-latest (push) Has been cancelled
Tests / Tests: Python 3.13 on ubuntu-latest (push) Has been cancelled
Tests / Tests: Python 3.10 on windows-latest (push) Has been cancelled
Tests / Tests with lowest-direct dependencies (push) Has been cancelled
Tests / MCP conformance tests (push) Has been cancelled
Tests / Integration tests (push) Has been cancelled
Tests / Package install smoke (push) Has been cancelled
Upgrade checks / Static analysis (push) Has been cancelled
Upgrade checks / Tests: Python 3.10 on ubuntu-latest (push) Has been cancelled
Upgrade checks / Tests: Python 3.13 on ubuntu-latest (push) Has been cancelled
Upgrade checks / Tests: Python 3.10 on windows-latest (push) Has been cancelled
Upgrade checks / Integration tests (push) Has been cancelled
Update MCPServerConfig Schema / update-config-schema (push) Has been cancelled
Update SDK Documentation / update-sdk-docs (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:39:59 +08:00

626 lines
23 KiB
Python

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
import mcp_types
import pytest
from mcp.server.context import ServerRequestContext
from fastmcp import Client, FastMCP
from fastmcp.server.context import Context
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from fastmcp.tools.base import ToolResult
@dataclass
class Recording:
# the hook is the name of the hook that was called, e.g. "on_list_tools"
hook: str
context: MiddlewareContext
result: mcp_types.ServerResult | None
class RecordingMiddleware(Middleware):
"""A middleware that automatically records all method calls."""
def __init__(self, name: str | None = None):
super().__init__()
self.calls: list[Recording] = []
self.name = name
def __getattribute__(self, name: str) -> Callable:
"""Dynamically create recording methods for any on_* method."""
if name.startswith("on_"):
async def record_and_call(
context: MiddlewareContext, call_next: Callable
) -> Any:
result = await call_next(context)
self.calls.append(Recording(hook=name, context=context, result=result))
return result
return record_and_call
return super().__getattribute__(name)
def get_calls(
self, method: str | None = None, hook: str | None = None
) -> list[Recording]:
"""
Get all recorded calls for a specific method or hook.
Args:
method: The method to filter by (e.g. "tools/list")
hook: The hook to filter by (e.g. "on_list_tools")
Returns:
A list of recorded calls.
"""
calls = []
for recording in self.calls:
if method and hook:
if recording.context.method == method and recording.hook == hook:
calls.append(recording)
elif method:
if recording.context.method == method:
calls.append(recording)
elif hook:
if recording.hook == hook:
calls.append(recording)
else:
calls.append(recording)
return calls
def assert_called(
self,
hook: str | None = None,
method: str | None = None,
times: int | None = None,
at_least: int | None = None,
) -> bool:
"""Assert that a hook was called a specific number of times."""
if times is not None and at_least is not None:
raise ValueError("Cannot specify both times and at_least")
elif times is None and at_least is None:
times = 1
calls = self.get_calls(hook=hook, method=method)
actual_times = len(calls)
identifier = dict(hook=hook, method=method)
if times is not None:
assert actual_times == times, (
f"Expected {times} calls for {identifier}, "
f"but was called {actual_times} times"
)
elif at_least is not None:
assert actual_times >= at_least, (
f"Expected at least {at_least} calls for {identifier}, "
f"but was called {actual_times} times"
)
return True
def assert_not_called(self, hook: str | None = None, method: str | None = None):
"""Assert that a hook was not called."""
calls = self.get_calls(hook=hook, method=method)
assert len(calls) == 0, f"Expected {hook!r} to not be called"
return True
def reset(self):
"""Clear all recorded calls."""
self.calls.clear()
@pytest.fixture
def recording_middleware():
"""Fixture that provides a recording middleware instance."""
middleware = RecordingMiddleware(name="recording_middleware")
yield middleware
@pytest.fixture
def mcp_server(recording_middleware):
mcp = FastMCP()
@mcp.tool(tags={"add-tool"})
def add(a: int, b: int) -> int:
return a + b
@mcp.resource("resource://test")
def test_resource() -> str:
return "test resource"
@mcp.resource("resource://test-template/{x}")
def test_resource_with_path(x: int) -> str:
return f"test resource with {x}"
@mcp.prompt
def test_prompt(x: str) -> str:
return f"test prompt with {x}"
@mcp.tool
async def progress_tool(context: Context) -> None:
await context.report_progress(progress=1, total=10, message="test")
@mcp.tool
async def log_tool(context: Context) -> None:
await context.info(message="test log")
@mcp.tool
async def sample_tool(context: Context) -> None:
await context.sample("hello")
mcp.add_middleware(recording_middleware)
# Register a progress notification handler (v2 API: (ctx, params)).
async def handle_progress(
_ctx: ServerRequestContext,
_params: mcp_types.ProgressNotificationParams,
) -> None:
pass
mcp._mcp_server.add_notification_handler(
"notifications/progress",
mcp_types.ProgressNotificationParams,
handle_progress,
)
return mcp
class TestMiddlewareHooks:
async def test_call_tool(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.call_tool("add", {"a": 1, "b": 2})
assert recording_middleware.assert_called(at_least=9)
assert recording_middleware.assert_called(method="tools/call", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_call_tool", at_least=1)
async def test_read_resource(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.read_resource("resource://test")
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(method="resources/read", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_read_resource", at_least=1)
async def test_read_resource_template(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.read_resource("resource://test-template/1")
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(method="resources/read", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_read_resource", at_least=1)
async def test_get_prompt(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.get_prompt("test_prompt", {"x": "test"})
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(method="prompts/get", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_get_prompt", at_least=1)
async def test_list_tools(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_tools()
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(method="tools/list", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_list_tools", at_least=1)
# Verify the middleware receives a list of tools
list_tools_calls = recording_middleware.get_calls(hook="on_list_tools")
assert len(list_tools_calls) > 0
result = list_tools_calls[0].result
assert isinstance(result, list)
async def test_list_resources(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_resources()
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(method="resources/list", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_list_resources", at_least=1)
# Verify the middleware receives a list of resources
list_resources_calls = recording_middleware.get_calls(hook="on_list_resources")
assert len(list_resources_calls) > 0
result = list_resources_calls[0].result
assert isinstance(result, list)
async def test_list_resource_templates(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_resource_templates()
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(
method="resources/templates/list", at_least=3
)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(
hook="on_list_resource_templates", at_least=1
)
# Verify the middleware receives a list of resource templates
list_templates_calls = recording_middleware.get_calls(
hook="on_list_resource_templates"
)
assert len(list_templates_calls) > 0
result = list_templates_calls[0].result
assert isinstance(result, list)
async def test_list_prompts(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_prompts()
assert recording_middleware.assert_called(at_least=3)
assert recording_middleware.assert_called(method="prompts/list", at_least=3)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_list_prompts", at_least=1)
# Verify the middleware receives a list of prompts
list_prompts_calls = recording_middleware.get_calls(hook="on_list_prompts")
assert len(list_prompts_calls) > 0
result = list_prompts_calls[0].result
assert isinstance(result, list)
async def test_initialize(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.ping()
assert recording_middleware.assert_called(at_least=1)
assert recording_middleware.assert_called(hook="on_message", at_least=1)
assert recording_middleware.assert_called(hook="on_request", at_least=1)
assert recording_middleware.assert_called(hook="on_initialize", at_least=1)
async def test_list_tools_filtering_middleware(self):
"""Test that middleware can filter tools."""
class FilteringMiddleware(Middleware):
async def on_list_tools(self, context: MiddlewareContext, call_next):
result = await call_next(context)
# Filter out tools with "private" tag - simple list filtering
filtered_tools = [tool for tool in result if "private" not in tool.tags]
return filtered_tools
server = FastMCP("TestServer")
@server.tool
def public_tool(name: str) -> str:
return f"Hello {name}"
@server.tool(tags={"private"})
def private_tool(secret: str) -> str:
return f"Secret: {secret}"
server.add_middleware(FilteringMiddleware())
async with Client(server) as client:
tools = await client.list_tools()
assert len(tools) == 1
assert tools[0].name == "public_tool"
async def test_list_resources_filtering_middleware(self):
"""Test that middleware can filter resources."""
class FilteringMiddleware(Middleware):
async def on_list_resources(self, context: MiddlewareContext, call_next):
result = await call_next(context)
# Filter out resources with "private" tag
filtered_resources = [
resource for resource in result if "private" not in resource.tags
]
return filtered_resources
server = FastMCP("TestServer")
@server.resource("resource://public")
def public_resource() -> str:
return "public data"
@server.resource("resource://private", tags={"private"})
def private_resource() -> str:
return "private data"
server.add_middleware(FilteringMiddleware())
async with Client(server) as client:
resources = await client.list_resources()
assert len(resources) == 1
assert str(resources[0].uri) == "resource://public"
async def test_list_resource_templates_filtering_middleware(self):
"""Test that middleware can filter resource templates."""
class FilteringMiddleware(Middleware):
async def on_list_resource_templates(
self, context: MiddlewareContext, call_next
):
result = await call_next(context)
# Filter out templates with "private" tag
filtered_templates = [
template for template in result if "private" not in template.tags
]
return filtered_templates
server = FastMCP("TestServer")
@server.resource("resource://public/{x}")
def public_template(x: str) -> str:
return f"public {x}"
@server.resource("resource://private/{x}", tags={"private"})
def private_template(x: str) -> str:
return f"private {x}"
server.add_middleware(FilteringMiddleware())
async with Client(server) as client:
templates = await client.list_resource_templates()
assert len(templates) == 1
assert str(templates[0].uri_template) == "resource://public/{x}"
async def test_list_prompts_filtering_middleware(self):
"""Test that middleware can filter prompts."""
class FilteringMiddleware(Middleware):
async def on_list_prompts(self, context: MiddlewareContext, call_next):
result = await call_next(context)
# Filter out prompts with "private" tag
filtered_prompts = [
prompt for prompt in result if "private" not in prompt.tags
]
return filtered_prompts
server = FastMCP("TestServer")
@server.prompt
def public_prompt(name: str) -> str:
return f"Hello {name}"
@server.prompt(tags={"private"})
def private_prompt(secret: str) -> str:
return f"Secret: {secret}"
server.add_middleware(FilteringMiddleware())
async with Client(server) as client:
prompts = await client.list_prompts()
assert len(prompts) == 1
assert prompts[0].name == "public_prompt"
async def test_call_tool_middleware(self):
server = FastMCP()
@server.tool
def add(a: int, b: int) -> int:
return a + b
class CallToolMiddleware(Middleware):
async def on_call_tool(
self,
context: MiddlewareContext[mcp_types.CallToolRequestParams],
call_next: CallNext[mcp_types.CallToolRequestParams, ToolResult],
):
# modify argument
if context.message.name == "add":
assert context.message.arguments is not None
args = context.message.arguments
assert isinstance(args["a"], int)
args["a"] += 100
result = await call_next(context)
# modify result
if context.message.name == "add":
assert result.structured_content is not None
content = result.structured_content
assert isinstance(content["result"], int)
content["result"] += 5
return result
server.add_middleware(CallToolMiddleware())
async with Client(server) as client:
result = await client.call_tool("add", {"a": 1, "b": 2})
assert isinstance(result.structured_content["result"], int)
assert result.structured_content["result"] == 108
class TestApplyMiddlewareParameter:
"""Tests for run_middleware parameter on execution methods."""
async def test_call_tool_with_run_middleware_true(self):
"""Middleware is applied when run_middleware=True (default)."""
recording = RecordingMiddleware()
server = FastMCP()
@server.tool
def add(a: int, b: int) -> int:
return a + b
server.add_middleware(recording)
result = await server.call_tool("add", {"a": 1, "b": 2})
assert result.structured_content["result"] == 3 # type: ignore[union-attr,index] # ty:ignore[not-subscriptable]
assert recording.assert_called(hook="on_call_tool", times=1)
async def test_call_tool_with_run_middleware_false(self):
"""Middleware is NOT applied when run_middleware=False."""
recording = RecordingMiddleware()
server = FastMCP()
@server.tool
def add(a: int, b: int) -> int:
return a + b
server.add_middleware(recording)
result = await server.call_tool("add", {"a": 1, "b": 2}, run_middleware=False)
assert result.structured_content["result"] == 3 # type: ignore[union-attr,index] # ty:ignore[not-subscriptable]
# Middleware should not have been called
assert len(recording.calls) == 0
async def test_read_resource_with_run_middleware_true(self):
"""Middleware is applied when run_middleware=True (default)."""
recording = RecordingMiddleware()
server = FastMCP()
@server.resource("resource://test")
def test_resource() -> str:
return "test content"
server.add_middleware(recording)
result = await server.read_resource("resource://test")
assert len(result.contents) == 1
assert result.contents[0].content == "test content"
assert recording.assert_called(hook="on_read_resource", times=1)
async def test_read_resource_with_run_middleware_false(self):
"""Middleware is NOT applied when run_middleware=False."""
recording = RecordingMiddleware()
server = FastMCP()
@server.resource("resource://test")
def test_resource() -> str:
return "test content"
server.add_middleware(recording)
result = await server.read_resource("resource://test", run_middleware=False)
assert len(result.contents) == 1
assert result.contents[0].content == "test content"
# Middleware should not have been called
assert len(recording.calls) == 0
async def test_read_resource_template_with_run_middleware_false(self):
"""Templates also skip middleware when run_middleware=False."""
recording = RecordingMiddleware()
server = FastMCP()
@server.resource("resource://items/{item_id}")
def get_item(item_id: int) -> str:
return f"item {item_id}"
server.add_middleware(recording)
result = await server.read_resource("resource://items/42", run_middleware=False)
assert len(result.contents) == 1
assert result.contents[0].content == "item 42"
assert len(recording.calls) == 0
async def test_render_prompt_with_run_middleware_true(self):
"""Middleware is applied when run_middleware=True (default)."""
recording = RecordingMiddleware()
server = FastMCP()
@server.prompt
def greet(name: str) -> str:
return f"Hello, {name}!"
server.add_middleware(recording)
result = await server.render_prompt("greet", {"name": "World"})
assert len(result.messages) == 1
# content is TextContent | EmbeddedResource, but we know it's TextContent from the test
assert isinstance(result.messages[0].content, mcp_types.TextContent)
assert result.messages[0].content.text == "Hello, World!"
assert recording.assert_called(hook="on_get_prompt", times=1)
async def test_render_prompt_with_run_middleware_false(self):
"""Middleware is NOT applied when run_middleware=False."""
recording = RecordingMiddleware()
server = FastMCP()
@server.prompt
def greet(name: str) -> str:
return f"Hello, {name}!"
server.add_middleware(recording)
result = await server.render_prompt(
"greet", {"name": "World"}, run_middleware=False
)
assert len(result.messages) == 1
# content is TextContent | EmbeddedResource, but we know it's TextContent from the test
assert isinstance(result.messages[0].content, mcp_types.TextContent)
assert result.messages[0].content.text == "Hello, World!"
# Middleware should not have been called
assert len(recording.calls) == 0
async def test_middleware_modification_skipped_when_run_middleware_false(self):
"""Middleware that modifies args/results is skipped."""
class ModifyingMiddleware(Middleware):
async def on_call_tool(self, context: MiddlewareContext, call_next):
# Double the 'a' argument
assert context.message.arguments is not None
context.message.arguments["a"] *= 2
return await call_next(context)
server = FastMCP()
@server.tool
def add(a: int, b: int) -> int:
return a + b
server.add_middleware(ModifyingMiddleware())
# With middleware: a=5 becomes a=10, result = 10 + 3 = 13
result_with = await server.call_tool("add", {"a": 5, "b": 3})
assert result_with.structured_content["result"] == 13 # type: ignore[union-attr,index] # ty:ignore[not-subscriptable]
# Without middleware: a=5 stays a=5, result = 5 + 3 = 8
result_without = await server.call_tool(
"add", {"a": 5, "b": 3}, run_middleware=False
)
assert result_without.structured_content["result"] == 8 # type: ignore[union-attr,index] # ty:ignore[not-subscriptable]