Cancel session background processes

This commit is contained in:
wuyang6
2026-06-12 16:24:25 +08:00
parent 953126e1d3
commit 0a3e044c9d
4 changed files with 125 additions and 7 deletions
+45 -4
View File
@@ -2322,6 +2322,11 @@ def create_app(state: AgentState) -> FastAPI:
if payload.run_id and state.run_manager.cancel_run(payload.run_id): if payload.run_id and state.run_manager.cancel_run(payload.run_id):
cancelled_ids.add(payload.run_id) cancelled_ids.add(payload.run_id)
cancelled_ids.update(state.run_manager.cancel_session(account_key, safe_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( stored_cancelled = state.run_state_store.finish_active_for_session(
account_key, account_key,
safe_id, safe_id,
@@ -2329,7 +2334,7 @@ def create_app(state: AgentState) -> FastAPI:
stage='用户已取消', stage='用户已取消',
) )
cancelled_ids.update(stored_cancelled) cancelled_ids.update(stored_cancelled)
cancelled = bool(cancelled_ids) cancelled = bool(cancelled_ids or cancelled_bg_task_ids)
if cancelled: if cancelled:
_mark_session_interrupted( _mark_session_interrupted(
state.account_paths(payload.account_id)['sessions'], state.account_paths(payload.account_id)['sessions'],
@@ -2340,6 +2345,7 @@ def create_app(state: AgentState) -> FastAPI:
'session_id': safe_id, 'session_id': safe_id,
'cancelled': cancelled, 'cancelled': cancelled,
'cancelled_run_ids': sorted(cancelled_ids), 'cancelled_run_ids': sorted(cancelled_ids),
'cancelled_bg_task_ids': sorted(cancelled_bg_task_ids),
} }
def _run_chat_payload( def _run_chat_payload(
@@ -5832,6 +5838,26 @@ class _BashBackgroundManager:
existing.cancel() existing.cancel()
return True 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( def status_query(
self, state: 'AgentState', task_id: str self, state: 'AgentState', task_id: str
) -> 'BgTaskStatus | None': ) -> 'BgTaskStatus | None':
@@ -6025,9 +6051,16 @@ class _BashBackgroundManager:
pid = int(row['pid']) pid = int(row['pid'])
if runtime is not None: if runtime is not None:
cmd = ( cmd = (
f'kill -- -$(ps -o pgid= -p {pid} 2>/dev/null ' f'pgid=$(ps -o pgid= -p {pid} 2>/dev/null | tr -d " "); '
f'| tr -d " ") 2>/dev/null; ' 'if [ -n "$pgid" ]; then '
f'kill {pid} 2>/dev/null; true' '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: try:
await asyncio.to_thread( await asyncio.to_thread(
@@ -6047,6 +6080,14 @@ class _BashBackgroundManager:
os.kill(pid, signal.SIGTERM) os.kill(pid, signal.SIGTERM)
except (ProcessLookupError, PermissionError, OSError): except (ProcessLookupError, PermissionError, OSError):
pass 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( async def _maybe_auto_resume(
self, state: 'AgentState', spec: BgTaskSpec, exit_code: int, self, state: 'AgentState', spec: BgTaskSpec, exit_code: int,
+18
View File
@@ -156,6 +156,24 @@ class BashBgStore:
).fetchall() ).fetchall()
return [dict(row) for row in rows] 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: def get(self, task_id: str) -> dict[str, Any] | None:
with self._connect() as conn: with self._connect() as conn:
row = conn.execute( row = conn.execute(
+18 -3
View File
@@ -536,13 +536,28 @@ PY
raise JupyterRuntimeError('Remote command is empty.') raise JupyterRuntimeError('Remote command is empty.')
marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__' marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__'
workspace = shlex.quote(self.binding.workspace_cwd) 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() env_prefix = self._remote_env_prefix()
inner_command = f'{env_prefix}{command}' inner_command = f'{env_prefix}{command}'
wrapped = ( wrapped = (
f'mkdir -p {workspace} && ' f'mkdir -p {workspace} {quoted_run_dir} && '
f'cd {workspace} && ' f'cd {workspace} && '
f'bash -lc {shlex.quote(inner_command)}; ' f'( setsid bash -lc {shlex.quote(inner_command)} ) & '
f'printf "\\n{marker}:%s\\n" "$?"' '__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: with self._lock:
name = self._ensure_terminal() name = self._ensure_terminal()
+44
View File
@@ -20,6 +20,7 @@ from fastapi.testclient import TestClient
from backend.api.server import ( from backend.api.server import (
AgentState, AgentState,
create_app, create_app,
_bash_bg_manager,
_derive_initial_session_title, _derive_initial_session_title,
_ensure_session_title, _ensure_session_title,
_generate_session_title, _generate_session_title,
@@ -28,6 +29,7 @@ from backend.api.server import (
_sanitize_stored_session_for_resume, _sanitize_stored_session_for_resume,
) )
from src.agent_types import AgentRunResult, AssistantTurn 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 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.status_code, 200)
self.assertEqual(latest.json()['status'], 'cancelled') 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: def test_latest_run_prefers_running_over_queued_for_session(self) -> None:
with tempfile.TemporaryDirectory() as d: with tempfile.TemporaryDirectory() as d:
client, state = _build_client(Path(d)) client, state = _build_client(Path(d))