From 129f94bc68977dc686e892a7a7d2301449de5fc1 Mon Sep 17 00:00:00 2001 From: zzstoatzz Date: Tue, 7 Oct 2025 15:04:25 -0500 Subject: [PATCH] fix: prevent race condition in proxy component fetching When mounting remote MCP servers via ProxyClient, concurrent list_tools/ list_resources/list_prompts calls would race to fetch from the backend, causing duplicate connections and request cancellations. This led to timeouts when calling tools on servers with remote proxies mounted. This change adds caching with async locks to ensure only one fetch occurs per component type, with subsequent concurrent calls waiting for and reusing the cached results. The first fetch still takes time (HTTP connection + MCP initialization), but subsequent operations are fast. Fixes #2019 --- src/fastmcp/server/proxy.py | 129 ++++++++++++++++++++++++++---------- 1 file changed, 95 insertions(+), 34 deletions(-) diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 6847befcd..eef047b22 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import inspect import warnings from collections.abc import Awaitable, Callable @@ -72,25 +73,24 @@ class ProxyToolManager(ToolManager, ProxyManagerMixin): def __init__(self, client_factory: ClientFactoryT, **kwargs: Any): super().__init__(**kwargs) self.client_factory = client_factory + self._proxy_tools_cache: dict[str, Tool] | None = None + self._fetch_lock = asyncio.Lock() 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: - client = await self._get_client() - async with client: - client_tools = await client.list_tools() - for tool in client_tools: - if tool.name not in all_tools: - all_tools[tool.name] = ProxyTool.from_mcp_tool(client, tool) - except McpError as e: - if e.error.code == METHOD_NOT_FOUND: - pass # No tools available from proxy - else: - raise e + # Then add proxy tools from cache or fetch them (with lock to prevent concurrent fetches) + if self._proxy_tools_cache is None: + async with self._fetch_lock: + # Double-check after acquiring lock in case another task already fetched + if self._proxy_tools_cache is None: + self._proxy_tools_cache = await self._fetch_proxy_tools() + + for tool_name, tool in self._proxy_tools_cache.items(): + if tool_name not in all_tools: + all_tools[tool_name] = tool transformed_tools = apply_transformations_to_tools( tools=all_tools, @@ -99,6 +99,22 @@ class ProxyToolManager(ToolManager, ProxyManagerMixin): return transformed_tools + async def _fetch_proxy_tools(self) -> dict[str, Tool]: + """Fetch tools from the remote proxy client.""" + proxy_tools: dict[str, Tool] = {} + try: + client = await self._get_client() + async with client: + client_tools = await client.list_tools() + for tool in client_tools: + proxy_tools[tool.name] = ProxyTool.from_mcp_tool(client, tool) + except McpError as e: + if e.error.code == METHOD_NOT_FOUND: + pass # No tools available from proxy + else: + raise e + return proxy_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() @@ -126,52 +142,82 @@ class ProxyResourceManager(ResourceManager, ProxyManagerMixin): def __init__(self, client_factory: ClientFactoryT, **kwargs: Any): super().__init__(**kwargs) self.client_factory = client_factory + self._proxy_resources_cache: dict[str, Resource] | None = None + self._proxy_templates_cache: dict[str, ResourceTemplate] | None = None + self._fetch_resources_lock = asyncio.Lock() + self._fetch_templates_lock = asyncio.Lock() 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 + # Then add proxy resources from cache or fetch them (with lock to prevent concurrent fetches) + if self._proxy_resources_cache is None: + async with self._fetch_resources_lock: + # Double-check after acquiring lock + if self._proxy_resources_cache is None: + self._proxy_resources_cache = await self._fetch_proxy_resources() + + for uri, resource in self._proxy_resources_cache.items(): + if uri not in all_resources: + all_resources[uri] = resource + + return all_resources + + async def _fetch_proxy_resources(self) -> dict[str, Resource]: + """Fetch resources from the remote proxy client.""" + proxy_resources: dict[str, Resource] = {} try: client = await self._get_client() async with client: client_resources = await 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(client, resource) - ) + proxy_resources[str(resource.uri)] = ( + ProxyResource.from_mcp_resource(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 + return proxy_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 + # Then add proxy templates from cache or fetch them (with lock to prevent concurrent fetches) + if self._proxy_templates_cache is None: + async with self._fetch_templates_lock: + # Double-check after acquiring lock + if self._proxy_templates_cache is None: + self._proxy_templates_cache = await self._fetch_proxy_templates() + + for uri_template, template in self._proxy_templates_cache.items(): + if uri_template not in all_templates: + all_templates[uri_template] = template + + return all_templates + + async def _fetch_proxy_templates(self) -> dict[str, ResourceTemplate]: + """Fetch resource templates from the remote proxy client.""" + proxy_templates: dict[str, ResourceTemplate] = {} try: client = await self._get_client() async with client: client_templates = await client.list_resource_templates() for template in client_templates: - if template.uriTemplate not in all_templates: - all_templates[template.uriTemplate] = ( - ProxyTemplate.from_mcp_template(client, template) - ) + proxy_templates[template.uriTemplate] = ( + ProxyTemplate.from_mcp_template(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 + return proxy_templates async def list_resources(self) -> list[Resource]: """Gets the filtered list of resources including local, mounted, and proxy resources.""" @@ -207,29 +253,44 @@ class ProxyPromptManager(PromptManager, ProxyManagerMixin): def __init__(self, client_factory: ClientFactoryT, **kwargs: Any): super().__init__(**kwargs) self.client_factory = client_factory + self._proxy_prompts_cache: dict[str, Prompt] | None = None + self._fetch_lock = asyncio.Lock() 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 + # Then add proxy prompts from cache or fetch them (with lock to prevent concurrent fetches) + if self._proxy_prompts_cache is None: + async with self._fetch_lock: + # Double-check after acquiring lock + if self._proxy_prompts_cache is None: + self._proxy_prompts_cache = await self._fetch_proxy_prompts() + + for prompt_name, prompt in self._proxy_prompts_cache.items(): + if prompt_name not in all_prompts: + all_prompts[prompt_name] = prompt + + return all_prompts + + async def _fetch_proxy_prompts(self) -> dict[str, Prompt]: + """Fetch prompts from the remote proxy client.""" + proxy_prompts: dict[str, Prompt] = {} try: client = await self._get_client() async with client: client_prompts = await client.list_prompts() for prompt in client_prompts: - if prompt.name not in all_prompts: - all_prompts[prompt.name] = ProxyPrompt.from_mcp_prompt( - client, prompt - ) + proxy_prompts[prompt.name] = ProxyPrompt.from_mcp_prompt( + 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 + return proxy_prompts async def list_prompts(self) -> list[Prompt]: """Gets the filtered list of prompts including local, mounted, and proxy prompts."""