|
1 | 1 | import os
|
| 2 | +import signal |
| 3 | +import subprocess |
2 | 4 | import sys
|
3 | 5 | from contextlib import asynccontextmanager
|
4 | 6 | from pathlib import Path
|
5 | 7 | from typing import Literal, TextIO
|
6 | 8 |
|
7 | 9 | import anyio
|
8 | 10 | import anyio.lowlevel
|
| 11 | +from anyio.abc import Process |
9 | 12 | from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
10 | 13 | from anyio.streams.text import TextReceiveStream
|
11 | 14 | from pydantic import BaseModel, Field
|
|
14 | 17 | from mcp.shared.message import SessionMessage
|
15 | 18 |
|
16 | 19 | from .win32 import (
|
| 20 | + FallbackProcess, |
17 | 21 | create_windows_process,
|
18 | 22 | get_windows_executable_command,
|
19 | 23 | )
|
@@ -193,18 +197,9 @@ async def stdin_writer():
|
193 | 197 | with anyio.fail_after(2.0):
|
194 | 198 | await process.wait()
|
195 | 199 | except TimeoutError:
|
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 |
| 200 | + # Process didn't exit from stdin closure, use our termination function |
| 201 | + # that handles child processes properly |
| 202 | + await _terminate_process_with_children(process) |
208 | 203 | except ProcessLookupError:
|
209 | 204 | # Process already exited, which is fine
|
210 | 205 | pass
|
@@ -245,6 +240,79 @@ async def _create_platform_compatible_process(
|
245 | 240 | if sys.platform == "win32":
|
246 | 241 | process = await create_windows_process(command, args, env, errlog, cwd)
|
247 | 242 | else:
|
248 |
| - process = await anyio.open_process([command, *args], env=env, stderr=errlog, cwd=cwd) |
| 243 | + # POSIX: Create new process group for easier cleanup |
| 244 | + process = await anyio.open_process( |
| 245 | + [command, *args], |
| 246 | + env=env, |
| 247 | + stderr=errlog, |
| 248 | + cwd=cwd, |
| 249 | + start_new_session=True, # Creates new process group |
| 250 | + ) |
249 | 251 |
|
250 | 252 | return process
|
| 253 | + |
| 254 | + |
| 255 | +async def _terminate_process_with_children(process: Process | FallbackProcess, timeout: float = 2.0) -> None: |
| 256 | + """ |
| 257 | + Terminate a process and all its children across platforms. |
| 258 | +
|
| 259 | + There's no cross-platform way in the stdlib to kill a process AND its children, |
| 260 | + so we need platform-specific handling: |
| 261 | + - POSIX: Use process groups and os.killpg() |
| 262 | + - Windows: Use taskkill with /T flag for tree termination |
| 263 | +
|
| 264 | + Args: |
| 265 | + process: The process to terminate |
| 266 | + timeout: Time to wait for graceful termination before force killing |
| 267 | + """ |
| 268 | + if sys.platform != "win32": |
| 269 | + # POSIX: Kill entire process group to avoid orphaning children |
| 270 | + pid = getattr(process, "pid", None) |
| 271 | + if not pid: |
| 272 | + return |
| 273 | + |
| 274 | + try: |
| 275 | + # Try graceful termination of the group first |
| 276 | + pgid = os.getpgid(pid) |
| 277 | + os.killpg(pgid, signal.SIGTERM) |
| 278 | + |
| 279 | + # Wait for termination |
| 280 | + with anyio.fail_after(timeout): |
| 281 | + await process.wait() |
| 282 | + except TimeoutError: |
| 283 | + # Force kill the process group |
| 284 | + try: |
| 285 | + pgid = os.getpgid(pid) |
| 286 | + os.killpg(pgid, signal.SIGKILL) |
| 287 | + except (OSError, ProcessLookupError): |
| 288 | + pass |
| 289 | + except (OSError, ProcessLookupError): |
| 290 | + # Process or group already dead |
| 291 | + pass |
| 292 | + else: |
| 293 | + # Windows: Extract PID from FallbackProcess if needed |
| 294 | + pid = getattr(process, "pid", None) |
| 295 | + if pid is None: |
| 296 | + popen = getattr(process, "popen", None) |
| 297 | + if popen: |
| 298 | + pid = getattr(popen, "pid", None) |
| 299 | + |
| 300 | + if not pid: |
| 301 | + return |
| 302 | + |
| 303 | + # Try graceful termination first |
| 304 | + try: |
| 305 | + process.terminate() |
| 306 | + with anyio.fail_after(timeout): |
| 307 | + await process.wait() |
| 308 | + except TimeoutError: |
| 309 | + # Force kill using taskkill for tree termination |
| 310 | + await anyio.to_thread.run_sync( |
| 311 | + subprocess.run, |
| 312 | + ["taskkill", "/F", "/T", "/PID", str(pid)], |
| 313 | + capture_output=True, |
| 314 | + shell=False, |
| 315 | + check=False, |
| 316 | + ) |
| 317 | + except ProcessLookupError: |
| 318 | + pass |
0 commit comments