Fix stdio client shutdown bugs and rebuild the stdio test suite (#2773)
This commit is contained in:
parent
19fe9faec8
commit
bdc48e98b1
3 changed files with 2645 additions and 1130 deletions
|
|
@ -1,21 +1,33 @@
|
|||
"""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 contextlib import asynccontextmanager
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager, suppress
|
||||
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.abc import AsyncResource, Process
|
||||
from anyio.streams.text import TextReceiveStream
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from mcp import types
|
||||
from mcp.client._transport import TransportStreams
|
||||
from mcp.os.posix.utilities import terminate_posix_process_tree
|
||||
from mcp.os.win32.utilities import (
|
||||
FallbackProcess,
|
||||
ServerProcess,
|
||||
close_process_job,
|
||||
create_windows_process,
|
||||
get_windows_executable_command,
|
||||
terminate_windows_process_tree,
|
||||
|
|
@ -44,14 +56,24 @@ DEFAULT_INHERITED_ENV_VARS = (
|
|||
else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"]
|
||||
)
|
||||
|
||||
# Timeout for process termination before falling back to force kill
|
||||
# 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 a default environment object including only environment variables deemed
|
||||
safe to inherit.
|
||||
"""
|
||||
"""Returns only the environment variables that are safe to inherit."""
|
||||
env: dict[str, str] = {}
|
||||
|
||||
for key in DEFAULT_INHERITED_ENV_VARS:
|
||||
|
|
@ -76,150 +98,227 @@ class StdioServerParameters(BaseModel):
|
|||
"""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.
|
||||
"""
|
||||
"""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"
|
||||
"""
|
||||
The text encoding used when sending/receiving messages to the server.
|
||||
|
||||
Defaults to utf-8.
|
||||
"""
|
||||
"""Text encoding for messages to and from the server."""
|
||||
|
||||
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.
|
||||
"""
|
||||
"""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):
|
||||
"""Client transport for stdio: this will connect to a server by spawning a
|
||||
process and communicating with it over stdin/stdout.
|
||||
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).
|
||||
"""
|
||||
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
|
||||
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
|
||||
command = _get_executable_command(server.command)
|
||||
|
||||
write_stream: MemoryObjectSendStream[SessionMessage]
|
||||
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
|
||||
process = await _create_platform_compatible_process(
|
||||
command=command,
|
||||
args=server.args,
|
||||
env=get_default_environment() | (server.env or {}),
|
||||
errlog=errlog,
|
||||
cwd=server.cwd,
|
||||
)
|
||||
|
||||
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
|
||||
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
|
||||
# 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)
|
||||
|
||||
try:
|
||||
command = _get_executable_command(server.command)
|
||||
shutting_down = False
|
||||
writer_done = anyio.Event()
|
||||
|
||||
# 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():
|
||||
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:
|
||||
buffer = ""
|
||||
async for chunk in TextReceiveStream(
|
||||
process.stdout,
|
||||
encoding=server.encoding,
|
||||
errors=server.encoding_error_handler,
|
||||
):
|
||||
lines = (buffer + chunk).split("\n")
|
||||
buffer = lines.pop()
|
||||
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")
|
||||
|
||||
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: lax no cover
|
||||
await anyio.lowlevel.checkpoint()
|
||||
|
||||
async def stdin_writer():
|
||||
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)
|
||||
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()
|
||||
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 with anyio.create_task_group() as tg, process:
|
||||
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:
|
||||
# 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
|
||||
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()
|
||||
|
||||
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 _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()
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
"""Normalizes the command for the current platform."""
|
||||
if sys.platform == "win32": # pragma: no cover
|
||||
return get_windows_executable_command(command)
|
||||
else: # pragma: lax no cover
|
||||
|
|
@ -232,16 +331,15 @@ async def _create_platform_compatible_process(
|
|||
env: dict[str, str] | None = None,
|
||||
errlog: TextIO = sys.stderr,
|
||||
cwd: Path | str | None = None,
|
||||
):
|
||||
"""Creates a subprocess in a platform-compatible way.
|
||||
) -> ServerProcess:
|
||||
"""Spawns the server in its own kill scope.
|
||||
|
||||
Unix: Creates process in a new session/process group for killpg support
|
||||
Windows: Creates process in a Job Object for reliable child termination
|
||||
A new session/process group on POSIX, a Job Object on Windows.
|
||||
"""
|
||||
if sys.platform == "win32": # pragma: no cover
|
||||
process = await create_windows_process(command, args, env, errlog, cwd)
|
||||
return await create_windows_process(command, args, env, errlog, cwd)
|
||||
else: # pragma: lax no cover
|
||||
process = await anyio.open_process(
|
||||
return await anyio.open_process(
|
||||
[command, *args],
|
||||
env=env,
|
||||
stderr=errlog,
|
||||
|
|
@ -249,22 +347,8 @@ async def _create_platform_compatible_process(
|
|||
start_new_session=True,
|
||||
)
|
||||
|
||||
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: lax no cover
|
||||
# FallbackProcess should only be used for Windows compatibility
|
||||
assert isinstance(process, Process)
|
||||
await terminate_posix_process_tree(process, timeout_seconds)
|
||||
async def _aclose_all(*streams: AsyncResource) -> None:
|
||||
"""Closes every given stream."""
|
||||
for stream in streams:
|
||||
await stream.aclose()
|
||||
|
|
|
|||
|
|
@ -3,55 +3,61 @@
|
|||
import logging
|
||||
import os
|
||||
import signal
|
||||
from contextlib import suppress
|
||||
|
||||
import anyio
|
||||
from anyio.abc import Process
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# How often to probe for surviving group members between SIGTERM and SIGKILL.
|
||||
_GROUP_POLL_INTERVAL = 0.01
|
||||
|
||||
|
||||
async def terminate_posix_process_tree(process: Process, timeout_seconds: float = 2.0) -> None:
|
||||
"""Terminate a process and all its children on POSIX systems.
|
||||
"""Terminates a process and all its descendants on POSIX.
|
||||
|
||||
Uses os.killpg() for atomic process group termination.
|
||||
|
||||
Args:
|
||||
process: The process to terminate
|
||||
timeout_seconds: Timeout in seconds before force killing (default: 2.0)
|
||||
SIGTERMs the process group, waits up to timeout_seconds for it to
|
||||
disappear, then SIGKILLs whatever remains. killpg reaches every descendant
|
||||
atomically, even ones whose parent already exited; daemonizers that left
|
||||
the group escape by design. A group only disappears once every member is
|
||||
dead and reaped, so a client running as PID 1 should reap orphans (e.g.
|
||||
docker run --init) or the wait below runs its full timeout.
|
||||
"""
|
||||
pid = getattr(process, "pid", None) or getattr(getattr(process, "popen", None), "pid", None)
|
||||
if not pid:
|
||||
# No PID means there's no process to terminate - it either never started,
|
||||
# already exited, or we have an invalid process object
|
||||
return
|
||||
# The leader's pid is the pgid (start_new_session). Never use getpgid():
|
||||
# it fails once the leader is reaped, even with live members left.
|
||||
pgid = process.pid
|
||||
|
||||
try:
|
||||
pgid = os.getpgid(pid)
|
||||
os.killpg(pgid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
return # the whole group is already gone
|
||||
except PermissionError:
|
||||
# EPERM never proves the group is gone (macOS raises it for zombie or
|
||||
# foreign-euid members), so keep waiting and escalating.
|
||||
logger.warning(
|
||||
"No permission to signal some of process group %d; waiting for it to exit anyway", pgid, exc_info=True
|
||||
)
|
||||
|
||||
with anyio.move_on_after(timeout_seconds):
|
||||
while True:
|
||||
try:
|
||||
# Check if process group still exists (signal 0 = check only)
|
||||
os.killpg(pgid, 0)
|
||||
await anyio.sleep(0.1)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
with anyio.move_on_after(timeout_seconds):
|
||||
while _group_alive(pgid):
|
||||
# Reading returncode reaps the leader on trio; a zombie leader would
|
||||
# otherwise keep the group alive for the full timeout.
|
||||
_ = process.returncode
|
||||
await anyio.sleep(_GROUP_POLL_INTERVAL)
|
||||
return
|
||||
|
||||
try:
|
||||
os.killpg(pgid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
# ESRCH: died since the last probe. EPERM: we killed what we were allowed to.
|
||||
with suppress(ProcessLookupError, PermissionError):
|
||||
os.killpg(pgid, signal.SIGKILL)
|
||||
|
||||
except (ProcessLookupError, PermissionError, OSError) as e:
|
||||
logger.warning(f"Process group termination failed for PID {pid}: {e}, falling back to simple terminate")
|
||||
try:
|
||||
process.terminate()
|
||||
with anyio.fail_after(timeout_seconds):
|
||||
await process.wait()
|
||||
except Exception:
|
||||
logger.warning(f"Process termination failed for PID {pid}, attempting force kill")
|
||||
try:
|
||||
process.kill()
|
||||
except Exception:
|
||||
logger.exception(f"Failed to kill process {pid}")
|
||||
|
||||
def _group_alive(pgid: int) -> bool:
|
||||
"""Probes the group with signal 0; only ESRCH proves it is gone."""
|
||||
try:
|
||||
os.killpg(pgid, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
pass # unsignalable survivors or unreaped zombies; EPERM is ambiguous
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -4,16 +4,16 @@ import logging
|
|||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import weakref
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO, TextIO, cast
|
||||
from typing import BinaryIO, TextIO, TypeAlias, cast
|
||||
|
||||
import anyio
|
||||
from anyio import to_thread
|
||||
from anyio.abc import Process
|
||||
from anyio.streams.file import FileReadStream, FileWriteStream
|
||||
from typing_extensions import deprecated
|
||||
|
||||
logger = logging.getLogger("client.stdio.win32")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Windows-specific imports for Job Objects
|
||||
if sys.platform == "win32":
|
||||
|
|
@ -28,110 +28,86 @@ else:
|
|||
win32job = None
|
||||
pywintypes = None
|
||||
|
||||
JobHandle = int
|
||||
# How often FallbackProcess polls the underlying Popen for exit.
|
||||
_EXIT_POLL_INTERVAL = 0.01
|
||||
|
||||
# Job Object handle per spawned process, for tree termination at shutdown.
|
||||
# Values stay pywin32 PyHANDLEs: if no pop site ever runs, the dying weak entry
|
||||
# drops the last reference and the PyHANDLE destructor closes the handle, which
|
||||
# is what makes KILL_ON_JOB_CLOSE reap an abandoned tree.
|
||||
_process_jobs: "weakref.WeakKeyDictionary[Process | FallbackProcess, object]" = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
def get_windows_executable_command(command: str) -> str:
|
||||
"""Get the correct executable command normalized for Windows.
|
||||
"""Resolves the command to a Windows executable path.
|
||||
|
||||
On Windows, commands might exist with specific extensions (.exe, .cmd, etc.)
|
||||
that need to be located for proper execution.
|
||||
|
||||
Args:
|
||||
command: Base command (e.g., 'uvx', 'npx')
|
||||
|
||||
Returns:
|
||||
str: Windows-appropriate command path
|
||||
Tries the bare name first, then the common script extensions (.cmd, .bat,
|
||||
.exe, .ps1).
|
||||
"""
|
||||
try:
|
||||
# First check if command exists in PATH as-is
|
||||
if command_path := shutil.which(command):
|
||||
return command_path
|
||||
|
||||
# Check for Windows-specific extensions
|
||||
for ext in [".cmd", ".bat", ".exe", ".ps1"]:
|
||||
ext_version = f"{command}{ext}"
|
||||
if ext_path := shutil.which(ext_version):
|
||||
return ext_path
|
||||
|
||||
# For regular commands or if we couldn't find special versions
|
||||
return command
|
||||
except OSError:
|
||||
# Handle file system errors during path resolution
|
||||
# (permissions, broken symlinks, etc.)
|
||||
return command
|
||||
return command # path probing failed (permissions, broken symlinks)
|
||||
|
||||
|
||||
class FallbackProcess:
|
||||
"""A fallback process wrapper for Windows to handle async I/O
|
||||
when using subprocess.Popen, which provides sync-only FileIO objects.
|
||||
"""Async wrapper around subprocess.Popen for SelectorEventLoop.
|
||||
|
||||
This wraps stdin and stdout into async-compatible
|
||||
streams (FileReadStream, FileWriteStream),
|
||||
so that MCP clients expecting async streams can work properly.
|
||||
Windows event loops without async subprocess support get this Popen-backed
|
||||
fallback, with anyio file streams wrapping the pipes.
|
||||
"""
|
||||
|
||||
def __init__(self, popen_obj: subprocess.Popen[bytes]):
|
||||
def __init__(self, popen_obj: subprocess.Popen[bytes]) -> None:
|
||||
self.popen: subprocess.Popen[bytes] = popen_obj
|
||||
self.stdin_raw = popen_obj.stdin # type: ignore[assignment]
|
||||
self.stdout_raw = popen_obj.stdout # type: ignore[assignment]
|
||||
self.stderr = popen_obj.stderr # type: ignore[assignment]
|
||||
stdin = popen_obj.stdin
|
||||
stdout = popen_obj.stdout
|
||||
|
||||
self.stdin = FileWriteStream(cast(BinaryIO, self.stdin_raw)) if self.stdin_raw else None
|
||||
self.stdout = FileReadStream(cast(BinaryIO, self.stdout_raw)) if self.stdout_raw else None
|
||||
self.stdin = FileWriteStream(cast(BinaryIO, stdin)) if stdin else None
|
||||
self.stdout = FileReadStream(cast(BinaryIO, stdout)) if stdout else None
|
||||
|
||||
async def __aenter__(self):
|
||||
"""Support async context manager entry."""
|
||||
return self
|
||||
async def wait(self) -> int:
|
||||
"""Waits for exit by polling the Popen.
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: BaseException | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: object | None,
|
||||
) -> None:
|
||||
"""Terminate and wait on process exit inside a thread."""
|
||||
A thread blocked in Popen.wait() cannot be cancelled by anyio, which
|
||||
would defeat every timeout placed around this call.
|
||||
"""
|
||||
while (returncode := self.popen.poll()) is None:
|
||||
await anyio.sleep(_EXIT_POLL_INTERVAL)
|
||||
return returncode
|
||||
|
||||
def terminate(self) -> None:
|
||||
"""Terminates the subprocess."""
|
||||
self.popen.terminate()
|
||||
await to_thread.run_sync(self.popen.wait)
|
||||
|
||||
# Close the file handles to prevent ResourceWarning
|
||||
if self.stdin:
|
||||
await self.stdin.aclose()
|
||||
if self.stdout:
|
||||
await self.stdout.aclose()
|
||||
if self.stdin_raw:
|
||||
self.stdin_raw.close()
|
||||
if self.stdout_raw:
|
||||
self.stdout_raw.close()
|
||||
if self.stderr:
|
||||
self.stderr.close()
|
||||
|
||||
async def wait(self):
|
||||
"""Async wait for process completion."""
|
||||
return await to_thread.run_sync(self.popen.wait)
|
||||
|
||||
def terminate(self):
|
||||
"""Terminate the subprocess immediately."""
|
||||
return self.popen.terminate()
|
||||
|
||||
def kill(self) -> None:
|
||||
"""Kill the subprocess immediately (alias for terminate)."""
|
||||
self.terminate()
|
||||
"""Kills the subprocess (on Windows the same hard kill as terminate)."""
|
||||
self.popen.kill()
|
||||
|
||||
@property
|
||||
def pid(self) -> int:
|
||||
"""Return the process ID."""
|
||||
"""Returns the process ID."""
|
||||
return self.popen.pid
|
||||
|
||||
@property
|
||||
def returncode(self) -> int | None:
|
||||
"""Return the exit code, or ``None`` if the process has not yet terminated."""
|
||||
return self.popen.returncode
|
||||
"""The exit code, or None while the process is still running.
|
||||
|
||||
Polls the Popen so death is observable without anyone calling wait().
|
||||
"""
|
||||
return self.popen.poll()
|
||||
|
||||
|
||||
# ------------------------
|
||||
# Updated function
|
||||
# ------------------------
|
||||
# The process handle stdio_client drives: anyio's Process, or the Popen-backed
|
||||
# fallback used on Windows event loops without async subprocess support.
|
||||
ServerProcess: TypeAlias = Process | FallbackProcess
|
||||
|
||||
|
||||
async def create_windows_process(
|
||||
|
|
@ -141,53 +117,35 @@ async def create_windows_process(
|
|||
errlog: TextIO | None = sys.stderr,
|
||||
cwd: Path | str | None = None,
|
||||
) -> Process | FallbackProcess:
|
||||
"""Creates a subprocess in a Windows-compatible way with Job Object support.
|
||||
"""Creates a subprocess with Job Object support for tree termination.
|
||||
|
||||
Attempts to use anyio's open_process for async subprocess creation.
|
||||
In some cases this will throw NotImplementedError on Windows, e.g.,
|
||||
when using the SelectorEventLoop, which does not support async subprocesses.
|
||||
In that case, we fall back to using subprocess.Popen.
|
||||
|
||||
The process is automatically added to a Job Object to ensure all child
|
||||
processes are terminated when the parent is terminated.
|
||||
|
||||
Args:
|
||||
command (str): The executable to run
|
||||
args (list[str]): List of command line arguments
|
||||
env (dict[str, str] | None): Environment variables
|
||||
errlog (TextIO | None): Where to send stderr output (defaults to sys.stderr)
|
||||
cwd (Path | str | None): Working directory for the subprocess
|
||||
Spawns via anyio's open_process; event loops without async subprocess
|
||||
support (notably the SelectorEventLoop) raise NotImplementedError, in which
|
||||
case the spawn falls back to a Popen-backed FallbackProcess. Either way the
|
||||
process is then assigned to a Job Object so its children can be terminated
|
||||
with it; children spawned before the assignment completes are not captured
|
||||
(see the inline note below).
|
||||
|
||||
Returns:
|
||||
Process | FallbackProcess: Async-compatible subprocess with stdin and stdout streams
|
||||
Process | FallbackProcess: The spawned process with async stdin/stdout streams.
|
||||
"""
|
||||
job = _create_job_object()
|
||||
process = None
|
||||
|
||||
try:
|
||||
# First try using anyio with Windows-specific flags to hide console window
|
||||
process = await anyio.open_process(
|
||||
[command, *args],
|
||||
env=env,
|
||||
# Ensure we don't create console windows for each process
|
||||
creationflags=subprocess.CREATE_NO_WINDOW # type: ignore
|
||||
if hasattr(subprocess, "CREATE_NO_WINDOW")
|
||||
else 0,
|
||||
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
||||
stderr=errlog,
|
||||
cwd=cwd,
|
||||
)
|
||||
except NotImplementedError:
|
||||
# If Windows doesn't support async subprocess creation, use fallback
|
||||
# Windows event loops without async subprocess support (SelectorEventLoop)
|
||||
process = await _create_windows_fallback_process(command, args, env, errlog, cwd)
|
||||
except Exception:
|
||||
# Try again without creation flags
|
||||
process = await anyio.open_process(
|
||||
[command, *args],
|
||||
env=env,
|
||||
stderr=errlog,
|
||||
cwd=cwd,
|
||||
)
|
||||
|
||||
# Children spawned before the assignment completes land outside the job
|
||||
# (membership is inherited at CreateProcess, never acquired retroactively);
|
||||
# if that ever bites, the fix is a CREATE_SUSPENDED spawn -> assign -> resume.
|
||||
job = _create_job_object()
|
||||
_maybe_assign_process_to_job(process, job)
|
||||
return process
|
||||
|
||||
|
|
@ -199,41 +157,26 @@ async def _create_windows_fallback_process(
|
|||
errlog: TextIO | None = sys.stderr,
|
||||
cwd: Path | str | None = None,
|
||||
) -> FallbackProcess:
|
||||
"""Create a subprocess using subprocess.Popen as a fallback when anyio fails.
|
||||
|
||||
This function wraps the sync subprocess.Popen in an async-compatible interface.
|
||||
"""
|
||||
try:
|
||||
# Try launching with creationflags to avoid opening a new console window
|
||||
popen_obj = subprocess.Popen(
|
||||
[command, *args],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=errlog,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
bufsize=0, # Unbuffered output
|
||||
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
||||
)
|
||||
except Exception:
|
||||
# If creationflags failed, fallback without them
|
||||
popen_obj = subprocess.Popen(
|
||||
[command, *args],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=errlog,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
bufsize=0,
|
||||
)
|
||||
"""Spawns via subprocess.Popen and wraps it in FallbackProcess."""
|
||||
popen_obj = subprocess.Popen(
|
||||
[command, *args],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=errlog,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
bufsize=0, # Unbuffered output
|
||||
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
||||
)
|
||||
return FallbackProcess(popen_obj)
|
||||
|
||||
|
||||
def _create_job_object() -> int | None:
|
||||
"""Create a Windows Job Object configured to terminate all processes when closed."""
|
||||
if sys.platform != "win32" or not win32job:
|
||||
def _create_job_object() -> object | None:
|
||||
"""Creates a Windows Job Object configured to terminate all its processes when closed."""
|
||||
if sys.platform != "win32" or not win32api or not win32job:
|
||||
return None
|
||||
|
||||
job = None
|
||||
try:
|
||||
job = win32job.CreateJobObject(None, "")
|
||||
extended_info = win32job.QueryInformationJobObject(job, win32job.JobObjectExtendedLimitInformation)
|
||||
|
|
@ -241,17 +184,20 @@ def _create_job_object() -> int | None:
|
|||
extended_info["BasicLimitInformation"]["LimitFlags"] |= win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
||||
win32job.SetInformationJobObject(job, win32job.JobObjectExtendedLimitInformation, extended_info)
|
||||
return job
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to create Job Object for process tree management: {e}")
|
||||
except pywintypes.error:
|
||||
logger.warning("Failed to create Job Object for process tree management", exc_info=True)
|
||||
# If creation succeeded but configuration failed, close the handle now.
|
||||
if job is not None:
|
||||
_close_job_handle(job)
|
||||
return None
|
||||
|
||||
|
||||
def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: JobHandle | None) -> None:
|
||||
"""Try to assign a process to a job object.
|
||||
def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: object | None) -> None:
|
||||
"""Assigns the process to the job and records it for tree termination.
|
||||
|
||||
If assignment fails for any reason, the job handle is closed.
|
||||
On any failure the job handle is closed instead.
|
||||
"""
|
||||
if not job:
|
||||
if job is None:
|
||||
return
|
||||
|
||||
if sys.platform != "win32" or not win32api or not win32con or not win32job:
|
||||
|
|
@ -262,72 +208,62 @@ def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: JobHan
|
|||
win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid
|
||||
)
|
||||
if not process_handle:
|
||||
raise Exception("Failed to open process handle")
|
||||
raise pywintypes.error(0, "OpenProcess", "Failed to open process handle")
|
||||
|
||||
try:
|
||||
win32job.AssignProcessToJobObject(job, process_handle)
|
||||
process._job_object = job
|
||||
finally:
|
||||
win32api.CloseHandle(process_handle)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to assign process {process.pid} to Job Object: {e}")
|
||||
if win32api:
|
||||
win32api.CloseHandle(job)
|
||||
# Record only after the CloseHandle above succeeded: had it failed, the
|
||||
# except below would close the job and KILL_ON_JOB_CLOSE takes the server.
|
||||
_process_jobs[process] = job
|
||||
except pywintypes.error:
|
||||
logger.warning("Failed to assign process %d to Job Object", process.pid, exc_info=True)
|
||||
_close_job_handle(job)
|
||||
|
||||
|
||||
async def terminate_windows_process_tree(process: Process | FallbackProcess, timeout_seconds: float = 2.0) -> None:
|
||||
"""Terminate a process and all its children on Windows.
|
||||
def close_process_job(process: Process | FallbackProcess) -> None:
|
||||
"""Closes the process's Job Object handle, if it still has one.
|
||||
|
||||
If the process has an associated job object, it will be terminated.
|
||||
Otherwise, falls back to basic process termination.
|
||||
|
||||
Args:
|
||||
process: The process to terminate
|
||||
timeout_seconds: Timeout in seconds before force killing (default: 2.0)
|
||||
KILL_ON_JOB_CLOSE makes the close also kill any members still alive,
|
||||
deterministically rather than at GC time; a deliberate divergence from
|
||||
POSIX, where a graceful server's children are left alive.
|
||||
"""
|
||||
if sys.platform != "win32":
|
||||
return
|
||||
|
||||
job = getattr(process, "_job_object", None)
|
||||
if job and win32job:
|
||||
try:
|
||||
win32job.TerminateJobObject(job, 1)
|
||||
except Exception:
|
||||
# Job might already be terminated
|
||||
pass
|
||||
finally:
|
||||
if win32api:
|
||||
try:
|
||||
win32api.CloseHandle(job)
|
||||
except Exception:
|
||||
pass
|
||||
job = _process_jobs.pop(process, None)
|
||||
if job is not None:
|
||||
_close_job_handle(job)
|
||||
|
||||
# Always try to terminate the process itself as well
|
||||
|
||||
async def terminate_windows_process_tree(process: Process | FallbackProcess) -> None:
|
||||
"""Terminates the process's job, or just the process if it has no job.
|
||||
|
||||
Job termination is an immediate hard kill of every member. Windows has no
|
||||
tree-wide SIGTERM; the stdin-close grace period is the server's chance to
|
||||
exit cleanly.
|
||||
"""
|
||||
if sys.platform != "win32":
|
||||
return
|
||||
|
||||
job = _process_jobs.pop(process, None)
|
||||
if job is not None and win32job:
|
||||
try:
|
||||
with suppress(pywintypes.error): # the job might already be terminated
|
||||
win32job.TerminateJobObject(job, 1)
|
||||
finally:
|
||||
_close_job_handle(job)
|
||||
|
||||
# The process may have no job (creation or assignment failed); kill it directly too.
|
||||
try:
|
||||
process.terminate()
|
||||
except Exception:
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
@deprecated(
|
||||
"terminate_windows_process is deprecated and will be removed in a future version. "
|
||||
"Process termination is now handled internally by the stdio_client context manager."
|
||||
)
|
||||
async def terminate_windows_process(process: Process | FallbackProcess):
|
||||
"""Terminate a Windows process.
|
||||
|
||||
Note: On Windows, terminating a process with process.terminate() doesn't
|
||||
always guarantee immediate process termination.
|
||||
If the process does not exit within 2 seconds, process.kill() is called
|
||||
to send a SIGKILL-equivalent signal.
|
||||
|
||||
Args:
|
||||
process: The process to terminate
|
||||
"""
|
||||
try:
|
||||
process.terminate()
|
||||
with anyio.fail_after(2.0):
|
||||
await process.wait()
|
||||
except TimeoutError:
|
||||
# Force kill if it doesn't terminate
|
||||
process.kill()
|
||||
def _close_job_handle(job: object) -> None:
|
||||
"""Closes a Job Object handle, tolerating one that is already closed."""
|
||||
if win32api and pywintypes:
|
||||
with suppress(pywintypes.error):
|
||||
win32api.CloseHandle(job)
|
||||
|
|
|
|||
Loading…
Reference in a new issue