From 0a3e044c9daf4c8fa447df4aa572754beda843fb Mon Sep 17 00:00:00 2001 From: wuyang6 Date: Fri, 12 Jun 2026 16:24:25 +0800 Subject: [PATCH] Cancel session background processes --- backend/api/server.py | 49 ++++++++++++++++++++++++++++++++++++---- src/bash_bg_store.py | 18 +++++++++++++++ src/jupyter_runtime.py | 21 ++++++++++++++--- tests/test_gui_server.py | 44 ++++++++++++++++++++++++++++++++++++ 4 files changed, 125 insertions(+), 7 deletions(-) diff --git a/backend/api/server.py b/backend/api/server.py index 41b2acb..7f40d26 100644 --- a/backend/api/server.py +++ b/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, diff --git a/src/bash_bg_store.py b/src/bash_bg_store.py index ef055d7..5a86709 100644 --- a/src/bash_bg_store.py +++ b/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( diff --git a/src/jupyter_runtime.py b/src/jupyter_runtime.py index 7632980..ecefd9c 100644 --- a/src/jupyter_runtime.py +++ b/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() diff --git a/tests/test_gui_server.py b/tests/test_gui_server.py index fedabe1..11cc38c 100644 --- a/tests/test_gui_server.py +++ b/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))