|
6 | 6 |
|
7 | 7 | import anyio
|
8 | 8 | import anyio.lowlevel
|
| 9 | +import anyio.to_thread |
| 10 | +import psutil |
| 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,86 @@ 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_tree(pid: int, timeout: float = 2.0) -> None: |
| 256 | + """ |
| 257 | + Terminate a process and all its children using psutil. |
| 258 | +
|
| 259 | + This provides consistent behavior across platforms and properly |
| 260 | + handles process trees without shell commands. |
| 261 | +
|
| 262 | + Platform behavior: |
| 263 | + - On Unix: psutil.terminate() sends SIGTERM, allowing graceful shutdown |
| 264 | + - On Windows: psutil.terminate() calls TerminateProcess() which is immediate |
| 265 | + and doesn't allow cleanup handlers to run. This can cause ResourceWarnings |
| 266 | + for subprocess.Popen objects that don't get to clean up. |
| 267 | + """ |
| 268 | + try: |
| 269 | + parent = psutil.Process(pid) |
| 270 | + children = parent.children(recursive=True) |
| 271 | + |
| 272 | + # First, try graceful termination for all processes |
| 273 | + for child in children: |
| 274 | + try: |
| 275 | + child.terminate() |
| 276 | + except psutil.NoSuchProcess: |
| 277 | + pass |
| 278 | + |
| 279 | + try: |
| 280 | + parent.terminate() |
| 281 | + except psutil.NoSuchProcess: |
| 282 | + return # Parent already dead |
| 283 | + |
| 284 | + # Wait for processes to exit gracefully |
| 285 | + all_procs = children + [parent] |
| 286 | + _, alive = await anyio.to_thread.run_sync(lambda: psutil.wait_procs(all_procs, timeout=timeout)) |
| 287 | + |
| 288 | + # Force kill any remaining processes |
| 289 | + for proc in alive: |
| 290 | + try: |
| 291 | + proc.kill() |
| 292 | + except psutil.NoSuchProcess: |
| 293 | + pass |
| 294 | + |
| 295 | + # Wait a bit more for force-killed processes |
| 296 | + if alive: |
| 297 | + await anyio.to_thread.run_sync(lambda: psutil.wait_procs(alive, timeout=0.5)) |
| 298 | + |
| 299 | + except psutil.NoSuchProcess: |
| 300 | + # Process already terminated |
| 301 | + pass |
| 302 | + |
| 303 | + |
| 304 | +async def _terminate_process_with_children(process: Process | FallbackProcess, timeout: float = 2.0) -> None: |
| 305 | + """ |
| 306 | + Terminate a process and all its children across platforms using psutil. |
| 307 | +
|
| 308 | + This provides consistent behavior across all platforms and properly |
| 309 | + handles process trees. |
| 310 | +
|
| 311 | + Args: |
| 312 | + process: The process to terminate |
| 313 | + timeout: Time to wait for graceful termination before force killing |
| 314 | + """ |
| 315 | + # Extract PID from any process type |
| 316 | + pid = getattr(process, "pid", None) |
| 317 | + if pid is None: |
| 318 | + popen = getattr(process, "popen", None) |
| 319 | + if popen: |
| 320 | + pid = getattr(popen, "pid", None) |
| 321 | + |
| 322 | + if not pid: |
| 323 | + return |
| 324 | + |
| 325 | + await _terminate_process_tree(pid, timeout) |
0 commit comments