diff --git a/hud/environment/namespace.py b/hud/environment/namespace.py index e6ccd27b9..a9c31cff1 100644 --- a/hud/environment/namespace.py +++ b/hud/environment/namespace.py @@ -503,40 +503,92 @@ async def _spawn( *command_prefix, *request["argv"], ] + process: ProcessGroup | None = None if channel.term_type: - master_fd, slave_fd = pty.openpty() - process = await create_process_group_exec( - *argv, - stdin=slave_fd, - stdout=slave_fd, - stderr=slave_fd, - cwd=request["cwd"], - env=request["env"], - ) - os.close(slave_fd) - await channel.redirect( - stdin=os.dup(master_fd), - stdout=os.dup(master_fd), - send_eof=False, - ) - os.close(master_fd) + master_fd = slave_fd = stdin_fd = stdout_fd = -1 + try: + master_fd, slave_fd = pty.openpty() + process = await create_process_group_exec( + *argv, + stdin=slave_fd, + stdout=slave_fd, + stderr=slave_fd, + cwd=request["cwd"], + env=request["env"], + ) + os.close(slave_fd) + slave_fd = -1 + stdin_fd = os.dup(master_fd) + stdout_fd = os.dup(master_fd) + os.close(master_fd) + master_fd = -1 + + descriptor, stdin_fd = stdin_fd, -1 + await channel.redirect_stdin(descriptor) + descriptor, stdout_fd = stdout_fd, -1 + await channel.redirect_stdout(descriptor, send_eof=False) + except BaseException as exc: + if process is not None: + try: + await process.terminate() + except Exception as cleanup_exc: + exc.add_note(f"failed to terminate spawned process: {cleanup_exc}") + raise + finally: + for descriptor in (master_fd, slave_fd, stdin_fd, stdout_fd): + if descriptor != -1: + os.close(descriptor) else: - process = await create_process_group_exec( - *argv, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=request["cwd"], - env=request["env"], - ) - assert process.process.stdin is not None - assert process.stdout is not None and process.stderr is not None - await channel.redirect( - stdin=process.process.stdin, - stdout=process.stdout, - stderr=process.stderr, - send_eof=False, - ) + stdin_read = stdin_write = -1 + stdout_read = stdout_write = -1 + stderr_read = stderr_write = -1 + try: + # AsyncSSH closes raw pipe transports with the channel, even + # when a background descendant retains the child end. + stdin_read, stdin_write = os.pipe() + stdout_read, stdout_write = os.pipe() + stderr_read, stderr_write = os.pipe() + process = await create_process_group_exec( + *argv, + stdin=stdin_read, + stdout=stdout_write, + stderr=stderr_write, + cwd=request["cwd"], + env=request["env"], + ) + + os.close(stdin_read) + stdin_read = -1 + os.close(stdout_write) + stdout_write = -1 + os.close(stderr_write) + stderr_write = -1 + + descriptor, stdin_write = stdin_write, -1 + await channel.redirect_stdin(descriptor) + descriptor, stdout_read = stdout_read, -1 + await channel.redirect_stdout(descriptor, send_eof=False) + descriptor, stderr_read = stderr_read, -1 + await channel.redirect_stderr(descriptor, send_eof=False) + except BaseException as exc: + if process is not None: + try: + await process.terminate() + except Exception as cleanup_exc: + exc.add_note(f"failed to terminate spawned process: {cleanup_exc}") + raise + finally: + for descriptor in ( + stdin_read, + stdin_write, + stdout_read, + stdout_write, + stderr_read, + stderr_write, + ): + if descriptor != -1: + os.close(descriptor) + assert process is not None wait_task = asyncio.create_task(process.wait()) closed_task = asyncio.create_task(channel.channel.wait_closed()) try: @@ -550,6 +602,12 @@ async def _spawn( if not request["persistent"]: await process.terminate() return returncode + except BaseException as exc: + try: + await process.terminate() + except Exception as cleanup_exc: + exc.add_note(f"failed to terminate spawned process: {cleanup_exc}") + raise finally: wait_task.cancel() closed_task.cancel() diff --git a/hud/environment/tests/test_workspace.py b/hud/environment/tests/test_workspace.py index bbc1c753d..ed9975432 100644 --- a/hud/environment/tests/test_workspace.py +++ b/hud/environment/tests/test_workspace.py @@ -31,10 +31,10 @@ from hud.environment import workspace as workspace_mod from hud.environment.egress import Peer, _field, _Unrelayable from hud.environment.workspace import Bubblewrap, Mount, Workspace -from hud.utils.process import ProcessResult +from hud.utils.process import ProcessGroup, ProcessResult if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import AsyncIterator, Mapping pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="POSIX workspace semantics") @@ -52,6 +52,67 @@ async def _connect(ws: Workspace) -> asyncssh.SSHClientConnection: ) +@contextlib.asynccontextmanager +async def _connected_namespace_host( + monkeypatch: pytest.MonkeyPatch, +) -> AsyncIterator[namespace_mod.NamespaceHost]: + socket_dir = tempfile.TemporaryDirectory(prefix="hud-namespace-", dir="/tmp") + socket_path = Path(socket_dir.name) / "host.sock" + server_host = namespace_mod._NamespaceHost( + socket_path, + setup_loopback=False, + holder_argv=[], + bwrap="", + launcher_depth=0, + map_identities=False, + ports=frozenset(), + ) + holders = [SimpleNamespace(terminate=AsyncMock()) for _ in range(2)] + monkeypatch.setattr( + server_host, + "_start_holder", + AsyncMock(side_effect=[(holder, os.getpid()) for holder in holders]), + ) + + create_process_group_exec = namespace_mod.create_process_group_exec + + async def spawn_without_nsenter(*argv: str, **kwargs: Any) -> ProcessGroup: + command = argv[argv.index("--") + 1 :] + return await create_process_group_exec(*command, **kwargs) + + monkeypatch.setattr(namespace_mod, "create_process_group_exec", spawn_without_nsenter) + listen = asyncssh.listen + listening = asyncio.Event() + + async def listen_and_signal(*args: Any, **kwargs: Any) -> asyncssh.SSHAcceptor: + server = await listen(*args, **kwargs) + listening.set() + return server + + monkeypatch.setattr(asyncssh, "listen", listen_and_signal) + serve_task = asyncio.create_task(server_host.serve()) + client: namespace_mod.NamespaceHost | None = None + try: + await asyncio.wait_for(listening.wait(), 1.0) + client = namespace_mod.NamespaceHost(socket_path) + await client.connect() + yield client + finally: + if client is not None: + await client.close() + serve_task.cancel() + await asyncio.gather(serve_task, return_exceptions=True) + socket_dir.cleanup() + + +def _wait_for_path(path: Path, timeout: float = 5.0) -> None: + deadline = time.monotonic() + timeout + while not path.exists(): + if time.monotonic() >= deadline: + raise TimeoutError(f"timed out waiting for {path}") + time.sleep(0.01) + + @pytest.mark.asyncio async def test_start_waits_until_the_ssh_acceptor_is_ready( tmp_path: Path, @@ -156,6 +217,94 @@ async def test_output_arrives_while_the_command_is_still_running(tmp_path: Path) assert elapsed < 2.0, f"first line took {elapsed:.1f}s — output is not streaming" +@pytest.mark.parametrize( + ("command", "marker"), + [ + ( + "sh -c 'sleep .1; exec >/dev/null 2>&1; : > silent-done' & printf started", + "silent-done", + ), + ( + "sh -c \"sleep .1; trap '' PIPE; printf late-out || :; " + "printf late-err >&2 || :; exec >/dev/null 2>&1; " + ': > late-done" & printf started', + "late-done", + ), + ( + "cd . && nohup sh -c 'sleep .1' " + "stdin.log 2>&1 && exec >/dev/null 2>&1 && " + ": > stdin-done & printf started", + "stdin-done", + ), + ( + "(cd . && exec sh -c 'sleep .1; : > redirected-done') " + "redirected.log 2>&1 & printf started", + "redirected-done", + ), + ], + ids=("inherited-eof", "late-output", "stdin-detached", "fully-redirected"), +) +@pytest.mark.asyncio +async def test_background_completion_preserves_subsequent_namespace_commands( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + command: str, + marker: str, +) -> None: + shell = shutil.which("bash") or "sh" + async with _connected_namespace_host(monkeypatch) as namespace: + launched = await namespace.spawn( + [shell, "-c", command], + cwd=tmp_path, + env=dict(os.environ), + persistent=True, + ) + launch_result = await asyncio.wait_for(launched.complete(), 5.0) + + assert launch_result.returncode == 0 + assert launch_result.stdout == b"started" + + await asyncio.to_thread(_wait_for_path, tmp_path / marker) + await asyncio.sleep(0) + + probe = await namespace.spawn( + [shell, "-c", "printf healthy"], + cwd=tmp_path, + env=dict(os.environ), + persistent=True, + ) + probe_result = await asyncio.wait_for(probe.complete(), 5.0) + + assert probe_result.returncode == 0 + assert probe_result.stdout == b"healthy" + + +@pytest.mark.asyncio +async def test_namespace_process_forwards_standard_streams( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + shell = shutil.which("bash") or "sh" + async with _connected_namespace_host(monkeypatch) as namespace: + process = await namespace.spawn( + [ + shell, + "-c", + "IFS= read -r line; printf 'out:%s' \"$line\"; printf err >&2", + ], + cwd=tmp_path, + env=dict(os.environ), + persistent=True, + ) + process.stdin.write(b"input\n") + process.stdin.write_eof() + result = await asyncio.wait_for(process.complete(), 5.0) + + assert result.returncode == 0 + assert result.stdout == b"out:input" + assert result.stderr == b"err" + + @pytest.mark.asyncio async def test_a_session_that_asks_for_a_terminal_gets_one(tmp_path: Path) -> None: """Programs branch on isatty: without a pty they take their batch path, so