mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
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
This commit is contained in:
Jeremiah Lowin
GitHub
parent
6cc9559f84
commit
9c861b232b
3 changed files with 230 additions and 28 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue