-
Notifications
You must be signed in to change notification settings - Fork 2k
Forward proxy initialize as bridge behavior #4228
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -31,9 +31,10 @@ | |
| from pydantic.networks import AnyUrl | ||
|
|
||
| from fastmcp.client.client import Client, FastMCP1Server | ||
| from fastmcp.client.elicitation import ElicitResult | ||
| from fastmcp.client.logging import LogMessage | ||
| from fastmcp.client.roots import RootsList | ||
| from fastmcp.client.elicitation import ElicitResult, create_elicitation_callback | ||
| from fastmcp.client.logging import LogMessage, create_log_callback | ||
| from fastmcp.client.roots import RootsList, create_roots_callback | ||
| from fastmcp.client.sampling import create_sampling_callback | ||
| from fastmcp.client.telemetry import client_span | ||
| from fastmcp.client.transports import ClientTransportT | ||
| from fastmcp.exceptions import ResourceError | ||
|
|
@@ -88,8 +89,15 @@ async def on_initialize( | |
| ) -> mcp.types.InitializeResult | None: | ||
| client = await self.proxy._get_client() | ||
| try: | ||
| if isinstance(client, StatefulProxyClient): | ||
| ctx = context.fastmcp_context | ||
| if ctx is not None: | ||
| client._proxy_rc_ref[0] = ( | ||
| ctx.request_context, | ||
| ctx._fastmcp, | ||
| ) | ||
| async with client: | ||
| pass | ||
| await client.initialize() | ||
| except McpError: | ||
| raise | ||
| except ( | ||
|
|
@@ -881,7 +889,6 @@ def __init__( | |
| *, | ||
| client_factory: ClientFactoryT, | ||
| provider_error_strategy: ProviderErrorStrategy = "warn", | ||
| validate_on_initialize: bool = False, | ||
| **kwargs, | ||
| ): | ||
| """Initialize the proxy server. | ||
|
|
@@ -896,17 +903,14 @@ def __init__( | |
| provider_error_strategy: How provider errors should affect aggregate | ||
| operations. Defaults to ``"warn"`` for compatibility; use | ||
| ``"raise"`` when the proxy should surface upstream failures. | ||
| validate_on_initialize: If true, connect to the upstream server during | ||
| the incoming MCP initialize request. | ||
| **kwargs: Additional settings for the FastMCP server. | ||
| """ | ||
| super().__init__(**kwargs) | ||
| self.provider_error_strategy = provider_error_strategy | ||
| self.client_factory = client_factory | ||
| provider: Provider = ProxyProvider(client_factory) | ||
| self.add_provider(provider) | ||
| if validate_on_initialize: | ||
| self.middleware.append(ProxyInitializeMiddleware(self)) | ||
| self.middleware.append(ProxyInitializeMiddleware(self)) | ||
| self._setup_proxy_ping_handler() | ||
|
|
||
| async def _get_client(self) -> Client: | ||
|
|
@@ -1140,11 +1144,13 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]): | |
| # would resolve stale values in the receive loop. The restore helper | ||
| # constructs a fresh Context from the weakref after setting request_ctx. | ||
| _proxy_rc_ref: list[Any] | ||
| _proxy_restoring_handler_keys: set[str] | ||
|
|
||
| def __init__(self, *args: Any, **kwargs: Any): | ||
| # Install context-restoring handler wrappers BEFORE super().__init__ | ||
| # registers them with the Client's session kwargs. | ||
| self._proxy_rc_ref = [None] | ||
| self._proxy_restoring_handler_keys = set() | ||
| for key, default_fn in ( | ||
| ("roots", default_proxy_roots_handler), | ||
| ("sampling_handler", default_proxy_sampling_handler), | ||
|
|
@@ -1154,10 +1160,43 @@ def __init__(self, *args: Any, **kwargs: Any): | |
| ): | ||
| if key not in kwargs: | ||
| kwargs[key] = _make_restoring_handler(default_fn, self._proxy_rc_ref) | ||
| self._proxy_restoring_handler_keys.add(key) | ||
|
|
||
| super().__init__(*args, **kwargs) | ||
| self._caches: dict[ServerSession, Client[ClientTransportT]] = {} | ||
|
|
||
| def _bind_restoring_handlers(self) -> None: | ||
| if "roots" in self._proxy_restoring_handler_keys: | ||
| self._session_kwargs["list_roots_callback"] = create_roots_callback( | ||
| _make_restoring_handler(default_proxy_roots_handler, self._proxy_rc_ref) | ||
| ) | ||
|
Comment on lines
+1169
to
+1172
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Useful? React with 👍 / 👎. |
||
| if "sampling_handler" in self._proxy_restoring_handler_keys: | ||
| self._session_kwargs["sampling_callback"] = create_sampling_callback( | ||
| _make_restoring_handler( | ||
| default_proxy_sampling_handler, self._proxy_rc_ref | ||
| ) | ||
| ) | ||
| if "elicitation_handler" in self._proxy_restoring_handler_keys: | ||
| self._session_kwargs["elicitation_callback"] = create_elicitation_callback( | ||
| _make_restoring_handler( | ||
| default_proxy_elicitation_handler, self._proxy_rc_ref | ||
| ) | ||
| ) | ||
| if "log_handler" in self._proxy_restoring_handler_keys: | ||
| self._session_kwargs["logging_callback"] = create_log_callback( | ||
| _make_restoring_handler(default_proxy_log_handler, self._proxy_rc_ref) | ||
| ) | ||
| if "progress_handler" in self._proxy_restoring_handler_keys: | ||
| self._progress_handler = _make_restoring_handler( | ||
| default_proxy_progress_handler, self._proxy_rc_ref | ||
| ) | ||
|
|
||
| def new(self) -> StatefulProxyClient[ClientTransportT]: | ||
| new_client = cast(StatefulProxyClient[ClientTransportT], super().new()) | ||
| new_client._proxy_rc_ref = [None] | ||
| new_client._bind_restoring_handlers() | ||
| return new_client | ||
|
|
||
| async def __aexit__(self, exc_type, exc_value, traceback) -> None: # type: ignore[override] # ty:ignore[invalid-method-override] | ||
| """The stateful proxy client will be forced disconnected when the session is exited. | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This change makes
FastMCPProxyalways installProxyInitializeMiddleware, so every downstreaminitializenow forces an upstream connect/initialize handshake with no way to disable it. That regresses the prior lazy behavior for directFastMCPProxy(...)usage (and any caller that previously usedvalidate_on_initialize=False), causing startup to fail immediately when upstream is temporarily unavailable instead of deferring failure to the first proxied operation; this is especially disruptive for stateful/customclient_factorymounting flows that rely on lazy or degraded startup.Useful? React with 👍 / 👎.