Cancel session background processes

backend/api/server.py

@@ -2322,6 +2322,11 @@ def create_app(state: AgentState) -> FastAPI:
         if payload.run_id and state.run_manager.cancel_run(payload.run_id):
             cancelled_ids.add(payload.run_id)
         cancelled_ids.update(state.run_manager.cancel_session(account_key, safe_id))
+        cancelled_bg_task_ids = await _bash_bg_manager.cancel_session(
+            state,
+            account_key,
+            safe_id,
+        )
         stored_cancelled = state.run_state_store.finish_active_for_session(
             account_key,
             safe_id,
@@ -2329,7 +2334,7 @@ def create_app(state: AgentState) -> FastAPI:
             stage='用户已取消',
         )
         cancelled_ids.update(stored_cancelled)
-        cancelled = bool(cancelled_ids)
+        cancelled = bool(cancelled_ids or cancelled_bg_task_ids)
         if cancelled:
             _mark_session_interrupted(
                 state.account_paths(payload.account_id)['sessions'],
@@ -2340,6 +2345,7 @@ def create_app(state: AgentState) -> FastAPI:
             'session_id': safe_id,
             'cancelled': cancelled,
             'cancelled_run_ids': sorted(cancelled_ids),
+            'cancelled_bg_task_ids': sorted(cancelled_bg_task_ids),
         }
 
     def _run_chat_payload(
@@ -5832,6 +5838,26 @@ class _BashBackgroundManager:
             existing.cancel()
         return True
 
+    async def cancel_session(
+        self,
+        state: 'AgentState',
+        account_key: str,
+        session_id: str,
+    ) -> list[str]:
+        rows = state.bash_bg_store.list_active_for_session(account_key, session_id)
+        cancelled: list[str] = []
+        for row in rows:
+            task_id = str(row.get('task_id') or '')
+            if not task_id:
+                continue
+            await self._kill_remote_or_local(state, row)
+            state.bash_bg_store.mark_killed(task_id)
+            existing = self._tasks.pop(task_id, None)
+            if existing is not None:
+                existing.cancel()
+            cancelled.append(task_id)
+        return cancelled
+
     def status_query(
         self, state: 'AgentState', task_id: str
     ) -> 'BgTaskStatus | None':
@@ -6025,9 +6051,16 @@ class _BashBackgroundManager:
         pid = int(row['pid'])
         if runtime is not None:
             cmd = (
-                f'kill -- -$(ps -o pgid= -p {pid} 2>/dev/null '
-                f'| tr -d " ") 2>/dev/null; '
-                f'kill {pid} 2>/dev/null; true'
+                f'pgid=$(ps -o pgid= -p {pid} 2>/dev/null | tr -d " "); '
+                'if [ -n "$pgid" ]; then '
+                'kill -TERM -- "-$pgid" 2>/dev/null || true; '
+                'fi; '
+                f'kill -TERM {pid} 2>/dev/null || true; '
+                'sleep 0.5; '
+                'if [ -n "$pgid" ]; then '
+                'kill -KILL -- "-$pgid" 2>/dev/null || true; '
+                'fi; '
+                f'kill -KILL {pid} 2>/dev/null || true'
             )
             try:
                 await asyncio.to_thread(
@@ -6047,6 +6080,14 @@ class _BashBackgroundManager:
                 os.kill(pid, signal.SIGTERM)
             except (ProcessLookupError, PermissionError, OSError):
                 pass
+        time.sleep(0.2)
+        try:
+            os.killpg(os.getpgid(pid), signal.SIGKILL)
+        except (ProcessLookupError, PermissionError, OSError):
+            try:
+                os.kill(pid, signal.SIGKILL)
+            except (ProcessLookupError, PermissionError, OSError):
+                pass
 
     async def _maybe_auto_resume(
         self, state: 'AgentState', spec: BgTaskSpec, exit_code: int,

src/bash_bg_store.py

@@ -156,6 +156,24 @@ class BashBgStore:
             ).fetchall()
             return [dict(row) for row in rows]
 
+    def list_active_for_session(
+        self,
+        account_key: str,
+        session_id: str,
+    ) -> list[dict[str, Any]]:
+        with self._connect() as conn:
+            rows = conn.execute(
+                """
+                select * from bash_bg_tasks
+                where account_key = ?
+                  and session_id = ?
+                  and status = 'running'
+                order by started_at desc
+                """,
+                (account_key, session_id),
+            ).fetchall()
+            return [dict(row) for row in rows]
+
     def get(self, task_id: str) -> dict[str, Any] | None:
         with self._connect() as conn:
             row = conn.execute(

src/jupyter_runtime.py

@@ -536,13 +536,28 @@ PY
             raise JupyterRuntimeError('Remote command is empty.')
         marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__'
         workspace = shlex.quote(self.binding.workspace_cwd)
+        run_dir = f'{self.binding.workspace_cwd}/scratchpad/.run_pids'
+        quoted_run_dir = shlex.quote(run_dir)
+        pid_file = shlex.quote(f'{run_dir}/{marker}.pid')
         env_prefix = self._remote_env_prefix()
         inner_command = f'{env_prefix}{command}'
         wrapped = (
-            f'mkdir -p {workspace} && '
+            f'mkdir -p {workspace} {quoted_run_dir} && '
             f'cd {workspace} && '
-            f'bash -lc {shlex.quote(inner_command)}; '
-            f'printf "\\n{marker}:%s\\n" "$?"'
+            f'( setsid bash -lc {shlex.quote(inner_command)} ) & '
+            '__zk_pid=$!; '
+            f'printf "%s\\n" "$__zk_pid" > {pid_file}; '
+            'cleanup() { '
+            'kill -TERM -- -"${__zk_pid}" 2>/dev/null '
+            '|| kill -TERM "${__zk_pid}" 2>/dev/null || true; '
+            'sleep 0.2; '
+            'kill -KILL -- -"${__zk_pid}" 2>/dev/null || true; '
+            '}; '
+            'trap "cleanup" INT TERM; '
+            'wait "$__zk_pid"; __zk_status=$?; '
+            'trap - INT TERM; '
+            f'rm -f {pid_file}; '
+            f'printf "\\n{marker}:%s\\n" "$__zk_status"'
         )
         with self._lock:
             name = self._ensure_terminal()

tests/test_gui_server.py

@@ -20,6 +20,7 @@ from fastapi.testclient import TestClient
 from backend.api.server import (
     AgentState,
     create_app,
+    _bash_bg_manager,
     _derive_initial_session_title,
     _ensure_session_title,
     _generate_session_title,
@@ -28,6 +29,7 @@ from backend.api.server import (
     _sanitize_stored_session_for_resume,
 )
 from src.agent_types import AgentRunResult, AssistantTurn
+from src.bash_bg_store import BgTaskSpec
 from src.session_store import StoredAgentSession, load_agent_session, save_agent_session
 
 
@@ -305,6 +307,48 @@ class GuiServerTests(unittest.TestCase):
             self.assertEqual(latest.status_code, 200)
             self.assertEqual(latest.json()['status'], 'cancelled')
 
+    def test_cancel_run_kills_active_background_tasks_for_session(self) -> None:
+        with tempfile.TemporaryDirectory() as d:
+            client, state = _build_client(Path(d))
+            spec = BgTaskSpec(
+                task_id='bg_cancel_me',
+                account_key='alice',
+                session_id='thread-1',
+                run_id='old-run',
+                pid=12345,
+                task_dir='/tmp/bg_cancel_me',
+                output_path='/tmp/bg_cancel_me/output',
+                pid_path='/tmp/bg_cancel_me/pid',
+                exit_code_path='/tmp/bg_cancel_me/exit_code',
+                command='sleep 3600',
+                started_at=123.0,
+                wait_for_completion=True,
+            )
+            state.bash_bg_store.record_start(spec)
+            killed: list[str] = []
+
+            async def fake_kill(_state: AgentState, row: dict[str, object]) -> None:
+                killed.append(str(row['task_id']))
+
+            original = _bash_bg_manager._kill_remote_or_local
+            _bash_bg_manager._kill_remote_or_local = fake_kill  # type: ignore[method-assign]
+            try:
+                response = client.post(
+                    '/api/runs/cancel',
+                    json={'account_id': 'alice', 'session_id': 'thread-1'},
+                )
+            finally:
+                _bash_bg_manager._kill_remote_or_local = original  # type: ignore[method-assign]
+
+            self.assertEqual(response.status_code, 200)
+            payload = response.json()
+            self.assertTrue(payload['cancelled'])
+            self.assertEqual(payload['cancelled_bg_task_ids'], ['bg_cancel_me'])
+            self.assertEqual(killed, ['bg_cancel_me'])
+            row = state.bash_bg_store.get('bg_cancel_me')
+            self.assertIsNotNone(row)
+            self.assertEqual(row['status'], 'cancelled')
+
     def test_latest_run_prefers_running_over_queued_for_session(self) -> None:
         with tempfile.TemporaryDirectory() as d:
             client, state = _build_client(Path(d))