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
- # Add the exception handling context
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, initialized ClientSession.
 
 
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
- An initialized mcp.ClientSession instance.
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 the client connects and disconnects properly."""
254
  client = Client(transport=FastMCPTransport(fastmcp_server))
255
 
256
- # Before connection
 
 
 
 
257
  assert not client.is_connected()
258
 
259
- # During connection
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
260
  async with client:
261
  assert client.is_connected()
262
 
263
- # After connection
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 == "Resource template: test 1"
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: