From 9c861b232bd43e51519ddb4a0edf4d3cf38be1b6 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Mon, 3 Nov 2025 13:01:30 -0800 Subject: [PATCH] Add manual initialization control to Client (#2355) * Add manual initialization control to Client - Add auto_initialize parameter (default True) to control automatic initialization - Make initialize() method public with idempotent caching - Add comprehensive test suite for initialization behavior * Document client initialization control and server instructions - Expand documentation to cover auto_initialize parameter - Show manual initialization for advanced use cases - Document accessing server instructions via initialize_result * Update client.mdx --- docs/clients/client.mdx | 49 ++++++++++++++ src/fastmcp/client/client.py | 84 ++++++++++++++++++----- tests/client/test_client.py | 125 ++++++++++++++++++++++++++++++++--- 3 files changed, 230 insertions(+), 28 deletions(-) diff --git a/docs/clients/client.mdx b/docs/clients/client.mdx index 0f5dc0f1e..36f300030 100644 --- a/docs/clients/client.mdx +++ b/docs/clients/client.mdx @@ -223,6 +223,55 @@ async with client: print("Server is reachable") ``` +### Initialization and Server Information + +When you enter the client context manager, the client automatically performs an MCP initialization handshake with the server. This handshake exchanges capabilities, server metadata, and instructions. The result is available through the `initialize_result` property. + +```python +from fastmcp import Client, FastMCP + +mcp = FastMCP(name="MyServer", instructions="Use the greet tool to say hello!") + +@mcp.tool +def greet(name: str) -> str: + """Greet a user by name.""" + return f"Hello, {name}!" + +async with Client(mcp) as client: + # Initialization already happened automatically + print(f"Server: {client.initialize_result.serverInfo.name}") + print(f"Version: {client.initialize_result.serverInfo.version}") + print(f"Instructions: {client.initialize_result.instructions}") + print(f"Capabilities: {client.initialize_result.capabilities.tools}") +``` + +#### Manual Initialization Control + +In advanced scenarios, you might want precise control over when initialization happens. For example, you may need custom error handling, want to defer initialization until after other setup, or need to measure initialization timing separately. + +Disable automatic initialization and call `initialize()` manually: + +```python +from fastmcp import Client + +# Disable automatic initialization +client = Client("my_mcp_server.py", auto_initialize=False) + +async with client: + # Connection established, but not initialized yet + print(f"Connected: {client.is_connected()}") + print(f"Initialized: {client.initialize_result is not None}") # False + + # Initialize manually with custom timeout + result = await client.initialize(timeout=10.0) + print(f"Server: {result.serverInfo.name}") + + # Now ready for operations + tools = await client.list_tools() +``` + +The `initialize()` method is idempotent - calling it multiple times returns the cached result from the first successful call. + ## Client Configuration Clients can be configured with additional handlers and settings for specialized use cases. diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index 9c5201b94..9c7840d52 100644 --- a/src/fastmcp/client/client.py +++ b/src/fastmcp/client/client.py @@ -77,6 +77,16 @@ logger = get_logger(__name__) T = TypeVar("T", bound="ClientTransport") +def _timeout_to_seconds( + timeout: datetime.timedelta | float | int | None, +) -> float | None: + if timeout is None: + return None + if isinstance(timeout, datetime.timedelta): + return timeout.total_seconds() + return float(timeout) + + @dataclass class ClientSessionState: """Holds all session-related state for a Client instance. @@ -222,6 +232,7 @@ class Client(Generic[ClientTransportT]): message_handler: MessageHandlerT | MessageHandler | None = None, progress_handler: ProgressHandler | None = None, timeout: datetime.timedelta | float | int | None = None, + auto_initialize: bool = True, init_timeout: datetime.timedelta | float | int | None = None, client_info: mcp.types.Implementation | None = None, auth: httpx.Auth | Literal["oauth"] | str | None = None, @@ -247,13 +258,9 @@ class Client(Generic[ClientTransportT]): # handle init handshake timeout if init_timeout is None: init_timeout = fastmcp.settings.client_init_timeout - if isinstance(init_timeout, datetime.timedelta): - init_timeout = init_timeout.total_seconds() - elif not init_timeout: - init_timeout = None - else: - init_timeout = float(init_timeout) - self._init_timeout = init_timeout + self._init_timeout = _timeout_to_seconds(init_timeout) + + self.auto_initialize = auto_initialize self._session_kwargs: SessionKwargs = { "sampling_callback": None, @@ -291,12 +298,8 @@ class Client(Generic[ClientTransportT]): return self._session_state.session @property - def initialize_result(self) -> mcp.types.InitializeResult: + def initialize_result(self) -> mcp.types.InitializeResult | None: """Get the result of the initialization request.""" - if self._session_state.initialize_result is None: - raise RuntimeError( - "Client is not connected. Use the 'async with client:' context manager first." - ) return self._session_state.initialize_result def set_roots(self, roots: RootsList | RootsHandler) -> None: @@ -358,15 +361,11 @@ class Client(Generic[ClientTransportT]): self._session_state.session = session # Initialize the session try: - with anyio.fail_after(self._init_timeout): - self._session_state.initialize_result = ( - await self._session_state.session.initialize() - ) + if self.auto_initialize: + await self.initialize() yield except anyio.ClosedResourceError as e: raise RuntimeError("Server session was closed unexpectedly") from e - except TimeoutError as e: - raise RuntimeError("Failed to initialize server session") from e finally: self._session_state.session = None self._session_state.initialize_result = None @@ -494,6 +493,55 @@ class Client(Generic[ClientTransportT]): # --- MCP Client Methods --- + async def initialize( + self, + timeout: datetime.timedelta | float | int | None = None, + ) -> mcp.types.InitializeResult: + """Send an initialize request to the server. + + This method performs the MCP initialization handshake with the server, + exchanging capabilities and server information. It is idempotent - calling + it multiple times returns the cached result from the first call. + + The initialization happens automatically when entering the client context + manager unless `auto_initialize=False` was set during client construction. + Manual calls to this method are only needed when auto-initialization is disabled. + + Args: + timeout: Optional timeout for the initialization request (seconds or timedelta). + If None, uses the client's init_timeout setting. + + Returns: + InitializeResult: The server's initialization response containing server info, + capabilities, protocol version, and optional instructions. + + Raises: + RuntimeError: If the client is not connected or initialization times out. + + Example: + ```python + # With auto-initialization disabled + client = Client(server, auto_initialize=False) + async with client: + result = await client.initialize() + print(f"Server: {result.serverInfo.name}") + print(f"Instructions: {result.instructions}") + ``` + """ + + if self.initialize_result is not None: + return self.initialize_result + + if timeout is None: + timeout = self._init_timeout + try: + with anyio.fail_after(_timeout_to_seconds(timeout)): + initialize_result = await self.session.initialize() + self._session_state.initialize_result = initialize_result + return initialize_result + except TimeoutError as e: + raise RuntimeError("Failed to initialize server session") from e + async def ping(self) -> bool: """Send a ping request.""" result = await self.session.send_ping() diff --git a/tests/client/test_client.py b/tests/client/test_client.py index a32e92073..8097b2965 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -386,9 +386,8 @@ 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 + # Initialize result should be None before connection + assert client.initialize_result is None async with client: # Once connected, initialize_result should be available @@ -401,21 +400,19 @@ async def test_initialize_result_connected(fastmcp_server): async def test_initialize_result_disconnected(fastmcp_server): - """Test that initialize_result raises an error when not connected.""" + """Test that initialize_result is None 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 + # Initialize result should be None before connection + assert client.initialize_result is None # Connect and then disconnect async with client: assert client.is_connected() - # After disconnection, initialize_result should raise an error + # After disconnection, initialize_result should be None again assert not client.is_connected() - with pytest.raises(RuntimeError, match="Client is not connected"): - _ = client.initialize_result + assert client.initialize_result is None async def test_server_info_custom_version(): @@ -1026,3 +1023,111 @@ class TestAuth: assert isinstance(client.transport, SSETransport) assert isinstance(client.transport.auth, BearerAuth) assert client.transport.auth.token.get_secret_value() == "test_token" + + +class TestInitialize: + """Tests for client initialization behavior.""" + + async def test_auto_initialize_default(self, fastmcp_server): + """Test that auto_initialize=True is the default and works automatically.""" + client = Client(fastmcp_server) + + async with client: + # Should be automatically initialized + assert client.initialize_result is not None + assert client.initialize_result.serverInfo.name == "TestServer" + assert client.initialize_result.instructions is None + + async def test_auto_initialize_explicit_true(self, fastmcp_server): + """Test explicit auto_initialize=True.""" + client = Client(fastmcp_server, auto_initialize=True) + + async with client: + assert client.initialize_result is not None + assert client.initialize_result.serverInfo.name == "TestServer" + + async def test_auto_initialize_false(self, fastmcp_server): + """Test that auto_initialize=False prevents automatic initialization.""" + client = Client(fastmcp_server, auto_initialize=False) + + async with client: + # Should not be automatically initialized + assert client.initialize_result is None + + async def test_manual_initialize(self, fastmcp_server): + """Test manual initialization when auto_initialize=False.""" + client = Client(fastmcp_server, auto_initialize=False) + + async with client: + # Manually initialize + result = await client.initialize() + + assert result is not None + assert result.serverInfo.name == "TestServer" + assert client.initialize_result is result + + async def test_initialize_idempotent(self, fastmcp_server): + """Test that calling initialize() multiple times returns cached result.""" + client = Client(fastmcp_server, auto_initialize=False) + + async with client: + result1 = await client.initialize() + result2 = await client.initialize() + result3 = await client.initialize() + + # All should return the same cached result + assert result1 is result2 + assert result2 is result3 + + async def test_initialize_with_instructions(self): + """Test that server instructions are available via initialize_result.""" + server = FastMCP("InstructionsServer", instructions="Use the greet tool!") + + @server.tool + def greet(name: str) -> str: + return f"Hello, {name}!" + + client = Client(server) + + async with client: + assert client.initialize_result.instructions == "Use the greet tool!" + + async def test_initialize_timeout_custom(self, fastmcp_server): + """Test custom timeout for initialize().""" + client = Client(fastmcp_server, auto_initialize=False) + + async with client: + # Should succeed with reasonable timeout + result = await client.initialize(timeout=5.0) + assert result is not None + + async def test_initialize_property_after_auto_init(self, fastmcp_server): + """Test accessing initialize_result property after auto-initialization.""" + client = Client(fastmcp_server, auto_initialize=True) + + async with client: + # Access via property + result = client.initialize_result + assert result.serverInfo.name == "TestServer" + + # Call method - should return cached + result2 = await client.initialize() + assert result is result2 + + async def test_initialize_property_before_connect(self, fastmcp_server): + """Test that initialize_result property is None before connection.""" + client = Client(fastmcp_server) + + # Not yet connected + assert client.initialize_result is None + + async def test_manual_initialize_can_call_tools(self, fastmcp_server): + """Test that manually initialized client can call tools.""" + client = Client(fastmcp_server, auto_initialize=False) + + async with client: + await client.initialize() + + # Should be able to call tools after manual initialization + result = await client.call_tool("greet", {"name": "World"}) + assert "Hello, World!" in str(result.content)