From f3f99a3c7fea33c7fe74262f36323efc448c35d1 Mon Sep 17 00:00:00 2001 From: Jake Kaplan Date: Wed, 5 Aug 2026 19:48:58 -0400 Subject: [PATCH] Forward proxy negotiation metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit πŸ€– Generated with OpenAI Codex --- docs/servers/middleware.mdx | 25 ++ docs/servers/providers/proxy.mdx | 32 ++ fastmcp_slim/fastmcp/server/_negotiation.py | 12 + fastmcp_slim/fastmcp/server/low_level.py | 65 +++- .../fastmcp/server/middleware/__init__.py | 16 + .../fastmcp/server/middleware/middleware.py | 9 + .../fastmcp/server/middleware/proxy.py | 209 +++++++++++++ .../fastmcp/server/providers/proxy.py | 204 ++++-------- .../middleware/test_discovery_middleware.py | 31 ++ .../proxy/test_negotiation_metadata.py | 292 ++++++++++++++++++ .../providers/proxy/test_proxy_server.py | 67 ++-- 11 files changed, 771 insertions(+), 191 deletions(-) create mode 100644 fastmcp_slim/fastmcp/server/_negotiation.py create mode 100644 fastmcp_slim/fastmcp/server/middleware/proxy.py create mode 100644 tests/server/middleware/test_discovery_middleware.py create mode 100644 tests/server/providers/proxy/test_negotiation_metadata.py diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx index dd841f00d..ffdd1ae3c 100644 --- a/docs/servers/middleware.mdx +++ b/docs/servers/middleware.mdx @@ -310,6 +310,31 @@ async def on_initialize(self, context: MiddlewareContext, call_next): Rejection works only **before** `call_next()`. Raising `McpError` afterward logs the error without sending it β€” the client still receives a successful initialize response. +#### on_discover + +Called when a modern client negotiates through `server/discover`. Like `on_initialize`, this hook receives the typed request and can safely inspect or transform the typed `DiscoverResult` returned by `call_next()`. + +```python +import mcp_types + +from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext + + +class DiscoveryMiddleware(Middleware): + async def on_discover( + self, + context: MiddlewareContext[mcp_types.DiscoverRequest], + call_next: CallNext[ + mcp_types.DiscoverRequest, + mcp_types.DiscoverResult, + ], + ) -> mcp_types.DiscoverResult: + result = await call_next(context) + return result.model_copy(update={"instructions": "Custom instructions"}) +``` + +**Returns:** `DiscoverResult` β€” the transformed value is serialized to the client. Fields such as `supported_versions`, `capabilities`, `ttl_ms`, and `cache_scope` describe the server's public behavior, so middleware should only change them when it also changes that behavior. + ### Raw Handler For complete control over all messages, override `__call__` instead of individual hooks: diff --git a/docs/servers/providers/proxy.mdx b/docs/servers/providers/proxy.mdx index df09b1c80..0cff53f78 100644 --- a/docs/servers/providers/proxy.mdx +++ b/docs/servers/providers/proxy.mdx @@ -388,6 +388,38 @@ Only reuse sessions when you know the backend is stateless (e.g. stateless HTTP) ## Advanced Usage +### Controlled Gateways with Negotiation Metadata + +`ProxyProvider` deliberately owns only remote components. To also forward optional server negotiation metadata, add `ProxyNegotiationMetadataMiddleware` explicitly. This is useful for gateways that need transforms, local components, or other middleware without adopting all of `FastMCPProxy`: + +```python +from fastmcp import FastMCP +from fastmcp.server.middleware import ProxyNegotiationMetadataMiddleware +from fastmcp.server.providers.proxy import ProxyClient, ProxyProvider + + +def create_backend_client() -> ProxyClient: + return ProxyClient("http://backend:8000/mcp", mode="auto") + + +backend = ProxyProvider(create_backend_client) +metadata = ProxyNegotiationMetadataMiddleware( + backend, + identity="proxy", +) + +gateway = FastMCP( + "Controlled Gateway", + instructions="Use only approved gateway operations.", + providers=[backend], + middleware=[metadata], +) +``` + +The middleware works across both the legacy `initialize` handshake and modern `server/discover`. It forwards upstream instructions, namespaced `_meta`, unknown extension fields, and optionally the full upstream `serverInfo`. The frontend remains authoritative for protocol versions, capabilities, cache policy, and `resultType`, since those fields must describe the gateway after its filtering and transforms. + +`identity="proxy"` (the default) retains the gateway's `serverInfo`; use `identity="upstream"` to expose the backend's complete implementation identity. Explicit frontend instructions and frontend metadata win on collision. Metadata is fetched lazily during negotiation, and an unavailable backend does not fail negotiation solely because optional metadata could not be readβ€”the first proxied operation reports that failure. + ### FastMCPProxy Class For explicit session control, use `FastMCPProxy` directly: diff --git a/fastmcp_slim/fastmcp/server/_negotiation.py b/fastmcp_slim/fastmcp/server/_negotiation.py new file mode 100644 index 000000000..523a26cd6 --- /dev/null +++ b/fastmcp_slim/fastmcp/server/_negotiation.py @@ -0,0 +1,12 @@ +"""Typed negotiation results that retain protocol extension fields.""" + +import mcp_types +from pydantic import ConfigDict + + +class _ExtensibleInitializeResult(mcp_types.InitializeResult): + model_config = ConfigDict(extra="allow") + + +class _ExtensibleDiscoverResult(mcp_types.DiscoverResult): + model_config = ConfigDict(extra="allow") diff --git a/fastmcp_slim/fastmcp/server/low_level.py b/fastmcp_slim/fastmcp/server/low_level.py index 851bd74a7..471f571f3 100644 --- a/fastmcp_slim/fastmcp/server/low_level.py +++ b/fastmcp_slim/fastmcp/server/low_level.py @@ -29,6 +29,10 @@ from mcp.shared.exceptions import MCPError from pydantic import ValidationError from fastmcp.apps.config import UI_EXTENSION_ID +from fastmcp.server._negotiation import ( + _ExtensibleDiscoverResult, + _ExtensibleInitializeResult, +) from fastmcp.server.telemetry import seam_span from fastmcp.utilities.logging import get_logger @@ -38,6 +42,7 @@ if TYPE_CHECKING: logger = get_logger(__name__) + # The request methods that FastMCP serves through a handler adapter, each of which # runs the FastMCP middleware chain interior (see MCPOperationsMixin). The root # dispatch leaves these to the interior dispatch and only observes them if they fail before @@ -153,10 +158,11 @@ class FastMCPServerMiddleware: Dispatch shapes: - - ``initialize`` runs the *whole* FastMCP chain here (``on_message`` -> - ``on_request`` -> ``on_initialize``) because there is no interior handler - adapter for it: the SDK builds the ``InitializeResult`` directly, so this is - the only place ``on_initialize`` can observe it or veto with ``MCPError``. + - Negotiation runs the *whole* FastMCP chain here: ``initialize`` dispatches + through ``on_initialize`` and ``server/discover`` through ``on_discover``. + Neither has an interior FastMCP handler adapter, and the SDK serializes both + results before returning through its middleware seam, so this root adapter + restores the typed result before FastMCP middleware observes it. - The component methods (``tools/call``, ``tools/list``, ``resources/read``, ...) still run their FastMCP chain *interior*, in the handler adapter, where ``on_call_tool`` receives the typed component result and a tool exception @@ -192,6 +198,8 @@ class FastMCPServerMiddleware: return await call_next(ctx) if ctx.method == "initialize" and ctx.request_id is not None: return await self._run_initialize_mw(fastmcp, ctx, call_next) + if ctx.method == "server/discover" and ctx.request_id is not None: + return await self._run_discover_mw(fastmcp, ctx, call_next) if ctx.request_id is not None and ctx.method in _INTERIOR_METHODS: return await self._dispatch_component(fastmcp, ctx, call_next) return await self._run_outer_mw(fastmcp, ctx, call_next, _raise=None) @@ -318,6 +326,49 @@ class FastMCPServerMiddleware: for var, token in reversed(tokens): var.reset(token) + async def _run_discover_mw( + self, + fastmcp: FastMCP, + ctx: ServerRequestContext, + call_next: CallNext, + ) -> HandlerResult: + """Run discovery through the typed FastMCP middleware hook.""" + from fastmcp.server.context import Context + from fastmcp.server.middleware.middleware import MiddlewareContext + + params = ctx.params if isinstance(ctx.params, dict) else {} + discover_message = mcp_types.DiscoverRequest.model_validate( + {"method": "server/discover", "params": params}, by_name=False + ) + + async def call_original_handler( + _mw_ctx: MiddlewareContext, + ) -> mcp_types.DiscoverResult: + raw = await call_next(ctx) + if isinstance(raw, mcp_types.DiscoverResult): + return _ExtensibleDiscoverResult.model_validate( + raw.model_dump(by_alias=True) + ) + if isinstance(raw, Mapping): + return _ExtensibleDiscoverResult.model_validate(dict(raw)) + raise TypeError( + "server/discover handler returned " + f"{type(raw).__name__}; expected DiscoverResult or mapping" + ) + + async with Context(fastmcp=fastmcp, session=ctx.session) as fastmcp_ctx: + mw_context = MiddlewareContext( + message=discover_message, + source="client", + type="request", + method="server/discover", + fastmcp_context=fastmcp_ctx, + ) + return await fastmcp._run_middleware( + mw_context, + cast("FastMCPCallNext[Any, Any]", call_original_handler), + ) + async def _run_initialize_mw( self, fastmcp: FastMCP, @@ -357,9 +408,11 @@ class FastMCPServerMiddleware: nonlocal captured_result, call_next_completed raw = await call_next(ctx) if isinstance(raw, mcp_types.InitializeResult): - captured_result = raw + captured_result = _ExtensibleInitializeResult.model_validate( + raw.model_dump(by_alias=True) + ) elif isinstance(raw, Mapping): - captured_result = mcp_types.InitializeResult.model_validate(dict(raw)) + captured_result = _ExtensibleInitializeResult.model_validate(dict(raw)) call_next_completed = True return captured_result if raw is not None else None diff --git a/fastmcp_slim/fastmcp/server/middleware/__init__.py b/fastmcp_slim/fastmcp/server/middleware/__init__.py index 8df6962bd..4c2cf33f7 100644 --- a/fastmcp_slim/fastmcp/server/middleware/__init__.py +++ b/fastmcp_slim/fastmcp/server/middleware/__init__.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + from .authorization import AuthMiddleware from .middleware import ( CallNext, @@ -6,10 +8,24 @@ from .middleware import ( ) from .ping import PingMiddleware +if TYPE_CHECKING: + from .proxy import ( + ProxyNegotiationMetadataMiddleware as ProxyNegotiationMetadataMiddleware, + ) + __all__ = [ "AuthMiddleware", "CallNext", "Middleware", "MiddlewareContext", "PingMiddleware", + "ProxyNegotiationMetadataMiddleware", ] + + +def __getattr__(name: str) -> object: + if name == "ProxyNegotiationMetadataMiddleware": + from .proxy import ProxyNegotiationMetadataMiddleware + + return ProxyNegotiationMetadataMiddleware + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/fastmcp_slim/fastmcp/server/middleware/middleware.py b/fastmcp_slim/fastmcp/server/middleware/middleware.py index 2a112aa5f..60ccd3c9f 100644 --- a/fastmcp_slim/fastmcp/server/middleware/middleware.py +++ b/fastmcp_slim/fastmcp/server/middleware/middleware.py @@ -170,6 +170,8 @@ class Middleware: match context.method: case "initialize": handler = make_handler_wrapper(self.on_initialize, handler) + case "server/discover": + handler = make_handler_wrapper(self.on_discover, handler) case "tools/call": handler = make_handler_wrapper(self.on_call_tool, handler) case "resources/read": @@ -227,6 +229,13 @@ class Middleware: ) -> mt.InitializeResult | None: return await call_next(context) + async def on_discover( + self, + context: MiddlewareContext[mt.DiscoverRequest], + call_next: CallNext[mt.DiscoverRequest, mt.DiscoverResult], + ) -> mt.DiscoverResult: + return await call_next(context) + async def on_call_tool( self, context: MiddlewareContext[mt.CallToolRequestParams], diff --git a/fastmcp_slim/fastmcp/server/middleware/proxy.py b/fastmcp_slim/fastmcp/server/middleware/proxy.py new file mode 100644 index 000000000..0152454d7 --- /dev/null +++ b/fastmcp_slim/fastmcp/server/middleware/proxy.py @@ -0,0 +1,209 @@ +"""Negotiation metadata forwarding for MCP proxy gateways.""" + +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast + +import mcp_types +from mcp.shared.exceptions import MCPError +from mcp_types.version import MODERN_PROTOCOL_VERSIONS + +from fastmcp.client.client import Client +from fastmcp.server._negotiation import ( + _ExtensibleDiscoverResult, + _ExtensibleInitializeResult, +) +from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext +from fastmcp.utilities.logging import get_logger + +if TYPE_CHECKING: + from fastmcp.server.providers.proxy import ProxyProvider + +logger = get_logger(__name__) + +ProxyIdentity = Literal["proxy", "upstream"] + +# Claims describing the connection or the result envelope belong to the public +# gateway. Some are unknown fields when they cross eras, so filter aliases and +# Python field names rather than relying only on the source model's typed fields. +_GATEWAY_OWNED_FIELDS = frozenset( + { + "protocolVersion", + "protocol_version", + "supportedVersions", + "supported_versions", + "capabilities", + "ttlMs", + "ttl_ms", + "cacheScope", + "cache_scope", + "resultType", + "result_type", + "serverInfo", + "server_info", + } +) + + +NegotiationResultT = TypeVar( + "NegotiationResultT", + mcp_types.InitializeResult, + mcp_types.DiscoverResult, +) + + +@dataclass(frozen=True) +class _NegotiationMetadata: + instructions: str | None + server_info: mcp_types.Implementation | None + meta: dict[str, Any] + extensions: dict[str, Any] + + @classmethod + def from_client(cls, client: Client) -> _NegotiationMetadata | None: + result = client.session.initialize_result or client.session.discover_result + if result is None: + return None + return cls( + instructions=result.instructions, + server_info=client.session.server_info, + meta=dict(result.meta or {}), + extensions={ + key: value + for key, value in (result.model_extra or {}).items() + if key not in _GATEWAY_OWNED_FIELDS + }, + ) + + +class ProxyNegotiationMetadataMiddleware(Middleware): + """Forward optional negotiation metadata from a ``ProxyProvider`` backend. + + The frontend always owns protocol versions, capabilities, cache policy, and + result type. Instructions, namespaced metadata, and extension fields are + filled from the backend only where the frontend has no value. ``identity`` + controls whether server identity remains the gateway's or is replaced by + the backend's. + """ + + def __init__( + self, + provider: ProxyProvider, + *, + identity: ProxyIdentity = "proxy", + ) -> None: + if identity not in ("proxy", "upstream"): + raise ValueError("identity must be 'proxy' or 'upstream'") + 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, + ) + + try: + client = self.client_factory() + if inspect.isawaitable(client): + client = cast(Client, await client) + + 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. + if client.mode in MODERN_PROTOCOL_VERSIONS: + 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 _merge_result( + self, + result: NegotiationResultT, + upstream: _NegotiationMetadata, + ) -> NegotiationResultT: + dumped = result.model_dump(by_alias=True) + if isinstance(result, mcp_types.InitializeResult): + forwarded = _ExtensibleInitializeResult.model_validate(dumped) + else: + forwarded = _ExtensibleDiscoverResult.model_validate(dumped) + + frontend_extensions = forwarded.model_extra or {} + updates = { + key: value + for key, value in upstream.extensions.items() + if key not in frontend_extensions + } + updates["meta"] = self._merge_meta(dict(forwarded.meta or {}), upstream) or None + if forwarded.instructions is None and upstream.instructions is not None: + updates["instructions"] = upstream.instructions + if ( + isinstance(forwarded, mcp_types.InitializeResult) + and self.identity == "upstream" + and upstream.server_info is not None + ): + updates["server_info"] = upstream.server_info + return cast("NegotiationResultT", forwarded.model_copy(update=updates)) + + 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) + if result is None: + return None + upstream = await self._read_upstream(context) + if upstream is None: + return result + return self._merge_result(result, upstream) + + async def on_discover( + self, + context: MiddlewareContext[mcp_types.DiscoverRequest], + call_next: CallNext[mcp_types.DiscoverRequest, mcp_types.DiscoverResult], + ) -> mcp_types.DiscoverResult: + result = await call_next(context) + upstream = await self._read_upstream(context) + if upstream is None: + return result + return self._merge_result(result, upstream) diff --git a/fastmcp_slim/fastmcp/server/providers/proxy.py b/fastmcp_slim/fastmcp/server/providers/proxy.py index a304eacc3..7bcdc91b2 100644 --- a/fastmcp_slim/fastmcp/server/providers/proxy.py +++ b/fastmcp_slim/fastmcp/server/providers/proxy.py @@ -50,9 +50,12 @@ from fastmcp.resources.base import ( ResourceResult, ) from fastmcp.resources.template import expand_uri_template +from fastmcp.server._negotiation import ( + _ExtensibleDiscoverResult, + _ExtensibleInitializeResult, +) from fastmcp.server.context import Context from fastmcp.server.dependencies import fastmcp_request_ctx, get_context -from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.server.providers.aggregate import ProviderErrorStrategy from fastmcp.server.providers.base import Provider from fastmcp.server.server import FastMCP @@ -66,7 +69,13 @@ from fastmcp.utilities.versions import VersionSpec, version_sort_key if TYPE_CHECKING: from pathlib import Path + from mcp.client.session import ReceiveResultT + from mcp.shared.dispatcher import ProgressFnT + from mcp.shared.message import ClientMessageMetadata + from pydantic import TypeAdapter + from fastmcp.client.transports import ClientTransport + from fastmcp.server.middleware.proxy import ProxyIdentity logger = get_logger(__name__) @@ -75,15 +84,56 @@ ClientFactoryT = Callable[[], Client] | Callable[[], Awaitable[Client]] class _ForwardingClientSession(ClientSession): - """A session that does not enforce the backend's declared output schema. + """A proxy session that relays backend results without consuming them. - `ClientSession.call_tool` normally validates a tool's structured content - against the output schema the backend advertised, raising if they disagree. - That check belongs to whoever consumes the result. A proxy only relays it, - and the end client runs the same check for itself, so enforcing it mid-path - turns a backend's schema bug into a proxy error and hides the real response. + Tool results skip output-schema validation because the end client performs + that validation itself. Negotiation results retain unknown extension fields + so the proxy middleware can forward them rather than discarding vocabulary + it does not understand. """ + _raw_discover_result: dict[str, Any] | None = None + + async def send_request( + self, + request: mcp_types.ClientRequest | mcp_types.Request[Any, Any], + result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT], + request_read_timeout_seconds: float | None = None, + metadata: ClientMessageMetadata | None = None, + progress_callback: ProgressFnT | None = None, + ) -> ReceiveResultT: + if result_type is mcp_types.InitializeResult: + result = await super().send_request( + request, + _ExtensibleInitializeResult, + request_read_timeout_seconds=request_read_timeout_seconds, + metadata=metadata, + progress_callback=progress_callback, + ) + return cast("ReceiveResultT", result) + return await super().send_request( + request, + result_type, + request_read_timeout_seconds=request_read_timeout_seconds, + metadata=metadata, + progress_callback=progress_callback, + ) + + async def send_discover(self, version: str) -> dict[str, Any]: + raw = await super().send_discover(version) + self._raw_discover_result = raw + return raw + + def adopt( + self, result: mcp_types.InitializeResult | mcp_types.DiscoverResult + ) -> None: + if ( + isinstance(result, mcp_types.DiscoverResult) + and self._raw_discover_result is not None + ): + result = _ExtensibleDiscoverResult.model_validate(self._raw_discover_result) + super().adopt(result) + async def validate_tool_result( self, name: str, result: mcp_types.CallToolResult ) -> None: @@ -228,65 +278,6 @@ def _stash_proxy_request_context(client: Client, ctx: Context) -> None: ) -class ProxyInitializeMiddleware(Middleware): - def __init__(self, proxy: FastMCPProxy) -> None: - self.proxy = proxy - - async def on_initialize( - self, - context: MiddlewareContext[mcp_types.InitializeRequest], - call_next: CallNext[ - mcp_types.InitializeRequest, - mcp_types.InitializeResult | None, - ], - ) -> mcp_types.InitializeResult | None: - client = await self.proxy._get_client() - upstream_instructions: str | None = None - try: - if isinstance(client, ProxyClient): - ctx = context.fastmcp_context - if ctx is not None: - client._proxy_rc_ref[0] = ( - ctx.request_context, - ctx._fastmcp, - ) - async with client: - # Entering the context already ran connect-time negotiation. - # `initialize()` returns the handshake result on a legacy backend, - # but raises on a modern (server/discover) backend, which has no - # InitializeResult. That mismatch only arises when an explicit - # `mode=` pins the backend to a different era than this legacy - # front (the era-mirroring default keeps the two in lockstep, so - # a legacy front always reaches a legacy backend here). Skip the - # handshake-only call when the backend negotiated the modern era. - if client.protocol_version not in MODERN_PROTOCOL_VERSIONS: - await client.initialize() - # Capture the upstream's instructions while the session is - # live; `initialize_result` clears once the context exits. - init_result = client.initialize_result - if init_result is not None: - upstream_instructions = init_result.instructions - except MCPError: - raise - except _PROXY_TRANSPORT_ERRORS as error: - raise _proxy_upstream_error(error) from error - - result = await call_next(context) - - # Forward the upstream server's instructions unless the proxy defines its - # own. `instructions` is part of the MCP InitializeResult and is meant to - # steer the model, so a proxy that dropped it would silently degrade any - # downstream consumer relying on upstream guidance. - if ( - result is not None - and self.proxy.instructions is None - and upstream_instructions is not None - ): - result.instructions = upstream_instructions - - return result - - # ----------------------------------------------------------------------------- # Proxy Component Classes # ----------------------------------------------------------------------------- @@ -1266,6 +1257,7 @@ class FastMCPProxy(FastMCP): *, client_factory: ClientFactoryT, provider_error_strategy: ProviderErrorStrategy = "warn", + identity: ProxyIdentity = "proxy", **kwargs, ): """Initialize the proxy server. @@ -1280,16 +1272,23 @@ class FastMCPProxy(FastMCP): provider_error_strategy: How provider errors should affect aggregate operations. Defaults to ``"warn"`` for compatibility; use ``"raise"`` when the proxy should surface upstream failures. + identity: Whether negotiation exposes the proxy's server identity or + the upstream server's. Defaults to ``"proxy"`` for compatibility. **kwargs: Additional settings for the FastMCP server. """ + from fastmcp.server.middleware.proxy import ( + ProxyNegotiationMetadataMiddleware, + ) + super().__init__(**kwargs) self.provider_error_strategy = provider_error_strategy self.client_factory = client_factory - provider: Provider = ProxyProvider(client_factory) + provider = ProxyProvider(client_factory) self.add_provider(provider) - self.middleware.append(ProxyInitializeMiddleware(self)) + self.middleware.append( + ProxyNegotiationMetadataMiddleware(provider, identity=identity) + ) self._setup_proxy_ping_handler() - self._setup_proxy_discover_handler() async def _get_client(self) -> Client: client = self.client_factory() @@ -1311,73 +1310,6 @@ class FastMCPProxy(FastMCP): "ping", mcp_types.RequestParams, ping_remote ) - def _setup_proxy_discover_handler(self) -> None: - """Forward the backend's instructions on the modern (`server/discover`) path. - - `ProxyInitializeMiddleware` forwards upstream instructions by patching - the `InitializeResult`, but `on_initialize` only fires for the legacy - handshake. A modern client negotiates via `server/discover`, whose - default SDK handler reads `self.instructions` off the low-level server - directly, so a proxy would silently drop its upstream's instructions for - every modern client. - - The SDK sanctions replacing this handler wholesale, so we delegate to - its own implementation for the rest of the result (supported versions, - capabilities, server info) and only fill in the instructions we would - otherwise lose. Resolving them here β€” at request time, from a live - backend session β€” keeps the proxy's lazy-connect contract intact: the - backend is contacted when a client actually asks, never at construction. - """ - build_default_result = self._mcp_server._handle_discover - - async def discover_remote( - ctx: ServerRequestContext[Any, Any], - params: mcp_types.RequestParams | None, - ) -> mcp_types.DiscoverResult: - result = await build_default_result(ctx, params) - # A proxy with its own instructions keeps them, matching the - # precedence `ProxyInitializeMiddleware` applies on the legacy path. - if result.instructions is not None: - return result - client = await self._get_client() - # `session.instructions` is era-neutral: it reads the backend's - # `DiscoverResult` or `InitializeResult` depending on what the - # backend negotiated, so a modern front can proxy a legacy backend. - if client.is_connected(): - result.instructions = client.session.instructions - return result - # Era mirroring pins a modern backend to an exact version, and a - # pinned version adopts a synthesized `DiscoverResult` instead of - # probing the wire β€” so the pinned client would report no - # instructions at all. Instructions are metadata with no - # back-channel, so this read does not need the era consistency - # mirroring exists to protect; negotiate with "auto" instead, which - # probes `server/discover` and falls back to the handshake for a - # legacy-only backend. - client.mode = "auto" - try: - async with client: - result.instructions = client.session.instructions - except (MCPError, *_PROXY_TRANSPORT_ERRORS) as error: - # Instructions are optional metadata, so an unreachable backend - # must not fail negotiation itself. Failing here would surface - # as a confusing protocol error: the client's auto-negotiation - # reads any `server/discover` error as "not a modern server" - # and retries with the initialize handshake, which this - # modern-serving proxy then rejects β€” hiding the real cause. - # Answer without upstream instructions instead and let the - # backend failure surface on the first real operation, where - # the proxy reports it as an upstream connection error. - logger.debug( - "Could not read upstream instructions for server/discover: %r", - error, - ) - return result - - self._mcp_server.add_request_handler( - "server/discover", mcp_types.RequestParams, discover_remote - ) - # ----------------------------------------------------------------------------- # ProxyClient and Related diff --git a/tests/server/middleware/test_discovery_middleware.py b/tests/server/middleware/test_discovery_middleware.py new file mode 100644 index 000000000..d7af479c7 --- /dev/null +++ b/tests/server/middleware/test_discovery_middleware.py @@ -0,0 +1,31 @@ +"""Tests for typed middleware support during modern discovery.""" + +import mcp_types + +from fastmcp import Client, FastMCP +from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext + + +async def test_on_discover_receives_and_transforms_typed_result(): + class DiscoveryMiddleware(Middleware): + def __init__(self) -> None: + self.request: mcp_types.DiscoverRequest | None = None + self.result: mcp_types.DiscoverResult | None = None + + async def on_discover( + self, + context: MiddlewareContext[mcp_types.DiscoverRequest], + call_next: CallNext[mcp_types.DiscoverRequest, mcp_types.DiscoverResult], + ) -> mcp_types.DiscoverResult: + self.request = context.message + self.result = await call_next(context) + return self.result.model_copy(update={"instructions": "discovered"}) + + middleware = DiscoveryMiddleware() + server = FastMCP("typed-discovery", middleware=[middleware]) + + async with Client(server, mode="auto") as client: + assert client.instructions == "discovered" + + assert isinstance(middleware.request, mcp_types.DiscoverRequest) + assert isinstance(middleware.result, mcp_types.DiscoverResult) diff --git a/tests/server/providers/proxy/test_negotiation_metadata.py b/tests/server/providers/proxy/test_negotiation_metadata.py new file mode 100644 index 000000000..3a3d4a37b --- /dev/null +++ b/tests/server/providers/proxy/test_negotiation_metadata.py @@ -0,0 +1,292 @@ +"""Negotiation metadata forwarding across proxy protocol eras.""" + +from itertools import product +from typing import Any, Literal, TypeVar + +import mcp_types +import pytest +from mcp import MCPError +from mcp_types.version import MODERN_PROTOCOL_VERSIONS + +from fastmcp import Client, FastMCP +from fastmcp.client.transports import StreamableHttpTransport +from fastmcp.server import create_proxy +from fastmcp.server.middleware import ( + CallNext, + Middleware, + MiddlewareContext, + ProxyNegotiationMetadataMiddleware, +) +from fastmcp.server.providers.proxy import ProxyClient, ProxyProvider +from fastmcp.utilities.http import find_available_port + +ResultT = TypeVar("ResultT", bound=mcp_types.Result) + +UPSTREAM_INFO = mcp_types.Implementation( + name="upstream", + title="Upstream title", + version="1.2.3", + description="Upstream description", + website_url="https://upstream.example.com", + icons=[mcp_types.Icon(src="https://upstream.example.com/icon.png")], +) + + +class UpstreamMetadataMiddleware(Middleware): + """Advertise metadata that differs from the gateway's own claims.""" + + def _updates(self, result: mcp_types.Result) -> dict[str, Any]: + meta = { + **(result.meta or {}), + "com.example/upstream": {"enabled": True}, + "com.example/shared": "upstream", + } + updates: dict[str, Any] = { + "instructions": "upstream instructions", + "meta": meta, + "x-upstream-extension": {"enabled": True}, + "x-shared-extension": "upstream", + } + if isinstance(result, mcp_types.InitializeResult): + updates.update( + server_info=UPSTREAM_INFO, + capabilities=mcp_types.ServerCapabilities( + experimental={"upstream": {"claimed": True}} + ), + supportedVersions=["upstream-version"], + ttlMs=91_000, + cacheScope="public", + resultType="upstream-result", + ) + else: + meta[mcp_types.SERVER_INFO_META_KEY] = UPSTREAM_INFO.model_dump( + by_alias=True, mode="json", exclude_none=True + ) + updates.update( + ttl_ms=91_000, + cache_scope="public", + capabilities=mcp_types.ServerCapabilities( + experimental={"upstream": {"claimed": True}} + ), + protocolVersion="upstream-version", + serverInfo=UPSTREAM_INFO.model_dump( + by_alias=True, mode="json", exclude_none=True + ), + ) + return updates + + 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 result is not None + return result.model_copy(update=self._updates(result)) + + async def on_discover( + self, + context: MiddlewareContext[mcp_types.DiscoverRequest], + call_next: CallNext[mcp_types.DiscoverRequest, mcp_types.DiscoverResult], + ) -> mcp_types.DiscoverResult: + result = await call_next(context) + return result.model_copy(update=self._updates(result)) + + +class FrontendMetadataMiddleware(Middleware): + """Set frontend values that must win over the upstream on collision.""" + + def _update(self, result: ResultT) -> ResultT: + return result.model_copy( + update={ + "meta": { + **(result.meta or {}), + "com.example/shared": "frontend", + "com.example/frontend": {"enabled": True}, + }, + "x-shared-extension": "frontend", + "x-frontend-extension": {"enabled": True}, + } + ) + + 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 result is not None + return self._update(result) + + async def on_discover( + self, + context: MiddlewareContext[mcp_types.DiscoverRequest], + call_next: CallNext[mcp_types.DiscoverRequest, mcp_types.DiscoverResult], + ) -> mcp_types.DiscoverResult: + result = await call_next(context) + return self._update(result) + + +def make_upstream() -> FastMCP: + return FastMCP("unmodified-upstream", middleware=[UpstreamMetadataMiddleware()]) + + +def make_gateway( + upstream: FastMCP, + *, + backend_mode: str, + identity: Literal["proxy", "upstream"] = "proxy", + instructions: str | None = None, + frontend_metadata: bool = False, +) -> FastMCP: + provider = ProxyProvider(lambda: ProxyClient(upstream, mode=backend_mode)) + negotiation = ProxyNegotiationMetadataMiddleware(provider, identity=identity) + middleware: list[Middleware] = [negotiation] + if frontend_metadata: + middleware.append(FrontendMetadataMiddleware()) + gateway = FastMCP( + "gateway", + version="9.8.7", + instructions=instructions, + providers=[provider], + middleware=middleware, + cache_ttl=7, + cache_scope="private", + ) + gateway.provider_error_strategy = "raise" + return gateway + + +@pytest.mark.parametrize( + ("frontend_mode", "backend_mode"), + list(product(("legacy", "auto"), repeat=2)), +) +async def test_forwards_metadata_across_all_protocol_era_combinations( + frontend_mode: str, backend_mode: str +): + gateway = make_gateway(make_upstream(), backend_mode=backend_mode) + + # ProxyClient retains unknown extension fields so this test observes the + # complete typed result rather than only its core mcp-types vocabulary. + async with ProxyClient(gateway, mode=frontend_mode) as client: + result = client.session.initialize_result or client.session.discover_result + assert result is not None + assert client.instructions == "upstream instructions" + assert client.server_info is not None + assert client.server_info.name == "gateway" + assert result.meta is not None + assert result.meta["com.example/upstream"] == {"enabled": True} + stamped_info = result.meta.get(mcp_types.SERVER_INFO_META_KEY) + assert result.model_extra == { + "x-upstream-extension": {"enabled": True}, + "x-shared-extension": "upstream", + } + assert result.capabilities.experimental is None + + if isinstance(result, mcp_types.InitializeResult): + assert result.protocol_version not in MODERN_PROTOCOL_VERSIONS + assert stamped_info is None + else: + assert stamped_info is not None + assert stamped_info["name"] == "gateway" + assert result.supported_versions == list(MODERN_PROTOCOL_VERSIONS) + assert result.ttl_ms == 7_000 + assert result.cache_scope == "private" + assert result.result_type == "complete" + + +@pytest.mark.parametrize("frontend_mode", ["legacy", "auto"]) +@pytest.mark.parametrize("identity", ["proxy", "upstream"]) +async def test_identity_policy_forwards_full_implementation( + frontend_mode: str, identity: Literal["proxy", "upstream"] +): + gateway = make_gateway(make_upstream(), backend_mode="auto", identity=identity) + + async with Client(gateway, mode=frontend_mode) as client: + assert client.server_info is not None + if identity == "proxy": + assert client.server_info.name == "gateway" + assert client.server_info.version == "9.8.7" + else: + assert client.server_info == UPSTREAM_INFO + + +@pytest.mark.parametrize("frontend_mode", ["legacy", "auto"]) +async def test_frontend_values_take_precedence(frontend_mode: str): + gateway = make_gateway( + make_upstream(), + backend_mode="auto", + instructions="frontend instructions", + frontend_metadata=True, + ) + + async with ProxyClient(gateway, mode=frontend_mode) as client: + result = client.session.initialize_result or client.session.discover_result + assert result is not None + assert client.instructions == "frontend instructions" + assert result.meta is not None + assert result.meta["com.example/shared"] == "frontend" + assert result.meta["com.example/frontend"] == {"enabled": True} + assert result.meta["com.example/upstream"] == {"enabled": True} + assert result.model_extra == { + "x-upstream-extension": {"enabled": True}, + "x-shared-extension": "frontend", + "x-frontend-extension": {"enabled": True}, + } + + +@pytest.mark.parametrize("frontend_mode", ["legacy", "auto"]) +async def test_unavailable_backend_does_not_fail_negotiation(frontend_mode: str): + port = find_available_port() + provider = ProxyProvider( + lambda: ProxyClient( + StreamableHttpTransport(f"http://127.0.0.1:{port}/mcp"), mode="auto" + ), + cache_ttl=0, + ) + gateway = FastMCP( + "available-gateway", + providers=[provider], + middleware=[ProxyNegotiationMetadataMiddleware(provider)], + ) + gateway.provider_error_strategy = "raise" + + async with Client(gateway, mode=frontend_mode) as client: + assert client.server_info is not None + assert client.server_info.name == "available-gateway" + with pytest.raises(MCPError, match="Client failed to connect"): + await client.list_tools() + + +def test_gateway_construction_does_not_create_backend_client(): + calls = 0 + + def client_factory() -> ProxyClient: + nonlocal calls + calls += 1 + return ProxyClient(make_upstream()) + + provider = ProxyProvider(client_factory) + FastMCP( + "lazy-gateway", + providers=[provider], + middleware=[ProxyNegotiationMetadataMiddleware(provider)], + ) + + assert calls == 0 + + +async def test_fastmcp_proxy_uses_public_negotiation_middleware(): + proxy = create_proxy(make_upstream(), name="convenience", identity="upstream") + + assert any( + isinstance(middleware, ProxyNegotiationMetadataMiddleware) + for middleware in proxy.middleware + ) + async with Client(proxy, mode="auto") as client: + assert client.instructions == "upstream instructions" + assert client.server_info == UPSTREAM_INFO diff --git a/tests/server/providers/proxy/test_proxy_server.py b/tests/server/providers/proxy/test_proxy_server.py index fe489d56f..ca84ab4c5 100644 --- a/tests/server/providers/proxy/test_proxy_server.py +++ b/tests/server/providers/proxy/test_proxy_server.py @@ -204,13 +204,7 @@ async def test_create_proxy_with_transport(fastmcp_server): async def test_proxy_forwards_upstream_instructions(): - """A proxy should surface the upstream server's instructions in the handshake. - - `FastMCPProxy` registers a `server/discover` handler that forwards the - upstream's instructions, mirroring what `ProxyInitializeMiddleware.on_initialize` - already does for the legacy handshake, so `client.session.instructions` - (era-neutral) resolves the same way on both protocol eras. - """ + """The shared negotiation middleware forwards upstream instructions.""" upstream = FastMCP(name="upstream", instructions="USE_THIS_MARKER_123") proxy = create_proxy(upstream, name="proxy") @@ -274,35 +268,25 @@ async def test_proxy_ping_surfaces_wrong_remote_path(): async with run_server_async(remote, transport="http") as url: proxy = create_proxy(StreamableHttpTransport(url.removesuffix("/mcp"))) - # This asserts the error surfaces from merely *connecting* to the proxy, - # with no operation performed. That only happens on the legacy handshake: - # `ProxyInitializeMiddleware.on_initialize` eagerly probes the backend - # during the front's own `initialize` call. A modern front negotiates - # `server/discover` instead, which never runs that middleware hook, so - # connecting succeeds regardless of backend health and the failure would - # only surface on first real use. Pinned because the subject here is - # that eager, handshake-time probe. - # - # SDK v2 surfaces a wrong remote path as an HTTP "Not Found" rather than - # the v1 "Session terminated" message. - with pytest.raises(MCPError, match="Not Found"): - async with Client(proxy, mode="legacy"): - pass + # Optional metadata lookup is best-effort, so negotiation succeeds. The + # first real proxied operation reports the bad backend path instead. + async with Client(proxy, mode="legacy") as client: + with pytest.raises(MCPError, match="Not Found"): + await client.ping() -async def test_proxy_initialize_forwards_remote_connection_error(): +async def test_proxy_initialize_defers_remote_connection_error(): port = find_available_port() proxy = create_proxy( StreamableHttpTransport(f"http://127.0.0.1:{port}/mcp"), provider_error_strategy="raise", ) - # Same reasoning as test_proxy_ping_surfaces_wrong_remote_path above: the - # error surfaces from connecting alone only via the legacy handshake's - # eager backend probe in `ProxyInitializeMiddleware.on_initialize`. - with pytest.raises(MCPError, match="Client failed to connect"): - async with Client(proxy, mode="legacy"): - pass + # Negotiation succeeds without optional backend metadata; the first + # component operation reports the unavailable backend. + async with Client(proxy, mode="legacy") as client: + with pytest.raises(MCPError, match="Client failed to connect"): + await client.list_tools() async def test_proxy_list_tools_surfaces_remote_connection_error(): @@ -324,13 +308,10 @@ async def test_proxy_list_tools_surfaces_remote_connection_error(): async def test_proxy_list_tools_client_surfaces_remote_connection_error(): - """With a modern front, connecting succeeds (no eager backend probe β€” see - test_proxy_ping_surfaces_wrong_remote_path) and the failure only surfaces - once `list_tools()` actually hits the dead backend. `ProxyProvider._list_tools` - now normalizes the raw `httpx2.ConnectError` from the failed backend connect - into the `MCPError("Client failed to connect...")` this test expects, the - same way `ProxyInitializeMiddleware.on_initialize` and `ProxyTool.run` - already did. + """Connecting succeeds and the first component operation reports the backend. + + `ProxyProvider._list_tools` normalizes the raw transport failure into the + `MCPError("Client failed to connect...")` this test expects. """ port = find_available_port() proxy = create_proxy( @@ -1459,13 +1440,7 @@ class TestProxyForwardingAppliesToEveryBackendClient: class TestProxyModernEraInstructions: - """Upstream instructions must reach a client on the modern era too. - - `ProxyInitializeMiddleware.on_initialize` only fires for the legacy - handshake. A `mode="auto"` client negotiates via `server/discover`, which - the SDK builds from the low-level server's own `instructions`, so without a - discover-side hook the proxy drops its upstream's instructions entirely. - """ + """Upstream instructions must reach a client on the modern era too.""" async def test_proxy_forwards_upstream_instructions_on_modern_era(self): upstream = FastMCP(name="upstream", instructions="USE_THIS_MARKER_123") @@ -1491,13 +1466,7 @@ class TestProxyModernEraInstructions: class TestProxyProviderTransportErrors: - """A dead backend must surface as an MCPError, not a raw transport error. - - `ProxyTool.run` and `ProxyInitializeMiddleware.on_initialize` normalize - connection failures into `MCPError`; the provider's list methods caught - only `MCPError`, so an `httpx2.ConnectError` (or the `RuntimeError` the - client wraps a failed connect in) escaped unwrapped to the caller. - """ + """A dead backend must surface as an MCPError, not a raw transport error.""" @pytest.fixture def unreachable_provider(self) -> ProxyProvider: