Preserve running session replay state

frontend/app/app/api/claw/runs/cancel/route.ts

@@ -0,0 +1,37 @@
+import { getCurrentAccount } from "@/lib/claw-auth";
+
+const CLAW_API_URL = process.env.CLAW_API_URL ?? "http://127.0.0.1:8765";
+
+export async function POST(req: Request) {
+	const account = await getCurrentAccount();
+	if (!account)
+		return Response.json({ error: "请先登录账号" }, { status: 401 });
+
+	const body = (await req.json().catch(() => ({}))) as {
+		session_id?: unknown;
+		run_id?: unknown;
+	};
+	const sessionId =
+		typeof body.session_id === "string" ? body.session_id.trim() : "";
+	if (!sessionId)
+		return Response.json({ error: "Missing session_id" }, { status: 400 });
+
+	const runId = typeof body.run_id === "string" ? body.run_id.trim() : "";
+	const response = await fetch(`${CLAW_API_URL}/api/runs/cancel`, {
+		method: "POST",
+		headers: { "content-type": "application/json" },
+		body: JSON.stringify({
+			account_id: account.id,
+			session_id: sessionId,
+			...(runId ? { run_id: runId } : {}),
+		}),
+	});
+	const payload = await response.text();
+	return new Response(payload, {
+		status: response.status,
+		headers: {
+			"content-type":
+				response.headers.get("content-type") ?? "application/json",
+		},
+	});
+}

frontend/app/components/assistant-ui/thread-list.tsx

