fweinberger/align shutdown with spec (#1091)

Co-authored-by: davenpi <davenport.ianc@gmail.com>

Author: felixweinberger

src/mcp/client/stdio/__init__.py

@@ -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

tests/client/test_stdio.py

@@ -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
@@ -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):
@@ -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()
@@ -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()
@@ -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()
@@ -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()
@@ -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()
@@ -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)."
+    )