diff --git a/hud/environment/tests/test_capability_backing.py b/hud/environment/tests/test_capability_backing.py index 02e03f3b9..fed11b6f4 100644 --- a/hud/environment/tests/test_capability_backing.py +++ b/hud/environment/tests/test_capability_backing.py @@ -11,6 +11,8 @@ import asyncio import contextlib import os +import shlex +import shutil import signal import subprocess import sys @@ -18,6 +20,7 @@ import pytest +import hud.environment.workspace as workspace_module from hud.capabilities import Capability from hud.environment import Environment @@ -148,6 +151,176 @@ async def test_workspace_command_teardown_kills_background_children(tmp_path: Pa os.kill(pid, signal.SIGKILL) +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX process-group regression") +async def test_workspace_session_close_kills_running_children(tmp_path: Path) -> None: + env = Environment("ws-env") + env.workspace(tmp_path / "root", track_files=False) + pid: int | None = None + + try: + async with served(env) as client: + ssh = cast("SSHClient", await client.open("shell")) + remote = await ssh.conn.create_process( + "trap '' TERM; sleep 120 & echo $! > child.pid; wait" + ) + try: + result = await asyncio.wait_for( + ssh.conn.run( + "while [ ! -s child.pid ]; do sleep 0.01; done; cat child.pid", + check=True, + ), + 5.0, + ) + assert isinstance(result.stdout, str) + pid = int(result.stdout) + assert _pid_is_running(pid) + finally: + remote.close() + await asyncio.wait_for(remote.wait_closed(), 5.0) + + assert await _wait_for_pid_inactive(pid, max_wait=5.0) + finally: + if pid is not None and _pid_is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX process-group regression") +async def test_workspace_channel_close_during_teardown_keeps_connection_usable( + tmp_path: Path, +) -> None: + env = Environment("ws-env") + env.workspace(tmp_path / "root", track_files=False) + pid: int | None = None + + try: + async with served(env) as client: + ssh = cast("SSHClient", await client.open("shell")) + remote = await ssh.conn.create_process( + "sh -c 'trap \"echo started > teardown-started; " + 'while :; do sleep 1; done" TERM; ' + "echo $$ > close-race-child.pid; while :; do sleep 1; done' & " + "printf output" + ) + try: + result = await asyncio.wait_for( + ssh.conn.run( + "while [ ! -s teardown-started ]; do sleep 0.01; done; " + "cat close-race-child.pid", + check=True, + ), + 5.0, + ) + assert isinstance(result.stdout, str) + pid = int(result.stdout) + assert _pid_is_running(pid) + + remote.close() + await asyncio.wait_for(remote.wait_closed(), 5.0) + assert await _wait_for_pid_inactive(pid, max_wait=5.0) + await ssh.conn.run("true", check=True) + finally: + remote.close() + await asyncio.wait_for(remote.wait_closed(), 5.0) + finally: + if pid is not None and _pid_is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX process-group regression") +async def test_workspace_timeout_reports_exit_after_child_teardown( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(workspace_module, "_COMMAND_TIMEOUT", 0.5) + env = Environment("ws-env") + env.workspace(tmp_path / "root", track_files=False) + pid: int | None = None + + try: + async with served(env) as client: + ssh = cast("SSHClient", await client.open("shell")) + remote = await ssh.conn.create_process( + "trap '' TERM; sleep 120 & echo $! > timeout-child.pid; wait" + ) + try: + result = await asyncio.wait_for( + ssh.conn.run( + "while [ ! -s timeout-child.pid ]; do sleep 0.01; done; " + "cat timeout-child.pid", + check=True, + ), + 5.0, + ) + assert isinstance(result.stdout, str) + pid = int(result.stdout) + assert _pid_is_running(pid) + + completed = await remote.wait(timeout=5.0) + assert completed.exit_status == 1 + assert not _pid_is_running(pid) + finally: + remote.close() + await asyncio.wait_for(remote.wait_closed(), 5.0) + finally: + if pid is not None and _pid_is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX setsid") +async def test_workspace_command_bounds_detached_output_drain(tmp_path: Path) -> None: + env = Environment("ws-env") + env.workspace(tmp_path / "root", track_files=False) + pid: int | None = None + setsid = shutil.which("setsid") + if setsid is not None: + detached_command = ( + f"{shlex.quote(setsid)} sh -c " + "'echo $$ > detached.pid; printf ready; kill -KILL $PPID; sleep 120'" + ) + else: + code = ( + "import os,signal,time;" + "from pathlib import Path;" + "parent=os.getppid();" + "os.setsid();" + "Path('detached.pid').write_text(str(os.getpid()));" + "os.write(1,b'ready');" + "os.kill(parent,signal.SIGKILL);" + "time.sleep(120)" + ) + detached_command = f"{shlex.quote(sys.executable)} -c {shlex.quote(code)}" + + try: + async with served(env) as client: + ssh = cast("SSHClient", await client.open("shell")) + remote = await ssh.conn.create_process(f"{detached_command} & wait") + try: + result = await asyncio.wait_for( + ssh.conn.run( + "while [ ! -s detached.pid ]; do sleep 0.01; done; cat detached.pid", + check=True, + ), + 5.0, + ) + assert isinstance(result.stdout, str) + pid = int(result.stdout) + assert _pid_is_running(pid) + + completed = await remote.wait(timeout=5.0) + assert completed.stdout == "ready" + assert _pid_is_running(pid) + finally: + remote.close() + await asyncio.wait_for(remote.wait_closed(), 5.0) + finally: + if pid is not None and _pid_is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.killpg(pid, signal.SIGKILL) + + async def test_stop_tears_down_the_workspace(tmp_path: Path) -> None: import asyncio from urllib.parse import urlsplit diff --git a/hud/environment/workspace.py b/hud/environment/workspace.py index 31b92c847..10c24f4fa 100644 --- a/hud/environment/workspace.py +++ b/hud/environment/workspace.py @@ -27,6 +27,8 @@ LOGGER = logging.getLogger("hud.environment.workspace") +_COMMAND_TIMEOUT = 3600.0 + # Set once the first Workspace logs the missing-bwrap notice (avoid per-instance spam). _warned_no_bwrap = False @@ -675,44 +677,65 @@ async def relay_stdin() -> None: while chunk := await process.stdin.read(65536): stdin.write(chunk) await stdin.drain() - except (BrokenPipeError, ConnectionResetError): + except (asyncssh.Error, BrokenPipeError, ConnectionResetError): pass finally: stdin.close() + async def drain_output(reader: asyncio.StreamReader, output: bytearray) -> None: + while chunk := await reader.read(65536): + output.extend(chunk) + + stdout_data = bytearray() + stderr_data = bytearray() stdin_task = asyncio.create_task(relay_stdin()) - stdout_task = asyncio.create_task(stdout.read()) - stderr_task = asyncio.create_task(stderr.read()) + stdout_task = asyncio.create_task(drain_output(stdout, stdout_data)) + stderr_task = asyncio.create_task(drain_output(stderr, stderr_data)) + wait_task = asyncio.create_task(sub.wait()) + channel_closed_task = asyncio.create_task(process.channel.wait_closed()) + timed_out = False try: - async with asyncio.timeout(3600.0): - _, stdout_data, stderr_data = await asyncio.gather( - sub.wait(), - stdout_task, - stderr_task, + async with asyncio.timeout(_COMMAND_TIMEOUT): + done, _ = await asyncio.wait( + (wait_task, channel_closed_task), + return_when=asyncio.FIRST_COMPLETED, ) - await sub.terminate() - stdin_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await stdin_task + if channel_closed_task in done and not wait_task.done(): + return + await wait_task except TimeoutError: + timed_out = True + finally: await sub.terminate() stdin_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await stdin_task - process.stderr.write(b"workspace: command timed out after 3600s\n") + wait_task.cancel() + channel_closed_task.cancel() + await asyncio.gather( + stdin_task, + wait_task, + channel_closed_task, + return_exceptions=True, + ) + _, output_pending = await asyncio.wait( + (stdout_task, stderr_task), + timeout=1.0, + ) + for task in output_pending: + task.cancel() + await asyncio.gather(stdout_task, stderr_task, return_exceptions=True) + + if process.channel.is_closing(): + return + if timed_out: + process.stderr.write( + f"workspace: command timed out after {_COMMAND_TIMEOUT:g}s\n".encode() + ) process.exit(1) return - except BaseException: - await sub.terminate() - stdin_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await stdin_task - raise - if stdout_data: - process.stdout.write(stdout_data) + process.stdout.write(bytes(stdout_data)) if stderr_data: - process.stderr.write(stderr_data) + process.stderr.write(bytes(stderr_data)) process.exit(sub.returncode if sub.returncode is not None else 0) diff --git a/hud/utils/process.py b/hud/utils/process.py index 2fc45e797..fcbd79f6f 100644 --- a/hud/utils/process.py +++ b/hud/utils/process.py @@ -9,6 +9,8 @@ from dataclasses import dataclass from typing import Any +_PROCESS_EXIT_POLL_INTERVAL = 0.05 + @dataclass(slots=True) class ProcessGroup: @@ -37,7 +39,26 @@ def returncode(self) -> int | None: return self.process.returncode async def wait(self) -> int: - return await self.process.wait() + """Wait for the process leader without requiring inherited pipes to close.""" + returncode = self.process.returncode + if returncode is not None: + return returncode + + wait_task = asyncio.create_task(self.process.wait()) + try: + while True: + done, _ = await asyncio.wait( + (wait_task,), + timeout=_PROCESS_EXIT_POLL_INTERVAL, + ) + if done: + return await wait_task + returncode = self.process.returncode + if returncode is not None: + return returncode + finally: + wait_task.cancel() + await asyncio.gather(wait_task, return_exceptions=True) async def communicate( self, @@ -50,10 +71,8 @@ async def communicate( result = await self.process.communicate(input=input) else: result = await asyncio.wait_for(self.process.communicate(input=input), max_wait) - except BaseException: + finally: await self.terminate() - raise - await self.terminate() return result async def terminate(self) -> None: @@ -105,6 +124,8 @@ async def _terminate_process_group( await asyncio.wait_for(proc.wait(), kill_timeout) return + loop = asyncio.get_running_loop() + term_deadline = loop.time() + term_timeout + settle_time try: os.killpg(proc.pid, signal.SIGTERM) except ProcessLookupError: @@ -114,11 +135,14 @@ async def _terminate_process_group( return if proc.returncode is None: + remaining = max(0.0, term_deadline - loop.time()) with contextlib.suppress(TimeoutError): - await asyncio.wait_for(proc.wait(), term_timeout) + await asyncio.wait_for(proc.wait(), remaining) + + remaining = max(0.0, term_deadline - loop.time()) + if await _wait_for_process_group_exit(proc.pid, remaining): + return - if settle_time > 0: - await asyncio.sleep(settle_time) with contextlib.suppress(ProcessLookupError): os.killpg(proc.pid, signal.SIGKILL) @@ -128,3 +152,16 @@ async def _terminate_process_group( else: with contextlib.suppress(TimeoutError): await asyncio.wait_for(proc.wait(), kill_timeout) + + +async def _wait_for_process_group_exit(process_group: int, max_wait: float) -> bool: + loop = asyncio.get_running_loop() + deadline = loop.time() + max_wait + while True: + try: + os.killpg(process_group, 0) + except ProcessLookupError: + return True + if loop.time() >= deadline: + return False + await asyncio.sleep(min(_PROCESS_EXIT_POLL_INTERVAL, deadline - loop.time()))