Cancel session background processes
This commit is contained in:
+45
-4
@@ -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,
|
||||||
|
|||||||
@@ -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
@@ -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()
|
||||||
|
|||||||
@@ -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))
|
||||||
|
|||||||
Reference in New Issue
Block a user