Fix import prefix behavior

This commit is contained in:
Jeremiah Lowin 2025-06-19 15:19:31 -04:00
commit 725c061d66
11 changed files with 104 additions and 135 deletions

View file

@ -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(

View file

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

View file

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

View file

@ -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()

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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