@@ -50,6 +50,7 @@ export type ClawRunStatus = {
 	session_id?: string;
 	status?: "idle" | "queued" | "running" | "completed" | "failed" | "cancelled";
 	current_stage?: string;
+	started_at?: number;
 };
 
 type ClawStoredToolCall = {
@@ -431,7 +432,11 @@ export function toReplayRepository(
 			id,
 			role: "assistant",
 			parts,
-			metadata: { sessionId: session.session_id ?? fallbackSessionId },
+			metadata: {
+				sessionId: session.session_id ?? fallbackSessionId,
+				runId: runStatus.run_id,
+				runStartedAtMs: normalizeRunTimestampMs(runStatus.started_at),
+			},
 		} as UIMessage;
 		const threadMessage = {
 			id,
@@ -444,7 +449,11 @@ export function toReplayRepository(
 				unstable_annotations: [],
 				unstable_data: [],
 				steps: [],
-				custom: { sessionId: session.session_id ?? fallbackSessionId },
+				custom: {
+					sessionId: session.session_id ?? fallbackSessionId,
+					runId: runStatus.run_id,
+					runStartedAtMs: normalizeRunTimestampMs(runStatus.started_at),
+				},
 			},
 		} as ThreadMessage;
 		bindExternalStoreMessage(threadMessage, uiMessage);
@@ -463,6 +472,11 @@ function isActiveRunStatus(
 	return runStatus?.status === "queued" || runStatus?.status === "running";
 }
 
+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);
+}
+
 function findLastAssistantContentIndex(messages: readonly ClawStoredMessage[]) {
 	for (let index = messages.length - 1; index >= 0; index -= 1) {
 		const message = messages[index];

frontend/app/components/assistant-ui/thread.tsx

@@ -157,21 +157,11 @@ export const Thread: FC = () => {
 
 function useRefreshReplayedRun() {
 	const { replaySession } = useClawSessionReplay();
-	const isRunning = useAuiState((s) => s.thread.isRunning);
-	const hasReplayedRunMessage = useAuiState((s) =>
-		s.thread.messages.some((message) => message.id.includes("-run-")),
-	);
-	const latestSessionId = useAuiState((s) =>
-		findLatestMessageMetadataValue(s.thread.messages, "sessionId"),
-	);
+	const replayedRun = useReplayedRunState();
 
 	useEffect(() => {
-		if (!isRunning || !hasReplayedRunMessage) return;
-		const sessionId = normalizeSessionId(
-			latestSessionId ?? window.localStorage.getItem("claw.activeSessionId"),
-		);
-		if (!sessionId) return;
-		const activeSessionId = sessionId;
+		if (!replayedRun.active || !replayedRun.sessionId) return;
+		const activeSessionId = replayedRun.sessionId;
 		let cancelled = false;
 
 		async function refreshIfFinished() {
@@ -198,7 +188,41 @@ function useRefreshReplayedRun() {
 			cancelled = true;
 			window.clearInterval(interval);
 		};
-	}, [hasReplayedRunMessage, isRunning, latestSessionId, replaySession]);
+	}, [replayedRun.active, replayedRun.sessionId, replaySession]);
+}
+
+function useReplayedRunState() {
+	return useAuiState((s) => {
+		const replayRunMessage = s.thread.messages.find((message) =>
+			isReplayRunMessageId(message.id),
+		);
+		const sessionId = normalizeSessionId(
+			readMetadataValue(replayRunMessage?.metadata, "sessionId") ??
+				findLatestMessageMetadataValue(s.thread.messages, "sessionId") ??
+				(typeof window !== "undefined"
+					? window.localStorage.getItem("claw.activeSessionId")
+					: null),
+		);
+		const runId = normalizeSessionId(
+			readMetadataValue(replayRunMessage?.metadata, "runId"),
+		);
+		return {
+			active: Boolean(replayRunMessage),
+			sessionId,
+			runId,
+		};
+	});
+}
+
+async function cancelLatestRun(sessionId: string, runId?: string | null) {
+	await fetch("/api/claw/runs/cancel", {
+		method: "POST",
+		headers: { "content-type": "application/json" },
+		body: JSON.stringify({
+			session_id: sessionId,
+			...(runId ? { run_id: runId } : {}),
+		}),
+	});
 }
 
 const ChatTopActions: FC = () => {
@@ -315,6 +339,7 @@ const ThreadSuggestionItem: FC = () => {
 };
 
 const Composer: FC = () => {
+	const replayedRun = useReplayedRunState();
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
 			<ComposerPrimitive.AttachmentDropzone asChild>
@@ -329,6 +354,7 @@ const Composer: FC = () => {
 						rows={1}
 						autoFocus
 						aria-label="Message input"
+						disabled={replayedRun.active}
 					/>
 					<ComposerAction />
 				</div>
@@ -338,6 +364,8 @@ const Composer: FC = () => {
 };
 
 const ComposerAction: FC = () => {
+	const runtimeRunning = useAuiState((s) => s.thread.isRunning);
+	const replayedRun = useReplayedRunState();
 	return (
 		<div className="aui-composer-action-wrapper relative flex items-center justify-between">
 			<div className="flex items-center gap-1">
@@ -345,7 +373,7 @@ const ComposerAction: FC = () => {
 				<ComposerAssistButtons />
 			</div>
 			<ComposerContextStatus />
-			<AuiIf condition={(s) => !s.thread.isRunning}>
+			{!runtimeRunning && !replayedRun.active ? (
 				<ComposerPrimitive.Send asChild>
 					<TooltipIconButton
 						tooltip="Send message"
@@ -359,8 +387,8 @@ const ComposerAction: FC = () => {
 						<ArrowUpIcon className="aui-composer-send-icon size-4" />
 					</TooltipIconButton>
 				</ComposerPrimitive.Send>
-			</AuiIf>
-			<AuiIf condition={(s) => s.thread.isRunning}>
+			) : null}
+			{runtimeRunning ? (
 				<ComposerPrimitive.Cancel asChild>
 					<Button
 						type="button"
@@ -372,11 +400,58 @@ const ComposerAction: FC = () => {
 						<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
 					</Button>
 				</ComposerPrimitive.Cancel>
-			</AuiIf>
+			) : null}
+			{!runtimeRunning && replayedRun.active ? (
+				<ReplayedRunCancelButton
+					sessionId={replayedRun.sessionId}
+					runId={replayedRun.runId}
+				/>
+			) : null}
 		</div>
 	);
 };
 
+const ReplayedRunCancelButton: FC<{
+	sessionId: string | null;
+	runId: string | null;
+}> = ({ sessionId, runId }) => {
+	const { replaySession } = useClawSessionReplay();
+	const [cancelling, setCancelling] = useState(false);
+	return (
+		<Button
+			type="button"
+			variant="default"
+			size="icon"
+			className="aui-composer-cancel size-8 rounded-full"
+			aria-label="Stop generating"
+			disabled={!sessionId || cancelling}
+			onClick={async () => {
+				if (!sessionId) return;
+				setCancelling(true);
+				try {
+					await cancelLatestRun(sessionId, runId);
+					const runStatus = await fetchLatestRunStatus(sessionId);
+					const response = await fetch(`/api/claw/sessions/${sessionId}`, {
+						cache: "no-store",
+					});
+					if (response.ok) {
+						const payload = await response.json();
+						replaySession(
+							sessionId,
+							toReplayRepository(payload, sessionId, runStatus),
+						);
+					}
+					window.dispatchEvent(new Event("claw-sessions-changed"));
+				} finally {
+					setCancelling(false);
+				}
+			}}
+		>
+			<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
+		</Button>
+	);
+};
+
 const ComposerContextStatus: FC = () => {
 	const status = useComposerContextStatus();
 	const totalTokens = status.contextBudget?.projected_input_tokens ?? 0;
@@ -1130,27 +1205,34 @@ const ActivityChainSummary: FC<{
 	const { openItem } = useActivityPanel();
 	const startedAtRef = useRef(Date.now());
 	const [elapsedMs, setElapsedMs] = useState(0);
+	const runStartedAtMs = useAuiState((s) => {
+		const value = readMetadataValue(s.message.metadata, "runStartedAtMs");
+		return typeof value === "number" && Number.isFinite(value) ? value : null;
+	});
 	const label = useAuiState((s) => {
 		const message = s.message;
-		return summarizeActivityChainLabel(
-			message.content,
-			message.status?.type,
-			elapsedMs,
-		);
+		const status =
+			message.status?.type === "running" || isReplayRunMessageId(messageId)
+				? "running"
+				: message.status?.type;
+		return summarizeActivityChainLabel(message.content, status, elapsedMs);
 	});
 	const active = useAuiState((s) => {
 		const message = s.message;
-		return message.status?.type === "running";
+		return (
+			message.status?.type === "running" || isReplayRunMessageId(messageId)
+		);
 	});
 
 	useEffect(() => {
 		if (!active) return;
-		setElapsedMs(Date.now() - startedAtRef.current);
+		const startedAt = runStartedAtMs ?? startedAtRef.current;
+		setElapsedMs(Math.max(0, Date.now() - startedAt));
 		const timer = window.setInterval(() => {
-			setElapsedMs(Date.now() - startedAtRef.current);
+			setElapsedMs(Math.max(0, Date.now() - startedAt));
 		}, 1000);
 		return () => window.clearInterval(timer);
-	}, [active]);
+	}, [active, runStartedAtMs]);
 
 	return (
 		<button
@@ -1167,6 +1249,10 @@ const ActivityChainSummary: FC<{
 	);
 };
 
+function isReplayRunMessageId(messageId: string) {
+	return messageId.includes("-run-");
+}
+
 type ActivitySummaryPart = {
 	type: string;
 	text?: string;