stdio.py

"""stdio client transport.

Runs an MCP server as a subprocess and exchanges newline-delimited JSON-RPC
messages with it over stdin/stdout. Two pipe tasks bridge the server's pipes
to the session's in-memory streams; shutdown follows the MCP spec sequence
(close stdin, wait, then kill the process tree) inside a cancellation shield
with every wait bounded, so a cancelled caller can neither leak a live server
process nor hang on one.
"""

import logging
import os
import sys
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager, suppress
from pathlib import Path
from typing import Literal, TextIO

import anyio
import anyio.lowlevel
import anyio.to_thread
import mcp_types as types
from anyio.abc import AsyncResource, Process
from anyio.streams.text import TextReceiveStream
from pydantic import BaseModel, Field

from mcp.client._transport import TransportStreams
from mcp.os.posix.utilities import terminate_posix_process_tree
from mcp.os.win32.utilities import (
    ServerProcess,
    close_process_job,
    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"]
)

# Grace period for the server to exit on its own after its stdin closes.
PROCESS_TERMINATION_TIMEOUT = 2.0

# Extra time after SIGTERM before SIGKILL; POSIX only (Windows kills hard).
FORCE_KILL_TIMEOUT = 2.0

# Time for the event loop to observe a kill; only an unkillable process runs this out.
_KILL_REAP_TIMEOUT = 2.0

# Time for the writer to flush accepted messages before stdin closes.
_WRITER_FLUSH_TIMEOUT = 0.5

# How often to poll returncode while waiting for the process to die.
_EXIT_POLL_INTERVAL = 0.01


def get_default_environment() -> dict[str, str]:
    """Returns only the environment variables that are safe to inherit."""
    env: dict[str, str] = {}

    for key in DEFAULT_INHERITED_ENV_VARS:
        value = os.environ.get(key)
        if value is None:  # pragma: lax no cover
            continue

        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
    """Extra environment variables, merged over get_default_environment()."""

    cwd: str | Path | None = None
    """The working directory to use when spawning the process."""

    encoding: str = "utf-8"
    """Text encoding for messages to and from the server."""

    encoding_error_handler: Literal["strict", "ignore", "replace"] = "strict"
    """Encoding error handler; see https://docs.python.org/3/library/codecs.html#error-handlers."""


@asynccontextmanager
async def stdio_client(
    server: StdioServerParameters, errlog: TextIO = sys.stderr
) -> AsyncGenerator[TransportStreams, None]:
    """Spawns an MCP server subprocess and connects to it over stdin/stdout.

    Raises:
        OSError: If the server process cannot be spawned.
        ValueError: If the spawn parameters are invalid (embedded NUL bytes).
    """
    command = await _get_executable_command(server.command)

    process = await _create_platform_compatible_process(
        command=command,
        args=server.args,
        env=get_default_environment() | (server.env or {}),
        errlog=errlog,
        cwd=server.cwd,
    )

    # The spawn succeeded; no awaits until the task group is entered, or a
    # cancellation delivered in the gap would leak the live process.
    read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
    write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)

    shutting_down = False
    writer_done = anyio.Event()

    async def stdout_reader() -> None:
        assert process.stdout, "Opened process is missing stdout"

        stdout = TextReceiveStream(process.stdout, encoding=server.encoding, errors=server.encoding_error_handler)
        try:
            async with read_stream_writer:
                try:
                    # One line at a time; no read-ahead while a delivery is blocked.
                    buffer = ""
                    async for chunk in stdout:
                        lines = (buffer + chunk).split("\n")
                        buffer = lines.pop()
                        for line in lines:
                            try:
                                await read_stream_writer.send(_parse_line(line))
                            except (anyio.ClosedResourceError, anyio.BrokenResourceError):
                                return  # the session is gone; only the drain below remains
                finally:
                    await _drain_stdout(process)
        except anyio.ClosedResourceError:
            pass  # our own shutdown closed the stdout stream under the read
        except (anyio.BrokenResourceError, ConnectionError):
            # Teardown noise during shutdown, a real failure otherwise; either way
            # the session sees clean closure when the read stream closes.
            if not shutting_down:
                logger.exception("Reading from the MCP server's stdout failed mid-session")

    async def stdin_writer() -> None:
        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_unset=True)
                    data = (json + "\n").encode(encoding=server.encoding, errors=server.encoding_error_handler)
                    await process.stdin.send(data)
        except (anyio.ClosedResourceError, anyio.BrokenResourceError, OSError):
            # The server may still be alive: close the read stream so the session
            # sees the connection end instead of a request hanging forever.
            await read_stream_writer.aclose()
        finally:
            writer_done.set()

    async def shutdown() -> None:
        """Winds the transport down: stop traffic, flush, stop the server, release the streams."""
        # Unblock the reader into its drain: a server stuck writing stdout cannot
        # read its stdin, so draining is what lets the flush below complete.
        read_stream.close()
        # Bounded window for the writer to flush already-accepted messages.
        write_stream.close()
        with anyio.move_on_after(_WRITER_FLUSH_TIMEOUT) as flush_scope:
            await writer_done.wait()
        if flush_scope.cancelled_caught:
            await anyio.lowlevel.cancel_shielded_checkpoint()  # resync coverage on 3.11 (gh-106749)
        await _stop_server_process(process)
        await _aclose_all(read_stream, write_stream, read_stream_writer, write_stream_reader)
        # One pass so unblocked tasks exit via their except paths before the cancel.
        await anyio.lowlevel.checkpoint()

    async with anyio.create_task_group() as tg:
        tg.start_soon(stdout_reader)
        tg.start_soon(stdin_writer)
        try:
            yield read_stream, write_stream
        finally:
            shutting_down = True
            # Shutdown must finish even under caller cancellation, or the server
            # process would leak; every wait inside is bounded. (Native
            # task.cancel() and the fallback's worker threads can still defeat it.)
            with anyio.CancelScope(shield=True):
                await shutdown()
            # Unstick pipe tasks a kill survivor's open pipe end could still block.
            tg.cancel_scope.cancel()
    # The cancel lands via throw(); one yield resyncs 3.11 coverage (gh-106749).
    await anyio.lowlevel.cancel_shielded_checkpoint()


def _parse_line(line: str) -> SessionMessage | Exception:
    """Parses one stdout line, returning parse errors as values for the session to surface."""
    try:
        message = types.jsonrpc_message_adapter.validate_json(line, by_name=False)
    except ValueError as exc:
        logger.exception("Failed to parse JSONRPC message from server")
        return exc
    return SessionMessage(message)


async def _drain_stdout(process: ServerProcess) -> None:
    """Consumes and discards the server's remaining stdout.

    Keeps a server flushing buffered output from blocking on a full pipe and
    missing its chance to exit; shielded, raw bytes, ends when shutdown closes
    the pipe.
    """
    assert process.stdout
    with anyio.CancelScope(shield=True):
        with suppress(
            anyio.EndOfStream,
            anyio.ClosedResourceError,
            anyio.BrokenResourceError,
            ConnectionError,
            OSError,
        ):
            while True:
                await process.stdout.receive()


async def _stop_server_process(process: ServerProcess) -> None:
    """Closes stdin, waits out the grace period, then kills the whole tree.

    The escalation order is spec text; timeouts and tree-wide scope are SDK policy:
    https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#shutdown
    """
    assert process.stdin and process.stdout, "server process is spawned with pipes"

    await _close_pipe(process.stdin)
    if not await _wait_for_process_exit(process, PROCESS_TERMINATION_TIMEOUT):
        await _terminate_process_tree(process)
        # Until the event loop observes the death, the transport cannot close.
        if not await _wait_for_process_exit(process, _KILL_REAP_TIMEOUT):
            logger.warning("MCP server process %d is still alive after the kill escalation; abandoning it", process.pid)

    # Reaps surviving Windows job members now, not at GC; no-op on POSIX.
    close_process_job(process)
    # A kill survivor can hold the stdout pipe open; poison the reader anyway.
    await _close_pipe(process.stdout)
    _close_subprocess_transport(process)


async def _close_pipe(stream: AsyncResource) -> None:
    """Closes a pipe stream, tolerating one already closed, broken, or contended."""
    with suppress(OSError, anyio.BrokenResourceError, anyio.ClosedResourceError):
        await stream.aclose()


async def _wait_for_process_exit(process: ServerProcess, timeout: float) -> bool:
    """Returns whether the process died within the timeout, by polling returncode.

    Not process.wait(): on asyncio 3.11+ it also waits for pipe EOF, and a
    child that inherited the pipes makes an exited server look hung.
    """
    deadline = anyio.current_time() + timeout
    while process.returncode is None:
        if anyio.current_time() >= deadline:
            return False
        await anyio.sleep(_EXIT_POLL_INTERVAL)
    return True


async def _terminate_process_tree(process: ServerProcess) -> None:
    """Kills the process and all its descendants.

    POSIX: SIGTERM to the process group, SIGKILL after FORCE_KILL_TIMEOUT.
    Windows: immediate Job Object termination (already a hard kill).
    """
    if sys.platform == "win32":  # pragma: no cover
        await terminate_windows_process_tree(process)
    else:  # pragma: lax no cover
        # The Windows-only FallbackProcess never reaches the POSIX path.
        assert isinstance(process, Process)
        await terminate_posix_process_tree(process, FORCE_KILL_TIMEOUT)


def _close_subprocess_transport(process: ServerProcess) -> None:
    """Closes the asyncio subprocess transport, if there is one.

    The transport otherwise stays open (and warns at GC) while a surviving
    descendant holds a pipe end; nothing public exposes it, hence the attribute
    walk. No-op on trio and the Windows fallback.
    """
    transport = getattr(getattr(process, "_process", None), "_transport", None)
    # Duck-typed: uvloop's UVProcessTransport is not an asyncio.SubprocessTransport.
    close = getattr(transport, "close", None)
    if callable(close):
        # close() on <=3.12 can raise PermissionError re-killing a setuid child.
        with suppress(PermissionError):
            close()


async def _get_executable_command(command: str) -> str:
    """Normalizes the command for the current platform."""
    if sys.platform == "win32":
        return await anyio.to_thread.run_sync(get_windows_executable_command, command, abandon_on_cancel=True)
    else:  # pragma: lax no cover
        return command


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,
) -> ServerProcess:
    """Spawns the server in its own kill scope.

    A new session/process group on POSIX, a Job Object on Windows.
    """
    if sys.platform == "win32":  # pragma: no cover
        return await create_windows_process(command, args, env, errlog, cwd)
    else:  # pragma: lax no cover
        return await anyio.open_process(
            [command, *args],
            env=env,
            stderr=errlog,
            cwd=cwd,
            start_new_session=True,
        )


async def _aclose_all(*streams: AsyncResource) -> None:
    """Closes every given stream."""
    for stream in streams:
        await stream.aclose()