Limit forwarded proxy metadata

🤖 Generated with OpenAI Codex
This commit is contained in:
Jake Kaplan 2026-08-05 20:18:03 -04:00
commit 698870343a
6 changed files with 27 additions and 169 deletions

View file

@ -416,7 +416,7 @@ gateway = FastMCP(
)
```
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.
The middleware works across both the legacy `initialize` handshake and modern `server/discover`. It forwards upstream instructions, namespaced `_meta`, 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. Unknown top-level fields are not forwarded because the gateway cannot safely interpret their claims.
`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.

View file

@ -1,12 +0,0 @@
"""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")

View file

@ -29,10 +29,6 @@ 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
@ -346,11 +342,9 @@ class FastMCPServerMiddleware:
) -> mcp_types.DiscoverResult:
raw = await call_next(ctx)
if isinstance(raw, mcp_types.DiscoverResult):
return _ExtensibleDiscoverResult.model_validate(
raw.model_dump(by_alias=True)
)
return raw
if isinstance(raw, Mapping):
return _ExtensibleDiscoverResult.model_validate(dict(raw))
return mcp_types.DiscoverResult.model_validate(dict(raw))
raise TypeError(
"server/discover handler returned "
f"{type(raw).__name__}; expected DiscoverResult or mapping"
@ -408,11 +402,9 @@ class FastMCPServerMiddleware:
nonlocal captured_result, call_next_completed
raw = await call_next(ctx)
if isinstance(raw, mcp_types.InitializeResult):
captured_result = _ExtensibleInitializeResult.model_validate(
raw.model_dump(by_alias=True)
)
captured_result = raw
elif isinstance(raw, Mapping):
captured_result = _ExtensibleInitializeResult.model_validate(dict(raw))
captured_result = mcp_types.InitializeResult.model_validate(dict(raw))
call_next_completed = True
return captured_result if raw is not None else None

View file

@ -4,17 +4,13 @@ from __future__ import annotations
import inspect
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
from typing import TYPE_CHECKING, Any, Literal, 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
@ -25,41 +21,12 @@ 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:
@ -70,11 +37,6 @@ class _NegotiationMetadata:
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
},
)
@ -82,10 +44,9 @@ 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.
result type. Instructions and namespaced metadata 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__(
@ -154,33 +115,23 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
)
return merged
def _merge_result(
def _updates(
self,
result: NegotiationResultT,
result: mcp_types.InitializeResult | mcp_types.DiscoverResult,
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
) -> dict[str, Any]:
updates: dict[str, Any] = {
"meta": self._merge_meta(dict(result.meta or {}), upstream) or None
}
updates["meta"] = self._merge_meta(dict(forwarded.meta or {}), upstream) or None
if forwarded.instructions is None and upstream.instructions is not None:
if result.instructions is None and upstream.instructions is not None:
updates["instructions"] = upstream.instructions
if (
isinstance(forwarded, mcp_types.InitializeResult)
isinstance(result, 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))
return updates
async def on_initialize(
self,
@ -195,7 +146,7 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
upstream = await self._read_upstream(context)
if upstream is None:
return result
return self._merge_result(result, upstream)
return result.model_copy(update=self._updates(result, upstream))
async def on_discover(
self,
@ -206,4 +157,4 @@ class ProxyNegotiationMetadataMiddleware(Middleware):
upstream = await self._read_upstream(context)
if upstream is None:
return result
return self._merge_result(result, upstream)
return result.model_copy(update=self._updates(result, upstream))

View file

@ -50,10 +50,6 @@ 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.providers.aggregate import ProviderErrorStrategy
@ -69,11 +65,6 @@ 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
@ -84,56 +75,15 @@ ClientFactoryT = Callable[[], Client] | Callable[[], Awaitable[Client]]
class _ForwardingClientSession(ClientSession):
"""A proxy session that relays backend results without consuming them.
"""A session that does not enforce the backend's declared output schema.
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.
`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.
"""
_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:

View file

@ -44,8 +44,6 @@ class UpstreamMetadataMiddleware(Middleware):
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(
@ -53,10 +51,6 @@ class UpstreamMetadataMiddleware(Middleware):
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(
@ -68,10 +62,6 @@ class UpstreamMetadataMiddleware(Middleware):
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
@ -106,8 +96,6 @@ class FrontendMetadataMiddleware(Middleware):
"com.example/shared": "frontend",
"com.example/frontend": {"enabled": True},
},
"x-shared-extension": "frontend",
"x-frontend-extension": {"enabled": True},
}
)
@ -170,9 +158,7 @@ async def test_forwards_metadata_across_all_protocol_era_combinations(
):
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:
async with Client(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"
@ -181,10 +167,6 @@ async def test_forwards_metadata_across_all_protocol_era_combinations(
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):
@ -224,7 +206,7 @@ async def test_frontend_values_take_precedence(frontend_mode: str):
frontend_metadata=True,
)
async with ProxyClient(gateway, mode=frontend_mode) as client:
async with Client(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"
@ -232,11 +214,6 @@ async def test_frontend_values_take_precedence(frontend_mode: str):
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"])