ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-22 13:34:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Merge pull request #782 from jlowin/proxy

Browse source 
Ensure proxies can overwrite remote tools without falling back to the remote
This commit is contained in:
Jeremiah Lowin 2025-06-10 11:52:08 -04:00 • committed by GitHub
commit 6ea5d91c51
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 269 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

29
src/fastmcp/server/proxy.py
View file

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

248
tests/server/test_proxy.py
View file

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