Spaces:
Running
Running
Jeremiah Lowin commited on
Add test for proxy tags visibility (#1302)
Browse files
tests/server/middleware/test_middleware.py
CHANGED
|
@@ -122,7 +122,7 @@ def recording_middleware():
|
|
| 122 |
def mcp_server(recording_middleware):
|
| 123 |
mcp = FastMCP()
|
| 124 |
|
| 125 |
-
@mcp.tool
|
| 126 |
def add(a: int, b: int) -> int:
|
| 127 |
return a + b
|
| 128 |
|
|
@@ -761,3 +761,27 @@ class TestProxyServer:
|
|
| 761 |
assert recording_middleware.assert_called(hook="on_request", at_least=2)
|
| 762 |
assert recording_middleware.assert_called(hook="on_call_tool", at_least=1)
|
| 763 |
assert recording_middleware.assert_called(hook="on_list_tools", at_least=1)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 122 |
def mcp_server(recording_middleware):
|
| 123 |
mcp = FastMCP()
|
| 124 |
|
| 125 |
+
@mcp.tool(tags={"add-tool"})
|
| 126 |
def add(a: int, b: int) -> int:
|
| 127 |
return a + b
|
| 128 |
|
|
|
|
| 761 |
assert recording_middleware.assert_called(hook="on_request", at_least=2)
|
| 762 |
assert recording_middleware.assert_called(hook="on_call_tool", at_least=1)
|
| 763 |
assert recording_middleware.assert_called(hook="on_list_tools", at_least=1)
|
| 764 |
+
|
| 765 |
+
async def test_proxied_tags_are_visible_to_middleware(
|
| 766 |
+
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
|
| 767 |
+
):
|
| 768 |
+
"""Tests that tags on remote FastMCP servers are visible to middleware
|
| 769 |
+
via proxy. See https://github.com/jlowin/fastmcp/issues/1300"""
|
| 770 |
+
proxy_server = FastMCP.as_proxy(mcp_server, name="Proxy Server")
|
| 771 |
+
|
| 772 |
+
TAGS = []
|
| 773 |
+
|
| 774 |
+
class TagMiddleware(Middleware):
|
| 775 |
+
async def on_list_tools(self, context: MiddlewareContext, call_next):
|
| 776 |
+
nonlocal TAGS
|
| 777 |
+
result = await call_next(context)
|
| 778 |
+
for tool in result:
|
| 779 |
+
TAGS.append(tool.tags)
|
| 780 |
+
return result
|
| 781 |
+
|
| 782 |
+
proxy_server.add_middleware(TagMiddleware())
|
| 783 |
+
|
| 784 |
+
async with Client(proxy_server) as client:
|
| 785 |
+
await client.list_tools()
|
| 786 |
+
|
| 787 |
+
assert TAGS == [{"add-tool"}, set(), set(), set()]
|