Stop proxies from validating backend results or mutating shared transports - #4552
Conversation
A proxy relays tool results; the end client validates them. Enforcing the backend's declared output schema mid-path turned a backend schema bug into a proxy error and forced an extra tools/list on every proxied call.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8207c4543d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # A proxy relays results; it does not police the backend's output | ||
| # schema. See ForwardingClientSession. | ||
| self.transport.session_class = ForwardingClientSession |
There was a problem hiding this comment.
Apply the forwarding session to plain Client proxy targets
When the supported create_proxy(Client(...)) form is used, _create_client_factory returns the original client or client.new() (proxy.py:853-889) rather than constructing a ProxyClient, so this assignment never runs. A backend result that violates its schema is therefore still converted into a proxy error instead of being relayed; this also affects internal callers such as transforming MCP configs, which explicitly construct a plain Client before calling create_proxy. Configure the forwarding session for every proxy client-factory path, not only ProxyClient.__init__.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct, and thanks — create_proxy(Client(...)) never routes through ProxyClient.__init__, so plain-Client targets kept validating.
Fixed: the settings are no longer attached in the constructor at all, so _create_client_factory applies them to the copy it builds for every backend path. Verified a plain disconnected Client target now relays {"status": "weird"} against a schema declaring ["ok", "error"], with tests for the plain-Client, single-server config, and multi-server config forms.
| # The ClientSession class this transport instantiates. Overridable per | ||
| # instance for callers that need modified session behavior — e.g. proxies | ||
| # set a session that skips output-schema validation (see ProxyClient). | ||
| session_class: type[ClientSession] = ClientSession |
There was a problem hiding this comment.
Propagate the custom session class through MCPConfigTransport
For the supported ProxyClient(MCPConfig(...)) and dict inputs, this class attribute is set only on the outer MCPConfigTransport. Its single-server path delegates directly to self.transport.connect_session() (config.py:94-97) without copying session_class, and its multi-server path creates a new default FastMCPTransport (config.py:140-142), so both paths can still validate and reject the backend result inside the proxy. The composite transport must propagate the requested session behavior to the session it actually creates.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct for the commit you reviewed. Fixed, and MCPConfigTransport now forwards the setting explicitly on both its single-server delegate and its composite FastMCPTransport, so it's a typed argument rather than something the composite has to remember not to drop.
Verified against a backend violating its own output schema through both a single-server and a two-server config; both relay now, with tests for each.
ProxyClient configured its transport in place, so a transport shared with another client leaked proxy behavior into it — including header forwarding, which would send the caller's credentials upstream from a client that never opted in. create_proxy() did the same to a user-supplied Client. Both settings now travel as a frozen TransportOptions in the client's own session kwargs.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: acb6918605
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # Reusing the caller's live session, so its connection settings are | ||
| # already fixed and proxy options have nothing left to configure. |
There was a problem hiding this comment.
Reconnect ordinary clients with the forwarding session
Fresh evidence after the earlier plain-Client finding: this new connected-client branch explicitly leaves the live session unchanged. When create_proxy(client) is called inside async with client, that session is still a regular ClientSession, so a schema-violating backend result is converted into a proxy error instead of being relayed. This is a supported target path; it needs a forwarding-capable session rather than reusing the ordinary session as-is.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Accurate, and deliberate. That branch reuses the caller's already-live session, and connection settings only apply at connect time, so there is nothing left to configure — swapping the session would mean discarding the session the caller explicitly asked to reuse.
The log line for that path now says results stay validated and points at passing a disconnected client instead, which gets the forwarding behavior. Changing reuse itself is a separate design question from this fix.
|
|
||
| if isinstance(self.transport, StreamableHttpTransport | SSETransport): | ||
| self.transport.forward_incoming_headers = True | ||
| self._session_kwargs["transport_options"] = PROXY_TRANSPORT_OPTIONS |
There was a problem hiding this comment.
Avoid injecting transport options into custom session kwargs
When ProxyClient wraps a third-party ClientTransport, this newly injected key is passed to its public connect_session(**session_kwargs) implementation. Existing transports commonly follow the documented contract by forwarding those kwargs directly into ClientSession; because only the built-in transports know to call pop_transport_options, they now fail initialization with TypeError: ClientSession.__init__() got an unexpected keyword argument 'transport_options'. Transport-only state must be handled centrally or without changing the kwargs contract for existing transport extensions.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This was a real regression and the most useful comment of the batch — thank you. Confirmed a transport following the documented contract died with ClientSession.__init__() got an unexpected keyword argument 'transport_options', and it hit ordinary clients, not just proxies.
SessionKwargs is back to being exactly the ClientSession constructor's parameters. The settings are now a typed keyword-only argument on connect_session, and the client only sends it when it wants non-default behavior, so an ordinary client never passes an argument a custom transport might not accept. Verified a third-party transport that splats its kwargs works again.
| await stack.enter_async_context( | ||
| ClientSession(read_stream, write_stream, **session_kwargs) | ||
| options.session_class( | ||
| read_stream, write_stream, **client_session_kwargs | ||
| ) |
There was a problem hiding this comment.
Rebuild cached stdio sessions when client options change
With the default keep_alive=True, the session class selected here remains cached on StdioTransport; later connect_session calls return early when _connect_task already exists. If a shared stdio transport is first used by a ProxyClient and then by an ordinary Client, the ordinary client therefore reuses _ForwardingClientSession and silently stops validating tool output schemas; the reverse order also leaves the proxy validating results. This defeats the new client-scoped isolation specifically for the built-in keep-alive stdio transport, so cached sessions must be recreated or separated when their transport options differ.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed. StdioTransport now records the options its live session was built with and tears the session down instead of reusing it when a connect asks for different ones, so a kept-alive subprocess can't hand a proxy's non-validating session to an ordinary client (or vice versa). Tests cover both the reuse and reconnect cases.
create_proxy() accepts plain Clients and MCPConfigs, which don't route through ProxyClient.__init__, so those forms still rejected backend results. Settings now travel in session kwargs, which MCPConfigTransport already propagates to the transports it delegates to.
Carrying them inside session kwargs broke transports that forward those kwargs straight to ClientSession, which is what the parameter is documented to contain. They are now a typed connect_session argument, sent only when a client wants non-default settings, so ordinary clients are unaffected. Composite transports forward it explicitly, and a kept-alive stdio session is rebuilt rather than reused when a client asks for different options.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d689455b5d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if self._connect_task is not None and self._session_options != options: | ||
| await self.disconnect() |
There was a problem hiding this comment.
Preserve active stdio sessions when options differ
When two clients concurrently share a keep-alive StdioTransport with different options—for example, an ordinary client is connected while a proxy created from the same transport starts a request—this branch disconnects the subprocess and session still in use by the first client. Its subsequent or in-flight calls then fail because the dispatcher has been closed. Reconnecting is safe only after the cached session is idle; differing client options need isolated sessions/transports rather than tearing down an active connection.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
You're right, and this was strictly worse than what it replaced — a live connection torn down instead of a setting silently leaked.
The transport now tracks how many sessions are open on it. An idle cached session still gets rebuilt when a client asks for different options, but one that's in use raises instead, pointing at giving each client its own transport. Swapping the session under an active client isn't something this transport can do safely, so refusing beats either corrupting it or breaking it.
Tests cover all three: matching options reuse, differing options on an idle session rebuild, differing options on an in-use session raise and leave the connection intact.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48335e3d9e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| if isinstance(self.transport, StreamableHttpTransport | SSETransport): | ||
| self.transport.forward_incoming_headers = True | ||
| self._transport_options = PROXY_TRANSPORT_OPTIONS |
There was a problem hiding this comment.
Apply proxy options to disconnected plain factory clients
Because proxy options are assigned only by ProxyClient and the create_proxy(target) conversion path, the public FastMCPProxy(client_factory=...) and ProxyProvider(client_factory) APIs still use a normal session when their factory returns a disconnected plain Client—a supported return type also covered by test_proxy_with_async_client_factory. In that scenario, a schema-violating backend result is still replaced with a proxy error rather than relayed, and HTTP authorization forwarding is also skipped; configure disconnected plain clients returned by these direct factories before they connect.
Useful? React with 👍 / 👎.
Rebuilding on differing options is only safe while the cached session is idle; doing it under an active client closed the subprocess out from under them. Idle sessions still rebuild, in-use ones now raise instead.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d5f7bfd34a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if self._connect_task is not None and self._session_options != options: | ||
| if self._active_sessions: |
There was a problem hiding this comment.
Serialize stdio option-dependent reconnects
When an idle keep-alive session exists and two clients with different transport options connect concurrently, both callers can observe _active_sessions == 0 and enter disconnect(). Their continuations can then reset and overwrite _connect_task while the other caller is starting its replacement, launching multiple subprocesses and leaving one session untracked and uncloseable. Protect the option check and rebuild sequence with a lock so only one caller can replace the cached session.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed, and it was already breaking CI — tests/client/test_stdio.py::TestParallelCalls::test_parallel_calls was timing out on my last push for exactly this reason.
The race was slightly worse than described. Client.new() shares the transport instance, so a proxy over a stdio client fans ten concurrent connects onto one keep-alive transport. The options a connect was started with were only recorded after it completed, so while the first was in flight the recorded value was still unset — which every other caller read as a mismatch, so they each tore down the connection being established.
Options are now recorded before the connect task starts, so an in-flight session already belongs to them, and the whole check-and-rebuild runs under a lock as you suggested. Full suite passes with no deselections.
| class MyTransport(ClientTransport): | ||
| async def connect_session(self, *, transport_options=None, **session_kwargs): |
There was a problem hiding this comment.
Decorate the custom transport example as a context manager
A user copying this breaking-change migration example gets an async generator, not the async context manager that Client expects, because the override lacks @contextlib.asynccontextmanager. Entering the client then fails with an error such as 'async_generator' object does not support the asynchronous context manager protocol; include the decorator and its import in the example.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct, thanks — the example would have handed back an async generator. Added @contextlib.asynccontextmanager and the import.
The options a connect was started with were only recorded once it finished, so concurrent callers read the unset value as a mismatch and tore down the connection being established. They are now recorded up front and the check-and-rebuild is serialized.
A proxy relays tool results; it does not consume them. Two behaviors contradicted that.
It enforced the backend's output schema.
ClientSession.call_toolvalidates structured content against the schema the backend advertised and raises when they disagree. That check belongs to whoever consumes the result, and the end client already runs it. So when a backend ships a schema that doesn't quite match what it returns (an enum missing a variant, say), the proxy manufactured its own error and the real response never reached the client that could decide what to do with it.It also cost a round trip per call. That validation refreshes its schema cache with
tools/listwhen it doesn't recognize a tool, and because a proxy builds a fresh client per request the cache was always cold. Five proxied calls cost seven backendtools/listrequests. Now they cost two, from component discovery alone.It reconfigured transports it didn't own.
ProxyClientmutated its transport in place to enable header forwarding, so a transport shared with another client leaked that setting into it.create_proxy(my_client)did the same to a caller'sClient, meaning their own later connections would forward inbound authorization headers upstream without ever opting in.Connection settings now travel with the client that wants them, as a typed argument rather than transport state:
Clientpasses it toconnect_sessiononly when it wants non-default behavior, so an ordinary client is unaffected and transports that don't accept the argument keep working. A custom transport used as a proxy backend must accept and honor it, and a transport that wraps others must pass it along —MCPConfigTransportforwards it to both its single-server delegate and its composite server. This removesforward_incoming_headersas a settable transport attribute; it was only ever set by proxy internals. Documented in the v4 change register.Ordinary clients are unaffected: they still validate, and a schema-violating result still raises for whoever actually consumes it.