Harden proxy metadata reads

🤖 Generated with OpenAI Codex

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

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