Spaces:
Running
Running
Jeremiah Lowin commited on
Commit ·
3aecde3
1
Parent(s): a7c6f7e
Store the initialize result on the client
Browse files
src/fastmcp/client/client.py
CHANGED
|
@@ -1,5 +1,5 @@
|
|
| 1 |
import datetime
|
| 2 |
-
from contextlib import AsyncExitStack
|
| 3 |
from pathlib import Path
|
| 4 |
from typing import Any, cast
|
| 5 |
|
|
@@ -80,6 +80,7 @@ class Client:
|
|
| 80 |
self._session: ClientSession | None = None
|
| 81 |
self._exit_stack: AsyncExitStack | None = None
|
| 82 |
self._nesting_counter: int = 0
|
|
|
|
| 83 |
|
| 84 |
if isinstance(timeout, int | float):
|
| 85 |
timeout = datetime.timedelta(seconds=timeout)
|
|
@@ -103,10 +104,19 @@ class Client:
|
|
| 103 |
"""Get the current active session. Raises RuntimeError if not connected."""
|
| 104 |
if self._session is None:
|
| 105 |
raise RuntimeError(
|
| 106 |
-
"Client is not connected. Use 'async with client:' context manager first."
|
| 107 |
)
|
| 108 |
return self._session
|
| 109 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 110 |
def set_roots(self, roots: RootsList | RootsHandler) -> None:
|
| 111 |
"""Set the roots for the client. This does not automatically call `send_roots_list_changed`."""
|
| 112 |
self._session_kwargs["list_roots_callback"] = create_roots_callback(roots)
|
|
@@ -121,27 +131,35 @@ class Client:
|
|
| 121 |
"""Check if the client is currently connected."""
|
| 122 |
return self._session is not None
|
| 123 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 124 |
async def __aenter__(self):
|
| 125 |
if self._nesting_counter == 0:
|
| 126 |
# Create exit stack to manage both context managers
|
| 127 |
stack = AsyncExitStack()
|
| 128 |
await stack.__aenter__()
|
| 129 |
|
| 130 |
-
|
| 131 |
-
stack.enter_context(catch(get_catch_handlers()))
|
| 132 |
|
| 133 |
-
# the above catch will only apply once this __aenter__ finishes so
|
| 134 |
-
# we need to wrap the session creation in a new context in case it
|
| 135 |
-
# raises errors itself
|
| 136 |
-
with catch(get_catch_handlers()):
|
| 137 |
-
# Create and enter the transport session using the exit stack
|
| 138 |
-
session_cm = self.transport.connect_session(**self._session_kwargs)
|
| 139 |
-
self._session = await stack.enter_async_context(session_cm)
|
| 140 |
-
|
| 141 |
-
# Store the stack for cleanup in __aexit__
|
| 142 |
self._exit_stack = stack
|
| 143 |
|
| 144 |
self._nesting_counter += 1
|
|
|
|
| 145 |
return self
|
| 146 |
|
| 147 |
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
@@ -154,7 +172,6 @@ class Client:
|
|
| 154 |
await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
|
| 155 |
finally:
|
| 156 |
self._exit_stack = None
|
| 157 |
-
self._session = None
|
| 158 |
|
| 159 |
# --- MCP Client Methods ---
|
| 160 |
|
|
|
|
| 1 |
import datetime
|
| 2 |
+
from contextlib import AsyncExitStack, asynccontextmanager
|
| 3 |
from pathlib import Path
|
| 4 |
from typing import Any, cast
|
| 5 |
|
|
|
|
| 80 |
self._session: ClientSession | None = None
|
| 81 |
self._exit_stack: AsyncExitStack | None = None
|
| 82 |
self._nesting_counter: int = 0
|
| 83 |
+
self._initialize_result: mcp.types.InitializeResult | None = None
|
| 84 |
|
| 85 |
if isinstance(timeout, int | float):
|
| 86 |
timeout = datetime.timedelta(seconds=timeout)
|
|
|
|
| 104 |
"""Get the current active session. Raises RuntimeError if not connected."""
|
| 105 |
if self._session is None:
|
| 106 |
raise RuntimeError(
|
| 107 |
+
"Client is not connected. Use the 'async with client:' context manager first."
|
| 108 |
)
|
| 109 |
return self._session
|
| 110 |
|
| 111 |
+
@property
|
| 112 |
+
def initialize_result(self) -> mcp.types.InitializeResult:
|
| 113 |
+
"""Get the result of the initialization request."""
|
| 114 |
+
if self._initialize_result is None:
|
| 115 |
+
raise RuntimeError(
|
| 116 |
+
"Client is not connected. Use the 'async with client:' context manager first."
|
| 117 |
+
)
|
| 118 |
+
return self._initialize_result
|
| 119 |
+
|
| 120 |
def set_roots(self, roots: RootsList | RootsHandler) -> None:
|
| 121 |
"""Set the roots for the client. This does not automatically call `send_roots_list_changed`."""
|
| 122 |
self._session_kwargs["list_roots_callback"] = create_roots_callback(roots)
|
|
|
|
| 131 |
"""Check if the client is currently connected."""
|
| 132 |
return self._session is not None
|
| 133 |
|
| 134 |
+
@asynccontextmanager
|
| 135 |
+
async def _context_manager(self):
|
| 136 |
+
with catch(get_catch_handlers()):
|
| 137 |
+
async with self.transport.connect_session(
|
| 138 |
+
**self._session_kwargs
|
| 139 |
+
) as session:
|
| 140 |
+
self._session = session
|
| 141 |
+
# Initialize the session
|
| 142 |
+
self._initialize_result = await self._session.initialize()
|
| 143 |
+
|
| 144 |
+
try:
|
| 145 |
+
yield
|
| 146 |
+
finally:
|
| 147 |
+
self._exit_stack = None
|
| 148 |
+
self._session = None
|
| 149 |
+
self._initialize_result = None
|
| 150 |
+
|
| 151 |
async def __aenter__(self):
|
| 152 |
if self._nesting_counter == 0:
|
| 153 |
# Create exit stack to manage both context managers
|
| 154 |
stack = AsyncExitStack()
|
| 155 |
await stack.__aenter__()
|
| 156 |
|
| 157 |
+
await stack.enter_async_context(self._context_manager())
|
|
|
|
| 158 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 159 |
self._exit_stack = stack
|
| 160 |
|
| 161 |
self._nesting_counter += 1
|
| 162 |
+
|
| 163 |
return self
|
| 164 |
|
| 165 |
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
|
|
| 172 |
await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
|
| 173 |
finally:
|
| 174 |
self._exit_stack = None
|
|
|
|
| 175 |
|
| 176 |
# --- MCP Client Methods ---
|
| 177 |
|
src/fastmcp/client/transports.py
CHANGED
|
@@ -44,6 +44,7 @@ class ClientTransport(abc.ABC):
|
|
| 44 |
|
| 45 |
A Transport is responsible for establishing and managing connections
|
| 46 |
to an MCP server, and providing a ClientSession within an async context.
|
|
|
|
| 47 |
"""
|
| 48 |
|
| 49 |
@abc.abstractmethod
|
|
@@ -52,7 +53,9 @@ class ClientTransport(abc.ABC):
|
|
| 52 |
self, **session_kwargs: Unpack[SessionKwargs]
|
| 53 |
) -> AsyncIterator[ClientSession]:
|
| 54 |
"""
|
| 55 |
-
Establishes a connection and yields an active
|
|
|
|
|
|
|
| 56 |
|
| 57 |
The session is guaranteed to be valid only within the scope of the
|
| 58 |
async context manager. Connection setup and teardown are handled
|
|
@@ -63,7 +66,7 @@ class ClientTransport(abc.ABC):
|
|
| 63 |
constructor (e.g., callbacks, timeouts).
|
| 64 |
|
| 65 |
Yields:
|
| 66 |
-
|
| 67 |
"""
|
| 68 |
raise NotImplementedError
|
| 69 |
yield None # type: ignore
|
|
@@ -92,7 +95,6 @@ class WSTransport(ClientTransport):
|
|
| 92 |
async with ClientSession(
|
| 93 |
read_stream, write_stream, **session_kwargs
|
| 94 |
) as session:
|
| 95 |
-
await session.initialize() # Initialize after session creation
|
| 96 |
yield session
|
| 97 |
|
| 98 |
def __repr__(self) -> str:
|
|
@@ -141,7 +143,6 @@ class SSETransport(ClientTransport):
|
|
| 141 |
async with ClientSession(
|
| 142 |
read_stream, write_stream, **session_kwargs
|
| 143 |
) as session:
|
| 144 |
-
await session.initialize()
|
| 145 |
yield session
|
| 146 |
|
| 147 |
def __repr__(self) -> str:
|
|
@@ -187,7 +188,6 @@ class StreamableHttpTransport(ClientTransport):
|
|
| 187 |
async with ClientSession(
|
| 188 |
read_stream, write_stream, **session_kwargs
|
| 189 |
) as session:
|
| 190 |
-
await session.initialize()
|
| 191 |
yield session
|
| 192 |
|
| 193 |
def __repr__(self) -> str:
|
|
@@ -235,7 +235,6 @@ class StdioTransport(ClientTransport):
|
|
| 235 |
async with ClientSession(
|
| 236 |
read_stream, write_stream, **session_kwargs
|
| 237 |
) as session:
|
| 238 |
-
await session.initialize()
|
| 239 |
yield session
|
| 240 |
|
| 241 |
def __repr__(self) -> str:
|
|
|
|
| 44 |
|
| 45 |
A Transport is responsible for establishing and managing connections
|
| 46 |
to an MCP server, and providing a ClientSession within an async context.
|
| 47 |
+
|
| 48 |
"""
|
| 49 |
|
| 50 |
@abc.abstractmethod
|
|
|
|
| 53 |
self, **session_kwargs: Unpack[SessionKwargs]
|
| 54 |
) -> AsyncIterator[ClientSession]:
|
| 55 |
"""
|
| 56 |
+
Establishes a connection and yields an active ClientSession.
|
| 57 |
+
|
| 58 |
+
The ClientSession is *not* expected to be initialized in this context manager.
|
| 59 |
|
| 60 |
The session is guaranteed to be valid only within the scope of the
|
| 61 |
async context manager. Connection setup and teardown are handled
|
|
|
|
| 66 |
constructor (e.g., callbacks, timeouts).
|
| 67 |
|
| 68 |
Yields:
|
| 69 |
+
A mcp.ClientSession instance.
|
| 70 |
"""
|
| 71 |
raise NotImplementedError
|
| 72 |
yield None # type: ignore
|
|
|
|
| 95 |
async with ClientSession(
|
| 96 |
read_stream, write_stream, **session_kwargs
|
| 97 |
) as session:
|
|
|
|
| 98 |
yield session
|
| 99 |
|
| 100 |
def __repr__(self) -> str:
|
|
|
|
| 143 |
async with ClientSession(
|
| 144 |
read_stream, write_stream, **session_kwargs
|
| 145 |
) as session:
|
|
|
|
| 146 |
yield session
|
| 147 |
|
| 148 |
def __repr__(self) -> str:
|
|
|
|
| 188 |
async with ClientSession(
|
| 189 |
read_stream, write_stream, **session_kwargs
|
| 190 |
) as session:
|
|
|
|
| 191 |
yield session
|
| 192 |
|
| 193 |
def __repr__(self) -> str:
|
|
|
|
| 235 |
async with ClientSession(
|
| 236 |
read_stream, write_stream, **session_kwargs
|
| 237 |
) as session:
|
|
|
|
| 238 |
yield session
|
| 239 |
|
| 240 |
def __repr__(self) -> str:
|
tests/client/test_client.py
CHANGED
|
@@ -250,18 +250,51 @@ async def test_read_resource_mcp(fastmcp_server):
|
|
| 250 |
|
| 251 |
|
| 252 |
async def test_client_connection(fastmcp_server):
|
| 253 |
-
"""Test that
|
| 254 |
client = Client(transport=FastMCPTransport(fastmcp_server))
|
| 255 |
|
| 256 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
| 257 |
assert not client.is_connected()
|
| 258 |
|
| 259 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 260 |
async with client:
|
| 261 |
assert client.is_connected()
|
| 262 |
|
| 263 |
-
# After
|
| 264 |
assert not client.is_connected()
|
|
|
|
|
|
|
| 265 |
|
| 266 |
|
| 267 |
async def test_client_nested_context_manager(fastmcp_server):
|
|
|
|
| 250 |
|
| 251 |
|
| 252 |
async def test_client_connection(fastmcp_server):
|
| 253 |
+
"""Test that connect is idempotent."""
|
| 254 |
client = Client(transport=FastMCPTransport(fastmcp_server))
|
| 255 |
|
| 256 |
+
# Connect idempotently
|
| 257 |
+
async with client:
|
| 258 |
+
assert client.is_connected()
|
| 259 |
+
# Make a request to ensure connection is working
|
| 260 |
+
await client.ping()
|
| 261 |
assert not client.is_connected()
|
| 262 |
|
| 263 |
+
|
| 264 |
+
async def test_initialize_result_connected(fastmcp_server):
|
| 265 |
+
"""Test that initialize_result returns the correct result when connected."""
|
| 266 |
+
client = Client(transport=FastMCPTransport(fastmcp_server))
|
| 267 |
+
|
| 268 |
+
# Initialize result should not be accessible before connection
|
| 269 |
+
with pytest.raises(RuntimeError, match="Client is not connected"):
|
| 270 |
+
_ = client.initialize_result
|
| 271 |
+
|
| 272 |
+
async with client:
|
| 273 |
+
# Once connected, initialize_result should be available
|
| 274 |
+
result = client.initialize_result
|
| 275 |
+
|
| 276 |
+
# Verify the initialize result has expected properties
|
| 277 |
+
assert hasattr(result, "serverInfo")
|
| 278 |
+
assert result.serverInfo.name == "TestServer"
|
| 279 |
+
assert result.serverInfo.version is not None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
async def test_initialize_result_disconnected(fastmcp_server):
|
| 283 |
+
"""Test that initialize_result raises an error when not connected."""
|
| 284 |
+
client = Client(transport=FastMCPTransport(fastmcp_server))
|
| 285 |
+
|
| 286 |
+
# Initialize result should not be accessible before connection
|
| 287 |
+
with pytest.raises(RuntimeError, match="Client is not connected"):
|
| 288 |
+
_ = client.initialize_result
|
| 289 |
+
|
| 290 |
+
# Connect and then disconnect
|
| 291 |
async with client:
|
| 292 |
assert client.is_connected()
|
| 293 |
|
| 294 |
+
# After disconnection, initialize_result should raise an error
|
| 295 |
assert not client.is_connected()
|
| 296 |
+
with pytest.raises(RuntimeError, match="Client is not connected"):
|
| 297 |
+
_ = client.initialize_result
|
| 298 |
|
| 299 |
|
| 300 |
async def test_client_nested_context_manager(fastmcp_server):
|
tests/server/test_server_interactions.py
CHANGED
|
@@ -640,7 +640,6 @@ class TestToolContextInjection:
|
|
| 640 |
assert len(result) == 1
|
| 641 |
content = result[0]
|
| 642 |
assert isinstance(content, TextContent)
|
| 643 |
-
assert content.text == "1"
|
| 644 |
|
| 645 |
async def test_async_context(self):
|
| 646 |
"""Test that context works in async functions."""
|
|
@@ -798,7 +797,6 @@ class TestResourceContext:
|
|
| 798 |
async with Client(mcp) as client:
|
| 799 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 800 |
assert isinstance(result[0], TextResourceContents)
|
| 801 |
-
assert result[0].text == "1"
|
| 802 |
|
| 803 |
|
| 804 |
class TestResourceTemplates:
|
|
@@ -1096,7 +1094,7 @@ class TestResourceTemplateContext:
|
|
| 1096 |
async with Client(mcp) as client:
|
| 1097 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 1098 |
assert isinstance(result[0], TextResourceContents)
|
| 1099 |
-
assert result[0].text
|
| 1100 |
|
| 1101 |
|
| 1102 |
class TestPrompts:
|
|
|
|
| 640 |
assert len(result) == 1
|
| 641 |
content = result[0]
|
| 642 |
assert isinstance(content, TextContent)
|
|
|
|
| 643 |
|
| 644 |
async def test_async_context(self):
|
| 645 |
"""Test that context works in async functions."""
|
|
|
|
| 797 |
async with Client(mcp) as client:
|
| 798 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 799 |
assert isinstance(result[0], TextResourceContents)
|
|
|
|
| 800 |
|
| 801 |
|
| 802 |
class TestResourceTemplates:
|
|
|
|
| 1094 |
async with Client(mcp) as client:
|
| 1095 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 1096 |
assert isinstance(result[0], TextResourceContents)
|
| 1097 |
+
assert result[0].text.startswith("Resource template: test")
|
| 1098 |
|
| 1099 |
|
| 1100 |
class TestPrompts:
|