mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 13:34:17 +02:00
Improve early error handling for stdio transport
This commit is contained in:
parent
44f6101606
commit
ed331e717b
3 changed files with 75 additions and 32 deletions
|
|
@ -286,6 +286,19 @@ class Client(Generic[ClientTransportT]):
|
|||
|
||||
async def __aenter__(self):
|
||||
await self._connect()
|
||||
|
||||
# Check if session task failed and raise error immediately
|
||||
if (
|
||||
self._session_task is not None
|
||||
and self._session_task.done()
|
||||
and not self._session_task.cancelled()
|
||||
):
|
||||
exception = self._session_task.exception()
|
||||
if exception is not None:
|
||||
raise RuntimeError(
|
||||
f"Client failed to connect: {exception}"
|
||||
) from exception
|
||||
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
|
|
@ -336,16 +349,21 @@ class Client(Generic[ClientTransportT]):
|
|||
self._initialize_result = None
|
||||
|
||||
async def _session_runner(self):
|
||||
async with AsyncExitStack() as stack:
|
||||
try:
|
||||
await stack.enter_async_context(self._context_manager())
|
||||
# Session/context is now ready
|
||||
self._ready_event.set()
|
||||
# Wait until disconnect/stop is requested
|
||||
await self._stop_event.wait()
|
||||
finally:
|
||||
# On exit, ensure ready event is set (idempotent)
|
||||
self._ready_event.set()
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
try:
|
||||
await stack.enter_async_context(self._context_manager())
|
||||
# Session/context is now ready
|
||||
self._ready_event.set()
|
||||
# Wait until disconnect/stop is requested
|
||||
await self._stop_event.wait()
|
||||
finally:
|
||||
# On exit, ensure ready event is set (idempotent)
|
||||
self._ready_event.set()
|
||||
except Exception:
|
||||
# Ensure ready event is set even if context manager entry fails
|
||||
self._ready_event.set()
|
||||
raise
|
||||
|
||||
async def close(self):
|
||||
await self._disconnect(force=True)
|
||||
|
|
|
|||
|
|
@ -361,34 +361,48 @@ class StdioTransport(ClientTransport):
|
|||
async def _connect_task():
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
||||
async with contextlib.AsyncExitStack() as stack:
|
||||
try:
|
||||
server_params = StdioServerParameters(
|
||||
command=self.command, args=self.args, env=self.env, cwd=self.cwd
|
||||
)
|
||||
transport = await stack.enter_async_context(
|
||||
stdio_client(server_params)
|
||||
)
|
||||
read_stream, write_stream = transport
|
||||
self._session = await stack.enter_async_context(
|
||||
ClientSession(read_stream, write_stream, **session_kwargs)
|
||||
)
|
||||
try:
|
||||
async with contextlib.AsyncExitStack() as stack:
|
||||
try:
|
||||
server_params = StdioServerParameters(
|
||||
command=self.command,
|
||||
args=self.args,
|
||||
env=self.env,
|
||||
cwd=self.cwd,
|
||||
)
|
||||
transport = await stack.enter_async_context(
|
||||
stdio_client(server_params)
|
||||
)
|
||||
read_stream, write_stream = transport
|
||||
self._session = await stack.enter_async_context(
|
||||
ClientSession(read_stream, write_stream, **session_kwargs)
|
||||
)
|
||||
|
||||
logger.debug("Stdio transport connected")
|
||||
self._ready_event.set()
|
||||
logger.debug("Stdio transport connected")
|
||||
self._ready_event.set()
|
||||
|
||||
# Wait until disconnect is requested (stop_event is set)
|
||||
await self._stop_event.wait()
|
||||
finally:
|
||||
# Clean up client on exit
|
||||
self._session = None
|
||||
logger.debug("Stdio transport disconnected")
|
||||
# Wait until disconnect is requested (stop_event is set)
|
||||
await self._stop_event.wait()
|
||||
finally:
|
||||
# Clean up client on exit
|
||||
self._session = None
|
||||
logger.debug("Stdio transport disconnected")
|
||||
except Exception:
|
||||
# Ensure ready event is set even if connection fails
|
||||
self._ready_event.set()
|
||||
raise
|
||||
|
||||
# start the connection task
|
||||
self._connect_task = asyncio.create_task(_connect_task())
|
||||
# wait for the client to be ready before returning
|
||||
await self._ready_event.wait()
|
||||
|
||||
# Check if connect task completed with an exception (early failure)
|
||||
if self._connect_task.done():
|
||||
exception = self._connect_task.exception()
|
||||
if exception is not None:
|
||||
raise exception
|
||||
|
||||
async def disconnect(self):
|
||||
if self._connect_task is None:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -115,3 +115,14 @@ class TestKeepAlive:
|
|||
await client.close()
|
||||
with pytest.raises(RuntimeError, match="Client is not connected"):
|
||||
await client.call_tool("pid")
|
||||
|
||||
async def test_session_task_failure_raises_immediately_on_enter(self):
|
||||
# Use a command that will fail to start
|
||||
client = Client(
|
||||
transport=StdioTransport(command="nonexistent_command", args=[])
|
||||
)
|
||||
|
||||
# Should raise RuntimeError immediately, not defer until first use
|
||||
with pytest.raises(RuntimeError, match="Client failed to connect"):
|
||||
async with client:
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue