fix: terminate workspace shell process trees

hud/environment/tests/test_capability_backing.py

@@ -11,13 +11,16 @@
 import asyncio
 import contextlib
 import os
+import shlex
+import shutil
 import signal
 import subprocess
 import sys
 from typing import TYPE_CHECKING, cast
 
 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

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)
 
 

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()))