From c07277fa136f536b2e7d5022e7df7e9692322e59 Mon Sep 17 00:00:00 2001 From: Jake Kaplan Date: Thu, 6 Aug 2026 08:29:16 -0400 Subject: [PATCH] Harden proxy metadata reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🤖 Generated with OpenAI Codex --- .../fastmcp/server/providers/proxy.py | 43 ++++++---- .../providers/proxy/test_server_metadata.py | 80 +++++++++++++++++++ 2 files changed, 106 insertions(+), 17 deletions(-) diff --git a/fastmcp_slim/fastmcp/server/providers/proxy.py b/fastmcp_slim/fastmcp/server/providers/proxy.py index 01e64321e..5ad740c7b 100644 --- a/fastmcp_slim/fastmcp/server/providers/proxy.py +++ b/fastmcp_slim/fastmcp/server/providers/proxy.py @@ -1065,22 +1065,31 @@ class ProxyMetadataMiddleware(Middleware): self.client_factory = provider.client_factory self.identity = identity - async def _read_upstream(self) -> _UpstreamServerMetadata | None: + async def _read_upstream( + self, context: Context | None + ) -> _UpstreamServerMetadata | None: + client = self.client_factory() + if inspect.isawaitable(client): + client = cast(Client, await client) + + connected = client.is_connected() + # A pinned modern client without a prior discovery result adopts a + # synthesized result without contacting the server. Probe with a copy so + # metadata comes from the backend without changing the configured mode. + if ( + not connected + and client.mode in MODERN_PROTOCOL_VERSIONS + and client._prior_discover is None + ): + client = client.new() + client.mode = "auto" + + if context is not None: + _stash_proxy_request_context(client, context) + if connected: + return _UpstreamServerMetadata.from_client(client) + try: - client = self.client_factory() - if inspect.isawaitable(client): - client = cast(Client, await client) - - if client.is_connected(): - return _UpstreamServerMetadata.from_client(client) - - # A pinned modern client adopts a synthesized DiscoverResult without - # contacting the server. Probe with a copy so metadata comes from the - # backend without changing the factory's configured mode. - if client.mode in MODERN_PROTOCOL_VERSIONS: - client = client.new() - client.mode = "auto" - async with client: return _UpstreamServerMetadata.from_client(client) except (MCPError, *_PROXY_TRANSPORT_ERRORS) as error: @@ -1122,7 +1131,7 @@ class ProxyMetadataMiddleware(Middleware): result = await call_next(context) if result is None: return None - upstream = await self._read_upstream() + upstream = await self._read_upstream(context.fastmcp_context) if upstream is None: return result return result.model_copy(update=self._updates(result, upstream)) @@ -1133,7 +1142,7 @@ class ProxyMetadataMiddleware(Middleware): call_next: CallNext[mcp_types.DiscoverRequest, mcp_types.DiscoverResult], ) -> mcp_types.DiscoverResult: result = await call_next(context) - upstream = await self._read_upstream() + upstream = await self._read_upstream(context.fastmcp_context) if upstream is None: return result return result.model_copy(update=self._updates(result, upstream)) diff --git a/tests/server/providers/proxy/test_server_metadata.py b/tests/server/providers/proxy/test_server_metadata.py index 18fe6f95f..db96f111b 100644 --- a/tests/server/providers/proxy/test_server_metadata.py +++ b/tests/server/providers/proxy/test_server_metadata.py @@ -9,6 +9,7 @@ from mcp import MCPError from mcp_types.version import MODERN_PROTOCOL_VERSIONS from fastmcp import Client, FastMCP +from fastmcp.client.logging import LogMessage from fastmcp.client.transports import StreamableHttpTransport from fastmcp.server import create_proxy from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext @@ -222,6 +223,85 @@ async def test_frontend_values_take_precedence(frontend_mode: str): assert result.meta["com.example/upstream"] == {"enabled": True} +async def test_forwards_backend_logs_while_reading_metadata(): + messages: list[str] = [] + + class LogOnInitialize(Middleware): + async def on_initialize( + self, + context: MiddlewareContext[mcp_types.InitializeRequest], + call_next: CallNext[ + mcp_types.InitializeRequest, mcp_types.InitializeResult | None + ], + ) -> mcp_types.InitializeResult | None: + result = await call_next(context) + assert context.fastmcp_context is not None + await context.fastmcp_context.log("metadata connection") + return result + + async def capture_log(message: LogMessage) -> None: + messages.append(message.data["msg"]) + + upstream = FastMCP("upstream", middleware=[LogOnInitialize()]) + proxy = create_proxy(upstream) + + async with Client(proxy, mode="legacy", log_handler=capture_log): + pass + + assert messages == ["metadata connection"] + + +async def test_pinned_client_uses_prior_discover_metadata(): + prior_info = mcp_types.Implementation(name="prior", version="1.0") + prior = mcp_types.DiscoverResult( + supported_versions=[MODERN_PROTOCOL_VERSIONS[0]], + capabilities=mcp_types.ServerCapabilities(), + instructions="prior instructions", + meta={ + mcp_types.SERVER_INFO_META_KEY: prior_info.model_dump( + by_alias=True, mode="json" + ), + "com.example/prior": True, + }, + ) + provider = ProxyProvider( + lambda: ProxyClient( + make_upstream(), + mode=MODERN_PROTOCOL_VERSIONS[0], + prior_discover=prior, + ) + ) + gateway = FastMCP( + "gateway", + providers=[provider], + middleware=[ProxyMetadataMiddleware(provider, identity="upstream")], + ) + + async with Client(gateway, mode="auto") as client: + result = client.session.discover_result + assert result is not None + assert client.instructions == "prior instructions" + assert client.server_info == prior_info + assert result.meta is not None + assert result.meta["com.example/prior"] is True + + +async def test_client_factory_errors_are_not_swallowed(): + def broken_factory() -> Client: + raise RuntimeError("broken client factory") + + provider = ProxyProvider(broken_factory) + gateway = FastMCP( + "gateway", + providers=[provider], + middleware=[ProxyMetadataMiddleware(provider)], + ) + + with pytest.raises(MCPError): + async with Client(gateway, mode="legacy"): + pass + + @pytest.mark.parametrize("frontend_mode", ["legacy", "auto"]) async def test_unavailable_backend_does_not_block_connection(frontend_mode: str): port = find_available_port()