Simplify proxy metadata client lifecycle

🤖 Generated with OpenAI Codex

fastmcp_slim/fastmcp/server/providers/proxy.py

@@ -1166,6 +1166,19 @@ class ProxyMetadataMiddleware(Middleware):
         self.client_factory = provider.client_factory
         self.identity = identity
 
+    async def _read_connected(self, client: Client) -> _UpstreamServerMetadata | None:
+        if client.mode in MODERN_PROTOCOL_VERSIONS and client.prior_discover is None:
+            raw = await client.session.send_discover(client.mode)
+            result_type = raw.get("resultType")
+            if (
+                isinstance(result_type, str)
+                and result_type not in mcp_types.CORE_RESULT_TYPES
+            ):
+                return None
+            result = mcp_types.DiscoverResult.model_validate(raw)
+            return _UpstreamServerMetadata.from_discover(result)
+        return _UpstreamServerMetadata.from_client(client)
+
     async def _read_upstream(
         self, context: Context | None
     ) -> _UpstreamServerMetadata | None:
@@ -1173,35 +1186,14 @@ class ProxyMetadataMiddleware(Middleware):
         if inspect.isawaitable(client):
             client = cast(Client, await client)
 
-        connected = client.is_connected()
-        synthesized_discover = (
-            client.mode in MODERN_PROTOCOL_VERSIONS and client.prior_discover is None
-        )
-        # A disconnected pinned client can probe with a copy. A connected client
-        # must keep using its existing transport, so query discovery directly
-        # without replacing the result already adopted by its session.
-        if synthesized_discover and not connected:
-            client = client.new()
-            client.mode = "auto"
-
         if context is not None:
             _stash_proxy_request_context(client, context)
 
         try:
-            if synthesized_discover and connected:
-                raw = await client.session.send_discover(client.mode)
-                result_type = raw.get("resultType")
-                if (
-                    isinstance(result_type, str)
-                    and result_type not in mcp_types.CORE_RESULT_TYPES
-                ):
-                    return None
-                result = mcp_types.DiscoverResult.model_validate(raw)
-                return _UpstreamServerMetadata.from_discover(result)
-            if connected:
-                return _UpstreamServerMetadata.from_client(client)
+            if client.is_connected():
+                return await self._read_connected(client)
             async with client:
-                return _UpstreamServerMetadata.from_client(client)
+                return await self._read_connected(client)
         except (MCPError, *_PROXY_TRANSPORT_ERRORS) as error:
             logger.debug("Could not read upstream server metadata: %r", error)
             return None

tests/server/providers/proxy/test_server_metadata.py

@@ -19,6 +19,7 @@ from fastmcp.server.providers.proxy import (
     ProxyInitializeMiddleware,
     ProxyMetadataMiddleware,
     ProxyProvider,
+    StatefulProxyClient,
 )
 from fastmcp.utilities.http import find_available_port
 
@@ -272,7 +273,7 @@ async def test_forwards_backend_logs_while_reading_metadata():
     assert messages == ["metadata connection"]
 
 
-async def test_pinned_clients_use_available_metadata():
+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]],
@@ -306,6 +307,8 @@ async def test_pinned_clients_use_available_metadata():
         assert result.meta is not None
         assert result.meta["com.example/prior"] is True
 
+
+async def test_connected_pinned_client_probes_without_adopting_metadata():
     version = MODERN_PROTOCOL_VERSIONS[0]
     upstream = make_upstream()
     async with Client(upstream, mode=version) as backend_client:
@@ -320,11 +323,54 @@ async def test_pinned_clients_use_available_metadata():
             assert result.meta is not None
             assert result.meta["com.example/upstream"] == {"enabled": True}
 
-        # The metadata probe must not replace the connected client's adopted
-        # synthetic result.
         assert backend_client.instructions is None
 
 
+async def test_disconnected_pinned_client_is_not_cloned():
+    class UnclonableProxyClient(ProxyClient):
+        def new(self) -> ProxyClient:
+            raise AssertionError("metadata client must not be cloned")
+
+    version = MODERN_PROTOCOL_VERSIONS[0]
+    provider = ProxyProvider(
+        lambda: UnclonableProxyClient(make_upstream(), mode=version)
+    )
+    gateway = FastMCP(
+        "gateway",
+        providers=[provider],
+        middleware=[ProxyMetadataMiddleware(provider, identity="upstream")],
+    )
+
+    async with Client(gateway, mode="auto") as client:
+        assert client.instructions == "upstream instructions"
+        assert client.server_info == UPSTREAM_INFO
+
+
+async def test_stateful_pinned_metadata_uses_registered_client_lifecycle():
+    created: list[StatefulProxyClient] = []
+
+    class TrackingStatefulProxyClient(StatefulProxyClient):
+        def new(self) -> StatefulProxyClient:
+            client = super().new()
+            created.append(client)
+            return client
+
+    version = MODERN_PROTOCOL_VERSIONS[0]
+    stateful_client = TrackingStatefulProxyClient(make_upstream(), mode=version)
+    proxy = FastMCPProxy(
+        name="stateful-proxy",
+        client_factory=stateful_client.new_stateful,
+        identity="upstream",
+    )
+
+    async with Client(proxy, mode="auto") as client:
+        assert client.instructions == "upstream instructions"
+        assert client.server_info == UPSTREAM_INFO
+
+    assert len(created) == 1
+    assert not created[0].is_connected()
+
+
 async def test_client_factory_errors_are_not_swallowed():
     def broken_factory() -> Client:
         raise RuntimeError("broken client factory")