Merge pull request #220 from yihuang/main

src/fastmcp/client/client.py

@@ -45,7 +45,8 @@ class Client:
     ):
         self.transport = infer_transport(transport)
         self._session: ClientSession | None = None
-        self._session_cms: list[AbstractAsyncContextManager[ClientSession]] = []
+        self._session_cm: AbstractAsyncContextManager[ClientSession] | None = None
+        self._nesting_counter: int = 0
 
         self._session_kwargs: SessionKwargs = {
             "sampling_callback": None,
@@ -85,29 +86,21 @@ class Client:
         return self._session is not None
 
     async def __aenter__(self):
-        if self.is_connected():
-            # We're already connected, no need to add None to the session_cms list
-            return self
-
-        try:
-            session_cm = self.transport.connect_session(**self._session_kwargs)
-            self._session_cms.append(session_cm)
-            self._session = await self._session_cms[-1].__aenter__()
-            return self
-        except Exception as e:
-            # Ensure cleanup if __aenter__ fails partially
-            self._session = None
-            if self._session_cms:
-                self._session_cms.pop()
-            raise ConnectionError(
-                f"Failed to connect using {self.transport}: {e}"
-            ) from e
+        if self._nesting_counter == 0:
+            # create new session
+            self._session_cm = self.transport.connect_session(**self._session_kwargs)
+            self._session = await self._session_cm.__aenter__()
+
+        self._nesting_counter += 1
+        return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        if self._session_cms:
-            await self._session_cms[-1].__aexit__(exc_type, exc_val, exc_tb)
+        self._nesting_counter -= 1
+
+        if self._nesting_counter == 0 and self._session_cm is not None:
+            await self._session_cm.__aexit__(exc_type, exc_val, exc_tb)
+            self._session_cm = None
             self._session = None
-            self._session_cms.pop()
 
     # --- MCP Client Methods ---
     async def ping(self) -> None:

src/fastmcp/server/proxy.py

@@ -3,7 +3,9 @@ from urllib.parse import quote
 
 import mcp.types
 from mcp.server.lowlevel.helper_types import ReadResourceContents
+from mcp.shared.exceptions import McpError
 from mcp.types import (
+    METHOD_NOT_FOUND,
     BlobResourceContents,
     EmbeddedResource,
     GetPromptResult,
@@ -173,7 +175,14 @@ class FastMCPProxy(FastMCP):
         tools = await super().get_tools()
 
         async with self.client:
-            for tool in await self.client.list_tools():
+            try:
+                client_tools = await self.client.list_tools()
+            except McpError as e:
+                if e.error.code == METHOD_NOT_FOUND:
+                    client_tools = []
+                else:
+                    raise e
+            for tool in client_tools:
                 tool_proxy = await ProxyTool.from_client(self.client, tool)
                 tools[tool_proxy.name] = tool_proxy
 
@@ -183,7 +192,14 @@ class FastMCPProxy(FastMCP):
         resources = await super().get_resources()
 
         async with self.client:
-            for resource in await self.client.list_resources():
+            try:
+                client_resources = await self.client.list_resources()
+            except McpError as e:
+                if e.error.code == METHOD_NOT_FOUND:
+                    client_resources = []
+                else:
+                    raise e
+            for resource in client_resources:
                 resource_proxy = await ProxyResource.from_client(self.client, resource)
                 resources[str(resource_proxy.uri)] = resource_proxy
 
@@ -193,7 +209,14 @@ class FastMCPProxy(FastMCP):
         templates = await super().get_resource_templates()
 
         async with self.client:
-            for template in await self.client.list_resource_templates():
+            try:
+                client_templates = await self.client.list_resource_templates()
+            except McpError as e:
+                if e.error.code == METHOD_NOT_FOUND:
+                    client_templates = []
+                else:
+                    raise e
+            for template in client_templates:
                 template_proxy = await ProxyTemplate.from_client(self.client, template)
                 templates[template_proxy.uri_template] = template_proxy
 
@@ -203,7 +226,14 @@ class FastMCPProxy(FastMCP):
         prompts = await super().get_prompts()
 
         async with self.client:
-            for prompt in await self.client.list_prompts():
+            try:
+                client_prompts = await self.client.list_prompts()
+            except McpError as e:
+                if e.error.code == METHOD_NOT_FOUND:
+                    client_prompts = []
+                else:
+                    raise e
+            for prompt in client_prompts:
                 prompt_proxy = await ProxyPrompt.from_client(self.client, prompt)
                 prompts[prompt_proxy.name] = prompt_proxy
         return prompts

tests/client/test_client.py

@@ -1,10 +1,15 @@
+import contextlib
+from collections.abc import AsyncIterator
 from typing import cast
 
 import pytest
+from mcp import ClientSession
+from mcp.shared.memory import create_client_server_memory_streams
 from pydantic import AnyUrl
+from typing_extensions import Unpack
 
 from fastmcp.client import Client
-from fastmcp.client.transports import FastMCPTransport
+from fastmcp.client.transports import ClientTransport, FastMCPTransport, SessionKwargs
 from fastmcp.server.server import FastMCP
 
 
@@ -160,6 +165,45 @@ async def test_client_connection(fastmcp_server):
     assert not client.is_connected()
 
 
+async def test_client_nested_context_manager(fastmcp_server):
+    """Test that the client connects and disconnects once in nested context manager."""
+
+    class MockTransport(ClientTransport):
+        def __init__(self):
+            self._connected = False
+
+        @contextlib.asynccontextmanager
+        async def connect_session(
+            self,
+            **session_kwargs: Unpack[SessionKwargs],
+        ) -> AsyncIterator[ClientSession]:
+            assert not self._connected, "Transport is connected multiple times"
+            self._connected = True
+            async with create_client_server_memory_streams() as (
+                _,
+                server_streams,
+            ):
+                yield ClientSession(*server_streams)
+
+    client = Client(transport=MockTransport())
+
+    # Before connection
+    assert not client.is_connected()
+
+    # During connection
+    async with client:
+        assert client.is_connected()
+
+        async with client:
+            assert client.is_connected()
+
+        async with client:
+            assert client.is_connected()
+
+    # After connection
+    assert not client.is_connected()
+
+
 async def test_resource_template(fastmcp_server):
     """Test using a resource template with InMemoryClient."""
     client = Client(transport=FastMCPTransport(fastmcp_server))