Merge pull request #782 from jlowin/proxy

Ensure proxies can overwrite remote tools without falling back to the remote

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(

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,