|
| 1 | +""" |
| 2 | +Regression test for issue #1027: Ensure cleanup procedures run properly during shutdown |
| 3 | +
|
| 4 | +Issue #1027 reported that cleanup code after "yield" in lifespan was unreachable when |
| 5 | +processes were terminated. This has been fixed by implementing the MCP spec-compliant |
| 6 | +stdio shutdown sequence that closes stdin first, allowing graceful exit. |
| 7 | +
|
| 8 | +These tests verify the fix continues to work correctly across all platforms. |
| 9 | +""" |
| 10 | + |
| 11 | +import asyncio |
| 12 | +import sys |
| 13 | +import tempfile |
| 14 | +import textwrap |
| 15 | +from pathlib import Path |
| 16 | + |
| 17 | +import anyio |
| 18 | +import pytest |
| 19 | + |
| 20 | +from mcp import ClientSession, StdioServerParameters |
| 21 | +from mcp.client.stdio import _create_platform_compatible_process, stdio_client |
| 22 | + |
| 23 | + |
| 24 | +@pytest.mark.anyio |
| 25 | +async def test_lifespan_cleanup_executed(): |
| 26 | + """ |
| 27 | + Regression test ensuring MCP server cleanup code runs during shutdown. |
| 28 | +
|
| 29 | + This test verifies that the fix for issue #1027 works correctly by: |
| 30 | + 1. Starting an MCP server that writes a marker file on startup |
| 31 | + 2. Shutting down the server normally via stdio_client |
| 32 | + 3. Verifying the cleanup code (after yield) executed and wrote its marker file |
| 33 | +
|
| 34 | + The fix implements proper stdin closure before termination, giving servers |
| 35 | + time to run their cleanup handlers. |
| 36 | + """ |
| 37 | + |
| 38 | + # Create marker files to track server lifecycle |
| 39 | + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: |
| 40 | + startup_marker = f.name |
| 41 | + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: |
| 42 | + cleanup_marker = f.name |
| 43 | + |
| 44 | + # Remove the files so we can detect when they're created |
| 45 | + Path(startup_marker).unlink() |
| 46 | + Path(cleanup_marker).unlink() |
| 47 | + |
| 48 | + # Create a minimal MCP server using FastMCP that tracks lifecycle |
| 49 | + server_code = textwrap.dedent(f""" |
| 50 | + import asyncio |
| 51 | + import sys |
| 52 | + from pathlib import Path |
| 53 | + from contextlib import asynccontextmanager |
| 54 | + from mcp.server.fastmcp import FastMCP |
| 55 | + |
| 56 | + STARTUP_MARKER = {repr(startup_marker)} |
| 57 | + CLEANUP_MARKER = {repr(cleanup_marker)} |
| 58 | + |
| 59 | + @asynccontextmanager |
| 60 | + async def lifespan(server): |
| 61 | + # Write startup marker |
| 62 | + Path(STARTUP_MARKER).write_text("started") |
| 63 | + try: |
| 64 | + yield {{"started": True}} |
| 65 | + finally: |
| 66 | + # This cleanup code now runs properly during shutdown |
| 67 | + Path(CLEANUP_MARKER).write_text("cleaned up") |
| 68 | + |
| 69 | + mcp = FastMCP("test-server", lifespan=lifespan) |
| 70 | + |
| 71 | + @mcp.tool() |
| 72 | + def echo(text: str) -> str: |
| 73 | + return text |
| 74 | + |
| 75 | + if __name__ == "__main__": |
| 76 | + mcp.run() |
| 77 | + """) |
| 78 | + |
| 79 | + # Write the server script to a temporary file |
| 80 | + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as f: |
| 81 | + server_script = f.name |
| 82 | + f.write(server_code) |
| 83 | + |
| 84 | + try: |
| 85 | + # Launch the MCP server |
| 86 | + params = StdioServerParameters(command=sys.executable, args=[server_script]) |
| 87 | + |
| 88 | + async with stdio_client(params) as (read, write): |
| 89 | + async with ClientSession(read, write) as session: |
| 90 | + # Initialize the session |
| 91 | + result = await session.initialize() |
| 92 | + assert result.protocolVersion in ["2024-11-05", "2025-06-18"] |
| 93 | + |
| 94 | + # Verify startup marker was created |
| 95 | + assert Path(startup_marker).exists(), "Server startup marker not created" |
| 96 | + assert Path(startup_marker).read_text() == "started" |
| 97 | + |
| 98 | + # Make a test request to ensure server is working |
| 99 | + response = await session.call_tool("echo", {"text": "hello"}) |
| 100 | + assert response.content[0].type == "text" |
| 101 | + assert getattr(response.content[0], "text") == "hello" |
| 102 | + |
| 103 | + # Session will be closed when exiting the context manager |
| 104 | + |
| 105 | + # Give server a moment to complete cleanup |
| 106 | + await asyncio.sleep(0.5) |
| 107 | + |
| 108 | + # Verify cleanup marker was created - this works now that stdio_client |
| 109 | + # properly closes stdin before termination, allowing graceful shutdown |
| 110 | + assert Path(cleanup_marker).exists(), "Server cleanup marker not created - regression in issue #1027 fix" |
| 111 | + assert Path(cleanup_marker).read_text() == "cleaned up" |
| 112 | + |
| 113 | + finally: |
| 114 | + # Clean up files |
| 115 | + for path in [server_script, startup_marker, cleanup_marker]: |
| 116 | + try: |
| 117 | + Path(path).unlink() |
| 118 | + except FileNotFoundError: |
| 119 | + pass |
| 120 | + |
| 121 | + |
| 122 | +@pytest.mark.anyio |
| 123 | +@pytest.mark.filterwarnings("ignore::ResourceWarning" if sys.platform == "win32" else "default") |
| 124 | +async def test_stdin_close_triggers_cleanup(): |
| 125 | + """ |
| 126 | + Regression test verifying the stdin-based graceful shutdown mechanism. |
| 127 | +
|
| 128 | + This test ensures the core fix for issue #1027 continues to work by: |
| 129 | + 1. Manually managing a server process |
| 130 | + 2. Closing stdin to trigger graceful shutdown |
| 131 | + 3. Verifying cleanup handlers run before the process exits |
| 132 | +
|
| 133 | + This mimics the behavior now implemented in stdio_client's shutdown sequence. |
| 134 | +
|
| 135 | + Note on Windows ResourceWarning: |
| 136 | + On Windows, we may see ResourceWarning about unclosed file descriptors. |
| 137 | + This is expected behavior because: |
| 138 | + - We're manually managing the process lifecycle |
| 139 | + - Windows file handle cleanup works differently than Unix |
| 140 | + - The warning doesn't indicate a real issue - cleanup still works |
| 141 | + We filter this warning on Windows only to avoid test noise. |
| 142 | + """ |
| 143 | + |
| 144 | + # Create marker files to track server lifecycle |
| 145 | + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: |
| 146 | + startup_marker = f.name |
| 147 | + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: |
| 148 | + cleanup_marker = f.name |
| 149 | + |
| 150 | + # Remove the files so we can detect when they're created |
| 151 | + Path(startup_marker).unlink() |
| 152 | + Path(cleanup_marker).unlink() |
| 153 | + |
| 154 | + # Create an MCP server that handles stdin closure gracefully |
| 155 | + server_code = textwrap.dedent(f""" |
| 156 | + import asyncio |
| 157 | + import sys |
| 158 | + from pathlib import Path |
| 159 | + from contextlib import asynccontextmanager |
| 160 | + from mcp.server.fastmcp import FastMCP |
| 161 | + |
| 162 | + STARTUP_MARKER = {repr(startup_marker)} |
| 163 | + CLEANUP_MARKER = {repr(cleanup_marker)} |
| 164 | + |
| 165 | + @asynccontextmanager |
| 166 | + async def lifespan(server): |
| 167 | + # Write startup marker |
| 168 | + Path(STARTUP_MARKER).write_text("started") |
| 169 | + try: |
| 170 | + yield {{"started": True}} |
| 171 | + finally: |
| 172 | + # This cleanup code runs when stdin closes, enabling graceful shutdown |
| 173 | + Path(CLEANUP_MARKER).write_text("cleaned up") |
| 174 | + |
| 175 | + mcp = FastMCP("test-server", lifespan=lifespan) |
| 176 | + |
| 177 | + @mcp.tool() |
| 178 | + def echo(text: str) -> str: |
| 179 | + return text |
| 180 | + |
| 181 | + if __name__ == "__main__": |
| 182 | + # The server should exit gracefully when stdin closes |
| 183 | + try: |
| 184 | + mcp.run() |
| 185 | + except Exception: |
| 186 | + # Server might get EOF or other errors when stdin closes |
| 187 | + pass |
| 188 | + """) |
| 189 | + |
| 190 | + # Write the server script to a temporary file |
| 191 | + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as f: |
| 192 | + server_script = f.name |
| 193 | + f.write(server_code) |
| 194 | + |
| 195 | + try: |
| 196 | + # This test manually manages the process to verify stdin-based shutdown |
| 197 | + # Start the server process |
| 198 | + process = await _create_platform_compatible_process( |
| 199 | + command=sys.executable, args=[server_script], env=None, errlog=sys.stderr, cwd=None |
| 200 | + ) |
| 201 | + |
| 202 | + # Wait for server to start |
| 203 | + await asyncio.sleep(1.0) # Give more time on Windows |
| 204 | + |
| 205 | + # Check if process is still running |
| 206 | + if hasattr(process, "returncode") and process.returncode is not None: |
| 207 | + pytest.fail(f"Server process exited with code {process.returncode}") |
| 208 | + |
| 209 | + assert Path(startup_marker).exists(), "Server startup marker not created" |
| 210 | + |
| 211 | + # Close stdin to signal shutdown |
| 212 | + if process.stdin: |
| 213 | + await process.stdin.aclose() |
| 214 | + |
| 215 | + # Wait for process to exit gracefully |
| 216 | + try: |
| 217 | + with anyio.fail_after(2.0): |
| 218 | + await process.wait() |
| 219 | + except TimeoutError: |
| 220 | + # If it doesn't exit after stdin close, terminate it |
| 221 | + process.terminate() |
| 222 | + await process.wait() |
| 223 | + |
| 224 | + # Check if cleanup ran |
| 225 | + await asyncio.sleep(0.5) |
| 226 | + |
| 227 | + # Verify the cleanup ran - stdin closure enables graceful shutdown |
| 228 | + assert Path(cleanup_marker).exists(), "Server cleanup marker not created - stdin-based shutdown failed" |
| 229 | + assert Path(cleanup_marker).read_text() == "cleaned up" |
| 230 | + |
| 231 | + finally: |
| 232 | + # Clean up files |
| 233 | + for path in [server_script, startup_marker, cleanup_marker]: |
| 234 | + try: |
| 235 | + Path(path).unlink() |
| 236 | + except FileNotFoundError: |
| 237 | + pass |
0 commit comments