Fix interrupted run replay state

This commit is contained in:
武阳
2026-05-07 17:07:57 +08:00
parent 47d3678a18
commit 5b4f7c2c8d
3 changed files with 409 additions and 12 deletions
+192 -4
View File
@@ -257,8 +257,7 @@ class RunManager:
def cancel_latest(self, account_key: str, session_id: str) -> bool:
with self._lock:
run_id = self._latest_by_session.get((account_key, session_id))
record = self._runs.get(run_id or '')
record = self._select_session_record(account_key, session_id)
return self.cancel_record(record)
def cancel_run(self, run_id: str) -> bool:
@@ -782,7 +781,20 @@ def create_app(state: AgentState) -> FastAPI:
state._account_key(account_id),
safe_id,
)
return snapshot or {'session_id': safe_id, 'status': 'idle'}
if snapshot is not None:
return snapshot
directory = state.account_paths(account_id)['sessions']
try:
stored = load_agent_session(safe_id, directory=directory)
except FileNotFoundError:
return {'session_id': safe_id, 'status': 'idle'}
if _stored_session_has_incomplete_tail(stored):
return {
'session_id': safe_id,
'status': 'interrupted',
'current_stage': '上次运行已中断,后台没有正在执行的进程',
}
return {'session_id': safe_id, 'status': 'idle'}
@app.post('/api/runs/cancel')
async def cancel_run(payload: RunCancelRequest) -> dict[str, Any]:
@@ -797,6 +809,12 @@ def create_app(state: AgentState) -> FastAPI:
safe_id,
)
)
if cancelled:
_mark_session_interrupted(
state.account_paths(payload.account_id)['sessions'],
safe_id,
status='cancelled',
)
return {'session_id': safe_id, 'cancelled': cancelled}
def _run_chat_payload(
@@ -875,6 +893,7 @@ def create_app(state: AgentState) -> FastAPI:
status_code=404,
detail='Session to resume not found',
)
stored = _sanitize_stored_session_for_resume(stored)
result = agent.resume(
request.prompt.strip(),
stored,
@@ -889,9 +908,23 @@ def create_app(state: AgentState) -> FastAPI:
event_sink=event_sink,
)
except Exception as exc:
_mark_session_interrupted(
session_directory,
requested_session_id,
status=(
'cancelled'
if run_record.cancel_event.is_set()
else 'failed'
),
detail=str(exc),
)
state.run_manager.update(
run_record.run_id,
status='failed',
status=(
'cancelled'
if run_record.cancel_event.is_set()
else 'failed'
),
error=str(exc),
)
raise
@@ -923,6 +956,12 @@ def create_app(state: AgentState) -> FastAPI:
run_record.run_id,
'cancelled' if run_record.cancel_event.is_set() else 'completed',
)
if run_record.cancel_event.is_set() and result.session_id:
_mark_session_interrupted(
session_directory,
result.session_id,
status='cancelled',
)
return payload
# ------------- chat ------------------------------------------------------
@@ -1048,6 +1087,9 @@ def _normalize_transcript_entry(entry: dict[str, Any]) -> dict[str, Any]:
for key in ('name', 'tool_call_id', 'tool_calls', 'metadata', 'message_id'):
if key in entry and entry[key] not in (None, '', [], {}):
out[key] = entry[key]
for key in ('state', 'stop_reason'):
if key in entry and entry[key] not in (None, ''):
out[key] = entry[key]
return out
@@ -1063,6 +1105,152 @@ def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
}
def _stored_session_has_incomplete_tail(stored: StoredAgentSession) -> bool:
_, changed = _trim_incomplete_message_tail(stored.messages)
if changed:
return True
budget_state = stored.budget_state if isinstance(stored.budget_state, dict) else {}
return budget_state.get('status') in {'queued', 'running'}
def _mark_session_interrupted(
directory: Path,
session_id: str,
*,
status: str,
detail: str = '',
) -> None:
try:
stored = load_agent_session(session_id, directory=directory)
except (FileNotFoundError, OSError, json.JSONDecodeError):
return
trimmed, changed = _trim_incomplete_message_tail(stored.messages)
budget_state = (
dict(stored.budget_state)
if isinstance(stored.budget_state, dict)
else {}
)
if not changed and budget_state.get('status') not in {'queued', 'running'}:
return
messages = list(trimmed)
if not _last_message_is_run_status(messages, status):
messages.append(
{
'role': 'assistant',
'content': _interrupted_session_message(status),
'state': 'final',
'stop_reason': status,
'metadata': {
'kind': 'run_status',
'status': status,
'detail': detail[:500],
'created_at_ms': int(time.time() * 1000),
},
}
)
budget_state['status'] = status
budget_state['interrupted_at'] = int(time.time())
save_agent_session(
replace(
stored,
messages=tuple(messages),
budget_state=budget_state,
),
directory=directory,
)
def _sanitize_stored_session_for_resume(
stored: StoredAgentSession,
) -> StoredAgentSession:
trimmed, changed = _trim_incomplete_message_tail(stored.messages)
if not changed:
return stored
messages = list(trimmed)
if not _last_message_is_run_status(messages, 'interrupted'):
messages.append(
{
'role': 'assistant',
'content': _interrupted_session_message('interrupted'),
'state': 'final',
'stop_reason': 'interrupted',
'metadata': {
'kind': 'run_status',
'status': 'interrupted',
'created_at_ms': int(time.time() * 1000),
},
}
)
budget_state = (
dict(stored.budget_state)
if isinstance(stored.budget_state, dict)
else {}
)
budget_state['status'] = 'interrupted'
return replace(
stored,
messages=tuple(messages),
budget_state=budget_state,
)
def _trim_incomplete_message_tail(
messages: tuple[dict[str, Any], ...] | tuple[Any, ...],
) -> tuple[tuple[dict[str, Any], ...], bool]:
trimmed = [dict(message) for message in messages if isinstance(message, dict)]
changed = False
while trimmed and _is_incomplete_session_message(trimmed[-1]):
trimmed.pop()
changed = True
while trimmed and _assistant_has_trailing_tool_call(trimmed[-1]):
trimmed.pop()
changed = True
return tuple(trimmed), changed
def _is_incomplete_session_message(message: dict[str, Any]) -> bool:
state = message.get('state')
if isinstance(state, str) and state not in {'', 'final'}:
return True
metadata = message.get('metadata')
if isinstance(metadata, dict):
if (
message.get('role') != 'user'
and metadata.get('placeholder') is True
and metadata.get('status') == 'running'
):
return True
if metadata.get('phase') in {'starting', 'running'} and message.get('role') == 'tool':
return True
return False
def _assistant_has_trailing_tool_call(message: dict[str, Any]) -> bool:
if message.get('role') != 'assistant':
return False
tool_calls = message.get('tool_calls')
return isinstance(tool_calls, list) and bool(tool_calls)
def _last_message_is_run_status(messages: list[dict[str, Any]], status: str) -> bool:
if not messages:
return False
metadata = messages[-1].get('metadata')
return (
isinstance(metadata, dict)
and metadata.get('kind') == 'run_status'
and metadata.get('status') == status
)
def _interrupted_session_message(status: str) -> str:
if status == 'cancelled':
return '上一次任务已取消,后台没有正在执行的进程。你可以继续回复,或重新发起任务。'
if status == 'failed':
return '上一次任务异常中断,后台没有正在执行的进程。你可以继续回复,或重新发起任务。'
return '上一次任务已中断,后台没有正在执行的进程。你可以继续回复,或重新发起任务。'
def _emit_runtime_event(
event_sink: Any | None,
event: dict[str, object],