diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index 803b5a205..01cfc287c 100644 --- a/src/fastmcp/client/client.py +++ b/src/fastmcp/client/client.py @@ -1,5 +1,5 @@ import datetime -from contextlib import AsyncExitStack +from contextlib import AsyncExitStack, asynccontextmanager from pathlib import Path from typing import Any, cast @@ -84,6 +84,7 @@ class Client: self._session: ClientSession | None = None self._exit_stack: AsyncExitStack | None = None self._nesting_counter: int = 0 + self._initialize_result: mcp.types.InitializeResult | None = None if log_handler is None: log_handler = default_log_handler @@ -117,10 +118,19 @@ class Client: """Get the current active session. Raises RuntimeError if not connected.""" if self._session is None: raise RuntimeError( - "Client is not connected. Use 'async with client:' context manager first." + "Client is not connected. Use the 'async with client:' context manager first." ) return self._session + @property + def initialize_result(self) -> mcp.types.InitializeResult: + """Get the result of the initialization request.""" + if self._initialize_result is None: + raise RuntimeError( + "Client is not connected. Use the 'async with client:' context manager first." + ) + return self._initialize_result + def set_roots(self, roots: RootsList | RootsHandler) -> None: """Set the roots for the client. This does not automatically call `send_roots_list_changed`.""" self._session_kwargs["list_roots_callback"] = create_roots_callback(roots) @@ -135,27 +145,35 @@ class Client: """Check if the client is currently connected.""" return self._session is not None + @asynccontextmanager + async def _context_manager(self): + with catch(get_catch_handlers()): + async with self.transport.connect_session( + **self._session_kwargs + ) as session: + self._session = session + # Initialize the session + self._initialize_result = await self._session.initialize() + + try: + yield + finally: + self._exit_stack = None + self._session = None + self._initialize_result = None + async def __aenter__(self): if self._nesting_counter == 0: # Create exit stack to manage both context managers stack = AsyncExitStack() await stack.__aenter__() - # Add the exception handling context - stack.enter_context(catch(get_catch_handlers())) + await stack.enter_async_context(self._context_manager()) - # the above catch will only apply once this __aenter__ finishes so - # we need to wrap the session creation in a new context in case it - # raises errors itself - with catch(get_catch_handlers()): - # Create and enter the transport session using the exit stack - session_cm = self.transport.connect_session(**self._session_kwargs) - self._session = await stack.enter_async_context(session_cm) - - # Store the stack for cleanup in __aexit__ self._exit_stack = stack self._nesting_counter += 1 + return self async def __aexit__(self, exc_type, exc_val, exc_tb): @@ -168,7 +186,6 @@ class Client: await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb) finally: self._exit_stack = None - self._session = None # --- MCP Client Methods --- diff --git a/src/fastmcp/client/transports.py b/src/fastmcp/client/transports.py index debc2629d..a5f0f95e0 100644 --- a/src/fastmcp/client/transports.py +++ b/src/fastmcp/client/transports.py @@ -46,6 +46,7 @@ class ClientTransport(abc.ABC): A Transport is responsible for establishing and managing connections to an MCP server, and providing a ClientSession within an async context. + """ @abc.abstractmethod @@ -54,7 +55,9 @@ class ClientTransport(abc.ABC): self, **session_kwargs: Unpack[SessionKwargs] ) -> AsyncIterator[ClientSession]: """ - Establishes a connection and yields an active, initialized ClientSession. + Establishes a connection and yields an active ClientSession. + + The ClientSession is *not* expected to be initialized in this context manager. The session is guaranteed to be valid only within the scope of the async context manager. Connection setup and teardown are handled @@ -65,7 +68,7 @@ class ClientTransport(abc.ABC): constructor (e.g., callbacks, timeouts). Yields: - An initialized mcp.ClientSession instance. + A mcp.ClientSession instance. """ raise NotImplementedError yield None # type: ignore @@ -94,7 +97,6 @@ class WSTransport(ClientTransport): async with ClientSession( read_stream, write_stream, **session_kwargs ) as session: - await session.initialize() # Initialize after session creation yield session def __repr__(self) -> str: @@ -143,7 +145,6 @@ class SSETransport(ClientTransport): async with ClientSession( read_stream, write_stream, **session_kwargs ) as session: - await session.initialize() yield session def __repr__(self) -> str: @@ -189,7 +190,6 @@ class StreamableHttpTransport(ClientTransport): async with ClientSession( read_stream, write_stream, **session_kwargs ) as session: - await session.initialize() yield session def __repr__(self) -> str: @@ -237,7 +237,6 @@ class StdioTransport(ClientTransport): async with ClientSession( read_stream, write_stream, **session_kwargs ) as session: - await session.initialize() yield session def __repr__(self) -> str: diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 7de6f13d0..efaa7f516 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -256,18 +256,51 @@ async def test_read_resource_mcp(fastmcp_server): async def test_client_connection(fastmcp_server): - """Test that the client connects and disconnects properly.""" + """Test that connect is idempotent.""" client = Client(transport=FastMCPTransport(fastmcp_server)) - # Before connection + # Connect idempotently + async with client: + assert client.is_connected() + # Make a request to ensure connection is working + await client.ping() assert not client.is_connected() - # During connection + +async def test_initialize_result_connected(fastmcp_server): + """Test that initialize_result returns the correct result when connected.""" + client = Client(transport=FastMCPTransport(fastmcp_server)) + + # Initialize result should not be accessible before connection + with pytest.raises(RuntimeError, match="Client is not connected"): + _ = client.initialize_result + + async with client: + # Once connected, initialize_result should be available + result = client.initialize_result + + # Verify the initialize result has expected properties + assert hasattr(result, "serverInfo") + assert result.serverInfo.name == "TestServer" + assert result.serverInfo.version is not None + + +async def test_initialize_result_disconnected(fastmcp_server): + """Test that initialize_result raises an error when not connected.""" + client = Client(transport=FastMCPTransport(fastmcp_server)) + + # Initialize result should not be accessible before connection + with pytest.raises(RuntimeError, match="Client is not connected"): + _ = client.initialize_result + + # Connect and then disconnect async with client: assert client.is_connected() - # After connection + # After disconnection, initialize_result should raise an error assert not client.is_connected() + with pytest.raises(RuntimeError, match="Client is not connected"): + _ = client.initialize_result async def test_client_nested_context_manager(fastmcp_server): diff --git a/tests/server/test_server_interactions.py b/tests/server/test_server_interactions.py index 1f9ef6f2b..2e739c2b8 100644 --- a/tests/server/test_server_interactions.py +++ b/tests/server/test_server_interactions.py @@ -640,7 +640,6 @@ class TestToolContextInjection: assert len(result) == 1 content = result[0] assert isinstance(content, TextContent) - assert content.text == "1" async def test_async_context(self): """Test that context works in async functions.""" @@ -656,8 +655,7 @@ class TestToolContextInjection: assert len(result) == 1 content = result[0] assert isinstance(content, TextContent) - assert "Async request" in content.text - assert "42" in content.text + assert content.text == "Async request 2: 42" async def test_optional_context(self): """Test that context is optional.""" @@ -798,7 +796,7 @@ class TestResourceContext: async with Client(mcp) as client: result = await client.read_resource(AnyUrl("resource://test")) assert isinstance(result[0], TextResourceContents) - assert result[0].text == "1" + assert result[0].text == "2" class TestResourceTemplates: @@ -1096,7 +1094,7 @@ class TestResourceTemplateContext: async with Client(mcp) as client: result = await client.read_resource(AnyUrl("resource://test")) assert isinstance(result[0], TextResourceContents) - assert result[0].text == "Resource template: test 1" + assert result[0].text.startswith("Resource template: test 2") class TestPrompts: