diff --git a/fastmcp_slim/fastmcp/server/providers/proxy.py b/fastmcp_slim/fastmcp/server/providers/proxy.py index 47ed7f5f6..180dbc371 100644 --- a/fastmcp_slim/fastmcp/server/providers/proxy.py +++ b/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 diff --git a/tests/server/providers/proxy/test_server_metadata.py b/tests/server/providers/proxy/test_server_metadata.py index 9e9b88b1a..ec5fd9dcf 100644 --- a/tests/server/providers/proxy/test_server_metadata.py +++ b/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")