mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 07:09:11 +02:00
Stop proxies from validating backend results or mutating shared transports (#4552)
This commit is contained in:
parent
b9b1deacb6
commit
1e529ee27d
12 changed files with 624 additions and 80 deletions
|
|
@ -205,6 +205,31 @@ client = Client("my_mcp_server.py", timeout=30.0) # also works
|
|||
|
||||
*Verify:* `fastmcp_slim/fastmcp/client/transports/base.py` (`SessionKwargs.read_timeout_seconds: float | None`), `client/client.py`.
|
||||
|
||||
### Connection settings passed to `connect_session` — Breaking (custom transports)
|
||||
|
||||
`ClientTransport.connect_session` takes a new keyword-only `transport_options: TransportOptions | None`, describing how the connecting client wants its session built: which `ClientSession` class to instantiate, and whether to forward the caller's authorization header upstream. Proxies use it to relay backend results without enforcing their output schema (see [Proxy Servers](/servers/providers/proxy#tool-results-are-relayed-not-inspected)).
|
||||
|
||||
These settings previously lived on the transport instance, so a transport shared between clients leaked one client's configuration into another — including credential forwarding, which `create_proxy(some_client)` would silently enable on the caller's own client. They now travel with the client that wants them, and `forward_incoming_headers` is no longer a settable transport attribute.
|
||||
|
||||
A client only passes the argument when it wants non-default settings, so an ordinary `Client` is unaffected and transports that don't accept it keep working. A custom `ClientTransport` used as a *proxy backend* must accept and honor it:
|
||||
|
||||
```python
|
||||
import contextlib
|
||||
|
||||
from fastmcp.client.transports.base import ClientTransport, TransportOptions
|
||||
|
||||
class MyTransport(ClientTransport):
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(self, *, transport_options=None, **session_kwargs):
|
||||
options = transport_options or TransportOptions()
|
||||
async with options.session_class(read, write, **session_kwargs) as session:
|
||||
yield session
|
||||
```
|
||||
|
||||
A transport that wraps others must pass it along; `MCPConfigTransport` forwards it to both its single-server delegate and its composite server.
|
||||
|
||||
*Verify:* `fastmcp_slim/fastmcp/client/transports/base.py` (`TransportOptions`), the four built-in transports, `transports/config.py`, and `tests/server/providers/proxy/test_proxy_server.py`.
|
||||
|
||||
### `get_session_id` via header sniff — Bridged
|
||||
|
||||
The SDK dropped `get_session_id` from the streamable-HTTP transport with no replacement (the SDK source has an author TODO acknowledging it breaks the Transport protocol). FastMCP reconstructs it by registering an httpx2 response event hook on the client it owns, capturing the `mcp-session-id` response header (httpx2 preserves httpx's `event_hooks` API). The removal trigger is the upstream TODO.
|
||||
|
|
|
|||
|
|
@ -170,6 +170,26 @@ backend = ProxyClient(
|
|||
)
|
||||
```
|
||||
|
||||
### Tool Results Are Relayed, Not Inspected
|
||||
|
||||
A proxy passes a backend's tool results through untouched, including results that don't match the output schema the backend advertised. Deciding whether a server honored its own contract belongs to the client consuming the result, and that client validates for itself.
|
||||
|
||||
This matters when a backend's declared schema is subtly wrong — an enum missing a variant it actually returns, say. A proxy that enforced the schema would replace the backend's working response with an error of its own, and the client would never see what the backend actually said.
|
||||
|
||||
```python
|
||||
from fastmcp import Client
|
||||
from fastmcp.server import create_proxy
|
||||
|
||||
proxy = create_proxy("backend_server.py")
|
||||
|
||||
async with Client(proxy) as client:
|
||||
# The backend's response arrives as the backend sent it. If it violates
|
||||
# the backend's own output schema, this client raises — its decision.
|
||||
result = await client.call_tool("get_status")
|
||||
```
|
||||
|
||||
Skipping the check also avoids a `tools/list` round trip to the backend on every proxied call, since validation would need the backend's schemas and a proxy builds a fresh connection per request.
|
||||
|
||||
## Configuration-Based Proxies
|
||||
|
||||
<VersionBadge version="2.4.0" />
|
||||
|
|
|
|||
|
|
@ -98,6 +98,7 @@ from .transports import (
|
|||
StreamableHttpTransport,
|
||||
infer_transport,
|
||||
)
|
||||
from .transports.base import TransportOptions
|
||||
|
||||
__all__ = [
|
||||
"Client",
|
||||
|
|
@ -490,6 +491,7 @@ class Client(
|
|||
|
||||
# Session context management - see class docstring for detailed explanation
|
||||
self._session_state = ClientSessionState()
|
||||
self._transport_options: TransportOptions | None = None
|
||||
|
||||
# Track task IDs submitted by this client (for list_tasks support)
|
||||
self._submitted_task_ids: set[str] = set()
|
||||
|
|
@ -656,6 +658,7 @@ class Client(
|
|||
# Always reset session state so cloned clients start disconnected and do not
|
||||
# share lifecycle state with the original instance.
|
||||
new_client._session_state = ClientSessionState()
|
||||
new_client._transport_options = self._transport_options
|
||||
|
||||
# Reset mutable task tracking state so new client is independent
|
||||
new_client._task_registry = {}
|
||||
|
|
@ -695,10 +698,17 @@ class Client(
|
|||
|
||||
@asynccontextmanager
|
||||
async def _context_manager(self):
|
||||
# Only passed when this client actually wants non-default settings, so an
|
||||
# ordinary client never sends an argument a transport might not accept.
|
||||
if self._transport_options is not None:
|
||||
connection = self.transport.connect_session(
|
||||
transport_options=self._transport_options, **self._session_kwargs
|
||||
)
|
||||
else:
|
||||
connection = self.transport.connect_session(**self._session_kwargs)
|
||||
|
||||
with catch(get_catch_handlers()):
|
||||
async with self.transport.connect_session(
|
||||
**self._session_kwargs
|
||||
) as session:
|
||||
async with connection as session:
|
||||
self._session_state.session = session
|
||||
# Initialize the session if auto_initialize is enabled
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import abc
|
||||
import contextlib
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, TypeVar
|
||||
|
||||
import httpx2
|
||||
|
|
@ -20,7 +21,7 @@ from typing_extensions import TypedDict, Unpack
|
|||
ClientTransportT = TypeVar("ClientTransportT", bound="ClientTransport")
|
||||
|
||||
|
||||
class SessionKwargs(TypedDict, total=False):
|
||||
class ClientSessionKwargs(TypedDict, total=False):
|
||||
"""Keyword arguments for the MCP ClientSession constructor."""
|
||||
|
||||
read_timeout_seconds: float | None
|
||||
|
|
@ -34,6 +35,32 @@ class SessionKwargs(TypedDict, total=False):
|
|||
notification_bindings: Sequence[NotificationBinding[Any]] | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TransportOptions:
|
||||
"""How one client wants its connection built.
|
||||
|
||||
These belong to the client rather than to the transport, so a transport
|
||||
shared between clients doesn't leak one client's settings to another.
|
||||
|
||||
Attributes:
|
||||
session_class: The ClientSession class to instantiate. Proxies supply a
|
||||
session that skips output-schema validation, since they relay
|
||||
results rather than consume them.
|
||||
forward_incoming_headers: Whether to forward the inbound request's
|
||||
authorization header upstream. Only appropriate for proxies, where
|
||||
the caller's credentials are meant to be propagated. Honored by the
|
||||
HTTP and SSE transports; ignored by the others.
|
||||
"""
|
||||
|
||||
session_class: type[ClientSession] = ClientSession
|
||||
forward_incoming_headers: bool = False
|
||||
|
||||
|
||||
# SessionKwargs stays exactly the ClientSession constructor's parameters, so a
|
||||
# transport can splat it into ClientSession without filtering.
|
||||
SessionKwargs = ClientSessionKwargs
|
||||
|
||||
|
||||
class ClientTransport(abc.ABC):
|
||||
"""
|
||||
Abstract base class for different MCP client transport mechanisms.
|
||||
|
|
@ -46,7 +73,10 @@ class ClientTransport(abc.ABC):
|
|||
@abc.abstractmethod
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
"""
|
||||
Establishes a connection and yields an active ClientSession.
|
||||
|
|
@ -58,6 +88,9 @@ class ClientTransport(abc.ABC):
|
|||
within this context.
|
||||
|
||||
Args:
|
||||
transport_options: How the connecting client wants this connection
|
||||
built. Defaults apply when omitted. A transport
|
||||
that wraps others must pass this along.
|
||||
**session_kwargs: Keyword arguments to pass to the ClientSession
|
||||
constructor (e.g., callbacks, timeouts).
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,11 @@ from mcp import ClientSession
|
|||
from typing_extensions import Unpack
|
||||
|
||||
from fastmcp import _install_hints
|
||||
from fastmcp.client.transports.base import ClientTransport, SessionKwargs
|
||||
from fastmcp.client.transports.base import (
|
||||
ClientTransport,
|
||||
SessionKwargs,
|
||||
TransportOptions,
|
||||
)
|
||||
from fastmcp.client.transports.memory import FastMCPTransport
|
||||
from fastmcp.mcp_config import (
|
||||
MCPConfig,
|
||||
|
|
@ -89,11 +93,16 @@ class MCPConfigTransport(ClientTransport):
|
|||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
# Single server - delegate directly to pre-created transport
|
||||
if len(self.config.mcpServers) == 1:
|
||||
async with self.transport.connect_session(**session_kwargs) as session:
|
||||
async with self.transport.connect_session(
|
||||
transport_options=transport_options, **session_kwargs
|
||||
) as session:
|
||||
yield session
|
||||
return
|
||||
|
||||
|
|
@ -138,7 +147,7 @@ class MCPConfigTransport(ClientTransport):
|
|||
raise ConnectionError("All MCP servers failed to connect")
|
||||
|
||||
async with FastMCPTransport(mcp=composite).connect_session(
|
||||
**session_kwargs
|
||||
transport_options=transport_options, **session_kwargs
|
||||
) as session:
|
||||
yield session
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,11 @@ from typing_extensions import Unpack
|
|||
from fastmcp.client.auth.bearer import BearerAuth
|
||||
from fastmcp.client.auth.oauth import OAuth
|
||||
from fastmcp.client.dependencies import get_http_headers
|
||||
from fastmcp.client.transports.base import ClientTransport, SessionKwargs
|
||||
from fastmcp.client.transports.base import (
|
||||
ClientTransport,
|
||||
SessionKwargs,
|
||||
TransportOptions,
|
||||
)
|
||||
|
||||
|
||||
class StreamableHttpTransport(ClientTransport):
|
||||
|
|
@ -74,8 +78,6 @@ class StreamableHttpTransport(ClientTransport):
|
|||
|
||||
self._set_auth(auth)
|
||||
|
||||
self.forward_incoming_headers: bool = False
|
||||
|
||||
# SDK v2's streamable_http_client no longer exposes a get_session_id
|
||||
# callback. We recover the session id ourselves by capturing the
|
||||
# `mcp-session-id` response header via an httpx event hook on the
|
||||
|
|
@ -143,13 +145,18 @@ class StreamableHttpTransport(ClientTransport):
|
|||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
options = transport_options or TransportOptions()
|
||||
|
||||
# When used in a proxy, forward the inbound request's authorization
|
||||
# header to the upstream server. This is off by default so that a
|
||||
# plain Client used inside a server tool handler doesn't accidentally
|
||||
# leak the caller's credentials to an unrelated remote server.
|
||||
if self.forward_incoming_headers:
|
||||
if options.forward_incoming_headers:
|
||||
headers = get_http_headers(include={"authorization"}) | self.headers
|
||||
else:
|
||||
headers = dict(self.headers)
|
||||
|
|
@ -202,7 +209,9 @@ class StreamableHttpTransport(ClientTransport):
|
|||
read_stream,
|
||||
write_stream,
|
||||
),
|
||||
ClientSession(read_stream, write_stream, **session_kwargs) as session,
|
||||
options.session_class(
|
||||
read_stream, write_stream, **session_kwargs
|
||||
) as session,
|
||||
):
|
||||
yield session
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,11 @@ from mcp.shared.memory import create_client_server_memory_streams
|
|||
from typing_extensions import Unpack
|
||||
|
||||
from fastmcp import _install_hints
|
||||
from fastmcp.client.transports.base import ClientTransport, SessionKwargs
|
||||
from fastmcp.client.transports.base import (
|
||||
ClientTransport,
|
||||
SessionKwargs,
|
||||
TransportOptions,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastmcp.server.server import FastMCP
|
||||
|
|
@ -52,8 +56,12 @@ class FastMCPTransport(ClientTransport):
|
|||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
options = transport_options or TransportOptions()
|
||||
async with create_client_server_memory_streams() as (
|
||||
client_streams,
|
||||
server_streams,
|
||||
|
|
@ -88,7 +96,7 @@ class FastMCPTransport(ClientTransport):
|
|||
)
|
||||
|
||||
try:
|
||||
async with ClientSession(
|
||||
async with options.session_class(
|
||||
read_stream=client_read,
|
||||
write_stream=client_write,
|
||||
**session_kwargs,
|
||||
|
|
|
|||
|
|
@ -18,7 +18,11 @@ from typing_extensions import Unpack
|
|||
from fastmcp.client.auth.bearer import BearerAuth
|
||||
from fastmcp.client.auth.oauth import OAuth
|
||||
from fastmcp.client.dependencies import get_http_headers
|
||||
from fastmcp.client.transports.base import ClientTransport, SessionKwargs
|
||||
from fastmcp.client.transports.base import (
|
||||
ClientTransport,
|
||||
SessionKwargs,
|
||||
TransportOptions,
|
||||
)
|
||||
from fastmcp.utilities.timeout import normalize_timeout_to_timedelta
|
||||
|
||||
|
||||
|
|
@ -61,8 +65,6 @@ class SSETransport(ClientTransport):
|
|||
|
||||
self._set_auth(auth)
|
||||
|
||||
self.forward_incoming_headers: bool = False
|
||||
|
||||
self.sse_read_timeout = normalize_timeout_to_timedelta(sse_read_timeout)
|
||||
|
||||
def _set_auth(self, auth: httpx2.Auth | Literal["oauth"] | str | None):
|
||||
|
|
@ -115,15 +117,19 @@ class SSETransport(ClientTransport):
|
|||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
options = transport_options or TransportOptions()
|
||||
client_kwargs: dict[str, Any] = {}
|
||||
|
||||
# When used in a proxy, forward the inbound request's authorization
|
||||
# header to the upstream server. This is off by default so that a
|
||||
# plain Client used inside a server tool handler doesn't accidentally
|
||||
# leak the caller's credentials to an unrelated remote server.
|
||||
if self.forward_incoming_headers:
|
||||
if options.forward_incoming_headers:
|
||||
client_kwargs["headers"] = (
|
||||
get_http_headers(include={"authorization"}) | self.headers
|
||||
)
|
||||
|
|
@ -149,7 +155,7 @@ class SSETransport(ClientTransport):
|
|||
|
||||
async with sse_client(self.url, auth=self.auth, **client_kwargs) as transport:
|
||||
read_stream, write_stream = transport
|
||||
async with ClientSession(
|
||||
async with options.session_class(
|
||||
read_stream, write_stream, **session_kwargs
|
||||
) as session:
|
||||
yield session
|
||||
|
|
|
|||
|
|
@ -12,7 +12,11 @@ from mcp import ClientSession, StdioServerParameters
|
|||
from mcp.client.stdio import stdio_client
|
||||
from typing_extensions import Unpack
|
||||
|
||||
from fastmcp.client.transports.base import ClientTransport, SessionKwargs
|
||||
from fastmcp.client.transports.base import (
|
||||
ClientTransport,
|
||||
SessionKwargs,
|
||||
TransportOptions,
|
||||
)
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
|
@ -63,17 +67,27 @@ class StdioTransport(ClientTransport):
|
|||
self.log_file = log_file
|
||||
|
||||
self._session: ClientSession | None = None
|
||||
self._session_options: TransportOptions | None = None
|
||||
self._active_sessions = 0
|
||||
self._connect_lock = anyio.Lock()
|
||||
self._connect_task: asyncio.Task | None = None
|
||||
self._ready_event = anyio.Event()
|
||||
self._stop_event = anyio.Event()
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_session(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
try:
|
||||
await self.connect(**session_kwargs)
|
||||
yield cast(ClientSession, self._session)
|
||||
await self.connect(transport_options=transport_options, **session_kwargs)
|
||||
self._active_sessions += 1
|
||||
try:
|
||||
yield cast(ClientSession, self._session)
|
||||
finally:
|
||||
self._active_sessions -= 1
|
||||
finally:
|
||||
if not self.keep_alive:
|
||||
await self.disconnect()
|
||||
|
|
@ -81,47 +95,76 @@ class StdioTransport(ClientTransport):
|
|||
logger.debug("Stdio transport has keep_alive=True, not disconnecting")
|
||||
|
||||
async def connect(
|
||||
self, **session_kwargs: Unpack[SessionKwargs]
|
||||
self,
|
||||
*,
|
||||
transport_options: TransportOptions | None = None,
|
||||
**session_kwargs: Unpack[SessionKwargs],
|
||||
) -> ClientSession | None:
|
||||
# If the connect task completed or the session's streams are dead,
|
||||
# the subprocess has exited. Tear down so we can start fresh.
|
||||
if self._connect_task is not None and (
|
||||
self._connect_task.done() or self._is_session_dead()
|
||||
):
|
||||
await self.disconnect()
|
||||
options = transport_options or TransportOptions()
|
||||
|
||||
if self._connect_task is not None:
|
||||
return
|
||||
# Serialized so concurrent callers can't each decide to replace the
|
||||
# session and race to spawn competing subprocesses.
|
||||
async with self._connect_lock:
|
||||
# A kept-alive session was built for one client's options; handing it
|
||||
# to a client that wants different ones would silently give it the
|
||||
# first client's behavior. Rebuild it when it's idle; refuse when
|
||||
# another client is using it, since tearing it down would break them.
|
||||
if self._connect_task is not None and self._session_options != options:
|
||||
if self._active_sessions:
|
||||
raise RuntimeError(
|
||||
"This stdio transport has a live session built for different "
|
||||
"connection options and another client is still using it. "
|
||||
"Sharing one transport across clients that need different "
|
||||
"sessions is not supported; give each client its own transport."
|
||||
)
|
||||
await self.disconnect()
|
||||
|
||||
session_future: asyncio.Future[ClientSession] = asyncio.Future()
|
||||
# If the connect task completed or the session's streams are dead,
|
||||
# the subprocess has exited. Tear down so we can start fresh.
|
||||
if self._connect_task is not None and (
|
||||
self._connect_task.done() or self._is_session_dead()
|
||||
):
|
||||
await self.disconnect()
|
||||
|
||||
# start the connection task
|
||||
self._connect_task = asyncio.create_task(
|
||||
_stdio_transport_connect_task(
|
||||
command=self.command,
|
||||
args=self.args,
|
||||
env=self.env,
|
||||
cwd=self.cwd,
|
||||
log_file=self.log_file,
|
||||
# TODO(ty): remove when ty supports Unpack[TypedDict] inference
|
||||
session_kwargs=session_kwargs, # type: ignore[arg-type]
|
||||
ready_event=self._ready_event,
|
||||
stop_event=self._stop_event,
|
||||
session_future=session_future,
|
||||
if self._connect_task is not None:
|
||||
return
|
||||
|
||||
session_future: asyncio.Future[ClientSession] = asyncio.Future()
|
||||
|
||||
# Recorded before the connect completes: while it is in flight the
|
||||
# session already belongs to these options, and a concurrent caller
|
||||
# comparing against an unset value would read it as a mismatch and
|
||||
# tear down the connection being established.
|
||||
self._session_options = options
|
||||
|
||||
# start the connection task
|
||||
self._connect_task = asyncio.create_task(
|
||||
_stdio_transport_connect_task(
|
||||
command=self.command,
|
||||
args=self.args,
|
||||
env=self.env,
|
||||
cwd=self.cwd,
|
||||
log_file=self.log_file,
|
||||
# TODO(ty): remove when ty supports Unpack[TypedDict] inference
|
||||
session_kwargs=session_kwargs, # type: ignore[arg-type]
|
||||
transport_options=options,
|
||||
ready_event=self._ready_event,
|
||||
stop_event=self._stop_event,
|
||||
session_future=session_future,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# wait for the client to be ready before returning
|
||||
await self._ready_event.wait()
|
||||
# wait for the client to be ready before returning
|
||||
await self._ready_event.wait()
|
||||
|
||||
# Check if connect task completed with an exception (early failure)
|
||||
if self._connect_task.done():
|
||||
exception = self._connect_task.exception()
|
||||
if exception is not None:
|
||||
raise exception
|
||||
# Check if connect task completed with an exception (early failure)
|
||||
if self._connect_task.done():
|
||||
exception = self._connect_task.exception()
|
||||
if exception is not None:
|
||||
raise exception
|
||||
|
||||
self._session = await session_future
|
||||
return self._session
|
||||
self._session = await session_future
|
||||
return self._session
|
||||
|
||||
async def disconnect(self):
|
||||
if self._connect_task is None:
|
||||
|
|
@ -137,6 +180,7 @@ class StdioTransport(ClientTransport):
|
|||
# reset variables and events for potential future reconnects
|
||||
self._connect_task = None
|
||||
self._session = None
|
||||
self._session_options = None
|
||||
self._stop_event = anyio.Event()
|
||||
self._ready_event = anyio.Event()
|
||||
|
||||
|
|
@ -181,6 +225,7 @@ async def _stdio_transport_connect_task(
|
|||
cwd: str | None,
|
||||
log_file: Path | TextIO | None,
|
||||
session_kwargs: SessionKwargs,
|
||||
transport_options: TransportOptions,
|
||||
ready_event: anyio.Event,
|
||||
stop_event: anyio.Event,
|
||||
session_future: asyncio.Future[ClientSession],
|
||||
|
|
@ -212,7 +257,9 @@ async def _stdio_transport_connect_task(
|
|||
read_stream, write_stream = transport
|
||||
session_future.set_result(
|
||||
await stack.enter_async_context(
|
||||
ClientSession(read_stream, write_stream, **session_kwargs)
|
||||
transport_options.session_class(
|
||||
read_stream, write_stream, **session_kwargs
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, cast
|
|||
import anyio
|
||||
import httpx2
|
||||
import mcp_types
|
||||
from mcp import ClientSession
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.shared.exceptions import MCPError
|
||||
|
|
@ -35,6 +36,7 @@ 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.client.transports.base import TransportOptions
|
||||
from fastmcp.exceptions import ResourceError
|
||||
from fastmcp.mcp_config import MCPConfig
|
||||
from fastmcp.prompts import Message, Prompt, PromptResult
|
||||
|
|
@ -66,6 +68,31 @@ logger = get_logger(__name__)
|
|||
ClientFactoryT = Callable[[], Client] | Callable[[], Awaitable[Client]]
|
||||
|
||||
|
||||
class _ForwardingClientSession(ClientSession):
|
||||
"""A session that does not enforce the backend's declared output schema.
|
||||
|
||||
`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.
|
||||
"""
|
||||
|
||||
async def validate_tool_result(
|
||||
self, name: str, result: mcp_types.CallToolResult
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
# Settings every proxy-backend connection uses: relay results without policing
|
||||
# the backend's output schema, and forward the caller's authorization header
|
||||
# upstream (appropriate for a proxy, where credentials are meant to propagate).
|
||||
PROXY_TRANSPORT_OPTIONS = TransportOptions(
|
||||
session_class=_ForwardingClientSession,
|
||||
forward_incoming_headers=True,
|
||||
)
|
||||
|
||||
|
||||
def _proxy_upstream_error(error: Exception) -> MCPError:
|
||||
return MCPError(
|
||||
code=mcp_types.INTERNAL_ERROR,
|
||||
|
|
@ -878,13 +905,16 @@ def _create_client_factory(
|
|||
if isinstance(target, Client):
|
||||
client = target
|
||||
|
||||
# Plain Clients used as proxy backends also need header forwarding,
|
||||
# same as ProxyClient (which sets this in __init__).
|
||||
from fastmcp.client.transports.http import StreamableHttpTransport
|
||||
from fastmcp.client.transports.sse import SSETransport
|
||||
def as_proxy_backend(c: Client) -> Client:
|
||||
"""Apply proxy connection settings to a copy we own.
|
||||
|
||||
if isinstance(client.transport, StreamableHttpTransport | SSETransport):
|
||||
client.transport.forward_incoming_headers = True
|
||||
The caller handed us their Client; configuring it in place would
|
||||
change how their own connections behave, including whether their
|
||||
credentials get forwarded upstream.
|
||||
"""
|
||||
fresh = c.new()
|
||||
fresh._transport_options = PROXY_TRANSPORT_OPTIONS
|
||||
return fresh
|
||||
|
||||
if client.is_connected() and type(client) is ProxyClient:
|
||||
logger.info(
|
||||
|
|
@ -893,23 +923,29 @@ def _create_client_factory(
|
|||
)
|
||||
|
||||
def fresh_client_factory() -> Client:
|
||||
return client.new()
|
||||
return as_proxy_backend(client)
|
||||
|
||||
return fresh_client_factory
|
||||
|
||||
if client.is_connected():
|
||||
logger.info(
|
||||
"Proxy detected connected client - reusing existing session for all requests. "
|
||||
"This may cause context mixing in concurrent scenarios."
|
||||
"This may cause context mixing in concurrent scenarios, and the session's "
|
||||
"existing settings apply, so backend results are validated against their "
|
||||
"declared output schema rather than relayed as-is. Pass a disconnected "
|
||||
"client to avoid both."
|
||||
)
|
||||
|
||||
# The caller's session is already built, so there are no connection
|
||||
# settings left to apply — proxy options only take effect at connect
|
||||
# time. Reuse is opt-in via passing an already-connected client.
|
||||
def reuse_client_factory() -> Client:
|
||||
return client
|
||||
|
||||
return reuse_client_factory
|
||||
|
||||
def fresh_client_factory() -> Client:
|
||||
return client.new()
|
||||
return as_proxy_backend(client)
|
||||
|
||||
return fresh_client_factory
|
||||
else:
|
||||
|
|
@ -1201,14 +1237,7 @@ class ProxyClient(Client[ClientTransportT]):
|
|||
self._proxy_restoring_handler_keys.add(key)
|
||||
super().__init__(transport=transport, **kwargs) # ty: ignore[no-matching-overload]
|
||||
|
||||
# Enable forwarding of inbound HTTP headers (e.g. authorization) to
|
||||
# the upstream server. This is only appropriate for proxy clients,
|
||||
# where the caller's credentials should be propagated.
|
||||
from fastmcp.client.transports.http import StreamableHttpTransport
|
||||
from fastmcp.client.transports.sse import SSETransport
|
||||
|
||||
if isinstance(self.transport, StreamableHttpTransport | SSETransport):
|
||||
self.transport.forward_incoming_headers = True
|
||||
self._transport_options = PROXY_TRANSPORT_OPTIONS
|
||||
|
||||
def _bind_restoring_handlers(self) -> None:
|
||||
if "roots" in self._proxy_restoring_handler_keys:
|
||||
|
|
|
|||
|
|
@ -1,10 +1,16 @@
|
|||
"""Client session and task error propagation tests."""
|
||||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import pytest
|
||||
from mcp import ClientSession
|
||||
from mcp_types import TextContent
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.client.transports import PythonStdioTransport
|
||||
from fastmcp.client.transports.base import TransportOptions
|
||||
|
||||
|
||||
class TestSessionTaskErrorPropagation:
|
||||
|
|
@ -135,3 +141,112 @@ class TestSessionTaskErrorPropagation:
|
|||
|
||||
# Restore for cleanup
|
||||
client._session_state.session_task = original_task
|
||||
|
||||
|
||||
class TestCustomSessionClass:
|
||||
"""Transports build the session class the client asks for."""
|
||||
|
||||
async def test_session_class_is_used_when_provided(self):
|
||||
built: list[str] = []
|
||||
|
||||
class RecordingClientSession(ClientSession):
|
||||
def __init__(self, *args, **kwargs):
|
||||
built.append("yes")
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
server = FastMCP("Server")
|
||||
|
||||
@server.tool
|
||||
def ping() -> str:
|
||||
return "pong"
|
||||
|
||||
client = Client(server)
|
||||
client._transport_options = TransportOptions(
|
||||
session_class=RecordingClientSession
|
||||
)
|
||||
async with client:
|
||||
await client.call_tool("ping")
|
||||
|
||||
assert built == ["yes"]
|
||||
|
||||
async def test_default_session_class_is_client_session(self):
|
||||
server = FastMCP("Server")
|
||||
|
||||
@server.tool
|
||||
def ping() -> str:
|
||||
return "pong"
|
||||
|
||||
client = Client(server)
|
||||
assert TransportOptions().session_class is ClientSession
|
||||
async with client:
|
||||
result = await client.call_tool("ping")
|
||||
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "pong"
|
||||
|
||||
|
||||
class TestKeepAliveSessionsRespectClientOptions:
|
||||
"""A cached stdio session must not be handed to a client wanting different options.
|
||||
|
||||
`StdioTransport` keeps its subprocess and session alive between connections
|
||||
by default. Serving that cached session to a second client would give it the
|
||||
first client's behavior — e.g. an ordinary client silently inheriting a
|
||||
proxy's non-validating session. Rebuilding it is only safe while nobody else
|
||||
is using it.
|
||||
"""
|
||||
|
||||
class Reconnected(Exception):
|
||||
"""Raised in place of a real teardown so the test stops at the guard."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def cached_connection(self, monkeypatch, options: TransportOptions):
|
||||
"""A transport that believes it already holds a session built for `options`."""
|
||||
transport = PythonStdioTransport(script_path=__file__, keep_alive=True)
|
||||
|
||||
async def never_finishes():
|
||||
await asyncio.sleep(60)
|
||||
|
||||
task = asyncio.create_task(never_finishes())
|
||||
transport._connect_task = task
|
||||
transport._session_options = options
|
||||
|
||||
async def fake_disconnect():
|
||||
raise TestKeepAliveSessionsRespectClientOptions.Reconnected
|
||||
|
||||
monkeypatch.setattr(transport, "disconnect", fake_disconnect)
|
||||
try:
|
||||
yield transport
|
||||
finally:
|
||||
task.cancel()
|
||||
transport._connect_task = None
|
||||
|
||||
async def test_matching_options_reuse_the_cached_session(self, monkeypatch):
|
||||
options = TransportOptions()
|
||||
async with self.cached_connection(monkeypatch, options) as transport:
|
||||
assert await transport.connect(transport_options=options) is None
|
||||
|
||||
async def test_differing_options_rebuild_an_idle_session(self, monkeypatch):
|
||||
class OtherSession(ClientSession):
|
||||
pass
|
||||
|
||||
async with self.cached_connection(monkeypatch, TransportOptions()) as transport:
|
||||
with pytest.raises(self.Reconnected):
|
||||
await transport.connect(
|
||||
transport_options=TransportOptions(session_class=OtherSession)
|
||||
)
|
||||
|
||||
async def test_differing_options_do_not_disturb_a_session_in_use(self, monkeypatch):
|
||||
"""Tearing down a live session would break the client already on it."""
|
||||
|
||||
class OtherSession(ClientSession):
|
||||
pass
|
||||
|
||||
async with self.cached_connection(monkeypatch, TransportOptions()) as transport:
|
||||
transport._active_sessions = 1
|
||||
|
||||
with pytest.raises(RuntimeError, match="still using it"):
|
||||
await transport.connect(
|
||||
transport_options=TransportOptions(session_class=OtherSession)
|
||||
)
|
||||
|
||||
assert transport._connect_task is not None
|
||||
|
|
|
|||
|
|
@ -15,13 +15,17 @@ from pydantic import AnyUrl
|
|||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.client.transports import FastMCPTransport, StreamableHttpTransport
|
||||
from fastmcp.client.transports.base import TransportOptions
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.mcp_config import MCPConfig
|
||||
from fastmcp.resources import ResourceContent, ResourceResult
|
||||
from fastmcp.server import create_proxy
|
||||
from fastmcp.server.middleware import Middleware
|
||||
from fastmcp.server.providers.proxy import (
|
||||
FastMCPProxy,
|
||||
ProxyClient,
|
||||
ProxyProvider,
|
||||
_ForwardingClientSession,
|
||||
)
|
||||
from fastmcp.tools.base import ToolResult
|
||||
from fastmcp.tools.tool_transform import (
|
||||
|
|
@ -1152,3 +1156,232 @@ class TestProxySpanAttributes:
|
|||
assert all(v is not None for v in attrs.values()), (
|
||||
f"OpenTelemetry rejects None attribute values; got {attrs!r}"
|
||||
)
|
||||
|
||||
|
||||
class TestProxyOutputSchemaEnforcement:
|
||||
"""A proxy relays tool results; it does not police the backend's schema.
|
||||
|
||||
`ClientSession.call_tool` validates structured content against the output
|
||||
schema the backend advertised. For a proxy that check is misplaced: it
|
||||
turns a backend's schema bug into a proxy error and hides the real
|
||||
response from the client that actually consumes it.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def backend_violating_its_schema(self) -> FastMCP:
|
||||
mcp = FastMCP("SchemaViolator")
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"status": {"enum": ["ok", "error"]}},
|
||||
"required": ["status"],
|
||||
}
|
||||
|
||||
@mcp.tool(output_schema=schema)
|
||||
def undeclared_status() -> dict:
|
||||
return {"status": "weird"}
|
||||
|
||||
@mcp.tool(output_schema=schema)
|
||||
def declared_status() -> dict:
|
||||
return {"status": "ok"}
|
||||
|
||||
return mcp
|
||||
|
||||
async def _call_without_validating(self, server: FastMCP, tool: str):
|
||||
"""Call through a client that does not enforce the schema itself."""
|
||||
client = Client(server)
|
||||
client._transport_options = TransportOptions(
|
||||
session_class=_ForwardingClientSession
|
||||
)
|
||||
async with client:
|
||||
return await client.call_tool_mcp(tool, {})
|
||||
|
||||
async def test_proxy_forwards_result_violating_backend_schema(
|
||||
self, backend_violating_its_schema
|
||||
):
|
||||
proxy = FastMCP("Proxy")
|
||||
proxy.add_provider(
|
||||
ProxyProvider(lambda: ProxyClient(backend_violating_its_schema))
|
||||
)
|
||||
|
||||
result = await self._call_without_validating(proxy, "undeclared_status")
|
||||
|
||||
assert result.is_error is False
|
||||
assert result.structured_content == {"status": "weird"}
|
||||
|
||||
async def test_proxy_forwards_conforming_result_unchanged(
|
||||
self, backend_violating_its_schema
|
||||
):
|
||||
proxy = FastMCP("Proxy")
|
||||
proxy.add_provider(
|
||||
ProxyProvider(lambda: ProxyClient(backend_violating_its_schema))
|
||||
)
|
||||
|
||||
result = await self._call_without_validating(proxy, "declared_status")
|
||||
|
||||
assert result.is_error is False
|
||||
assert result.structured_content == {"status": "ok"}
|
||||
|
||||
async def test_end_client_still_enforces_the_schema(
|
||||
self, backend_violating_its_schema
|
||||
):
|
||||
"""Skipping validation in the proxy doesn't disarm the real client."""
|
||||
proxy = FastMCP("Proxy")
|
||||
proxy.add_provider(
|
||||
ProxyProvider(lambda: ProxyClient(backend_violating_its_schema))
|
||||
)
|
||||
|
||||
async with Client(proxy) as client:
|
||||
with pytest.raises(RuntimeError, match="Invalid structured content"):
|
||||
await client.call_tool_mcp("undeclared_status", {})
|
||||
|
||||
async def test_direct_client_still_enforces_the_schema(
|
||||
self, backend_violating_its_schema
|
||||
):
|
||||
"""The behavior change is scoped to proxies, not clients generally."""
|
||||
async with Client(backend_violating_its_schema) as client:
|
||||
with pytest.raises(RuntimeError, match="Invalid structured content"):
|
||||
await client.call_tool_mcp("undeclared_status", {})
|
||||
|
||||
async def test_proxied_calls_do_not_refetch_the_backend_tool_list(self):
|
||||
"""Validation used to force a `tools/list` on every proxied call.
|
||||
|
||||
The proxy builds a fresh client per request, so the SDK's output-schema
|
||||
cache was always cold and each call paid an extra backend round trip.
|
||||
"""
|
||||
counts = {"list": 0, "call": 0}
|
||||
|
||||
class CountingMiddleware(Middleware):
|
||||
async def on_list_tools(self, context, call_next):
|
||||
counts["list"] += 1
|
||||
return await call_next(context)
|
||||
|
||||
async def on_call_tool(self, context, call_next):
|
||||
counts["call"] += 1
|
||||
return await call_next(context)
|
||||
|
||||
backend = FastMCP("Backend")
|
||||
backend.add_middleware(CountingMiddleware())
|
||||
|
||||
@backend.tool(
|
||||
output_schema={
|
||||
"type": "object",
|
||||
"properties": {"n": {"type": "integer"}},
|
||||
"required": ["n"],
|
||||
}
|
||||
)
|
||||
def echo(n: int) -> dict:
|
||||
return {"n": n}
|
||||
|
||||
proxy = FastMCP("Proxy")
|
||||
proxy.add_provider(ProxyProvider(lambda: ProxyClient(backend)))
|
||||
|
||||
async with Client(proxy) as client:
|
||||
await client.call_tool("echo", {"n": 1})
|
||||
lists_after_first = counts["list"]
|
||||
|
||||
for n in range(2, 5):
|
||||
await client.call_tool("echo", {"n": n})
|
||||
|
||||
assert counts["call"] == 4
|
||||
assert counts["list"] == lists_after_first
|
||||
|
||||
|
||||
class TestProxySettingsAreNotSharedBetweenClients:
|
||||
"""Proxy connection settings belong to the client, not to the transport.
|
||||
|
||||
Configuring a shared transport in place used to leak proxy behavior into
|
||||
unrelated clients — including header forwarding, which would send the
|
||||
caller's credentials to a server the user never meant to authorize.
|
||||
"""
|
||||
|
||||
def test_building_a_proxy_client_does_not_reconfigure_a_shared_transport(self):
|
||||
shared = StreamableHttpTransport("http://example.com/mcp/")
|
||||
plain = Client(shared)
|
||||
|
||||
ProxyClient(shared)
|
||||
|
||||
assert plain._transport_options is None
|
||||
|
||||
def test_proxy_client_carries_its_own_options(self):
|
||||
proxy_client = ProxyClient(StreamableHttpTransport("http://example.com/mcp/"))
|
||||
|
||||
options = proxy_client._transport_options
|
||||
assert options is not None
|
||||
assert options.forward_incoming_headers is True
|
||||
assert options.session_class is _ForwardingClientSession
|
||||
|
||||
def test_options_survive_the_per_request_client_copy(self):
|
||||
"""The proxy builds a fresh client per request via `new()`."""
|
||||
proxy_client = ProxyClient(StreamableHttpTransport("http://example.com/mcp/"))
|
||||
|
||||
assert proxy_client.new()._transport_options is proxy_client._transport_options
|
||||
|
||||
def test_a_user_supplied_client_is_not_reconfigured(self):
|
||||
"""`create_proxy(client)` must not change how the caller's client behaves."""
|
||||
user_client = Client(StreamableHttpTransport("http://example.com/mcp/"))
|
||||
|
||||
create_proxy(user_client)
|
||||
|
||||
assert user_client._transport_options is None
|
||||
|
||||
|
||||
class TestProxyForwardingAppliesToEveryBackendClient:
|
||||
"""Every path that builds a proxy backend gets the forwarding session.
|
||||
|
||||
`create_proxy` accepts plain Clients and MCPConfigs, none of which route
|
||||
through `ProxyClient.__init__`, so configuring only that constructor would
|
||||
leave those forms still rejecting backend results.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def backend(self) -> FastMCP:
|
||||
mcp = FastMCP("SchemaViolator")
|
||||
|
||||
@mcp.tool(
|
||||
output_schema={
|
||||
"type": "object",
|
||||
"properties": {"status": {"enum": ["ok", "error"]}},
|
||||
"required": ["status"],
|
||||
}
|
||||
)
|
||||
def status() -> dict:
|
||||
return {"status": "weird"}
|
||||
|
||||
return mcp
|
||||
|
||||
async def _forwarded(self, server: FastMCP, tool: str = "status"):
|
||||
client = Client(server)
|
||||
client._transport_options = TransportOptions(
|
||||
session_class=_ForwardingClientSession
|
||||
)
|
||||
async with client:
|
||||
return await client.call_tool_mcp(tool, {})
|
||||
|
||||
async def test_plain_client_target_forwards(self, backend):
|
||||
result = await self._forwarded(create_proxy(Client(backend)))
|
||||
|
||||
assert result.is_error is False
|
||||
assert result.structured_content == {"status": "weird"}
|
||||
|
||||
async def test_single_server_config_target_forwards(self, backend):
|
||||
port = find_available_port()
|
||||
async with run_server_async(backend, port=port):
|
||||
config = MCPConfig.from_dict(
|
||||
{"mcpServers": {"a": {"url": f"http://127.0.0.1:{port}/mcp/"}}}
|
||||
)
|
||||
result = await self._forwarded(create_proxy(Client(config)))
|
||||
|
||||
assert result.is_error is False
|
||||
assert result.structured_content == {"status": "weird"}
|
||||
|
||||
async def test_multi_server_config_target_forwards(self, backend):
|
||||
port = find_available_port()
|
||||
async with run_server_async(backend, port=port):
|
||||
url = f"http://127.0.0.1:{port}/mcp/"
|
||||
config = MCPConfig.from_dict(
|
||||
{"mcpServers": {"a": {"url": url}, "b": {"url": url}}}
|
||||
)
|
||||
result = await self._forwarded(create_proxy(Client(config)), "a_status")
|
||||
|
||||
assert result.is_error is False
|
||||
assert result.structured_content == {"status": "weird"}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue