Skip to content

Commit 3abefee

Browse files
fweinberger/align shutdown with spec (#1091)
Co-authored-by: davenpi <davenport.ianc@gmail.com>
1 parent 9b4c2dd commit 3abefee

File tree

4 files changed

+390
-15
lines changed

4 files changed

+390
-15
lines changed

src/mcp/client/stdio/__init__.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -186,13 +186,24 @@ async def stdin_writer():
186186
try:
187187
yield read_stream, write_stream
188188
finally:
189-
# Clean up process to prevent any dangling orphaned processes
189+
# MCP spec: stdio shutdown sequence
190+
# 1. Close input stream to server
191+
# 2. Wait for server to exit, or send SIGTERM if it doesn't exit in time
192+
# 3. Send SIGKILL if still not exited
193+
if process.stdin:
194+
try:
195+
await process.stdin.aclose()
196+
except Exception:
197+
# stdin might already be closed, which is fine
198+
pass
199+
190200
try:
191-
process.terminate()
201+
# Give the process time to exit gracefully after stdin closes
192202
with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT):
193203
await process.wait()
194204
except TimeoutError:
195-
# If process doesn't terminate in time, force kill it
205+
# Process didn't exit from stdin closure, use platform-specific termination
206+
# which handles SIGTERM -> SIGKILL escalation
196207
await _terminate_process_tree(process)
197208
except ProcessLookupError:
198209
# Process already exited, which is fine

tests/client/test_stdio.py

Lines changed: 123 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from mcp.shared.exceptions import McpError
1818
from mcp.shared.message import SessionMessage
1919
from mcp.types import CONNECTION_CLOSED, JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
20+
from tests.shared.test_win32_utils import escape_path_for_python
2021

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

