diff --git a/fastmcp_slim/fastmcp/server/middleware/proxy.py b/fastmcp_slim/fastmcp/server/middleware/proxy.py index d36187f52..a00899f19 100644 --- a/fastmcp_slim/fastmcp/server/middleware/proxy.py +++ b/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)) diff --git a/tests/server/providers/proxy/test_negotiation_metadata.py b/tests/server/providers/proxy/test_negotiation_metadata.py index 9f62ae2f8..98caeba68 100644 --- a/tests/server/providers/proxy/test_negotiation_metadata.py +++ b/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"])