wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix run cancellation and live activity stability

Browse Source
This commit is contained in:
wuyang6
2026-05-14 12:23:53 +08:00
parent 105d287ebd
commit 22960c5e02
6 changed files with 117 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+39 -27
View File
@@ -443,6 +443,28 @@ class RunManager:
record = self._runs.get(run_id)
return self.cancel_record(record)
def cancel_session(self, account_key: str, session_id: str) -> list[str]:
"""Cancel every non-terminal run for a session.
The UI can hold a stale run_id after refresh/replay while a newer
request is queued behind the same session lock. Cancelling only the
stale id leaves the queued run alive and makes the composer look
impossible to stop. Session-level cancel is the user-facing intent.
"""
with self._lock:
records = [
record
for record in self._runs.values()
if record.account_key == account_key
and record.session_id == session_id
and record.status not in {'completed', 'failed', 'cancelled'}
]
cancelled: list[str] = []
for record in records:
if self.cancel_record(record):
cancelled.append(record.run_id)
return cancelled
def cancel_record(self, record: RunRecord | None) -> bool:
if record is None or record.status in {'completed', 'failed', 'cancelled'}:
return False
@@ -1635,40 +1657,30 @@ def create_app(state: AgentState) -> FastAPI:
safe_id = _safe_session_id(payload.session_id)
if safe_id is None:
raise HTTPException(status_code=400, detail='Invalid session id')
cancelled = (
state.run_manager.cancel_run(payload.run_id)
if payload.run_id
else state.run_manager.cancel_latest(
state._account_key(payload.account_id),
safe_id,
)
account_key = state._account_key(payload.account_id)
cancelled_ids: set[str] = set()
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))
stored_cancelled = state.run_state_store.finish_active_for_session(
account_key,
safe_id,
status='cancelled',
stage='用户已取消',
)
cancelled_ids.update(stored_cancelled)
cancelled = bool(cancelled_ids)
if cancelled:
run_id = payload.run_id
if not run_id:
snapshot = state.run_state_store.snapshot_latest(
state._account_key(payload.account_id),
safe_id,
)
run_id = str(snapshot.get('run_id') or '') if snapshot else ''
if run_id:
stored = state.run_state_store.snapshot_run(run_id)
elapsed_ms = None
started_at = stored.get('started_at') if stored else None
if isinstance(started_at, (int, float)):
elapsed_ms = max(0, int((time.time() - float(started_at)) * 1000))
state.run_state_store.finish(
run_id,
status='cancelled',
elapsed_ms=elapsed_ms,
stage='用户已取消',
)
_mark_session_interrupted(
state.account_paths(payload.account_id)['sessions'],
safe_id,
status='cancelled',
)
return {'session_id': safe_id, 'cancelled': cancelled}
return {
'session_id': safe_id,
'cancelled': cancelled,
'cancelled_run_ids': sorted(cancelled_ids),
}
def _run_chat_payload(
request: ChatRequest,