feat: add session management to shell commands and improve task cancellation handling

Signed-off-by: Frost Ming <me@frostming.com>

Author: frostming

src/bub/builtin/shell_manager.py

@@ -13,6 +13,7 @@ class ManagedShell:
     shell_id: str
     cmd: str
     cwd: str | None
+    session_id: str | None
     process: asyncio.subprocess.Process
     output_chunks: list[str] = field(default_factory=list)
     read_tasks: list[asyncio.Task[None]] = field(default_factory=list)
@@ -36,15 +37,21 @@ class ShellManager:
     def __init__(self) -> None:
         self._shells: dict[str, ManagedShell] = {}
 
-    async def start(self, *, cmd: str, cwd: str | None) -> ManagedShell:
+    async def start(self, *, cmd: str, cwd: str | None, session_id: str | None = None) -> ManagedShell:
         process = await asyncio.create_subprocess_shell(
             cmd,
             cwd=cwd,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
             executable=self.SHELL,
         )
-        shell = ManagedShell(shell_id=f"bash-{uuid.uuid4().hex[:8]}", cmd=cmd, cwd=cwd, process=process)
+        shell = ManagedShell(
+            shell_id=f"bash-{uuid.uuid4().hex[:8]}",
+            cmd=cmd,
+            cwd=cwd,
+            session_id=session_id,
+            process=process,
+        )
         shell.read_tasks.extend([
             asyncio.create_task(self._drain_stream(shell, process.stdout)),
             asyncio.create_task(self._drain_stream(shell, process.stderr)),
@@ -77,6 +84,13 @@ async def terminate(self, shell_id: str) -> ManagedShell:
         await self._finalize_shell(shell)
         return shell
 
+    async def terminate_session(self, session_id: str) -> int:
+        shell_ids = [shell.shell_id for shell in self._shells.values() if shell.session_id == session_id]
+        for shell_id in shell_ids:
+            with contextlib.suppress(KeyError):
+                await self.terminate(shell_id)
+        return len(shell_ids)
+
     async def wait_closed(self, shell_id: str) -> ManagedShell:
         shell = self.get(shell_id)
         if shell.returncode is None:

src/bub/builtin/tools.py

@@ -79,12 +79,17 @@ async def bash(
     """Run a shell command. Use background=true to keep it running and fetch output later via bash_output."""
     workspace = context.state.get("_runtime_workspace")
     target_cwd = cwd or workspace
-    shell = await shell_manager.start(cmd=cmd, cwd=target_cwd)
+    raw_session_id = context.state.get("session_id")
+    session_id = str(raw_session_id) if raw_session_id is not None else None
+    shell = await shell_manager.start(cmd=cmd, cwd=target_cwd, session_id=session_id)
     if background:
         return f"started: {shell.shell_id}"
     try:
         async with asyncio.timeout(timeout_seconds):
             shell = await shell_manager.wait_closed(shell.shell_id)
+    except asyncio.CancelledError:
+        await shell_manager.terminate(shell.shell_id)
+        raise
     except TimeoutError:
         await shell_manager.terminate(shell.shell_id)
         return f"command timed out after {timeout_seconds} seconds and was terminated"
@@ -309,7 +314,8 @@ def show_help() -> str:
 async def quit_tool(*, context: ToolContext) -> str:
     """Quit the tasks of the current session."""
     agent = _get_agent(context)
-    session_id = context.state.get("session_id", "temp/unknown")
+    session_id = str(context.state.get("session_id", "temp/unknown"))
+    await shell_manager.terminate_session(session_id)
     await agent.framework.quit_via_router(session_id)
     return "Session tasks stopped."
 

src/bub/channels/manager.py

@@ -118,11 +118,16 @@ def wrap_stream(self, message: Envelope, stream: AsyncIterable[StreamEvent]) ->
 
     async def quit(self, session_id: str) -> None:
         tasks = self._ongoing_tasks.pop(session_id, set())
+        current_task = asyncio.current_task()
+        cancelled_count = 0
         for task in tasks:
+            if task is current_task:
+                continue
             task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await task
-        logger.info(f"channel.manager quit session_id={session_id}, cancelled {len(tasks)} tasks")
+            cancelled_count += 1
+        logger.info(f"channel.manager quit session_id={session_id}, cancelled {cancelled_count} tasks")
 
     def enabled_channels(self) -> list[Channel]:
         if "all" in self._enabled_channels:
@@ -133,7 +138,10 @@ def enabled_channels(self) -> list[Channel]:
         ]
 
     def _on_task_done(self, session_id: str, task: asyncio.Task) -> None:
-        task.exception()  # to log any exception
+        if task.cancelled():
+            logger.info("channel.manager task cancelled session_id={}", session_id)
+        else:
+            task.exception()  # to log any exception
         tasks = self._ongoing_tasks.get(session_id, set())
         tasks.discard(task)
         if not tasks:

tests/test_builtin_tools.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import shlex
 import sys
+from types import SimpleNamespace
 
 import pytest
 from republic import ToolContext
@@ -11,11 +13,11 @@
 
 import bub.builtin.tools as builtin_tools
 from bub.builtin.shell_manager import ShellManager
-from bub.builtin.tools import bash, bash_output, kill_bash
+from bub.builtin.tools import bash, bash_output, kill_bash, quit_tool
 
 
-def _tool_context(tmp_path) -> ToolContext:
-    return ToolContext(tape="test-tape", run_id="test-run", state={"_runtime_workspace": str(tmp_path)})
+def _tool_context(tmp_path, **state) -> ToolContext:
+    return ToolContext(tape="test-tape", run_id="test-run", state={"_runtime_workspace": str(tmp_path), **state})
 
 
 def _python_shell(code: str) -> str:
@@ -51,6 +53,26 @@ async def test_foreground_bash_releases_shell_when_command_fails(tmp_path, monke
     assert manager._shells == {}
 
 
+@pytest.mark.asyncio
+async def test_foreground_bash_terminates_shell_when_cancelled(tmp_path, monkeypatch) -> None:
+    manager = ShellManager()
+    monkeypatch.setattr(builtin_tools, "shell_manager", manager)
+
+    task = asyncio.create_task(
+        bash.run(
+            cmd=_python_shell("import time; time.sleep(10)"),
+            context=_tool_context(tmp_path, session_id="session:target"),
+        )
+    )
+    await asyncio.sleep(0.1)
+
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+    assert manager._shells == {}
+
+
 @pytest.mark.asyncio
 async def test_bash_non_zero_exit_is_returned_as_tool_error(tmp_path) -> None:
     command = _python_shell("import sys; print('boom'); sys.exit(7)")
@@ -124,3 +146,46 @@ async def test_kill_bash_returns_status_when_process_already_finished(tmp_path)
     result = await kill_bash.run(shell_id=shell_id)
 
     assert result == f"id: {shell_id}\nstatus: exited\nexit_code: 0"
+
+
+@pytest.mark.asyncio
+async def test_quit_tool_terminates_background_shells_for_current_session(tmp_path, monkeypatch) -> None:
+    manager = ShellManager()
+    monkeypatch.setattr(builtin_tools, "shell_manager", manager)
+
+    target_started = await bash.run(
+        cmd=_python_shell("import time; time.sleep(10)"),
+        background=True,
+        context=_tool_context(tmp_path, session_id="session:target"),
+    )
+    target_shell_id = target_started.removeprefix("started: ").strip()
+    other_started = await bash.run(
+        cmd=_python_shell("import time; time.sleep(10)"),
+        background=True,
+        context=_tool_context(tmp_path, session_id="session:other"),
+    )
+    other_shell_id = other_started.removeprefix("started: ").strip()
+
+    class FakeFramework:
+        def __init__(self) -> None:
+            self.quit_sessions: list[str] = []
+
+        async def quit_via_router(self, session_id: str) -> None:
+            self.quit_sessions.append(session_id)
+
+    framework = FakeFramework()
+    context = _tool_context(
+        tmp_path,
+        session_id="session:target",
+        _runtime_agent=SimpleNamespace(framework=framework),
+    )
+
+    result = await quit_tool.run(context=context)
+
+    assert result == "Session tasks stopped."
+    assert framework.quit_sessions == ["session:target"]
+    with pytest.raises(KeyError, match="unknown shell id"):
+        await bash_output.run(shell_id=target_shell_id)
+    assert manager.get(other_shell_id).returncode is None
+
+    await kill_bash.run(shell_id=other_shell_id)

tests/test_channels.py

@@ -301,6 +301,45 @@ async def never_finish() -> None:
         await other_task
 
 
+@pytest.mark.asyncio
+async def test_channel_manager_quit_skips_current_task(load_config) -> None:
+    _load_channel_config(load_config, enabled_channels="telegram")
+    manager = ChannelManager(FakeFramework({"telegram": FakeChannel("telegram")}), enabled_channels=["telegram"])
+
+    async def never_finish() -> None:
+        await asyncio.sleep(10)
+
+    current_task = asyncio.current_task()
+    assert current_task is not None
+    target_task = asyncio.create_task(never_finish())
+    manager._ongoing_tasks["session:target"] = {current_task, target_task}
+
+    await manager.quit("session:target")
+
+    assert current_task.cancelled() is False
+    assert target_task.cancelled()
+    assert "session:target" not in manager._ongoing_tasks
+
+
+@pytest.mark.asyncio
+async def test_channel_manager_done_callback_handles_cancelled_task(load_config) -> None:
+    _load_channel_config(load_config, enabled_channels="telegram")
+    manager = ChannelManager(FakeFramework({"telegram": FakeChannel("telegram")}), enabled_channels=["telegram"])
+
+    async def never_finish() -> None:
+        await asyncio.sleep(10)
+
+    task = asyncio.create_task(never_finish())
+    manager._ongoing_tasks["session:target"] = {task}
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+    manager._on_task_done("session:target", task)
+
+    assert "session:target" not in manager._ongoing_tasks
+
+
 def test_cli_channel_normalize_input_prefixes_shell_commands() -> None:
     channel = CliChannel.__new__(CliChannel)
     channel._mode = "shell"