Add startup wait time to stdio client

Author: mattzh72

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "sse-starlette>=1.6.1",
     "pydantic-settings>=2.5.2",
     "uvicorn>=0.23.1",
+    "exceptiongroup>=1.2.2",
 ]
 
 [project.optional-dependencies]

src/mcp/client/stdio.py

@@ -83,7 +83,7 @@ class StdioServerParameters(BaseModel):
 
 
 @asynccontextmanager
-async def stdio_client(server: StdioServerParameters):
+async def stdio_client(server: StdioServerParameters, startup_wait_time: float = 0.0):
     """
     Client transport for stdio: this will connect to a server by spawning a
     process and communicating with it over stdin/stdout.
@@ -97,11 +97,17 @@ async def stdio_client(server: StdioServerParameters):
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
-    process = await anyio.open_process(
-        [server.command, *server.args],
-        env=server.env if server.env is not None else get_default_environment(),
-        stderr=sys.stderr,
-    )
+    try:
+        process = await anyio.open_process(
+            [server.command, *server.args],
+            env=server.env or get_default_environment(),
+            stderr=sys.stderr,
+        )
+    except OSError as exc:
+        raise RuntimeError(
+            f"Failed to spawn process: {server.command} {server.args}. "
+            f"Check that the binary exists and is executable."
+        ) from exc
 
     async def stdout_reader():
         assert process.stdout, "Opened process is missing stdout"
@@ -144,10 +150,21 @@ async def stdin_writer():
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
 
-    async with (
-        anyio.create_task_group() as tg,
-        process,
-    ):
+    async def watch_process_exit():
+        returncode = await process.wait()
+        if returncode != 0:
+            raise RuntimeError(
+                f"Subprocess exited with code {returncode}. "
+                f"Command: {server.command}, {server.args}"
+            )
+
+    async with anyio.create_task_group() as tg, process:
         tg.start_soon(stdout_reader)
         tg.start_soon(stdin_writer)
+        tg.start_soon(watch_process_exit)
+
+        if startup_wait_time > 0:
+            with anyio.move_on_after(startup_wait_time):
+                await anyio.sleep_forever()
+
         yield read_stream, write_stream

tests/client/test_stdio.py

@@ -1,4 +1,11 @@
+import re
 import shutil
+import sys
+
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+else:
+    ExceptionGroup = ExceptionGroup
 
 import pytest
 
@@ -14,7 +21,6 @@ async def test_stdio_client():
     server_parameters = StdioServerParameters(command=tee)
 
     async with stdio_client(server_parameters) as (read_stream, write_stream):
-        # Test sending and receiving messages
         messages = [
             JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
             JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
@@ -31,13 +37,31 @@ async def test_stdio_client():
                     raise message
 
                 read_messages.append(message)
-                if len(read_messages) == 2:
+                if len(read_messages) == len(messages):
                     break
 
-        assert len(read_messages) == 2
-        assert read_messages[0] == JSONRPCMessage(
-            root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
-        )
-        assert read_messages[1] == JSONRPCMessage(
-            root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})
-        )
+        assert read_messages == messages
+
+
+@pytest.mark.anyio
+async def test_stdio_client_spawn_failure():
+    server_parameters = StdioServerParameters(command="/does/not/exist")
+
+    with pytest.raises(RuntimeError, match="Failed to spawn process"):
+        async with stdio_client(server_parameters):
+            pytest.fail("Should never be reached.")
+
+
+@pytest.mark.anyio
+async def test_stdio_client_nonzero_exit():
+    server_parameters = StdioServerParameters(
+        command="python", args=["-c", "import sys; sys.exit(2)"]
+    )
+
+    with pytest.raises(ExceptionGroup) as eg_info:
+        async with stdio_client(server_parameters, startup_wait_time=0.2):
+            pytest.fail("Should never be reached.")
+
+    exc = eg_info.value.exceptions[0]
+    assert isinstance(exc, RuntimeError)
+    assert re.search(r"exited with code 2", str(exc))