Fix proxy servers

This commit is contained in:
Jeremiah Lowin 2025-06-19 14:50:58 -04:00
commit 2ae5a800e7
2 changed files with 268 additions and 157 deletions

View file

@ -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:
"""

View file

@ -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.
#