From 0702bad1c2aa9509b89a8c01c985b63db08c64d8 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 19:47:19 -0400 Subject: [PATCH 01/17] Add single-component methods --- src/fastmcp/server/context.py | 11 +- src/fastmcp/server/middleware/__init__.py | 0 src/fastmcp/server/middleware/middleware.py | 174 +++++++++++++ src/fastmcp/server/server.py | 255 +++++++++++++++++--- tests/server/middleware/__init__.py | 0 tests/server/middleware/test_middleware.py | 196 +++++++++++++++ 6 files changed, 601 insertions(+), 35 deletions(-) create mode 100644 src/fastmcp/server/middleware/__init__.py create mode 100644 src/fastmcp/server/middleware/middleware.py create mode 100644 tests/server/middleware/__init__.py create mode 100644 tests/server/middleware/test_middleware.py diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 145b127fe..081e08d49 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from mcp import LoggingLevel from mcp.server.lowlevel.helper_types import ReadResourceContents +from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext from mcp.types import ( CreateMessageResult, @@ -95,8 +96,14 @@ class Context: @property def request_context(self) -> RequestContext: - """Access to the underlying request context.""" - return self.fastmcp._mcp_server.request_context + """Access to the underlying request context. + + If called outside of a request context, this will raise a ValueError. + """ + try: + return request_ctx.get() + except LookupError: + raise ValueError("Context is not available outside of a request") async def report_progress( self, progress: float, total: float | None = None, message: str | None = None diff --git a/src/fastmcp/server/middleware/__init__.py b/src/fastmcp/server/middleware/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/fastmcp/server/middleware/middleware.py b/src/fastmcp/server/middleware/middleware.py new file mode 100644 index 000000000..a1030ebaa --- /dev/null +++ b/src/fastmcp/server/middleware/middleware.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import logging +from collections.abc import Awaitable +from dataclasses import dataclass, field, replace +from datetime import datetime, timezone +from functools import partial +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Literal, + Protocol, + TypeVar, + runtime_checkable, +) + +import mcp.types as mt + +if TYPE_CHECKING: + from fastmcp.server.context import Context + +logger = logging.getLogger(__name__) + + +T = TypeVar("T") +R = TypeVar("R", covariant=True) + + +@runtime_checkable +class CallNext(Protocol[T, R]): + def __call__(self, context: MiddlewareContext[T]) -> Awaitable[R]: ... + + +ServerResultT = TypeVar( + "ServerResultT", + bound=mt.EmptyResult + | mt.InitializeResult + | mt.CompleteResult + | mt.GetPromptResult + | mt.ListPromptsResult + | mt.ListResourcesResult + | mt.ListResourceTemplatesResult + | mt.ReadResourceResult + | mt.CallToolResult + | mt.ListToolsResult, +) + + +@runtime_checkable +class ServerResultProtocol(Protocol[ServerResultT]): + root: ServerResultT + + +@dataclass(kw_only=True, frozen=True) +class MiddlewareContext(Generic[T]): + """ + Unified context for all middleware operations. + """ + + message: T + + fastmcp_context: Context | None = None + + # Common metadata + source: Literal["client", "server"] = "client" + type: Literal["request", "notification"] = "request" + method: str | None = None + timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + def copy(self, **kwargs: Any) -> MiddlewareContext[T]: + return replace(self, **kwargs) + + +def make_middleware_wrapper( + middleware: MCPMiddleware, call_next: CallNext[T, R] +) -> CallNext[T, R]: + """Create a wrapper that applies a single middleware to a context. The + closure bakes in the middleware and call_next function, so it can be + passed to other functions that expect a call_next function.""" + + async def wrapper(context: MiddlewareContext[T]) -> R: + return await middleware(context, call_next) + + return wrapper + + +class MCPMiddleware: + """Base class for FastMCP middleware with dispatching hooks.""" + + async def __call__( + self, + context: MiddlewareContext[T], + call_next: CallNext[T, Any], + ) -> Any: + """Main entry point that orchestrates the pipeline.""" + handler_chain = await self._dispatch_handler( + context.message, + call_next=call_next, + ) + return await handler_chain(context) + + async def _dispatch_handler( + self, message: Any, call_next: CallNext[Any, Any] + ) -> CallNext[Any, Any]: + """Builds a chain of handlers for a given message.""" + handler = call_next + + match message: + case mt.CallToolRequest(): + handler = partial(self.on_call_tool, call_next=handler) + case mt.ReadResourceRequest(): + handler = partial(self.on_read_resource, call_next=handler) + case mt.GetPromptRequest(): + handler = partial(self.on_get_prompt, call_next=handler) + + match message: + case mt.Request(): + handler = partial(self.on_request, call_next=handler) + case mt.Notification(): + handler = partial(self.on_notification, call_next=handler) + + handler = partial(self.on_message, call_next=handler) + + return handler + + async def on_message( + self, + context: MiddlewareContext[Any], + call_next: CallNext[Any, Any], + ) -> Any: + return await call_next(context) + + async def on_request( + self, + context: MiddlewareContext[mt.Request], + call_next: CallNext[mt.Request, Any], + ) -> Any: + return await call_next(context) + + async def on_notification( + self, + context: MiddlewareContext[mt.Notification], + call_next: CallNext[mt.Notification, Any], + ) -> Any: + return await call_next(context) + + async def on_call_tool( + self, + context: MiddlewareContext[mt.CallToolRequest], + call_next: CallNext[mt.CallToolRequest, mt.CallToolResult], + ) -> mt.CallToolResult: + return await call_next(context) + + async def on_read_resource( + self, + context: MiddlewareContext[mt.ReadResourceRequest], + call_next: CallNext[mt.ReadResourceRequest, mt.ReadResourceResult], + ) -> mt.ReadResourceResult: + return await call_next(context) + + async def on_get_prompt( + self, + context: MiddlewareContext[mt.GetPromptRequest], + call_next: CallNext[mt.GetPromptRequest, mt.GetPromptResult], + ) -> mt.GetPromptResult: + return await call_next(context) + + async def on_list_tools( + self, + context: MiddlewareContext[mt.ListToolsRequest], + call_next: CallNext[mt.ListToolsRequest, mt.ListToolsResult], + ) -> mt.ListToolsResult: + return await call_next(context) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 47bd001eb..6d05db36a 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload import anyio import httpx +import mcp.types import uvicorn from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.server.lowlevel.server import LifespanResultT, NotificationOptions @@ -53,6 +54,7 @@ from fastmcp.server.http import ( create_sse_app, create_streamable_http_app, ) +from fastmcp.server.middleware.middleware import MCPMiddleware, MiddlewareContext from fastmcp.settings import Settings from fastmcp.tools import ToolManager from fastmcp.tools.tool import FunctionTool, Tool @@ -115,6 +117,7 @@ class FastMCP(Generic[LifespanResultT]): *, version: str | None = None, auth: OAuthProvider | None = None, + middleware: list[MCPMiddleware] | None = None, lifespan: ( Callable[ [FastMCP[LifespanResultT]], @@ -196,6 +199,8 @@ class FastMCP(Generic[LifespanResultT]): self.include_tags = include_tags self.exclude_tags = exclude_tags + self.middleware = middleware or [] + # Set up MCP protocol handlers self._setup_handlers() self.dependencies = dependencies or fastmcp.settings.server_dependencies @@ -319,6 +324,20 @@ class FastMCP(Generic[LifespanResultT]): self._mcp_server.read_resource()(self._mcp_read_resource) self._mcp_server.get_prompt()(self._mcp_get_prompt) + async def _apply_middleware( + self, + context: MiddlewareContext[Any], + call_next: Callable[[MiddlewareContext[Any]], Awaitable[Any]], + ) -> Any: + """Builds and executes the middleware chain.""" + chain = call_next + for mw in reversed(self.middleware): + chain = partial(mw, call_next=chain) + return await chain(context) + + def add_middleware(self, middleware: MCPMiddleware) -> None: + self.middleware.append(middleware) + async def get_tools(self) -> dict[str, Tool]: """Get all registered tools, indexed by registered key.""" if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND: @@ -508,14 +527,33 @@ class FastMCP(Generic[LifespanResultT]): server. """ - tools = await self.get_tools() + logger.debug("List tools") - mcp_tools: list[MCPTool] = [] - for key, tool in tools.items(): - if self._should_enable_component(tool): - mcp_tools.append(tool.to_mcp_tool(name=key)) + async def _final_handler( + context: MiddlewareContext[dict[str, Any]], + ) -> list[MCPTool]: + # Call the business logic method + tools = await self.get_tools() - return mcp_tools + mcp_tools: list[MCPTool] = [] + for key, tool in tools.items(): + if self._should_enable_component(tool): + mcp_tools.append(tool.to_mcp_tool(name=key)) + + return mcp_tools + + with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: + # Create the middleware context. + mw_context = MiddlewareContext( + message={}, # List tools doesn't have parameters + source="client", + type="request", + method="tools/list", + fastmcp_context=fastmcp_ctx, + ) + + # Apply the middleware chain. + return await self._apply_middleware(mw_context, _final_handler) async def _mcp_list_resources(self) -> list[MCPResource]: """ @@ -523,12 +561,31 @@ class FastMCP(Generic[LifespanResultT]): server. """ - resources = await self.get_resources() - mcp_resources: list[MCPResource] = [] - for key, resource in resources.items(): - if self._should_enable_component(resource): - mcp_resources.append(resource.to_mcp_resource(uri=key)) - return mcp_resources + logger.debug("List resources") + + async def _final_handler( + context: MiddlewareContext[dict[str, Any]], + ) -> list[MCPResource]: + # Call the business logic method + resources = await self.get_resources() + mcp_resources: list[MCPResource] = [] + for key, resource in resources.items(): + if self._should_enable_component(resource): + mcp_resources.append(resource.to_mcp_resource(uri=key)) + return mcp_resources + + with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: + # Create the middleware context. + mw_context = MiddlewareContext( + message={}, # List resources doesn't have parameters + source="client", + type="request", + method="resources/list", + fastmcp_context=fastmcp_ctx, + ) + + # Apply the middleware chain. + return await self._apply_middleware(mw_context, _final_handler) async def _mcp_list_resource_templates(self) -> list[MCPResourceTemplate]: """ @@ -536,12 +593,31 @@ class FastMCP(Generic[LifespanResultT]): MCP server. """ - templates = await self.get_resource_templates() - mcp_templates: list[MCPResourceTemplate] = [] - for key, template in templates.items(): - if self._should_enable_component(template): - mcp_templates.append(template.to_mcp_template(uriTemplate=key)) - return mcp_templates + logger.debug("List resource templates") + + async def _final_handler( + context: MiddlewareContext[dict[str, Any]], + ) -> list[MCPResourceTemplate]: + # Call the business logic method + templates = await self.get_resource_templates() + mcp_templates: list[MCPResourceTemplate] = [] + for key, template in templates.items(): + if self._should_enable_component(template): + mcp_templates.append(template.to_mcp_template(uriTemplate=key)) + return mcp_templates + + with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: + # Create the middleware context. + mw_context = MiddlewareContext( + message={}, # List resource templates doesn't have parameters + source="client", + type="request", + method="resources/list_templates", + fastmcp_context=fastmcp_ctx, + ) + + # Apply the middleware chain. + return await self._apply_middleware(mw_context, _final_handler) async def _mcp_list_prompts(self) -> list[MCPPrompt]: """ @@ -549,12 +625,31 @@ class FastMCP(Generic[LifespanResultT]): server. """ - prompts = await self.get_prompts() - mcp_prompts: list[MCPPrompt] = [] - for key, prompt in prompts.items(): - if self._should_enable_component(prompt): - mcp_prompts.append(prompt.to_mcp_prompt(name=key)) - return mcp_prompts + logger.debug("List prompts") + + async def _final_handler( + context: MiddlewareContext[dict[str, Any]], + ) -> list[MCPPrompt]: + # Call the business logic method + prompts = await self.get_prompts() + mcp_prompts: list[MCPPrompt] = [] + for key, prompt in prompts.items(): + if self._should_enable_component(prompt): + mcp_prompts.append(prompt.to_mcp_prompt(name=key)) + return mcp_prompts + + with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: + # Create the middleware context. + mw_context = MiddlewareContext( + message={}, # List prompts doesn't have parameters + source="client", + type="request", + method="prompts/list", + fastmcp_context=fastmcp_ctx, + ) + + # Apply the middleware chain. + return await self._apply_middleware(mw_context, _final_handler) async def _mcp_call_tool( self, key: str, arguments: dict[str, Any] @@ -573,17 +668,43 @@ class FastMCP(Generic[LifespanResultT]): """ logger.debug("Call tool: %s with %s", key, arguments) - # Create and use context for the entire call with fastmcp.server.context.Context(fastmcp=self): try: - return await self._call_tool(key, arguments) + return await self._middleware_call_tool(key, arguments) except DisabledError: - # convert to NotFoundError to avoid leaking tool presence raise NotFoundError(f"Unknown tool: {key}") except NotFoundError: - # standardize NotFound message raise NotFoundError(f"Unknown tool: {key}") + async def _middleware_call_tool( + self, + key: str, + arguments: dict[str, Any], + ) -> list[MCPContent]: + """ + Call a tool with middleware. + """ + + async def _handler( + context: MiddlewareContext[mcp.types.CallToolRequest], + ) -> list[MCPContent]: + return await self._call_tool( + key=context.message.params.name, + arguments=context.message.params.arguments or {}, + ) + + mw_context = MiddlewareContext( + message=mcp.types.CallToolRequest( + method="tools/call", + params=mcp.types.CallToolRequestParams(name=key, arguments=arguments), + ), + source="client", + type="request", + method="tools/call", + fastmcp_context=fastmcp.server.dependencies.get_context(), + ) + return await self._apply_middleware(mw_context, _handler) + async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: """ Call a tool with raw MCP arguments. FastMCP subclasses should override @@ -615,7 +736,9 @@ class FastMCP(Generic[LifespanResultT]): tool_key = tool_key.removeprefix(f"{mounted_server.prefix}_") else: continue - return await mounted_server.server._call_tool(tool_key, arguments) + return await mounted_server.server._middleware_call_tool( + tool_key, arguments + ) except NotFoundError: # Tool not found on this server, try the next one continue @@ -632,7 +755,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self): try: - return await self._read_resource(uri) + return await self._middleware_read_resource(uri) except DisabledError: # convert to NotFoundError to avoid leaking resource presence raise NotFoundError(f"Unknown resource: {str(uri)!r}") @@ -640,6 +763,41 @@ class FastMCP(Generic[LifespanResultT]): # standardize NotFound message raise NotFoundError(f"Unknown resource: {str(uri)!r}") + async def _middleware_read_resource( + self, + uri: AnyUrl | str, + ) -> list[ReadResourceContents]: + """ + Read a resource with middleware. + """ + + async def _handler( + context: MiddlewareContext[mcp.types.ReadResourceRequest], + ) -> list[ReadResourceContents]: + return await self._read_resource( + uri=context.message.params.uri, + ) + + # Convert string URI to AnyUrl if needed + if isinstance(uri, str): + from pydantic import AnyUrl + + uri_param = AnyUrl(uri) + else: + uri_param = uri + + mw_context = MiddlewareContext( + message=mcp.types.ReadResourceRequest( + method="resources/read", + params=mcp.types.ReadResourceRequestParams(uri=uri_param), + ), + source="client", + type="request", + method="resources/read", + fastmcp_context=fastmcp.server.dependencies.get_context(), + ) + return await self._apply_middleware(mw_context, _handler) + async def _read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]: """ Read a resource by URI, in the format expected by the low-level MCP @@ -675,7 +833,9 @@ class FastMCP(Generic[LifespanResultT]): ) else: continue - return await mounted_server.server._mcp_read_resource(resource_uri) + return await mounted_server.server._middleware_read_resource( + resource_uri + ) except NotFoundError: # Resource not found on this server, try the next one continue @@ -694,7 +854,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self): try: - return await self._get_prompt(name, arguments) + return await self._middleware_get_prompt(name, arguments) except DisabledError: # convert to NotFoundError to avoid leaking prompt presence raise NotFoundError(f"Unknown prompt: {name}") @@ -702,6 +862,35 @@ class FastMCP(Generic[LifespanResultT]): # standardize NotFound message raise NotFoundError(f"Unknown prompt: {name}") + async def _middleware_get_prompt( + self, + name: str, + arguments: dict[str, Any] | None = None, + ) -> GetPromptResult: + """ + Get a prompt with middleware. + """ + + async def _handler( + context: MiddlewareContext[mcp.types.GetPromptRequest], + ) -> GetPromptResult: + return await self._get_prompt( + name=context.message.params.name, + arguments=context.message.params.arguments, + ) + + mw_context = MiddlewareContext( + message=mcp.types.GetPromptRequest( + method="prompts/get", + params=mcp.types.GetPromptRequestParams(name=name, arguments=arguments), + ), + source="client", + type="request", + method="prompts/get", + fastmcp_context=fastmcp.server.dependencies.get_context(), + ) + return await self._apply_middleware(mw_context, _handler) + async def _get_prompt( self, name: str, arguments: dict[str, Any] | None = None ) -> GetPromptResult: @@ -736,7 +925,7 @@ class FastMCP(Generic[LifespanResultT]): ) else: continue - return await mounted_server.server._mcp_get_prompt( + return await mounted_server.server._middleware_get_prompt( prompt_name, arguments ) except NotFoundError: diff --git a/tests/server/middleware/__init__.py b/tests/server/middleware/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py new file mode 100644 index 000000000..c21591788 --- /dev/null +++ b/tests/server/middleware/test_middleware.py @@ -0,0 +1,196 @@ +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.middleware import MCPMiddleware, MiddlewareContext + + +@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(MCPMiddleware): + """A middleware that automatically records all method calls.""" + + def __init__(self): + super().__init__() + self.calls: list[Recording] = [] + + 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 = 1 + ) -> bool: + """Assert that a hook was called a specific number of times.""" + calls = self.get_calls(hook=hook, method=method) + actual_times = len(calls) + assert actual_times == times, ( + f"Expected {hook!r} to be called {times} times" + f"{f' for method {method!r}' if method else ''}, " + 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() + yield middleware + + +@pytest.fixture +def mcp_server(recording_middleware): + mcp = FastMCP() + + @mcp.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(times=3) + assert recording_middleware.assert_called(method="tools/call", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_call_tool", times=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(times=3) + assert recording_middleware.assert_called(method="resources/read", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_read_resource", times=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(times=3) + assert recording_middleware.assert_called(method="prompts/get", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_get_prompt", times=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(times=3) + assert recording_middleware.assert_called(method="tools/list", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_list_tools", times=1) From e11417c3c2174096b7c9158a71633c2637e345d6 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 21:10:25 -0400 Subject: [PATCH 02/17] Merge main --- .../server/{middleware => }/middleware.py | 65 ++++++--- src/fastmcp/server/middleware/__init__.py | 0 src/fastmcp/server/server.py | 125 +++++++++--------- 3 files changed, 114 insertions(+), 76 deletions(-) rename src/fastmcp/server/{middleware => }/middleware.py (73%) delete mode 100644 src/fastmcp/server/middleware/__init__.py diff --git a/src/fastmcp/server/middleware/middleware.py b/src/fastmcp/server/middleware.py similarity index 73% rename from src/fastmcp/server/middleware/middleware.py rename to src/fastmcp/server/middleware.py index a1030ebaa..3b0ba369e 100644 --- a/src/fastmcp/server/middleware/middleware.py +++ b/src/fastmcp/server/middleware.py @@ -17,6 +17,11 @@ from typing import ( import mcp.types as mt +from fastmcp.prompts.prompt import Prompt +from fastmcp.resources.resource import Resource +from fastmcp.resources.template import ResourceTemplate +from fastmcp.tools.tool import Tool + if TYPE_CHECKING: from fastmcp.server.context import Context @@ -47,6 +52,32 @@ ServerResultT = TypeVar( ) +@dataclass(kw_only=True) +class CallToolResult: + content: list[mt.Content] + isError: bool = False + + +@dataclass(kw_only=True) +class ListToolsResult: + tools: dict[str, Tool] + + +@dataclass(kw_only=True) +class ListResourcesResult: + resources: list[Resource] + + +@dataclass(kw_only=True) +class ListResourceTemplatesResult: + resource_templates: list[ResourceTemplate] + + +@dataclass(kw_only=True) +class ListPromptsResult: + prompts: list[Prompt] + + @runtime_checkable class ServerResultProtocol(Protocol[ServerResultT]): root: ServerResultT @@ -95,29 +126,29 @@ class MCPMiddleware: ) -> Any: """Main entry point that orchestrates the pipeline.""" handler_chain = await self._dispatch_handler( - context.message, + context, call_next=call_next, ) return await handler_chain(context) async def _dispatch_handler( - self, message: Any, call_next: CallNext[Any, Any] + self, context: MiddlewareContext[Any], call_next: CallNext[Any, Any] ) -> CallNext[Any, Any]: """Builds a chain of handlers for a given message.""" handler = call_next - match message: - case mt.CallToolRequest(): + match context.method: + case "tools/call": handler = partial(self.on_call_tool, call_next=handler) - case mt.ReadResourceRequest(): + case "resources/read": handler = partial(self.on_read_resource, call_next=handler) - case mt.GetPromptRequest(): + case "prompts/get": handler = partial(self.on_get_prompt, call_next=handler) - match message: - case mt.Request(): + match context.type: + case "request": handler = partial(self.on_request, call_next=handler) - case mt.Notification(): + case "notification": handler = partial(self.on_notification, call_next=handler) handler = partial(self.on_message, call_next=handler) @@ -147,28 +178,28 @@ class MCPMiddleware: async def on_call_tool( self, - context: MiddlewareContext[mt.CallToolRequest], - call_next: CallNext[mt.CallToolRequest, mt.CallToolResult], + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult], ) -> mt.CallToolResult: return await call_next(context) async def on_read_resource( self, - context: MiddlewareContext[mt.ReadResourceRequest], - call_next: CallNext[mt.ReadResourceRequest, mt.ReadResourceResult], + context: MiddlewareContext[mt.ReadResourceRequestParams], + call_next: CallNext[mt.ReadResourceRequestParams, mt.ReadResourceResult], ) -> mt.ReadResourceResult: return await call_next(context) async def on_get_prompt( self, - context: MiddlewareContext[mt.GetPromptRequest], - call_next: CallNext[mt.GetPromptRequest, mt.GetPromptResult], + context: MiddlewareContext[mt.GetPromptRequestParams], + call_next: CallNext[mt.GetPromptRequestParams, mt.GetPromptResult], ) -> mt.GetPromptResult: return await call_next(context) async def on_list_tools( self, context: MiddlewareContext[mt.ListToolsRequest], - call_next: CallNext[mt.ListToolsRequest, mt.ListToolsResult], - ) -> mt.ListToolsResult: + call_next: CallNext[mt.ListToolsRequest, ListToolsResult], + ) -> ListToolsResult: return await call_next(context) diff --git a/src/fastmcp/server/middleware/__init__.py b/src/fastmcp/server/middleware/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 2fdb28e23..c8d93f29f 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -42,6 +42,7 @@ from starlette.routing import BaseRoute, Route import fastmcp import fastmcp.server +import fastmcp.server.middleware from fastmcp.exceptions import DisabledError, NotFoundError from fastmcp.prompts import Prompt, PromptManager from fastmcp.prompts.prompt import FunctionPrompt @@ -54,7 +55,7 @@ from fastmcp.server.http import ( create_sse_app, create_streamable_http_app, ) -from fastmcp.server.middleware.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext from fastmcp.settings import Settings from fastmcp.tools import ToolManager from fastmcp.tools.tool import FunctionTool, Tool @@ -340,28 +341,7 @@ class FastMCP(Generic[LifespanResultT]): async def get_tools(self) -> dict[str, Tool]: """Get all registered tools, indexed by registered key.""" - if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND: - tools: dict[str, Tool] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - server_tools = await mounted_server.server.get_tools() - # Apply prefix to each tool key if prefix exists and is not empty - if mounted_server.prefix: - for tool in server_tools.values(): - tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}") - tools[tool.key] = tool - else: - tools.update(server_tools) - except Exception as e: - logger.warning( - f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}" - ) - continue - tools.update(self._tool_manager.get_tools()) - self._cache.set("tools", tools) - return tools + return await self._list_tools(apply_middleware=False) async def get_tool(self, key: str) -> Tool: tools = await self.get_tools() @@ -529,18 +509,23 @@ class FastMCP(Generic[LifespanResultT]): return decorator async def _mcp_list_tools(self) -> list[MCPTool]: + logger.debug("Handler called: list_tools") + + with fastmcp.server.context.Context(fastmcp=self): + tools = await self._middleware_list_tools() + return [tool.to_mcp_tool(name=tool.name) for tool in tools] + + async def _middleware_list_tools(self) -> dict[str, Tool]: """ List all available tools, in the format expected by the low-level MCP server. """ - logger.debug("List tools") - async def _final_handler( - context: MiddlewareContext[dict[str, Any]], + async def _handler( + context: MiddlewareContext[mcp.types.ListToolsRequest], ) -> list[MCPTool]: - # Call the business logic method - tools = await self.get_tools() + tools = await self._list_tools() mcp_tools: list[MCPTool] = [] for key, tool in tools.items(): @@ -552,7 +537,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: # Create the middleware context. mw_context = MiddlewareContext( - message={}, # List tools doesn't have parameters + message=mcp.types.ListToolsRequest(method="tools/list"), source="client", type="request", method="tools/list", @@ -560,7 +545,41 @@ class FastMCP(Generic[LifespanResultT]): ) # Apply the middleware chain. - return await self._apply_middleware(mw_context, _final_handler) + return await self._apply_middleware(mw_context, _handler) + + async def _list_tools(self, apply_middleware: bool = True) -> dict[str, Tool]: + """ + List all available tools. + """ + + if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND: + tools: dict[str, Tool] = {} + + # iterate such that new mounts overwrite older ones + for mounted_server in self._mounted_servers: + try: + if apply_middleware: + server_tools = ( + await mounted_server.server._middleware_list_tools() + ) + else: + server_tools = await mounted_server.server._list_tools() + # Apply prefix to each tool key if prefix exists and is not empty + if mounted_server.prefix: + for tool in server_tools.values(): + tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}") + tools[tool.key] = tool + else: + tools.update(server_tools) + tools.update(server_tools) + except Exception as e: + logger.warning( + f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}" + ) + continue + tools.update(self._tool_manager.get_tools()) + self._cache.set("tools", tools) + return tools async def _mcp_list_resources(self) -> list[MCPResource]: """ @@ -568,12 +587,11 @@ class FastMCP(Generic[LifespanResultT]): server. """ - logger.debug("List resources") + logger.debug("Handler called: list_resources") async def _final_handler( context: MiddlewareContext[dict[str, Any]], ) -> list[MCPResource]: - # Call the business logic method resources = await self.get_resources() mcp_resources: list[MCPResource] = [] for key, resource in resources.items(): @@ -600,12 +618,11 @@ class FastMCP(Generic[LifespanResultT]): MCP server. """ - logger.debug("List resource templates") + logger.debug("Handler called: list_resource_templates") async def _final_handler( context: MiddlewareContext[dict[str, Any]], ) -> list[MCPResourceTemplate]: - # Call the business logic method templates = await self.get_resource_templates() mcp_templates: list[MCPResourceTemplate] = [] for key, template in templates.items(): @@ -632,12 +649,11 @@ class FastMCP(Generic[LifespanResultT]): server. """ - logger.debug("List prompts") + logger.debug("Handler called: list_prompts") async def _final_handler( context: MiddlewareContext[dict[str, Any]], ) -> list[MCPPrompt]: - # Call the business logic method prompts = await self.get_prompts() mcp_prompts: list[MCPPrompt] = [] for key, prompt in prompts.items(): @@ -673,7 +689,7 @@ class FastMCP(Generic[LifespanResultT]): Returns: List of MCP Content objects containing the tool results """ - logger.debug("Call tool: %s with %s", key, arguments) + logger.debug("Handler called: call_tool %s with %s", key, arguments) with fastmcp.server.context.Context(fastmcp=self): try: @@ -693,18 +709,15 @@ class FastMCP(Generic[LifespanResultT]): """ async def _handler( - context: MiddlewareContext[mcp.types.CallToolRequest], + context: MiddlewareContext[mcp.types.CallToolRequestParams], ) -> list[MCPContent]: return await self._call_tool( - key=context.message.params.name, - arguments=context.message.params.arguments or {}, + key=context.message.name, + arguments=context.message.arguments or {}, ) mw_context = MiddlewareContext( - message=mcp.types.CallToolRequest( - method="tools/call", - params=mcp.types.CallToolRequestParams(name=key, arguments=arguments), - ), + message=mcp.types.CallToolRequestParams(name=key, arguments=arguments), source="client", type="request", method="tools/call", @@ -758,7 +771,7 @@ class FastMCP(Generic[LifespanResultT]): Delegates to _read_resource, which should be overridden by FastMCP subclasses. """ - logger.debug("Read resource: %s", uri) + logger.debug("Handler called: read_resource %s", uri) with fastmcp.server.context.Context(fastmcp=self): try: @@ -779,10 +792,10 @@ class FastMCP(Generic[LifespanResultT]): """ async def _handler( - context: MiddlewareContext[mcp.types.ReadResourceRequest], + context: MiddlewareContext[mcp.types.ReadResourceRequestParams], ) -> list[ReadResourceContents]: return await self._read_resource( - uri=context.message.params.uri, + uri=context.message.uri, ) # Convert string URI to AnyUrl if needed @@ -794,10 +807,7 @@ class FastMCP(Generic[LifespanResultT]): uri_param = uri mw_context = MiddlewareContext( - message=mcp.types.ReadResourceRequest( - method="resources/read", - params=mcp.types.ReadResourceRequestParams(uri=uri_param), - ), + message=mcp.types.ReadResourceRequestParams(uri=uri_param), source="client", type="request", method="resources/read", @@ -857,7 +867,7 @@ class FastMCP(Generic[LifespanResultT]): Delegates to _get_prompt, which should be overridden by FastMCP subclasses. """ - logger.debug("Get prompt: %s with %s", name, arguments) + logger.debug("Handler called: get_prompt %s with %s", name, arguments) with fastmcp.server.context.Context(fastmcp=self): try: @@ -879,18 +889,15 @@ class FastMCP(Generic[LifespanResultT]): """ async def _handler( - context: MiddlewareContext[mcp.types.GetPromptRequest], + context: MiddlewareContext[mcp.types.GetPromptRequestParams], ) -> GetPromptResult: return await self._get_prompt( - name=context.message.params.name, - arguments=context.message.params.arguments, + name=context.message.name, + arguments=context.message.arguments, ) mw_context = MiddlewareContext( - message=mcp.types.GetPromptRequest( - method="prompts/get", - params=mcp.types.GetPromptRequestParams(name=name, arguments=arguments), - ), + message=mcp.types.GetPromptRequestParams(name=name, arguments=arguments), source="client", type="request", method="prompts/get", From c183e3a99cc2977dabaa56d03d5ba1cc3d87ffbd Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 21:17:38 -0400 Subject: [PATCH 03/17] Updated list_tools --- src/fastmcp/server/middleware.py | 31 ++++++++++++++++ src/fastmcp/server/server.py | 26 +++++++------- tests/server/middleware/test_middleware.py | 42 +++++++++++++++++++++- 3 files changed, 85 insertions(+), 14 deletions(-) diff --git a/src/fastmcp/server/middleware.py b/src/fastmcp/server/middleware.py index 3b0ba369e..3f73bfe07 100644 --- a/src/fastmcp/server/middleware.py +++ b/src/fastmcp/server/middleware.py @@ -144,6 +144,14 @@ class MCPMiddleware: handler = partial(self.on_read_resource, call_next=handler) case "prompts/get": handler = partial(self.on_get_prompt, call_next=handler) + case "tools/list": + handler = partial(self.on_list_tools, call_next=handler) + case "resources/list": + handler = partial(self.on_list_resources, call_next=handler) + case "resource-templates/list": + handler = partial(self.on_list_resource_templates, call_next=handler) + case "prompts/list": + handler = partial(self.on_list_prompts, call_next=handler) match context.type: case "request": @@ -203,3 +211,26 @@ class MCPMiddleware: call_next: CallNext[mt.ListToolsRequest, ListToolsResult], ) -> ListToolsResult: return await call_next(context) + + async def on_list_resources( + self, + context: MiddlewareContext[mt.ListResourcesRequest], + call_next: CallNext[mt.ListResourcesRequest, ListResourcesResult], + ) -> ListResourcesResult: + return await call_next(context) + + async def on_list_resource_templates( + self, + context: MiddlewareContext[mt.ListResourceTemplatesRequest], + call_next: CallNext[ + mt.ListResourceTemplatesRequest, ListResourceTemplatesResult + ], + ) -> ListResourceTemplatesResult: + return await call_next(context) + + async def on_list_prompts( + self, + context: MiddlewareContext[mt.ListPromptsRequest], + call_next: CallNext[mt.ListPromptsRequest, ListPromptsResult], + ) -> ListPromptsResult: + return await call_next(context) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index c8d93f29f..86298d5e1 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -341,7 +341,8 @@ class FastMCP(Generic[LifespanResultT]): async def get_tools(self) -> dict[str, Tool]: """Get all registered tools, indexed by registered key.""" - return await self._list_tools(apply_middleware=False) + tools = await self._list_tools(apply_middleware=False) + return {tool.key: tool for tool in tools} async def get_tool(self, key: str) -> Tool: tools = await self.get_tools() @@ -515,7 +516,7 @@ class FastMCP(Generic[LifespanResultT]): tools = await self._middleware_list_tools() return [tool.to_mcp_tool(name=tool.name) for tool in tools] - async def _middleware_list_tools(self) -> dict[str, Tool]: + async def _middleware_list_tools(self) -> list[Tool]: """ List all available tools, in the format expected by the low-level MCP server. @@ -524,13 +525,13 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[mcp.types.ListToolsRequest], - ) -> list[MCPTool]: + ) -> list[Tool]: tools = await self._list_tools() - mcp_tools: list[MCPTool] = [] - for key, tool in tools.items(): + mcp_tools: list[Tool] = [] + for tool in tools: if self._should_enable_component(tool): - mcp_tools.append(tool.to_mcp_tool(name=key)) + mcp_tools.append(tool) return mcp_tools @@ -547,13 +548,13 @@ class FastMCP(Generic[LifespanResultT]): # Apply the middleware chain. return await self._apply_middleware(mw_context, _handler) - async def _list_tools(self, apply_middleware: bool = True) -> dict[str, Tool]: + async def _list_tools(self, apply_middleware: bool = True) -> list[Tool]: """ List all available tools. """ if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND: - tools: dict[str, Tool] = {} + tools: list[Tool] = [] # iterate such that new mounts overwrite older ones for mounted_server in self._mounted_servers: @@ -566,18 +567,17 @@ class FastMCP(Generic[LifespanResultT]): server_tools = await mounted_server.server._list_tools() # Apply prefix to each tool key if prefix exists and is not empty if mounted_server.prefix: - for tool in server_tools.values(): + for tool in server_tools: tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}") - tools[tool.key] = tool + tools.append(tool) else: - tools.update(server_tools) - tools.update(server_tools) + tools.extend(server_tools) except Exception as e: logger.warning( f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}" ) continue - tools.update(self._tool_manager.get_tools()) + tools.extend(self._tool_manager.get_tools().values()) self._cache.set("tools", tools) return tools diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index c21591788..3a2bf7f2c 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -7,7 +7,7 @@ import pytest from fastmcp import Client, FastMCP from fastmcp.server.context import Context -from fastmcp.server.middleware.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext @dataclass @@ -194,3 +194,43 @@ class TestMiddlewareHooks: assert recording_middleware.assert_called(hook="on_message", times=1) assert recording_middleware.assert_called(hook="on_request", times=1) assert recording_middleware.assert_called(hook="on_list_tools", times=1) + + 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(times=3) + assert recording_middleware.assert_called(method="resources/list", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_list_resources", times=1) + + 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(times=3) + assert recording_middleware.assert_called( + method="resource-templates/list", times=3 + ) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called( + hook="on_list_resource_templates", times=1 + ) + + 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(times=3) + assert recording_middleware.assert_called(method="prompts/list", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_list_prompts", times=1) From a42c0c40b04ef932e5b0f4a5ad2a7890841fda45 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 21:41:28 -0400 Subject: [PATCH 04/17] Add middleware for all current handlers --- src/fastmcp/server/middleware.py | 2 +- src/fastmcp/server/server.py | 296 +++++++++++++-------- tests/server/middleware/test_middleware.py | 262 +++++++++++++++++- 3 files changed, 442 insertions(+), 118 deletions(-) diff --git a/src/fastmcp/server/middleware.py b/src/fastmcp/server/middleware.py index 3f73bfe07..2363d41be 100644 --- a/src/fastmcp/server/middleware.py +++ b/src/fastmcp/server/middleware.py @@ -148,7 +148,7 @@ class MCPMiddleware: handler = partial(self.on_list_tools, call_next=handler) case "resources/list": handler = partial(self.on_list_resources, call_next=handler) - case "resource-templates/list": + case "resources/templates/list": handler = partial(self.on_list_resource_templates, call_next=handler) case "prompts/list": handler = partial(self.on_list_prompts, call_next=handler) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 86298d5e1..8eabe815d 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -352,34 +352,8 @@ class FastMCP(Generic[LifespanResultT]): async def get_resources(self) -> dict[str, Resource]: """Get all registered resources, indexed by registered key.""" - if (resources := self._cache.get("resources")) is self._cache.NOT_FOUND: - resources: dict[str, Resource] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - server_resources = await mounted_server.server.get_resources() - # Apply prefix to each resource key if prefix exists - if mounted_server.prefix: - for resource in server_resources.values(): - resource = resource.with_key( - add_resource_prefix( - resource.key, - mounted_server.prefix, - self.resource_prefix_format, - ) - ) - resources[resource.key] = resource - else: - resources.update(server_resources) - except Exception as e: - logger.warning( - f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}" - ) - continue - resources.update(self._resource_manager.get_resources()) - self._cache.set("resources", resources) - return resources + resources = await self._list_resources(apply_middleware=False) + return {resource.key: resource for resource in resources} async def get_resource(self, key: str) -> Resource: resources = await self.get_resources() @@ -389,39 +363,8 @@ class FastMCP(Generic[LifespanResultT]): async def get_resource_templates(self) -> dict[str, ResourceTemplate]: """Get all registered resource templates, indexed by registered key.""" - if ( - templates := self._cache.get("resource_templates") - ) is self._cache.NOT_FOUND: - templates: dict[str, ResourceTemplate] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - server_templates = ( - await mounted_server.server.get_resource_templates() - ) - # Apply prefix to each template key if prefix exists - if mounted_server.prefix: - for template in server_templates.values(): - template = template.with_key( - add_resource_prefix( - template.key, - mounted_server.prefix, - self.resource_prefix_format, - ) - ) - templates[template.key] = template - else: - templates.update(server_templates) - except Exception as e: - logger.warning( - "Failed to get resource templates from mounted server " - f"'{mounted_server.prefix}': {e}" - ) - continue - templates.update(self._resource_manager.get_templates()) - self._cache.set("resource_templates", templates) - return templates + templates = await self._list_resource_templates(apply_middleware=False) + return {template.key: template for template in templates} async def get_resource_template(self, key: str) -> ResourceTemplate: templates = await self.get_resource_templates() @@ -434,30 +377,8 @@ class FastMCP(Generic[LifespanResultT]): List all available prompts. """ - if (prompts := self._cache.get("prompts")) is self._cache.NOT_FOUND: - prompts: dict[str, Prompt] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - server_prompts = await mounted_server.server.get_prompts() - # Apply prefix to each prompt key if prefix exists - if mounted_server.prefix: - for prompt in server_prompts.values(): - prompt = prompt.with_key( - f"{mounted_server.prefix}_{prompt.key}" - ) - prompts[prompt.key] = prompt - else: - prompts.update(server_prompts) - except Exception as e: - logger.warning( - f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}" - ) - continue - prompts.update(self._prompt_manager.get_prompts()) - self._cache.set("prompts", prompts) - return prompts + prompts = await self._list_prompts(apply_middleware=False) + return {prompt.key: prompt for prompt in prompts} async def get_prompt(self, key: str) -> Prompt: prompts = await self.get_prompts() @@ -582,21 +503,31 @@ class FastMCP(Generic[LifespanResultT]): return tools async def _mcp_list_resources(self) -> list[MCPResource]: + logger.debug("Handler called: list_resources") + + with fastmcp.server.context.Context(fastmcp=self): + resources = await self._middleware_list_resources() + return [ + resource.to_mcp_resource(uri=resource.key) for resource in resources + ] + + async def _middleware_list_resources(self) -> list[Resource]: """ List all available resources, in the format expected by the low-level MCP server. """ - logger.debug("Handler called: list_resources") - async def _final_handler( + async def _handler( context: MiddlewareContext[dict[str, Any]], - ) -> list[MCPResource]: - resources = await self.get_resources() - mcp_resources: list[MCPResource] = [] - for key, resource in resources.items(): + ) -> list[Resource]: + resources = await self._list_resources() + + mcp_resources: list[Resource] = [] + for resource in resources: if self._should_enable_component(resource): - mcp_resources.append(resource.to_mcp_resource(uri=key)) + mcp_resources.append(resource) + return mcp_resources with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: @@ -610,24 +541,74 @@ class FastMCP(Generic[LifespanResultT]): ) # Apply the middleware chain. - return await self._apply_middleware(mw_context, _final_handler) + return await self._apply_middleware(mw_context, _handler) + + async def _list_resources(self, apply_middleware: bool = True) -> list[Resource]: + """ + List all available resources. + """ + + if (resources := self._cache.get("resources")) is self._cache.NOT_FOUND: + resources: list[Resource] = [] + + # iterate such that new mounts overwrite older ones + for mounted_server in self._mounted_servers: + try: + if apply_middleware: + server_resources = ( + await mounted_server.server._middleware_list_resources() + ) + else: + server_resources = await mounted_server.server._list_resources() + # Apply prefix to each resource key if prefix exists + if mounted_server.prefix: + for resource in server_resources: + resource = resource.with_key( + add_resource_prefix( + resource.key, + mounted_server.prefix, + self.resource_prefix_format, + ) + ) + resources.append(resource) + else: + resources.extend(server_resources) + except Exception as e: + logger.warning( + f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}" + ) + continue + resources.extend(self._resource_manager.get_resources().values()) + self._cache.set("resources", resources) + return resources async def _mcp_list_resource_templates(self) -> list[MCPResourceTemplate]: - """ - List all available resource templates, in the format expected by the low-level - MCP server. - - """ logger.debug("Handler called: list_resource_templates") - async def _final_handler( + with fastmcp.server.context.Context(fastmcp=self): + templates = await self._middleware_list_resource_templates() + return [ + template.to_mcp_template(uriTemplate=template.key) + for template in templates + ] + + async def _middleware_list_resource_templates(self) -> list[ResourceTemplate]: + """ + List all available resource templates, in the format expected by the low-level MCP + server. + + """ + + async def _handler( context: MiddlewareContext[dict[str, Any]], - ) -> list[MCPResourceTemplate]: - templates = await self.get_resource_templates() - mcp_templates: list[MCPResourceTemplate] = [] - for key, template in templates.items(): + ) -> list[ResourceTemplate]: + templates = await self._list_resource_templates() + + mcp_templates: list[ResourceTemplate] = [] + for template in templates: if self._should_enable_component(template): - mcp_templates.append(template.to_mcp_template(uriTemplate=key)) + mcp_templates.append(template) + return mcp_templates with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: @@ -636,29 +617,81 @@ class FastMCP(Generic[LifespanResultT]): message={}, # List resource templates doesn't have parameters source="client", type="request", - method="resources/list_templates", + method="resources/templates/list", fastmcp_context=fastmcp_ctx, ) # Apply the middleware chain. - return await self._apply_middleware(mw_context, _final_handler) + return await self._apply_middleware(mw_context, _handler) + + async def _list_resource_templates( + self, apply_middleware: bool = True + ) -> list[ResourceTemplate]: + """ + List all available resource templates. + """ + + if ( + templates := self._cache.get("resource_templates") + ) is self._cache.NOT_FOUND: + templates: list[ResourceTemplate] = [] + + # iterate such that new mounts overwrite older ones + for mounted_server in self._mounted_servers: + try: + if apply_middleware: + server_templates = await mounted_server.server._middleware_list_resource_templates() + else: + server_templates = ( + await mounted_server.server._list_resource_templates() + ) + # Apply prefix to each template key if prefix exists + if mounted_server.prefix: + for template in server_templates: + template = template.with_key( + add_resource_prefix( + template.key, + mounted_server.prefix, + self.resource_prefix_format, + ) + ) + templates.append(template) + else: + templates.extend(server_templates) + except Exception as e: + logger.warning( + "Failed to get resource templates from mounted server " + f"'{mounted_server.prefix}': {e}" + ) + continue + templates.extend(self._resource_manager.get_templates().values()) + self._cache.set("resource_templates", templates) + return templates async def _mcp_list_prompts(self) -> list[MCPPrompt]: + logger.debug("Handler called: list_prompts") + + with fastmcp.server.context.Context(fastmcp=self): + prompts = await self._middleware_list_prompts() + return [prompt.to_mcp_prompt(name=prompt.key) for prompt in prompts] + + async def _middleware_list_prompts(self) -> list[Prompt]: """ List all available prompts, in the format expected by the low-level MCP server. """ - logger.debug("Handler called: list_prompts") - async def _final_handler( + async def _handler( context: MiddlewareContext[dict[str, Any]], - ) -> list[MCPPrompt]: - prompts = await self.get_prompts() - mcp_prompts: list[MCPPrompt] = [] - for key, prompt in prompts.items(): + ) -> list[Prompt]: + prompts = await self._list_prompts() + + mcp_prompts: list[Prompt] = [] + for prompt in prompts: if self._should_enable_component(prompt): - mcp_prompts.append(prompt.to_mcp_prompt(name=key)) + mcp_prompts.append(prompt) + return mcp_prompts with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: @@ -672,7 +705,42 @@ class FastMCP(Generic[LifespanResultT]): ) # Apply the middleware chain. - return await self._apply_middleware(mw_context, _final_handler) + return await self._apply_middleware(mw_context, _handler) + + async def _list_prompts(self, apply_middleware: bool = True) -> list[Prompt]: + """ + List all available prompts. + """ + + if (prompts := self._cache.get("prompts")) is self._cache.NOT_FOUND: + prompts: list[Prompt] = [] + + # iterate such that new mounts overwrite older ones + for mounted_server in self._mounted_servers: + try: + if apply_middleware: + server_prompts = ( + await mounted_server.server._middleware_list_prompts() + ) + else: + server_prompts = await mounted_server.server._list_prompts() + # Apply prefix to each prompt key if prefix exists + if mounted_server.prefix: + for prompt in server_prompts: + prompt = prompt.with_key( + f"{mounted_server.prefix}_{prompt.key}" + ) + prompts.append(prompt) + else: + prompts.extend(server_prompts) + except Exception as e: + logger.warning( + f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}" + ) + continue + prompts.extend(self._prompt_manager.get_prompts().values()) + self._cache.set("prompts", prompts) + return prompts async def _mcp_call_tool( self, key: str, arguments: dict[str, Any] diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index 3a2bf7f2c..67f15138c 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -21,9 +21,10 @@ class Recording: class RecordingMiddleware(MCPMiddleware): """A middleware that automatically records all method calls.""" - def __init__(self): + 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.""" @@ -95,7 +96,7 @@ class RecordingMiddleware(MCPMiddleware): @pytest.fixture def recording_middleware(): """Fixture that provides a recording middleware instance.""" - middleware = RecordingMiddleware() + middleware = RecordingMiddleware(name="recording_middleware") yield middleware @@ -215,7 +216,7 @@ class TestMiddlewareHooks: assert recording_middleware.assert_called(times=3) assert recording_middleware.assert_called( - method="resource-templates/list", times=3 + method="resources/templates/list", times=3 ) assert recording_middleware.assert_called(hook="on_message", times=1) assert recording_middleware.assert_called(hook="on_request", times=1) @@ -234,3 +235,258 @@ class TestMiddlewareHooks: assert recording_middleware.assert_called(hook="on_message", times=1) assert recording_middleware.assert_called(hook="on_request", times=1) assert recording_middleware.assert_called(hook="on_list_prompts", times=1) + + +class TestNestedMiddlewareHooks: + @pytest.fixture + @staticmethod + def nested_middleware(): + return RecordingMiddleware(name="nested_middleware") + + @pytest.fixture + def nested_mcp_server(self, nested_middleware: RecordingMiddleware): + mcp = FastMCP(name="Nested MCP") + + @mcp.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(nested_middleware) + + return mcp + + async def test_call_tool_on_parent_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.call_tool("add", {"a": 1, "b": 2}) + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="tools/call", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_call_tool", times=1) + + assert nested_middleware.assert_called(times=0) + + async def test_call_tool_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.call_tool("nested_add", {"a": 1, "b": 2}) + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="tools/call", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_call_tool", times=1) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called(method="tools/call", times=3) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called(hook="on_call_tool", times=1) + + async def test_read_resource_on_parent_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.read_resource("resource://test") + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="resources/read", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_read_resource", times=1) + + assert nested_middleware.assert_called(times=0) + + async def test_read_resource_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.read_resource("resource://nested/test") + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="resources/read", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_read_resource", times=1) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called(method="resources/read", times=3) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called(hook="on_read_resource", times=1) + + async def test_get_prompt_on_parent_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.get_prompt("test_prompt", {"x": "test"}) + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="prompts/get", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_get_prompt", times=1) + + assert nested_middleware.assert_called(times=0) + + async def test_get_prompt_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.get_prompt("nested_test_prompt", {"x": "test"}) + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="prompts/get", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_get_prompt", times=1) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called(method="prompts/get", times=3) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called(hook="on_get_prompt", times=1) + + async def test_list_tools_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.list_tools() + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="tools/list", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_list_tools", times=1) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called(method="tools/list", times=3) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called(hook="on_list_tools", times=1) + + async def test_list_resources_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.list_resources() + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="resources/list", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_list_resources", times=1) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called(method="resources/list", times=3) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called(hook="on_list_resources", times=1) + + async def test_list_resource_templates_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.list_resource_templates() + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called( + method="resources/templates/list", times=3 + ) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called( + hook="on_list_resource_templates", times=1 + ) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called( + method="resources/templates/list", times=3 + ) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called( + hook="on_list_resource_templates", times=1 + ) From dbb6a76f636cf5c104ec475225edd9d709035de0 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 21:53:03 -0400 Subject: [PATCH 05/17] Update proxy behavior and conflict resolution --- src/fastmcp/server/proxy.py | 41 +++++++++++++++++++----------- src/fastmcp/server/server.py | 48 ++++++++++++++++++++---------------- 2 files changed, 53 insertions(+), 36 deletions(-) diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 91a03c89e..0d82c9cb3 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -172,8 +172,11 @@ class FastMCPProxy(FastMCP): super().__init__(**kwargs) self.client = client - async def get_tools(self) -> dict[str, Tool]: - tools = await super().get_tools() + async def _list_tools(self, apply_middleware: bool = True) -> list[Tool]: + tools = { + tool.name: tool + for tool in await super()._list_tools(apply_middleware=apply_middleware) + } async with self.client: try: @@ -187,12 +190,17 @@ class FastMCPProxy(FastMCP): # don't overwrite tools defined in the server if tool.name not in tools: tool_proxy = await ProxyTool.from_client(self.client, tool) - tools[tool_proxy.name] = tool_proxy + tools[tool_proxy.key] = tool_proxy - return tools + return list(tools.values()) - async def get_resources(self) -> dict[str, Resource]: - resources = await super().get_resources() + async def _list_resources(self, apply_middleware: bool = True) -> list[Resource]: + resources = { + resource.uri: resource + for resource in await super()._list_resources( + apply_middleware=apply_middleware + ) + } async with self.client: try: @@ -204,16 +212,19 @@ class FastMCPProxy(FastMCP): raise e for resource in client_resources: # don't overwrite resources defined in the server - if str(resource.uri) not in resources: + if resource.uri not in resources: resource_proxy = await ProxyResource.from_client( self.client, resource ) - resources[str(resource_proxy.uri)] = resource_proxy + resources[resource_proxy.uri] = resource_proxy - return resources + return list(resources.values()) - async def get_resource_templates(self) -> dict[str, ResourceTemplate]: - templates = await super().get_resource_templates() + async def _list_resource_templates(self) -> list[ResourceTemplate]: + templates = { + template.uri_template: template + for template in await super()._list_resource_templates() + } async with self.client: try: @@ -231,10 +242,10 @@ class FastMCPProxy(FastMCP): ) templates[template_proxy.uri_template] = template_proxy - return templates + return list(templates.values()) - async def get_prompts(self) -> dict[str, Prompt]: - prompts = await super().get_prompts() + async def _list_prompts(self) -> list[Prompt]: + prompts = {prompt.name: prompt for prompt in await super()._list_prompts()} async with self.client: try: @@ -250,7 +261,7 @@ class FastMCPProxy(FastMCP): prompt_proxy = await ProxyPrompt.from_client(self.client, prompt) prompts[prompt_proxy.name] = prompt_proxy - return prompts + return list(prompts.values()) async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: try: diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 8eabe815d..0805f98e1 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -435,7 +435,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self): tools = await self._middleware_list_tools() - return [tool.to_mcp_tool(name=tool.name) for tool in tools] + return [tool.to_mcp_tool(name=tool.key) for tool in tools] async def _middleware_list_tools(self) -> list[Tool]: """ @@ -475,7 +475,7 @@ class FastMCP(Generic[LifespanResultT]): """ if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND: - tools: list[Tool] = [] + tools: dict[str, Tool] = {} # iterate such that new mounts overwrite older ones for mounted_server in self._mounted_servers: @@ -490,17 +490,17 @@ class FastMCP(Generic[LifespanResultT]): if mounted_server.prefix: for tool in server_tools: tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}") - tools.append(tool) + tools[tool.key] = tool else: - tools.extend(server_tools) + tools.update({tool.key: tool for tool in server_tools}) except Exception as e: logger.warning( f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}" ) continue - tools.extend(self._tool_manager.get_tools().values()) + tools.update(self._tool_manager.get_tools()) self._cache.set("tools", tools) - return tools + return list(tools.values()) async def _mcp_list_resources(self) -> list[MCPResource]: logger.debug("Handler called: list_resources") @@ -549,7 +549,7 @@ class FastMCP(Generic[LifespanResultT]): """ if (resources := self._cache.get("resources")) is self._cache.NOT_FOUND: - resources: list[Resource] = [] + resources: dict[str, Resource] = {} # iterate such that new mounts overwrite older ones for mounted_server in self._mounted_servers: @@ -570,17 +570,19 @@ class FastMCP(Generic[LifespanResultT]): self.resource_prefix_format, ) ) - resources.append(resource) + resources[resource.key] = resource else: - resources.extend(server_resources) + resources.update( + {resource.key: resource for resource in server_resources} + ) except Exception as e: logger.warning( f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}" ) continue - resources.extend(self._resource_manager.get_resources().values()) + resources.update(self._resource_manager.get_resources()) self._cache.set("resources", resources) - return resources + return list(resources.values()) async def _mcp_list_resource_templates(self) -> list[MCPResourceTemplate]: logger.debug("Handler called: list_resource_templates") @@ -634,7 +636,7 @@ class FastMCP(Generic[LifespanResultT]): if ( templates := self._cache.get("resource_templates") ) is self._cache.NOT_FOUND: - templates: list[ResourceTemplate] = [] + templates: dict[str, ResourceTemplate] = {} # iterate such that new mounts overwrite older ones for mounted_server in self._mounted_servers: @@ -655,18 +657,20 @@ class FastMCP(Generic[LifespanResultT]): self.resource_prefix_format, ) ) - templates.append(template) + templates[template.key] = template else: - templates.extend(server_templates) + templates.update( + {template.key: template for template in server_templates} + ) except Exception as e: logger.warning( "Failed to get resource templates from mounted server " f"'{mounted_server.prefix}': {e}" ) continue - templates.extend(self._resource_manager.get_templates().values()) + templates.update(self._resource_manager.get_templates()) self._cache.set("resource_templates", templates) - return templates + return list(templates.values()) async def _mcp_list_prompts(self) -> list[MCPPrompt]: logger.debug("Handler called: list_prompts") @@ -713,7 +717,7 @@ class FastMCP(Generic[LifespanResultT]): """ if (prompts := self._cache.get("prompts")) is self._cache.NOT_FOUND: - prompts: list[Prompt] = [] + prompts: dict[str, Prompt] = {} # iterate such that new mounts overwrite older ones for mounted_server in self._mounted_servers: @@ -730,17 +734,19 @@ class FastMCP(Generic[LifespanResultT]): prompt = prompt.with_key( f"{mounted_server.prefix}_{prompt.key}" ) - prompts.append(prompt) + prompts[prompt.key] = prompt else: - prompts.extend(server_prompts) + prompts.update( + {prompt.key: prompt for prompt in server_prompts} + ) except Exception as e: logger.warning( f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}" ) continue - prompts.extend(self._prompt_manager.get_prompts().values()) + prompts.update(self._prompt_manager.get_prompts()) self._cache.set("prompts", prompts) - return prompts + return list(prompts.values()) async def _mcp_call_tool( self, key: str, arguments: dict[str, Any] From cafa11d3adde1b6c888592e1d25cc26f04e19e41 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 21:57:19 -0400 Subject: [PATCH 06/17] Fix missing proxy arg --- src/fastmcp/server/proxy.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 0d82c9cb3..a65f17b5b 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -220,10 +220,14 @@ class FastMCPProxy(FastMCP): return list(resources.values()) - async def _list_resource_templates(self) -> list[ResourceTemplate]: + async def _list_resource_templates( + self, apply_middleware: bool = True + ) -> list[ResourceTemplate]: templates = { template.uri_template: template - for template in await super()._list_resource_templates() + for template in await super()._list_resource_templates( + apply_middleware=apply_middleware + ) } async with self.client: @@ -244,8 +248,11 @@ class FastMCPProxy(FastMCP): return list(templates.values()) - async def _list_prompts(self) -> list[Prompt]: - prompts = {prompt.name: prompt for prompt in await super()._list_prompts()} + async def _list_prompts(self, apply_middleware: bool = True) -> list[Prompt]: + prompts = { + prompt.name: prompt + for prompt in await super()._list_prompts(apply_middleware=apply_middleware) + } async with self.client: try: From 607125abf8a5d55092aec8d552254dae4ddb8445 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 11:04:50 -0400 Subject: [PATCH 07/17] Update tool manager --- src/fastmcp/server/server.py | 123 ++++++-------------------- src/fastmcp/tools/tool_manager.py | 142 ++++++++++++++++++++++-------- tests/tools/test_tool_manager.py | 82 +++++++++-------- 3 files changed, 179 insertions(+), 168 deletions(-) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 0805f98e1..e4cf623f5 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -434,10 +434,10 @@ class FastMCP(Generic[LifespanResultT]): logger.debug("Handler called: list_tools") with fastmcp.server.context.Context(fastmcp=self): - tools = await self._middleware_list_tools() + tools = await self._list_tools() return [tool.to_mcp_tool(name=tool.key) for tool in tools] - async def _middleware_list_tools(self) -> list[Tool]: + async def _list_tools(self) -> list[Tool]: """ List all available tools, in the format expected by the low-level MCP server. @@ -447,7 +447,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[mcp.types.ListToolsRequest], ) -> list[Tool]: - tools = await self._list_tools() + tools = await self._tool_manager.list_tools() mcp_tools: list[Tool] = [] for tool in tools: @@ -469,39 +469,6 @@ class FastMCP(Generic[LifespanResultT]): # Apply the middleware chain. return await self._apply_middleware(mw_context, _handler) - async def _list_tools(self, apply_middleware: bool = True) -> list[Tool]: - """ - List all available tools. - """ - - if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND: - tools: dict[str, Tool] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - if apply_middleware: - server_tools = ( - await mounted_server.server._middleware_list_tools() - ) - else: - server_tools = await mounted_server.server._list_tools() - # Apply prefix to each tool key if prefix exists and is not empty - if mounted_server.prefix: - for tool in server_tools: - tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}") - tools[tool.key] = tool - else: - tools.update({tool.key: tool for tool in server_tools}) - except Exception as e: - logger.warning( - f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}" - ) - continue - tools.update(self._tool_manager.get_tools()) - self._cache.set("tools", tools) - return list(tools.values()) - async def _mcp_list_resources(self) -> list[MCPResource]: logger.debug("Handler called: list_resources") @@ -580,7 +547,11 @@ class FastMCP(Generic[LifespanResultT]): f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}" ) continue - resources.update(self._resource_manager.get_resources()) + ( + local_resources, + _, + ) = await self._resource_manager.get_resources_and_templates() + resources.update(local_resources) self._cache.set("resources", resources) return list(resources.values()) @@ -668,7 +639,11 @@ class FastMCP(Generic[LifespanResultT]): f"'{mounted_server.prefix}': {e}" ) continue - templates.update(self._resource_manager.get_templates()) + ( + _, + local_templates, + ) = await self._resource_manager.get_resources_and_templates() + templates.update(local_templates) self._cache.set("resource_templates", templates) return list(templates.values()) @@ -744,7 +719,7 @@ class FastMCP(Generic[LifespanResultT]): f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}" ) continue - prompts.update(self._prompt_manager.get_prompts()) + prompts.update(await self._prompt_manager.get_prompts()) self._cache.set("prompts", prompts) return list(prompts.values()) @@ -767,27 +742,26 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self): try: - return await self._middleware_call_tool(key, arguments) + return await self._call_tool(key, arguments) except DisabledError: raise NotFoundError(f"Unknown tool: {key}") except NotFoundError: raise NotFoundError(f"Unknown tool: {key}") - async def _middleware_call_tool( - self, - key: str, - arguments: dict[str, Any], - ) -> list[MCPContent]: + async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: """ - Call a tool with middleware. + Applies this server's middleware and delegates the filtered call to the manager. """ async def _handler( context: MiddlewareContext[mcp.types.CallToolRequestParams], ) -> list[MCPContent]: - return await self._call_tool( - key=context.message.name, - arguments=context.message.arguments or {}, + tool = await self._tool_manager.get_tool(context.message.name) + if not self._should_enable_component(tool): + raise NotFoundError(f"Unknown tool: {context.message.name!r}") + + return await self._tool_manager.call_tool( + key=context.message.name, arguments=context.message.arguments or {} ) mw_context = MiddlewareContext( @@ -799,46 +773,6 @@ class FastMCP(Generic[LifespanResultT]): ) return await self._apply_middleware(mw_context, _handler) - async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: - """ - Call a tool with raw MCP arguments. FastMCP subclasses should override - this method, not _mcp_call_tool. - - Args: - key: The name of the tool to call arguments: Arguments to pass to - the tool - - Returns: - List of MCP Content objects containing the tool results - """ - - # Get tool, checking first from our tools, then from the mounted servers - if self._tool_manager.has_tool(key): - tool = self._tool_manager.get_tool(key) - if not self._should_enable_component(tool): - raise DisabledError(f"Tool {key!r} is disabled") - return await self._tool_manager.call_tool(key, arguments) - - # Check mounted servers to see if they have the tool - # iterate such that new mounts take precedence over older ones - for mounted_server in reversed(self._mounted_servers): - tool_key = key - try: - # If server has a prefix, check if key matches and strip prefix - if mounted_server.prefix: - if tool_key.startswith(f"{mounted_server.prefix}_"): - tool_key = tool_key.removeprefix(f"{mounted_server.prefix}_") - else: - continue - return await mounted_server.server._middleware_call_tool( - tool_key, arguments - ) - except NotFoundError: - # Tool not found on this server, try the next one - continue - - raise NotFoundError(f"Unknown tool: {key!r}") - async def _mcp_read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]: """ Handle MCP 'readResource' requests. @@ -1853,11 +1787,12 @@ class FastMCP(Generic[LifespanResultT]): if as_proxy and not isinstance(server, FastMCPProxy): server = FastMCPProxy(Client(transport=FastMCPTransport(server))) - mounted_server = MountedServer( - server=server, - prefix=prefix, - ) - self._mounted_servers.append(mounted_server) + # Delegate mounting to all three managers + mounted_server = MountedServer(prefix=prefix, server=server) + self._tool_manager.mount(mounted_server) + self._resource_manager.mount(mounted_server) + self._prompt_manager.mount(mounted_server) + self._cache.clear() async def import_server( diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index fd8d45085..89696d642 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -1,8 +1,8 @@ -from __future__ import annotations as _annotations +from __future__ import annotations import warnings from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from mcp.types import ToolAnnotations @@ -14,7 +14,7 @@ from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import MCPContent if TYPE_CHECKING: - pass + from fastmcp.server.server import MountedServer logger = get_logger(__name__) @@ -28,6 +28,7 @@ class ToolManager: mask_error_details: bool | None = None, ): self._tools: dict[str, Tool] = {} + self._mounted_sources: list[MountedServer] = [] self.mask_error_details = mask_error_details or settings.mask_error_details # Default to "warn" if None is provided @@ -42,23 +43,78 @@ class ToolManager: self.duplicate_behavior = duplicate_behavior - def has_tool(self, key: str) -> bool: + def mount(self, server: MountedServer) -> None: + """Adds a mounted server as a source for tools.""" + self._mounted_sources.append(server) + + async def _load_tools( + self, *, mode: Literal["inventory", "protocol"] + ) -> dict[str, Tool]: + """ + The single, consolidated recursive method for fetching tools. The 'mode' + parameter determines the communication path. + + - mode="inventory": Manager-to-manager path for complete, unfiltered inventory + - mode="protocol": Server-to-server path for filtered MCP requests + """ + all_tools: dict[str, Tool] = {} + + for mounted in self._mounted_sources: + try: + if mode == "protocol": + # PATH 2: Use the server-to-server filtered path + child_results = await mounted.server._list_tools() + else: # mode == "inventory" + # PATH 1: Use the manager-to-manager unfiltered path + child_results = await mounted.server._tool_manager.get_tools() + + # The combination logic is the same for both paths + child_dict = ( + {t.key: t for t in child_results} + if isinstance(child_results, list) + else child_results + ) + if mounted.prefix: + for tool in child_dict.values(): + prefixed_tool = tool.with_key(f"{mounted.prefix}_{tool.key}") + all_tools[prefixed_tool.key] = prefixed_tool + else: + all_tools.update(child_dict) + except Exception as e: + # Skip failed mounts silently, matches existing behavior + logger.warning( + f"Failed to get tools from mounted server '{mounted.prefix}': {e}" + ) + continue + + # Finally, add local tools, which always take precedence + all_tools.update(self._tools) + return all_tools + + async def has_tool(self, key: str) -> bool: """Check if a tool exists.""" - return key in self._tools + tools = await self.get_tools() + return key in tools - def get_tool(self, key: str) -> Tool: + async def get_tool(self, key: str) -> Tool: """Get tool by key.""" - if key in self._tools: - return self._tools[key] - raise NotFoundError(f"Unknown tool: {key}") + tools = await self.get_tools() + if key in tools: + return tools[key] + raise NotFoundError(f"Tool {key!r} not found") - def get_tools(self) -> dict[str, Tool]: - """Get all registered tools, indexed by registered key.""" - return self._tools + async def get_tools(self) -> dict[str, Tool]: + """ + Gets the complete, unfiltered inventory of all tools. + """ + return await self._load_tools(mode="inventory") - def list_tools(self) -> list[Tool]: - """List all registered tools.""" - return list(self.get_tools().values()) + async def list_tools(self) -> list[Tool]: + """ + Lists all tools, applying protocol filtering. + """ + tools_dict = await self._load_tools(mode="protocol") + return list(tools_dict.values()) def add_tool_from_fn( self, @@ -119,28 +175,44 @@ class ToolManager: if key in self._tools: del self._tools[key] else: - raise NotFoundError(f"Unknown tool: {key}") + raise NotFoundError(f"Tool {key!r} not found") async def call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: - """Call a tool by name with arguments.""" - tool = self.get_tool(key) - if not tool: - raise NotFoundError(f"Unknown tool: {key}") + """ + Internal API for servers: Finds and calls a tool, respecting the + filtered protocol path. + """ + # 1. Check local tools first. The server will have already applied its filter. + if key in self._tools: + tool = await self.get_tool(key) + if not tool: + raise NotFoundError(f"Tool {key!r} not found") - try: - return await tool.run(arguments) + try: + return await tool.run(arguments) - # raise ToolErrors as-is - except ToolError as e: - logger.exception(f"Error calling tool {key!r}: {e}") - raise e + # raise ToolErrors as-is + except ToolError as e: + logger.exception(f"Error calling tool {key!r}: {e}") + raise e - # Handle other exceptions - except Exception as e: - logger.exception(f"Error calling tool {key!r}: {e}") - if self.mask_error_details: - # Mask internal details - raise ToolError(f"Error calling tool {key!r}") from e - else: - # Include original error details - raise ToolError(f"Error calling tool {key!r}: {e}") from e + # Handle other exceptions + except Exception as e: + logger.exception(f"Error calling tool {key!r}: {e}") + if self.mask_error_details: + # Mask internal details + raise ToolError(f"Error calling tool {key!r}") from e + else: + # Include original error details + raise ToolError(f"Error calling tool {key!r}: {e}") from e + + # 2. Check mounted servers using the filtered protocol path. + for mounted in reversed(self._mounted_sources): + if mounted.prefix and key.startswith(f"{mounted.prefix}_"): + key_on_child = key.removeprefix(f"{mounted.prefix}_") + try: + return await mounted.server._call_tool(key_on_child, arguments) + except NotFoundError: + continue + + raise NotFoundError(f"Tool {key!r} not found.") diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index b5cb15781..47923d61a 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -17,7 +17,7 @@ from fastmcp.utilities.types import Image class TestAddTools: - def test_basic_function(self): + async def test_basic_function(self): """Test registering and running a basic function.""" def add(a: int, b: int) -> int: @@ -28,7 +28,7 @@ class TestAddTools: tool = Tool.from_function(add) manager.add_tool(tool) - tool = manager.get_tool("add") + tool = await manager.get_tool("add") assert tool is not None assert tool.name == "add" assert tool.description == "Add two numbers." @@ -46,13 +46,13 @@ class TestAddTools: tool = Tool.from_function(fetch_data) manager.add_tool(tool) - tool = manager.get_tool("fetch_data") + tool = await manager.get_tool("fetch_data") assert tool is not None assert tool.name == "fetch_data" assert tool.description == "Fetch data from URL." assert tool.parameters["properties"]["url"]["type"] == "string" - def test_pydantic_model_function(self): + async def test_pydantic_model_function(self): """Test registering a function that takes a Pydantic model.""" class UserInput(BaseModel): @@ -67,7 +67,7 @@ class TestAddTools: tool = Tool.from_function(create_user) manager.add_tool(tool) - tool = manager.get_tool("create_user") + tool = await manager.get_tool("create_user") assert tool is not None assert tool.name == "create_user" assert tool.description == "Create a new user." @@ -75,7 +75,7 @@ class TestAddTools: assert "age" in tool.parameters["$defs"]["UserInput"]["properties"] assert "flag" in tool.parameters["properties"] - def test_callable_object(self): + async def test_callable_object(self): class Adder: """Adds two numbers.""" @@ -87,7 +87,7 @@ class TestAddTools: tool = Tool.from_function(Adder()) manager.add_tool(tool) - tool = manager.get_tool("Adder") + tool = await manager.get_tool("Adder") assert tool is not None assert tool.name == "Adder" assert tool.description == "Adds two numbers." @@ -95,7 +95,7 @@ class TestAddTools: assert tool.parameters["properties"]["x"]["type"] == "integer" assert tool.parameters["properties"]["y"]["type"] == "integer" - def test_async_callable_object(self): + async def test_async_callable_object(self): class Adder: """Adds two numbers.""" @@ -107,7 +107,7 @@ class TestAddTools: tool = Tool.from_function(Adder()) manager.add_tool(tool) - tool = manager.get_tool("Adder") + tool = await manager.get_tool("Adder") assert tool is not None assert tool.name == "Adder" assert tool.description == "Adds two numbers." @@ -123,7 +123,7 @@ class TestAddTools: tool = Tool.from_function(image_tool) manager.add_tool(tool) - tool = manager.get_tool("image_tool") + tool = await manager.get_tool("image_tool") result = await tool.run({"data": "test.png"}) assert tool.parameters["properties"]["data"]["type"] == "string" assert isinstance(result[0], ImageContent) @@ -148,7 +148,7 @@ class TestAddTools: tool = Tool.from_function(lambda x: x) manager.add_tool(tool) - def test_remove_tool_successfully(self): + async def test_remove_tool_successfully(self): """Test removing an added tool by key.""" manager = ToolManager() @@ -157,19 +157,19 @@ class TestAddTools: tool = Tool.from_function(add) manager.add_tool(tool) - assert manager.get_tool("add") is not None + assert await manager.get_tool("add") is not None manager.remove_tool("add") with pytest.raises(NotFoundError): - manager.get_tool("add") + await manager.get_tool("add") def test_remove_tool_missing_key(self): """Test removing a tool that does not exist raises NotFoundError.""" manager = ToolManager() - with pytest.raises(NotFoundError, match=f"Unknown tool: {'missing'}"): + with pytest.raises(NotFoundError, match="Tool 'missing' not found"): manager.remove_tool("missing") - def test_warn_on_duplicate_tools(self, caplog): + async def test_warn_on_duplicate_tools(self, caplog): """Test warning on duplicate tools.""" manager = ToolManager(duplicate_behavior="warn") @@ -183,7 +183,7 @@ class TestAddTools: assert "Tool already exists: test_tool" in caplog.text # Should have the tool - assert manager.get_tool("test_tool") is not None + assert await manager.get_tool("test_tool") is not None def test_disable_warn_on_duplicate_tools(self, caplog): """Test disabling warning on duplicate tools.""" @@ -213,7 +213,7 @@ class TestAddTools: tool2 = Tool.from_function(test_fn, name="test_tool") manager.add_tool(tool2) - def test_replace_duplicate_tools(self): + async def test_replace_duplicate_tools(self): """Test replacing duplicate tools.""" manager = ToolManager(duplicate_behavior="replace") @@ -229,12 +229,12 @@ class TestAddTools: manager.add_tool(result) # Should have replaced with the new tool - tool = manager.get_tool("test_tool") + tool = await manager.get_tool("test_tool") assert tool is not None assert isinstance(tool, FunctionTool) assert tool.fn.__name__ == "replacement_fn" - def test_ignore_duplicate_tools(self): + async def test_ignore_duplicate_tools(self): """Test ignoring duplicate tools.""" manager = ToolManager(duplicate_behavior="ignore") @@ -250,7 +250,7 @@ class TestAddTools: manager.add_tool(result) # Should keep the original - tool = manager.get_tool("test_tool") + tool = await manager.get_tool("test_tool") assert tool is not None assert isinstance(tool, FunctionTool) assert tool.fn.__name__ == "original_fn" @@ -262,7 +262,7 @@ class TestAddTools: class TestToolTags: """Test functionality related to tool tags.""" - def test_add_tool_with_tags(self): + async def test_add_tool_with_tags(self): """Test adding tags to a tool.""" def example_tool(x: int) -> int: @@ -274,11 +274,11 @@ class TestToolTags: manager.add_tool(tool) assert tool.tags == {"math", "utility"} - tool = manager.get_tool("example_tool") + tool = await manager.get_tool("example_tool") assert tool is not None assert tool.tags == {"math", "utility"} - def test_add_tool_with_empty_tags(self): + async def test_add_tool_with_empty_tags(self): """Test adding a tool with empty tags set.""" def example_tool(x: int) -> int: @@ -291,7 +291,7 @@ class TestToolTags: assert tool.tags == set() - def test_add_tool_with_none_tags(self): + async def test_add_tool_with_none_tags(self): """Test adding a tool with None tags.""" def example_tool(x: int) -> int: @@ -304,7 +304,7 @@ class TestToolTags: assert tool.tags == set() - def test_list_tools_with_tags(self): + async def test_list_tools_with_tags(self): """Test listing tools with specific tags.""" def math_tool(x: int) -> int: @@ -328,12 +328,16 @@ class TestToolTags: manager.add_tool(tool3) # Check if we can filter by tags when listing tools - math_tools = [tool for tool in manager.list_tools() if "math" in tool.tags] + math_tools = [ + tool for tool in (await manager.get_tools()).values() if "math" in tool.tags + ] assert len(math_tools) == 2 assert {tool.name for tool in math_tools} == {"math_tool", "mixed_tool"} utility_tools = [ - tool for tool in manager.list_tools() if "utility" in tool.tags + tool + for tool in (await manager.get_tools()).values() + if "utility" in tool.tags ] assert len(utility_tools) == 2 assert {tool.name for tool in utility_tools} == {"string_tool", "mixed_tool"} @@ -416,7 +420,7 @@ class TestCallTools: async def test_call_unknown_tool(self): manager = ToolManager() - with pytest.raises(NotFoundError, match="Unknown tool: unknown"): + with pytest.raises(NotFoundError, match="Tool 'unknown' not found"): await manager.call_tool("unknown", {"a": 1}) async def test_call_tool_with_list_int_input(self): @@ -728,7 +732,7 @@ class TestContextHandling: class TestCustomToolNames: """Test adding tools with custom names that differ from their function names.""" - def test_add_tool_with_custom_name(self): + async def test_add_tool_with_custom_name(self): """Test adding a tool with a custom name parameter using add_tool_from_fn.""" def original_fn(x: int) -> int: @@ -739,15 +743,15 @@ class TestCustomToolNames: manager.add_tool(tool) # The tool is stored under the custom name and its .name is also set to custom_name - assert manager.get_tool("custom_name") is not None + assert await manager.get_tool("custom_name") is not None assert tool.name == "custom_name" assert isinstance(tool, FunctionTool) assert tool.fn.__name__ == "original_fn" # The tool should not be accessible via its original function name - with pytest.raises(NotFoundError, match="Unknown tool: original_fn"): - manager.get_tool("original_fn") + with pytest.raises(NotFoundError, match="Tool 'original_fn' not found"): + await manager.get_tool("original_fn") - def test_add_tool_object_with_custom_key(self): + async def test_add_tool_object_with_custom_key(self): """Test adding a Tool object with a custom key using add_tool().""" def fn(x: int) -> int: @@ -759,13 +763,13 @@ class TestCustomToolNames: # Store it under a different name manager.add_tool(tool, key="proxy_tool") # The tool is accessible under the key - stored = manager.get_tool("proxy_tool") + stored = await manager.get_tool("proxy_tool") assert stored is not None # But the tool's .name is unchanged assert stored.name == "my_tool" # The tool is not accessible under its original name - with pytest.raises(NotFoundError, match="Unknown tool: my_tool"): - manager.get_tool("my_tool") + with pytest.raises(NotFoundError, match="Tool 'my_tool' not found"): + await manager.get_tool("my_tool") async def test_call_tool_with_custom_name(self): """Test calling a tool added with a custom name.""" @@ -783,10 +787,10 @@ class TestCustomToolNames: assert result[0].text == "15" # type: ignore[attr-defined] # Original name should not be registered - with pytest.raises(NotFoundError, match="Unknown tool: multiply"): + with pytest.raises(NotFoundError, match="Tool 'multiply' not found"): await manager.call_tool("multiply", {"a": 5, "b": 3}) - def test_replace_tool_keeps_original_name(self): + async def test_replace_tool_keeps_original_name(self): """Test that replacing a tool with "replace" keeps the original name.""" def original_fn(x: int) -> int: @@ -808,7 +812,7 @@ class TestCustomToolNames: manager.add_tool(replacement_tool) # The tool object should have been replaced - stored_tool = manager.get_tool("test_tool") + stored_tool = await manager.get_tool("test_tool") assert stored_tool is not None assert stored_tool == replacement_tool From 6b1ba049c7c981163cbc25ea495388e84cec43ad Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 11:33:46 -0400 Subject: [PATCH 08/17] Update prompt manager --- src/fastmcp/prompts/prompt_manager.py | 142 +++++++++++++++++++------ src/fastmcp/server/server.py | 145 ++++---------------------- 2 files changed, 130 insertions(+), 157 deletions(-) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index 10805f5e6..bde804a3d 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations as _annotations import warnings from collections.abc import Awaitable, Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from mcp import GetPromptResult @@ -13,7 +13,7 @@ from fastmcp.settings import DuplicateBehavior from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: - pass + from fastmcp.server.server import MountedServer logger = get_logger(__name__) @@ -27,6 +27,7 @@ class PromptManager: mask_error_details: bool | None = None, ): self._prompts: dict[str, Prompt] = {} + self._mounted_sources: list[MountedServer] = [] self.mask_error_details = mask_error_details or settings.mask_error_details # Default to "warn" if None is provided @@ -41,15 +42,80 @@ class PromptManager: self.duplicate_behavior = duplicate_behavior - def get_prompt(self, key: str) -> Prompt: + def mount(self, server: MountedServer) -> None: + """Adds a mounted server as a source for prompts.""" + self._mounted_sources.append(server) + + async def _load_prompts( + self, *, mode: Literal["inventory", "protocol"] + ) -> dict[str, Prompt]: + """ + The single, consolidated recursive method for fetching prompts. The 'mode' + parameter determines the communication path. + + - mode="inventory": Manager-to-manager path for complete, unfiltered inventory + - mode="protocol": Server-to-server path for filtered MCP requests + """ + all_prompts: dict[str, Prompt] = {} + + for mounted in self._mounted_sources: + try: + if mode == "protocol": + # PATH 2: Use the server-to-server filtered path + child_results = await mounted.server._list_prompts() + else: # mode == "inventory" + # PATH 1: Use the manager-to-manager unfiltered path + child_results = await mounted.server._prompt_manager.get_prompts() + + # The combination logic is the same for both paths + child_dict = ( + {p.key: p for p in child_results} + if isinstance(child_results, list) + else child_results + ) + if mounted.prefix: + for prompt in child_dict.values(): + prefixed_prompt = prompt.with_key( + f"{mounted.prefix}_{prompt.key}" + ) + all_prompts[prefixed_prompt.key] = prefixed_prompt + else: + all_prompts.update(child_dict) + except Exception as e: + # Skip failed mounts silently, matches existing behavior + logger.warning( + f"Failed to get prompts from mounted server '{mounted.prefix}': {e}" + ) + continue + + # Finally, add local prompts, which always take precedence + all_prompts.update(self._prompts) + return all_prompts + + async def has_prompt(self, key: str) -> bool: + """Check if a prompt exists.""" + prompts = await self.get_prompts() + return key in prompts + + async def get_prompt(self, key: str) -> Prompt: """Get prompt by key.""" - if key in self._prompts: - return self._prompts[key] + prompts = await self.get_prompts() + if key in prompts: + return prompts[key] raise NotFoundError(f"Unknown prompt: {key}") - def get_prompts(self) -> dict[str, Prompt]: - """Get all registered prompts, indexed by registered key.""" - return self._prompts + async def get_prompts(self) -> dict[str, Prompt]: + """ + Gets the complete, unfiltered inventory of all prompts. + """ + return await self._load_prompts(mode="inventory") + + async def list_prompts(self) -> list[Prompt]: + """ + Lists all prompts, applying protocol filtering. + """ + prompts_dict = await self._load_prompts(mode="protocol") + return list(prompts_dict.values()) def add_prompt_from_fn( self, @@ -96,30 +162,44 @@ class PromptManager: name: str, arguments: dict[str, Any] | None = None, ) -> GetPromptResult: - """Render a prompt by name with arguments.""" - prompt = self.get_prompt(name) - if not prompt: - raise NotFoundError(f"Unknown prompt: {name}") + """ + Internal API for servers: Finds and renders a prompt, respecting the + filtered protocol path. + """ + # 1. Check local prompts first. The server will have already applied its filter. + if name in self._prompts: + prompt = await self.get_prompt(name) + if not prompt: + raise NotFoundError(f"Unknown prompt: {name}") - try: - messages = await prompt.render(arguments) - return GetPromptResult(description=prompt.description, messages=messages) + try: + messages = await prompt.render(arguments) + return GetPromptResult( + description=prompt.description, messages=messages + ) - # Pass through PromptErrors as-is - except PromptError as e: - logger.exception(f"Error rendering prompt {name!r}: {e}") - raise e + # Pass through PromptErrors as-is + except PromptError as e: + logger.exception(f"Error rendering prompt {name!r}: {e}") + raise e - # Handle other exceptions - except Exception as e: - logger.exception(f"Error rendering prompt {name!r}: {e}") - if self.mask_error_details: - # Mask internal details - raise PromptError(f"Error rendering prompt {name!r}") - else: - # Include original error details - raise PromptError(f"Error rendering prompt {name!r}: {e}") + # Handle other exceptions + except Exception as e: + logger.exception(f"Error rendering prompt {name!r}: {e}") + if self.mask_error_details: + # Mask internal details + raise PromptError(f"Error rendering prompt {name!r}") from e + else: + # Include original error details + raise PromptError(f"Error rendering prompt {name!r}: {e}") from e - def has_prompt(self, key: str) -> bool: - """Check if a prompt exists.""" - return key in self._prompts + # 2. Check mounted servers using the filtered protocol path. + for mounted in reversed(self._mounted_sources): + if mounted.prefix and name.startswith(f"{mounted.prefix}_"): + name_on_child = name.removeprefix(f"{mounted.prefix}_") + try: + return await mounted.server._get_prompt(name_on_child, arguments) + except NotFoundError: + continue + + raise NotFoundError(f"Unknown prompt: {name}") diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index e4cf623f5..38b5b668b 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -341,8 +341,7 @@ class FastMCP(Generic[LifespanResultT]): async def get_tools(self) -> dict[str, Tool]: """Get all registered tools, indexed by registered key.""" - tools = await self._list_tools(apply_middleware=False) - return {tool.key: tool for tool in tools} + return await self._tool_manager.get_tools() async def get_tool(self, key: str) -> Tool: tools = await self.get_tools() @@ -376,9 +375,7 @@ class FastMCP(Generic[LifespanResultT]): """ List all available prompts. """ - - prompts = await self._list_prompts(apply_middleware=False) - return {prompt.key: prompt for prompt in prompts} + return await self._prompt_manager.get_prompts() async def get_prompt(self, key: str) -> Prompt: prompts = await self.get_prompts() @@ -651,10 +648,10 @@ class FastMCP(Generic[LifespanResultT]): logger.debug("Handler called: list_prompts") with fastmcp.server.context.Context(fastmcp=self): - prompts = await self._middleware_list_prompts() + prompts = await self._list_prompts() return [prompt.to_mcp_prompt(name=prompt.key) for prompt in prompts] - async def _middleware_list_prompts(self) -> list[Prompt]: + async def _list_prompts(self) -> list[Prompt]: """ List all available prompts, in the format expected by the low-level MCP server. @@ -662,9 +659,9 @@ class FastMCP(Generic[LifespanResultT]): """ async def _handler( - context: MiddlewareContext[dict[str, Any]], + context: MiddlewareContext[mcp.types.ListPromptsRequest], ) -> list[Prompt]: - prompts = await self._list_prompts() + prompts = await self._prompt_manager.list_prompts() mcp_prompts: list[Prompt] = [] for prompt in prompts: @@ -676,7 +673,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx: # Create the middleware context. mw_context = MiddlewareContext( - message={}, # List prompts doesn't have parameters + message=mcp.types.ListPromptsRequest(method="prompts/list"), source="client", type="request", method="prompts/list", @@ -686,43 +683,6 @@ class FastMCP(Generic[LifespanResultT]): # Apply the middleware chain. return await self._apply_middleware(mw_context, _handler) - async def _list_prompts(self, apply_middleware: bool = True) -> list[Prompt]: - """ - List all available prompts. - """ - - if (prompts := self._cache.get("prompts")) is self._cache.NOT_FOUND: - prompts: dict[str, Prompt] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - if apply_middleware: - server_prompts = ( - await mounted_server.server._middleware_list_prompts() - ) - else: - server_prompts = await mounted_server.server._list_prompts() - # Apply prefix to each prompt key if prefix exists - if mounted_server.prefix: - for prompt in server_prompts: - prompt = prompt.with_key( - f"{mounted_server.prefix}_{prompt.key}" - ) - prompts[prompt.key] = prompt - else: - prompts.update( - {prompt.key: prompt for prompt in server_prompts} - ) - except Exception as e: - logger.warning( - f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}" - ) - continue - prompts.update(await self._prompt_manager.get_prompts()) - self._cache.set("prompts", prompts) - return list(prompts.values()) - async def _mcp_call_tool( self, key: str, arguments: dict[str, Any] ) -> list[MCPContent]: @@ -828,7 +788,7 @@ class FastMCP(Generic[LifespanResultT]): Read a resource by URI, in the format expected by the low-level MCP server. """ - if self._resource_manager.has_resource(uri): + if await self._resource_manager.has_resource(uri): resource = await self._resource_manager.get_resource(uri) if not self._should_enable_component(resource): raise DisabledError(f"Resource {str(uri)!r} is disabled") @@ -840,32 +800,7 @@ class FastMCP(Generic[LifespanResultT]): ) ] else: - # iterate such that new mounts take precedence over older ones - for mounted_server in reversed(self._mounted_servers): - resource_uri = uri - try: - if mounted_server.prefix: - # If server has a prefix, check if URI matches and strip prefix - if has_resource_prefix( - str(resource_uri), - mounted_server.prefix, - self.resource_prefix_format, - ): - resource_uri = remove_resource_prefix( - str(resource_uri), - mounted_server.prefix, - self.resource_prefix_format, - ) - else: - continue - return await mounted_server.server._middleware_read_resource( - resource_uri - ) - except NotFoundError: - # Resource not found on this server, try the next one - continue - else: - raise NotFoundError(f"Unknown resource: {uri}") + raise NotFoundError(f"Unknown resource: {uri}") async def _mcp_get_prompt( self, name: str, arguments: dict[str, Any] | None = None @@ -879,7 +814,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self): try: - return await self._middleware_get_prompt(name, arguments) + return await self._get_prompt(name, arguments) except DisabledError: # convert to NotFoundError to avoid leaking prompt presence raise NotFoundError(f"Unknown prompt: {name}") @@ -887,21 +822,22 @@ class FastMCP(Generic[LifespanResultT]): # standardize NotFound message raise NotFoundError(f"Unknown prompt: {name}") - async def _middleware_get_prompt( - self, - name: str, - arguments: dict[str, Any] | None = None, + async def _get_prompt( + self, name: str, arguments: dict[str, Any] | None = None ) -> GetPromptResult: """ - Get a prompt with middleware. + Applies this server's middleware and delegates the filtered call to the manager. """ async def _handler( context: MiddlewareContext[mcp.types.GetPromptRequestParams], ) -> GetPromptResult: - return await self._get_prompt( - name=context.message.name, - arguments=context.message.arguments, + prompt = await self._prompt_manager.get_prompt(context.message.name) + if not self._should_enable_component(prompt): + raise NotFoundError(f"Unknown prompt: {context.message.name!r}") + + return await self._prompt_manager.render_prompt( + name=context.message.name, arguments=context.message.arguments ) mw_context = MiddlewareContext( @@ -913,49 +849,6 @@ class FastMCP(Generic[LifespanResultT]): ) return await self._apply_middleware(mw_context, _handler) - async def _get_prompt( - self, name: str, arguments: dict[str, Any] | None = None - ) -> GetPromptResult: - """Handle MCP 'getPrompt' requests. - - Args: - name: The name of the prompt to render - arguments: Arguments to pass to the prompt - - Returns: - GetPromptResult containing the rendered prompt messages - """ - logger.debug("Get prompt: %s with %s", name, arguments) - - # Get prompt, checking first from our prompts, then from the mounted servers - if self._prompt_manager.has_prompt(name): - prompt = self._prompt_manager.get_prompt(name) - if not self._should_enable_component(prompt): - raise DisabledError(f"Prompt {name!r} is disabled") - return await self._prompt_manager.render_prompt(name, arguments) - - # Check mounted servers to see if they have the prompt - # iterate such that new mounts take precedence over older ones - for mounted_server in reversed(self._mounted_servers): - prompt_name = name - try: - if mounted_server.prefix: - # If server has a prefix, check if name matches and strip prefix - if prompt_name.startswith(f"{mounted_server.prefix}_"): - prompt_name = prompt_name.removeprefix( - f"{mounted_server.prefix}_" - ) - else: - continue - return await mounted_server.server._middleware_get_prompt( - prompt_name, arguments - ) - except NotFoundError: - # Prompt not found on this server, try the next one - continue - - raise NotFoundError(f"Unknown prompt: {name}") - def add_tool(self, tool: Tool) -> None: """Add a tool to the server. From 35e13ef9bfa7a97036f8347f36b380dc4e20c35d Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 12:22:14 -0400 Subject: [PATCH 09/17] Update resource manager --- src/fastmcp/prompts/prompt_manager.py | 10 +- src/fastmcp/resources/resource_manager.py | 257 ++++++++++++++++++--- src/fastmcp/server/server.py | 168 +++----------- src/fastmcp/tools/tool_manager.py | 10 +- tests/resources/test_resource_manager.py | 8 +- tests/server/middleware/test_middleware.py | 56 +++++ tests/server/openapi/test_openapi.py | 42 ++-- tests/server/test_server.py | 2 +- tests/server/test_tool_annotations.py | 6 +- tests/server/test_tool_exclude_args.py | 4 +- 10 files changed, 349 insertions(+), 214 deletions(-) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index bde804a3d..ee1e6f3c9 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -65,14 +65,10 @@ class PromptManager: child_results = await mounted.server._list_prompts() else: # mode == "inventory" # PATH 1: Use the manager-to-manager unfiltered path - child_results = await mounted.server._prompt_manager.get_prompts() + child_results = await mounted.server._prompt_manager._list_prompts() # The combination logic is the same for both paths - child_dict = ( - {p.key: p for p in child_results} - if isinstance(child_results, list) - else child_results - ) + child_dict = {p.key: p for p in child_results} if mounted.prefix: for prompt in child_dict.values(): prefixed_prompt = prompt.with_key( @@ -110,7 +106,7 @@ class PromptManager: """ return await self._load_prompts(mode="inventory") - async def list_prompts(self) -> list[Prompt]: + async def _list_prompts(self) -> list[Prompt]: """ Lists all prompts, applying protocol filtering. """ diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 64a11aae1..8b307456e 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -1,9 +1,11 @@ """Resource manager functionality.""" +from __future__ import annotations + import inspect import warnings from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING, Any, Literal from pydantic import AnyUrl @@ -17,6 +19,9 @@ from fastmcp.resources.template import ( from fastmcp.settings import DuplicateBehavior from fastmcp.utilities.logging import get_logger +if TYPE_CHECKING: + from fastmcp.server.server import MountedServer + logger = get_logger(__name__) @@ -38,6 +43,7 @@ class ResourceManager: """ self._resources: dict[str, Resource] = {} self._templates: dict[str, ResourceTemplate] = {} + self._mounted_sources: list[MountedServer] = [] self.mask_error_details = mask_error_details or settings.mask_error_details # Default to "warn" if None is provided @@ -51,6 +57,128 @@ class ResourceManager: ) self.duplicate_behavior = duplicate_behavior + def mount(self, server: MountedServer) -> None: + """Adds a mounted server as a source for resources and templates.""" + self._mounted_sources.append(server) + + async def get_resources(self) -> dict[str, Resource]: + """Get all registered resources, keyed by URI.""" + return await self._load_resources(mode="inventory") + + async def get_resource_templates(self) -> dict[str, ResourceTemplate]: + """Get all registered templates, keyed by URI template.""" + return await self._load_resource_templates(mode="inventory") + + async def _load_resources( + self, *, mode: Literal["inventory", "protocol"] + ) -> dict[str, Resource]: + """ + The single, consolidated recursive method for fetching resources. The 'mode' + parameter determines the communication path. + + - mode="inventory": Manager-to-manager path for complete, unfiltered inventory + - mode="protocol": Server-to-server path for filtered MCP requests + """ + all_resources: dict[str, Resource] = {} + + for mounted in self._mounted_sources: + try: + if mode == "protocol": + # PATH 2: Use the server-to-server filtered path + child_resources_list = await mounted.server._list_resources() + child_resources = { + resource.key: resource for resource in child_resources_list + } + else: # mode == "inventory" + # PATH 1: Use the manager-to-manager unfiltered path + child_resources = ( + await mounted.server._resource_manager.get_resources() + ) + + # Apply prefix if needed + if mounted.prefix: + from fastmcp.server.server import add_resource_prefix + + for uri, resource in child_resources.items(): + prefixed_uri = add_resource_prefix( + uri, mounted.prefix, mounted.resource_prefix_format + ) + # Create a copy of the resource with the prefixed key + prefixed_resource = resource.with_key(prefixed_uri) + all_resources[prefixed_uri] = prefixed_resource + else: + all_resources.update(child_resources) + except Exception as e: + # Skip failed mounts silently, matches existing behavior + logger.warning( + f"Failed to get resources from mounted server '{mounted.prefix}': {e}" + ) + continue + + # Finally, add local resources, which always take precedence + all_resources.update(self._resources) + return all_resources + + async def _load_resource_templates( + self, *, mode: Literal["inventory", "protocol"] + ) -> dict[str, ResourceTemplate]: + """ + The single, consolidated recursive method for fetching templates. The 'mode' + parameter determines the communication path. + + - mode="inventory": Manager-to-manager path for complete, unfiltered inventory + - mode="protocol": Server-to-server path for filtered MCP requests + """ + all_templates: dict[str, ResourceTemplate] = {} + + for mounted in self._mounted_sources: + try: + if mode == "protocol": + # PATH 2: Use the server-to-server filtered path + child_templates = await mounted.server._list_resource_templates() + else: # mode == "inventory" + # PATH 1: Use the manager-to-manager unfiltered path + child_templates = await mounted.server._resource_manager._list_resource_templates() + child_dict = {template.key: template for template in child_templates} + + # Apply prefix if needed + if mounted.prefix: + from fastmcp.server.server import add_resource_prefix + + for uri_template, template in child_dict.items(): + prefixed_uri_template = add_resource_prefix( + uri_template, mounted.prefix, mounted.resource_prefix_format + ) + # Create a copy of the template with the prefixed key + prefixed_template = template.with_key(prefixed_uri_template) + all_templates[prefixed_uri_template] = prefixed_template + else: + all_templates.update(child_dict) + except Exception as e: + # Skip failed mounts silently, matches existing behavior + logger.warning( + f"Failed to get templates from mounted server '{mounted.prefix}': {e}" + ) + continue + + # Finally, add local templates, which always take precedence + all_templates.update(self._templates) + return all_templates + + async def _list_resources(self) -> list[Resource]: + """ + Lists all resources, applying protocol filtering. + """ + resources_dict = await self._load_resources(mode="protocol") + return list(resources_dict.values()) + + async def _list_resource_templates(self) -> list[ResourceTemplate]: + """ + Lists all templates, applying protocol filtering. + """ + templates_dict = await self._load_resource_templates(mode="protocol") + return list(templates_dict.values()) + def add_resource_or_template_from_fn( self, fn: Callable[..., Any], @@ -235,14 +363,21 @@ class ResourceManager: self._templates[storage_key] = template return template - def has_resource(self, uri: AnyUrl | str) -> bool: + async def has_resource(self, uri: AnyUrl | str) -> bool: """Check if a resource exists.""" uri_str = str(uri) - if uri_str in self._resources: + + # First check concrete resources (local and mounted) + resources = await self.get_resources() + if uri_str in resources: return True - for template_key in self._templates.keys(): + + # Then check templates (local and mounted) only if not found in concrete resources + templates = await self.get_resource_templates() + for template_key in templates.keys(): if match_uri_template(uri_str, template_key): return True + return False async def get_resource(self, uri: AnyUrl | str) -> Resource: @@ -257,12 +392,14 @@ class ResourceManager: uri_str = str(uri) logger.debug("Getting resource", extra={"uri": uri_str}) - # First check concrete resources - if resource := self._resources.get(uri_str): + # First check concrete resources (local and mounted) + resources = await self.get_resources() + if resource := resources.get(uri_str): return resource - # Then check templates - use the utility function to match against storage keys - for storage_key, template in self._templates.items(): + # Then check templates (local and mounted) - use the utility function to match against storage keys + templates = await self.get_resource_templates() + for storage_key, template in templates.items(): # Try to match against the storage key (which might be a custom key) if params := match_uri_template(uri_str, storage_key): try: @@ -289,31 +426,89 @@ class ResourceManager: raise NotFoundError(f"Unknown resource: {uri_str}") async def read_resource(self, uri: AnyUrl | str) -> str | bytes: - """Read a resource contents.""" - resource = await self.get_resource(uri) + """ + Internal API for servers: Finds and reads a resource, respecting the + filtered protocol path. + """ + uri_str = str(uri) - try: - return await resource.read() + # 1. Check local resources first. The server will have already applied its filter. + if uri_str in self._resources: + resource = await self.get_resource(uri_str) + if not resource: + raise NotFoundError(f"Resource {uri_str!r} not found") - # raise ResourceErrors as-is - except ResourceError as e: - logger.error(f"Error reading resource {uri!r}: {e}") - raise e + try: + return await resource.read() - # Handle other exceptions - except Exception as e: - logger.error(f"Error reading resource {uri!r}: {e}") - if self.mask_error_details: - # Mask internal details - raise ResourceError(f"Error reading resource {uri!r}") from e - else: - # Include original error details - raise ResourceError(f"Error reading resource {uri!r}: {e}") from e + # raise ResourceErrors as-is + except ResourceError as e: + logger.exception(f"Error reading resource {uri_str!r}: {e}") + raise e - def get_resources(self) -> dict[str, Resource]: - """Get all registered resources, keyed by URI.""" - return self._resources + # Handle other exceptions + except Exception as e: + logger.exception(f"Error reading resource {uri_str!r}: {e}") + if self.mask_error_details: + # Mask internal details + raise ResourceError(f"Error reading resource {uri_str!r}") from e + else: + # Include original error details + raise ResourceError( + f"Error reading resource {uri_str!r}: {e}" + ) from e - def get_templates(self) -> dict[str, ResourceTemplate]: - """Get all registered templates, keyed by URI template.""" - return self._templates + # 1b. Check local templates if not found in concrete resources + for template in self._templates.values(): + if params := match_uri_template(uri_str, template.uri_template): + try: + resource = await template.create_resource(uri_str, params=params) + return await resource.read() + except ResourceError as e: + logger.exception( + f"Error reading resource from template {uri_str!r}: {e}" + ) + raise e + except Exception as e: + logger.exception( + f"Error reading resource from template {uri_str!r}: {e}" + ) + if self.mask_error_details: + raise ResourceError( + f"Error reading resource from template {uri_str!r}" + ) from e + else: + raise ResourceError( + f"Error reading resource from template {uri_str!r}: {e}" + ) from e + + # 2. Check mounted servers using the filtered protocol path. + from fastmcp.server.server import has_resource_prefix, remove_resource_prefix + + for mounted in reversed(self._mounted_sources): + resource_uri = uri_str + try: + if mounted.prefix: + # If server has a prefix, check if URI matches and strip prefix + if has_resource_prefix( + resource_uri, + mounted.prefix, + mounted.resource_prefix_format, + ): + resource_uri = remove_resource_prefix( + resource_uri, + mounted.prefix, + mounted.resource_prefix_format, + ) + else: + continue + + result = await mounted.server._read_resource(resource_uri) + # Extract content from the first ReadResourceContents + if result and len(result) > 0: + return result[0].content + raise NotFoundError(f"Resource {uri_str!r} returned empty content") + except NotFoundError: + continue + + raise NotFoundError(f"Resource {uri_str!r} not found.") diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 38b5b668b..456572a8b 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -351,8 +351,7 @@ class FastMCP(Generic[LifespanResultT]): async def get_resources(self) -> dict[str, Resource]: """Get all registered resources, indexed by registered key.""" - resources = await self._list_resources(apply_middleware=False) - return {resource.key: resource for resource in resources} + return await self._resource_manager.get_resources() async def get_resource(self, key: str) -> Resource: resources = await self.get_resources() @@ -362,8 +361,7 @@ class FastMCP(Generic[LifespanResultT]): async def get_resource_templates(self) -> dict[str, ResourceTemplate]: """Get all registered resource templates, indexed by registered key.""" - templates = await self._list_resource_templates(apply_middleware=False) - return {template.key: template for template in templates} + return await self._resource_manager.get_resource_templates() async def get_resource_template(self, key: str) -> ResourceTemplate: templates = await self.get_resource_templates() @@ -444,7 +442,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[mcp.types.ListToolsRequest], ) -> list[Tool]: - tools = await self._tool_manager.list_tools() + tools = await self._tool_manager._list_tools() # type: ignore[reportPrivateUsage] mcp_tools: list[Tool] = [] for tool in tools: @@ -470,12 +468,12 @@ class FastMCP(Generic[LifespanResultT]): logger.debug("Handler called: list_resources") with fastmcp.server.context.Context(fastmcp=self): - resources = await self._middleware_list_resources() + resources = await self._list_resources() return [ resource.to_mcp_resource(uri=resource.key) for resource in resources ] - async def _middleware_list_resources(self) -> list[Resource]: + async def _list_resources(self) -> list[Resource]: """ List all available resources, in the format expected by the low-level MCP server. @@ -485,7 +483,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[dict[str, Any]], ) -> list[Resource]: - resources = await self._list_resources() + resources = await self._resource_manager._list_resources() # type: ignore[reportPrivateUsage] mcp_resources: list[Resource] = [] for resource in resources: @@ -507,62 +505,17 @@ class FastMCP(Generic[LifespanResultT]): # Apply the middleware chain. return await self._apply_middleware(mw_context, _handler) - async def _list_resources(self, apply_middleware: bool = True) -> list[Resource]: - """ - List all available resources. - """ - - if (resources := self._cache.get("resources")) is self._cache.NOT_FOUND: - resources: dict[str, Resource] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - if apply_middleware: - server_resources = ( - await mounted_server.server._middleware_list_resources() - ) - else: - server_resources = await mounted_server.server._list_resources() - # Apply prefix to each resource key if prefix exists - if mounted_server.prefix: - for resource in server_resources: - resource = resource.with_key( - add_resource_prefix( - resource.key, - mounted_server.prefix, - self.resource_prefix_format, - ) - ) - resources[resource.key] = resource - else: - resources.update( - {resource.key: resource for resource in server_resources} - ) - except Exception as e: - logger.warning( - f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}" - ) - continue - ( - local_resources, - _, - ) = await self._resource_manager.get_resources_and_templates() - resources.update(local_resources) - self._cache.set("resources", resources) - return list(resources.values()) - async def _mcp_list_resource_templates(self) -> list[MCPResourceTemplate]: logger.debug("Handler called: list_resource_templates") with fastmcp.server.context.Context(fastmcp=self): - templates = await self._middleware_list_resource_templates() + templates = await self._list_resource_templates() return [ template.to_mcp_template(uriTemplate=template.key) for template in templates ] - async def _middleware_list_resource_templates(self) -> list[ResourceTemplate]: + async def _list_resource_templates(self) -> list[ResourceTemplate]: """ List all available resource templates, in the format expected by the low-level MCP server. @@ -572,7 +525,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[dict[str, Any]], ) -> list[ResourceTemplate]: - templates = await self._list_resource_templates() + templates = await self._resource_manager._list_resource_templates() mcp_templates: list[ResourceTemplate] = [] for template in templates: @@ -594,56 +547,6 @@ class FastMCP(Generic[LifespanResultT]): # Apply the middleware chain. return await self._apply_middleware(mw_context, _handler) - async def _list_resource_templates( - self, apply_middleware: bool = True - ) -> list[ResourceTemplate]: - """ - List all available resource templates. - """ - - if ( - templates := self._cache.get("resource_templates") - ) is self._cache.NOT_FOUND: - templates: dict[str, ResourceTemplate] = {} - - # iterate such that new mounts overwrite older ones - for mounted_server in self._mounted_servers: - try: - if apply_middleware: - server_templates = await mounted_server.server._middleware_list_resource_templates() - else: - server_templates = ( - await mounted_server.server._list_resource_templates() - ) - # Apply prefix to each template key if prefix exists - if mounted_server.prefix: - for template in server_templates: - template = template.with_key( - add_resource_prefix( - template.key, - mounted_server.prefix, - self.resource_prefix_format, - ) - ) - templates[template.key] = template - else: - templates.update( - {template.key: template for template in server_templates} - ) - except Exception as e: - logger.warning( - "Failed to get resource templates from mounted server " - f"'{mounted_server.prefix}': {e}" - ) - continue - ( - _, - local_templates, - ) = await self._resource_manager.get_resources_and_templates() - templates.update(local_templates) - self._cache.set("resource_templates", templates) - return list(templates.values()) - async def _mcp_list_prompts(self) -> list[MCPPrompt]: logger.debug("Handler called: list_prompts") @@ -661,7 +564,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[mcp.types.ListPromptsRequest], ) -> list[Prompt]: - prompts = await self._prompt_manager.list_prompts() + prompts = await self._prompt_manager._list_prompts() # type: ignore[reportPrivateUsage] mcp_prompts: list[Prompt] = [] for prompt in prompts: @@ -743,7 +646,7 @@ class FastMCP(Generic[LifespanResultT]): with fastmcp.server.context.Context(fastmcp=self): try: - return await self._middleware_read_resource(uri) + return await self._read_resource(uri) except DisabledError: # convert to NotFoundError to avoid leaking resource presence raise NotFoundError(f"Unknown resource: {str(uri)!r}") @@ -751,25 +654,28 @@ class FastMCP(Generic[LifespanResultT]): # standardize NotFound message raise NotFoundError(f"Unknown resource: {str(uri)!r}") - async def _middleware_read_resource( - self, - uri: AnyUrl | str, - ) -> list[ReadResourceContents]: + async def _read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]: """ - Read a resource with middleware. + Applies this server's middleware and delegates the filtered call to the manager. """ async def _handler( context: MiddlewareContext[mcp.types.ReadResourceRequestParams], ) -> list[ReadResourceContents]: - return await self._read_resource( - uri=context.message.uri, - ) + resource = await self._resource_manager.get_resource(context.message.uri) + if not self._should_enable_component(resource): + raise NotFoundError(f"Unknown resource: {str(context.message.uri)!r}") + + content = await self._resource_manager.read_resource(context.message.uri) + return [ + ReadResourceContents( + content=content, + mime_type=resource.mime_type, + ) + ] # Convert string URI to AnyUrl if needed if isinstance(uri, str): - from pydantic import AnyUrl - uri_param = AnyUrl(uri) else: uri_param = uri @@ -783,25 +689,6 @@ class FastMCP(Generic[LifespanResultT]): ) return await self._apply_middleware(mw_context, _handler) - async def _read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]: - """ - Read a resource by URI, in the format expected by the low-level MCP - server. - """ - if await self._resource_manager.has_resource(uri): - resource = await self._resource_manager.get_resource(uri) - if not self._should_enable_component(resource): - raise DisabledError(f"Resource {str(uri)!r} is disabled") - content = await self._resource_manager.read_resource(uri) - return [ - ReadResourceContents( - content=content, - mime_type=resource.mime_type, - ) - ] - else: - raise NotFoundError(f"Unknown resource: {uri}") - async def _mcp_get_prompt( self, name: str, arguments: dict[str, Any] | None = None ) -> GetPromptResult: @@ -1681,7 +1568,11 @@ class FastMCP(Generic[LifespanResultT]): server = FastMCPProxy(Client(transport=FastMCPTransport(server))) # Delegate mounting to all three managers - mounted_server = MountedServer(prefix=prefix, server=server) + mounted_server = MountedServer( + prefix=prefix, + server=server, + resource_prefix_format=self.resource_prefix_format, + ) self._tool_manager.mount(mounted_server) self._resource_manager.mount(mounted_server) self._prompt_manager.mount(mounted_server) @@ -1979,6 +1870,7 @@ class FastMCP(Generic[LifespanResultT]): class MountedServer: prefix: str | None server: FastMCP[Any] + resource_prefix_format: Literal["protocol", "path"] | None = None def add_resource_prefix( diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index 89696d642..d067f72d7 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -66,14 +66,10 @@ class ToolManager: child_results = await mounted.server._list_tools() else: # mode == "inventory" # PATH 1: Use the manager-to-manager unfiltered path - child_results = await mounted.server._tool_manager.get_tools() + child_results = await mounted.server._tool_manager._list_tools() # The combination logic is the same for both paths - child_dict = ( - {t.key: t for t in child_results} - if isinstance(child_results, list) - else child_results - ) + child_dict = {t.key: t for t in child_results} if mounted.prefix: for tool in child_dict.values(): prefixed_tool = tool.with_key(f"{mounted.prefix}_{tool.key}") @@ -109,7 +105,7 @@ class ToolManager: """ return await self._load_tools(mode="inventory") - async def list_tools(self) -> list[Tool]: + async def _list_tools(self) -> list[Tool]: """ Lists all tools, applying protocol filtering. """ diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index 26e45b6e2..a64fded09 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -180,7 +180,7 @@ class TestResourceManager: assert "Template already exists" in caplog.text # Should have the template - assert manager.get_templates() == {"test://{id}": template} + assert manager.get_resource_templates() == {"test://{id}": template} def test_error_on_duplicate_templates(self): """Test error on duplicate templates.""" @@ -226,7 +226,7 @@ class TestResourceManager: manager.add_template(template2) # Should have replaced with the new template - templates = list(manager.get_templates().values()) + templates = list(manager.get_resource_templates().values()) assert len(templates) == 1 assert templates[0].name == "replacement" @@ -256,7 +256,7 @@ class TestResourceManager: result = manager.add_template(template2) # Should keep the original - templates = list(manager.get_templates().values()) + templates = list(manager.get_resource_templates().values()) assert len(templates) == 1 assert templates[0].name == "original" # Result should be the original template @@ -379,7 +379,7 @@ class TestResourceTags: ) manager.add_template(template) - templates = list(manager.get_templates().values()) + templates = list(manager.get_resource_templates().values()) assert len(templates) == 1 assert templates[0].tags == {"users", "template", "data"} diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index 67f15138c..c789f42be 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -172,6 +172,18 @@ class TestMiddlewareHooks: assert recording_middleware.assert_called(hook="on_request", times=1) assert recording_middleware.assert_called(hook="on_read_resource", times=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(times=3) + assert recording_middleware.assert_called(method="resources/read", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_read_resource", times=1) + async def test_get_prompt( self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware ): @@ -367,6 +379,50 @@ class TestNestedMiddlewareHooks: assert nested_middleware.assert_called(hook="on_request", times=1) assert nested_middleware.assert_called(hook="on_read_resource", times=1) + async def test_read_resource_template_on_parent_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.read_resource("resource://test-template/1") + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="resources/read", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_read_resource", times=1) + + assert nested_middleware.assert_called(times=0) + + async def test_read_resource_template_on_nested_server( + self, + mcp_server: FastMCP, + nested_mcp_server: FastMCP, + recording_middleware: RecordingMiddleware, + nested_middleware: RecordingMiddleware, + ): + mcp_server.mount(nested_mcp_server, prefix="nested") + + async with Client(mcp_server) as client: + await client.read_resource("resource://nested/test-template/1") + + assert recording_middleware.assert_called(times=3) + assert recording_middleware.assert_called(method="resources/read", times=3) + assert recording_middleware.assert_called(hook="on_message", times=1) + assert recording_middleware.assert_called(hook="on_request", times=1) + assert recording_middleware.assert_called(hook="on_read_resource", times=1) + + assert nested_middleware.assert_called(times=3) + assert nested_middleware.assert_called(method="resources/read", times=3) + assert nested_middleware.assert_called(hook="on_message", times=1) + assert nested_middleware.assert_called(hook="on_request", times=1) + assert nested_middleware.assert_called(hook="on_read_resource", times=1) + async def test_get_prompt_on_parent_server( self, mcp_server: FastMCP, diff --git a/tests/server/openapi/test_openapi.py b/tests/server/openapi/test_openapi.py index c7a0b74f6..5fc6611ae 100644 --- a/tests/server/openapi/test_openapi.py +++ b/tests/server/openapi/test_openapi.py @@ -478,7 +478,7 @@ class TestTagTransfer: ): """Test that tags from OpenAPI routes are correctly transferred to Tools.""" # Get internal tools directly (not the public API which returns MCP.Content) - tools = fastmcp_openapi_server_with_all_types._tool_manager.list_tools() + tools = fastmcp_openapi_server_with_all_types._tool_manager._list_tools() # Find the create_user and update_user_name tools create_user_tool = next( @@ -528,7 +528,7 @@ class TestTagTransfer: """Test that tags from OpenAPI routes are correctly transferred to ResourceTemplates.""" # Get internal resource templates directly templates = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_templates().values() + fastmcp_openapi_server_with_all_types._resource_manager.get_resource_templates().values() ) # Find the get_user template @@ -549,7 +549,7 @@ class TestTagTransfer: """Test that tags are preserved when creating resources from templates.""" # Get internal resource templates directly templates = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_templates().values() + fastmcp_openapi_server_with_all_types._resource_manager.get_resource_templates().values() ) # Find the get_user template @@ -1537,11 +1537,11 @@ class TestFastAPIDescriptionPropagation: print(f" Resource: {name}, Name attribute: {resource.name}") print("\nDEBUG - Templates created:") - for name, template in server._resource_manager.get_templates().items(): + for name, template in server._resource_manager.get_resource_templates().items(): print(f" Template: {name}, Name attribute: {template.name}") print("\nDEBUG - Tools created:") - for tool in server._tool_manager.list_tools(): + for tool in server._tool_manager._list_tools(): print(f" Tool: {tool.name}") return server @@ -1759,7 +1759,7 @@ class TestReprMethods: self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI ): """Test that OpenAPITool's __repr__ method works without recursion errors.""" - tools = fastmcp_openapi_server_with_all_types._tool_manager.list_tools() + tools = fastmcp_openapi_server_with_all_types._tool_manager._list_tools() tool = next(iter(tools)) # Verify repr doesn't cause recursion and contains expected elements @@ -1790,7 +1790,7 @@ class TestReprMethods: ): """Test that OpenAPIResourceTemplate's __repr__ method works without recursion errors.""" templates = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_templates().values() + fastmcp_openapi_server_with_all_types._resource_manager.get_resource_templates().values() ) template = next(iter(templates)) @@ -1836,7 +1836,7 @@ class TestEnumHandling: ) # Get the tools from the server - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() # Find the read_item tool read_item_tool = next((t for t in tools if t.name == "read_item_items"), None) @@ -1929,7 +1929,7 @@ class TestRouteMapWildcard: ) # All operations should be mapped to tools - tools = mcp._tool_manager.list_tools() + tools = mcp._tool_manager._list_tools() tool_names = {tool.name for tool in tools} # Check that all 4 operations became tools @@ -2246,12 +2246,12 @@ class TestMCPNames: ) # Check tools use custom names - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} assert "admin_create_user" in tool_names # Check resource templates use custom names - templates = list(server._resource_manager.get_templates().values()) + templates = list(server._resource_manager.get_resource_templates().values()) template_names = {template.name for template in templates} assert "user_detail" in template_names @@ -2276,10 +2276,10 @@ class TestMCPNames: route_maps=GET_ROUTE_MAPS, ) - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} - templates = list(server._resource_manager.get_templates().values()) + templates = list(server._resource_manager.get_resource_templates().values()) template_names = {template.name for template in templates} resources = list(server._resource_manager.get_resources().values()) @@ -2330,13 +2330,13 @@ class TestMCPNames: # Check all component types all_names = [] - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() all_names.extend(tool.name for tool in tools) resources = list(server._resource_manager.get_resources().values()) all_names.extend(resource.name for resource in resources) - templates = list(server._resource_manager.get_templates().values()) + templates = list(server._resource_manager.get_resource_templates().values()) all_names.extend(template.name for template in templates) # All names should be 56 characters or less @@ -2363,7 +2363,7 @@ class TestMCPNames: mcp_names=mcp_names, ) - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} assert "openapi_user_list" in tool_names @@ -2395,7 +2395,7 @@ class TestMCPNames: mcp_names=mcp_names, ) - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} assert "fastapi_create_user" in tool_names @@ -2496,7 +2496,7 @@ class TestRouteMapMCPTags: ) # Get the POST tool - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() create_user_tool = next((t for t in tools if "create_user" in t.name), None) assert create_user_tool is not None, "create_user tool not found" @@ -2564,7 +2564,7 @@ class TestRouteMapMCPTags: ) # Get the resource template - templates = list(server._resource_manager.get_templates().values()) + templates = list(server._resource_manager.get_resource_templates().values()) get_user_template = next((t for t in templates if "get_user" in t.name), None) assert get_user_template is not None, "get_user template not found" @@ -2610,14 +2610,14 @@ class TestRouteMapMCPTags: ) # Check tool tags - tools = server._tool_manager.list_tools() + tools = server._tool_manager._list_tools() create_tool = next((t for t in tools if "create_user" in t.name), None) assert create_tool is not None assert "write-operation" in create_tool.tags assert "mutation" in create_tool.tags # Check resource template tags - templates = list(server._resource_manager.get_templates().values()) + templates = list(server._resource_manager.get_resource_templates().values()) detail_template = next((t for t in templates if "get_user" in t.name), None) assert detail_template is not None assert "detail" in detail_template.tags diff --git a/tests/server/test_server.py b/tests/server/test_server.py index 403c94cae..21b13b841 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -282,7 +282,7 @@ class TestToolDecorator: return x * 2 # Verify the tags were set correctly - tools = mcp._tool_manager.list_tools() + tools = await mcp._tool_manager._list_tools() assert len(tools) == 1 assert tools[0].tags == {"example", "test-tag"} diff --git a/tests/server/test_tool_annotations.py b/tests/server/test_tool_annotations.py index dd569d9e4..84d031df9 100644 --- a/tests/server/test_tool_annotations.py +++ b/tests/server/test_tool_annotations.py @@ -22,7 +22,7 @@ async def test_tool_annotations_in_tool_manager(): return message # Check internal tool objects directly - tools = mcp._tool_manager.list_tools() + tools = mcp._tool_manager._list_tools() assert len(tools) == 1 assert tools[0].annotations is not None assert tools[0].annotations.title == "Echo Tool" @@ -124,7 +124,7 @@ async def test_direct_tool_annotations_in_tool_manager(): return {"modified": True, **data} # Check internal tool objects directly - tools = mcp._tool_manager.list_tools() + tools = mcp._tool_manager._list_tools() assert len(tools) == 1 assert tools[0].annotations is not None assert tools[0].annotations.title == "Direct Tool" @@ -183,7 +183,7 @@ async def test_add_tool_method_annotations(): mcp.add_tool(tool) # Check internal tool objects directly - tools = mcp._tool_manager.list_tools() + tools = mcp._tool_manager._list_tools() assert len(tools) == 1 assert tools[0].annotations is not None assert tools[0].annotations.title == "Create Item" diff --git a/tests/server/test_tool_exclude_args.py b/tests/server/test_tool_exclude_args.py index efe8349a7..cc57bf3c2 100644 --- a/tests/server/test_tool_exclude_args.py +++ b/tests/server/test_tool_exclude_args.py @@ -19,7 +19,7 @@ async def test_tool_exclude_args_in_tool_manager(): pass return message - tools = mcp._tool_manager.list_tools() + tools = mcp._tool_manager._list_tools() assert len(tools) == 1 assert "state" not in echo.parameters["properties"] @@ -60,7 +60,7 @@ async def test_add_tool_method_exclude_args(): mcp.add_tool(tool) # Check internal tool objects directly - tools = mcp._tool_manager.list_tools() + tools = mcp._tool_manager._list_tools() assert len(tools) == 1 assert "state" not in tools[0].parameters["properties"] From d8235e1af7f2eebb87dc682c64d1a3029bb4a6eb Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 12:44:34 -0400 Subject: [PATCH 10/17] Clean up mode kwarg --- src/fastmcp/prompts/prompt_manager.py | 24 ++++++------- src/fastmcp/resources/resource_manager.py | 44 +++++++++++------------ src/fastmcp/tools/tool_manager.py | 24 ++++++------- 3 files changed, 43 insertions(+), 49 deletions(-) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index ee1e6f3c9..a4aaaf1ae 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations as _annotations import warnings from collections.abc import Awaitable, Callable -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any from mcp import GetPromptResult @@ -46,25 +46,23 @@ class PromptManager: """Adds a mounted server as a source for prompts.""" self._mounted_sources.append(server) - async def _load_prompts( - self, *, mode: Literal["inventory", "protocol"] - ) -> dict[str, Prompt]: + async def _load_prompts(self, *, via_server: bool = False) -> dict[str, Prompt]: """ - The single, consolidated recursive method for fetching prompts. The 'mode' + The single, consolidated recursive method for fetching prompts. The 'via_server' parameter determines the communication path. - - mode="inventory": Manager-to-manager path for complete, unfiltered inventory - - mode="protocol": Server-to-server path for filtered MCP requests + - via_server=False: Manager-to-manager path for complete, unfiltered inventory + - via_server=True: Server-to-server path for filtered MCP requests """ all_prompts: dict[str, Prompt] = {} for mounted in self._mounted_sources: try: - if mode == "protocol": - # PATH 2: Use the server-to-server filtered path + if via_server: + # Use the server-to-server filtered path child_results = await mounted.server._list_prompts() - else: # mode == "inventory" - # PATH 1: Use the manager-to-manager unfiltered path + else: + # Use the manager-to-manager unfiltered path child_results = await mounted.server._prompt_manager._list_prompts() # The combination logic is the same for both paths @@ -104,13 +102,13 @@ class PromptManager: """ Gets the complete, unfiltered inventory of all prompts. """ - return await self._load_prompts(mode="inventory") + return await self._load_prompts(via_server=False) async def _list_prompts(self) -> list[Prompt]: """ Lists all prompts, applying protocol filtering. """ - prompts_dict = await self._load_prompts(mode="protocol") + prompts_dict = await self._load_prompts(via_server=True) return list(prompts_dict.values()) def add_prompt_from_fn( diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 8b307456e..a22e515dc 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -5,7 +5,7 @@ from __future__ import annotations import inspect import warnings from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any from pydantic import AnyUrl @@ -63,34 +63,32 @@ class ResourceManager: async def get_resources(self) -> dict[str, Resource]: """Get all registered resources, keyed by URI.""" - return await self._load_resources(mode="inventory") + return await self._load_resources(via_server=False) async def get_resource_templates(self) -> dict[str, ResourceTemplate]: """Get all registered templates, keyed by URI template.""" - return await self._load_resource_templates(mode="inventory") + return await self._load_resource_templates(via_server=False) - async def _load_resources( - self, *, mode: Literal["inventory", "protocol"] - ) -> dict[str, Resource]: + async def _load_resources(self, *, via_server: bool = False) -> dict[str, Resource]: """ - The single, consolidated recursive method for fetching resources. The 'mode' + The single, consolidated recursive method for fetching resources. The 'via_server' parameter determines the communication path. - - mode="inventory": Manager-to-manager path for complete, unfiltered inventory - - mode="protocol": Server-to-server path for filtered MCP requests + - via_server=False: Manager-to-manager path for complete, unfiltered inventory + - via_server=True: Server-to-server path for filtered MCP requests """ all_resources: dict[str, Resource] = {} for mounted in self._mounted_sources: try: - if mode == "protocol": - # PATH 2: Use the server-to-server filtered path + if via_server: + # Use the server-to-server filtered path child_resources_list = await mounted.server._list_resources() child_resources = { resource.key: resource for resource in child_resources_list } - else: # mode == "inventory" - # PATH 1: Use the manager-to-manager unfiltered path + else: + # Use the manager-to-manager unfiltered path child_resources = ( await mounted.server._resource_manager.get_resources() ) @@ -120,24 +118,24 @@ class ResourceManager: return all_resources async def _load_resource_templates( - self, *, mode: Literal["inventory", "protocol"] + self, *, via_server: bool = False ) -> dict[str, ResourceTemplate]: """ - The single, consolidated recursive method for fetching templates. The 'mode' + The single, consolidated recursive method for fetching templates. The 'via_server' parameter determines the communication path. - - mode="inventory": Manager-to-manager path for complete, unfiltered inventory - - mode="protocol": Server-to-server path for filtered MCP requests + - via_server=False: Manager-to-manager path for complete, unfiltered inventory + - via_server=True: Server-to-server path for filtered MCP requests """ all_templates: dict[str, ResourceTemplate] = {} for mounted in self._mounted_sources: try: - if mode == "protocol": - # PATH 2: Use the server-to-server filtered path + if via_server: + # Use the server-to-server filtered path child_templates = await mounted.server._list_resource_templates() - else: # mode == "inventory" - # PATH 1: Use the manager-to-manager unfiltered path + else: + # Use the manager-to-manager unfiltered path child_templates = await mounted.server._resource_manager._list_resource_templates() child_dict = {template.key: template for template in child_templates} @@ -169,14 +167,14 @@ class ResourceManager: """ Lists all resources, applying protocol filtering. """ - resources_dict = await self._load_resources(mode="protocol") + resources_dict = await self._load_resources(via_server=True) return list(resources_dict.values()) async def _list_resource_templates(self) -> list[ResourceTemplate]: """ Lists all templates, applying protocol filtering. """ - templates_dict = await self._load_resource_templates(mode="protocol") + templates_dict = await self._load_resource_templates(via_server=True) return list(templates_dict.values()) def add_resource_or_template_from_fn( diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index d067f72d7..e8fed9def 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations import warnings from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any from mcp.types import ToolAnnotations @@ -47,25 +47,23 @@ class ToolManager: """Adds a mounted server as a source for tools.""" self._mounted_sources.append(server) - async def _load_tools( - self, *, mode: Literal["inventory", "protocol"] - ) -> dict[str, Tool]: + async def _load_tools(self, *, via_server: bool = False) -> dict[str, Tool]: """ - The single, consolidated recursive method for fetching tools. The 'mode' + The single, consolidated recursive method for fetching tools. The 'via_server' parameter determines the communication path. - - mode="inventory": Manager-to-manager path for complete, unfiltered inventory - - mode="protocol": Server-to-server path for filtered MCP requests + - via_server=False: Manager-to-manager path for complete, unfiltered inventory + - via_server=True: Server-to-server path for filtered MCP requests """ all_tools: dict[str, Tool] = {} for mounted in self._mounted_sources: try: - if mode == "protocol": - # PATH 2: Use the server-to-server filtered path + if via_server: + # Use the server-to-server filtered path child_results = await mounted.server._list_tools() - else: # mode == "inventory" - # PATH 1: Use the manager-to-manager unfiltered path + else: + # Use the manager-to-manager unfiltered path child_results = await mounted.server._tool_manager._list_tools() # The combination logic is the same for both paths @@ -103,13 +101,13 @@ class ToolManager: """ Gets the complete, unfiltered inventory of all tools. """ - return await self._load_tools(mode="inventory") + return await self._load_tools(via_server=False) async def _list_tools(self) -> list[Tool]: """ Lists all tools, applying protocol filtering. """ - tools_dict = await self._load_tools(mode="protocol") + tools_dict = await self._load_tools(via_server=True) return list(tools_dict.values()) def add_tool_from_fn( From 2e31b2704b62882a63f09c1609f64298d71644e8 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 14:05:19 -0400 Subject: [PATCH 11/17] Update tests --- docs/servers/middleware.mdx | 510 +++++++++++++++++++ src/fastmcp/prompts/prompt_manager.py | 22 +- src/fastmcp/resources/resource_manager.py | 31 +- src/fastmcp/resources/template.py | 3 + src/fastmcp/tools/tool_manager.py | 22 +- tests/prompts/test_prompt_manager.py | 47 +- tests/resources/test_resource_manager.py | 110 ++-- tests/resources/test_resource_template.py | 1 + tests/server/openapi/test_openapi.py | 350 +++++++------ tests/server/test_mount.py | 5 +- tests/server/test_resource_prefix_formats.py | 4 +- tests/server/test_tool_annotations.py | 9 +- tests/server/test_tool_exclude_args.py | 8 +- 13 files changed, 843 insertions(+), 279 deletions(-) create mode 100644 docs/servers/middleware.mdx diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx new file mode 100644 index 000000000..690026e5f --- /dev/null +++ b/docs/servers/middleware.mdx @@ -0,0 +1,510 @@ +--- +title: MCP Middleware +sidebarTitle: Middleware +description: Add cross-cutting functionality to your MCP server with middleware that can inspect, modify, and respond to all MCP requests and responses. +icon: layers +--- + +import { VersionBadge } from "/snippets/version-badge.mdx" + + + +MCP middleware is a powerful concept that allows you to add cross-cutting functionality to your FastMCP server. Unlike traditional web middleware, MCP middleware is designed specifically for the Model Context Protocol, providing hooks for different types of MCP operations like tool calls, resource reads, and prompt requests. + + +MCP middleware is a FastMCP-specific concept and is not part of the official MCP protocol specification. This middleware system is designed to work with FastMCP servers and may not be compatible with other MCP implementations. + + +## What is MCP Middleware? + +MCP middleware lets you intercept and modify MCP requests and responses as they flow through your server. Unlike traditional HTTP middleware that operates on request/response pairs, MCP middleware is aware of the specific MCP protocol operations and can provide targeted hooks for different types of interactions. + +Common use cases for MCP middleware include: +- **Authentication and Authorization**: Verify client permissions before executing operations +- **Logging and Monitoring**: Track usage patterns and performance metrics +- **Rate Limiting**: Control request frequency per client or operation type +- **Request/Response Transformation**: Modify data before it reaches tools or after it leaves +- **Caching**: Store frequently requested data to improve performance +- **Error Handling**: Provide consistent error responses across your server + +## How MCP Middleware Works + +MCP middleware operates on a pipeline model where each middleware can: + +1. **Inspect the incoming request** and its context +2. **Modify the request** before passing it to the next middleware or handler +3. **Execute the next middleware/handler** in the chain +4. **Inspect and modify the response** before returning it +5. **Handle errors** that occur during processing + +The middleware system provides specialized hooks for different MCP operations: + +- `on_message`: Called for all MCP messages (requests and notifications) +- `on_request`: Called specifically for MCP requests (that expect responses) +- `on_notification`: Called specifically for MCP notifications (fire-and-forget) +- `on_call_tool`: Called when tools are being executed +- `on_read_resource`: Called when resources are being read +- `on_get_prompt`: Called when prompts are being retrieved +- `on_list_tools`: Called when listing available tools +- `on_list_resources`: Called when listing available resources +- `on_list_resource_templates`: Called when listing resource templates +- `on_list_prompts`: Called when listing available prompts + + +The middleware hook system is designed to be extensible. As FastMCP evolves and new MCP operations are added, additional hooks will be introduced to provide fine-grained control over new functionality. + + +## Creating Middleware + +### Basic Middleware Structure + +MCP middleware is implemented by subclassing the `MCPMiddleware` base class and overriding the hooks you need: + +```python +from fastmcp import FastMCP +from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext + +class LoggingMiddleware(MCPMiddleware): + """Middleware that logs all MCP operations.""" + + async def on_message(self, context: MiddlewareContext, call_next): + """Called for all MCP messages.""" + print(f"Processing {context.method} from {context.source}") + + # Call the next middleware/handler in the chain + result = await call_next(context) + + print(f"Completed {context.method}") + return result + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Called specifically for tool calls.""" + tool_name = context.message.name + print(f"Calling tool: {tool_name}") + + result = await call_next(context) + + print(f"Tool {tool_name} completed") + return result + +# Add middleware to your server +mcp = FastMCP("MyServer") +mcp.add_middleware(LoggingMiddleware()) +``` + +### Middleware Context + +The `MiddlewareContext` object provides access to information about the current request: + +```python +class InspectionMiddleware(MCPMiddleware): + async def on_request(self, context: MiddlewareContext, call_next): + # Access request information + method = context.method # e.g., "tools/call" + source = context.source # "client" or "server" + message_type = context.type # "request" or "notification" + timestamp = context.timestamp # When the request was received + message = context.message # The actual MCP message + fastmcp_context = context.fastmcp_context # FastMCP Context object (if available) + + # Continue processing + return await call_next(context) +``` + +### Middleware Hooks + +Each middleware hook receives a `MiddlewareContext` and a `call_next` function. The hooks are organized in a hierarchy: + +1. **`on_message`**: The broadest hook, called for all MCP messages +2. **`on_request`** / **`on_notification`**: Called based on message type +3. **Operation-specific hooks**: Called for specific MCP operations + +```python +class ComprehensiveMiddleware(MCPMiddleware): + async def on_message(self, context: MiddlewareContext, call_next): + """Called for ALL messages (requests and notifications).""" + print(f"Message: {context.method}") + return await call_next(context) + + async def on_request(self, context: MiddlewareContext, call_next): + """Called only for requests (messages that expect responses).""" + print(f"Request: {context.method}") + return await call_next(context) + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Called only for tool execution requests.""" + tool_name = context.message.name + print(f"Executing tool: {tool_name}") + return await call_next(context) +``` + +## Middleware Examples + +### Authentication Middleware + +```python +from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.exceptions import ToolError + +class AuthenticationMiddleware(MCPMiddleware): + def __init__(self, required_token: str): + self.required_token = required_token + + async def on_request(self, context: MiddlewareContext, call_next): + """Verify authentication for all requests.""" + + # Check if this is an HTTP request with headers + if hasattr(context, 'fastmcp_context') and context.fastmcp_context: + try: + # Access HTTP request if available + request = context.fastmcp_context.get_http_request() + auth_header = request.headers.get("Authorization") + + if not auth_header or not auth_header.startswith("Bearer "): + raise ToolError("Missing or invalid authorization header") + + token = auth_header.split(" ", 1)[1] + if token != self.required_token: + raise ToolError("Invalid authentication token") + + except Exception: + # If HTTP request is not available, continue without auth + pass + + return await call_next(context) + +# Usage +mcp = FastMCP("SecureServer") +mcp.add_middleware(AuthenticationMiddleware("secret-token-123")) +``` + +### Performance Monitoring Middleware + +```python +import time +import logging + +class PerformanceMiddleware(MCPMiddleware): + def __init__(self): + self.logger = logging.getLogger("performance") + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Monitor tool execution performance.""" + tool_name = context.message.name + start_time = time.time() + + try: + result = await call_next(context) + execution_time = time.time() - start_time + + self.logger.info( + f"Tool {tool_name} completed in {execution_time:.3f}s" + ) + + return result + + except Exception as e: + execution_time = time.time() - start_time + self.logger.error( + f"Tool {tool_name} failed after {execution_time:.3f}s: {e}" + ) + raise +``` + +### Request/Response Transformation Middleware + +```python +class TransformationMiddleware(MCPMiddleware): + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Transform tool arguments and results.""" + + # Access and modify tool arguments + if hasattr(context.message, 'arguments'): + args = context.message.arguments or {} + + # Example: Add a timestamp to all tool calls + args['_middleware_timestamp'] = context.timestamp.isoformat() + + # Create a modified context + modified_context = context.copy( + message=context.message.model_copy(update={'arguments': args}) + ) + else: + modified_context = context + + # Execute with modified context + result = await call_next(modified_context) + + # Transform the result if needed + if hasattr(result, 'content'): + # Example: Add metadata to tool results + if result.content and len(result.content) > 0: + original_content = result.content[0].text + enhanced_content = f"[Processed at {context.timestamp}]\n{original_content}" + result.content[0].text = enhanced_content + + return result +``` + +### Rate Limiting Middleware + +```python +import asyncio +from collections import defaultdict +from datetime import datetime, timedelta + +class RateLimitMiddleware(MCPMiddleware): + def __init__(self, max_requests: int = 100, window_minutes: int = 1): + self.max_requests = max_requests + self.window = timedelta(minutes=window_minutes) + self.requests = defaultdict(list) # client_id -> [timestamps] + + async def on_request(self, context: MiddlewareContext, call_next): + """Implement rate limiting per client.""" + + # Get client identifier (you may need to implement this based on your auth) + client_id = getattr(context.fastmcp_context, 'client_id', 'anonymous') + + now = datetime.now() + + # Clean old requests + self.requests[client_id] = [ + timestamp for timestamp in self.requests[client_id] + if now - timestamp < self.window + ] + + # Check rate limit + if len(self.requests[client_id]) >= self.max_requests: + raise ToolError( + f"Rate limit exceeded: {self.max_requests} requests per " + f"{self.window.total_seconds()/60:.0f} minutes" + ) + + # Record this request + self.requests[client_id].append(now) + + return await call_next(context) +``` + +## Adding Middleware to Your Server + +### Single Middleware + +```python +from fastmcp import FastMCP + +mcp = FastMCP("MyServer") + +# Add a single middleware instance +logging_middleware = LoggingMiddleware() +mcp.add_middleware(logging_middleware) +``` + +### Multiple Middleware + +Middleware is executed in the order it's added to the server: + +```python +mcp = FastMCP("MyServer") + +# Add multiple middleware - they execute in order +mcp.add_middleware(AuthenticationMiddleware("secret-token")) +mcp.add_middleware(PerformanceMiddleware()) +mcp.add_middleware(LoggingMiddleware()) + +# Request flow: +# 1. AuthenticationMiddleware.on_request() +# 2. PerformanceMiddleware.on_request() +# 3. LoggingMiddleware.on_request() +# 4. Actual tool/resource handler +# 5. LoggingMiddleware response processing +# 6. PerformanceMiddleware response processing +# 7. AuthenticationMiddleware response processing +``` + +## Advanced Patterns + +### Conditional Middleware + +```python +class ConditionalMiddleware(MCPMiddleware): + def __init__(self, condition_func): + self.should_process = condition_func + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Only process certain tools.""" + + if not self.should_process(context.message.name): + # Skip processing for this tool + return await call_next(context) + + # Apply middleware logic + print(f"Processing tool: {context.message.name}") + return await call_next(context) + +# Usage +def only_expensive_tools(tool_name: str) -> bool: + return tool_name in ["complex_analysis", "heavy_computation"] + +mcp.add_middleware(ConditionalMiddleware(only_expensive_tools)) +``` + +### Middleware with State + +```python +class StatefulMiddleware(MCPMiddleware): + def __init__(self): + self.call_count = 0 + self.tools_used = set() + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Track usage statistics.""" + self.call_count += 1 + self.tools_used.add(context.message.name) + + print(f"Total calls: {self.call_count}, Unique tools: {len(self.tools_used)}") + + return await call_next(context) + + def get_stats(self): + return { + "total_calls": self.call_count, + "unique_tools": len(self.tools_used), + "tools_used": list(self.tools_used) + } +``` + +### Error Handling Middleware + +```python +class ErrorHandlingMiddleware(MCPMiddleware): + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Provide consistent error handling.""" + try: + return await call_next(context) + except ToolError: + # Re-raise ToolErrors as-is + raise + except Exception as e: + # Log the error and convert to a user-friendly message + logging.error(f"Tool {context.message.name} failed: {e}") + raise ToolError(f"Tool execution failed: {str(e)}") +``` + +## Best Practices + +### Performance Considerations + +1. **Keep middleware lightweight**: Avoid heavy computations in middleware +2. **Use async operations**: Don't block the event loop with synchronous operations +3. **Cache when possible**: Store frequently accessed data to avoid repeated lookups + +```python +class EfficientMiddleware(MCPMiddleware): + def __init__(self): + self._cache = {} + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Example of efficient middleware with caching.""" + + # Check cache first + cache_key = f"{context.message.name}:{hash(str(context.message.arguments))}" + + if cache_key in self._cache: + print("Returning cached result") + return self._cache[cache_key] + + # Execute and cache result + result = await call_next(context) + self._cache[cache_key] = result + + return result +``` + +### Error Handling + +1. **Always call `call_next`**: Unless you're intentionally stopping the chain +2. **Handle exceptions appropriately**: Don't let middleware errors break the entire request +3. **Use `ToolError` for client-facing errors**: Keep internal errors internal + +```python +class RobustMiddleware(MCPMiddleware): + async def on_request(self, context: MiddlewareContext, call_next): + """Robust error handling example.""" + try: + # Middleware logic here + return await call_next(context) + except ToolError: + # Client-facing errors should be re-raised + raise + except Exception as e: + # Log internal errors but don't expose details + logging.error(f"Middleware error: {e}") + # Optionally continue without middleware processing + return await call_next(context) +``` + +### Testing Middleware + +```python +import pytest +from fastmcp import FastMCP, Client + +@pytest.mark.asyncio +async def test_logging_middleware(): + """Test middleware functionality.""" + + # Create server with middleware + mcp = FastMCP("TestServer") + logging_middleware = LoggingMiddleware() + mcp.add_middleware(logging_middleware) + + @mcp.tool + def test_tool(x: int) -> int: + return x * 2 + + # Test with client + async with Client(mcp) as client: + result = await client.call_tool("test_tool", {"x": 5}) + assert result == 10 + + # Verify middleware was called + # (You'll need to add tracking to your middleware for testing) +``` + +## Server Composition and Middleware + +When using [Server Composition](/servers/composition) with `mount` or `import_server`, middleware behavior follows these rules: + +1. **Parent server middleware** runs for all requests, including those routed to mounted servers +2. **Mounted server middleware** only runs for requests handled by that specific server +3. **Middleware order** is preserved within each server + +```python +# Parent server with middleware +parent = FastMCP("Parent") +parent.add_middleware(AuthenticationMiddleware("token")) + +# Child server with its own middleware +child = FastMCP("Child") +child.add_middleware(LoggingMiddleware()) + +@child.tool +def child_tool() -> str: + return "from child" + +# Mount the child server +parent.mount(child, prefix="child") + +# Request to "child_tool" will: +# 1. Run parent's AuthenticationMiddleware +# 2. Route to child server +# 3. Run child's LoggingMiddleware +# 4. Execute child_tool +``` + +This allows you to create layered middleware architectures where parent servers handle cross-cutting concerns like authentication, while child servers focus on domain-specific middleware. + + +When debugging middleware issues in composed servers, remember that both parent and child middleware may be executing. Use detailed logging to trace the middleware execution path. + \ No newline at end of file diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index a4aaaf1ae..b237dc0fc 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -27,7 +27,7 @@ class PromptManager: mask_error_details: bool | None = None, ): self._prompts: dict[str, Prompt] = {} - self._mounted_sources: list[MountedServer] = [] + self._mounted_servers: list[MountedServer] = [] self.mask_error_details = mask_error_details or settings.mask_error_details # Default to "warn" if None is provided @@ -44,7 +44,7 @@ class PromptManager: def mount(self, server: MountedServer) -> None: """Adds a mounted server as a source for prompts.""" - self._mounted_sources.append(server) + self._mounted_servers.append(server) async def _load_prompts(self, *, via_server: bool = False) -> dict[str, Prompt]: """ @@ -56,7 +56,7 @@ class PromptManager: """ all_prompts: dict[str, Prompt] = {} - for mounted in self._mounted_sources: + for mounted in self._mounted_servers: try: if via_server: # Use the server-to-server filtered path @@ -188,12 +188,16 @@ class PromptManager: raise PromptError(f"Error rendering prompt {name!r}: {e}") from e # 2. Check mounted servers using the filtered protocol path. - for mounted in reversed(self._mounted_sources): - if mounted.prefix and name.startswith(f"{mounted.prefix}_"): - name_on_child = name.removeprefix(f"{mounted.prefix}_") - try: - return await mounted.server._get_prompt(name_on_child, arguments) - except NotFoundError: + for mounted in reversed(self._mounted_servers): + prompt_key = name + if mounted.prefix: + if name.startswith(f"{mounted.prefix}_"): + prompt_key = name.removeprefix(f"{mounted.prefix}_") + else: continue + try: + return await mounted.server._get_prompt(prompt_key, arguments) + except NotFoundError: + continue raise NotFoundError(f"Unknown prompt: {name}") diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index a22e515dc..4bc084181 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -43,7 +43,7 @@ class ResourceManager: """ self._resources: dict[str, Resource] = {} self._templates: dict[str, ResourceTemplate] = {} - self._mounted_sources: list[MountedServer] = [] + self._mounted_servers: list[MountedServer] = [] self.mask_error_details = mask_error_details or settings.mask_error_details # Default to "warn" if None is provided @@ -59,7 +59,7 @@ class ResourceManager: def mount(self, server: MountedServer) -> None: """Adds a mounted server as a source for resources and templates.""" - self._mounted_sources.append(server) + self._mounted_servers.append(server) async def get_resources(self) -> dict[str, Resource]: """Get all registered resources, keyed by URI.""" @@ -79,7 +79,7 @@ class ResourceManager: """ all_resources: dict[str, Resource] = {} - for mounted in self._mounted_sources: + for mounted in self._mounted_servers: try: if via_server: # Use the server-to-server filtered path @@ -129,7 +129,7 @@ class ResourceManager: """ all_templates: dict[str, ResourceTemplate] = {} - for mounted in self._mounted_sources: + for mounted in self._mounted_servers: try: if via_server: # Use the server-to-server filtered path @@ -457,8 +457,8 @@ class ResourceManager: ) from e # 1b. Check local templates if not found in concrete resources - for template in self._templates.values(): - if params := match_uri_template(uri_str, template.uri_template): + for key, template in self._templates.items(): + if params := match_uri_template(uri_str, key): try: resource = await template.create_resource(uri_str, params=params) return await resource.read() @@ -483,29 +483,28 @@ class ResourceManager: # 2. Check mounted servers using the filtered protocol path. from fastmcp.server.server import has_resource_prefix, remove_resource_prefix - for mounted in reversed(self._mounted_sources): - resource_uri = uri_str + for mounted in reversed(self._mounted_servers): + key = uri_str try: if mounted.prefix: - # If server has a prefix, check if URI matches and strip prefix if has_resource_prefix( - resource_uri, + key, mounted.prefix, mounted.resource_prefix_format, ): - resource_uri = remove_resource_prefix( - resource_uri, + key = remove_resource_prefix( + key, mounted.prefix, mounted.resource_prefix_format, ) else: continue - result = await mounted.server._read_resource(resource_uri) - # Extract content from the first ReadResourceContents - if result and len(result) > 0: + try: + result = await mounted.server._read_resource(key) return result[0].content - raise NotFoundError(f"Resource {uri_str!r} returned empty content") + except NotFoundError: + continue except NotFoundError: continue diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index 728a12764..8ea97c73f 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -62,6 +62,9 @@ class ResourceTemplate(FastMCPComponent): description="JSON schema for function parameters" ) + def __repr__(self) -> str: + return f"{self.__class__.__name__}(uri_template={self.uri_template!r}, name={self.name!r}, description={self.description!r}, tags={self.tags})" + @staticmethod def from_function( fn: Callable[..., Any], diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index e8fed9def..438480882 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -28,7 +28,7 @@ class ToolManager: mask_error_details: bool | None = None, ): self._tools: dict[str, Tool] = {} - self._mounted_sources: list[MountedServer] = [] + self._mounted_servers: list[MountedServer] = [] self.mask_error_details = mask_error_details or settings.mask_error_details # Default to "warn" if None is provided @@ -45,7 +45,7 @@ class ToolManager: def mount(self, server: MountedServer) -> None: """Adds a mounted server as a source for tools.""" - self._mounted_sources.append(server) + self._mounted_servers.append(server) async def _load_tools(self, *, via_server: bool = False) -> dict[str, Tool]: """ @@ -57,7 +57,7 @@ class ToolManager: """ all_tools: dict[str, Tool] = {} - for mounted in self._mounted_sources: + for mounted in self._mounted_servers: try: if via_server: # Use the server-to-server filtered path @@ -201,12 +201,16 @@ class ToolManager: raise ToolError(f"Error calling tool {key!r}: {e}") from e # 2. Check mounted servers using the filtered protocol path. - for mounted in reversed(self._mounted_sources): - if mounted.prefix and key.startswith(f"{mounted.prefix}_"): - key_on_child = key.removeprefix(f"{mounted.prefix}_") - try: - return await mounted.server._call_tool(key_on_child, arguments) - except NotFoundError: + for mounted in reversed(self._mounted_servers): + tool_key = key + if mounted.prefix: + if key.startswith(f"{mounted.prefix}_"): + tool_key = key.removeprefix(f"{mounted.prefix}_") + else: continue + try: + return await mounted.server._call_tool(tool_key, arguments) + except NotFoundError: + continue raise NotFoundError(f"Tool {key!r} not found.") diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index c33751061..d69789972 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/tests/prompts/test_prompt_manager.py @@ -10,7 +10,7 @@ from fastmcp.prompts.prompt_manager import PromptManager class TestPromptManager: - def test_add_prompt(self): + async def test_add_prompt(self): """Test adding a prompt to the manager.""" def fn() -> str: @@ -20,9 +20,9 @@ class TestPromptManager: prompt = Prompt.from_function(fn) added = manager.add_prompt(prompt) assert added == prompt - assert manager.get_prompt("fn") == prompt + assert await manager.get_prompt("fn") == prompt - def test_add_duplicate_prompt(self, caplog): + async def test_add_duplicate_prompt(self, caplog): """Test adding the same prompt twice.""" def fn() -> str: @@ -35,7 +35,7 @@ class TestPromptManager: assert first == second assert "Prompt already exists" in caplog.text - def test_disable_warn_on_duplicate_prompts(self, caplog): + async def test_disable_warn_on_duplicate_prompts(self, caplog): """Test disabling warning on duplicate prompts.""" def fn() -> str: @@ -48,7 +48,7 @@ class TestPromptManager: assert first == second assert "Prompt already exists" not in caplog.text - def test_warn_on_duplicate_prompts(self, caplog): + async def test_warn_on_duplicate_prompts(self, caplog): """Test warning on duplicate prompts.""" manager = PromptManager(duplicate_behavior="warn") @@ -62,9 +62,9 @@ class TestPromptManager: assert "Prompt already exists: test_prompt" in caplog.text # Should have the prompt - assert manager.get_prompt("test_prompt") is not None + assert await manager.get_prompt("test_prompt") is not None - def test_error_on_duplicate_prompts(self): + async def test_error_on_duplicate_prompts(self): """Test error on duplicate prompts.""" manager = PromptManager(duplicate_behavior="error") @@ -78,7 +78,7 @@ class TestPromptManager: with pytest.raises(ValueError, match="Prompt already exists: test_prompt"): manager.add_prompt(prompt) - def test_replace_duplicate_prompts(self): + async def test_replace_duplicate_prompts(self): """Test replacing duplicate prompts.""" manager = PromptManager(duplicate_behavior="replace") @@ -95,12 +95,12 @@ class TestPromptManager: manager.add_prompt(prompt2) # Should have replaced with the new prompt - prompt = manager.get_prompt("test_prompt") + prompt = await manager.get_prompt("test_prompt") assert prompt is not None assert isinstance(prompt, FunctionPrompt) assert prompt.fn.__name__ == "replacement_fn" - def test_ignore_duplicate_prompts(self): + async def test_ignore_duplicate_prompts(self): """Test ignoring duplicate prompts.""" manager = PromptManager(duplicate_behavior="ignore") @@ -117,7 +117,7 @@ class TestPromptManager: result = manager.add_prompt(prompt2) # Should keep the original - prompt = manager.get_prompt("test_prompt") + prompt = await manager.get_prompt("test_prompt") assert prompt is not None assert isinstance(prompt, FunctionPrompt) assert prompt.fn.__name__ == "original_fn" @@ -125,7 +125,7 @@ class TestPromptManager: assert isinstance(result, FunctionPrompt) assert result.fn.__name__ == "original_fn" - def test_get_prompts(self): + async def test_get_prompts(self): """Test retrieving all prompts.""" def fn1() -> str: @@ -139,7 +139,7 @@ class TestPromptManager: prompt2 = Prompt.from_function(fn2) manager.add_prompt(prompt1) manager.add_prompt(prompt2) - prompts = manager.get_prompts() + prompts = await manager.get_prompts() assert len(prompts) == 2 assert prompts["fn1"] == prompt1 assert prompts["fn2"] == prompt2 @@ -270,7 +270,7 @@ class TestRenderPrompt: class TestPromptTags: """Test functionality related to prompt tags.""" - def test_add_prompt_with_tags(self): + async def test_add_prompt_with_tags(self): """Test adding a prompt with tags.""" def greeting() -> str: @@ -280,11 +280,11 @@ class TestPromptTags: prompt = Prompt.from_function(greeting, tags={"greeting", "simple"}) manager.add_prompt(prompt) - prompt = manager.get_prompt("greeting") + prompt = await manager.get_prompt("greeting") assert prompt is not None assert prompt.tags == {"greeting", "simple"} - def test_add_prompt_with_empty_tags(self): + async def test_add_prompt_with_empty_tags(self): """Test adding a prompt with empty tags.""" def greeting() -> str: @@ -294,11 +294,11 @@ class TestPromptTags: prompt = Prompt.from_function(greeting, tags=set()) manager.add_prompt(prompt) - prompt = manager.get_prompt("greeting") + prompt = await manager.get_prompt("greeting") assert prompt is not None assert prompt.tags == set() - def test_add_prompt_with_none_tags(self): + async def test_add_prompt_with_none_tags(self): """Test adding a prompt with None tags.""" def greeting() -> str: @@ -308,11 +308,11 @@ class TestPromptTags: prompt = Prompt.from_function(greeting, tags=None) manager.add_prompt(prompt) - prompt = manager.get_prompt("greeting") + prompt = await manager.get_prompt("greeting") assert prompt is not None assert prompt.tags == set() - def test_list_prompts_with_tags(self): + async def test_list_prompts_with_tags(self): """Test listing prompts with specific tags.""" def greeting() -> str: @@ -332,13 +332,12 @@ class TestPromptTags: ) # Filter prompts by tags - simple_prompts = [ - p for p in manager.get_prompts().values() if "simple" in p.tags - ] + prompts = await manager.get_prompts() + simple_prompts = [p for p in prompts.values() if "simple" in p.tags] assert len(simple_prompts) == 2 assert {p.name for p in simple_prompts} == {"greeting", "summary"} - nlp_prompts = [p for p in manager.get_prompts().values() if "nlp" in p.tags] + nlp_prompts = [p for p in prompts.values() if "nlp" in p.tags] assert len(nlp_prompts) == 1 assert nlp_prompts[0].name == "summary" diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index a64fded09..7c1f02066 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -33,7 +33,7 @@ def temp_file(): class TestResourceManager: """Test ResourceManager functionality.""" - def test_add_resource(self, temp_file: Path): + async def test_add_resource(self, temp_file: Path): """Test adding a resource.""" manager = ResourceManager() file_url = "file://test-resource" @@ -45,10 +45,11 @@ class TestResourceManager: added = manager.add_resource(resource) assert added == resource # Get the actual key from the resource manager - assert len(manager.get_resources()) == 1 - assert resource in manager.get_resources().values() + resources = await manager.get_resources() + assert len(resources) == 1 + assert resource in resources.values() - def test_add_duplicate_resource(self, temp_file: Path): + async def test_add_duplicate_resource(self, temp_file: Path): """Test adding the same resource twice.""" manager = ResourceManager() file_url = "file://test-resource" @@ -61,10 +62,11 @@ class TestResourceManager: second = manager.add_resource(resource) assert first == second # Check the resource is there - assert len(manager.get_resources()) == 1 - assert resource in manager.get_resources().values() + resources = await manager.get_resources() + assert len(resources) == 1 + assert resource in resources.values() - def test_warn_on_duplicate_resources(self, temp_file: Path, caplog): + async def test_warn_on_duplicate_resources(self, temp_file: Path, caplog): """Test warning on duplicate resources.""" manager = ResourceManager(duplicate_behavior="warn") @@ -80,10 +82,11 @@ class TestResourceManager: assert "Resource already exists" in caplog.text # Should have the resource - assert len(manager.get_resources()) == 1 - assert resource in manager.get_resources().values() + resources = await manager.get_resources() + assert len(resources) == 1 + assert resource in resources.values() - def test_disable_warn_on_duplicate_resources(self, temp_file: Path, caplog): + async def test_disable_warn_on_duplicate_resources(self, temp_file: Path, caplog): """Test disabling warning on duplicate resources.""" manager = ResourceManager(duplicate_behavior="ignore") resource = FileResource( @@ -95,7 +98,7 @@ class TestResourceManager: manager.add_resource(resource) assert "Resource already exists" not in caplog.text - def test_error_on_duplicate_resources(self, temp_file: Path): + async def test_error_on_duplicate_resources(self, temp_file: Path): """Test error on duplicate resources.""" manager = ResourceManager(duplicate_behavior="error") @@ -110,7 +113,7 @@ class TestResourceManager: with pytest.raises(ValueError, match="Resource already exists"): manager.add_resource(resource) - def test_replace_duplicate_resources(self, temp_file: Path): + async def test_replace_duplicate_resources(self, temp_file: Path): """Test replacing duplicate resources.""" manager = ResourceManager(duplicate_behavior="replace") @@ -131,11 +134,12 @@ class TestResourceManager: manager.add_resource(resource2) # Should have replaced with the new resource - resources = list(manager.get_resources().values()) - assert len(resources) == 1 - assert resources[0].name == "replacement" + resources = await manager.get_resources() + resource_list = list(resources.values()) + assert len(resource_list) == 1 + assert resource_list[0].name == "replacement" - def test_ignore_duplicate_resources(self, temp_file: Path): + async def test_ignore_duplicate_resources(self, temp_file: Path): """Test ignoring duplicate resources.""" manager = ResourceManager(duplicate_behavior="ignore") @@ -156,13 +160,14 @@ class TestResourceManager: result = manager.add_resource(resource2) # Should keep the original - resources = list(manager.get_resources().values()) - assert len(resources) == 1 - assert resources[0].name == "original" + resources = await manager.get_resources() + resource_list = list(resources.values()) + assert len(resource_list) == 1 + assert resource_list[0].name == "original" # Result should be the original resource assert result.name == "original" - def test_warn_on_duplicate_templates(self, caplog): + async def test_warn_on_duplicate_templates(self, caplog): """Test warning on duplicate templates.""" manager = ResourceManager(duplicate_behavior="warn") @@ -180,9 +185,10 @@ class TestResourceManager: assert "Template already exists" in caplog.text # Should have the template - assert manager.get_resource_templates() == {"test://{id}": template} + templates = await manager.get_resource_templates() + assert templates == {"test://{id}": template} - def test_error_on_duplicate_templates(self): + async def test_error_on_duplicate_templates(self): """Test error on duplicate templates.""" manager = ResourceManager(duplicate_behavior="error") @@ -200,7 +206,7 @@ class TestResourceManager: with pytest.raises(ValueError, match="Template already exists"): manager.add_template(template) - def test_replace_duplicate_templates(self): + async def test_replace_duplicate_templates(self): """Test replacing duplicate templates.""" manager = ResourceManager(duplicate_behavior="replace") @@ -226,11 +232,12 @@ class TestResourceManager: manager.add_template(template2) # Should have replaced with the new template - templates = list(manager.get_resource_templates().values()) + templates_dict = await manager.get_resource_templates() + templates = list(templates_dict.values()) assert len(templates) == 1 assert templates[0].name == "replacement" - def test_ignore_duplicate_templates(self): + async def test_ignore_duplicate_templates(self): """Test ignoring duplicate templates.""" manager = ResourceManager(duplicate_behavior="ignore") @@ -256,7 +263,8 @@ class TestResourceManager: result = manager.add_template(template2) # Should keep the original - templates = list(manager.get_resource_templates().values()) + templates_dict = await manager.get_resource_templates() + templates = list(templates_dict.values()) assert len(templates) == 1 assert templates[0].name == "original" # Result should be the original template @@ -299,7 +307,7 @@ class TestResourceManager: with pytest.raises(NotFoundError, match="Unknown resource"): await manager.get_resource(AnyUrl("unknown://test")) - def test_get_resources(self, temp_file: Path): + async def test_get_resources(self, temp_file: Path): """Test retrieving all resources.""" manager = ResourceManager() file_url1 = "file://test-resource1" @@ -316,7 +324,7 @@ class TestResourceManager: ) manager.add_resource(resource1) manager.add_resource(resource2) - resources = manager.get_resources() + resources = await manager.get_resources() assert len(resources) == 2 values = list(resources.values()) assert resource1 in values @@ -326,7 +334,7 @@ class TestResourceManager: class TestResourceTags: """Test functionality related to resource tags.""" - def test_add_resource_with_tags(self, temp_file: Path): + async def test_add_resource_with_tags(self, temp_file: Path): """Test adding a resource with tags.""" manager = ResourceManager() resource = FileResource( @@ -338,11 +346,12 @@ class TestResourceTags: manager.add_resource(resource) # Check that tags are preserved - resources = list(manager.get_resources().values()) + resources_dict = await manager.get_resources() + resources = list(resources_dict.values()) assert len(resources) == 1 assert resources[0].tags == {"weather", "data"} - def test_add_function_resource_with_tags(self): + async def test_add_function_resource_with_tags(self): """Test adding a function resource with tags.""" manager = ResourceManager() @@ -359,11 +368,12 @@ class TestResourceTags: ) manager.add_resource(resource) - resources = list(manager.get_resources().values()) + resources_dict = await manager.get_resources() + resources = list(resources_dict.values()) assert len(resources) == 1 assert resources[0].tags == {"sample", "test", "data"} - def test_add_template_with_tags(self): + async def test_add_template_with_tags(self): """Test adding a resource template with tags.""" manager = ResourceManager() @@ -379,11 +389,12 @@ class TestResourceTags: ) manager.add_template(template) - templates = list(manager.get_resource_templates().values()) + templates_dict = await manager.get_resource_templates() + templates = list(templates_dict.values()) assert len(templates) == 1 assert templates[0].tags == {"users", "template", "data"} - def test_filter_resources_by_tags(self, temp_file: Path): + async def test_filter_resources_by_tags(self, temp_file: Path): """Test filtering resources by tags.""" manager = ResourceManager() @@ -392,7 +403,7 @@ class TestResourceTags: uri=FileUrl("file://weather-data"), name="weather_data", path=temp_file, - tags={"weather", "external"}, + tags={"weather", "data"}, ) async def get_user_data(): @@ -401,8 +412,10 @@ class TestResourceTags: resource2 = FunctionResource( uri=AnyUrl("data://users"), name="user_data", + description="User data resource", + mime_type="text/plain", fn=get_user_data, - tags={"users", "internal"}, + tags={"users", "data"}, ) async def get_system_data(): @@ -411,26 +424,25 @@ class TestResourceTags: resource3 = FunctionResource( uri=AnyUrl("data://system"), name="system_data", + description="System data resource", + mime_type="text/plain", fn=get_system_data, - tags={"system", "internal"}, + tags={"system", "admin"}, ) manager.add_resource(resource1) manager.add_resource(resource2) manager.add_resource(resource3) - # Filter resources by tags - internal_resources = [ - r for r in manager.get_resources().values() if "internal" in r.tags - ] - assert len(internal_resources) == 2 - assert {r.name for r in internal_resources} == {"user_data", "system_data"} + # Filter by tags + resources_dict = await manager.get_resources() + data_resources = [r for r in resources_dict.values() if "data" in r.tags] + assert len(data_resources) == 2 + assert {r.name for r in data_resources} == {"weather_data", "user_data"} - external_resources = [ - r for r in manager.get_resources().values() if "external" in r.tags - ] - assert len(external_resources) == 1 - assert external_resources[0].name == "weather_data" + admin_resources = [r for r in resources_dict.values() if "admin" in r.tags] + assert len(admin_resources) == 1 + assert admin_resources[0].name == "system_data" class TestCustomResourceKeys: diff --git a/tests/resources/test_resource_template.py b/tests/resources/test_resource_template.py index efb9dea65..d8668c1ca 100644 --- a/tests/resources/test_resource_template.py +++ b/tests/resources/test_resource_template.py @@ -405,6 +405,7 @@ class TestMatchUriTemplate: ("test://a/b/c", None), ("test://a/x/b", {"x": "x"}), ("test://a/x/y/b", None), + ("test://a/1-2/b", {"x": "1-2"}), ], ) def test_match_uri_template_single_param( diff --git a/tests/server/openapi/test_openapi.py b/tests/server/openapi/test_openapi.py index 5fc6611ae..393dec27e 100644 --- a/tests/server/openapi/test_openapi.py +++ b/tests/server/openapi/test_openapi.py @@ -136,7 +136,7 @@ def api_client(fastapi_app: FastAPI) -> AsyncClient: @pytest.fixture -async def fastmcp_openapi_server_with_all_types( +async def fastmcp_openapi_server( fastapi_app: FastAPI, api_client: httpx.AsyncClient ) -> FastMCPOpenAPI: openapi_spec = fastapi_app.openapi() @@ -213,13 +213,11 @@ class TestTools: assert len(await server.get_resources()) == 0 assert len(await server.get_resource_templates()) == 0 - async def test_list_tools( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI - ): + async def test_list_tools(self, fastmcp_openapi_server: FastMCPOpenAPI): """ By default, tools exclude GET methods """ - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: tools = await client.list_tools() assert len(tools) == 2 @@ -254,13 +252,13 @@ class TestTools: async def test_call_create_user_tool( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, ): """ The tool created by the OpenAPI server should be the same as the original """ - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: tool_response = await client.call_tool( "create_user_users_post", {"name": "David", "active": False} ) @@ -274,7 +272,7 @@ class TestTools: assert len(response.json()) == 4 # Check that the user was created via MCP - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: user_response = await client.read_resource("resource://get_user_users/4") response_text = user_response[0].text # type: ignore[attr-defined] user = json.loads(response_text) @@ -282,13 +280,13 @@ class TestTools: async def test_call_update_user_name_tool( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, ): """ The tool created by the OpenAPI server should be the same as the original """ - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: tool_response = await client.call_tool( "update_user_name_users", {"user_id": 1, "name": "XYZ"}, @@ -303,7 +301,7 @@ class TestTools: assert expected_data in response.json() # Check that the user was updated via MCP - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: user_response = await client.read_resource("resource://get_user_users/1") response_text = user_response[0].text # type: ignore[attr-defined] user = json.loads(response_text) @@ -335,13 +333,11 @@ class TestTools: class TestResources: - async def test_list_resources( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI - ): + async def test_list_resources(self, fastmcp_openapi_server: FastMCPOpenAPI): """ By default, resources exclude GET methods without parameters """ - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resources = await client.list_resources() assert len(resources) == 4 assert resources[0].uri == AnyUrl("resource://get_users_users_get") @@ -349,7 +345,7 @@ class TestResources: async def test_get_resource( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, users_db: dict[int, User], ): @@ -360,7 +356,7 @@ class TestResources: json_users = TypeAdapter(list[User]).dump_python( sorted(users_db.values(), key=lambda x: x.id) ) - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resource_response = await client.read_resource( "resource://get_users_users_get" ) @@ -372,11 +368,11 @@ class TestResources: async def test_get_bytes_resource( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, ): """Test reading a resource that returns bytes.""" - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resource_response = await client.read_resource( "resource://ping_bytes_ping_bytes_get" ) @@ -385,23 +381,23 @@ class TestResources: async def test_get_str_resource( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, ): """Test reading a resource that returns a string.""" - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resource_response = await client.read_resource("resource://ping_ping_get") assert resource_response[0].text == "pong" # type: ignore[attr-defined] class TestResourceTemplates: async def test_list_resource_templates( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI + self, fastmcp_openapi_server: FastMCPOpenAPI ): """ By default, resource templates exclude GET methods without parameters """ - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resource_templates = await client.list_resource_templates() assert len(resource_templates) == 2 assert resource_templates[0].name == "get_user_users" @@ -416,7 +412,7 @@ class TestResourceTemplates: async def test_get_resource_template( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, users_db: dict[int, User], ): @@ -424,7 +420,7 @@ class TestResourceTemplates: The resource template created by the OpenAPI server should be the same as the original """ user_id = 2 - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resource_response = await client.read_resource( f"resource://get_user_users/{user_id}" ) @@ -437,7 +433,7 @@ class TestResourceTemplates: async def test_get_resource_template_multi_param( self, - fastmcp_openapi_server_with_all_types: FastMCPOpenAPI, + fastmcp_openapi_server: FastMCPOpenAPI, api_client, users_db: dict[int, User], ): @@ -446,7 +442,7 @@ class TestResourceTemplates: """ user_id = 2 is_active = True - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: resource_response = await client.read_resource( f"resource://get_user_active_state_users/{is_active}/{user_id}" ) @@ -459,13 +455,11 @@ class TestResourceTemplates: class TestPrompts: - async def test_list_prompts( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI - ): + async def test_list_prompts(self, fastmcp_openapi_server: FastMCPOpenAPI): """ By default, there are no prompts. """ - async with Client(fastmcp_openapi_server_with_all_types) as client: + async with Client(fastmcp_openapi_server) as client: prompts = await client.list_prompts() assert len(prompts) == 0 @@ -474,11 +468,11 @@ class TestTagTransfer: """Tests for transferring tags from OpenAPI routes to MCP objects.""" async def test_tags_transferred_to_tools( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI + self, fastmcp_openapi_server: FastMCPOpenAPI ): """Test that tags from OpenAPI routes are correctly transferred to Tools.""" # Get internal tools directly (not the public API which returns MCP.Content) - tools = fastmcp_openapi_server_with_all_types._tool_manager._list_tools() + tools = await fastmcp_openapi_server._tool_manager._list_tools() # Find the create_user and update_user_name tools create_user_tool = next( @@ -502,13 +496,12 @@ class TestTagTransfer: assert len(update_user_tool.tags) == 2 async def test_tags_transferred_to_resources( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI + self, fastmcp_openapi_server: FastMCPOpenAPI ): """Test that tags from OpenAPI routes are correctly transferred to Resources.""" # Get internal resources directly - resources = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_resources().values() - ) + resources_dict = await fastmcp_openapi_server._resource_manager.get_resources() + resources = list(resources_dict.values()) # Find the get_users resource get_users_resource = next( @@ -523,13 +516,14 @@ class TestTagTransfer: assert len(get_users_resource.tags) == 2 async def test_tags_transferred_to_resource_templates( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI + self, fastmcp_openapi_server: FastMCPOpenAPI ): """Test that tags from OpenAPI routes are correctly transferred to ResourceTemplates.""" # Get internal resource templates directly - templates = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_resource_templates().values() + templates_dict = ( + await fastmcp_openapi_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) # Find the get_user template get_user_template = next( @@ -544,13 +538,14 @@ class TestTagTransfer: assert len(get_user_template.tags) == 2 async def test_tags_preserved_in_resources_created_from_templates( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI + self, fastmcp_openapi_server: FastMCPOpenAPI ): """Test that tags are preserved when creating resources from templates.""" # Get internal resource templates directly - templates = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_resource_templates().values() + templates_dict = ( + await fastmcp_openapi_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) # Find the get_user template get_user_template = next( @@ -1167,7 +1162,7 @@ class TestDescriptionPropagation: return httpx.AsyncClient(transport=transport, base_url="http://test") @pytest.fixture - async def simple_server_with_all_types(self, simple_openapi_spec, mock_client): + async def simple_mcp_server(self, simple_openapi_spec, mock_client): """Create a FastMCPOpenAPI server with the simple test spec.""" return FastMCPOpenAPI( openapi_spec=simple_openapi_spec, @@ -1179,11 +1174,11 @@ class TestDescriptionPropagation: # --- RESOURCE TESTS --- async def test_resource_includes_route_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a Resource includes the route description.""" resources = list( - simple_server_with_all_types._resource_manager.get_resources().values() + (await simple_mcp_server._resource_manager.get_resources()).values() ) list_resource = next((r for r in resources if r.name == "listItems"), None) @@ -1193,11 +1188,11 @@ class TestDescriptionPropagation: ) async def test_resource_includes_response_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a Resource includes the response description.""" resources = list( - simple_server_with_all_types._resource_manager.get_resources().values() + (await simple_mcp_server._resource_manager.get_resources()).values() ) list_resource = next((r for r in resources if r.name == "listItems"), None) @@ -1207,11 +1202,11 @@ class TestDescriptionPropagation: ) async def test_resource_includes_response_model_fields( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a Resource description includes response model field descriptions.""" resources = list( - simple_server_with_all_types._resource_manager.get_resources().values() + (await simple_mcp_server._resource_manager.get_resources()).values() ) list_resource = next((r for r in resources if r.name == "listItems"), None) @@ -1230,12 +1225,13 @@ class TestDescriptionPropagation: # --- RESOURCE TEMPLATE TESTS --- async def test_template_includes_route_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a ResourceTemplate includes the route description.""" - templates = list( - simple_server_with_all_types._resource_manager.get_templates().values() + templates_dict = ( + await simple_mcp_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) get_template = next((t for t in templates if t.name == "getItem"), None) assert get_template is not None, "getItem template wasn't created" @@ -1244,12 +1240,13 @@ class TestDescriptionPropagation: ) async def test_template_includes_function_docstring( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a ResourceTemplate includes the function docstring.""" - templates = list( - simple_server_with_all_types._resource_manager.get_templates().values() + templates_dict = ( + await simple_mcp_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) get_template = next((t for t in templates if t.name == "getItem"), None) assert get_template is not None, "getItem template wasn't created" @@ -1258,12 +1255,13 @@ class TestDescriptionPropagation: ) async def test_template_includes_path_parameter_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a ResourceTemplate includes path parameter descriptions.""" - templates = list( - simple_server_with_all_types._resource_manager.get_templates().values() + templates_dict = ( + await simple_mcp_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) get_template = next((t for t in templates if t.name == "getItem"), None) assert get_template is not None, "getItem template wasn't created" @@ -1272,12 +1270,13 @@ class TestDescriptionPropagation: ) async def test_template_includes_query_parameter_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a ResourceTemplate includes query parameter descriptions.""" - templates = list( - simple_server_with_all_types._resource_manager.get_templates().values() + templates_dict = ( + await simple_mcp_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) get_template = next((t for t in templates if t.name == "getItem"), None) assert get_template is not None, "getItem template wasn't created" @@ -1286,12 +1285,13 @@ class TestDescriptionPropagation: ) async def test_template_parameter_schema_includes_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a ResourceTemplate's parameter schema includes parameter descriptions.""" - templates = list( - simple_server_with_all_types._resource_manager.get_templates().values() + templates_dict = ( + await simple_mcp_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) get_template = next((t for t in templates if t.name == "getItem"), None) assert get_template is not None, "getItem template wasn't created" @@ -1311,9 +1311,10 @@ class TestDescriptionPropagation: # --- TOOL TESTS --- - async def test_tool_includes_route_description(self, simple_server_with_all_types): + async def test_tool_includes_route_description(self, simple_mcp_server: FastMCP): """Test that a Tool includes the route description.""" - tools = simple_server_with_all_types._tool_manager.list_tools() + tools_dict = await simple_mcp_server._tool_manager.get_tools() + tools = list(tools_dict.values()) create_tool = next((t for t in tools if t.name == "createItem"), None) assert create_tool is not None, "createItem tool wasn't created" @@ -1321,9 +1322,10 @@ class TestDescriptionPropagation: "Route description missing from Tool" ) - async def test_tool_includes_function_docstring(self, simple_server_with_all_types): + async def test_tool_includes_function_docstring(self, simple_mcp_server: FastMCP): """Test that a Tool includes the function docstring.""" - tools = simple_server_with_all_types._tool_manager.list_tools() + tools_dict = await simple_mcp_server._tool_manager.get_tools() + tools = list(tools_dict.values()) create_tool = next((t for t in tools if t.name == "createItem"), None) assert create_tool is not None, "createItem tool wasn't created" @@ -1333,10 +1335,11 @@ class TestDescriptionPropagation: ) async def test_tool_parameter_schema_includes_property_description( - self, simple_server_with_all_types + self, simple_mcp_server: FastMCP ): """Test that a Tool's parameter schema includes property descriptions from request model.""" - tools = simple_server_with_all_types._tool_manager.list_tools() + tools_dict = await simple_mcp_server._tool_manager.get_tools() + tools = list(tools_dict.values()) create_tool = next((t for t in tools if t.name == "createItem"), None) assert create_tool is not None, "createItem tool wasn't created" @@ -1356,9 +1359,9 @@ class TestDescriptionPropagation: # --- CLIENT API TESTS --- - async def test_client_api_resource_description(self, simple_server_with_all_types): + async def test_client_api_resource_description(self, simple_mcp_server: FastMCP): """Test that Resource descriptions are accessible via the client API.""" - async with Client(simple_server_with_all_types) as client: + async with Client(simple_mcp_server) as client: resources = await client.list_resources() list_resource = next((r for r in resources if r.name == "listItems"), None) @@ -1370,9 +1373,9 @@ class TestDescriptionPropagation: "Route description missing in Resource from client API" ) - async def test_client_api_template_description(self, simple_server_with_all_types): + async def test_client_api_template_description(self, simple_mcp_server: FastMCP): """Test that ResourceTemplate descriptions are accessible via the client API.""" - async with Client(simple_server_with_all_types) as client: + async with Client(simple_mcp_server) as client: templates = await client.list_resource_templates() get_template = next((t for t in templates if t.name == "getItem"), None) @@ -1384,9 +1387,9 @@ class TestDescriptionPropagation: "Route description missing in ResourceTemplate from client API" ) - async def test_client_api_tool_description(self, simple_server_with_all_types): + async def test_client_api_tool_description(self, simple_mcp_server: FastMCP): """Test that Tool descriptions are accessible via the client API.""" - async with Client(simple_server_with_all_types) as client: + async with Client(simple_mcp_server) as client: tools = await client.list_tools() create_tool = next((t for t in tools if t.name == "createItem"), None) @@ -1398,9 +1401,9 @@ class TestDescriptionPropagation: "Function docstring missing in Tool from client API" ) - async def test_client_api_tool_parameter_schema(self, simple_server_with_all_types): + async def test_client_api_tool_parameter_schema(self, simple_mcp_server: FastMCP): """Test that Tool parameter schemas are accessible via the client API.""" - async with Client(simple_server_with_all_types) as client: + async with Client(simple_mcp_server) as client: tools = await client.list_tools() create_tool = next((t for t in tools if t.name == "createItem"), None) @@ -1533,22 +1536,26 @@ class TestFastAPIDescriptionPropagation: # Debug: print all components created print("\nDEBUG - Resources created:") - for name, resource in server._resource_manager.get_resources().items(): + resources_dict = await server._resource_manager.get_resources() + for name, resource in resources_dict.items(): print(f" Resource: {name}, Name attribute: {resource.name}") print("\nDEBUG - Templates created:") - for name, template in server._resource_manager.get_resource_templates().items(): + templates_dict = await server._resource_manager.get_resource_templates() + for name, template in templates_dict.items(): print(f" Template: {name}, Name attribute: {template.name}") print("\nDEBUG - Tools created:") - for tool in server._tool_manager._list_tools(): + tools = await server._tool_manager._list_tools() + for tool in tools: print(f" Tool: {tool.name}") return server - async def test_resource_includes_function_docstring(self, fastapi_server): + async def test_resource_includes_function_docstring(self, fastapi_server: FastMCP): """Test that a Resource includes the function docstring.""" - resources = list(fastapi_server._resource_manager.get_resources().values()) + resources_dict = await fastapi_server._resource_manager.get_resources() + resources = list(resources_dict.values()) # Now checking for the get_items operation ID rather than list_items list_resource = next((r for r in resources if "items_get" in r.name), None) @@ -1559,13 +1566,16 @@ class TestFastAPIDescriptionPropagation: "Function docstring missing from Resource" ) - async def test_resource_includes_response_model_fields(self, fastapi_server): + async def test_resource_includes_response_model_fields( + self, fastapi_server: FastMCP + ): """Test that a Resource description includes basic response information. Note: FastAPI doesn't reliably include Pydantic field descriptions in the OpenAPI schema, so we can only check for basic response information being present. """ - resources = list(fastapi_server._resource_manager.get_resources().values()) + resources_dict = await fastapi_server._resource_manager.get_resources() + resources = list(resources_dict.values()) list_resource = next((r for r in resources if "items_get" in r.name), None) assert list_resource is not None, "GET /items resource wasn't created" @@ -1579,9 +1589,10 @@ class TestFastAPIDescriptionPropagation: # We've already verified in TestDescriptionPropagation that when descriptions # are present in the OpenAPI schema, they are properly included in the component description - async def test_template_includes_function_docstring(self, fastapi_server): + async def test_template_includes_function_docstring(self, fastapi_server: FastMCP): """Test that a ResourceTemplate includes the function docstring.""" - templates = list(fastapi_server._resource_manager.get_templates().values()) + templates_dict = await fastapi_server._resource_manager.get_resource_templates() + templates = list(templates_dict.values()) get_template = next((t for t in templates if "get_item_items" in t.name), None) assert get_template is not None, "GET /items/{item_id} template wasn't created" @@ -1590,13 +1601,16 @@ class TestFastAPIDescriptionPropagation: "Function docstring missing from ResourceTemplate" ) - async def test_template_includes_path_parameter_description(self, fastapi_server): + async def test_template_includes_path_parameter_description( + self, fastapi_server: FastMCP + ): """Test that a ResourceTemplate includes path parameter descriptions. Note: Currently, FastAPI parameter descriptions using Annotated[type, Field(description=...)] are not properly propagated to the OpenAPI schema. The parameters appear but without the description. """ - templates = list(fastapi_server._resource_manager.get_templates().values()) + templates_dict = await fastapi_server._resource_manager.get_resource_templates() + templates = list(templates_dict.values()) get_template = next((t for t in templates if "get_item_items" in t.name), None) assert get_template is not None, "GET /items/{item_id} template wasn't created" @@ -1610,13 +1624,16 @@ class TestFastAPIDescriptionPropagation: "item_id parameter missing from ResourceTemplate description" ) - async def test_template_includes_query_parameter_description(self, fastapi_server): + async def test_template_includes_query_parameter_description( + self, fastapi_server: FastMCP + ): """Test that a ResourceTemplate includes query parameter descriptions. Note: Currently, FastAPI parameter descriptions using Annotated[type, Field(description=...)] are not properly propagated to the OpenAPI schema. The parameters appear but without the description. """ - templates = list(fastapi_server._resource_manager.get_templates().values()) + templates_dict = await fastapi_server._resource_manager.get_resource_templates() + templates = list(templates_dict.values()) get_template = next((t for t in templates if "get_item_items" in t.name), None) assert get_template is not None, "GET /items/{item_id} template wasn't created" @@ -1630,9 +1647,12 @@ class TestFastAPIDescriptionPropagation: "fields parameter missing from ResourceTemplate description" ) - async def test_template_parameter_schema_includes_description(self, fastapi_server): + async def test_template_parameter_schema_includes_description( + self, fastapi_server: FastMCP + ): """Test that a ResourceTemplate's parameter schema includes parameter descriptions.""" - templates = list(fastapi_server._resource_manager.get_templates().values()) + templates_dict = await fastapi_server._resource_manager.get_resource_templates() + templates = list(templates_dict.values()) get_template = next((t for t in templates if "get_item_items" in t.name), None) assert get_template is not None, "GET /items/{item_id} template wasn't created" @@ -1650,9 +1670,10 @@ class TestFastAPIDescriptionPropagation: in get_template.parameters["properties"]["item_id"]["description"] ), "Path parameter description incorrect in schema" - async def test_tool_includes_function_docstring(self, fastapi_server): + async def test_tool_includes_function_docstring(self, fastapi_server: FastMCP): """Test that a Tool includes the function docstring.""" - tools = fastapi_server._tool_manager.list_tools() + tools_dict = await fastapi_server._tool_manager.get_tools() + tools = list(tools_dict.values()) create_tool = next( (t for t in tools if "create_item_items_post" == t.name), None ) @@ -1664,7 +1685,7 @@ class TestFastAPIDescriptionPropagation: ) async def test_tool_parameter_schema_includes_property_description( - self, fastapi_server + self, fastapi_server: FastMCP ): """Test that a Tool's parameter schema includes property descriptions from request model. @@ -1672,7 +1693,8 @@ class TestFastAPIDescriptionPropagation: may not be consistently propagated into the FastAPI OpenAPI schema and thus not into the tool's parameter schema. """ - tools = fastapi_server._tool_manager.list_tools() + tools_dict = await fastapi_server._tool_manager.get_tools() + tools = list(tools_dict.values()) create_tool = next( (t for t in tools if "create_item_items_post" == t.name), None ) @@ -1686,7 +1708,7 @@ class TestFastAPIDescriptionPropagation: ) # We don't test for the description field content as it may not be consistently propagated - async def test_client_api_resource_description(self, fastapi_server): + async def test_client_api_resource_description(self, fastapi_server: FastMCP): """Test that Resource descriptions are accessible via the client API.""" async with Client(fastapi_server) as client: resources = await client.list_resources() @@ -1700,7 +1722,7 @@ class TestFastAPIDescriptionPropagation: "Function docstring missing in Resource from client API" ) - async def test_client_api_template_description(self, fastapi_server): + async def test_client_api_template_description(self, fastapi_server: FastMCP): """Test that ResourceTemplate descriptions are accessible via the client API.""" async with Client(fastapi_server) as client: templates = await client.list_resource_templates() @@ -1716,7 +1738,7 @@ class TestFastAPIDescriptionPropagation: "Function docstring missing in ResourceTemplate from client API" ) - async def test_client_api_tool_description(self, fastapi_server): + async def test_client_api_tool_description(self, fastapi_server: FastMCP): """Test that Tool descriptions are accessible via the client API.""" async with Client(fastapi_server) as client: tools = await client.list_tools() @@ -1732,7 +1754,7 @@ class TestFastAPIDescriptionPropagation: "Function docstring missing in Tool from client API" ) - async def test_client_api_tool_parameter_schema(self, fastapi_server): + async def test_client_api_tool_parameter_schema(self, fastapi_server: FastMCP): """Test that Tool parameter schemas are accessible via the client API.""" async with Client(fastapi_server) as client: tools = await client.list_tools() @@ -1755,11 +1777,9 @@ class TestFastAPIDescriptionPropagation: class TestReprMethods: """Tests for the custom __repr__ methods of OpenAPI objects.""" - async def test_openapi_tool_repr( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI - ): + async def test_openapi_tool_repr(self, fastmcp_openapi_server: FastMCPOpenAPI): """Test that OpenAPITool's __repr__ method works without recursion errors.""" - tools = fastmcp_openapi_server_with_all_types._tool_manager._list_tools() + tools = await fastmcp_openapi_server._tool_manager._list_tools() tool = next(iter(tools)) # Verify repr doesn't cause recursion and contains expected elements @@ -1769,13 +1789,10 @@ class TestReprMethods: assert "method=" in tool_repr assert "path=" in tool_repr - async def test_openapi_resource_repr( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI - ): + async def test_openapi_resource_repr(self, fastmcp_openapi_server: FastMCPOpenAPI): """Test that OpenAPIResource's __repr__ method works without recursion errors.""" - resources = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_resources().values() - ) + resources_dict = await fastmcp_openapi_server._resource_manager.get_resources() + resources = list(resources_dict.values()) resource = next(iter(resources)) # Verify repr doesn't cause recursion and contains expected elements @@ -1786,12 +1803,13 @@ class TestReprMethods: assert "path=" in resource_repr async def test_openapi_resource_template_repr( - self, fastmcp_openapi_server_with_all_types: FastMCPOpenAPI + self, fastmcp_openapi_server: FastMCPOpenAPI ): """Test that OpenAPIResourceTemplate's __repr__ method works without recursion errors.""" - templates = list( - fastmcp_openapi_server_with_all_types._resource_manager.get_resource_templates().values() + templates_dict = ( + await fastmcp_openapi_server._resource_manager.get_resource_templates() ) + templates = list(templates_dict.values()) template = next(iter(templates)) # Verify repr doesn't cause recursion and contains expected elements @@ -1836,7 +1854,7 @@ class TestEnumHandling: ) # Get the tools from the server - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() # Find the read_item tool read_item_tool = next((t for t in tools if t.name == "read_item_items"), None) @@ -1929,7 +1947,7 @@ class TestRouteMapWildcard: ) # All operations should be mapped to tools - tools = mcp._tool_manager._list_tools() + tools = await mcp._tool_manager._list_tools() tool_names = {tool.name for tool in tools} # Check that all 4 operations became tools @@ -2007,11 +2025,11 @@ class TestRouteMapTags: ) # Check that admin-tagged routes are tools - tools = server._tool_manager.get_tools() - tool_names = {t.name for t in tools.values()} + tools_dict = await server._tool_manager.get_tools() + tool_names = {t.name for t in tools_dict.values()} - resources = server._resource_manager.get_resources() - resource_names = {r.name for r in resources.values()} + resources_dict = await server._resource_manager.get_resources() + resource_names = {r.name for r in resources_dict.values()} # Routes with "admin" tag should be tools assert "createUser" in tool_names @@ -2040,11 +2058,11 @@ class TestRouteMapTags: ) # Check that internal-tagged routes are excluded - resources = server._resource_manager.get_resources() - resource_names = {r.name for r in resources.values()} + resources_dict = await server._resource_manager.get_resources() + resource_names = {r.name for r in resources_dict.values()} - tools = server._tool_manager.get_tools() - tool_names = {t.name for t in tools.values()} + tools_dict = await server._tool_manager.get_tools() + tool_names = {t.name for t in tools_dict.values()} # Internal-tagged route should be excluded assert "getAdminStats" not in resource_names @@ -2075,11 +2093,11 @@ class TestRouteMapTags: route_maps=route_maps, ) - tools = server._tool_manager.get_tools() - tool_names = {t.name for t in tools.values()} + tools_dict = await server._tool_manager.get_tools() + tool_names = {t.name for t in tools_dict.values()} - resources = server._resource_manager.get_resources() - resource_names = {r.name for r in resources.values()} + resources_dict = await server._resource_manager.get_resources() + resource_names = {r.name for r in resources_dict.values()} # Only createUser has both "users" AND "admin" tags assert "createUser" in tool_names @@ -2110,11 +2128,11 @@ class TestRouteMapTags: route_maps=route_maps, ) - tools = server._tool_manager.get_tools() - tool_names = {t.name for t in tools.values()} + tools_dict = await server._tool_manager.get_tools() + tool_names = {t.name for t in tools_dict.values()} - resources = server._resource_manager.get_resources() - resource_names = {r.name for r in resources.values()} + resources_dict = await server._resource_manager.get_resources() + resource_names = {r.name for r in resources_dict.values()} # Only getAdminStats matches both /admin/ pattern AND "admin" tag assert "getAdminStats" in tool_names @@ -2140,8 +2158,8 @@ class TestRouteMapTags: route_maps=route_maps, ) - tools = server._tool_manager.get_tools() - tool_names = {t.name for t in tools.values()} + tools_dict = await server._tool_manager.get_tools() + tool_names = {t.name for t in tools_dict.values()} # All routes should be tools since empty tags matches everything expected_tools = { @@ -2246,18 +2264,18 @@ class TestMCPNames: ) # Check tools use custom names - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} assert "admin_create_user" in tool_names # Check resource templates use custom names - templates = list(server._resource_manager.get_resource_templates().values()) - template_names = {template.name for template in templates} + templates_dict = await server._resource_manager.get_resource_templates() + template_names = {template.name for template in templates_dict.values()} assert "user_detail" in template_names # Check resources use custom names - resources = list(server._resource_manager.get_resources().values()) - resource_names = {resource.name for resource in resources} + resources_dict = await server._resource_manager.get_resources() + resource_names = {resource.name for resource in resources_dict.values()} assert "user_list" in resource_names async def test_mcp_names_fallback_to_operation_id_short( @@ -2276,14 +2294,14 @@ class TestMCPNames: route_maps=GET_ROUTE_MAPS, ) - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} - templates = list(server._resource_manager.get_resource_templates().values()) - template_names = {template.name for template in templates} + templates_dict = await server._resource_manager.get_resource_templates() + template_names = {template.name for template in templates_dict.values()} - resources = list(server._resource_manager.get_resources().values()) - resource_names = {resource.name for resource in resources} + resources_dict = await server._resource_manager.get_resources() + resource_names = {resource.name for resource in resources_dict.values()} # Custom mapped name should be used assert "custom_user_list" in resource_names @@ -2300,9 +2318,11 @@ class TestMCPNames: route_maps=GET_ROUTE_MAPS, ) - resources = list(server._resource_manager.get_resources().values()) + resources_dict = await server._resource_manager.get_resources() resource_names = { - resource.name for resource in resources if resource.name is not None + resource.name + for resource in resources_dict.values() + if resource.name is not None } # Special chars and spaces should be slugified @@ -2330,14 +2350,14 @@ class TestMCPNames: # Check all component types all_names = [] - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() all_names.extend(tool.name for tool in tools) - resources = list(server._resource_manager.get_resources().values()) - all_names.extend(resource.name for resource in resources) + resources_dict = await server._resource_manager.get_resources() + all_names.extend(resource.name for resource in resources_dict.values()) - templates = list(server._resource_manager.get_resource_templates().values()) - all_names.extend(template.name for template in templates) + templates_dict = await server._resource_manager.get_resource_templates() + all_names.extend(template.name for template in templates_dict.values()) # All names should be 56 characters or less for name in all_names: @@ -2363,7 +2383,7 @@ class TestMCPNames: mcp_names=mcp_names, ) - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} assert "openapi_user_list" in tool_names @@ -2395,7 +2415,7 @@ class TestMCPNames: mcp_names=mcp_names, ) - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() tool_names = {tool.name for tool in tools} assert "fastapi_create_user" in tool_names @@ -2419,9 +2439,11 @@ class TestMCPNames: route_maps=GET_ROUTE_MAPS, ) - resources = list(server._resource_manager.get_resources().values()) + resources_dict = await server._resource_manager.get_resources() resource_names = { - resource.name for resource in resources if resource.name is not None + resource.name + for resource in resources_dict.values() + if resource.name is not None } # Find the resource that should have the custom name @@ -2496,7 +2518,7 @@ class TestRouteMapMCPTags: ) # Get the POST tool - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() create_user_tool = next((t for t in tools if "create_user" in t.name), None) assert create_user_tool is not None, "create_user tool not found" @@ -2530,7 +2552,8 @@ class TestRouteMapMCPTags: ) # Get the resource - resources = list(server._resource_manager.get_resources().values()) + resources_dict = await server._resource_manager.get_resources() + resources = list(resources_dict.values()) get_users_resource = next((r for r in resources if "get_users" in r.name), None) assert get_users_resource is not None, "get_users resource not found" @@ -2564,7 +2587,8 @@ class TestRouteMapMCPTags: ) # Get the resource template - templates = list(server._resource_manager.get_resource_templates().values()) + templates_dict = await server._resource_manager.get_resource_templates() + templates = list(templates_dict.values()) get_user_template = next((t for t in templates if "get_user" in t.name), None) assert get_user_template is not None, "get_user template not found" @@ -2610,21 +2634,23 @@ class TestRouteMapMCPTags: ) # Check tool tags - tools = server._tool_manager._list_tools() + tools = await server._tool_manager._list_tools() create_tool = next((t for t in tools if "create_user" in t.name), None) assert create_tool is not None assert "write-operation" in create_tool.tags assert "mutation" in create_tool.tags # Check resource template tags - templates = list(server._resource_manager.get_resource_templates().values()) + templates_dict = await server._resource_manager.get_resource_templates() + templates = list(templates_dict.values()) detail_template = next((t for t in templates if "get_user" in t.name), None) assert detail_template is not None assert "detail" in detail_template.tags assert "single-item" in detail_template.tags # Check resource tags - resources = list(server._resource_manager.get_resources().values()) + resources_dict = await server._resource_manager.get_resources() + resources = list(resources_dict.values()) list_resource = next((r for r in resources if "get_users" in r.name), None) assert list_resource is not None assert "list" in list_resource.tags diff --git a/tests/server/test_mount.py b/tests/server/test_mount.py index d32d83c45..befaccfdd 100644 --- a/tests/server/test_mount.py +++ b/tests/server/test_mount.py @@ -71,7 +71,8 @@ class TestBasicMount: # Mount without deprecated parameters main_app.mount(api_app, "api") - async def test_mount_with_no_prefix(self): + @pytest.mark.parametrize("prefix", ["", None]) + async def test_mount_with_no_prefix(self, prefix): main_app = FastMCP("MainApp") sub_app = FastMCP("SubApp") @@ -80,7 +81,7 @@ class TestBasicMount: return "This is from the sub app" # Mount with empty prefix but without deprecated separators - main_app.mount(sub_app, prefix="") + main_app.mount(sub_app, prefix=prefix) tools = await main_app.get_tools() # With empty prefix, the tool should keep its original name diff --git a/tests/server/test_resource_prefix_formats.py b/tests/server/test_resource_prefix_formats.py index cc13dd45e..c45731ccd 100644 --- a/tests/server/test_resource_prefix_formats.py +++ b/tests/server/test_resource_prefix_formats.py @@ -56,8 +56,8 @@ async def test_resource_prefix_format_in_import_server(): await main_server_protocol.import_server(server, "sub") # Check that the resources are prefixed correctly - path_resources = main_server_path._resource_manager.get_resources() - protocol_resources = main_server_protocol._resource_manager.get_resources() + path_resources = await main_server_path._resource_manager.get_resources() + protocol_resources = await main_server_protocol._resource_manager.get_resources() # Path format should be resource://sub/test assert "resource://sub/test" in path_resources diff --git a/tests/server/test_tool_annotations.py b/tests/server/test_tool_annotations.py index 84d031df9..07a21077d 100644 --- a/tests/server/test_tool_annotations.py +++ b/tests/server/test_tool_annotations.py @@ -22,7 +22,8 @@ async def test_tool_annotations_in_tool_manager(): return message # Check internal tool objects directly - tools = mcp._tool_manager._list_tools() + tools_dict = await mcp._tool_manager.get_tools() + tools = list(tools_dict.values()) assert len(tools) == 1 assert tools[0].annotations is not None assert tools[0].annotations.title == "Echo Tool" @@ -124,7 +125,8 @@ async def test_direct_tool_annotations_in_tool_manager(): return {"modified": True, **data} # Check internal tool objects directly - tools = mcp._tool_manager._list_tools() + tools_dict = await mcp._tool_manager.get_tools() + tools = list(tools_dict.values()) assert len(tools) == 1 assert tools[0].annotations is not None assert tools[0].annotations.title == "Direct Tool" @@ -183,7 +185,8 @@ async def test_add_tool_method_annotations(): mcp.add_tool(tool) # Check internal tool objects directly - tools = mcp._tool_manager._list_tools() + tools_dict = await mcp._tool_manager.get_tools() + tools = list(tools_dict.values()) assert len(tools) == 1 assert tools[0].annotations is not None assert tools[0].annotations.title == "Create Item" diff --git a/tests/server/test_tool_exclude_args.py b/tests/server/test_tool_exclude_args.py index cc57bf3c2..025555ebc 100644 --- a/tests/server/test_tool_exclude_args.py +++ b/tests/server/test_tool_exclude_args.py @@ -19,9 +19,10 @@ async def test_tool_exclude_args_in_tool_manager(): pass return message - tools = mcp._tool_manager._list_tools() + tools_dict = await mcp._tool_manager.get_tools() + tools = list(tools_dict.values()) assert len(tools) == 1 - assert "state" not in echo.parameters["properties"] + assert "state" not in tools[0].parameters["properties"] async def test_tool_exclude_args_without_default_value_raises_error(): @@ -60,7 +61,8 @@ async def test_add_tool_method_exclude_args(): mcp.add_tool(tool) # Check internal tool objects directly - tools = mcp._tool_manager._list_tools() + tools_dict = await mcp._tool_manager.get_tools() + tools = list(tools_dict.values()) assert len(tools) == 1 assert "state" not in tools[0].parameters["properties"] From 2ae5a800e76a6dca853233a50d8502a633854006 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 14:50:58 -0400 Subject: [PATCH 12/17] Fix proxy servers --- src/fastmcp/resources/template.py | 13 + src/fastmcp/server/proxy.py | 412 ++++++++++++++++++------------ 2 files changed, 268 insertions(+), 157 deletions(-) diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index 8ea97c73f..ca4f3a3ee 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -131,6 +131,19 @@ class ResourceTemplate(FastMCPComponent): } return MCPResourceTemplate(**kwargs | overrides) + @classmethod + def from_mcp_template(cls, mcp_template: MCPResourceTemplate) -> ResourceTemplate: + """Creates a FastMCP ResourceTemplate from a raw MCP ResourceTemplate object.""" + # Note: This creates a simple ResourceTemplate instance. For function-based templates, + # the original function is lost, which is expected for remote templates. + return cls( + uri_template=mcp_template.uriTemplate, + name=mcp_template.name, + description=mcp_template.description, + mime_type=mcp_template.mimeType or "text/plain", + parameters={}, # Remote templates don't have local parameters + ) + @property def key(self) -> str: """ diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index a65f17b5b..7958e8234 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any, cast from urllib.parse import quote import mcp.types -from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.shared.exceptions import McpError from mcp.types import ( METHOD_NOT_FOUND, @@ -17,10 +16,14 @@ from pydantic.networks import AnyUrl from fastmcp.client import Client from fastmcp.exceptions import NotFoundError, ResourceError, ToolError from fastmcp.prompts import Prompt, PromptMessage +from fastmcp.prompts.prompt import PromptArgument +from fastmcp.prompts.prompt_manager import PromptManager from fastmcp.resources import Resource, ResourceTemplate +from fastmcp.resources.resource_manager import ResourceManager from fastmcp.server.context import Context from fastmcp.server.server import FastMCP from fastmcp.tools.tool import Tool +from fastmcp.tools.tool_manager import ToolManager from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import MCPContent @@ -30,18 +33,197 @@ if TYPE_CHECKING: logger = get_logger(__name__) +class ProxyToolManager(ToolManager): + """A ToolManager that sources its tools from a remote client in addition to local and mounted tools.""" + + def __init__(self, client: Client, **kwargs): + super().__init__(**kwargs) + self.client = client + + async def get_tools(self) -> dict[str, Tool]: + """Gets the unfiltered tool inventory including local, mounted, and proxy tools.""" + # First get local and mounted tools from parent + all_tools = await super().get_tools() + + # Then add proxy tools, but don't overwrite existing ones + try: + async with self.client: + client_tools = await self.client.list_tools() + for tool in client_tools: + if tool.name not in all_tools: + all_tools[tool.name] = ProxyTool.from_mcp_tool( + self.client, tool + ) + except McpError as e: + if e.error.code == METHOD_NOT_FOUND: + pass # No tools available from proxy + else: + raise e + + return all_tools + + async def _list_tools(self) -> list[Tool]: + """Gets the filtered list of tools including local, mounted, and proxy tools.""" + tools_dict = await self.get_tools() + return list(tools_dict.values()) + + async def call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: + """Calls a tool, trying local/mounted first, then proxy if not found.""" + try: + # First try local and mounted tools + return await super().call_tool(key, arguments) + except NotFoundError: + # If not found locally, try proxy + async with self.client: + return await self.client.call_tool(key, arguments) + + +class ProxyResourceManager(ResourceManager): + """A ResourceManager that sources its resources from a remote client in addition to local and mounted resources.""" + + def __init__(self, client: Client, **kwargs): + super().__init__(**kwargs) + self.client = client + + async def get_resources(self) -> dict[str, Resource]: + """Gets the unfiltered resource inventory including local, mounted, and proxy resources.""" + # First get local and mounted resources from parent + all_resources = await super().get_resources() + + # Then add proxy resources, but don't overwrite existing ones + try: + async with self.client: + client_resources = await self.client.list_resources() + for resource in client_resources: + if str(resource.uri) not in all_resources: + all_resources[str(resource.uri)] = ( + ProxyResource.from_mcp_resource(self.client, resource) + ) + except McpError as e: + if e.error.code == METHOD_NOT_FOUND: + pass # No resources available from proxy + else: + raise e + + return all_resources + + async def get_resource_templates(self) -> dict[str, ResourceTemplate]: + """Gets the unfiltered template inventory including local, mounted, and proxy templates.""" + # First get local and mounted templates from parent + all_templates = await super().get_resource_templates() + + # Then add proxy templates, but don't overwrite existing ones + try: + async with self.client: + client_templates = await self.client.list_resource_templates() + for template in client_templates: + if template.uriTemplate not in all_templates: + all_templates[template.uriTemplate] = ( + ProxyTemplate.from_mcp_template(self.client, template) + ) + except McpError as e: + if e.error.code == METHOD_NOT_FOUND: + pass # No templates available from proxy + else: + raise e + + return all_templates + + async def _list_resources(self) -> list[Resource]: + """Gets the filtered list of resources including local, mounted, and proxy resources.""" + resources_dict = await self.get_resources() + return list(resources_dict.values()) + + async def _list_resource_templates(self) -> list[ResourceTemplate]: + """Gets the filtered list of templates including local, mounted, and proxy templates.""" + templates_dict = await self.get_resource_templates() + return list(templates_dict.values()) + + async def read_resource(self, uri: AnyUrl | str) -> str | bytes: + """Reads a resource, trying local/mounted first, then proxy if not found.""" + try: + # First try local and mounted resources + return await super().read_resource(uri) + except NotFoundError: + # If not found locally, try proxy + async with self.client: + result = await self.client.read_resource(uri) + if isinstance(result[0], TextResourceContents): + return result[0].text + elif isinstance(result[0], BlobResourceContents): + return result[0].blob + else: + raise ResourceError(f"Unsupported content type: {type(result[0])}") + + +class ProxyPromptManager(PromptManager): + """A PromptManager that sources its prompts from a remote client in addition to local and mounted prompts.""" + + def __init__(self, client: Client, **kwargs): + super().__init__(**kwargs) + self.client = client + + async def get_prompts(self) -> dict[str, Prompt]: + """Gets the unfiltered prompt inventory including local, mounted, and proxy prompts.""" + # First get local and mounted prompts from parent + all_prompts = await super().get_prompts() + + # Then add proxy prompts, but don't overwrite existing ones + try: + async with self.client: + client_prompts = await self.client.list_prompts() + for prompt in client_prompts: + if prompt.name not in all_prompts: + all_prompts[prompt.name] = ProxyPrompt.from_mcp_prompt( + self.client, prompt + ) + except McpError as e: + if e.error.code == METHOD_NOT_FOUND: + pass # No prompts available from proxy + else: + raise e + + return all_prompts + + async def _list_prompts(self) -> list[Prompt]: + """Gets the filtered list of prompts including local, mounted, and proxy prompts.""" + prompts_dict = await self.get_prompts() + return list(prompts_dict.values()) + + async def render_prompt( + self, + name: str, + arguments: dict[str, Any] | None = None, + ) -> GetPromptResult: + """Renders a prompt, trying local/mounted first, then proxy if not found.""" + try: + # First try local and mounted prompts + return await super().render_prompt(name, arguments) + except NotFoundError: + # If not found locally, try proxy + async with self.client: + result = await self.client.get_prompt(name, arguments) + return result + + class ProxyTool(Tool): + """ + A Tool that represents and executes a tool on a remote server. + """ + def __init__(self, client: Client, **kwargs): super().__init__(**kwargs) self._client = client @classmethod - async def from_client(cls, client: Client, tool: mcp.types.Tool) -> ProxyTool: + def from_mcp_tool(cls, client: Client, mcp_tool: mcp.types.Tool) -> ProxyTool: + """Factory method to create a ProxyTool from a raw MCP tool schema.""" return cls( client=client, - name=tool.name, - description=tool.description, - parameters=tool.inputSchema, + name=mcp_tool.name, + description=mcp_tool.description, + parameters=mcp_tool.inputSchema, + annotations=mcp_tool.annotations, ) async def run( @@ -49,8 +231,8 @@ class ProxyTool(Tool): arguments: dict[str, Any], context: Context | None = None, ) -> list[MCPContent]: - # the client context manager will swallow any exceptions inside a TaskGroup - # so we return the raw result and raise an exception ourselves + """Executes the tool by making a call through the client.""" + # This is where the remote execution logic lives. async with self._client: result = await self._client.call_tool_mcp( name=self.name, @@ -62,6 +244,10 @@ class ProxyTool(Tool): class ProxyResource(Resource): + """ + A Resource that represents and reads a resource from a remote server. + """ + _client: Client _value: str | bytes | None = None @@ -71,18 +257,20 @@ class ProxyResource(Resource): self._value = _value @classmethod - async def from_client( - cls, client: Client, resource: mcp.types.Resource + def from_mcp_resource( + cls, client: Client, mcp_resource: mcp.types.Resource ) -> ProxyResource: + """Factory method to create a ProxyResource from a raw MCP resource schema.""" return cls( client=client, - uri=resource.uri, - name=resource.name, - description=resource.description, - mime_type=resource.mimeType, + uri=mcp_resource.uri, + name=mcp_resource.name, + description=mcp_resource.description, + mime_type=mcp_resource.mimeType or "text/plain", ) async def read(self) -> str | bytes: + """Read the resource content from the remote server.""" if self._value is not None: return self._value @@ -97,20 +285,26 @@ class ProxyResource(Resource): class ProxyTemplate(ResourceTemplate): + """ + A ResourceTemplate that represents and creates resources from a remote server template. + """ + def __init__(self, client: Client, **kwargs): super().__init__(**kwargs) self._client = client @classmethod - async def from_client( - cls, client: Client, template: mcp.types.ResourceTemplate + def from_mcp_template( + cls, client: Client, mcp_template: mcp.types.ResourceTemplate ) -> ProxyTemplate: + """Factory method to create a ProxyTemplate from a raw MCP template schema.""" return cls( client=client, - uri_template=template.uriTemplate, - name=template.name, - description=template.description, - parameters={}, + uri_template=mcp_template.uriTemplate, + name=mcp_template.name, + description=mcp_template.description, + mime_type=mcp_template.mimeType or "text/plain", + parameters={}, # Remote templates don't have local parameters ) async def create_resource( @@ -119,6 +313,7 @@ class ProxyTemplate(ResourceTemplate): params: dict[str, Any], context: Context | None = None, ) -> ProxyResource: + """Create a resource from the template by calling the remote server.""" # don't use the provided uri, because it may not be the same as the # uri_template on the remote server. # quote params to ensure they are valid for the uri_template @@ -146,6 +341,10 @@ class ProxyTemplate(ResourceTemplate): class ProxyPrompt(Prompt): + """ + A Prompt that represents and renders a prompt from a remote server. + """ + _client: Client def __init__(self, client: Client, **kwargs): @@ -153,157 +352,56 @@ class ProxyPrompt(Prompt): self._client = client @classmethod - async def from_client(cls, client: Client, prompt: mcp.types.Prompt) -> ProxyPrompt: + def from_mcp_prompt( + cls, client: Client, mcp_prompt: mcp.types.Prompt + ) -> ProxyPrompt: + """Factory method to create a ProxyPrompt from a raw MCP prompt schema.""" + arguments = [ + PromptArgument( + name=arg.name, + description=arg.description, + required=arg.required or False, + ) + for arg in mcp_prompt.arguments or [] + ] return cls( client=client, - name=prompt.name, - description=prompt.description, - arguments=[a.model_dump() for a in prompt.arguments or []], + name=mcp_prompt.name, + description=mcp_prompt.description, + arguments=arguments, ) async def render(self, arguments: dict[str, Any]) -> list[PromptMessage]: + """Render the prompt by making a call through the client.""" async with self._client: result = await self._client.get_prompt(self.name, arguments) return result.messages class FastMCPProxy(FastMCP): + """ + A FastMCP server that acts as a proxy to a remote MCP-compliant server. + It uses specialized managers that fulfill requests via an HTTP client. + """ + def __init__(self, client: Client, **kwargs): + """ + Initializes the proxy server. + + Args: + client: The FastMCP client connected to the backend server. + **kwargs: Additional settings for the FastMCP server. + """ super().__init__(**kwargs) self.client = client - async def _list_tools(self, apply_middleware: bool = True) -> list[Tool]: - tools = { - tool.name: tool - for tool in await super()._list_tools(apply_middleware=apply_middleware) - } + # Replace the default managers with our specialized proxy managers. + self._tool_manager = ProxyToolManager(client=self.client) + self._resource_manager = ProxyResourceManager(client=self.client) + self._prompt_manager = ProxyPromptManager(client=self.client) - async with self.client: - try: - client_tools = await self.client.list_tools() - except McpError as e: - if e.error.code == METHOD_NOT_FOUND: - client_tools = [] - else: - raise e - for tool in client_tools: - # don't overwrite tools defined in the server - if tool.name not in tools: - tool_proxy = await ProxyTool.from_client(self.client, tool) - tools[tool_proxy.key] = tool_proxy - - return list(tools.values()) - - async def _list_resources(self, apply_middleware: bool = True) -> list[Resource]: - resources = { - resource.uri: resource - for resource in await super()._list_resources( - apply_middleware=apply_middleware - ) - } - - async with self.client: - try: - client_resources = await self.client.list_resources() - except McpError as e: - if e.error.code == METHOD_NOT_FOUND: - client_resources = [] - else: - raise e - for resource in client_resources: - # don't overwrite resources defined in the server - if resource.uri not in resources: - resource_proxy = await ProxyResource.from_client( - self.client, resource - ) - resources[resource_proxy.uri] = resource_proxy - - return list(resources.values()) - - async def _list_resource_templates( - self, apply_middleware: bool = True - ) -> list[ResourceTemplate]: - templates = { - template.uri_template: template - for template in await super()._list_resource_templates( - apply_middleware=apply_middleware - ) - } - - async with self.client: - try: - client_templates = await self.client.list_resource_templates() - except McpError as e: - if e.error.code == METHOD_NOT_FOUND: - client_templates = [] - else: - raise e - for template in client_templates: - # don't overwrite templates defined in the server - if template.uriTemplate not in templates: - template_proxy = await ProxyTemplate.from_client( - self.client, template - ) - templates[template_proxy.uri_template] = template_proxy - - return list(templates.values()) - - async def _list_prompts(self, apply_middleware: bool = True) -> list[Prompt]: - prompts = { - prompt.name: prompt - for prompt in await super()._list_prompts(apply_middleware=apply_middleware) - } - - async with self.client: - try: - client_prompts = await self.client.list_prompts() - except McpError as e: - if e.error.code == METHOD_NOT_FOUND: - client_prompts = [] - else: - raise e - for prompt in client_prompts: - # don't overwrite prompts defined in the server - if prompt.name not in prompts: - prompt_proxy = await ProxyPrompt.from_client(self.client, prompt) - prompts[prompt_proxy.name] = prompt_proxy - - return list(prompts.values()) - - async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]: - try: - result = await super()._call_tool(key, arguments) - return result - except NotFoundError: - async with self.client: - result = await self.client.call_tool(key, arguments) - return result - - async def _read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]: - try: - result = await super()._read_resource(uri) - return result - except NotFoundError: - async with self.client: - resource = await self.client.read_resource(uri) - if isinstance(resource[0], TextResourceContents): - content = resource[0].text - elif isinstance(resource[0], BlobResourceContents): - content = resource[0].blob - else: - raise ValueError(f"Unsupported content type: {type(resource[0])}") - - return [ - ReadResourceContents(content=content, mime_type=resource[0].mimeType) - ] - - async def _get_prompt( - self, name: str, arguments: dict[str, Any] | None = None - ) -> GetPromptResult: - try: - result = await super()._get_prompt(name, arguments) - return result - except NotFoundError: - async with self.client: - result = await self.client.get_prompt(name, arguments) - return result + # + # NO MORE METHOD OVERRIDES ARE NEEDED! + # The inherited _list_tools, _call_tool, etc., from FastMCP will now + # correctly delegate to the proxy managers, which in turn call the client. + # From 703589e6f592cf77d6c238482eab84b81cc4ad2f Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 15:00:11 -0400 Subject: [PATCH 13/17] Add proxy middleware test --- tests/server/middleware/test_middleware.py | 23 ++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index c789f42be..cd8b256bf 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -75,9 +75,9 @@ class RecordingMiddleware(MCPMiddleware): """Assert that a hook was called a specific number of times.""" calls = self.get_calls(hook=hook, method=method) actual_times = len(calls) + identifier = dict(hook=hook, method=method) assert actual_times == times, ( - f"Expected {hook!r} to be called {times} times" - f"{f' for method {method!r}' if method else ''}, " + f"Expected {times} calls for {identifier}, " f"but was called {actual_times} times" ) return True @@ -546,3 +546,22 @@ class TestNestedMiddlewareHooks: assert nested_middleware.assert_called( hook="on_list_resource_templates", times=1 ) + + +class TestProxyServer: + async def test_call_tool( + self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware + ): + # proxy server will have its tools listed as well as called in order to + # run the `should_enable_component` hook prior to the call. + proxy_server = FastMCP.as_proxy(mcp_server, name="Proxy Server") + async with Client(proxy_server) as client: + await client.call_tool("add", {"a": 1, "b": 2}) + + assert recording_middleware.assert_called(times=6) + assert recording_middleware.assert_called(method="tools/call", times=3) + assert recording_middleware.assert_called(method="tools/list", times=3) + assert recording_middleware.assert_called(hook="on_message", times=2) + assert recording_middleware.assert_called(hook="on_request", times=2) + assert recording_middleware.assert_called(hook="on_call_tool", times=1) + assert recording_middleware.assert_called(hook="on_list_tools", times=1) From 725c061d668bbd549232f6f973d84c2d02acdca1 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 15:19:31 -0400 Subject: [PATCH 14/17] Fix import prefix behavior --- src/fastmcp/prompts/prompt_manager.py | 20 +++---- src/fastmcp/resources/resource_manager.py | 69 ++++++++-------------- src/fastmcp/server/proxy.py | 14 ++--- src/fastmcp/server/server.py | 47 ++++++--------- src/fastmcp/tools/tool_manager.py | 19 +++--- tests/client/test_client.py | 7 ++- tests/deprecated/test_resource_prefixes.py | 2 +- tests/server/openapi/test_openapi.py | 24 ++++---- tests/server/test_import_server.py | 8 +-- tests/server/test_mount.py | 27 +++++---- tests/server/test_server.py | 2 +- 11 files changed, 104 insertions(+), 135 deletions(-) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index b237dc0fc..3436e71c4 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -63,7 +63,7 @@ class PromptManager: child_results = await mounted.server._list_prompts() else: # Use the manager-to-manager unfiltered path - child_results = await mounted.server._prompt_manager._list_prompts() + child_results = await mounted.server._prompt_manager.list_prompts() # The combination logic is the same for both paths child_dict = {p.key: p for p in child_results} @@ -104,7 +104,7 @@ class PromptManager: """ return await self._load_prompts(via_server=False) - async def _list_prompts(self) -> list[Prompt]: + async def list_prompts(self) -> list[Prompt]: """ Lists all prompts, applying protocol filtering. """ @@ -131,24 +131,22 @@ class PromptManager: ) return self.add_prompt(prompt) # type: ignore - def add_prompt(self, prompt: Prompt, key: str | None = None) -> Prompt: + def add_prompt(self, prompt: Prompt) -> Prompt: """Add a prompt to the manager.""" - key = key or prompt.name - # Check for duplicates - existing = self._prompts.get(key) + existing = self._prompts.get(prompt.key) if existing: if self.duplicate_behavior == "warn": - logger.warning(f"Prompt already exists: {key}") - self._prompts[key] = prompt + logger.warning(f"Prompt already exists: {prompt.key}") + self._prompts[prompt.key] = prompt elif self.duplicate_behavior == "replace": - self._prompts[key] = prompt + self._prompts[prompt.key] = prompt elif self.duplicate_behavior == "error": - raise ValueError(f"Prompt already exists: {key}") + raise ValueError(f"Prompt already exists: {prompt.key}") elif self.duplicate_behavior == "ignore": return existing else: - self._prompts[key] = prompt + self._prompts[prompt.key] = prompt return prompt async def render_prompt( diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 4bc084181..8741837ba 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -136,7 +136,9 @@ class ResourceManager: child_templates = await mounted.server._list_resource_templates() else: # Use the manager-to-manager unfiltered path - child_templates = await mounted.server._resource_manager._list_resource_templates() + child_templates = ( + await mounted.server._resource_manager.list_resource_templates() + ) child_dict = {template.key: template for template in child_templates} # Apply prefix if needed @@ -163,14 +165,14 @@ class ResourceManager: all_templates.update(self._templates) return all_templates - async def _list_resources(self) -> list[Resource]: + async def list_resources(self) -> list[Resource]: """ Lists all resources, applying protocol filtering. """ resources_dict = await self._load_resources(via_server=True) return list(resources_dict.values()) - async def _list_resource_templates(self) -> list[ResourceTemplate]: + async def list_resource_templates(self) -> list[ResourceTemplate]: """ Lists all templates, applying protocol filtering. """ @@ -265,35 +267,26 @@ class ResourceManager: ) return self.add_resource(resource) - def add_resource(self, resource: Resource, key: str | None = None) -> Resource: + def add_resource(self, resource: Resource) -> Resource: """Add a resource to the manager. Args: - resource: A Resource instance to add - key: Optional URI to use as the storage key (if different from resource.uri) + resource: A Resource instance to add. The resource's .key attribute + will be used as the storage key. To overwrite it, call + Resource.with_key() before calling this method. """ - storage_key = key or str(resource.uri) - logger.debug( - "Adding resource", - extra={ - "uri": resource.uri, - "storage_key": storage_key, - "type": type(resource).__name__, - "resource_name": resource.name, - }, - ) - existing = self._resources.get(storage_key) + existing = self._resources.get(resource.key) if existing: if self.duplicate_behavior == "warn": - logger.warning(f"Resource already exists: {storage_key}") - self._resources[storage_key] = resource + logger.warning(f"Resource already exists: {resource.key}") + self._resources[resource.key] = resource elif self.duplicate_behavior == "replace": - self._resources[storage_key] = resource + self._resources[resource.key] = resource elif self.duplicate_behavior == "error": - raise ValueError(f"Resource already exists: {storage_key}") + raise ValueError(f"Resource already exists: {resource.key}") elif self.duplicate_behavior == "ignore": return existing - self._resources[storage_key] = resource + self._resources[resource.key] = resource return resource def add_template_from_fn( @@ -323,42 +316,30 @@ class ResourceManager: ) return self.add_template(template) - def add_template( - self, template: ResourceTemplate, key: str | None = None - ) -> ResourceTemplate: + def add_template(self, template: ResourceTemplate) -> ResourceTemplate: """Add a template to the manager. Args: - template: A ResourceTemplate instance to add - key: Optional URI template to use as the storage key (if different from template.uri_template) + template: A ResourceTemplate instance to add. The template's .key attribute + will be used as the storage key. To overwrite it, call + ResourceTemplate.with_key() before calling this method. Returns: The added template. If a template with the same URI already exists, returns the existing template. """ - uri_template_str = str(template.uri_template) - storage_key = key or uri_template_str - logger.debug( - "Adding template", - extra={ - "uri_template": uri_template_str, - "storage_key": storage_key, - "type": type(template).__name__, - "template_name": template.name, - }, - ) - existing = self._templates.get(storage_key) + existing = self._templates.get(template.key) if existing: if self.duplicate_behavior == "warn": - logger.warning(f"Template already exists: {storage_key}") - self._templates[storage_key] = template + logger.warning(f"Template already exists: {template.key}") + self._templates[template.key] = template elif self.duplicate_behavior == "replace": - self._templates[storage_key] = template + self._templates[template.key] = template elif self.duplicate_behavior == "error": - raise ValueError(f"Template already exists: {storage_key}") + raise ValueError(f"Template already exists: {template.key}") elif self.duplicate_behavior == "ignore": return existing - self._templates[storage_key] = template + self._templates[template.key] = template return template async def has_resource(self, uri: AnyUrl | str) -> bool: diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 7958e8234..aeabf499f 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -62,7 +62,7 @@ class ProxyToolManager(ToolManager): return all_tools - async def _list_tools(self) -> list[Tool]: + async def list_tools(self) -> list[Tool]: """Gets the filtered list of tools including local, mounted, and proxy tools.""" tools_dict = await self.get_tools() return list(tools_dict.values()) @@ -129,12 +129,12 @@ class ProxyResourceManager(ResourceManager): return all_templates - async def _list_resources(self) -> list[Resource]: + async def list_resources(self) -> list[Resource]: """Gets the filtered list of resources including local, mounted, and proxy resources.""" resources_dict = await self.get_resources() return list(resources_dict.values()) - async def _list_resource_templates(self) -> list[ResourceTemplate]: + async def list_resource_templates(self) -> list[ResourceTemplate]: """Gets the filtered list of templates including local, mounted, and proxy templates.""" templates_dict = await self.get_resource_templates() return list(templates_dict.values()) @@ -185,7 +185,7 @@ class ProxyPromptManager(PromptManager): return all_prompts - async def _list_prompts(self) -> list[Prompt]: + async def list_prompts(self) -> list[Prompt]: """Gets the filtered list of prompts including local, mounted, and proxy prompts.""" prompts_dict = await self.get_prompts() return list(prompts_dict.values()) @@ -399,9 +399,3 @@ class FastMCPProxy(FastMCP): self._tool_manager = ProxyToolManager(client=self.client) self._resource_manager = ProxyResourceManager(client=self.client) self._prompt_manager = ProxyPromptManager(client=self.client) - - # - # NO MORE METHOD OVERRIDES ARE NEEDED! - # The inherited _list_tools, _call_tool, etc., from FastMCP will now - # correctly delegate to the proxy managers, which in turn call the client. - # diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 456572a8b..a6c6321ab 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -159,7 +159,6 @@ class FastMCP(Generic[LifespanResultT]): self._cache = TimedCache( expiration=datetime.timedelta(seconds=cache_expiration_seconds or 0) ) - self._mounted_servers: list[MountedServer] = [] self._additional_http_routes: list[BaseRoute] = [] self._tool_manager = ToolManager( duplicate_behavior=on_duplicate_tools, @@ -442,7 +441,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[mcp.types.ListToolsRequest], ) -> list[Tool]: - tools = await self._tool_manager._list_tools() # type: ignore[reportPrivateUsage] + tools = await self._tool_manager.list_tools() # type: ignore[reportPrivateUsage] mcp_tools: list[Tool] = [] for tool in tools: @@ -483,7 +482,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[dict[str, Any]], ) -> list[Resource]: - resources = await self._resource_manager._list_resources() # type: ignore[reportPrivateUsage] + resources = await self._resource_manager.list_resources() # type: ignore[reportPrivateUsage] mcp_resources: list[Resource] = [] for resource in resources: @@ -525,7 +524,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[dict[str, Any]], ) -> list[ResourceTemplate]: - templates = await self._resource_manager._list_resource_templates() + templates = await self._resource_manager.list_resource_templates() mcp_templates: list[ResourceTemplate] = [] for template in templates: @@ -564,7 +563,7 @@ class FastMCP(Generic[LifespanResultT]): async def _handler( context: MiddlewareContext[mcp.types.ListPromptsRequest], ) -> list[Prompt]: - prompts = await self._prompt_manager._list_prompts() # type: ignore[reportPrivateUsage] + prompts = await self._prompt_manager.list_prompts() # type: ignore[reportPrivateUsage] mcp_prompts: list[Prompt] = [] for prompt in prompts: @@ -902,23 +901,23 @@ class FastMCP(Generic[LifespanResultT]): enabled=enabled, ) - def add_resource(self, resource: Resource, key: str | None = None) -> None: + def add_resource(self, resource: Resource) -> None: """Add a resource to the server. Args: resource: A Resource instance to add """ - self._resource_manager.add_resource(resource, key=key) + self._resource_manager.add_resource(resource) self._cache.clear() - def add_template(self, template: ResourceTemplate, key: str | None = None) -> None: + def add_template(self, template: ResourceTemplate) -> None: """Add a resource template to the server. Args: template: A ResourceTemplate instance to add """ - self._resource_manager.add_template(template, key=key) + self._resource_manager.add_template(template) def add_resource_fn( self, @@ -1670,10 +1669,8 @@ class FastMCP(Generic[LifespanResultT]): # Import tools from the server for key, tool in (await server.get_tools()).items(): if prefix: - tool_key = f"{prefix}_{key}" - else: - tool_key = key - self._tool_manager.add_tool(tool, key=tool_key) + tool = tool.with_key(f"{prefix}_{key}") + self._tool_manager.add_tool(tool) # Import resources and templates from the server for key, resource in (await server.get_resources()).items(): @@ -1681,35 +1678,27 @@ class FastMCP(Generic[LifespanResultT]): resource_key = add_resource_prefix( key, prefix, self.resource_prefix_format ) - else: - resource_key = key - self._resource_manager.add_resource(resource, key=resource_key) + resource = resource.with_key(resource_key) + self._resource_manager.add_resource(resource) for key, template in (await server.get_resource_templates()).items(): if prefix: template_key = add_resource_prefix( key, prefix, self.resource_prefix_format ) - else: - template_key = key - self._resource_manager.add_template(template, key=template_key) + template = template.with_key(template_key) + self._resource_manager.add_template(template) # Import prompts from the server for key, prompt in (await server.get_prompts()).items(): if prefix: - prompt_key = f"{prefix}_{key}" - else: - prompt_key = key - self._prompt_manager.add_prompt(prompt, key=prompt_key) + prompt = prompt.with_key(f"{prefix}_{key}") + self._prompt_manager.add_prompt(prompt) if prefix: - logger.info(f"Imported server {server.name} with prefix '{prefix}'") - logger.debug(f"Imported tools with prefix '{prefix}_'") - logger.debug(f"Imported resources and templates with prefix '{prefix}/'") - logger.debug(f"Imported prompts with prefix '{prefix}_'") + logger.debug(f"Imported server {server.name} with prefix '{prefix}'") else: - logger.info(f"Imported server {server.name}") - logger.debug("Imported tools, resources, templates, and prompts") + logger.debug(f"Imported server {server.name}") self._cache.clear() diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index 438480882..0d51ba32e 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -64,7 +64,7 @@ class ToolManager: child_results = await mounted.server._list_tools() else: # Use the manager-to-manager unfiltered path - child_results = await mounted.server._tool_manager._list_tools() + child_results = await mounted.server._tool_manager.list_tools() # The combination logic is the same for both paths child_dict = {t.key: t for t in child_results} @@ -103,7 +103,7 @@ class ToolManager: """ return await self._load_tools(via_server=False) - async def _list_tools(self) -> list[Tool]: + async def list_tools(self) -> list[Tool]: """ Lists all tools, applying protocol filtering. """ @@ -139,22 +139,21 @@ class ToolManager: ) return self.add_tool(tool) - def add_tool(self, tool: Tool, key: str | None = None) -> Tool: + def add_tool(self, tool: Tool) -> Tool: """Register a tool with the server.""" - key = key or tool.name - existing = self._tools.get(key) + existing = self._tools.get(tool.key) if existing: if self.duplicate_behavior == "warn": - logger.warning(f"Tool already exists: {key}") - self._tools[key] = tool + logger.warning(f"Tool already exists: {tool.key}") + self._tools[tool.key] = tool elif self.duplicate_behavior == "replace": - self._tools[key] = tool + self._tools[tool.key] = tool elif self.duplicate_behavior == "error": - raise ValueError(f"Tool already exists: {key}") + raise ValueError(f"Tool already exists: {tool.key}") elif self.duplicate_behavior == "ignore": return existing else: - self._tools[key] = tool + self._tools[tool.key] = tool return tool def remove_tool(self, key: str) -> None: diff --git a/tests/client/test_client.py b/tests/client/test_client.py index fcb0e14ad..d975ed77b 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -833,7 +833,12 @@ class TestInferTransport: transport = infer_transport(config) assert isinstance(transport, MCPConfigTransport) assert isinstance(transport.transport, FastMCPTransport) - assert len(cast(FastMCP, transport.transport.server)._mounted_servers) == 2 + assert ( + len( + cast(FastMCP, transport.transport.server)._tool_manager._mounted_servers + ) + == 2 + ) def test_infer_fastmcp_server(self, fastmcp_server): """FastMCP server instances should infer to FastMCPTransport.""" diff --git a/tests/deprecated/test_resource_prefixes.py b/tests/deprecated/test_resource_prefixes.py index 85ffdcf7d..44390f80f 100644 --- a/tests/deprecated/test_resource_prefixes.py +++ b/tests/deprecated/test_resource_prefixes.py @@ -99,7 +99,7 @@ async def test_import_server_with_legacy_prefixes(): await main_server.import_server("sub", sub_server) # type: ignore[arg-type] # Check that the resource is prefixed using the legacy format - resources = main_server._resource_manager.get_resources() + resources = await main_server.get_resources() # In legacy format, the key would be "sub+resource://test" assert "sub+resource://test" in resources diff --git a/tests/server/openapi/test_openapi.py b/tests/server/openapi/test_openapi.py index 393dec27e..9afb1d656 100644 --- a/tests/server/openapi/test_openapi.py +++ b/tests/server/openapi/test_openapi.py @@ -472,7 +472,7 @@ class TestTagTransfer: ): """Test that tags from OpenAPI routes are correctly transferred to Tools.""" # Get internal tools directly (not the public API which returns MCP.Content) - tools = await fastmcp_openapi_server._tool_manager._list_tools() + tools = await fastmcp_openapi_server._tool_manager.list_tools() # Find the create_user and update_user_name tools create_user_tool = next( @@ -1546,7 +1546,7 @@ class TestFastAPIDescriptionPropagation: print(f" Template: {name}, Name attribute: {template.name}") print("\nDEBUG - Tools created:") - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() for tool in tools: print(f" Tool: {tool.name}") @@ -1779,7 +1779,7 @@ class TestReprMethods: async def test_openapi_tool_repr(self, fastmcp_openapi_server: FastMCPOpenAPI): """Test that OpenAPITool's __repr__ method works without recursion errors.""" - tools = await fastmcp_openapi_server._tool_manager._list_tools() + tools = await fastmcp_openapi_server._tool_manager.list_tools() tool = next(iter(tools)) # Verify repr doesn't cause recursion and contains expected elements @@ -1854,7 +1854,7 @@ class TestEnumHandling: ) # Get the tools from the server - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() # Find the read_item tool read_item_tool = next((t for t in tools if t.name == "read_item_items"), None) @@ -1947,7 +1947,7 @@ class TestRouteMapWildcard: ) # All operations should be mapped to tools - tools = await mcp._tool_manager._list_tools() + tools = await mcp._tool_manager.list_tools() tool_names = {tool.name for tool in tools} # Check that all 4 operations became tools @@ -2264,7 +2264,7 @@ class TestMCPNames: ) # Check tools use custom names - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() tool_names = {tool.name for tool in tools} assert "admin_create_user" in tool_names @@ -2294,7 +2294,7 @@ class TestMCPNames: route_maps=GET_ROUTE_MAPS, ) - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() tool_names = {tool.name for tool in tools} templates_dict = await server._resource_manager.get_resource_templates() @@ -2350,7 +2350,7 @@ class TestMCPNames: # Check all component types all_names = [] - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() all_names.extend(tool.name for tool in tools) resources_dict = await server._resource_manager.get_resources() @@ -2383,7 +2383,7 @@ class TestMCPNames: mcp_names=mcp_names, ) - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() tool_names = {tool.name for tool in tools} assert "openapi_user_list" in tool_names @@ -2415,7 +2415,7 @@ class TestMCPNames: mcp_names=mcp_names, ) - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() tool_names = {tool.name for tool in tools} assert "fastapi_create_user" in tool_names @@ -2518,7 +2518,7 @@ class TestRouteMapMCPTags: ) # Get the POST tool - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() create_user_tool = next((t for t in tools if "create_user" in t.name), None) assert create_user_tool is not None, "create_user tool not found" @@ -2634,7 +2634,7 @@ class TestRouteMapMCPTags: ) # Check tool tags - tools = await server._tool_manager._list_tools() + tools = await server._tool_manager.list_tools() create_tool = next((t for t in tools if "create_user" in t.name), None) assert create_tool is not None assert "write-operation" in create_tool.tags diff --git a/tests/server/test_import_server.py b/tests/server/test_import_server.py index 3db05da93..958fbdc83 100644 --- a/tests/server/test_import_server.py +++ b/tests/server/test_import_server.py @@ -25,7 +25,7 @@ async def test_import_basic_functionality(): assert "sub_tool" in sub_app._tool_manager._tools # Verify the original tool still exists in the sub-app - tool = main_app._tool_manager.get_tool("sub_sub_tool") + tool = await main_app._tool_manager.get_tool("sub_sub_tool") assert tool is not None assert tool.name == "sub_tool" assert isinstance(tool, FunctionTool) @@ -203,7 +203,7 @@ async def test_tool_custom_name_preserved_when_imported(): await main_app.import_server(api_app, "api") # Check that the tool is accessible by its prefixed name - tool = main_app._tool_manager.get_tool("api_get_data") + tool = await main_app._tool_manager.get_tool("api_get_data") assert tool is not None # Check that the function name is preserved @@ -239,7 +239,7 @@ async def test_first_level_importing_with_custom_name(): await service_app.import_server(provider_app, "provider") # Tool is accessible in the service app with the first prefix - tool = service_app._tool_manager.get_tool("provider_compute") + tool = await service_app._tool_manager.get_tool("provider_compute") assert tool is not None assert isinstance(tool, FunctionTool) assert tool.fn.__name__ == "calculate_value" @@ -259,7 +259,7 @@ async def test_nested_importing_preserves_prefixes(): await main_app.import_server(service_app, "service") # Tool is accessible in the main app with both prefixes - tool = main_app._tool_manager.get_tool("service_provider_compute") + tool = await main_app._tool_manager.get_tool("service_provider_compute") assert tool is not None diff --git a/tests/server/test_mount.py b/tests/server/test_mount.py index befaccfdd..46658395b 100644 --- a/tests/server/test_mount.py +++ b/tests/server/test_mount.py @@ -862,7 +862,7 @@ class TestAsProxyKwarg: sub = FastMCP("Sub") mcp.mount(sub, "sub") - assert mcp._mounted_servers[0].server is sub + assert mcp._tool_manager._mounted_servers[0].server is sub async def test_as_proxy_false(self): mcp = FastMCP("Main") @@ -870,7 +870,7 @@ class TestAsProxyKwarg: mcp.mount(sub, "sub", as_proxy=False) - assert mcp._mounted_servers[0].server is sub + assert mcp._tool_manager._mounted_servers[0].server is sub async def test_as_proxy_true(self): mcp = FastMCP("Main") @@ -878,8 +878,8 @@ class TestAsProxyKwarg: mcp.mount(sub, "sub", as_proxy=True) - assert mcp._mounted_servers[0].server is not sub - assert isinstance(mcp._mounted_servers[0].server, FastMCPProxy) + assert mcp._tool_manager._mounted_servers[0].server is not sub + assert isinstance(mcp._tool_manager._mounted_servers[0].server, FastMCPProxy) async def test_as_proxy_defaults_true_if_lifespan(self): @asynccontextmanager @@ -891,8 +891,8 @@ class TestAsProxyKwarg: mcp.mount(sub, "sub") - assert mcp._mounted_servers[0].server is not sub - assert isinstance(mcp._mounted_servers[0].server, FastMCPProxy) + assert mcp._tool_manager._mounted_servers[0].server is not sub + assert isinstance(mcp._tool_manager._mounted_servers[0].server, FastMCPProxy) async def test_as_proxy_ignored_for_proxy_mounts_default(self): mcp = FastMCP("Main") @@ -901,7 +901,7 @@ class TestAsProxyKwarg: mcp.mount(sub_proxy, "sub") - assert mcp._mounted_servers[0].server is sub_proxy + assert mcp._tool_manager._mounted_servers[0].server is sub_proxy async def test_as_proxy_ignored_for_proxy_mounts_false(self): mcp = FastMCP("Main") @@ -910,7 +910,7 @@ class TestAsProxyKwarg: mcp.mount(sub_proxy, "sub", as_proxy=False) - assert mcp._mounted_servers[0].server is sub_proxy + assert mcp._tool_manager._mounted_servers[0].server is sub_proxy async def test_as_proxy_ignored_for_proxy_mounts_true(self): mcp = FastMCP("Main") @@ -919,7 +919,7 @@ class TestAsProxyKwarg: mcp.mount(sub_proxy, "sub", as_proxy=True) - assert mcp._mounted_servers[0].server is sub_proxy + assert mcp._tool_manager._mounted_servers[0].server is sub_proxy async def test_as_proxy_mounts_still_have_live_link(self): mcp = FastMCP("Main") @@ -950,11 +950,14 @@ class TestAsProxyKwarg: def hello(): return "hi" - mcp.mount(sub, "sub", as_proxy=True) + mcp.mount(sub, as_proxy=True) assert lifespan_check == [] async with Client(mcp) as client: - await client.call_tool("sub_hello", {}) + await client.call_tool("hello", {}) - assert lifespan_check == ["start"] + assert len(lifespan_check) > 0 + # in the present implementation the sub server will be invoked 3 times + # to call its tool + assert lifespan_check == ["start", "start", "start"] diff --git a/tests/server/test_server.py b/tests/server/test_server.py index 21b13b841..d255ad76c 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -282,7 +282,7 @@ class TestToolDecorator: return x * 2 # Verify the tags were set correctly - tools = await mcp._tool_manager._list_tools() + tools = await mcp._tool_manager.list_tools() assert len(tools) == 1 assert tools[0].tags == {"example", "test-tag"} From 0beb79d11a1d3b48eee1d09866f6653ff6453b75 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 15:28:55 -0400 Subject: [PATCH 15/17] Minor updates --- docs/docs.json | 1 + docs/servers/fastmcp.mdx | 2 +- docs/servers/middleware.mdx | 40 ++++++++++++---------- src/fastmcp/server/middleware.py | 4 +-- src/fastmcp/server/server.py | 18 +++++----- tests/server/middleware/test_middleware.py | 4 +-- 6 files changed, 37 insertions(+), 32 deletions(-) diff --git a/docs/docs.json b/docs/docs.json index 4b0b177d2..0c0b83e53 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -78,6 +78,7 @@ "servers/auth/bearer" ] }, + "servers/middleware", "servers/openapi", "servers/proxy", "servers/composition", diff --git a/docs/servers/fastmcp.mdx b/docs/servers/fastmcp.mdx index ce262e535..12aa08fd8 100644 --- a/docs/servers/fastmcp.mdx +++ b/docs/servers/fastmcp.mdx @@ -193,7 +193,7 @@ def hello(): return "hi" # Mount directly -main.mount("sub", sub) +main.mount(sub, prefix="sub") ``` ## Proxying Servers diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx index 690026e5f..e0be4b311 100644 --- a/docs/servers/middleware.mdx +++ b/docs/servers/middleware.mdx @@ -1,18 +1,22 @@ --- title: MCP Middleware sidebarTitle: Middleware -description: Add cross-cutting functionality to your MCP server with middleware that can inspect, modify, and respond to all MCP requests and responses. +description: Add custom functionality to your MCP server with middleware that can inspect, modify, and respond to all MCP requests and responses. icon: layers --- import { VersionBadge } from "/snippets/version-badge.mdx" - + MCP middleware is a powerful concept that allows you to add cross-cutting functionality to your FastMCP server. Unlike traditional web middleware, MCP middleware is designed specifically for the Model Context Protocol, providing hooks for different types of MCP operations like tool calls, resource reads, and prompt requests. - + MCP middleware is a FastMCP-specific concept and is not part of the official MCP protocol specification. This middleware system is designed to work with FastMCP servers and may not be compatible with other MCP implementations. + + + +MCP middleware is a brand new concept and may be subject to breaking changes in future versions. ## What is MCP Middleware? @@ -58,13 +62,13 @@ The middleware hook system is designed to be extensible. As FastMCP evolves and ### Basic Middleware Structure -MCP middleware is implemented by subclassing the `MCPMiddleware` base class and overriding the hooks you need: +MCP middleware is implemented by subclassing the `Middleware` base class and overriding the hooks you need: ```python from fastmcp import FastMCP -from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.server.middleware import Middleware, MiddlewareContext -class LoggingMiddleware(MCPMiddleware): +class LoggingMiddleware(Middleware): """Middleware that logs all MCP operations.""" async def on_message(self, context: MiddlewareContext, call_next): @@ -97,7 +101,7 @@ mcp.add_middleware(LoggingMiddleware()) The `MiddlewareContext` object provides access to information about the current request: ```python -class InspectionMiddleware(MCPMiddleware): +class InspectionMiddleware(Middleware): async def on_request(self, context: MiddlewareContext, call_next): # Access request information method = context.method # e.g., "tools/call" @@ -120,7 +124,7 @@ Each middleware hook receives a `MiddlewareContext` and a `call_next` function. 3. **Operation-specific hooks**: Called for specific MCP operations ```python -class ComprehensiveMiddleware(MCPMiddleware): +class ComprehensiveMiddleware(Middleware): async def on_message(self, context: MiddlewareContext, call_next): """Called for ALL messages (requests and notifications).""" print(f"Message: {context.method}") @@ -143,10 +147,10 @@ class ComprehensiveMiddleware(MCPMiddleware): ### Authentication Middleware ```python -from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.exceptions import ToolError -class AuthenticationMiddleware(MCPMiddleware): +class AuthenticationMiddleware(Middleware): def __init__(self, required_token: str): self.required_token = required_token @@ -184,7 +188,7 @@ mcp.add_middleware(AuthenticationMiddleware("secret-token-123")) import time import logging -class PerformanceMiddleware(MCPMiddleware): +class PerformanceMiddleware(Middleware): def __init__(self): self.logger = logging.getLogger("performance") @@ -214,7 +218,7 @@ class PerformanceMiddleware(MCPMiddleware): ### Request/Response Transformation Middleware ```python -class TransformationMiddleware(MCPMiddleware): +class TransformationMiddleware(Middleware): async def on_call_tool(self, context: MiddlewareContext, call_next): """Transform tool arguments and results.""" @@ -253,7 +257,7 @@ import asyncio from collections import defaultdict from datetime import datetime, timedelta -class RateLimitMiddleware(MCPMiddleware): +class RateLimitMiddleware(Middleware): def __init__(self, max_requests: int = 100, window_minutes: int = 1): self.max_requests = max_requests self.window = timedelta(minutes=window_minutes) @@ -327,7 +331,7 @@ mcp.add_middleware(LoggingMiddleware()) ### Conditional Middleware ```python -class ConditionalMiddleware(MCPMiddleware): +class ConditionalMiddleware(Middleware): def __init__(self, condition_func): self.should_process = condition_func @@ -352,7 +356,7 @@ mcp.add_middleware(ConditionalMiddleware(only_expensive_tools)) ### Middleware with State ```python -class StatefulMiddleware(MCPMiddleware): +class StatefulMiddleware(Middleware): def __init__(self): self.call_count = 0 self.tools_used = set() @@ -377,7 +381,7 @@ class StatefulMiddleware(MCPMiddleware): ### Error Handling Middleware ```python -class ErrorHandlingMiddleware(MCPMiddleware): +class ErrorHandlingMiddleware(Middleware): async def on_call_tool(self, context: MiddlewareContext, call_next): """Provide consistent error handling.""" try: @@ -400,7 +404,7 @@ class ErrorHandlingMiddleware(MCPMiddleware): 3. **Cache when possible**: Store frequently accessed data to avoid repeated lookups ```python -class EfficientMiddleware(MCPMiddleware): +class EfficientMiddleware(Middleware): def __init__(self): self._cache = {} @@ -428,7 +432,7 @@ class EfficientMiddleware(MCPMiddleware): 3. **Use `ToolError` for client-facing errors**: Keep internal errors internal ```python -class RobustMiddleware(MCPMiddleware): +class RobustMiddleware(Middleware): async def on_request(self, context: MiddlewareContext, call_next): """Robust error handling example.""" try: diff --git a/src/fastmcp/server/middleware.py b/src/fastmcp/server/middleware.py index 2363d41be..d31a2a820 100644 --- a/src/fastmcp/server/middleware.py +++ b/src/fastmcp/server/middleware.py @@ -104,7 +104,7 @@ class MiddlewareContext(Generic[T]): def make_middleware_wrapper( - middleware: MCPMiddleware, call_next: CallNext[T, R] + middleware: Middleware, call_next: CallNext[T, R] ) -> CallNext[T, R]: """Create a wrapper that applies a single middleware to a context. The closure bakes in the middleware and call_next function, so it can be @@ -116,7 +116,7 @@ def make_middleware_wrapper( return wrapper -class MCPMiddleware: +class Middleware: """Base class for FastMCP middleware with dispatching hooks.""" async def __call__( diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index a6c6321ab..d2df75e6a 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -35,7 +35,7 @@ from mcp.types import Resource as MCPResource from mcp.types import ResourceTemplate as MCPResourceTemplate from mcp.types import Tool as MCPTool from pydantic import AnyUrl -from starlette.middleware import Middleware +from starlette.middleware import Middleware as ASGIMiddleware from starlette.requests import Request from starlette.responses import Response from starlette.routing import BaseRoute, Route @@ -55,7 +55,7 @@ from fastmcp.server.http import ( create_sse_app, create_streamable_http_app, ) -from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.settings import Settings from fastmcp.tools import ToolManager from fastmcp.tools.tool import FunctionTool, Tool @@ -118,7 +118,7 @@ class FastMCP(Generic[LifespanResultT]): *, version: str | None = None, auth: OAuthProvider | None = None, - middleware: list[MCPMiddleware] | None = None, + middleware: list[Middleware] | None = None, lifespan: ( Callable[ [FastMCP[LifespanResultT]], @@ -335,7 +335,7 @@ class FastMCP(Generic[LifespanResultT]): chain = partial(mw, call_next=chain) return await chain(context) - def add_middleware(self, middleware: MCPMiddleware) -> None: + def add_middleware(self, middleware: Middleware) -> None: self.middleware.append(middleware) async def get_tools(self) -> dict[str, Tool]: @@ -917,7 +917,7 @@ class FastMCP(Generic[LifespanResultT]): Args: template: A ResourceTemplate instance to add """ - self._resource_manager.add_template(template) + self._resource_manager.add_template(template, key=key) def add_resource_fn( self, @@ -1260,7 +1260,7 @@ class FastMCP(Generic[LifespanResultT]): log_level: str | None = None, path: str | None = None, uvicorn_config: dict[str, Any] | None = None, - middleware: list[Middleware] | None = None, + middleware: list[ASGIMiddleware] | None = None, ) -> None: """Run the server using HTTP transport. @@ -1331,7 +1331,7 @@ class FastMCP(Generic[LifespanResultT]): self, path: str | None = None, message_path: str | None = None, - middleware: list[Middleware] | None = None, + middleware: list[ASGIMiddleware] | None = None, ) -> StarletteWithLifespan: """ Create a Starlette app for the SSE server. @@ -1361,7 +1361,7 @@ class FastMCP(Generic[LifespanResultT]): def streamable_http_app( self, path: str | None = None, - middleware: list[Middleware] | None = None, + middleware: list[ASGIMiddleware] | None = None, ) -> StarletteWithLifespan: """ Create a Starlette app for the StreamableHTTP server. @@ -1382,7 +1382,7 @@ class FastMCP(Generic[LifespanResultT]): def http_app( self, path: str | None = None, - middleware: list[Middleware] | None = None, + middleware: list[ASGIMiddleware] | None = None, json_response: bool | None = None, stateless_http: bool | None = None, transport: Literal["streamable-http", "sse"] = "streamable-http", diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index cd8b256bf..7327e9ad8 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -7,7 +7,7 @@ import pytest from fastmcp import Client, FastMCP from fastmcp.server.context import Context -from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext +from fastmcp.server.middleware import Middleware, MiddlewareContext @dataclass @@ -18,7 +18,7 @@ class Recording: result: mcp.types.ServerResult | None -class RecordingMiddleware(MCPMiddleware): +class RecordingMiddleware(Middleware): """A middleware that automatically records all method calls.""" def __init__(self, name: str | None = None): From 085c5b765b0278e6e3b81cead96846a4171d0ad6 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 17:06:54 -0400 Subject: [PATCH 16/17] Update docs and tests --- docs/servers/middleware.mdx | 604 ++++++++++------------- src/fastmcp/server/server.py | 3 +- tests/resources/test_resource_manager.py | 16 +- tests/tools/test_tool_manager.py | 5 +- 4 files changed, 271 insertions(+), 357 deletions(-) diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx index e0be4b311..59835a44c 100644 --- a/docs/servers/middleware.mdx +++ b/docs/servers/middleware.mdx @@ -1,7 +1,7 @@ --- title: MCP Middleware sidebarTitle: Middleware -description: Add custom functionality to your MCP server with middleware that can inspect, modify, and respond to all MCP requests and responses. +description: Add cross-cutting functionality to your MCP server with middleware that can inspect, modify, and respond to all MCP requests and responses. icon: layers --- @@ -21,7 +21,7 @@ MCP middleware is a brand new concept and may be subject to breaking changes in ## What is MCP Middleware? -MCP middleware lets you intercept and modify MCP requests and responses as they flow through your server. Unlike traditional HTTP middleware that operates on request/response pairs, MCP middleware is aware of the specific MCP protocol operations and can provide targeted hooks for different types of interactions. +MCP middleware lets you intercept and modify MCP requests and responses as they flow through your server. Think of it as a pipeline where each piece of middleware can inspect what's happening, make changes, and then pass control to the next middleware in the chain. Common use cases for MCP middleware include: - **Authentication and Authorization**: Verify client permissions before executing operations @@ -31,17 +31,60 @@ Common use cases for MCP middleware include: - **Caching**: Store frequently requested data to improve performance - **Error Handling**: Provide consistent error responses across your server -## How MCP Middleware Works +## How Middleware Works -MCP middleware operates on a pipeline model where each middleware can: +FastMCP middleware operates on a pipeline model. When a request comes in, it flows through your middleware in the order they were added to the server. Each middleware can: 1. **Inspect the incoming request** and its context 2. **Modify the request** before passing it to the next middleware or handler -3. **Execute the next middleware/handler** in the chain +3. **Execute the next middleware/handler** in the chain by calling `call_next()` 4. **Inspect and modify the response** before returning it 5. **Handle errors** that occur during processing -The middleware system provides specialized hooks for different MCP operations: +The key insight is that middleware forms a chain where each piece decides whether to continue processing or stop the chain entirely. + +If you're familiar with ASGI middleware, the basic structure of FastMCP middleware will feel familiar. At its core, middleware is a callable class that receives a context object containing information about the current JSON-RPC message and a handler function to continue the middleware chain. + +It's important to understand that MCP operates on the [JSON-RPC specification](https://spec.modelcontextprotocol.io/specification/basic/transports/). While FastMCP presents requests and responses in a familiar way, these are fundamentally JSON-RPC messages, not HTTP request/response pairs like you might be used to in web applications. FastMCP middleware works with all [transport types](/clients/transports), including local stdio transport and HTTP transports, though not all middleware implementations are compatible across all transports (e.g., middleware that inspects HTTP headers won't work with stdio transport). + +The most fundamental way to implement middleware is by overriding the `__call__` method on the `Middleware` base class: + +```python +from fastmcp.server.middleware import Middleware, MiddlewareContext + +class RawMiddleware(Middleware): + async def __call__(self, context: MiddlewareContext, call_next): + # This method receives ALL messages regardless of type + print(f"Raw middleware processing: {context.method}") + result = await call_next(context) + print(f"Raw middleware completed: {context.method}") + return result +``` + +This gives you complete control over every message that flows through your server, but requires you to handle all message types manually. + +## Middleware Hooks + +To make it easier for users to target specific types of messages, FastMCP middleware provides a variety of specialized hooks. Instead of implementing the raw `__call__` method, you can override specific hook methods that are called only for certain types of operations, allowing you to target exactly the level of specificity you need for your middleware logic. + +### Hook Hierarchy and Execution Order + +FastMCP provides multiple hooks that are called with varying levels of specificity. Understanding this hierarchy is crucial for effective middleware design. + +When a request comes in, **multiple hooks may be called for the same request**, going from general to specific: + +1. **`on_message`** - Called for ALL MCP messages (both requests and notifications) +2. **`on_request` or `on_notification`** - Called based on the message type +3. **Operation-specific hooks** - Called for specific MCP operations like `on_call_tool` + +For example, when a client calls a tool, your middleware will receive **three separate hook calls**: +1. First: `on_message` (because it's any MCP message) +2. Second: `on_request` (because tool calls expect responses) +3. Third: `on_call_tool` (because it's specifically a tool execution) + +This hierarchy allows you to target your middleware logic with the right level of specificity. Use `on_message` for broad concerns like logging, `on_request` for authentication, and `on_call_tool` for tool-specific logic like performance monitoring. + +### Available Hooks - `on_message`: Called for all MCP messages (requests and notifications) - `on_request`: Called specifically for MCP requests (that expect responses) @@ -54,15 +97,152 @@ The middleware system provides specialized hooks for different MCP operations: - `on_list_resource_templates`: Called when listing resource templates - `on_list_prompts`: Called when listing available prompts - -The middleware hook system is designed to be extensible. As FastMCP evolves and new MCP operations are added, additional hooks will be introduced to provide fine-grained control over new functionality. - +## Component Access in Middleware + +Understanding how to access component information (tools, resources, prompts) in middleware is crucial for building powerful middleware functionality. The access patterns differ significantly between listing operations and execution operations. + +### Listing Operations vs Execution Operations + +FastMCP middleware handles two types of operations differently: + +**Listing Operations** (`on_list_tools`, `on_list_resources`, `on_list_prompts`, etc.): +- Middleware receives **FastMCP component objects** with full metadata +- These objects include FastMCP-specific properties like `tags` that aren't part of the MCP specification +- The result contains complete component information before it's converted to MCP format +- Tags and other metadata are stripped when finally returned to the MCP client + +**Execution Operations** (`on_call_tool`, `on_read_resource`, `on_get_prompt`): +- Middleware runs **before** the component is executed +- The middleware result is either the execution result or an error if the component wasn't found +- Component metadata isn't directly available in the hook parameters + +### Accessing Component Metadata During Execution + +If you need to check component properties (like tags) during execution operations, use the FastMCP server instance available through the context: + +```python +from fastmcp.server.middleware import Middleware, MiddlewareContext +from fastmcp.exceptions import ToolError + +class TagBasedMiddleware(Middleware): + async def on_call_tool(self, context: MiddlewareContext, call_next): + # Access the tool object to check its metadata + if context.fastmcp_context: + try: + tool = await context.fastmcp_context.fastmcp.get_tool(context.message.name) + + # Check if this tool has a "private" tag + if "private" in tool.tags: + raise ToolError("Access denied: private tool") + + # Check if tool is enabled + if not tool.enabled: + raise ToolError("Tool is currently disabled") + + except Exception: + # Tool not found or other error - let execution continue + # and handle the error naturally + pass + + return await call_next(context) +``` + +The same pattern works for resources and prompts: + +```python +from fastmcp.server.middleware import Middleware, MiddlewareContext +from fastmcp.exceptions import ResourceError, PromptError + +class ComponentAccessMiddleware(Middleware): + async def on_read_resource(self, context: MiddlewareContext, call_next): + if context.fastmcp_context: + try: + resource = await context.fastmcp_context.fastmcp.get_resource(context.message.uri) + if "restricted" in resource.tags: + raise ResourceError("Access denied: restricted resource") + except Exception: + pass + return await call_next(context) + + async def on_get_prompt(self, context: MiddlewareContext, call_next): + if context.fastmcp_context: + try: + prompt = await context.fastmcp_context.fastmcp.get_prompt(context.message.name) + if not prompt.enabled: + raise PromptError("Prompt is currently disabled") + except Exception: + pass + return await call_next(context) +``` + +### Working with Listing Results + +For listing operations, you can inspect and modify the FastMCP components directly: + +```python +from fastmcp.server.middleware import Middleware, MiddlewareContext, ListToolsResult + +class ListingFilterMiddleware(Middleware): + async def on_list_tools(self, context: MiddlewareContext, call_next): + result = await call_next(context) + + # Filter out tools with "private" tag + filtered_tools = { + name: tool for name, tool in result.tools.items() + if "private" not in tool.tags + } + + # Return modified result + return ListToolsResult(tools=filtered_tools) +``` + +This filtering happens before the components are converted to MCP format and returned to the client, so the tags (which are FastMCP-specific) are naturally stripped in the final response. + +### Anatomy of a Hook + +Every middleware hook follows the same pattern. Let's examine the `on_message` hook to understand the structure: + +```python +async def on_message(self, context: MiddlewareContext, call_next): + # 1. Pre-processing: Inspect and optionally modify the request + print(f"Processing {context.method}") + + # 2. Chain continuation: Call the next middleware/handler + result = await call_next(context) + + # 3. Post-processing: Inspect and optionally modify the response + print(f"Completed {context.method}") + + # 4. Return the result (potentially modified) + return result +``` + +### Hook Parameters + +Every hook receives two parameters: + +1. **`context: MiddlewareContext`** - Contains information about the current request: + - `context.method` - The MCP method name (e.g., "tools/call") + - `context.source` - Where the request came from ("client" or "server") + - `context.type` - Message type ("request" or "notification") + - `context.message` - The MCP message data + - `context.timestamp` - When the request was received + - `context.fastmcp_context` - FastMCP Context object (if available) + +2. **`call_next`** - A function that continues the middleware chain. You **must** call this to proceed, unless you want to stop processing entirely. + +### Control Flow + +You have complete control over the request flow: +- **Continue processing**: Call `await call_next(context)` to proceed +- **Modify the request**: Change the context before calling `call_next` +- **Modify the response**: Change the result after calling `call_next` +- **Stop the chain**: Don't call `call_next` (rarely needed) +- **Handle errors**: Wrap `call_next` in try/catch blocks ## Creating Middleware -### Basic Middleware Structure - -MCP middleware is implemented by subclassing the `Middleware` base class and overriding the hooks you need: +FastMCP middleware is implemented by subclassing the `Middleware` base class and overriding the hooks you need. You only need to implement the hooks that are relevant to your use case. ```python from fastmcp import FastMCP @@ -75,77 +255,85 @@ class LoggingMiddleware(Middleware): """Called for all MCP messages.""" print(f"Processing {context.method} from {context.source}") - # Call the next middleware/handler in the chain result = await call_next(context) print(f"Completed {context.method}") return result - - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Called specifically for tool calls.""" - tool_name = context.message.name - print(f"Calling tool: {tool_name}") - - result = await call_next(context) - - print(f"Tool {tool_name} completed") - return result # Add middleware to your server mcp = FastMCP("MyServer") mcp.add_middleware(LoggingMiddleware()) ``` -### Middleware Context +This creates a basic logging middleware that will print information about every request that flows through your server. -The `MiddlewareContext` object provides access to information about the current request: +## Adding Middleware to Your Server + +### Single Middleware + +Adding middleware to your server is straightforward: ```python -class InspectionMiddleware(Middleware): - async def on_request(self, context: MiddlewareContext, call_next): - # Access request information - method = context.method # e.g., "tools/call" - source = context.source # "client" or "server" - message_type = context.type # "request" or "notification" - timestamp = context.timestamp # When the request was received - message = context.message # The actual MCP message - fastmcp_context = context.fastmcp_context # FastMCP Context object (if available) - - # Continue processing - return await call_next(context) +mcp = FastMCP("MyServer") +mcp.add_middleware(LoggingMiddleware()) ``` -### Middleware Hooks +### Multiple Middleware -Each middleware hook receives a `MiddlewareContext` and a `call_next` function. The hooks are organized in a hierarchy: - -1. **`on_message`**: The broadest hook, called for all MCP messages -2. **`on_request`** / **`on_notification`**: Called based on message type -3. **Operation-specific hooks**: Called for specific MCP operations +Middleware executes in the order it's added to the server. The first middleware added runs first on the way in, and last on the way out: ```python -class ComprehensiveMiddleware(Middleware): - async def on_message(self, context: MiddlewareContext, call_next): - """Called for ALL messages (requests and notifications).""" - print(f"Message: {context.method}") - return await call_next(context) - - async def on_request(self, context: MiddlewareContext, call_next): - """Called only for requests (messages that expect responses).""" - print(f"Request: {context.method}") - return await call_next(context) - - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Called only for tool execution requests.""" - tool_name = context.message.name - print(f"Executing tool: {tool_name}") - return await call_next(context) +mcp = FastMCP("MyServer") + +mcp.add_middleware(AuthenticationMiddleware("secret-token")) +mcp.add_middleware(PerformanceMiddleware()) +mcp.add_middleware(LoggingMiddleware()) ``` -## Middleware Examples +This creates the following execution flow: +1. AuthenticationMiddleware (pre-processing) +2. PerformanceMiddleware (pre-processing) +3. LoggingMiddleware (pre-processing) +4. Actual tool/resource handler +5. LoggingMiddleware (post-processing) +6. PerformanceMiddleware (post-processing) +7. AuthenticationMiddleware (post-processing) + +## Server Composition and Middleware + +When using [Server Composition](/servers/composition) with `mount` or `import_server`, middleware behavior follows these rules: + +1. **Parent server middleware** runs for all requests, including those routed to mounted servers +2. **Mounted server middleware** only runs for requests handled by that specific server +3. **Middleware order** is preserved within each server + +This allows you to create layered middleware architectures where parent servers handle cross-cutting concerns like authentication, while child servers focus on domain-specific middleware. + +```python +# Parent server with middleware +parent = FastMCP("Parent") +parent.add_middleware(AuthenticationMiddleware("token")) + +# Child server with its own middleware +child = FastMCP("Child") +child.add_middleware(LoggingMiddleware()) + +@child.tool +def child_tool() -> str: + return "from child" + +# Mount the child server +parent.mount(child, prefix="child") +``` + +When a client calls "child_tool", the request will flow through the parent's authentication middleware first, then route to the child server where it will go through the child's logging middleware. + +## Examples ### Authentication Middleware +This middleware checks for a valid authorization token on all requests: + ```python from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.exceptions import ToolError @@ -155,12 +343,8 @@ class AuthenticationMiddleware(Middleware): self.required_token = required_token async def on_request(self, context: MiddlewareContext, call_next): - """Verify authentication for all requests.""" - - # Check if this is an HTTP request with headers if hasattr(context, 'fastmcp_context') and context.fastmcp_context: try: - # Access HTTP request if available request = context.fastmcp_context.get_http_request() auth_header = request.headers.get("Authorization") @@ -172,7 +356,6 @@ class AuthenticationMiddleware(Middleware): raise ToolError("Invalid authentication token") except Exception: - # If HTTP request is not available, continue without auth pass return await call_next(context) @@ -184,6 +367,8 @@ mcp.add_middleware(AuthenticationMiddleware("secret-token-123")) ### Performance Monitoring Middleware +This middleware tracks how long tools take to execute: + ```python import time import logging @@ -193,7 +378,6 @@ class PerformanceMiddleware(Middleware): self.logger = logging.getLogger("performance") async def on_call_tool(self, context: MiddlewareContext, call_next): - """Monitor tool execution performance.""" tool_name = context.message.name start_time = time.time() @@ -215,300 +399,22 @@ class PerformanceMiddleware(Middleware): raise ``` -### Request/Response Transformation Middleware +### Request Transformation Middleware + +This middleware adds metadata to tool calls: ```python class TransformationMiddleware(Middleware): async def on_call_tool(self, context: MiddlewareContext, call_next): - """Transform tool arguments and results.""" - - # Access and modify tool arguments if hasattr(context.message, 'arguments'): args = context.message.arguments or {} - - # Example: Add a timestamp to all tool calls args['_middleware_timestamp'] = context.timestamp.isoformat() - # Create a modified context modified_context = context.copy( message=context.message.model_copy(update={'arguments': args}) ) else: modified_context = context - # Execute with modified context - result = await call_next(modified_context) - - # Transform the result if needed - if hasattr(result, 'content'): - # Example: Add metadata to tool results - if result.content and len(result.content) > 0: - original_content = result.content[0].text - enhanced_content = f"[Processed at {context.timestamp}]\n{original_content}" - result.content[0].text = enhanced_content - - return result -``` - -### Rate Limiting Middleware - -```python -import asyncio -from collections import defaultdict -from datetime import datetime, timedelta - -class RateLimitMiddleware(Middleware): - def __init__(self, max_requests: int = 100, window_minutes: int = 1): - self.max_requests = max_requests - self.window = timedelta(minutes=window_minutes) - self.requests = defaultdict(list) # client_id -> [timestamps] - - async def on_request(self, context: MiddlewareContext, call_next): - """Implement rate limiting per client.""" - - # Get client identifier (you may need to implement this based on your auth) - client_id = getattr(context.fastmcp_context, 'client_id', 'anonymous') - - now = datetime.now() - - # Clean old requests - self.requests[client_id] = [ - timestamp for timestamp in self.requests[client_id] - if now - timestamp < self.window - ] - - # Check rate limit - if len(self.requests[client_id]) >= self.max_requests: - raise ToolError( - f"Rate limit exceeded: {self.max_requests} requests per " - f"{self.window.total_seconds()/60:.0f} minutes" - ) - - # Record this request - self.requests[client_id].append(now) - - return await call_next(context) -``` - -## Adding Middleware to Your Server - -### Single Middleware - -```python -from fastmcp import FastMCP - -mcp = FastMCP("MyServer") - -# Add a single middleware instance -logging_middleware = LoggingMiddleware() -mcp.add_middleware(logging_middleware) -``` - -### Multiple Middleware - -Middleware is executed in the order it's added to the server: - -```python -mcp = FastMCP("MyServer") - -# Add multiple middleware - they execute in order -mcp.add_middleware(AuthenticationMiddleware("secret-token")) -mcp.add_middleware(PerformanceMiddleware()) -mcp.add_middleware(LoggingMiddleware()) - -# Request flow: -# 1. AuthenticationMiddleware.on_request() -# 2. PerformanceMiddleware.on_request() -# 3. LoggingMiddleware.on_request() -# 4. Actual tool/resource handler -# 5. LoggingMiddleware response processing -# 6. PerformanceMiddleware response processing -# 7. AuthenticationMiddleware response processing -``` - -## Advanced Patterns - -### Conditional Middleware - -```python -class ConditionalMiddleware(Middleware): - def __init__(self, condition_func): - self.should_process = condition_func - - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Only process certain tools.""" - - if not self.should_process(context.message.name): - # Skip processing for this tool - return await call_next(context) - - # Apply middleware logic - print(f"Processing tool: {context.message.name}") - return await call_next(context) - -# Usage -def only_expensive_tools(tool_name: str) -> bool: - return tool_name in ["complex_analysis", "heavy_computation"] - -mcp.add_middleware(ConditionalMiddleware(only_expensive_tools)) -``` - -### Middleware with State - -```python -class StatefulMiddleware(Middleware): - def __init__(self): - self.call_count = 0 - self.tools_used = set() - - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Track usage statistics.""" - self.call_count += 1 - self.tools_used.add(context.message.name) - - print(f"Total calls: {self.call_count}, Unique tools: {len(self.tools_used)}") - - return await call_next(context) - - def get_stats(self): - return { - "total_calls": self.call_count, - "unique_tools": len(self.tools_used), - "tools_used": list(self.tools_used) - } -``` - -### Error Handling Middleware - -```python -class ErrorHandlingMiddleware(Middleware): - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Provide consistent error handling.""" - try: - return await call_next(context) - except ToolError: - # Re-raise ToolErrors as-is - raise - except Exception as e: - # Log the error and convert to a user-friendly message - logging.error(f"Tool {context.message.name} failed: {e}") - raise ToolError(f"Tool execution failed: {str(e)}") -``` - -## Best Practices - -### Performance Considerations - -1. **Keep middleware lightweight**: Avoid heavy computations in middleware -2. **Use async operations**: Don't block the event loop with synchronous operations -3. **Cache when possible**: Store frequently accessed data to avoid repeated lookups - -```python -class EfficientMiddleware(Middleware): - def __init__(self): - self._cache = {} - - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Example of efficient middleware with caching.""" - - # Check cache first - cache_key = f"{context.message.name}:{hash(str(context.message.arguments))}" - - if cache_key in self._cache: - print("Returning cached result") - return self._cache[cache_key] - - # Execute and cache result - result = await call_next(context) - self._cache[cache_key] = result - - return result -``` - -### Error Handling - -1. **Always call `call_next`**: Unless you're intentionally stopping the chain -2. **Handle exceptions appropriately**: Don't let middleware errors break the entire request -3. **Use `ToolError` for client-facing errors**: Keep internal errors internal - -```python -class RobustMiddleware(Middleware): - async def on_request(self, context: MiddlewareContext, call_next): - """Robust error handling example.""" - try: - # Middleware logic here - return await call_next(context) - except ToolError: - # Client-facing errors should be re-raised - raise - except Exception as e: - # Log internal errors but don't expose details - logging.error(f"Middleware error: {e}") - # Optionally continue without middleware processing - return await call_next(context) -``` - -### Testing Middleware - -```python -import pytest -from fastmcp import FastMCP, Client - -@pytest.mark.asyncio -async def test_logging_middleware(): - """Test middleware functionality.""" - - # Create server with middleware - mcp = FastMCP("TestServer") - logging_middleware = LoggingMiddleware() - mcp.add_middleware(logging_middleware) - - @mcp.tool - def test_tool(x: int) -> int: - return x * 2 - - # Test with client - async with Client(mcp) as client: - result = await client.call_tool("test_tool", {"x": 5}) - assert result == 10 - - # Verify middleware was called - # (You'll need to add tracking to your middleware for testing) -``` - -## Server Composition and Middleware - -When using [Server Composition](/servers/composition) with `mount` or `import_server`, middleware behavior follows these rules: - -1. **Parent server middleware** runs for all requests, including those routed to mounted servers -2. **Mounted server middleware** only runs for requests handled by that specific server -3. **Middleware order** is preserved within each server - -```python -# Parent server with middleware -parent = FastMCP("Parent") -parent.add_middleware(AuthenticationMiddleware("token")) - -# Child server with its own middleware -child = FastMCP("Child") -child.add_middleware(LoggingMiddleware()) - -@child.tool -def child_tool() -> str: - return "from child" - -# Mount the child server -parent.mount(child, prefix="child") - -# Request to "child_tool" will: -# 1. Run parent's AuthenticationMiddleware -# 2. Route to child server -# 3. Run child's LoggingMiddleware -# 4. Execute child_tool -``` - -This allows you to create layered middleware architectures where parent servers handle cross-cutting concerns like authentication, while child servers focus on domain-specific middleware. - - -When debugging middleware issues in composed servers, remember that both parent and child middleware may be executing. Use detailed logging to trace the middleware execution path. - \ No newline at end of file + return await call_next(modified_context) +``` \ No newline at end of file diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index d2df75e6a..d87aacd41 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -42,7 +42,6 @@ from starlette.routing import BaseRoute, Route import fastmcp import fastmcp.server -import fastmcp.server.middleware from fastmcp.exceptions import DisabledError, NotFoundError from fastmcp.prompts import Prompt, PromptManager from fastmcp.prompts.prompt import FunctionPrompt @@ -917,7 +916,7 @@ class FastMCP(Generic[LifespanResultT]): Args: template: A ResourceTemplate instance to add """ - self._resource_manager.add_template(template, key=key) + self._resource_manager.add_template(template) def add_resource_fn( self, diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index 7c1f02066..c458c49b5 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -464,7 +464,9 @@ class TestCustomResourceKeys: fn=get_data, ) - manager.add_resource(resource, key=custom_key) + # Use with_key to create a new resource with the custom key + resource_with_custom_key = resource.with_key(custom_key) + manager.add_resource(resource_with_custom_key) # Resource should be accessible via custom key assert custom_key in manager._resources @@ -489,7 +491,9 @@ class TestCustomResourceKeys: name="test_template", ) - manager.add_template(template, key=custom_key) + # Use with_key to create a new template with the custom key + template_with_custom_key = template.with_key(custom_key) + manager.add_template(template_with_custom_key) # Template should be accessible via custom key assert custom_key in manager._templates @@ -514,7 +518,9 @@ class TestCustomResourceKeys: fn=get_data, ) - manager.add_resource(resource, key=custom_key) + # Use with_key to create a new resource with the custom key + resource_with_custom_key = resource.with_key(custom_key) + manager.add_resource(resource_with_custom_key) # Should be retrievable by the custom key retrieved = await manager.get_resource(custom_key) @@ -541,7 +547,9 @@ class TestCustomResourceKeys: name="custom_greeter", ) - manager.add_template(template, key=custom_key) + # Use with_key to create a new template with the custom key + template_with_custom_key = template.with_key(custom_key) + manager.add_template(template_with_custom_key) # Using a URI that matches the custom key pattern resource = await manager.get_resource("custom://greet/world") diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index 47923d61a..6d6a19952 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -760,8 +760,9 @@ class TestCustomToolNames: # Create a tool with a specific name tool = Tool.from_function(fn, name="my_tool") manager = ToolManager() - # Store it under a different name - manager.add_tool(tool, key="proxy_tool") + # Use with_key to create a new tool with the custom key + tool_with_custom_key = tool.with_key("proxy_tool") + manager.add_tool(tool_with_custom_key) # The tool is accessible under the key stored = await manager.get_tool("proxy_tool") assert stored is not None From dee2f33a8103070c20c2f1a1a964b9b490e094e7 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 19 Jun 2025 17:28:00 -0400 Subject: [PATCH 17/17] Update doc icons --- docs/patterns/tool-transformation.mdx | 1 + docs/servers/middleware.mdx | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/patterns/tool-transformation.mdx b/docs/patterns/tool-transformation.mdx index 49f9f71e0..f736f791c 100644 --- a/docs/patterns/tool-transformation.mdx +++ b/docs/patterns/tool-transformation.mdx @@ -3,6 +3,7 @@ title: Tool Transformation sidebarTitle: Tool Transformation description: Create enhanced tool variants with modified schemas, argument mappings, and custom behavior. icon: wand-magic-sparkles +tag: NEW --- import { VersionBadge } from '/snippets/version-badge.mdx' diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx index 59835a44c..51ddd9327 100644 --- a/docs/servers/middleware.mdx +++ b/docs/servers/middleware.mdx @@ -2,7 +2,8 @@ title: MCP Middleware sidebarTitle: Middleware description: Add cross-cutting functionality to your MCP server with middleware that can inspect, modify, and respond to all MCP requests and responses. -icon: layers +icon: layer-group +tag: NEW --- import { VersionBadge } from "/snippets/version-badge.mdx"