|
1 | 1 | import shutil
|
| 2 | +import sys |
| 3 | +import textwrap |
| 4 | +import time |
2 | 5 |
|
| 6 | +import anyio |
3 | 7 | import pytest
|
4 | 8 |
|
5 | 9 | from mcp.client.session import ClientSession
|
|
11 | 15 | from mcp.shared.message import SessionMessage
|
12 | 16 | from mcp.types import CONNECTION_CLOSED, JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
|
13 | 17 |
|
| 18 | +# Timeout for cleanup of processes that ignore SIGTERM |
| 19 | +# This timeout ensures the test fails quickly if the cleanup logic doesn't have |
| 20 | +# proper fallback mechanisms (SIGINT/SIGKILL) for processes that ignore SIGTERM |
| 21 | +SIGTERM_IGNORING_PROCESS_TIMEOUT = 5.0 |
| 22 | + |
14 | 23 | tee: str = shutil.which("tee") # type: ignore
|
15 | 24 | python: str = shutil.which("python") # type: ignore
|
16 | 25 |
|
@@ -90,3 +99,123 @@ async def test_stdio_client_nonexistent_command():
|
90 | 99 | or "not found" in error_message.lower()
|
91 | 100 | or "cannot find the file" in error_message.lower() # Windows error message
|
92 | 101 | )
|
| 102 | + |
| 103 | + |
| 104 | +@pytest.mark.anyio |
| 105 | +async def test_stdio_client_universal_cleanup(): |
| 106 | + """ |
| 107 | + Test that stdio_client completes cleanup within reasonable time |
| 108 | + even when connected to processes that exit slowly. |
| 109 | + """ |
| 110 | + |
| 111 | + # Use a Python script that simulates a long-running process |
| 112 | + # This ensures consistent behavior across platforms |
| 113 | + long_running_script = textwrap.dedent( |
| 114 | + """ |
| 115 | + import time |
| 116 | + import sys |
| 117 | + |
| 118 | + # Simulate a long-running process |
| 119 | + for i in range(100): |
| 120 | + time.sleep(0.1) |
| 121 | + # Flush to ensure output is visible |
| 122 | + sys.stdout.flush() |
| 123 | + sys.stderr.flush() |
| 124 | + """ |
| 125 | + ) |
| 126 | + |
| 127 | + server_params = StdioServerParameters( |
| 128 | + command=sys.executable, |
| 129 | + args=["-c", long_running_script], |
| 130 | + ) |
| 131 | + |
| 132 | + start_time = time.time() |
| 133 | + |
| 134 | + with anyio.move_on_after(8.0) as cancel_scope: |
| 135 | + async with stdio_client(server_params) as (read_stream, write_stream): |
| 136 | + # Immediately exit - this triggers cleanup while process is still running |
| 137 | + pass |
| 138 | + |
| 139 | + end_time = time.time() |
| 140 | + elapsed = end_time - start_time |
| 141 | + |
| 142 | + # On Windows: 2s (stdin wait) + 2s (terminate wait) + overhead = ~5s expected |
| 143 | + assert elapsed < 6.0, ( |
| 144 | + f"stdio_client cleanup took {elapsed:.1f} seconds, expected < 6.0 seconds. " |
| 145 | + f"This suggests the timeout mechanism may not be working properly." |
| 146 | + ) |
| 147 | + |
| 148 | + # Check if we timed out |
| 149 | + if cancel_scope.cancelled_caught: |
| 150 | + pytest.fail( |
| 151 | + "stdio_client cleanup timed out after 8.0 seconds. " |
| 152 | + "This indicates the cleanup mechanism is hanging and needs fixing." |
| 153 | + ) |
| 154 | + |
| 155 | + |
| 156 | +@pytest.mark.anyio |
| 157 | +@pytest.mark.skipif(sys.platform == "win32", reason="Windows signal handling is different") |
| 158 | +async def test_stdio_client_sigint_only_process(): |
| 159 | + """ |
| 160 | + Test cleanup with a process that ignores SIGTERM but responds to SIGINT. |
| 161 | + """ |
| 162 | + # Create a Python script that ignores SIGTERM but handles SIGINT |
| 163 | + script_content = textwrap.dedent( |
| 164 | + """ |
| 165 | + import signal |
| 166 | + import sys |
| 167 | + import time |
| 168 | +
|
| 169 | + # Ignore SIGTERM (what process.terminate() sends) |
| 170 | + signal.signal(signal.SIGTERM, signal.SIG_IGN) |
| 171 | +
|
| 172 | + # Handle SIGINT (Ctrl+C signal) by exiting cleanly |
| 173 | + def sigint_handler(signum, frame): |
| 174 | + sys.exit(0) |
| 175 | +
|
| 176 | + signal.signal(signal.SIGINT, sigint_handler) |
| 177 | +
|
| 178 | + # Keep running until SIGINT received |
| 179 | + while True: |
| 180 | + time.sleep(0.1) |
| 181 | + """ |
| 182 | + ) |
| 183 | + |
| 184 | + server_params = StdioServerParameters( |
| 185 | + command=sys.executable, |
| 186 | + args=["-c", script_content], |
| 187 | + ) |
| 188 | + |
| 189 | + start_time = time.time() |
| 190 | + |
| 191 | + try: |
| 192 | + # Use anyio timeout to prevent test from hanging forever |
| 193 | + with anyio.move_on_after(5.0) as cancel_scope: |
| 194 | + async with stdio_client(server_params) as (read_stream, write_stream): |
| 195 | + # Let the process start and begin ignoring SIGTERM |
| 196 | + await anyio.sleep(0.5) |
| 197 | + # Exit context triggers cleanup - this should not hang |
| 198 | + pass |
| 199 | + |
| 200 | + if cancel_scope.cancelled_caught: |
| 201 | + raise TimeoutError("Test timed out") |
| 202 | + |
| 203 | + end_time = time.time() |
| 204 | + elapsed = end_time - start_time |
| 205 | + |
| 206 | + # Should complete quickly even with SIGTERM-ignoring process |
| 207 | + # This will fail if cleanup only uses process.terminate() without fallback |
| 208 | + assert elapsed < SIGTERM_IGNORING_PROCESS_TIMEOUT, ( |
| 209 | + f"stdio_client cleanup took {elapsed:.1f} seconds with SIGTERM-ignoring process. " |
| 210 | + f"Expected < {SIGTERM_IGNORING_PROCESS_TIMEOUT} seconds. " |
| 211 | + "This suggests the cleanup needs SIGINT/SIGKILL fallback." |
| 212 | + ) |
| 213 | + except (TimeoutError, Exception) as e: |
| 214 | + if isinstance(e, TimeoutError) or "timed out" in str(e): |
| 215 | + pytest.fail( |
| 216 | + f"stdio_client cleanup timed out after {SIGTERM_IGNORING_PROCESS_TIMEOUT} seconds " |
| 217 | + "with SIGTERM-ignoring process. " |
| 218 | + "This confirms the cleanup needs SIGINT/SIGKILL fallback for processes that ignore SIGTERM." |
| 219 | + ) |
| 220 | + else: |
| 221 | + raise |
0 commit comments