fweinberger/align shutdown with spec (#1091)

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

src/mcp/client/stdio/__init__.py

@@ -186,13 +186,24 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
         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.client.stdio import (
 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 @@ class TestChildProcessCleanup:
                 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 @@ class TestChildProcessCleanup:
 
                 # 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 @@ class TestChildProcessCleanup:
                 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 @@ class TestChildProcessCleanup:
                 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 @@ class TestChildProcessCleanup:
 
                 # 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 @@ class TestChildProcessCleanup:
                 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)."
+    )

tests/issues/test_1027_win_unreachable_cleanup.py

@@ -0,0 +1,243 @@
+"""
+Regression test for issue #1027: Ensure cleanup procedures run properly during shutdown
+
+Issue #1027 reported that cleanup code after "yield" in lifespan was unreachable when
+processes were terminated. This has been fixed by implementing the MCP spec-compliant
+stdio shutdown sequence that closes stdin first, allowing graceful exit.
+
+These tests verify the fix continues to work correctly across all platforms.
+"""
+
+import sys
+import tempfile
+import textwrap
+from pathlib import Path
+
+import anyio
+import pytest
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import _create_platform_compatible_process, stdio_client
+from tests.shared.test_win32_utils import escape_path_for_python
+
+
+@pytest.mark.anyio
+async def test_lifespan_cleanup_executed():
+    """
+    Regression test ensuring MCP server cleanup code runs during shutdown.
+
+    This test verifies that the fix for issue #1027 works correctly by:
+    1. Starting an MCP server that writes a marker file on startup
+    2. Shutting down the server normally via stdio_client
+    3. Verifying the cleanup code (after yield) executed and wrote its marker file
+
+    The fix implements proper stdin closure before termination, giving servers
+    time to run their cleanup handlers.
+    """
+
+    # Create marker files to track server lifecycle
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f:
+        startup_marker = f.name
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f:
+        cleanup_marker = f.name
+
+    # Remove the files so we can detect when they're created
+    Path(startup_marker).unlink()
+    Path(cleanup_marker).unlink()
+
+    # Create a minimal MCP server using FastMCP that tracks lifecycle
+    server_code = textwrap.dedent(f"""
+        import asyncio
+        import sys
+        from pathlib import Path
+        from contextlib import asynccontextmanager
+        from mcp.server.fastmcp import FastMCP
+        
+        STARTUP_MARKER = {escape_path_for_python(startup_marker)}
+        CLEANUP_MARKER = {escape_path_for_python(cleanup_marker)}
+        
+        @asynccontextmanager
+        async def lifespan(server):
+            # Write startup marker
+            Path(STARTUP_MARKER).write_text("started")
+            try:
+                yield {{"started": True}}
+            finally:
+                # This cleanup code now runs properly during shutdown
+                Path(CLEANUP_MARKER).write_text("cleaned up")
+        
+        mcp = FastMCP("test-server", lifespan=lifespan)
+        
+        @mcp.tool()
+        def echo(text: str) -> str:
+            return text
+        
+        if __name__ == "__main__":
+            mcp.run()
+    """)
+
+    # Write the server script to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as f:
+        server_script = f.name
+        f.write(server_code)
+
+    try:
+        # Launch the MCP server
+        params = StdioServerParameters(command=sys.executable, args=[server_script])
+
+        async with stdio_client(params) as (read, write):
+            async with ClientSession(read, write) as session:
+                # Initialize the session
+                result = await session.initialize()
+                assert result.protocolVersion in ["2024-11-05", "2025-06-18"]
+
+                # Verify startup marker was created
+                assert Path(startup_marker).exists(), "Server startup marker not created"
+                assert Path(startup_marker).read_text() == "started"
+
+                # Make a test request to ensure server is working
+                response = await session.call_tool("echo", {"text": "hello"})
+                assert response.content[0].type == "text"
+                assert getattr(response.content[0], "text") == "hello"
+
+                # Session will be closed when exiting the context manager
+
+        # Give server a moment to complete cleanup
+        with anyio.move_on_after(5.0):
+            while not Path(cleanup_marker).exists():
+                await anyio.sleep(0.1)
+
+        # Verify cleanup marker was created - this works now that stdio_client
+        # properly closes stdin before termination, allowing graceful shutdown
+        assert Path(cleanup_marker).exists(), "Server cleanup marker not created - regression in issue #1027 fix"
+        assert Path(cleanup_marker).read_text() == "cleaned up"
+
+    finally:
+        # Clean up files
+        for path in [server_script, startup_marker, cleanup_marker]:
+            try:
+                Path(path).unlink()
+            except FileNotFoundError:
+                pass
+
+
+@pytest.mark.anyio
+@pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default")
+async def test_stdin_close_triggers_cleanup():
+    """
+    Regression test verifying the stdin-based graceful shutdown mechanism.
+
+    This test ensures the core fix for issue #1027 continues to work by:
+    1. Manually managing a server process
+    2. Closing stdin to trigger graceful shutdown
+    3. Verifying cleanup handlers run before the process exits
+
+    This mimics the behavior now implemented in stdio_client's shutdown sequence.
+
+    Note on Windows ResourceWarning:
+    On Windows, we may see ResourceWarning about unclosed file descriptors.
+    This is expected behavior because:
+    - We're manually managing the process lifecycle
+    - Windows file handle cleanup works differently than Unix
+    - The warning doesn't indicate a real issue - cleanup still works
+    We filter this warning on Windows only to avoid test noise.
+    """
+
+    # Create marker files to track server lifecycle
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f:
+        startup_marker = f.name
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f:
+        cleanup_marker = f.name
+
+    # Remove the files so we can detect when they're created
+    Path(startup_marker).unlink()
+    Path(cleanup_marker).unlink()
+
+    # Create an MCP server that handles stdin closure gracefully
+    server_code = textwrap.dedent(f"""
+        import asyncio
+        import sys
+        from pathlib import Path
+        from contextlib import asynccontextmanager
+        from mcp.server.fastmcp import FastMCP
+        
+        STARTUP_MARKER = {escape_path_for_python(startup_marker)}
+        CLEANUP_MARKER = {escape_path_for_python(cleanup_marker)}
+        
+        @asynccontextmanager
+        async def lifespan(server):
+            # Write startup marker
+            Path(STARTUP_MARKER).write_text("started")
+            try:
+                yield {{"started": True}}
+            finally:
+                # This cleanup code runs when stdin closes, enabling graceful shutdown
+                Path(CLEANUP_MARKER).write_text("cleaned up")
+        
+        mcp = FastMCP("test-server", lifespan=lifespan)
+        
+        @mcp.tool()
+        def echo(text: str) -> str:
+            return text
+        
+        if __name__ == "__main__":
+            # The server should exit gracefully when stdin closes
+            try:
+                mcp.run()
+            except Exception:
+                # Server might get EOF or other errors when stdin closes
+                pass
+    """)
+
+    # Write the server script to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as f:
+        server_script = f.name
+        f.write(server_code)
+
+    try:
+        # This test manually manages the process to verify stdin-based shutdown
+        # Start the server process
+        process = await _create_platform_compatible_process(
+            command=sys.executable, args=[server_script], env=None, errlog=sys.stderr, cwd=None
+        )
+
+        # Wait for server to start
+        with anyio.move_on_after(10.0):
+            while not Path(startup_marker).exists():
+                await anyio.sleep(0.1)
+
+        # Check if process is still running
+        if hasattr(process, "returncode") and process.returncode is not None:
+            pytest.fail(f"Server process exited with code {process.returncode}")
+
+        assert Path(startup_marker).exists(), "Server startup marker not created"
+
+        # Close stdin to signal shutdown
+        if process.stdin:
+            await process.stdin.aclose()
+
+        # Wait for process to exit gracefully
+        try:
+            with anyio.fail_after(5.0):  # Increased from 2.0 to 5.0
+                await process.wait()
+        except TimeoutError:
+            # If it doesn't exit after stdin close, terminate it
+            process.terminate()
+            await process.wait()
+
+        # Check if cleanup ran
+        with anyio.move_on_after(5.0):
+            while not Path(cleanup_marker).exists():
+                await anyio.sleep(0.1)
+
+        # Verify the cleanup ran - stdin closure enables graceful shutdown
+        assert Path(cleanup_marker).exists(), "Server cleanup marker not created - stdin-based shutdown failed"
+        assert Path(cleanup_marker).read_text() == "cleaned up"
+
+    finally:
+        # Clean up files
+        for path in [server_script, startup_marker, cleanup_marker]:
+            try:
+                Path(path).unlink()
+            except FileNotFoundError:
+                pass

tests/shared/test_win32_utils.py

@@ -0,0 +1,10 @@
+"""Windows-specific test utilities."""
+
+
+def escape_path_for_python(path: str) -> str:
+    """Escape a file path for use in Python code strings.
+
+    Converts backslashes to forward slashes which work on all platforms
+    and don't need escaping in Python strings.
+    """
+    return repr(path.replace("\\", "/"))