fastmcp/tests/server/middleware/test_middleware.py

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import mcp.types
import pytest

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 progress handler
    @mcp._mcp_server.progress_notification()
    async def handle_progress(
        progress_token: str | int,
        progress: float,
        total: float | None,
        message: str | None,
    ):
        print("HI")

    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].uriTemplate) == "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]