Jeremiah Lowin commited on
Commit
08f0e73
·
unverified ·
1 Parent(s): 2e693f4

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()]