Skip to content

Commit 721133e

Browse files
Implement MCP spec-compliant stdio shutdown sequence
Following the MCP spec recommendations, the shutdown sequence now: 1. Closes stdin first to allow graceful server exit 2. Waits 2 seconds for the server to exit on its own 3. Only escalates to SIGTERM if the server doesn't exit 4. Finally uses SIGKILL as a last resort This unified approach works consistently across all platforms and gives well-behaved servers a chance to exit cleanly without signals. Co-Authored-By: davenpi <[email protected]> Add tests for MCP spec-compliant stdio shutdown Added two tests to validate the stdin-first shutdown behavior: 1. test_stdio_client_graceful_stdin_exit: Verifies that a well-behaved server exits cleanly when stdin is closed, without needing signals 2. test_stdio_client_stdin_close_ignored: Tests proper escalation to SIGTERM when a process ignores stdin closure These tests ensure the MCP spec shutdown sequence works correctly and provides graceful exit opportunities before using forceful termination.
1 parent f7143b3 commit 721133e

File tree

3 files changed

+167
-14
lines changed

3 files changed

+167
-14
lines changed

src/mcp/client/stdio/__init__.py

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,17 +177,38 @@ async def stdin_writer():
177177
try:
178178
yield read_stream, write_stream
179179
finally:
180-
# Clean up process to prevent any dangling orphaned processes
180+
# MCP spec: stdio shutdown sequence
181+
# 1. Close input stream to server
182+
# 2. Wait for server to exit, or send SIGTERM if it doesn't exit in time
183+
# 3. Send SIGKILL if still not exited
184+
if process.stdin:
185+
try:
186+
await process.stdin.aclose()
187+
except Exception:
188+
# stdin might already be closed, which is fine
189+
pass
190+
181191
try:
182-
process.terminate()
192+
# Give the process time to exit gracefully after stdin closes
183193
with anyio.fail_after(2.0):
184194
await process.wait()
185195
except TimeoutError:
186-
# If process doesn't terminate in time, force kill it
187-
process.kill()
196+
# Process didn't exit from stdin closure, escalate to SIGTERM
197+
try:
198+
process.terminate()
199+
with anyio.fail_after(2.0):
200+
await process.wait()
201+
except TimeoutError:
202+
# Process didn't respond to SIGTERM, force kill it
203+
process.kill()
204+
await process.wait()
205+
except ProcessLookupError:
206+
# Process already exited, which is fine
207+
pass
188208
except ProcessLookupError:
189209
# Process already exited, which is fine
190210
pass
211+
191212
await read_stream.aclose()
192213
await write_stream.aclose()
193214
await read_stream_writer.aclose()

src/mcp/client/stdio/win32.py

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from pathlib import Path
99
from typing import BinaryIO, TextIO, cast
1010

11-
from anyio import to_thread
11+
import anyio
1212
from anyio.streams.file import FileReadStream, FileWriteStream
1313

1414

@@ -73,9 +73,20 @@ async def __aexit__(
7373
exc_val: BaseException | None,
7474
exc_tb: object | None,
7575
) -> None:
76-
"""Terminate and wait on process exit inside a thread."""
77-
self.popen.terminate()
78-
await to_thread.run_sync(self.popen.wait)
76+
"""Clean up process and streams.
77+
78+
Attempts to terminate the process, but doesn't fail if termination
79+
is not possible (e.g., process already dead or being handled elsewhere).
80+
"""
81+
# Try to terminate the process, but don't fail if it's already being handled
82+
try:
83+
self.popen.terminate()
84+
# Use our non-blocking wait with a short timeout
85+
with anyio.move_on_after(0.5):
86+
await self.wait()
87+
except (ProcessLookupError, OSError):
88+
# Process already dead or being handled elsewhere
89+
pass
7990

8091
# Close the file handles to prevent ResourceWarning
8192
if self.stdin:
@@ -91,7 +102,11 @@ async def __aexit__(
91102

92103
async def wait(self):
93104
"""Async wait for process completion."""
94-
return await to_thread.run_sync(self.popen.wait)
105+
# Poll the process status instead of blocking wait
106+
# This allows anyio timeouts to work properly
107+
while self.popen.poll() is None:
108+
await anyio.sleep(0.1)
109+
return self.popen.returncode
95110

96111
def terminate(self):
97112
"""Terminate the subprocess immediately."""

tests/client/test_stdio.py

Lines changed: 122 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -120,12 +120,13 @@ async def test_stdio_client_universal_cleanup():
120120
)
121121

122122
server_params = StdioServerParameters(
123-
command=sys.executable,
123+
command="python",
124124
args=["-c", long_running_script],
125125
)
126126

127127
start_time = time.time()
128128

129+
# Use move_on_after which is more reliable for cleanup scenarios
129130
with anyio.move_on_after(8.0) as cancel_scope:
130131
async with stdio_client(server_params) as (read_stream, write_stream):
131132
# Immediately exit - this triggers cleanup while process is still running
@@ -134,16 +135,16 @@ async def test_stdio_client_universal_cleanup():
134135
end_time = time.time()
135136
elapsed = end_time - start_time
136137

