Spaces:
Running
Running
Jeremiah Lowin commited on
Commit ·
445f86a
1
Parent(s): bbd582a
Improve early error handling for stdio transport
Browse files- src/fastmcp/client/client.py +28 -10
- src/fastmcp/client/transports.py +36 -22
- tests/client/test_stdio.py +11 -0
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 |
-
|
| 340 |
-
|
| 341 |
-
|
| 342 |
-
|
| 343 |
-
|
| 344 |
-
|
| 345 |
-
|
| 346 |
-
|
| 347 |
-
|
| 348 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
| 365 |
-
|
| 366 |
-
|
| 367 |
-
|
| 368 |
-
|
| 369 |
-
|
| 370 |
-
|
| 371 |
-
|
| 372 |
-
|
| 373 |
-
|
| 374 |
-
|
| 375 |
-
|
| 376 |
-
|
| 377 |
-
|
| 378 |
-
|
| 379 |
-
|
| 380 |
-
|
| 381 |
-
|
| 382 |
-
|
| 383 |
-
|
| 384 |
-
|
| 385 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|