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

Fix interrupted run replay state

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+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],
frontend/app/components/assistant-ui/thread-list.tsx
+114 -8
View File
@@ -35,8 +35,13 @@ type ClawStoredMessage = {
name?: string;
tool_call_id?: string;
tool_calls?: ClawStoredToolCall[];
state?: string;
stop_reason?: string;
metadata?: {
elapsed_ms?: unknown;
placeholder?: unknown;
phase?: unknown;
status?: unknown;
};
};
@@ -48,11 +53,22 @@ type ClawStoredSession = {
export type ClawRunStatus = {
run_id?: string;
session_id?: string;
status?: "idle" | "queued" | "running" | "completed" | "failed" | "cancelled";
status?:
| "idle"
| "queued"
| "running"
| "completed"
| "failed"
| "cancelled"
| "interrupted";
current_stage?: string;
started_at?: number;
};
type ClawActiveRunStatus = ClawRunStatus & {
status: "queued" | "running";
};
type ClawStoredToolCall = {
id?: string;
name?: string;
@@ -364,12 +380,12 @@ export function toReplayRepository(
runStatus?: ClawRunStatus | null,
): ExportedMessageRepository {
const messages: ExportedMessageRepository["messages"] = [];
const toolResults = collectToolResults(session.messages ?? []);
const lastAssistantContentIndex = findLastAssistantContentIndex(
session.messages ?? [],
);
const replayMessages = getReplayMessages(session.messages ?? [], runStatus);
const toolResults = collectToolResults(replayMessages);
const lastAssistantContentIndex =
findLastAssistantContentIndex(replayMessages);
let previousId: string | null = null;
for (const [index, message] of (session.messages ?? []).entries()) {
for (const [index, message] of replayMessages.entries()) {
if (message.role === "tool") continue;
if (message.role !== "user" && message.role !== "assistant") continue;
const content = cleanStoredContent(message.content ?? "");
@@ -421,6 +437,37 @@ export function toReplayRepository(
messages.push({ message: threadMessage, parentId: previousId });
previousId = id;
}
if (shouldShowInterruptedNotice(session.messages ?? [], runStatus)) {
const id = `${fallbackSessionId}-interrupted`;
const text =
runStatus?.status === "cancelled"
? "上一次任务已取消，后台没有正在执行的进程。你可以继续回复，或重新发起任务。"
: "上一次任务已中断，后台没有正在执行的进程。你可以继续回复，或重新发起任务。";
const parts = [{ type: "text", text }] as UIMessage["parts"];
const uiMessage: UIMessage = {
id,
role: "assistant",
parts,
metadata: { sessionId: session.session_id ?? fallbackSessionId },
} as UIMessage;
const threadMessage = {
id,
role: "assistant",
createdAt: new Date(),
content: toThreadMessageContent(parts),
status: { type: "complete", reason: "stop" },
metadata: {
unstable_state: null,
unstable_annotations: [],
unstable_data: [],
steps: [],
custom: { sessionId: session.session_id ?? fallbackSessionId },
},
} as ThreadMessage;
bindExternalStoreMessage(threadMessage, uiMessage);
messages.push({ message: threadMessage, parentId: previousId });
previousId = id;
}
if (isActiveRunStatus(runStatus)) {
const id = `${fallbackSessionId}-run-${runStatus.run_id ?? "active"}`;
const text =
@@ -468,10 +515,65 @@ export function toReplayRepository(
function isActiveRunStatus(
runStatus?: ClawRunStatus | null,
): runStatus is ClawRunStatus {
): runStatus is ClawActiveRunStatus {
return runStatus?.status === "queued" || runStatus?.status === "running";
}
function getReplayMessages(
messages: readonly ClawStoredMessage[],
runStatus?: ClawRunStatus | null,
) {
if (isActiveRunStatus(runStatus)) return [...messages];
return trimIncompleteReplayTail(messages);
}
function shouldShowInterruptedNotice(
messages: readonly ClawStoredMessage[],
runStatus?: ClawRunStatus | null,
) {
if (isActiveRunStatus(runStatus)) return false;
return (
runStatus?.status === "interrupted" ||
runStatus?.status === "cancelled" ||
trimIncompleteReplayTail(messages).length !== messages.length
);
}
function trimIncompleteReplayTail(messages: readonly ClawStoredMessage[]) {
const trimmed = [...messages];
while (trimmed.length && isIncompleteReplayMessage(trimmed.at(-1))) {
trimmed.pop();
}
while (trimmed.length && hasTrailingAssistantToolCall(trimmed.at(-1))) {
trimmed.pop();
}
return trimmed;
}
function isIncompleteReplayMessage(message?: ClawStoredMessage) {
if (!message) return false;
if (message.state && message.state !== "final") return true;
if (
message.role === "tool" &&
(message.metadata?.phase === "starting" ||
message.metadata?.phase === "running")
) {
return true;
}
if (
message.role !== "user" &&
message.metadata?.placeholder === true &&
message.metadata?.status === "running"
) {
return true;
}
return false;
}
function hasTrailingAssistantToolCall(message?: ClawStoredMessage) {
return message?.role === "assistant" && Boolean(message.tool_calls?.length);
}
function normalizeRunTimestampMs(value: unknown) {
if (typeof value !== "number" || !Number.isFinite(value)) return undefined;
return value < 10_000_000_000 ? Math.round(value * 1000) : Math.round(value);
@@ -499,7 +601,11 @@ function cleanStoredContent(content: string) {
function collectToolResults(messages: readonly ClawStoredMessage[]) {
const results = new Map<string, string>();
for (const message of messages) {
if (message.role === "tool" && message.tool_call_id) {
if (
message.role === "tool" &&
message.tool_call_id &&
!isIncompleteReplayMessage(message)
) {
results.set(message.tool_call_id, message.content ?? "");
}
}
tests/test_gui_server.py
+103
View File
@@ -257,6 +257,109 @@ class GuiServerTests(unittest.TestCase):
self.assertEqual(payload['status'], 'running')
self.assertEqual(payload['current_stage'], 'first stage')
def test_cancel_run_marks_streaming_session_tail_cancelled(self) -> None:
with tempfile.TemporaryDirectory() as d:
root = Path(d)
client, state = _build_client(root)
session_root = root / 'accounts' / 'alice' / 'sessions'
save_agent_session(
StoredAgentSession(
session_id='thread-1',
model_config={'model': 'test-model'},
runtime_config={'cwd': str(root)},
system_prompt_parts=('system prompt',),
user_context={},
system_context={},
messages=(
{'role': 'user', 'content': 'hello'},
{
'role': 'assistant',
'content': '我先看看文件',
'tool_calls': [
{
'id': 'call-1',
'function': {'name': 'list_dir', 'arguments': '{}'},
}
],
},
{
'role': 'tool',
'content': '',
'tool_call_id': 'call-1',
'state': 'streaming',
'metadata': {'phase': 'starting'},
},
),
turns=1,
tool_calls=1,
usage={},
total_cost_usd=0.0,
file_history=(),
budget_state={'status': 'running'},
plugin_state={},
),
directory=session_root,
)
state.run_manager.start('alice', 'thread-1')
response = client.post(
'/api/runs/cancel',
json={'account_id': 'alice', 'session_id': 'thread-1'},
)
self.assertEqual(response.status_code, 200)
self.assertTrue(response.json()['cancelled'])
stored = load_agent_session('thread-1', directory=session_root)
self.assertEqual(stored.messages[-1]['role'], 'assistant')
self.assertEqual(stored.messages[-1]['stop_reason'], 'cancelled')
self.assertEqual(
stored.messages[-1]['metadata']['status'],
'cancelled',
)
self.assertNotEqual(stored.messages[-1].get('state'), 'streaming')
def test_latest_run_reports_interrupted_for_orphaned_streaming_session(self) -> None:
with tempfile.TemporaryDirectory() as d:
root = Path(d)
client, _ = _build_client(root)
session_root = root / 'accounts' / 'alice' / 'sessions'
save_agent_session(
StoredAgentSession(
session_id='thread-1',
model_config={'model': 'test-model'},
runtime_config={'cwd': str(root)},
system_prompt_parts=('system prompt',),
user_context={},
system_context={},
messages=(
{'role': 'user', 'content': 'hello'},
{
'role': 'tool',
'content': '',
'tool_call_id': 'call-1',
'state': 'streaming',
'metadata': {'phase': 'starting'},
},
),
turns=1,
tool_calls=1,
usage={},
total_cost_usd=0.0,
file_history=(),
budget_state={'status': 'running'},
plugin_state={},
),
directory=session_root,
)
response = client.get(
'/api/runs/latest',
params={'account_id': 'alice', 'session_id': 'thread-1'},
)
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json()['status'], 'interrupted')
def test_chat_uses_account_scoped_model_config(self) -> None:
with tempfile.TemporaryDirectory() as d:
client, _ = _build_client(Path(d))