diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 077e6509b..0fb846e11 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -186,8 +186,10 @@ class FastMCPProxy(FastMCP): else: raise e for tool in client_tools: - tool_proxy = await ProxyTool.from_client(self.client, tool) - tools[tool_proxy.name] = tool_proxy + # 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.name] = tool_proxy return tools @@ -203,8 +205,12 @@ class FastMCPProxy(FastMCP): else: raise e for resource in client_resources: - resource_proxy = await ProxyResource.from_client(self.client, resource) - resources[str(resource_proxy.uri)] = resource_proxy + # don't overwrite resources defined in the server + if str(resource.uri) not in resources: + resource_proxy = await ProxyResource.from_client( + self.client, resource + ) + resources[str(resource_proxy.uri)] = resource_proxy return resources @@ -220,8 +226,12 @@ class FastMCPProxy(FastMCP): else: raise e for template in client_templates: - template_proxy = await ProxyTemplate.from_client(self.client, template) - templates[template_proxy.uri_template] = template_proxy + # 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 templates @@ -237,8 +247,11 @@ class FastMCPProxy(FastMCP): else: raise e for prompt in client_prompts: - prompt_proxy = await ProxyPrompt.from_client(self.client, prompt) - prompts[prompt_proxy.name] = prompt_proxy + # 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 prompts async def _call_tool( diff --git a/tests/server/test_proxy.py b/tests/server/test_proxy.py index 164f33bf8..68524ad2f 100644 --- a/tests/server/test_proxy.py +++ b/tests/server/test_proxy.py @@ -144,6 +144,61 @@ class TestTools: async with Client(proxy_server) as client: await client.call_tool("error_tool", {}) + async def test_proxy_can_overwrite_proxied_tool(self, proxy_server): + """ + Test that a tool defined on the proxy can overwrite the proxied tool with the same name. + """ + + @proxy_server.tool + def greet(name: str, extra: str = "extra") -> str: + return f"Overwritten, {name}! {extra}" + + async with Client(proxy_server) as client: + result = await client.call_tool("greet", {"name": "Marvin", "extra": "abc"}) + assert result[0].text == "Overwritten, Marvin! abc" # type: ignore[attr-defined] + + async def test_proxy_errors_if_overwritten_tool_is_disabled(self, proxy_server): + """ + Test that a tool defined on the proxy is not listed if it is disabled, + and it doesn't fall back to the proxied tool with the same name + """ + + @proxy_server.tool(enabled=False) + def greet(name: str, extra: str = "extra") -> str: + return f"Overwritten, {name}! {extra}" + + async with Client(proxy_server) as client: + with pytest.raises(ToolError, match="Unknown tool"): + await client.call_tool("greet", {"name": "Marvin", "extra": "abc"}) + + async def test_proxy_can_list_overwritten_tool(self, proxy_server): + """ + Test that a tool defined on the proxy is listed instead of the proxied tool + """ + + @proxy_server.tool + def greet(name: str, extra: str = "extra") -> str: + return f"Overwritten, {name}! {extra}" + + async with Client(proxy_server) as client: + tools = await client.list_tools() + greet_tool = next(t for t in tools if t.name == "greet") + assert "extra" in greet_tool.inputSchema["properties"] + + async def test_proxy_can_list_overwritten_tool_if_disabled(self, proxy_server): + """ + Test that a tool defined on the proxy is not listed if it is disabled, + and it doesn't fall back to the proxied tool with the same name + """ + + @proxy_server.tool(enabled=False) + def greet(name: str, extra: str = "extra") -> str: + return f"Overwritten, {name}! {extra}" + + async with Client(proxy_server) as client: + tools = await client.list_tools() + assert not any(t.name == "greet" for t in tools) + class TestResources: async def test_get_resources(self, proxy_server): @@ -184,6 +239,64 @@ class TestResources: async with Client(proxy_server) as client: await client.read_resource("resource://nonexistent") + async def test_proxy_can_overwrite_proxied_resource(self, proxy_server): + """ + Test that a resource defined on the proxy can overwrite the proxied resource with the same URI. + """ + + @proxy_server.resource(uri="resource://wave") + def overwritten_wave() -> str: + return "Overwritten wave! 🌊" + + async with Client(proxy_server) as client: + result = await client.read_resource("resource://wave") + assert result[0].text == "Overwritten wave! 🌊" # type: ignore[attr-defined] + + async def test_proxy_errors_if_overwritten_resource_is_disabled(self, proxy_server): + """ + Test that a resource defined on the proxy is not accessible if it is disabled, + and it doesn't fall back to the proxied resource with the same URI + """ + + @proxy_server.resource(uri="resource://wave", enabled=False) + def overwritten_wave() -> str: + return "Overwritten wave! 🌊" + + async with Client(proxy_server) as client: + with pytest.raises(McpError, match="Unknown resource"): + await client.read_resource("resource://wave") + + async def test_proxy_can_list_overwritten_resource(self, proxy_server): + """ + Test that a resource defined on the proxy is listed instead of the proxied resource + """ + + @proxy_server.resource(uri="resource://wave", name="overwritten_wave") + def overwritten_wave() -> str: + return "Overwritten wave! 🌊" + + async with Client(proxy_server) as client: + resources = await client.list_resources() + wave_resource = next( + r for r in resources if str(r.uri) == "resource://wave" + ) + assert wave_resource.name == "overwritten_wave" + + async def test_proxy_can_list_overwritten_resource_if_disabled(self, proxy_server): + """ + Test that a resource defined on the proxy is not listed if it is disabled, + and it doesn't fall back to the proxied resource with the same URI + """ + + @proxy_server.resource(uri="resource://wave", enabled=False) + def overwritten_wave() -> str: + return "Overwritten wave! 🌊" + + async with Client(proxy_server) as client: + resources = await client.list_resources() + wave_resources = [r for r in resources if str(r.uri) == "resource://wave"] + assert len(wave_resources) == 0 + class TestResourceTemplates: async def test_get_resource_templates(self, proxy_server): @@ -212,6 +325,77 @@ class TestResourceTemplates: proxy_result = await client.read_resource("data://user/1") assert proxy_result == result + async def test_proxy_can_overwrite_proxied_resource_template(self, proxy_server): + """ + Test that a resource template defined on the proxy can overwrite the proxied template with the same URI template. + """ + + @proxy_server.resource(uri="data://user/{user_id}", name="overwritten_get_user") + def overwritten_get_user(user_id: str) -> dict[str, Any]: + return { + "id": user_id, + "name": "Overwritten User", + "active": True, + "extra": "data", + } + + async with Client(proxy_server) as client: + result = await client.read_resource("data://user/1") + user_data = json.loads(result[0].text) # type: ignore[attr-defined] + assert user_data["name"] == "Overwritten User" + assert user_data["extra"] == "data" + + async def test_proxy_errors_if_overwritten_resource_template_is_disabled( + self, proxy_server + ): + """ + Test that a resource template defined on the proxy is not accessible if it is disabled, + and it doesn't fall back to the proxied template with the same URI template + """ + + @proxy_server.resource(uri="data://user/{user_id}", enabled=False) + def overwritten_get_user(user_id: str) -> dict[str, Any]: + return {"id": user_id, "name": "Overwritten User", "active": True} + + async with Client(proxy_server) as client: + with pytest.raises(McpError, match="Unknown resource"): + await client.read_resource("data://user/1") + + async def test_proxy_can_list_overwritten_resource_template(self, proxy_server): + """ + Test that a resource template defined on the proxy is listed instead of the proxied template + """ + + @proxy_server.resource(uri="data://user/{user_id}", name="overwritten_get_user") + def overwritten_get_user(user_id: str) -> dict[str, Any]: + return {"id": user_id, "name": "Overwritten User", "active": True} + + async with Client(proxy_server) as client: + templates = await client.list_resource_templates() + user_template = next( + t for t in templates if t.uriTemplate == "data://user/{user_id}" + ) + assert user_template.name == "overwritten_get_user" + + async def test_proxy_can_list_overwritten_resource_template_if_disabled( + self, proxy_server + ): + """ + Test that a resource template defined on the proxy is not listed if it is disabled, + and it doesn't fall back to the proxied template with the same URI template + """ + + @proxy_server.resource(uri="data://user/{user_id}", enabled=False) + def overwritten_get_user(user_id: str) -> dict[str, Any]: + return {"id": user_id, "name": "Overwritten User", "active": True} + + async with Client(proxy_server) as client: + templates = await client.list_resource_templates() + user_templates = [ + t for t in templates if t.uriTemplate == "data://user/{user_id}" + ] + assert len(user_templates) == 0 + class TestPrompts: async def test_get_prompts_server_method(self, proxy_server: FastMCPProxy): @@ -240,6 +424,70 @@ class TestPrompts: assert result.messages[0].role == "user" assert result.messages[0].content.text == "Welcome to FastMCP, Alice!" # type: ignore[attr-defined] + async def test_proxy_can_overwrite_proxied_prompt(self, proxy_server): + """ + Test that a prompt defined on the proxy can overwrite the proxied prompt with the same name. + """ + + @proxy_server.prompt + def welcome(name: str, extra: str = "friend") -> str: + return f"Overwritten welcome, {name}! You are my {extra}." + + async with Client(proxy_server) as client: + result = await client.get_prompt( + "welcome", {"name": "Alice", "extra": "colleague"} + ) + assert result.messages[0].role == "user" + assert ( + result.messages[0].content.text # type: ignore[attr-defined] + == "Overwritten welcome, Alice! You are my colleague." + ) + + async def test_proxy_errors_if_overwritten_prompt_is_disabled(self, proxy_server): + """ + Test that a prompt defined on the proxy is not accessible if it is disabled, + and it doesn't fall back to the proxied prompt with the same name + """ + + @proxy_server.prompt(enabled=False) + def welcome(name: str, extra: str = "friend") -> str: + return f"Overwritten welcome, {name}! You are my {extra}." + + async with Client(proxy_server) as client: + with pytest.raises(McpError, match="Unknown prompt"): + await client.get_prompt("welcome", {"name": "Alice"}) + + async def test_proxy_can_list_overwritten_prompt(self, proxy_server): + """ + Test that a prompt defined on the proxy is listed instead of the proxied prompt + """ + + @proxy_server.prompt + def welcome(name: str, extra: str = "friend") -> str: + return f"Overwritten welcome, {name}! You are my {extra}." + + async with Client(proxy_server) as client: + prompts = await client.list_prompts() + welcome_prompt = next(p for p in prompts if p.name == "welcome") + # Check that the overwritten prompt has the additional 'extra' parameter + param_names = [arg.name for arg in welcome_prompt.arguments or []] + assert "extra" in param_names + + async def test_proxy_can_list_overwritten_prompt_if_disabled(self, proxy_server): + """ + Test that a prompt defined on the proxy is not listed if it is disabled, + and it doesn't fall back to the proxied prompt with the same name + """ + + @proxy_server.prompt(enabled=False) + def welcome(name: str, extra: str = "friend") -> str: + return f"Overwritten welcome, {name}! You are my {extra}." + + async with Client(proxy_server) as client: + prompts = await client.list_prompts() + welcome_prompts = [p for p in prompts if p.name == "welcome"] + assert len(welcome_prompts) == 0 + async def test_proxy_handles_multiple_concurrent_tasks_correctly( proxy_server: FastMCPProxy,