refactor: move stdio package to its own module (#1942)

src/mcp/client/__init__.py

@@ -3,7 +3,4 @@
 from mcp.client.client import Client
 from mcp.client.session import ClientSession
 
-__all__ = [
+__all__ = ["Client", "ClientSession"]
-    "Client",
-    "ClientSession",
-]

src/mcp/client/session_group.py

@@ -30,7 +30,7 @@ from mcp.shared.session import ProgressFnT
 
 
 class SseServerParameters(BaseModel):
-    """Parameters for intializing a sse_client."""
+    """Parameters for initializing a sse_client."""
 
     # The endpoint URL.
     url: str
@@ -46,7 +46,7 @@ class SseServerParameters(BaseModel):
 
 
 class StreamableHttpParameters(BaseModel):
-    """Parameters for intializing a streamable_http_client."""
+    """Parameters for initializing a streamable_http_client."""
 
     # The endpoint URL.
     url: str

src/mcp/client/stdio.py

@@ -0,0 +1,273 @@
+import logging
+import os
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Literal, TextIO
+
+import anyio
+import anyio.lowlevel
+from anyio.abc import Process
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from anyio.streams.text import TextReceiveStream
+from pydantic import BaseModel, Field
+
+import mcp.types as types
+from mcp.os.posix.utilities import terminate_posix_process_tree
+from mcp.os.win32.utilities import (
+    FallbackProcess,
+    create_windows_process,
+    get_windows_executable_command,
+    terminate_windows_process_tree,
+)
+from mcp.shared.message import SessionMessage
+
+logger = logging.getLogger(__name__)
+
+# Environment variables to inherit by default
+DEFAULT_INHERITED_ENV_VARS = (
+    [
+        "APPDATA",
+        "HOMEDRIVE",
+        "HOMEPATH",
+        "LOCALAPPDATA",
+        "PATH",
+        "PATHEXT",
+        "PROCESSOR_ARCHITECTURE",
+        "SYSTEMDRIVE",
+        "SYSTEMROOT",
+        "TEMP",
+        "USERNAME",
+        "USERPROFILE",
+    ]
+    if sys.platform == "win32"
+    else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"]
+)
+
+# Timeout for process termination before falling back to force kill
+PROCESS_TERMINATION_TIMEOUT = 2.0
+
+
+def get_default_environment() -> dict[str, str]:
+    """Returns a default environment object including only environment variables deemed
+    safe to inherit.
+    """
+    env: dict[str, str] = {}
+
+    for key in DEFAULT_INHERITED_ENV_VARS:
+        value = os.environ.get(key)
+        if value is None:
+            continue  # pragma: no cover
+
+        if value.startswith("()"):  # pragma: no cover
+            # Skip functions, which are a security risk
+            continue  # pragma: no cover
+
+        env[key] = value
+
+    return env
+
+
+class StdioServerParameters(BaseModel):
+    command: str
+    """The executable to run to start the server."""
+
+    args: list[str] = Field(default_factory=list)
+    """Command line arguments to pass to the executable."""
+
+    env: dict[str, str] | None = None
+    """
+    The environment to use when spawning the process.
+
+    If not specified, the result of get_default_environment() will be used.
+    """
+
+    cwd: str | Path | None = None
+    """The working directory to use when spawning the process."""
+
+    encoding: str = "utf-8"
+    """
+    The text encoding used when sending/receiving messages to the server
+
+    defaults to utf-8
+    """
+
+    encoding_error_handler: Literal["strict", "ignore", "replace"] = "strict"
+    """
+    The text encoding error handler.
+
+    See https://docs.python.org/3/library/codecs.html#codec-base-classes for
+    explanations of possible values
+    """
+
+
+@asynccontextmanager
+async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stderr):
+    """Client transport for stdio: this will connect to a server by spawning a
+    process and communicating with it over stdin/stdout.
+    """
+    read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[SessionMessage]
+    write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    try:
+        command = _get_executable_command(server.command)
+
+        # Open process with stderr piped for capture
+        process = await _create_platform_compatible_process(
+            command=command,
+            args=server.args,
+            env=({**get_default_environment(), **server.env} if server.env is not None else get_default_environment()),
+            errlog=errlog,
+            cwd=server.cwd,
+        )
+    except OSError:
+        # Clean up streams if process creation fails
+        await read_stream.aclose()
+        await write_stream.aclose()
+        await read_stream_writer.aclose()
+        await write_stream_reader.aclose()
+        raise
+
+    async def stdout_reader():
+        assert process.stdout, "Opened process is missing stdout"
+
+        try:
+            async with read_stream_writer:
+                buffer = ""
+                async for chunk in TextReceiveStream(
+                    process.stdout,
+                    encoding=server.encoding,
+                    errors=server.encoding_error_handler,
+                ):
+                    lines = (buffer + chunk).split("\n")
+                    buffer = lines.pop()
+
+                    for line in lines:
+                        try:
+                            message = types.jsonrpc_message_adapter.validate_json(line, by_name=False)
+                        except Exception as exc:  # pragma: no cover
+                            logger.exception("Failed to parse JSONRPC message from server")
+                            await read_stream_writer.send(exc)
+                            continue
+
+                        session_message = SessionMessage(message)
+                        await read_stream_writer.send(session_message)
+        except anyio.ClosedResourceError:  # pragma: no cover
+            await anyio.lowlevel.checkpoint()
+
+    async def stdin_writer():
+        assert process.stdin, "Opened process is missing stdin"
+
+        try:
+            async with write_stream_reader:
+                async for session_message in write_stream_reader:
+                    json = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
+                    await process.stdin.send(
+                        (json + "\n").encode(
+                            encoding=server.encoding,
+                            errors=server.encoding_error_handler,
+                        )
+                    )
+        except anyio.ClosedResourceError:  # pragma: no cover
+            await anyio.lowlevel.checkpoint()
+
+    async with (
+        anyio.create_task_group() as tg,
+        process,
+    ):
+        tg.start_soon(stdout_reader)
+        tg.start_soon(stdin_writer)
+        try:
+            yield read_stream, write_stream
+        finally:
+            # MCP spec: stdio shutdown sequence
+            # 1. Close input stream to server
+            # 2. Wait for server to exit, or send SIGTERM if it doesn't exit in time
+            # 3. Send SIGKILL if still not exited
+            if process.stdin:  # pragma: no branch
+                try:
+                    await process.stdin.aclose()
+                except Exception:  # pragma: no cover
+                    # stdin might already be closed, which is fine
+                    pass
+
+            try:
+                # Give the process time to exit gracefully after stdin closes
+                with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT):
+                    await process.wait()
+            except TimeoutError:
+                # Process didn't exit from stdin closure, use platform-specific termination
+                # which handles SIGTERM -> SIGKILL escalation
+                await _terminate_process_tree(process)
+            except ProcessLookupError:  # pragma: no cover
+                # Process already exited, which is fine
+                pass
+            await read_stream.aclose()
+            await write_stream.aclose()
+            await read_stream_writer.aclose()
+            await write_stream_reader.aclose()
+
+
+def _get_executable_command(command: str) -> str:
+    """Get the correct executable command normalized for the current platform.
+
+    Args:
+        command: Base command (e.g., 'uvx', 'npx')
+
+    Returns:
+        str: Platform-appropriate command
+    """
+    if sys.platform == "win32":  # pragma: no cover
+        return get_windows_executable_command(command)
+    else:
+        return command  # pragma: no cover
+
+
+async def _create_platform_compatible_process(
+    command: str,
+    args: list[str],
+    env: dict[str, str] | None = None,
+    errlog: TextIO = sys.stderr,
+    cwd: Path | str | None = None,
+):
+    """Creates a subprocess in a platform-compatible way.
+
+    Unix: Creates process in a new session/process group for killpg support
+    Windows: Creates process in a Job Object for reliable child termination
+    """
+    if sys.platform == "win32":  # pragma: no cover
+        process = await create_windows_process(command, args, env, errlog, cwd)
+    else:
+        process = await anyio.open_process(
+            [command, *args],
+            env=env,
+            stderr=errlog,
+            cwd=cwd,
+            start_new_session=True,
+        )  # pragma: no cover
+
+    return process
+
+
+async def _terminate_process_tree(process: Process | FallbackProcess, timeout_seconds: float = 2.0) -> None:
+    """Terminate a process and all its children using platform-specific methods.
+
+    Unix: Uses os.killpg() for atomic process group termination
+    Windows: Uses Job Objects via pywin32 for reliable child process cleanup
+
+    Args:
+        process: The process to terminate
+        timeout_seconds: Timeout in seconds before force killing (default: 2.0)
+    """
+    if sys.platform == "win32":  # pragma: no cover
+        await terminate_windows_process_tree(process, timeout_seconds)
+    else:  # pragma: no cover
+        # FallbackProcess should only be used for Windows compatibility
+        assert isinstance(process, Process)
+        await terminate_posix_process_tree(process, timeout_seconds)

src/mcp/client/websocket.py

@@ -8,7 +8,7 @@ from pydantic import ValidationError
 from websockets.asyncio.client import connect as ws_connect
 from websockets.typing import Subprotocol
 
-import mcp.types as types
+from mcp import types
 from mcp.shared.message import SessionMessage