252-
@staticmethod
253-
def _escape_path_for_python(path: str) -> str:
254-
"""Escape a file path for use in Python code strings."""
255-
# Use forward slashes which work on all platforms and don't need escaping
256-
return repr(path.replace("\\", "/"))
257-
258253
@pytest.mark.anyio
259254
@pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default")
260255
async def test_basic_child_process_cleanup(self):
@@ -280,13 +275,13 @@ async def test_basic_child_process_cleanup(self):
280275
import os
281276
282277
# Mark that parent started
283-
with open({self._escape_path_for_python(parent_marker)}, 'w') as f:
278+
with open({escape_path_for_python(parent_marker)}, 'w') as f:
284279
f.write('parent started\\n')
285280
286281
# Child script that writes continuously
287282
child_script = f'''
288283
import time
289-
with open({self._escape_path_for_python(marker_file)}, 'a') as f:
284+
with open({escape_path_for_python(marker_file)}, 'a') as f:
290285
while True:
291286
f.write(f"{time.time()}")
292287
f.flush()
@@ -381,7 +376,7 @@ async def test_nested_process_tree(self):
381376
382377
# Grandchild just writes to file
383378
grandchild_script = \"\"\"import time
384-
with open({self._escape_path_for_python(grandchild_file)}, 'a') as f:
379+
with open({escape_path_for_python(grandchild_file)}, 'a') as f:
385380
while True:
386381
f.write(f"gc {{time.time()}}")
387382
f.flush()
@@ -391,7 +386,7 @@ async def test_nested_process_tree(self):
391386
subprocess.Popen([sys.executable, '-c', grandchild_script])
392387
393388
# Child writes to its file
394-
with open({self._escape_path_for_python(child_file)}, 'a') as f:
389+
with open({escape_path_for_python(child_file)}, 'a') as f:
395390
while True:
396391
f.write(f"c {time.time()}")
397392
f.flush()
@@ -401,7 +396,7 @@ async def test_nested_process_tree(self):
401396
subprocess.Popen([sys.executable, '-c', child_script])
402397
403398
# Parent writes to its file
404-
with open({self._escape_path_for_python(parent_file)}, 'a') as f:
399+
with open({escape_path_for_python(parent_file)}, 'a') as f:
405400
while True:
406401
f.write(f"p {time.time()}")
407402
f.flush()
@@ -470,7 +465,7 @@ async def test_early_parent_exit(self):
470465
471466
# Child that continues running
472467
child_script = f'''import time
473-
with open({self._escape_path_for_python(marker_file)}, 'a') as f:
468+
with open({escape_path_for_python(marker_file)}, 'a') as f:
474469
while True:
475470
f.write(f"child {time.time()}")
476471
f.flush()
@@ -525,3 +520,119 @@ def handle_term(sig, frame):
525520
os.unlink(marker_file)
526521
except OSError:
527522
pass
523+
524+
525+
@pytest.mark.anyio
526+
async def test_stdio_client_graceful_stdin_exit():
527+
"""
528+
Test that a process exits gracefully when stdin is closed,
529+
without needing SIGTERM or SIGKILL.
530+
"""
531+
# Create a Python script that exits when stdin is closed
532+
script_content = textwrap.dedent(
533+
"""
534+
import sys
535+
536+
# Read from stdin until it's closed
537+
try:
538+
while True:
539+
line = sys.stdin.readline()
540+
if not line: # EOF/stdin closed
541+
break
542+
except:
543+
pass
544+
545+
# Exit gracefully
546+
sys.exit(0)
547+
"""
548+
)
549+
550+
server_params = StdioServerParameters(
551+
command=sys.executable,
552+
args=["-c", script_content],
553+
)
554+
555+
start_time = time.time()
556+
557+
# Use anyio timeout to prevent test from hanging forever
558+
with anyio.move_on_after(5.0) as cancel_scope:
559+
async with stdio_client(server_params) as (read_stream, write_stream):
560+
# Let the process start and begin reading stdin
561+
await anyio.sleep(0.2)
562+
# Exit context triggers cleanup - process should exit from stdin closure
563+
pass
564+
565+
if cancel_scope.cancelled_caught:
566+
pytest.fail(
567+
"stdio_client cleanup timed out after 5.0 seconds. "
568+
"Process should have exited gracefully when stdin was closed."
569+
)
570+
571+
end_time = time.time()
572+
elapsed = end_time - start_time
573+
574+
# Should complete quickly with just stdin closure (no signals needed)
575+
assert elapsed < 3.0, (
576+
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-aware process. "
577+
f"Expected < 3.0 seconds since process should exit on stdin closure."
578+
)
579+
580+
581+
@pytest.mark.anyio
582+
async def test_stdio_client_stdin_close_ignored():
583+
"""
584+
Test that when a process ignores stdin closure, the shutdown sequence
585+
properly escalates to SIGTERM.
586+
"""
587+
# Create a Python script that ignores stdin closure but responds to SIGTERM
588+
script_content = textwrap.dedent(
589+
"""
590+
import signal
591+
import sys
592+
import time
593+
594+
# Set up SIGTERM handler to exit cleanly
595+
def sigterm_handler(signum, frame):
596+
sys.exit(0)
597+
598+
signal.signal(signal.SIGTERM, sigterm_handler)
599+
600+
# Close stdin immediately to simulate ignoring it
601+
sys.stdin.close()
602+
603+
# Keep running until SIGTERM
604+
while True:
605+
time.sleep(0.1)
606+
"""
607+
)
608+
609+
server_params = StdioServerParameters(
610+
command=sys.executable,
611+
args=["-c", script_content],
612+
)
613+
614+
start_time = time.time()
615+
616+
# Use anyio timeout to prevent test from hanging forever
617+
with anyio.move_on_after(7.0) as cancel_scope:
618+
async with stdio_client(server_params) as (read_stream, write_stream):
619+
# Let the process start
620+
await anyio.sleep(0.2)
621+
# Exit context triggers cleanup
622+
pass
623+
624+
if cancel_scope.cancelled_caught:
625+
pytest.fail(
626+
"stdio_client cleanup timed out after 7.0 seconds. "
627+
"Process should have been terminated via SIGTERM escalation."
628+
)
629+
630+
end_time = time.time()
631+
elapsed = end_time - start_time
632+
633+
# Should take ~2 seconds (stdin close timeout) before SIGTERM is sent
634+
# Total time should be between 2-4 seconds
635+
assert 1.5 < elapsed < 4.5, (
636+
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-ignoring process. "
637+
f"Expected between 2-4 seconds (2s stdin timeout + termination time)."
638+
)

0 commit comments

Comments
 (0)