From f7dfd5be9c8252d377b5b8e139b5de2e5b0e9f2f Mon Sep 17 00:00:00 2001 From: hopeful0 Date: Sat, 5 Jul 2025 00:21:57 +0800 Subject: [PATCH] Proxy support advanced MCP features (#1022) * Proxy support advanced MCP features * Optimize ProxyClient testing * Add documentation for ProxyClient --- docs/servers/proxy.mdx | 6 + src/fastmcp/server/proxy.py | 116 +++++++++++ src/fastmcp/server/server.py | 27 +-- tests/server/proxy/test_proxy_client.py | 257 ++++++++++++++++++++++++ tests/server/test_proxy.py | 6 +- 5 files changed, 396 insertions(+), 16 deletions(-) create mode 100644 tests/server/proxy/test_proxy_client.py diff --git a/docs/servers/proxy.mdx b/docs/servers/proxy.mdx index 5ebff6a04..ede0c678b 100644 --- a/docs/servers/proxy.mdx +++ b/docs/servers/proxy.mdx @@ -162,6 +162,12 @@ composite_proxy = FastMCP.as_proxy(config, name="Composite Proxy") # - weather://weather/icons/sunny, calendar://calendar/events/today ``` +## Forwarding Interactions + +`ProxyClient` is a subclass of `Client` that implements a set of default handlers to forward advanced interactions between the backend server and the client connected to the proxy. These handlers receive requests or notifications from the backend server and forward them to the client through the related request context, relaying the response back to the backend server if needed. This setup enables the proxy to support advanced MCP features, including roots, sampling, elicitation, logging, and progress. + +To prevent the forwarding of any interaction, pass `None` to the corresponding handler when creating the `ProxyClient`. Typically, you should use `ProxyClient` to establish a proxy unless there is a specific reason not to. If the `transport` parameter is not an instance of `Client`, `FastMCP.as_proxy()` will automatically instantiate a `ProxyClient`. + ## `FastMCPProxy` Class Internally, `FastMCP.as_proxy()` uses the `FastMCPProxy` class. You generally don't need to interact with this class directly, but it's available if needed. diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 07650fc27..6d7a6613b 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -1,9 +1,12 @@ from __future__ import annotations +from pathlib import Path from typing import TYPE_CHECKING, Any, cast from urllib.parse import quote import mcp.types +from mcp.client.session import ClientSession +from mcp.shared.context import LifespanContextT, RequestContext from mcp.shared.exceptions import McpError from mcp.types import ( METHOD_NOT_FOUND, @@ -14,6 +17,10 @@ from mcp.types import ( from pydantic.networks import AnyUrl from fastmcp.client import Client +from fastmcp.client.elicitation import ElicitResult +from fastmcp.client.logging import LogMessage +from fastmcp.client.roots import RootsList +from fastmcp.client.transports import ClientTransportT from fastmcp.exceptions import NotFoundError, ResourceError, ToolError from fastmcp.prompts import Prompt, PromptMessage from fastmcp.prompts.prompt import PromptArgument @@ -21,10 +28,12 @@ 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.dependencies import get_context from fastmcp.server.server import FastMCP from fastmcp.tools.tool import Tool, ToolResult from fastmcp.tools.tool_manager import ToolManager from fastmcp.utilities.logging import get_logger +from fastmcp.utilities.mcp_config import MCPConfig if TYPE_CHECKING: from fastmcp.server import Context @@ -406,3 +415,110 @@ class FastMCPProxy(FastMCP): self._tool_manager = ProxyToolManager(client=self.client) self._resource_manager = ProxyResourceManager(client=self.client) self._prompt_manager = ProxyPromptManager(client=self.client) + + +async def default_proxy_roots_handler( + context: RequestContext[ClientSession, LifespanContextT], +) -> RootsList: + """ + A handler that forwards the list roots request from the remote server to the proxy's connected clients and relays the response back to the remote server. + """ + ctx = get_context() + return await ctx.list_roots() + + +class ProxyClient(Client[ClientTransportT]): + """ + A proxy client that forwards advanced interactions between a remote MCP server and the proxy's connected clients. + Supports forwarding roots, sampling, elicitation, logging, and progress. + """ + + def __init__( + self, + transport: ( + ClientTransportT + | FastMCP + | AnyUrl + | Path + | MCPConfig + | dict[str, Any] + | str + ), + **kwargs, + ): + if "roots" not in kwargs: + kwargs["roots"] = default_proxy_roots_handler + if "sampling_handler" not in kwargs: + kwargs["sampling_handler"] = ProxyClient.default_sampling_handler + if "elicitation_handler" not in kwargs: + kwargs["elicitation_handler"] = ProxyClient.default_elicitation_handler + if "log_handler" not in kwargs: + kwargs["log_handler"] = ProxyClient.default_log_handler + if "progress_handler" not in kwargs: + kwargs["progress_handler"] = ProxyClient.default_progress_handler + super().__init__(transport, **kwargs) + + @classmethod + async def default_sampling_handler( + cls, + messages: list[mcp.types.SamplingMessage], + params: mcp.types.CreateMessageRequestParams, + context: RequestContext[ClientSession, LifespanContextT], + ) -> mcp.types.CreateMessageResult: + """ + A handler that forwards the sampling request from the remote server to the proxy's connected clients and relays the response back to the remote server. + """ + ctx = get_context() + content = await ctx.sample( + [msg for msg in messages], + system_prompt=params.systemPrompt, + temperature=params.temperature, + max_tokens=params.maxTokens, + model_preferences=params.modelPreferences, + ) + if isinstance(content, mcp.types.ResourceLink | mcp.types.EmbeddedResource): + raise RuntimeError("Content is not supported") + return mcp.types.CreateMessageResult( + role="assistant", + model="fastmcp-client", + content=content, + ) + + @classmethod + async def default_elicitation_handler( + cls, + message: str, + response_type: type, + params: mcp.types.ElicitRequestParams, + context: RequestContext[ClientSession, LifespanContextT], + ) -> ElicitResult: + """ + A handler that forwards the elicitation request from the remote server to the proxy's connected clients and relays the response back to the remote server. + """ + ctx = get_context() + result = await ctx.elicit(message, response_type) + if result.action == "accept": + return result.data + else: + return ElicitResult(action=result.action) + + @classmethod + async def default_log_handler(cls, message: LogMessage) -> None: + """ + A handler that forwards the log notification from the remote server to the proxy's connected clients. + """ + ctx = get_context() + await ctx.log(message.data, level=message.level, logger_name=message.logger) + + @classmethod + async def default_progress_handler( + cls, + progress: float, + total: float | None, + message: str | None, + ) -> None: + """ + A handler that forwards the progress notification from the remote server to the proxy's connected clients. + """ + ctx = get_context() + await ctx.report_progress(progress, total, message) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 541533d6e..6d92e412f 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -1643,9 +1643,8 @@ class FastMCP(Generic[LifespanResultT]): resource_separator: Deprecated. Separator character for resource URIs. prompt_separator: Deprecated. Separator character for prompt names. """ - from fastmcp import Client from fastmcp.client.transports import FastMCPTransport - from fastmcp.server.proxy import FastMCPProxy + from fastmcp.server.proxy import FastMCPProxy, ProxyClient # Deprecated since 2.9.0 # Prior to 2.9.0, the first positional argument was the prefix and the @@ -1697,7 +1696,7 @@ class FastMCP(Generic[LifespanResultT]): as_proxy = server._has_lifespan if as_proxy and not isinstance(server, FastMCPProxy): - server = FastMCPProxy(Client(transport=FastMCPTransport(server))) + server = FastMCPProxy(ProxyClient(transport=FastMCPTransport(server))) # Delegate mounting to all three managers mounted_server = MountedServer( @@ -1908,14 +1907,16 @@ class FastMCP(Generic[LifespanResultT]): @classmethod def as_proxy( cls, - backend: Client[ClientTransportT] - | ClientTransport - | FastMCP[Any] - | AnyUrl - | Path - | MCPConfig - | dict[str, Any] - | str, + backend: ( + Client[ClientTransportT] + | ClientTransport + | FastMCP[Any] + | AnyUrl + | Path + | MCPConfig + | dict[str, Any] + | str + ), **settings: Any, ) -> FastMCPProxy: """Create a FastMCP proxy server for the given backend. @@ -1926,12 +1927,12 @@ class FastMCP(Generic[LifespanResultT]): `fastmcp.client.Client` constructor. """ from fastmcp.client.client import Client - from fastmcp.server.proxy import FastMCPProxy + from fastmcp.server.proxy import FastMCPProxy, ProxyClient if isinstance(backend, Client): client = backend else: - client = Client(backend) + client = ProxyClient(backend) return FastMCPProxy(client=client, **settings) diff --git a/tests/server/proxy/test_proxy_client.py b/tests/server/proxy/test_proxy_client.py new file mode 100644 index 000000000..cf524f9b4 --- /dev/null +++ b/tests/server/proxy/test_proxy_client.py @@ -0,0 +1,257 @@ +from dataclasses import dataclass +from typing import cast + +import pytest +from mcp.types import LoggingLevel, ModelHint, ModelPreferences, TextContent + +from fastmcp import Client, Context, FastMCP +from fastmcp.client.elicitation import ElicitRequestParams, ElicitResult +from fastmcp.client.logging import LogMessage +from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams +from fastmcp.exceptions import ToolError +from fastmcp.server.proxy import ProxyClient + + +@pytest.fixture +def fastmcp_server(): + mcp = FastMCP("TestServer") + + @mcp.tool + async def list_roots(context: Context) -> list[str]: + roots = await context.list_roots() + return [str(r.uri) for r in roots] + + @mcp.tool + async def sampling( + context: Context, + ) -> str: + result = await context.sample( + "Hello, world!", + system_prompt="You love FastMCP", + temperature=0.5, + max_tokens=100, + model_preferences="gpt-4o", + ) + return cast(TextContent, result).text + + @dataclass + class Person: + name: str + + @mcp.tool + async def elicit(context: Context) -> str: + result = await context.elicit( + message="What is your name?", + response_type=Person, + ) + if result.action == "accept": + return f"Hello, {result.data.name}!" + else: + return "No name provided." + + @mcp.tool + async def log( + message: str, level: LoggingLevel, logger: str, context: Context + ) -> None: + await context.log(message=message, level=level, logger_name=logger) + + @mcp.tool + async def report_progress(context: Context) -> int: + for i in range(3): + await context.report_progress( + progress=i + 1, + total=3, + message=f"{(i + 1) / 3 * 100:.2f}% complete", + ) + return 100 + + return mcp + + +@pytest.fixture +async def proxy_server(fastmcp_server: FastMCP): + """ + A proxy server that forwards interactions with the proxy client to the given fastmcp server. + """ + return FastMCP.as_proxy(ProxyClient(fastmcp_server)) + + +class TestProxyClient: + async def test_forward_error_response(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards an error response. + """ + async with Client(proxy_server) as client: + with pytest.raises(ToolError, match="Elicitation not supported"): + await client.call_tool("elicit", {}) + + async def test_forward_list_roots_request(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `list_roots` request. + """ + roots_handler_called = False + + async def roots_handler(ctx: RequestContext): + nonlocal roots_handler_called + roots_handler_called = True + return [] + + async with Client(proxy_server, roots=roots_handler) as client: + await client.call_tool("list_roots", {}) + + assert roots_handler_called + + async def test_forward_list_roots_response(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `list_roots` response. + """ + async with Client(proxy_server, roots=["file://x/y/z"]) as client: + result = await client.call_tool("list_roots", {}) + assert result.data == ["file://x/y/z"] + + async def test_forward_sampling_request(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `sampling` request. + """ + sampling_handler_called = False + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> str: + nonlocal sampling_handler_called + sampling_handler_called = True + assert messages == [ + SamplingMessage( + role="user", + content=TextContent(type="text", text="Hello, world!"), + ) + ] + assert params.systemPrompt == "You love FastMCP" + assert params.temperature == 0.5 + assert params.maxTokens == 100 + assert params.modelPreferences == ModelPreferences( + hints=[ModelHint(name="gpt-4o")] + ) + return "" + + async with Client(proxy_server, sampling_handler=sampling_handler) as client: + await client.call_tool("sampling", {}) + + assert sampling_handler_called + + async def test_forward_sampling_response(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `sampling` response. + """ + async with Client( + proxy_server, sampling_handler=lambda *args: "I love FastMCP" + ) as client: + result = await client.call_tool("sampling", {}) + assert result.data == "I love FastMCP" + + async def test_elicit_request(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `elicit` request. + """ + elicitation_handler_called = False + + async def elicitation_handler( + message, response_type, params: ElicitRequestParams, ctx + ): + nonlocal elicitation_handler_called + elicitation_handler_called = True + assert message == "What is your name?" + assert "Person" in str(response_type) + assert params.requestedSchema == { + "title": "Person", + "type": "object", + "properties": {"name": {"title": "Name", "type": "string"}}, + "required": ["name"], + } + return ElicitResult(action="accept", content=response_type(name="Alice")) + + async with Client( + proxy_server, elicitation_handler=elicitation_handler + ) as client: + await client.call_tool("elicit", {}) + + assert elicitation_handler_called + + async def test_elicit_accept_response(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `elicit` accept response. + """ + + async def elicitation_handler( + message, response_type, params: ElicitRequestParams, ctx + ): + return ElicitResult(action="accept", content=response_type(name="Alice")) + + async with Client( + proxy_server, + elicitation_handler=elicitation_handler, + ) as client: + result = await client.call_tool("elicit", {}) + assert result.data == "Hello, Alice!" + + async def test_elicit_decline_response(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `elicit` decline response. + """ + + async def elicitation_handler( + message, response_type, params: ElicitRequestParams, ctx + ): + return ElicitResult(action="decline") + + async with Client( + proxy_server, elicitation_handler=elicitation_handler + ) as client: + result = await client.call_tool("elicit", {}) + assert result.data == "No name provided." + + async def test_log_request(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `log` request. + """ + log_handler_called = False + + async def log_handler(message: LogMessage) -> None: + nonlocal log_handler_called + log_handler_called = True + assert message.data == "Hello, world!" + assert message.level == "info" + assert message.logger == "test" + + async with Client(proxy_server, log_handler=log_handler) as client: + await client.call_tool( + "log", {"message": "Hello, world!", "level": "info", "logger": "test"} + ) + + assert log_handler_called + + async def test_report_progress_request(self, proxy_server: FastMCP): + """ + Test that the proxy client correctly forwards the `report_progress` request. + """ + + EXPECTED_PROGRESS_MESSAGES = [ + dict(progress=1, total=3, message="33.33% complete"), + dict(progress=2, total=3, message="66.67% complete"), + dict(progress=3, total=3, message="100.00% complete"), + ] + PROGRESS_MESSAGES = [] + + async def progress_handler( + progress: float, total: float | None, message: str | None + ) -> None: + PROGRESS_MESSAGES.append( + dict(progress=progress, total=total, message=message) + ) + + async with Client(proxy_server, progress_handler=progress_handler) as client: + await client.call_tool("report_progress", {}) + + assert PROGRESS_MESSAGES == EXPECTED_PROGRESS_MESSAGES diff --git a/tests/server/test_proxy.py b/tests/server/test_proxy.py index e63c18551..5f2f20bb2 100644 --- a/tests/server/test_proxy.py +++ b/tests/server/test_proxy.py @@ -11,7 +11,7 @@ from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.client.transports import FastMCPTransport, StreamableHttpTransport from fastmcp.exceptions import ToolError -from fastmcp.server.proxy import FastMCPProxy +from fastmcp.server.proxy import FastMCPProxy, ProxyClient USERS = [ {"id": "1", "name": "Alice", "active": True}, @@ -71,13 +71,13 @@ def fastmcp_server(): @pytest.fixture async def proxy_server(fastmcp_server): """Fixture that creates a FastMCP proxy server.""" - return FastMCP.as_proxy(Client(transport=FastMCPTransport(fastmcp_server))) + return FastMCP.as_proxy(ProxyClient(transport=FastMCPTransport(fastmcp_server))) async def test_create_proxy(fastmcp_server): """Test that the proxy server properly forwards requests to the original server.""" # Create a client - client = Client(transport=FastMCPTransport(fastmcp_server)) + client = ProxyClient(transport=FastMCPTransport(fastmcp_server)) server = FastMCPProxy.as_proxy(client)