fastmcp/tests/server/middleware/test_tool_injection.py

"""Tests for tool injection middleware."""

import math
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager

import pytest
from inline_snapshot import snapshot
from mcp.types import Tool as SDKTool

from fastmcp import FastMCP
from fastmcp.apps import FastMCPApp
from fastmcp.client import Client
from fastmcp.client.client import CallToolResult
from fastmcp.client.transports import FastMCPTransport
from fastmcp.exceptions import ToolError
from fastmcp.server.auth import require_scopes
from fastmcp.server.auth.providers.jwt import StaticTokenVerifier
from fastmcp.server.middleware import AuthMiddleware, Middleware
from fastmcp.server.middleware.tool_injection import (
    ToolInjectionMiddleware,
)
from fastmcp.server.providers.addressing import hashed_backend_name
from fastmcp.tools.base import Tool
from fastmcp.tools.function_tool import FunctionTool
from fastmcp.utilities.tests import run_server_async


def multiply_fn(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b


def divide_fn(a: int, b: int) -> float:
    """Divide two numbers."""
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b


multiply_tool = Tool.from_function(fn=multiply_fn, name="multiply", tags={"math"})
divide_tool = Tool.from_function(fn=divide_fn, name="divide", tags={"math"})


class TestToolInjectionMiddleware:
    """Tests with real FastMCP server."""

    @pytest.fixture
    def base_server(self):
        """Create a base FastMCP server."""
        mcp = FastMCP("BaseServer")

        @mcp.tool
        def add(a: int, b: int) -> int:
            """Add two numbers."""
            return a + b

        @mcp.tool
        def subtract(a: int, b: int) -> int:
            """Subtract two numbers."""
            return a - b

        return mcp

    async def test_list_tools_includes_injected_tools(self, base_server: FastMCP):
        """Test that list_tools returns both base and injected tools."""

        injected_tools: list[FunctionTool] = [
            multiply_tool,
            divide_tool,
        ]
        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(
            tools=injected_tools
        )
        base_server.add_middleware(middleware)

        async with Client[FastMCPTransport](base_server) as client:
            tools: list[SDKTool] = await client.list_tools()

        # Should have all tools: multiply, divide, add, subtract
        assert len(tools) == 4
        tool_names: list[str] = [tool.name for tool in tools]
        assert "multiply" in tool_names
        assert "divide" in tool_names
        assert "add" in tool_names
        assert "subtract" in tool_names

    async def test_call_injected_tool(self, base_server: FastMCP):
        """Test that injected tools can be called successfully."""

        injected_tools: list[FunctionTool] = [multiply_tool]
        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(
            tools=injected_tools
        )
        base_server.add_middleware(middleware)

        async with Client[FastMCPTransport](base_server) as client:
            result: CallToolResult = await client.call_tool(
                name="multiply", arguments={"a": 7, "b": 6}
            )

        assert result.structured_content is not None
        assert isinstance(result.structured_content, dict)
        assert result.structured_content["result"] == 42

    async def test_call_base_tool_still_works(self, base_server: FastMCP):
        """Test that base server tools still work after injecting tools."""

        injected_tools: list[FunctionTool] = [multiply_tool]
        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(
            tools=injected_tools
        )
        base_server.add_middleware(middleware)

        async with Client[FastMCPTransport](base_server) as client:
            result: CallToolResult = await client.call_tool(
                name="add", arguments={"a": 10, "b": 5}
            )

        assert result.structured_content is not None
        assert isinstance(result.structured_content, dict)
        assert result.structured_content["result"] == 15

    async def test_injected_tool_error_handling(self, base_server: FastMCP):
        """Test that errors in injected tools are properly handled."""

        injected_tools: list[FunctionTool] = [divide_tool]
        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(
            tools=injected_tools
        )
        base_server.add_middleware(middleware)

        async with Client[FastMCPTransport](base_server) as client:
            with pytest.raises(Exception, match="Cannot divide by zero"):
                _ = await client.call_tool(name="divide", arguments={"a": 10, "b": 0})

    async def test_multiple_tool_injections(self, base_server: FastMCP):
        """Test multiple tool injection middlewares can be stacked."""

        def power(a: int, b: int) -> int:
            """Raise a to the power of b."""
            return int(math.pow(float(a), float(b)))

        def modulo(a: int, b: int) -> int:
            """Calculate a modulo b."""
            return a % b

        middleware1 = ToolInjectionMiddleware(
            tools=[Tool.from_function(fn=power, name="power")]
        )
        middleware2 = ToolInjectionMiddleware(
            tools=[Tool.from_function(fn=modulo, name="modulo")]
        )

        base_server.add_middleware(middleware1)
        base_server.add_middleware(middleware2)

        async with Client(base_server) as client:
            tools = await client.list_tools()

        # Should have all tools
        assert len(tools) == 4
        tool_names = [tool.name for tool in tools]
        assert "power" in tool_names
        assert "modulo" in tool_names
        assert "add" in tool_names
        assert "subtract" in tool_names

        # Test that both injected tools work
        async with Client(base_server) as client:
            power_result = await client.call_tool("power", {"a": 2, "b": 3})
            assert power_result.structured_content is not None
            assert isinstance(power_result.structured_content, dict)
            assert power_result.structured_content["result"] == 8

            modulo_result = await client.call_tool("modulo", {"a": 10, "b": 3})
            assert modulo_result.structured_content is not None
            assert isinstance(modulo_result.structured_content, dict)
            assert modulo_result.structured_content["result"] == 1

    async def test_injected_tool_with_complex_return_type(self, base_server: FastMCP):
        """Test injected tools with complex return types."""

        def calculate_stats(numbers: list[int]) -> dict[str, int | float]:
            """Calculate statistics for a list of numbers."""
            return {
                "sum": sum(numbers),
                "average": sum(numbers) / len(numbers),
                "min": min(numbers),
                "max": max(numbers),
                "count": len(numbers),
            }

        middleware = ToolInjectionMiddleware(
            tools=[Tool.from_function(fn=calculate_stats, name="calculate_stats")]
        )
        base_server.add_middleware(middleware)

        async with Client(base_server) as client:
            result = await client.call_tool(
                "calculate_stats", {"numbers": [1, 2, 3, 4, 5]}
            )

        assert result.structured_content is not None

        assert isinstance(result.structured_content, dict)

        assert result.structured_content == snapshot(
            {"sum": 15, "average": 3.0, "min": 1, "max": 5, "count": 5}
        )

    async def test_injected_tool_metadata_preserved(self, base_server: FastMCP):
        """Test that injected tool metadata is preserved."""

        def multiply(a: int, b: int) -> int:
            """Multiply two numbers."""
            return a * b

        injected_tools = [Tool.from_function(fn=multiply, name="multiply")]
        middleware = ToolInjectionMiddleware(tools=injected_tools)
        base_server.add_middleware(middleware)

        async with Client(base_server) as client:
            tools = await client.list_tools()

        multiply_tool = next(t for t in tools if t.name == "multiply")
        assert multiply_tool.description == "Multiply two numbers."
        assert "a" in multiply_tool.inputSchema["properties"]
        assert "b" in multiply_tool.inputSchema["properties"]

    async def test_injected_tool_does_not_conflict_with_base_tool(
        self, base_server: FastMCP
    ):
        """Test that injected tools with same name as base tools are called correctly."""

        def add(a: int, b: int) -> int:
            """Injected add that multiplies instead."""
            return a * b

        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(
            tools=[Tool.from_function(fn=add, name="add")]
        )
        base_server.add_middleware(middleware)

        async with Client[FastMCPTransport](base_server) as client:
            result: CallToolResult = await client.call_tool(
                name="add", arguments={"a": 5, "b": 3}
            )

        # Should use the injected tool (multiply behavior)
        assert result.structured_content is not None
        assert result.structured_content["result"] == 15

    async def test_injected_tool_bypass_filtering(self, base_server: FastMCP):
        """Test that injected tools bypass filtering."""
        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(
            tools=[multiply_tool]
        )
        base_server.add_middleware(middleware)
        base_server.disable(tags={"math"})

        async with Client[FastMCPTransport](base_server) as client:
            tools: list[SDKTool] = await client.list_tools()
            tool_names: list[str] = [tool.name for tool in tools]
            assert "multiply" in tool_names

    async def test_empty_tool_injection(self, base_server: FastMCP):
        """Test that middleware with no tools doesn't affect behavior."""
        middleware: ToolInjectionMiddleware = ToolInjectionMiddleware(tools=[])
        base_server.add_middleware(middleware)

        async with Client[FastMCPTransport](base_server) as client:
            tools: list[SDKTool] = await client.list_tools()
            result: CallToolResult = await client.call_tool(
                name="add", arguments={"a": 3, "b": 4}
            )

        # Should only have the base tools
        assert len(tools) == 2
        tool_names: list[str] = [tool.name for tool in tools]
        assert "add" in tool_names
        assert "subtract" in tool_names
        assert result.structured_content is not None
        assert isinstance(result.structured_content, dict)
        assert result.structured_content["result"] == 7


TOKENS = {
    "read-token": {"client_id": "reader", "scopes": ["read"]},
    "admin-token": {"client_id": "administrator", "scopes": ["read", "admin"]},
}


def admin_report() -> str:
    """Return the admin report."""
    return "Quarterly totals"


def public_report() -> str:
    """Return the public report."""
    return "Public summary"


def make_auth_server(*middleware: Middleware) -> FastMCP:
    return FastMCP(
        "AuthServer",
        auth=StaticTokenVerifier(TOKENS),
        middleware=list(middleware),
    )


@asynccontextmanager
async def http_client(server: FastMCP, token: str) -> AsyncIterator[Client]:
    async with run_server_async(server) as url:
        async with Client(url, auth=token) as client:
            yield client


def scoped_admin_tool() -> Tool:
    return Tool.from_function(
        admin_report,
        name="admin_report",
        auth=require_scopes("admin"),
    )


def unscoped_tool(fn: Callable[[], str], name: str) -> Tool:
    return Tool.from_function(fn, name=name)


class TestInjectedToolAuthorization:
    """Injected tools follow the same authorization as registered tools."""

    async def test_tool_auth_hides_injected_tool_without_scope(self):
        server = make_auth_server(ToolInjectionMiddleware([scoped_admin_tool()]))

        async with http_client(server, "read-token") as client:
            tools = await client.list_tools()

        assert "admin_report" not in [tool.name for tool in tools]

    async def test_tool_auth_denies_injected_tool_call_without_scope(self):
        server = make_auth_server(ToolInjectionMiddleware([scoped_admin_tool()]))

        async with http_client(server, "read-token") as client:
            with pytest.raises(ToolError, match="Unknown tool"):
                await client.call_tool("admin_report", {})

    async def test_tool_auth_allows_injected_tool_with_scope(self):
        server = make_auth_server(ToolInjectionMiddleware([scoped_admin_tool()]))

        async with http_client(server, "admin-token") as client:
            tools = await client.list_tools()
            result = await client.call_tool("admin_report", {})

        assert "admin_report" in [tool.name for tool in tools]
        assert result.data == "Quarterly totals"

    async def test_tool_auth_on_injected_tool_matches_registered_tool(self):
        """An injected tool and a registered tool with the same auth agree."""
        registered = make_auth_server()
        registered.add_tool(scoped_admin_tool())
        injected = make_auth_server(ToolInjectionMiddleware([scoped_admin_tool()]))

        outcomes: list[tuple[list[str], bool]] = []
        for server in (registered, injected):
            async with http_client(server, "read-token") as client:
                names = [tool.name for tool in await client.list_tools()]
                result = await client.call_tool(
                    "admin_report", {}, raise_on_error=False
                )
            outcomes.append((names, result.is_error))

        assert outcomes[0] == outcomes[1] == ([], True)

    async def test_injected_tool_without_auth_is_available_to_any_token(self):
        server = make_auth_server(
            ToolInjectionMiddleware([unscoped_tool(public_report, "public_report")])
        )

        async with http_client(server, "read-token") as client:
            tools = await client.list_tools()
            result = await client.call_tool("public_report", {})

        assert [tool.name for tool in tools] == ["public_report"]
        assert result.data == "Public summary"

    @pytest.mark.parametrize("injection_first", [True, False])
    async def test_auth_middleware_denies_injected_tool(self, injection_first: bool):
        injection = ToolInjectionMiddleware(
            [unscoped_tool(admin_report, "admin_report")]
        )
        auth = AuthMiddleware(auth=require_scopes("admin"))
        order = [injection, auth] if injection_first else [auth, injection]
        server = make_auth_server(*order)

        async with http_client(server, "read-token") as client:
            tools = await client.list_tools()
            result = await client.call_tool("admin_report", {}, raise_on_error=False)

        assert tools == []
        assert result.is_error
        assert "Quarterly totals" not in str(result.content)

    @pytest.mark.parametrize("injection_first", [True, False])
    async def test_auth_middleware_allows_injected_tool(self, injection_first: bool):
        injection = ToolInjectionMiddleware(
            [unscoped_tool(admin_report, "admin_report")]
        )
        auth = AuthMiddleware(auth=require_scopes("admin"))
        order = [injection, auth] if injection_first else [auth, injection]
        server = make_auth_server(*order)

        async with http_client(server, "admin-token") as client:
            tools = await client.list_tools()
            result = await client.call_tool("admin_report", {})

        assert [tool.name for tool in tools] == ["admin_report"]
        assert result.data == "Quarterly totals"

    async def test_injected_tool_on_mounted_server_follows_tool_auth(self):
        child = FastMCP(
            "Child",
            middleware=[ToolInjectionMiddleware([scoped_admin_tool()])],
        )
        parent = make_auth_server()
        parent.mount(child, namespace="child")

        async with run_server_async(parent) as url:
            async with Client(url, auth="read-token") as client:
                read_tools = await client.list_tools()
                read_result = await client.call_tool(
                    "child_admin_report", {}, raise_on_error=False
                )
            async with Client(url, auth="admin-token") as client:
                admin_tools = await client.list_tools()
                admin_result = await client.call_tool("child_admin_report", {})

        assert read_tools == []
        assert read_result.is_error
        assert [tool.name for tool in admin_tools] == ["child_admin_report"]
        assert admin_result.data == "Quarterly totals"

    async def test_hashed_call_reaches_app_tool_sharing_injected_name(self):
        """An injected tool owns its display name, not an app tool's hashed name."""
        app = FastMCPApp("contacts")

        @app.tool()
        def save(name: str) -> str:
            return f"app saved {name}"

        def injected_save(name: str) -> str:
            return f"injected saved {name}"

        server = FastMCP("Platform")
        server.add_provider(app)
        server.add_middleware(
            ToolInjectionMiddleware([Tool.from_function(injected_save, name="save")])
        )

        async with Client(server) as client:
            by_name = await client.call_tool("save", {"name": "alice"})
            by_hash = await client.call_tool(
                hashed_backend_name("contacts", "save"), {"name": "alice"}
            )

        assert by_name.data == "injected saved alice"
        assert by_hash.data == "app saved alice"

    async def test_shadowed_registered_tool_is_not_listed(self):
        """Listing and calling agree when an injected tool shadows a registered one."""
        server = make_auth_server(
            ToolInjectionMiddleware(
                [
                    Tool.from_function(
                        admin_report, name="report", auth=require_scopes("admin")
                    )
                ]
            )
        )
        server.add_tool(unscoped_tool(public_report, "report"))

        async with run_server_async(server) as url:
            async with Client(url, auth="read-token") as client:
                read_tools = await client.list_tools()
                read_result = await client.call_tool("report", {}, raise_on_error=False)
            async with Client(url, auth="admin-token") as client:
                admin_tools = await client.list_tools()
                admin_result = await client.call_tool("report", {})

        assert read_tools == []
        assert read_result.is_error
        assert [tool.name for tool in admin_tools] == ["report"]
        assert admin_result.data == "Quarterly totals"

    async def test_injected_task_tool_runs_as_task(self):
        async def square(n: int) -> int:
            return n * n

        server = FastMCP(
            "TaskServer",
            middleware=[
                ToolInjectionMiddleware([Tool.from_function(square, task=True)])
            ],
        )

        async with Client(server) as client:
            task = await client.call_tool("square", {"n": 7}, task=True)
            assert not task.returned_immediately
            result = await task.result()

        assert result.data == 49

    async def test_background_task_runs_the_injected_tool_when_its_name_is_shared(
        self,
    ):
        async def injected_report(n: int) -> str:
            return f"injected {n}"

        server = FastMCP(
            "TaskServer",
            middleware=[
                ToolInjectionMiddleware(
                    [Tool.from_function(injected_report, name="report", task=True)]
                )
            ],
        )

        @server.tool(task=True)
        async def report(n: int) -> str:
            return f"registered {n}"

        async with Client(server) as client:
            task = await client.call_tool("report", {"n": 3}, task=True)
            assert not task.returned_immediately
            result = await task.result()

        assert result.data == "injected 3"

    async def test_background_task_runs_the_injected_tool_when_signatures_differ(
        self,
    ):
        async def injected_report(n: int) -> str:
            return f"injected {n}"

        server = FastMCP(
            "TaskServer",
            middleware=[
                ToolInjectionMiddleware(
                    [Tool.from_function(injected_report, name="report", task=True)]
                )
            ],
        )

        @server.tool(task=True)
        async def report(label: str, flag: bool) -> str:
            return f"registered {label} {flag}"

        async with Client(server) as client:
            task = await client.call_tool("report", {"n": 5}, task=True)
            assert not task.returned_immediately
            result = await task.result()

        assert result.data == "injected 5"

    async def test_background_task_runs_the_injected_tool_over_a_mounted_tool(self):
        async def injected_report(n: int) -> str:
            return f"injected {n}"

        child = FastMCP("Child")

        @child.tool(task=True)
        async def report(n: int) -> str:
            return f"mounted {n}"

        @child.tool(task=True)
        async def summary(n: int) -> str:
            return f"summary {n}"

        parent = FastMCP(
            "Parent",
            middleware=[
                ToolInjectionMiddleware(
                    [
                        Tool.from_function(
                            injected_report, name="child_report", task=True
                        )
                    ]
                )
            ],
        )
        parent.mount(child, namespace="child")

        async with Client(parent) as client:
            shadowed = await client.call_tool("child_report", {"n": 4}, task=True)
            unshadowed = await client.call_tool("child_summary", {"n": 4}, task=True)
            shadowed_result = await shadowed.result()
            unshadowed_result = await unshadowed.result()

        assert shadowed_result.data == "injected 4"
        assert unshadowed_result.data == "summary 4"

    async def test_task_discovery_lists_one_tool_for_a_shared_name(self):
        async def injected_report(n: int) -> str:
            return f"injected {n}"

        injected = Tool.from_function(injected_report, name="report", task=True)
        server = FastMCP("TaskServer", middleware=[ToolInjectionMiddleware([injected])])

        @server.tool(task=True)
        async def report(n: int) -> str:
            return f"registered {n}"

        @server.tool(task=True)
        async def other(n: int) -> str:
            return f"other {n}"

        components = await server.get_tasks()

        assert [c for c in components if c.name == "report"] == [injected]
        assert {c.name for c in components} == {"report", "other"}

    async def test_task_discovery_skips_a_mounted_tool_shadowed_by_an_injected_tool(
        self,
    ):
        async def injected_report(n: int) -> str:
            return f"injected {n}"

        child = FastMCP("Child")

        @child.tool(task=True)
        async def report(n: int) -> str:
            return f"mounted {n}"

        @child.tool(task=True)
        async def summary(n: int) -> str:
            return f"summary {n}"

        injected = Tool.from_function(injected_report, name="child_report", task=True)
        parent = FastMCP("Parent", middleware=[ToolInjectionMiddleware([injected])])
        parent.mount(child, namespace="child")

        components = await parent.get_tasks()

        assert [c for c in components if c.name == "child_report"] == [injected]
        assert "child_summary" in {c.name for c in components}

    async def test_task_discovery_skips_a_registered_tool_shadowed_by_a_synchronous_injected_tool(
        self,
    ):
        async def injected_report(n: int) -> str:
            return f"injected {n}"

        injected = Tool.from_function(injected_report, name="report", task=False)
        server = FastMCP("TaskServer", middleware=[ToolInjectionMiddleware([injected])])

        @server.tool(task=True)
        async def report(n: int) -> str:
            return f"registered {n}"

        components = await server.get_tasks()

        assert [c for c in components if c.name == "report"] == []

    async def test_background_tasks_run_for_tools_with_distinct_names(self):
        async def injected_only(n: int) -> str:
            return f"injected {n}"

        server = FastMCP(
            "TaskServer",
            middleware=[
                ToolInjectionMiddleware(
                    [Tool.from_function(injected_only, name="injected_only", task=True)]
                )
            ],
        )

        @server.tool(task=True)
        async def registered_only(n: int) -> str:
            return f"registered {n}"

        async with Client(server) as client:
            first = await client.call_tool("injected_only", {"n": 1}, task=True)
            second = await client.call_tool("registered_only", {"n": 2}, task=True)
            first_result = await first.result()
            second_result = await second.result()

        assert first_result.data == "injected 1"
        assert second_result.data == "registered 2"