Jeremiah Lowin commited on
Commit
445f86a
·
1 Parent(s): bbd582a

Improve early error handling for stdio transport

Browse files
src/fastmcp/client/client.py CHANGED
@@ -286,6 +286,19 @@ class Client(Generic[ClientTransportT]):
286
 
287
  async def __aenter__(self):
288
  await self._connect()
 
 
 
 
 
 
 
 
 
 
 
 
 
289
  return self
290
 
291
  async def __aexit__(self, exc_type, exc_val, exc_tb):
@@ -336,16 +349,21 @@ class Client(Generic[ClientTransportT]):
336
  self._initialize_result = None
337
 
338
  async def _session_runner(self):
339
- async with AsyncExitStack() as stack:
340
- try:
341
- await stack.enter_async_context(self._context_manager())
342
- # Session/context is now ready
343
- self._ready_event.set()
344
- # Wait until disconnect/stop is requested
345
- await self._stop_event.wait()
346
- finally:
347
- # On exit, ensure ready event is set (idempotent)
348
- self._ready_event.set()
 
 
 
 
 
349
 
350
  async def close(self):
351
  await self._disconnect(force=True)
 
286
 
287
  async def __aenter__(self):
288
  await self._connect()
289
+
290
+ # Check if session task failed and raise error immediately
291
+ if (
292
+ self._session_task is not None
293
+ and self._session_task.done()
294
+ and not self._session_task.cancelled()
295
+ ):
296
+ exception = self._session_task.exception()
297
+ if exception is not None:
298
+ raise RuntimeError(
299
+ f"Client failed to connect: {exception}"
300
+ ) from exception
301
+
302
  return self
303
 
304
  async def __aexit__(self, exc_type, exc_val, exc_tb):
 
349
  self._initialize_result = None
350
 
351
  async def _session_runner(self):
352
+ try:
353
+ async with AsyncExitStack() as stack:
354
+ try:
355
+ await stack.enter_async_context(self._context_manager())
356
+ # Session/context is now ready
357
+ self._ready_event.set()
358
+ # Wait until disconnect/stop is requested
359
+ await self._stop_event.wait()
360
+ finally:
361
+ # On exit, ensure ready event is set (idempotent)
362
+ self._ready_event.set()
363
+ except Exception:
364
+ # Ensure ready event is set even if context manager entry fails
365
+ self._ready_event.set()
366
+ raise
367
 
368
  async def close(self):
369
  await self._disconnect(force=True)
src/fastmcp/client/transports.py CHANGED
@@ -361,34 +361,48 @@ class StdioTransport(ClientTransport):
361
  async def _connect_task():
362
  from mcp.client.stdio import stdio_client
363
 
364
- async with contextlib.AsyncExitStack() as stack:
365
- try:
366
- server_params = StdioServerParameters(
367
- command=self.command, args=self.args, env=self.env, cwd=self.cwd
368
- )
369
- transport = await stack.enter_async_context(
370
- stdio_client(server_params)
371
- )
372
- read_stream, write_stream = transport
373
- self._session = await stack.enter_async_context(
374
- ClientSession(read_stream, write_stream, **session_kwargs)
375
- )
376
-
377
- logger.debug("Stdio transport connected")
378
- self._ready_event.set()
379
-
380
- # Wait until disconnect is requested (stop_event is set)
381
- await self._stop_event.wait()
382
- finally:
383
- # Clean up client on exit
384
- self._session = None
385
- logger.debug("Stdio transport disconnected")
 
 
 
 
 
 
 
 
386
 
387
  # start the connection task
388
  self._connect_task = asyncio.create_task(_connect_task())
389
  # wait for the client to be ready before returning
390
  await self._ready_event.wait()
391
 
 
 
 
 
 
 
392
  async def disconnect(self):
393
  if self._connect_task is None:
394
  return
 
361
  async def _connect_task():
362
  from mcp.client.stdio import stdio_client
363
 
364
+ try:
365
+ async with contextlib.AsyncExitStack() as stack:
366
+ try:
367
+ server_params = StdioServerParameters(
368
+ command=self.command,
369
+ args=self.args,
370
+ env=self.env,
371
+ cwd=self.cwd,
372
+ )
373
+ transport = await stack.enter_async_context(
374
+ stdio_client(server_params)
375
+ )
376
+ read_stream, write_stream = transport
377
+ self._session = await stack.enter_async_context(
378
+ ClientSession(read_stream, write_stream, **session_kwargs)
379
+ )
380
+
381
+ logger.debug("Stdio transport connected")
382
+ self._ready_event.set()
383
+
384
+ # Wait until disconnect is requested (stop_event is set)
385
+ await self._stop_event.wait()
386
+ finally:
387
+ # Clean up client on exit
388
+ self._session = None
389
+ logger.debug("Stdio transport disconnected")
390
+ except Exception:
391
+ # Ensure ready event is set even if connection fails
392
+ self._ready_event.set()
393
+ raise
394
 
395
  # start the connection task
396
  self._connect_task = asyncio.create_task(_connect_task())
397
  # wait for the client to be ready before returning
398
  await self._ready_event.wait()
399
 
400
+ # Check if connect task completed with an exception (early failure)
401
+ if self._connect_task.done():
402
+ exception = self._connect_task.exception()
403
+ if exception is not None:
404
+ raise exception
405
+
406
  async def disconnect(self):
407
  if self._connect_task is None:
408
  return
tests/client/test_stdio.py CHANGED
@@ -115,3 +115,14 @@ class TestKeepAlive:
115
  await client.close()
116
  with pytest.raises(RuntimeError, match="Client is not connected"):
117
  await client.call_tool("pid")
 
 
 
 
 
 
 
 
 
 
 
 
115
  await client.close()
116
  with pytest.raises(RuntimeError, match="Client is not connected"):
117
  await client.call_tool("pid")
118
+
119
+ async def test_session_task_failure_raises_immediately_on_enter(self):
120
+ # Use a command that will fail to start
121
+ client = Client(
122
+ transport=StdioTransport(command="nonexistent_command", args=[])
123
+ )
124
+
125
+ # Should raise RuntimeError immediately, not defer until first use
126
+ with pytest.raises(RuntimeError, match="Client failed to connect"):
127
+ async with client:
128
+ pass