From 8f1e06f48dd50a670a4dc2e0c757fb2a0a9f8965 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 8 Apr 2025 13:28:48 -0400 Subject: [PATCH] Add proxy server --- pyproject.toml | 4 + src/fastmcp/clients/base.py | 12 +- src/fastmcp/server/proxy.py | 212 +++++++++++++++++++++++++++ src/fastmcp/server/server.py | 24 ++- tests/clients/test_fastmcp_client.py | 18 +-- tests/server/test_proxy.py | 179 ++++++++++++++++++++++ uv.lock | 13 +- 7 files changed, 444 insertions(+), 18 deletions(-) create mode 100644 src/fastmcp/server/proxy.py create mode 100644 tests/server/test_proxy.py diff --git a/pyproject.toml b/pyproject.toml index 2b13c2d2c..20e56e036 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,11 +32,15 @@ dev = [ "copychat>=0.5.2", "ipython>=8.12.3", "pdbpp>=0.10.3", + "dirty-equals>=0.9.0", ] [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "session" +filterwarnings = [ + "ignore:Accessing the 'model_fields' attribute on the instance is deprecated:DeprecationWarning", +] [tool.hatch.version] source = "vcs" diff --git a/src/fastmcp/clients/base.py b/src/fastmcp/clients/base.py index ba95b6b1d..b2a2a454d 100644 --- a/src/fastmcp/clients/base.py +++ b/src/fastmcp/clients/base.py @@ -165,16 +165,22 @@ class BaseClient(abc.ABC): """Send a resources/listResourceTemplates request.""" return await self.session.list_resource_templates() - async def read_resource(self, uri: AnyUrl) -> mcp.types.ReadResourceResult: + async def read_resource(self, uri: AnyUrl | str) -> mcp.types.ReadResourceResult: """Send a resources/read request.""" + if isinstance(uri, str): + uri = AnyUrl(uri) return await self.session.read_resource(uri) - async def subscribe_resource(self, uri: AnyUrl) -> None: + async def subscribe_resource(self, uri: AnyUrl | str) -> None: """Send a resources/subscribe request.""" + if isinstance(uri, str): + uri = AnyUrl(uri) await self.session.subscribe_resource(uri) - async def unsubscribe_resource(self, uri: AnyUrl) -> None: + async def unsubscribe_resource(self, uri: AnyUrl | str) -> None: """Send a resources/unsubscribe request.""" + if isinstance(uri, str): + uri = AnyUrl(uri) await self.session.unsubscribe_resource(uri) async def list_prompts(self) -> mcp.types.ListPromptsResult: diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py new file mode 100644 index 000000000..4dda5e1d7 --- /dev/null +++ b/src/fastmcp/server/proxy.py @@ -0,0 +1,212 @@ +from typing import Any, cast + +import mcp.types +from mcp.server.fastmcp.prompts import Prompt +from mcp.server.fastmcp.resources import Resource, ResourceTemplate +from mcp.server.fastmcp.tools.base import Tool +from mcp.server.fastmcp.utilities.func_metadata import func_metadata +from mcp.types import BlobResourceContents, PromptMessage, TextResourceContents + +from fastmcp.clients.base import BaseClient +from fastmcp.server.context import Context +from fastmcp.server.server import FastMCP +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) + + +def _proxy_passthrough(): + pass + + +class ProxyTool(Tool): + def __init__(self, client: "BaseClient", **kwargs): + super().__init__(**kwargs) + self._client = client + + @classmethod + async def from_client( + cls, client: "BaseClient", tool: mcp.types.Tool + ) -> "ProxyTool": + return cls( + client=client, + name=tool.name, + description=tool.description, + parameters=tool.inputSchema, + fn=_proxy_passthrough, + fn_metadata=func_metadata(_proxy_passthrough), + is_async=True, + ) + + async def run( + self, arguments: dict[str, Any], context: Context | None = None + ) -> Any: + async with self._client: + result = await self._client.call_tool(self.name, arguments) + if result.isError: + raise ValueError(cast(mcp.types.TextContent, result.content[0]).text) + return result.content[0] + + +class ProxyResource(Resource): + def __init__( + self, client: "BaseClient", *, _value: str | bytes | None = None, **kwargs + ): + super().__init__(**kwargs) + self._client = client + self._value = _value + + @classmethod + async def from_client( + cls, client: "BaseClient", resource: mcp.types.Resource + ) -> "ProxyResource": + return cls( + client=client, + uri=resource.uri, + name=resource.name, + description=resource.description, + mime_type=resource.mimeType, + ) + + async def read(self) -> str | bytes: + if self._value is not None: + return self._value + + async with self._client: + result = await self._client.read_resource(self.uri) + if isinstance(result.contents[0], TextResourceContents): + return result.contents[0].text + elif isinstance(result.contents[0], BlobResourceContents): + return result.contents[0].blob + else: + raise ValueError(f"Unsupported content type: {type(result.contents[0])}") + + +class ProxyTemplate(ResourceTemplate): + def __init__(self, client: "BaseClient", **kwargs): + super().__init__(**kwargs) + self._client = client + + @classmethod + async def from_client( + cls, client: "BaseClient", template: mcp.types.ResourceTemplate + ) -> "ProxyTemplate": + return cls( + client=client, + uri_template=template.uriTemplate, + name=template.name, + description=template.description, + fn=_proxy_passthrough, + parameters={}, + ) + + async def create_resource(self, uri: str, params: dict[str, Any]) -> ProxyResource: + async with self._client: + result = await self._client.read_resource(uri) + + if isinstance(result.contents[0], TextResourceContents): + value = result.contents[0].text + elif isinstance(result.contents[0], BlobResourceContents): + value = result.contents[0].blob + else: + raise ValueError(f"Unsupported content type: {type(result.contents[0])}") + + return ProxyResource( + client=self._client, + uri=uri, + name=self.name, + description=self.description, + mime_type=result.contents[0].mimeType, + contents=result.contents, + _value=value, + ) + + +class ProxyPrompt(Prompt): + def __init__(self, client: "BaseClient", **kwargs): + super().__init__(**kwargs) + self._client = client + + @classmethod + async def from_client( + cls, client: "BaseClient", prompt: mcp.types.Prompt + ) -> "ProxyPrompt": + return cls( + client=client, + name=prompt.name, + description=prompt.description, + arguments=[a.model_dump() for a in prompt.arguments or []], + fn=_proxy_passthrough, + ) + + async def render(self, arguments: dict[str, Any]) -> list[PromptMessage]: + async with self._client: + result = await self._client.get_prompt(self.name, arguments) + return result.messages + + +class FastMCPProxy(FastMCP): + def __init__(self, _async_constructor: bool, **kwargs): + if not _async_constructor: + raise ValueError( + "FastMCPProxy() was initialied unexpectedly. Please use a constructor like `FastMCPProxy.from_client()` instead." + ) + super().__init__(**kwargs) + + @classmethod + async def from_client( + cls, client: "BaseClient", name: str | None = None, **settings: Any + ) -> "FastMCPProxy": + """Create a FastMCP proxy server from a client. + + This method creates a new FastMCP server instance that proxies requests to the provided client. + It discovers the client's tools, resources, prompts, and templates, and creates corresponding + components in the server that forward requests to the client. + + Args: + client: The client to proxy requests to + name: Optional name for the new FastMCP server (defaults to client name if available) + **settings: Additional settings for the FastMCP server + + Returns: + A FastMCP server that proxies requests to the client + """ + server = cls(name=name, **settings, _async_constructor=True) + + async with client: + # Register proxies for client tools + tools_result = await client.list_tools() + for tool in tools_result.tools: + tool_proxy = await ProxyTool.from_client(client, tool) + server._tool_manager._tools[tool_proxy.name] = tool_proxy + logger.debug(f"Created proxy for tool: {tool_proxy.name}") + + # Register proxies for client resources + resources_result = await client.list_resources() + for resource in resources_result.resources: + resource_proxy = await ProxyResource.from_client(client, resource) + server._resource_manager._resources[str(resource_proxy.uri)] = ( + resource_proxy + ) + logger.debug(f"Created proxy for resource: {resource_proxy.uri}") + + # Register proxies for client resource templates + templates_result = await client.list_resource_templates() + for template in templates_result.resourceTemplates: + template_proxy = await ProxyTemplate.from_client(client, template) + server._resource_manager._templates[template_proxy.uri_template] = ( + template_proxy + ) + logger.debug( + f"Created proxy for template: {template_proxy.uri_template}" + ) + + # Register proxies for client prompts + prompts_result = await client.list_prompts() + for prompt in prompts_result.prompts: + prompt_proxy = await ProxyPrompt.from_client(client, prompt) + server._prompt_manager._prompts[prompt_proxy.name] = prompt_proxy + logger.debug(f"Created proxy for prompt: {prompt_proxy.name}") + + logger.info(f"Created server '{server.name}' proxying to client: {client}") + return server diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 61289f5fc..7d63f6d61 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -10,7 +10,9 @@ from fastmcp.tools.tool_manager import ToolManager from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: - pass + from fastmcp.clients.base import BaseClient + + from .proxy import FastMCPProxy logger = get_logger(__name__) @@ -83,3 +85,23 @@ class FastMCP(mcp.server.fastmcp.FastMCP): logger.debug(f"Imported resources with prefix '{resource_prefix}'") logger.debug(f"Imported templates with prefix '{resource_prefix}'") logger.debug(f"Imported prompts with prefix '{prompt_prefix}'") + + @classmethod + async def as_proxy(cls, client: "BaseClient", **settings: Any) -> "FastMCPProxy": + """ + Create a FastMCP proxy server from a client. + + This method creates a new FastMCP server instance that proxies requests to the provided client. + It discovers the client's tools, resources, prompts, and templates, and creates corresponding + components in the server that forward requests to the client. + + Args: + client: The client to proxy requests to + **settings: Additional settings for the FastMCP server + + Returns: + A FastMCP server that proxies requests to the client + """ + from .proxy import FastMCPProxy + + return await FastMCPProxy.from_client(client=client, **settings) diff --git a/tests/clients/test_fastmcp_client.py b/tests/clients/test_fastmcp_client.py index 0d4b88167..813f0df2e 100644 --- a/tests/clients/test_fastmcp_client.py +++ b/tests/clients/test_fastmcp_client.py @@ -7,12 +7,6 @@ from fastmcp.clients import FastMCPClient from fastmcp.server.server import FastMCP -class _TestException(Exception): - """Test exception for testing raise_exceptions behavior.""" - - pass - - @pytest.fixture def fastmcp_server(): """Fixture that creates a FastMCP server with tools, resources, and prompts.""" @@ -24,11 +18,11 @@ def fastmcp_server(): """Greet someone by name.""" return f"Hello, {name}!" - # Add a tool that raises an exception + # Add a second tool @server.tool() - def error_tool() -> str: - """A tool that always raises an exception.""" - raise _TestException("Deliberate test exception") + def add(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b # Add a resource @server.resource(uri="data://users") @@ -57,9 +51,7 @@ async def test_list_tools(fastmcp_server): # Check that our tools are available assert len(result.tools) == 2 - tool_names = [tool.name for tool in result.tools] - assert "greet" in tool_names - assert "error_tool" in tool_names + assert set(tool.name for tool in result.tools) == {"greet", "add"} async def test_call_tool(fastmcp_server): diff --git a/tests/server/test_proxy.py b/tests/server/test_proxy.py new file mode 100644 index 000000000..62d8b92b6 --- /dev/null +++ b/tests/server/test_proxy.py @@ -0,0 +1,179 @@ +import json +from typing import Any + +import pytest +from dirty_equals import Contains + +from fastmcp import FastMCP +from fastmcp.clients.fastmcp_client import FastMCPClient +from fastmcp.server.proxy import FastMCPProxy + +USERS = [ + {"id": "1", "name": "Alice", "active": True}, + {"id": "2", "name": "Bob", "active": True}, + {"id": "3", "name": "Charlie", "active": False}, +] + + +@pytest.fixture +def fastmcp_server(): + server = FastMCP("TestServer") + + # --- Tools --- + + @server.tool() + def greet(name: str) -> str: + """Greet someone by name.""" + return f"Hello, {name}!" + + @server.tool() + def add(a: int, b: int) -> int: + """Add two numbers together.""" + return a + b + + @server.tool() + def error_tool(): + """This tool always raises an error.""" + raise ValueError("This is a test error") + + # --- Resources --- + + @server.resource(uri="resource://wave") + def wave() -> str: + return "👋" + + @server.resource(uri="data://users") + async def get_users() -> list[dict[str, Any]]: + return USERS + + @server.resource(uri="data://user/{user_id}") + async def get_user(user_id: str) -> dict[str, Any] | None: + return next((user for user in USERS if user["id"] == user_id), None) + + # --- Prompts --- + + @server.prompt() + def welcome(name: str) -> str: + return f"Welcome to FastMCP, {name}!" + + return server + + +@pytest.fixture +async def proxy_server(fastmcp_server): + """Fixture that creates a FastMCP proxy server.""" + return await FastMCP.as_proxy(FastMCPClient(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 = FastMCPClient(fastmcp_server) + + server = await FastMCPProxy.from_client(client) + + assert isinstance(server, FastMCP) + assert server.name == "FastMCP" + + +class TestTools: + async def test_list_tools(self, proxy_server): + tools = await proxy_server.list_tools() + assert [t.name for t in tools] == Contains("greet", "add", "error_tool") + + async def test_list_tools_same_as_original(self, fastmcp_server, proxy_server): + assert await proxy_server.list_tools() == await fastmcp_server.list_tools() + + async def test_call_tool_result_same_as_original( + self, fastmcp_server: FastMCP, proxy_server: FastMCPProxy + ): + result = await fastmcp_server.call_tool("greet", {"name": "Alice"}) + proxy_result = await proxy_server.call_tool("greet", {"name": "Alice"}) + + assert result == proxy_result + + async def test_call_tool_calls_tool(self, proxy_server): + proxy_result = await proxy_server.call_tool("add", {"a": 1, "b": 2}) + + assert proxy_result[0].text == "3" + + async def test_error_tool_raises_error(self, proxy_server): + with pytest.raises(ValueError, match="This is a test error"): + await proxy_server.call_tool("error_tool", {}) + + +class TestResources: + async def test_list_resources(self, proxy_server): + resources = await proxy_server.list_resources() + assert [r.name for r in resources] == Contains( + "data://users", "resource://wave" + ) + + async def test_list_resources_same_as_original(self, fastmcp_server, proxy_server): + assert ( + await proxy_server.list_resources() == await fastmcp_server.list_resources() + ) + + async def test_read_resource(self, proxy_server: FastMCPProxy): + result = await proxy_server.read_resource("resource://wave") + assert result[0].content == "👋" # type: ignore + + async def test_read_resource_same_as_original(self, fastmcp_server, proxy_server): + result = await fastmcp_server.read_resource("resource://wave") + proxy_result = await proxy_server.read_resource("resource://wave") + assert proxy_result == result + + async def test_read_json_resource(self, proxy_server: FastMCPProxy): + result = await proxy_server.read_resource("data://users") + assert json.loads(result[0].content) == USERS # type: ignore + + async def test_read_resource_returns_none_if_not_found(self, proxy_server): + with pytest.raises( + ValueError, match="Unknown resource: resource://nonexistent" + ): + await proxy_server.read_resource("resource://nonexistent") + + +class TestResourceTemplates: + async def test_list_resource_templates(self, proxy_server): + templates = await proxy_server.list_resource_templates() + assert [t.name for t in templates] == Contains("get_user") + + async def test_list_resource_templates_same_as_original( + self, fastmcp_server, proxy_server + ): + result = await fastmcp_server.list_resource_templates() + proxy_result = await proxy_server.list_resource_templates() + assert proxy_result == result + + @pytest.mark.parametrize("id", [1, 2, 3]) + async def test_read_resource_template(self, proxy_server: FastMCPProxy, id: int): + result = await proxy_server.read_resource(f"data://user/{id}") + assert json.loads(result[0].content) == USERS[id - 1] # type: ignore + + async def test_read_resource_template_same_as_original( + self, fastmcp_server, proxy_server + ): + result = await fastmcp_server.read_resource("data://user/1") + proxy_result = await proxy_server.read_resource("data://user/1") + assert proxy_result == result + + +class TestPrompts: + async def test_list_prompts(self, proxy_server): + prompts = await proxy_server.list_prompts() + assert [p.name for p in prompts] == Contains("welcome") + + async def test_list_prompts_same_as_original(self, fastmcp_server, proxy_server): + assert await proxy_server.list_prompts() == await fastmcp_server.list_prompts() + + async def test_render_prompt_same_as_original( + self, fastmcp_server: FastMCP, proxy_server + ): + result = await fastmcp_server.get_prompt("welcome", {"name": "Alice"}) + proxy_result = await proxy_server.get_prompt("welcome", {"name": "Alice"}) + assert proxy_result == result + + async def test_render_prompt_calls_prompt(self, proxy_server): + result = await proxy_server.get_prompt("welcome", {"name": "Alice"}) + assert result.messages[0].content.text == "Welcome to FastMCP, Alice!" diff --git a/uv.lock b/uv.lock index 3ac545e5c..9444459d0 100644 --- a/uv.lock +++ b/uv.lock @@ -169,6 +169,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190 }, ] +[[package]] +name = "dirty-equals" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/99/133892f401ced5a27e641a473c547d5fbdb39af8f85dac8a9d633ea3e7a7/dirty_equals-0.9.0.tar.gz", hash = "sha256:17f515970b04ed7900b733c95fd8091f4f85e52f1fb5f268757f25c858eb1f7b", size = 50412 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/0c/03cc99bf3b6328604b10829de3460f2b2ad3373200c45665c38508e550c6/dirty_equals-0.9.0-py3-none-any.whl", hash = "sha256:ff4d027f5cfa1b69573af00f7ba9043ea652dbdce3fe5cbe828e478c7346db9c", size = 28226 }, +] + [[package]] name = "distlib" version = "0.3.9" @@ -220,7 +229,7 @@ wheels = [ [[package]] name = "fastmcp" -version = "0.4.2.dev21+g0b6e56c.d20250407" +version = "0.4.2.dev28+g728aeec.d20250408" source = { editable = "." } dependencies = [ { name = "mcp" }, @@ -232,6 +241,7 @@ dependencies = [ [package.dev-dependencies] dev = [ { name = "copychat" }, + { name = "dirty-equals" }, { name = "ipython" }, { name = "pdbpp" }, { name = "pre-commit" }, @@ -254,6 +264,7 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ { name = "copychat", specifier = ">=0.5.2" }, + { name = "dirty-equals", specifier = ">=0.9.0" }, { name = "ipython", specifier = ">=8.12.3" }, { name = "pdbpp", specifier = ">=0.10.3" }, { name = "pre-commit" },