Preserve running session replay state

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;