diff --git a/packages/kaos/src/kaos/local.py b/packages/kaos/src/kaos/local.py index 7ab2f7cd0a..815f05b151 100644 --- a/packages/kaos/src/kaos/local.py +++ b/packages/kaos/src/kaos/local.py @@ -2,10 +2,13 @@ import asyncio import os +import signal +import subprocess from asyncio.subprocess import Process as AsyncioProcess from collections.abc import AsyncGenerator +from contextlib import suppress from pathlib import Path, PurePath -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal, cast if os.name == "nt": import ntpath as pathmodule @@ -59,7 +62,36 @@ async def wait(self) -> int: return await self._process.wait() async def kill(self) -> None: - self._process.kill() + if self._process.returncode is not None: + return + + if os.name == "nt": + # The process has its own console process group, so Ctrl+Break + # cannot reach Kimi or another independently launched command. + try: + os.kill(self.pid, signal.CTRL_BREAK_EVENT) + await asyncio.sleep(0.1) + except OSError: + pass + killer = await asyncio.create_subprocess_exec( + "taskkill", + "/PID", + str(self.pid), + "/T", + "/F", + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + await killer.wait() + if killer.returncode != 0 and self._process.returncode is None: + self._process.kill() + else: + # exec() starts each command in a dedicated session. Its PID is + # therefore also the process-group ID and safe to target here. + with suppress(ProcessLookupError): + os_module = cast(Any, os) + signal_module = cast(Any, signal) + os_module.killpg(self.pid, signal_module.SIGKILL) def pathclass(self) -> type[PurePath]: return PurePathClass @@ -166,13 +198,26 @@ async def exec(self, *args: str, env: Mapping[str, str] | None = None) -> KaosPr if not args: raise ValueError("At least one argument (the program to execute) is required.") - process = await asyncio.create_subprocess_exec( - *args, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=env, - ) + if os.name == "nt": + process = await asyncio.create_subprocess_exec( + *args, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + # Isolate console control events used by Process.kill(). + creationflags=subprocess.CREATE_NEW_PROCESS_GROUP, + ) + else: + process = await asyncio.create_subprocess_exec( + *args, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + # Give Process.kill() an exact group containing only this tree. + start_new_session=True, + ) return self.Process(process) diff --git a/packages/kaos/tests/test_local_kaos.py b/packages/kaos/tests/test_local_kaos.py index 80709f4f56..c03f467956 100644 --- a/packages/kaos/tests/test_local_kaos.py +++ b/packages/kaos/tests/test_local_kaos.py @@ -196,3 +196,43 @@ async def test_exec_wait_timeout(local_kaos: LocalKaos): if process.returncode is None: await process.kill() await process.wait() + + +async def test_kill_terminates_descendant_processes(local_kaos: LocalKaos, tmp_path: Path): + marker = tmp_path / "descendant-survived" + if os.name == "nt": + bash = Path(os.environ["PROGRAMFILES"]) / "Git" / "bin" / "bash.exe" + if not bash.exists(): + pytest.skip("Git Bash is required for the Windows process-tree regression test") + marker_arg = marker.as_posix() + process = await local_kaos.exec( + str(bash), + "-c", + f'(sleep 0.5; echo alive > "{marker_arg}") & sleep 30', + ) + else: + descendant_code = ( + "import pathlib, sys, time; " + "time.sleep(0.5); " + "pathlib.Path(sys.argv[1]).write_text('alive', encoding='utf-8')" + ) + parent_code = ( + "import subprocess, sys, time; " + "subprocess.Popen([sys.executable, '-c', sys.argv[1], sys.argv[2]]); " + "time.sleep(30)" + ) + process = await local_kaos.exec( + sys.executable, + "-c", + parent_code, + descendant_code, + str(marker), + ) + + await asyncio.sleep(0.1) + await process.kill() + await asyncio.wait_for(process.wait(), timeout=2) + await asyncio.gather(process.stdout.read(), process.stderr.read()) + await asyncio.sleep(0.6) + + assert not marker.exists()