|
1 | 1 | import os |
2 | 2 | import sys |
3 | | -from contextlib import asynccontextmanager |
| 3 | +from contextlib import asynccontextmanager, suppress |
4 | 4 | from pathlib import Path |
5 | 5 | from typing import Literal, TextIO |
6 | 6 |
|
7 | 7 | import anyio |
8 | 8 | import anyio.lowlevel |
| 9 | +import psutil |
| 10 | +from anyio.abc import Process |
9 | 11 | from anyio.streams.memory import ( |
10 | 12 | MemoryObjectReceiveStream, |
11 | 13 | MemoryObjectSendStream, |
|
15 | 17 |
|
16 | 18 | import mcp.types as types |
17 | 19 |
|
18 | | -from .win32 import create_windows_process, get_windows_executable_command |
| 20 | +from .win32 import ( |
| 21 | + create_windows_process, |
| 22 | + get_windows_executable_command, |
| 23 | + terminate_windows_process, |
| 24 | +) |
19 | 25 |
|
20 | 26 | # Environment variables to inherit by default |
21 | 27 | DEFAULT_INHERITED_ENV_VARS = ( |
@@ -172,9 +178,7 @@ async def stdin_writer(): |
172 | 178 | yield read_stream, write_stream |
173 | 179 | finally: |
174 | 180 | # Clean up process to prevent any dangling orphaned processes |
175 | | - tg.cancel_scope.cancel() |
176 | | - process.terminate() |
177 | | - await process.wait() |
| 181 | + await _terminate_process(process) |
178 | 182 |
|
179 | 183 |
|
180 | 184 | def _get_executable_command(command: str) -> str: |
@@ -212,3 +216,28 @@ async def _create_platform_compatible_process( |
212 | 216 | ) |
213 | 217 |
|
214 | 218 | return process |
| 219 | + |
| 220 | + |
| 221 | +async def _terminate_process(process: Process): |
| 222 | + """ |
| 223 | + Terminate the process and its children. |
| 224 | +
|
| 225 | + Note: On Windows, `process.terminate()` calls the Win32 `TerminateProcess` API, |
| 226 | + which only terminates the specified process. Any child |
| 227 | + processes remain running, which can lead to unclosed resources and may cause the |
| 228 | + parent process to hang during shutdown. This function uses `psutil` to recursively |
| 229 | + terminate the full process tree, ensuring a cleaner and more reliable shutdown. |
| 230 | +
|
| 231 | + Args: |
| 232 | + process: The AnyIO process to terminate |
| 233 | + """ |
| 234 | + with suppress(psutil.NoSuchProcess): |
| 235 | + proc = psutil.Process(process.pid) |
| 236 | + children = proc.children(recursive=True) |
| 237 | + for child in children: |
| 238 | + child.kill() |
| 239 | + with suppress(ProcessLookupError): |
| 240 | + if sys.platform == "win32": |
| 241 | + await terminate_windows_process(process) |
| 242 | + else: |
| 243 | + process.terminate() |
0 commit comments