Fix stateful proxy client mixing in multi-proxies sessions (#1245)

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)
 

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"