137-
# On Windows: 2s (stdin wait) + 2s (terminate wait) + overhead = ~5s expected
138-
assert elapsed < 6.0, (
139-
f"stdio_client cleanup took {elapsed:.1f} seconds, expected < 6.0 seconds. "
138+
# Key assertion: Should complete quickly due to timeout mechanism
139+
assert elapsed < 5.0, (
140+
f"stdio_client cleanup took {elapsed:.1f} seconds, expected < 5.0 seconds. "
140141
f"This suggests the timeout mechanism may not be working properly."
141142
)
142143

143144
# Check if we timed out
144145
if cancel_scope.cancelled_caught:
145146
pytest.fail(
146-
"stdio_client cleanup timed out after 8.0 seconds. "
147+
"stdio_client cleanup timed out after 6.0 seconds. "
147148
"This indicates the cleanup mechanism is hanging and needs fixing."
148149
)
149150

@@ -212,3 +213,119 @@ def sigint_handler(signum, frame):
212213
)
213214
else:
214215
raise
216+
217+
218+
@pytest.mark.anyio
219+
async def test_stdio_client_graceful_stdin_exit():
220+
"""
221+
Test that a process exits gracefully when stdin is closed,
222+
without needing SIGTERM or SIGKILL.
223+
"""
224+
# Create a Python script that exits when stdin is closed
225+
script_content = textwrap.dedent(
226+
"""
227+
import sys
228+
229+
# Read from stdin until it's closed
230+
try:
231+
while True:
232+
line = sys.stdin.readline()
233+
if not line: # EOF/stdin closed
234+
break
235+
except:
236+
pass
237+
238+
# Exit gracefully
239+
sys.exit(0)
240+
"""
241+
)
242+
243+
server_params = StdioServerParameters(
244+
command=sys.executable,
245+
args=["-c", script_content],
246+
)
247+
248+
start_time = time.time()
249+
250+
# Use anyio timeout to prevent test from hanging forever
251+
with anyio.move_on_after(5.0) as cancel_scope:
252+
async with stdio_client(server_params) as (read_stream, write_stream):
253+
# Let the process start and begin reading stdin
254+
await anyio.sleep(0.2)
255+
# Exit context triggers cleanup - process should exit from stdin closure
256+
pass
257+
258+
if cancel_scope.cancelled_caught:
259+
pytest.fail(
260+
"stdio_client cleanup timed out after 5.0 seconds. "
261+
"Process should have exited gracefully when stdin was closed."
262+
)
263+
264+
end_time = time.time()
265+
elapsed = end_time - start_time
266+
267+
# Should complete quickly with just stdin closure (no signals needed)
268+
assert elapsed < 3.0, (
269+
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-aware process. "
270+
f"Expected < 3.0 seconds since process should exit on stdin closure."
271+
)
272+
273+
274+
@pytest.mark.anyio
275+
async def test_stdio_client_stdin_close_ignored():
276+
"""
277+
Test that when a process ignores stdin closure, the shutdown sequence
278+
properly escalates to SIGTERM.
279+
"""
280+
# Create a Python script that ignores stdin closure but responds to SIGTERM
281+
script_content = textwrap.dedent(
282+
"""
283+
import signal
284+
import sys
285+
import time
286+
287+
# Set up SIGTERM handler to exit cleanly
288+
def sigterm_handler(signum, frame):
289+
sys.exit(0)
290+
291+
signal.signal(signal.SIGTERM, sigterm_handler)
292+
293+
# Close stdin immediately to simulate ignoring it
294+
sys.stdin.close()
295+
296+
# Keep running until SIGTERM
297+
while True:
298+
time.sleep(0.1)
299+
"""
300+
)
301+
302+
server_params = StdioServerParameters(
303+
command=sys.executable,
304+
args=["-c", script_content],
305+
)
306+
307+
start_time = time.time()
308+
309+
# Use anyio timeout to prevent test from hanging forever
310+
with anyio.move_on_after(7.0) as cancel_scope:
311+
async with stdio_client(server_params) as (read_stream, write_stream):
312+
# Let the process start
313+
await anyio.sleep(0.2)
314+
# Exit context triggers cleanup
315+
pass
316+
317+
if cancel_scope.cancelled_caught:
318+
pytest.fail(
319+
"stdio_client cleanup timed out after 7.0 seconds. "
320+
"Process should have been terminated via SIGTERM escalation."
321+
)
322+
323+
end_time = time.time()
324+
elapsed = end_time - start_time
325+
326+
# Should take ~2 seconds (stdin close timeout) before SIGTERM is sent
327+
# Total time should be between 2-4 seconds
328+
assert 1.5 < elapsed < 4.5, (
329+
f"stdio_client cleanup took {elapsed:.1f} seconds for stdin-ignoring process. "
330+
f"Expected between 2-4 seconds (2s stdin timeout + termination time)."
331+
)

0 commit comments

Comments
 (0)