Tighten negotiation metadata forwarding

🤖 Generated with OpenAI Codex

fastmcp_slim/fastmcp/server/middleware/proxy.py

@@ -60,13 +60,8 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
         self.client_factory = provider.client_factory
         self.identity = identity
 
-    async def _read_upstream(
-        self, context: MiddlewareContext[Any]
-    ) -> _NegotiationMetadata | None:
-        from fastmcp.server.providers.proxy import (
-            _PROXY_TRANSPORT_ERRORS,
-            _stash_proxy_request_context,
-        )
+    async def _read_upstream(self) -> _NegotiationMetadata | None:
+        from fastmcp.server.providers.proxy import _PROXY_TRANSPORT_ERRORS
 
         try:
             client = self.client_factory()
@@ -76,61 +71,42 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
             if client.is_connected():
                 return _NegotiationMetadata.from_client(client)
 
-            # Metadata negotiation is independent of component operations. Use a
-            # fresh session so a factory that returns one reusable disconnected
-            # client is neither mutated nor made unavailable to a concurrent call.
-            client = client.new()
-
             # A pinned modern client adopts a synthesized DiscoverResult without
-            # contacting the server. Metadata reads need the real result, so probe
-            # modern discovery and retain its normal legacy fallback instead.
+            # 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"
 
-            if context.fastmcp_context is not None:
-                _stash_proxy_request_context(client, context.fastmcp_context)
             async with client:
                 return _NegotiationMetadata.from_client(client)
         except (MCPError, *_PROXY_TRANSPORT_ERRORS) as error:
             logger.debug("Could not read upstream negotiation metadata: %r", error)
             return None
 
-    def _merge_meta(
-        self,
-        frontend_meta: dict[str, Any],
-        upstream: _NegotiationMetadata,
-    ) -> dict[str, Any]:
-        # Identity is handled as a policy, not an ordinary `_meta` collision.
-        # In particular, a modern backend's identity stamp must not leak onto a
-        # legacy frontend while `identity="proxy"` keeps the canonical field.
-        merged = {
-            key: value
-            for key, value in upstream.meta.items()
-            if key != mcp_types.SERVER_INFO_META_KEY
-        }
-        merged.update(frontend_meta)
-        if self.identity == "upstream" and upstream.server_info is not None:
-            merged[mcp_types.SERVER_INFO_META_KEY] = upstream.server_info.model_dump(
-                by_alias=True, mode="json", exclude_none=True
-            )
-        return merged
-
     def _updates(
         self,
         result: mcp_types.InitializeResult | mcp_types.DiscoverResult,
         upstream: _NegotiationMetadata,
     ) -> dict[str, Any]:
-        updates: dict[str, Any] = {
-            "meta": self._merge_meta(dict(result.meta or {}), upstream) or None
+        meta = {
+            key: value
+            for key, value in upstream.meta.items()
+            if key != mcp_types.SERVER_INFO_META_KEY
         }
+        meta.update(result.meta or {})
+
+        updates: dict[str, Any] = {"meta": meta or None}
         if result.instructions is None and upstream.instructions is not None:
             updates["instructions"] = upstream.instructions
-        if (
-            isinstance(result, mcp_types.InitializeResult)
-            and self.identity == "upstream"
-            and upstream.server_info is not None
-        ):
-            updates["server_info"] = upstream.server_info
+        if self.identity == "upstream" and upstream.server_info is not None:
+            if isinstance(result, mcp_types.InitializeResult):
+                updates["server_info"] = upstream.server_info
+            else:
+                meta[mcp_types.SERVER_INFO_META_KEY] = upstream.server_info.model_dump(
+                    by_alias=True, mode="json", exclude_none=True
+                )
+                updates["meta"] = meta
         return updates
 
     async def on_initialize(
@@ -143,7 +119,7 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
         result = await call_next(context)
         if result is None:
             return None
-        upstream = await self._read_upstream(context)
+        upstream = await self._read_upstream()
         if upstream is None:
             return result
         return result.model_copy(update=self._updates(result, upstream))
@@ -154,7 +130,7 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
         call_next: CallNext[mcp_types.DiscoverRequest, mcp_types.DiscoverResult],
     ) -> mcp_types.DiscoverResult:
         result = await call_next(context)
-        upstream = await self._read_upstream(context)
+        upstream = await self._read_upstream()
         if upstream is None:
             return result
         return result.model_copy(update=self._updates(result, upstream))

tests/server/providers/proxy/test_negotiation_metadata.py

@@ -195,6 +195,13 @@ async def test_identity_policy_forwards_full_implementation(
             assert client.server_info.version == "9.8.7"
         else:
             assert client.server_info == UPSTREAM_INFO
+            result = client.session.initialize_result or client.session.discover_result
+            assert result is not None
+            if isinstance(result, mcp_types.InitializeResult):
+                assert mcp_types.SERVER_INFO_META_KEY not in (result.meta or {})
+            else:
+                assert result.meta is not None
+                assert result.meta[mcp_types.SERVER_INFO_META_KEY]["name"] == "upstream"
 
 
 @pytest.mark.parametrize("frontend_mode", ["legacy", "auto"])