From 5ff93f6bd6109c8391f78fd443df88558b94cd28 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 15 Apr 2025 21:03:00 -0400 Subject: [PATCH] Add mount tests --- src/fastmcp/prompts/prompt_manager.py | 4 + src/fastmcp/resources/resource_manager.py | 10 + src/fastmcp/server/server.py | 154 ++++---- src/fastmcp/settings.py | 2 +- src/fastmcp/tools/tool_manager.py | 10 +- tests/server/test_import_server.py | 11 +- tests/server/test_mount.py | 429 ++++++++++++++++++++++ tests/tools/test_tool_manager.py | 6 +- 8 files changed, 547 insertions(+), 79 deletions(-) create mode 100644 tests/server/test_mount.py diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index f223d1b32..a48ec91b9 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -77,3 +77,7 @@ class PromptManager: raise NotFoundError(f"Unknown prompt: {name}") return await prompt.render(arguments) + + def has_prompt(self, key: str) -> bool: + """Check if a prompt exists.""" + return key in self._prompts diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 4a83a2315..613cd728b 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -202,6 +202,16 @@ class ResourceManager: self._templates[storage_key] = template return template + def has_resource(self, uri: AnyUrl | str) -> bool: + """Check if a resource exists.""" + uri_str = str(uri) + if uri_str in self._resources: + return True + for template_key in self._templates.keys(): + if match_uri_template(uri_str, template_key): + return True + return False + async def get_resource(self, uri: AnyUrl | str) -> Resource: """Get resource by URI, checking concrete resources first, then templates. diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index dc52a9f3a..facd5ff18 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -67,7 +67,6 @@ class MountedServer: tool_separator: str | None = None, resource_separator: str | None = None, prompt_separator: str | None = None, - cache_expiration_seconds: int = 10, ): if tool_separator is None: tool_separator = "_" @@ -81,54 +80,53 @@ class MountedServer: self.tool_separator = tool_separator self.resource_separator = resource_separator self.prompt_separator = prompt_separator - self.cache = TimedCache( - expiration=datetime.timedelta(seconds=cache_expiration_seconds) - ) async def get_tools(self) -> dict[str, Tool]: - cached_tools = self.cache.get("tools") - if cached_tools is NOT_FOUND: - self.cache.set("tools", {}) - cached_tools = await self.server.get_tools() - self.cache.set("tools", cached_tools) + tools = await self.server.get_tools() return { f"{self.prefix}{self.tool_separator}{key}": tool - for key, tool in cached_tools.items() + for key, tool in tools.items() } async def get_resources(self) -> dict[str, Resource]: - cached_resources = self.cache.get("resources") - if cached_resources is NOT_FOUND: - self.cache.set("resources", {}) - cached_resources = await self.server.get_resources() - self.cache.set("resources", cached_resources) - return cached_resources + resources = await self.server.get_resources() + return { + f"{self.prefix}{self.resource_separator}{key}": resource + for key, resource in resources.items() + } async def get_resource_templates(self) -> dict[str, ResourceTemplate]: - cached_templates = self.cache.get("resource_templates") - if cached_templates is NOT_FOUND: - self.cache.set("resource_templates", {}) - cached_templates = await self.server.get_resource_templates() - self.cache.set("resource_templates", cached_templates) - return cached_templates + templates = await self.server.get_resource_templates() + return { + f"{self.prefix}{self.resource_separator}{key}": template + for key, template in templates.items() + } async def get_prompts(self) -> dict[str, Prompt]: - cached_prompts = self.cache.get("prompts") - if cached_prompts is NOT_FOUND: - self.cache.set("prompts", {}) - cached_prompts = await self.server.get_prompts() - self.cache.set("prompts", cached_prompts) - return cached_prompts + prompts = await self.server.get_prompts() + return { + f"{self.prefix}{self.prompt_separator}{key}": prompt + for key, prompt in prompts.items() + } - async def match_tool(self, key: str) -> bool: + def match_tool(self, key: str) -> bool: return key.startswith(f"{self.prefix}{self.tool_separator}") - async def match_resource(self, key: str) -> bool: + def strip_tool_prefix(self, key: str) -> str: + return key.removeprefix(f"{self.prefix}{self.tool_separator}") + + def match_resource(self, key: str) -> bool: return key.startswith(f"{self.prefix}{self.resource_separator}") - async def match_prompt(self, key: str) -> bool: + def strip_resource_prefix(self, key: str) -> str: + return key.removeprefix(f"{self.prefix}{self.resource_separator}") + + def match_prompt(self, key: str) -> bool: return key.startswith(f"{self.prefix}{self.prompt_separator}") + def strip_prompt_prefix(self, key: str) -> str: + return key.removeprefix(f"{self.prefix}{self.prompt_separator}") + class TimedCache: def __init__(self, expiration: datetime.timedelta): @@ -146,6 +144,9 @@ class TimedCache: else: return NOT_FOUND + def clear(self) -> None: + self.cache.clear() + @asynccontextmanager async def default_lifespan(server: "FastMCP") -> AsyncIterator[Any]: @@ -194,7 +195,7 @@ class FastMCP(Generic[LifespanResultT]): ) ) - self._server_mounts: dict[str, MountedServer] = {} + self._mounted_servers: dict[str, MountedServer] = {} if lifespan is None: lifespan = default_lifespan @@ -287,40 +288,48 @@ class FastMCP(Generic[LifespanResultT]): async def get_tools(self) -> dict[str, Tool]: """Get all registered tools, indexed by registered key.""" - tools = {} - for server in self._server_mounts.values(): - server_tools = await server.get_tools() - tools.update(server_tools) - tools.update(self._tool_manager.get_tools()) + if (tools := self._cache.get("tools")) is NOT_FOUND: + tools = {} + for server in self._mounted_servers.values(): + server_tools = await server.get_tools() + tools.update(server_tools) + tools.update(self._tool_manager.get_tools()) + self._cache.set("tools", tools) return tools async def get_resources(self) -> dict[str, Resource]: """Get all registered resources, indexed by registered key.""" - resources = {} - for server in self._server_mounts.values(): - server_resources = await server.get_resources() - resources.update(server_resources) - resources.update(self._resource_manager.get_resources()) + if (resources := self._cache.get("resources")) is NOT_FOUND: + resources = {} + for server in self._mounted_servers.values(): + server_resources = await server.get_resources() + resources.update(server_resources) + resources.update(self._resource_manager.get_resources()) + self._cache.set("resources", resources) return resources async def get_resource_templates(self) -> dict[str, ResourceTemplate]: """Get all registered resource templates, indexed by registered key.""" - templates = {} - for server in self._server_mounts.values(): - server_templates = await server.get_resource_templates() - templates.update(server_templates) - templates.update(self._resource_manager.get_templates()) + if (templates := self._cache.get("resource_templates")) is NOT_FOUND: + templates = {} + for server in self._mounted_servers.values(): + server_templates = await server.get_resource_templates() + templates.update(server_templates) + templates.update(self._resource_manager.get_templates()) + self._cache.set("resource_templates", templates) return templates async def get_prompts(self) -> dict[str, Prompt]: """ List all available prompts. """ - prompts = {} - for server in self._server_mounts.values(): - server_prompts = await server.get_prompts() - prompts.update(server_prompts) - prompts.update(self._prompt_manager.get_prompts()) + if (prompts := self._cache.get("prompts")) is NOT_FOUND: + prompts = {} + for server in self._mounted_servers.values(): + server_prompts = await server.get_prompts() + prompts.update(server_prompts) + prompts.update(self._prompt_manager.get_prompts()) + self._cache.set("prompts", prompts) return prompts async def _mcp_list_tools(self) -> list[MCPTool]: @@ -368,14 +377,15 @@ class FastMCP(Generic[LifespanResultT]): self, key: str, arguments: dict[str, Any] ) -> list[TextContent | ImageContent | EmbeddedResource]: """Call a tool by name with arguments.""" - if self._tool_manager.get_tool(key): + if self._tool_manager.has_tool(key): context = self.get_context() result = await self._tool_manager.call_tool(key, arguments, context=context) else: - for server in self._server_mounts.values(): + for server in self._mounted_servers.values(): if server.match_tool(key): - result = await server.server._mcp_call_tool(key, arguments) + new_key = server.strip_tool_prefix(key) + result = await server.server._mcp_call_tool(new_key, arguments) break else: raise NotFoundError(f"Unknown tool: {key}") @@ -385,10 +395,9 @@ class FastMCP(Generic[LifespanResultT]): """ Read a resource by URI, in the format expected by the low-level MCP server. - - See `read_resource` for a more ergonomic way to read resources. """ - if resource := await self._resource_manager.get_resource(uri): + if self._resource_manager.has_resource(uri): + resource = await self._resource_manager.get_resource(uri) try: content = await resource.read() return [ @@ -398,9 +407,10 @@ class FastMCP(Generic[LifespanResultT]): logger.error(f"Error reading resource {uri}: {e}") raise ResourceError(str(e)) else: - for server in self._server_mounts.values(): + for server in self._mounted_servers.values(): if server.match_resource(str(uri)): - return await server.server._mcp_read_resource(uri) + new_uri = server.strip_resource_prefix(str(uri)) + return await server.server._mcp_read_resource(new_uri) else: raise NotFoundError(f"Unknown resource: {uri}") @@ -411,15 +421,15 @@ class FastMCP(Generic[LifespanResultT]): Get a prompt by name with arguments, in the format expected by the low-level MCP server. - See `get_prompt` for a more ergonomic way to get prompts. """ - if prompt := self._prompt_manager.get_prompt(name): - messages = await prompt.render(arguments) + if self._prompt_manager.has_prompt(name): + messages = await self._prompt_manager.render_prompt(name, arguments) return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages)) else: - for server in self._server_mounts.values(): + for server in self._mounted_servers.values(): if server.match_prompt(name): - return await server.server._mcp_get_prompt(name, arguments) + new_key = server.strip_prompt_prefix(name) + return await server.server._mcp_get_prompt(new_key, arguments) else: raise NotFoundError(f"Unknown prompt: {name}") @@ -444,6 +454,7 @@ class FastMCP(Generic[LifespanResultT]): self._tool_manager.add_tool_from_fn( fn, name=name, description=description, tags=tags ) + self._cache.clear() def tool( self, @@ -499,6 +510,7 @@ class FastMCP(Generic[LifespanResultT]): """ self._resource_manager.add_resource(resource, key=key) + self._cache.clear() def add_resource_fn( self, @@ -530,6 +542,7 @@ class FastMCP(Generic[LifespanResultT]): mime_type=mime_type, tags=tags, ) + self._cache.clear() def resource( self, @@ -585,7 +598,7 @@ class FastMCP(Generic[LifespanResultT]): ) def decorator(fn: AnyFunction) -> AnyFunction: - self._resource_manager.add_resource_or_template_from_fn( + self.add_resource_fn( fn=fn, uri=uri, name=name, @@ -615,6 +628,7 @@ class FastMCP(Generic[LifespanResultT]): description=description, tags=tags, ) + self._cache.clear() def prompt( self, @@ -738,10 +752,12 @@ class FastMCP(Generic[LifespanResultT]): resource_separator=resource_separator, prompt_separator=prompt_separator, ) - self._server_mounts[prefix] = mounted_server + self._mounted_servers[prefix] = mounted_server + self._cache.clear() def unmount(self, prefix: str) -> None: - self._server_mounts.pop(prefix) + self._mounted_servers.pop(prefix) + self._cache.clear() async def import_server( self, @@ -815,6 +831,8 @@ class FastMCP(Generic[LifespanResultT]): logger.debug(f"Imported templates with prefix '{resource_prefix}'") logger.debug(f"Imported prompts with prefix '{prompt_prefix}'") + self._cache.clear() + @classmethod def from_openapi( cls, openapi_spec: dict[str, Any], client: httpx.AsyncClient, **settings: Any diff --git a/src/fastmcp/settings.py b/src/fastmcp/settings.py index 68401f68b..e82b75cc9 100644 --- a/src/fastmcp/settings.py +++ b/src/fastmcp/settings.py @@ -63,7 +63,7 @@ class ServerSettings(BaseSettings): ) # cache settings (for checking mounted servers) - cache_expiration_seconds: float = 5 + cache_expiration_seconds: float = 0 class ClientSettings(BaseSettings): diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index 6da7d0961..0675cfc65 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -37,9 +37,15 @@ class ToolManager: self.duplicate_behavior = duplicate_behavior - def get_tool(self, key: str) -> Tool | None: + def has_tool(self, key: str) -> bool: + """Check if a tool exists.""" + return key in self._tools + + def get_tool(self, key: str) -> Tool: """Get tool by key.""" - return self._tools.get(key) + if key in self._tools: + return self._tools[key] + raise NotFoundError(f"Unknown tool: {key}") def get_tools(self) -> dict[str, Tool]: """Get all registered tools, indexed by registered key.""" diff --git a/tests/server/test_import_server.py b/tests/server/test_import_server.py index 4583d4406..c9160e4b1 100644 --- a/tests/server/test_import_server.py +++ b/tests/server/test_import_server.py @@ -221,12 +221,10 @@ async def test_call_imported_custom_named_tool(): api_app.add_tool(fetch_data, name="get_data") await main_app.import_server("api", api_app) - context = main_app.get_context() - result = await main_app._tool_manager.call_tool( - "api_get_data", {"query": "test"}, context=context - ) - assert isinstance(result[0], TextContent) - assert result[0].text == "Data for query: test" + async with Client(main_app) as client: + result = await client.call_tool("api_get_data", {"query": "test"}) + assert isinstance(result[0], TextContent) + assert result[0].text == "Data for query: test" async def test_first_level_importing_with_custom_name(): @@ -382,6 +380,7 @@ async def test_import_with_proxy_resource_templates(): await main_app.import_server("api", proxy_app) # Instantiate the template through the main app with the prefixed key + quoted_name = quote("John Doe", safe="") quoted_email = quote("john@example.com", safe="") async with Client(main_app) as client: diff --git a/tests/server/test_mount.py b/tests/server/test_mount.py new file mode 100644 index 000000000..61d958868 --- /dev/null +++ b/tests/server/test_mount.py @@ -0,0 +1,429 @@ +import json + +import pytest +from mcp.server.lowlevel.helper_types import ReadResourceContents +from mcp.types import TextContent, TextResourceContents + +from fastmcp import FastMCP +from fastmcp.client import Client +from fastmcp.client.transports import FastMCPTransport +from fastmcp.exceptions import NotFoundError + + +class TestBasicMount: + """Test basic mounting functionality.""" + + async def test_mount_simple_server(self): + """Test mounting a simple server and accessing its tool.""" + # Create main app and sub-app + main_app = FastMCP("MainApp") + sub_app = FastMCP("SubApp") + + # Add a tool to the sub-app + @sub_app.tool() + def sub_tool() -> str: + return "This is from the sub app" + + # Mount the sub-app to the main app + main_app.mount("sub", sub_app) + + # Get tools from main app, should include sub_app's tools + tools = await main_app.get_tools() + assert "sub_sub_tool" in tools + + async with Client(main_app) as client: + result = await client.call_tool("sub_sub_tool", {}) + assert isinstance(result[0], TextContent) + assert result[0].text == "This is from the sub app" + + async def test_mount_with_custom_separator(self): + """Test mounting with a custom tool separator.""" + main_app = FastMCP("MainApp") + sub_app = FastMCP("SubApp") + + @sub_app.tool() + def greet(name: str) -> str: + return f"Hello, {name}!" + + # Mount with custom separator + main_app.mount("sub", sub_app, tool_separator="-") + + # Tool should be accessible with custom separator + tools = await main_app.get_tools() + assert "sub-greet" in tools + + # Call the tool + result = await main_app._mcp_call_tool("sub-greet", {"name": "World"}) + assert isinstance(result[0], TextContent) + assert result[0].text == "Hello, World!" + + async def test_unmount_server(self): + """Test unmounting a server removes access to its tools.""" + main_app = FastMCP("MainApp") + sub_app = FastMCP("SubApp") + + @sub_app.tool() + def sub_tool() -> str: + return "This is from the sub app" + + # Mount the sub-app + main_app.mount("sub", sub_app) + + # Verify it was mounted + tools = await main_app.get_tools() + assert "sub_sub_tool" in tools + + # Unmount the sub-app + main_app.unmount("sub") + + # Verify it was unmounted + tools = await main_app.get_tools() + assert "sub_sub_tool" not in tools + + # Calling the tool should fail + with pytest.raises(NotFoundError, match="Unknown tool: sub_sub_tool"): + await main_app._mcp_call_tool("sub_sub_tool", {}) + + +class TestMultipleServerMount: + """Test mounting multiple servers simultaneously.""" + + async def test_mount_multiple_servers(self): + """Test mounting multiple servers with different prefixes.""" + main_app = FastMCP("MainApp") + weather_app = FastMCP("WeatherApp") + news_app = FastMCP("NewsApp") + + @weather_app.tool() + def get_forecast() -> str: + return "Weather forecast" + + @news_app.tool() + def get_headlines() -> str: + return "News headlines" + + # Mount both apps + main_app.mount("weather", weather_app) + main_app.mount("news", news_app) + + # Check both are accessible + tools = await main_app.get_tools() + assert "weather_get_forecast" in tools + assert "news_get_headlines" in tools + + # Call tools from both mounted servers + result1 = await main_app._mcp_call_tool("weather_get_forecast", {}) + assert isinstance(result1[0], TextContent) + assert result1[0].text == "Weather forecast" + + result2 = await main_app._mcp_call_tool("news_get_headlines", {}) + assert isinstance(result2[0], TextContent) + assert result2[0].text == "News headlines" + + async def test_mount_same_prefix(self): + """Test that mounting with the same prefix replaces the previous mount.""" + main_app = FastMCP("MainApp") + first_app = FastMCP("FirstApp") + second_app = FastMCP("SecondApp") + + @first_app.tool() + def first_tool() -> str: + return "First app tool" + + @second_app.tool() + def second_tool() -> str: + return "Second app tool" + + # Mount first app + main_app.mount("api", first_app) + tools = await main_app.get_tools() + assert "api_first_tool" in tools + + # Mount second app with same prefix + main_app.mount("api", second_app) + tools = await main_app.get_tools() + + # First app's tool should no longer be accessible + assert "api_first_tool" not in tools + + # Second app's tool should be accessible + assert "api_second_tool" in tools + + +class TestDynamicChanges: + """Test that changes to mounted servers are reflected dynamically.""" + + async def test_adding_tool_after_mounting(self): + """Test that tools added after mounting are accessible.""" + main_app = FastMCP("MainApp") + sub_app = FastMCP("SubApp") + + # Mount the sub-app before adding any tools + main_app.mount("sub", sub_app) + + # Initially, there should be no tools from sub_app + tools = await main_app.get_tools() + assert not any(key.startswith("sub_") for key in tools) + + # Add a tool to the sub-app after mounting + @sub_app.tool() + def dynamic_tool() -> str: + return "Added after mounting" + + # The tool should be accessible through the main app + tools = await main_app.get_tools() + assert "sub_dynamic_tool" in tools + + # Call the dynamically added tool + result = await main_app._mcp_call_tool("sub_dynamic_tool", {}) + assert isinstance(result[0], TextContent) + assert result[0].text == "Added after mounting" + + async def test_removing_tool_after_mounting(self): + """Test that tools removed from mounted servers are no longer accessible.""" + main_app = FastMCP("MainApp") + sub_app = FastMCP("SubApp") + + @sub_app.tool() + def temp_tool() -> str: + return "Temporary tool" + + # Mount the sub-app + main_app.mount("sub", sub_app) + + # Initially, the tool should be accessible + tools = await main_app.get_tools() + assert "sub_temp_tool" in tools + + # Remove the tool from sub_app + sub_app._tool_manager._tools.pop("temp_tool") + + # The tool should no longer be accessible + # Refresh the cache by clearing it + main_app._cache.cache.clear() + tools = await main_app.get_tools() + assert "sub_temp_tool" not in tools + + +class TestResourcesAndTemplates: + """Test mounting with resources and resource templates.""" + + async def test_mount_with_resources(self): + """Test mounting a server with resources.""" + main_app = FastMCP("MainApp") + data_app = FastMCP("DataApp") + + @data_app.resource(uri="data://users") + async def get_users(): + return ["user1", "user2"] + + # Mount the data app + main_app.mount("data", data_app) + + # Resource should be accessible through main app + resources = await main_app.get_resources() + assert any("data+data://users" in str(uri) for uri in resources) + + async with Client(main_app) as client: + resource = await client.read_resource("data+data://users") + assert isinstance(resource[0], TextResourceContents) + assert resource[0].text == '["user1", "user2"]' + + async def test_mount_with_resource_templates(self): + """Test mounting a server with resource templates.""" + main_app = FastMCP("MainApp") + user_app = FastMCP("UserApp") + + @user_app.resource(uri="users://{user_id}/profile") + def get_user_profile(user_id: str) -> dict: + return {"id": user_id, "name": f"User {user_id}"} + + # Mount the user app + main_app.mount("api", user_app) + + # Template should be accessible through main app + templates = await main_app.get_resource_templates() + assert any("api+users://{user_id}/profile" in str(t) for t in templates) + + # Read from the template + result = await main_app._mcp_read_resource("api+users://123/profile") + assert isinstance(result[0], ReadResourceContents) + profile = json.loads(result[0].content) + assert profile["id"] == "123" + assert profile["name"] == "User 123" + + async def test_adding_resource_after_mounting(self): + """Test adding a resource after mounting.""" + main_app = FastMCP("MainApp") + data_app = FastMCP("DataApp") + + # Mount the data app before adding resources + main_app.mount("data", data_app) + + # Add a resource after mounting + @data_app.resource(uri="data://config") + def get_config(): + return {"version": "1.0"} + + # Resource should be accessible through main app + resources = await main_app.get_resources() + assert any("data+data://config" in str(uri) for uri in resources) + + # Read the resource + result = await main_app._mcp_read_resource("data+data://config") + assert isinstance(result[0], ReadResourceContents) + config = json.loads(result[0].content) + assert config["version"] == "1.0" + + +class TestPrompts: + """Test mounting with prompts.""" + + async def test_mount_with_prompts(self): + """Test mounting a server with prompts.""" + main_app = FastMCP("MainApp") + assistant_app = FastMCP("AssistantApp") + + @assistant_app.prompt() + def greeting(name: str) -> str: + return f"Hello, {name}!" + + # Mount the assistant app + main_app.mount("assistant", assistant_app) + + # Prompt should be accessible through main app + prompts = await main_app.get_prompts() + assert "assistant_greeting" in prompts + + # Render the prompt + result = await main_app._mcp_get_prompt("assistant_greeting", {"name": "World"}) + assert result.messages is not None + # The message should contain our greeting text + + async def test_adding_prompt_after_mounting(self): + """Test adding a prompt after mounting.""" + main_app = FastMCP("MainApp") + assistant_app = FastMCP("AssistantApp") + + # Mount the assistant app before adding prompts + main_app.mount("assistant", assistant_app) + + # Add a prompt after mounting + @assistant_app.prompt() + def farewell(name: str) -> str: + return f"Goodbye, {name}!" + + # Prompt should be accessible through main app + prompts = await main_app.get_prompts() + assert "assistant_farewell" in prompts + + # Render the prompt + result = await main_app._mcp_get_prompt("assistant_farewell", {"name": "World"}) + assert result.messages is not None + # The message should contain our farewell text + + +class TestProxyServer: + """Test mounting a proxy server.""" + + async def test_mount_proxy_server(self): + """Test mounting a proxy server.""" + # Create original server + original_server = FastMCP("OriginalServer") + + @original_server.tool() + def get_data(query: str) -> str: + return f"Data for {query}" + + # Create proxy server + proxy_server = FastMCP.from_client( + Client(transport=FastMCPTransport(original_server)) + ) + + # Mount proxy server + main_app = FastMCP("MainApp") + main_app.mount("proxy", proxy_server) + + # Tool should be accessible through main app + tools = await main_app.get_tools() + assert "proxy_get_data" in tools + + # Call the tool + result = await main_app._mcp_call_tool("proxy_get_data", {"query": "test"}) + assert isinstance(result[0], TextContent) + assert result[0].text == "Data for test" + + async def test_dynamically_adding_to_proxied_server(self): + """Test that changes to the original server are reflected in the mounted proxy.""" + # Create original server + original_server = FastMCP("OriginalServer") + + # Create proxy server + proxy_server = FastMCP.from_client( + Client(transport=FastMCPTransport(original_server)) + ) + + # Mount proxy server + main_app = FastMCP("MainApp") + main_app.mount("proxy", proxy_server) + + # Add a tool to the original server + @original_server.tool() + def dynamic_data() -> str: + return "Dynamic data" + + # Tool should be accessible through main app via proxy + tools = await main_app.get_tools() + assert "proxy_dynamic_data" in tools + + # Call the tool + result = await main_app._mcp_call_tool("proxy_dynamic_data", {}) + assert isinstance(result[0], TextContent) + assert result[0].text == "Dynamic data" + + async def test_proxy_server_with_resources(self): + """Test mounting a proxy server with resources.""" + # Create original server + original_server = FastMCP("OriginalServer") + + @original_server.resource(uri="config://settings") + def get_config(): + return {"api_key": "12345"} + + # Create proxy server + proxy_server = FastMCP.from_client( + Client(transport=FastMCPTransport(original_server)) + ) + + # Mount proxy server + main_app = FastMCP("MainApp") + main_app.mount("proxy", proxy_server) + + # Resource should be accessible through main app + result = await main_app._mcp_read_resource("proxy+config://settings") + assert isinstance(result[0], ReadResourceContents) + config = json.loads(result[0].content) + assert config["api_key"] == "12345" + + async def test_proxy_server_with_prompts(self): + """Test mounting a proxy server with prompts.""" + # Create original server + original_server = FastMCP("OriginalServer") + + @original_server.prompt() + def welcome(name: str) -> str: + return f"Welcome, {name}!" + + # Create proxy server + proxy_server = FastMCP.from_client( + Client(transport=FastMCPTransport(original_server)) + ) + + # Mount proxy server + main_app = FastMCP("MainApp") + main_app.mount("proxy", proxy_server) + + # Prompt should be accessible through main app + result = await main_app._mcp_get_prompt("proxy_welcome", {"name": "World"}) + assert result.messages is not None + # The message should contain our welcome text diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index de86b5192..8206fd645 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -523,7 +523,8 @@ class TestCustomToolNames: assert tool.name == "custom_name" assert tool.fn.__name__ == "original_fn" # The tool should not be accessible via its original function name - assert manager.get_tool("original_fn") is None + with pytest.raises(NotFoundError, match="Unknown tool: original_fn"): + manager.get_tool("original_fn") def test_add_tool_object_with_custom_key(self): """Test adding a Tool object with a custom key using add_tool().""" @@ -542,7 +543,8 @@ class TestCustomToolNames: # But the tool's .name is unchanged assert stored.name == "my_tool" # The tool is not accessible under its original name - assert manager.get_tool("my_tool") is None + with pytest.raises(NotFoundError, match="Unknown tool: my_tool"): + manager.get_tool("my_tool") async def test_call_tool_with_custom_name(self): """Test calling a tool added with a custom name."""