Skip to content

fweinberger/align shutdown with spec #1091

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 9, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions src/mcp/client/stdio/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,13 +186,24 @@ async def stdin_writer():
try:
yield read_stream, write_stream
finally:
# Clean up process to prevent any dangling orphaned processes
# 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:
try:
await process.stdin.aclose()
except Exception:
# stdin might already be closed, which is fine
pass

try:
process.terminate()
# Give the process time to exit gracefully after stdin closes
with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT):
await process.wait()
except TimeoutError:
# If process doesn't terminate in time, force kill it
# Process didn't exit from stdin closure, use platform-specific termination
# which handles SIGTERM -> SIGKILL escalation
await _terminate_process_tree(process)
except ProcessLookupError:
# Process already exited, which is fine
Expand Down
135 changes: 123 additions & 12 deletions tests/client/test_stdio.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from mcp.shared.exceptions import McpError
from mcp.shared.message import SessionMessage
from mcp.types import CONNECTION_CLOSED, JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
from tests.shared.test_win32_utils import escape_path_for_python

# Timeout for cleanup of processes that ignore SIGTERM
# This timeout ensures the test fails quickly if the cleanup logic doesn't have
Expand Down Expand Up @@ -249,12 +250,6 @@ class TestChildProcessCleanup:
This is a fundamental difference between Windows and Unix process termination.
"""

@staticmethod
def _escape_path_for_python(path: str) -> str:
"""Escape a file path for use in Python code strings."""
# Use forward slashes which work on all platforms and don't need escaping
return repr(path.replace("\\", "/"))

@pytest.mark.anyio
@pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default")
async def test_basic_child_process_cleanup(self):
Expand All @@ -280,13 +275,13 @@ async def test_basic_child_process_cleanup(self):
import os

# Mark that parent started
with open({self._escape_path_for_python(parent_marker)}, 'w') as f:
with open({escape_path_for_python(parent_marker)}, 'w') as f:
f.write('parent started\\n')

# Child script that writes continuously
child_script = f'''
import time
with open({self._escape_path_for_python(marker_file)}, 'a') as f:
with open({escape_path_for_python(marker_file)}, 'a') as f:
while True:
f.write(f"{time.time()}")
f.flush()
Expand Down Expand Up @@ -381,7 +376,7 @@ async def test_nested_process_tree(self):

# Grandchild just writes to file
grandchild_script = \"\"\"import time
with open({self._escape_path_for_python(grandchild_file)}, 'a') as f:
with open({escape_path_for_python(grandchild_file)}, 'a') as f:
while True:
f.write(f"gc {{time.time()}}")
f.flush()
Expand All @@ -391,7 +386,7 @@ async def test_nested_process_tree(self):
subprocess.Popen([sys.executable, '-c', grandchild_script])

# Child writes to its file
with open({self._escape_path_for_python(child_file)}, 'a') as f:
with open({escape_path_for_python(child_file)}, 'a') as f:
while True:
f.write(f"c {time.time()}")
f.flush()
Expand All @@ -401,7 +396,7 @@ async def test_nested_process_tree(self):
subprocess.Popen([sys.executable, '-c', child_script])

# Parent writes to its file
with open({self._escape_path_for_python(parent_file)}, 'a') as f:
with open({escape_path_for_python(parent_file)}, 'a') as f:
while True:
f.write(f"p {time.time()}")
f.flush()
Expand Down Expand Up @@ -470,7 +465,7 @@ async def test_early_parent_exit(self):

# Child that continues running
child_script = f'''import time
with open({self._escape_path_for_python(marker_file)}, 'a') as f:
with open({escape_path_for_python(marker_file)}, 'a') as f:
while True:
f.write(f"child {time.time()}")
f.flush()
Expand Down Expand Up @@ -525,3 +520,119 @@ def handle_term(sig, frame):
os.unlink(marker_file)
except OSError:
pass


@pytest.mark.anyio
async def test_stdio_client_graceful_stdin_exit():
"""
Test that a process exits gracefully when stdin is closed,
without needing SIGTERM or SIGKILL.
"""
# Create a Python script that exits when stdin is closed
script_content = textwrap.dedent(
"""
import sys

# Read from stdin until it's closed
try:
while True:
line = sys.stdin.readline()
if not line: # EOF/stdin closed
break
except:
pass

# Exit gracefully
sys.exit(0)
"""
)

server_params = StdioServerParameters(
command=sys.executable,
args=["-c", script_content],
)

start_time = time.time()

# Use anyio timeout to prevent test from hanging forever
with anyio.move_on_after(5.0) as cancel_scope:
async with stdio_client(server_params) as (read_stream, write_stream):
# Let the process start and begin reading stdin
await anyio.sleep(0.2)
# Exit context triggers cleanup - process should exit from stdin closure
pass

if cancel_scope.cancelled_caught:
pytest.fail(
"stdio_client cleanup timed out after 5.0 seconds. "
"Process should have exited gracefully when stdin was closed."
)

end_time = time.time()
elapsed = end_time - start_time

# Should complete quickly with just stdin closure (no signals needed)
assert elapsed < 3.0, (
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-aware process. "
f"Expected < 3.0 seconds since process should exit on stdin closure."
)


@pytest.mark.anyio
async def test_stdio_client_stdin_close_ignored():
"""
Test that when a process ignores stdin closure, the shutdown sequence
properly escalates to SIGTERM.
"""
# Create a Python script that ignores stdin closure but responds to SIGTERM
script_content = textwrap.dedent(
"""
import signal
import sys
import time

# Set up SIGTERM handler to exit cleanly
def sigterm_handler(signum, frame):
sys.exit(0)

signal.signal(signal.SIGTERM, sigterm_handler)

# Close stdin immediately to simulate ignoring it
sys.stdin.close()

# Keep running until SIGTERM
while True:
time.sleep(0.1)
"""
)

server_params = StdioServerParameters(
command=sys.executable,
args=["-c", script_content],
)

start_time = time.time()

# Use anyio timeout to prevent test from hanging forever
with anyio.move_on_after(7.0) as cancel_scope:
async with stdio_client(server_params) as (read_stream, write_stream):
# Let the process start
await anyio.sleep(0.2)
# Exit context triggers cleanup
pass

if cancel_scope.cancelled_caught:
pytest.fail(
"stdio_client cleanup timed out after 7.0 seconds. "
"Process should have been terminated via SIGTERM escalation."
)

end_time = time.time()
elapsed = end_time - start_time

# Should take ~2 seconds (stdin close timeout) before SIGTERM is sent
# Total time should be between 2-4 seconds
assert 1.5 < elapsed < 4.5, (
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-ignoring process. "
f"Expected between 2-4 seconds (2s stdin timeout + termination time)."
)
Loading
Loading