From 461cf53980d414670307f4b3823e332692b04c26 Mon Sep 17 00:00:00 2001
From: Rishit Sharma <rishitsharmar89@gmail.com>
Date: Mon, 14 Sep 2026 10:30:10 +0530
Subject: [PATCH] fix(apodex): kill the whole command session when run_shell
 times out
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

run_shell awaited communicate() under a bare wait_for, so a timeout — or a
cancelled tool call, since tool_exec wraps every invocation in its own
wait_for — left the shell and everything it had started still running. The
model was told the command had stopped while it kept burning CPU and writing
the workspace.

Start the command in its own session and SIGKILL that group in a finally, the
contract _CurrentCommands.run already uses on the bare host for the same
reason. Terminating only the shell is not enough: a grandchild keeps the
capture pipes open, asyncio wakes wait() only once those close, so that fix
both leaves the grandchild writing and blocks until it exits on its own. The
wait after the kill is bounded, because a setsid escapee is out of killpg's
reach and the caller still has to get its TimeoutError.

The regression test drives a backgrounded subshell, which survives a
shell-only kill and is the case a plain terminate/kill pair passes by luck.
---
 apodex/sandbox.py           | 16 +++++++++++++++-
 apodex/tests/test_native.py | 19 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/apodex/sandbox.py b/apodex/sandbox.py
index 4b0a933..eb6e402 100644
--- a/apodex/sandbox.py
+++ b/apodex/sandbox.py
@@ -41,9 +41,11 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 import os
 import shlex
+import signal
 import sys
 from dataclasses import dataclass
 from pathlib import Path
@@ -258,8 +260,20 @@ async def run_shell(
         cwd=cwd,
         stdout=asyncio.subprocess.PIPE,
         stderr=asyncio.subprocess.PIPE,
+        start_new_session=True,
     )
-    out, err = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+    try:
+        out, err = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+    finally:
+        if proc.returncode is None:
+            # Timed out or cancelled: kill the whole session, not just the
+            # shell, or its children keep running and holding the pipes. Same
+            # contract as ``_CurrentCommands.run`` in plugins.tools._sandbox,
+            # including the bounded wait for a setsid escapee killpg misses.
+            with contextlib.suppress(OSError):
+                os.killpg(proc.pid, signal.SIGKILL)
+            with contextlib.suppress(TimeoutError):
+                await asyncio.wait_for(proc.wait(), timeout=5)
     return (
         proc.returncode or 0,
         out.decode("utf-8", "replace"),
diff --git a/apodex/tests/test_native.py b/apodex/tests/test_native.py
index 20c1a88..51ae6e0 100644
--- a/apodex/tests/test_native.py
+++ b/apodex/tests/test_native.py
@@ -2,8 +2,11 @@
 
 import asyncio
 import os
+import time
 from pathlib import Path
 
+import pytest
+
 from apodex import cli, docker, sandbox
 from apodex.native import prepare_native_runtime
 from apodex.sandbox import BWRAP, CONTAINER, NATIVE, Strategy, resolve_strategy
@@ -199,6 +202,22 @@ def kill(self):
     assert second.binds == ((str(second_workspace.resolve()),) * 2 + (False,),)
 
 
+def test_run_shell_kills_the_whole_command_on_timeout(tmp_path) -> None:
+    """A timed-out command must not keep writing to the workspace.
+
+    The subshell is a grandchild holding the output pipes, so killing only the
+    shell would still leave it alive to write the marker.
+    """
+    with pytest.raises(TimeoutError):
+        asyncio.run(sandbox.run_shell(
+            "(sleep 2; touch marker) & wait", str(tmp_path), 1,
+            Strategy(NATIVE, "test"),
+        ))
+    time.sleep(2)
+
+    assert not (tmp_path / "marker").exists()
+
+
 def test_macos_falls_back_to_native_when_docker_is_unavailable(
     tmp_path, monkeypatch, capsys,
 ) -> None: