diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index ebcb0f80e..24fb6c680 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, cast from urllib.parse import quote import mcp.types +from mcp import ServerSession from mcp.client.session import ClientSession from mcp.shared.context import LifespanContextT, RequestContext from mcp.shared.exceptions import McpError @@ -606,6 +607,10 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]): Note that it is essential to ensure that the proxy server itself is also stateful. """ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._caches: dict[ServerSession, Client[ClientTransportT]] = {} + async def __aexit__(self, exc_type, exc_value, traceback) -> None: """ The stateful proxy client will be forced disconnected when the session is exited. @@ -613,6 +618,14 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]): """ pass + async def clear(self): + """ + Clear all cached clients and force disconnect them. + """ + while self._caches: + _, cache = self._caches.popitem() + await cache._disconnect(force=True) + def new_stateful(self) -> Client[ClientTransportT]: """ Create a new stateful proxy client instance with the same configuration. @@ -620,15 +633,15 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]): Use this method as the client factory for stateful proxy server. """ session = get_context().session - proxy_client = session.__dict__.get("_proxy_client", None) + proxy_client = self._caches.get(session, None) if proxy_client is None: proxy_client = self.new() logger.debug(f"{proxy_client} created for {session}") - session.__dict__["_proxy_client"] = proxy_client + self._caches[session] = proxy_client async def _on_session_exit(): - proxy_client: Client = session.__dict__.pop("_proxy_client") + self._caches.pop(session) logger.debug(f"{proxy_client} will be disconnect") await proxy_client._disconnect(force=True) diff --git a/tests/server/proxy/test_stateful_proxy_client.py b/tests/server/proxy/test_stateful_proxy_client.py index 52bf86f33..9db4114d6 100644 --- a/tests/server/proxy/test_stateful_proxy_client.py +++ b/tests/server/proxy/test_stateful_proxy_client.py @@ -118,3 +118,31 @@ class TestStatefulProxyClient: with pytest.raises(ToolError, match="Value not found"): await client.call_tool("stateful_get", {}) + + async def test_multi_proxies_no_mixing(self): + """Test that the stateful proxy client won't be mixed in multi-proxies sessions.""" + mcp_a, mcp_b = FastMCP(), FastMCP() + + @mcp_a.tool + def tool_a() -> str: + return "a" + + @mcp_b.tool + def tool_b() -> str: + return "b" + + proxy_mcp_a = FastMCPProxy( + client_factory=StatefulProxyClient(mcp_a).new_stateful + ) + proxy_mcp_b = FastMCPProxy( + client_factory=StatefulProxyClient(mcp_b).new_stateful + ) + multi_proxy_mcp = FastMCP() + multi_proxy_mcp.mount(proxy_mcp_a, prefix="a") + multi_proxy_mcp.mount(proxy_mcp_b, prefix="b") + + async with Client(multi_proxy_mcp) as client: + result_a = await client.call_tool("a_tool_a", {}) + result_b = await client.call_tool("b_tool_b", {}) + assert result_a.data == "a" + assert result_b.data == "b"