Restore active run state on session replay

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

@@ -42,6 +42,10 @@ import {
 } from "@/components/assistant-ui/attachment";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { Reasoning } from "@/components/assistant-ui/reasoning";
+import {
+	fetchLatestRunStatus,
+	toReplayRepository,
+} from "@/components/assistant-ui/thread-list";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { Button } from "@/components/ui/button";
@@ -58,6 +62,7 @@ import {
 	TooltipContent,
 	TooltipTrigger,
 } from "@/components/ui/tooltip";
+import { useClawSessionReplay } from "@/lib/claw-session-replay";
 import { cn } from "@/lib/utils";
 
 type ComposerInsertEvent = CustomEvent<{ text: string }>;
@@ -110,6 +115,7 @@ type ContextBudget = {
 };
 
 export const Thread: FC = () => {
+	useRefreshReplayedRun();
 	return (
 		<ThreadPrimitive.Root
 			className="aui-root aui-thread-root @container flex h-full flex-col bg-background"
@@ -149,6 +155,52 @@ 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"),
+	);
+
+	useEffect(() => {
+		if (!isRunning || !hasReplayedRunMessage) return;
+		const sessionId = normalizeSessionId(
+			latestSessionId ?? window.localStorage.getItem("claw.activeSessionId"),
+		);
+		if (!sessionId) return;
+		const activeSessionId = sessionId;
+		let cancelled = false;
+
+		async function refreshIfFinished() {
+			const runStatus = await fetchLatestRunStatus(activeSessionId);
+			if (cancelled) return;
+			if (runStatus?.status === "queued" || runStatus?.status === "running") {
+				return;
+			}
+			const response = await fetch(`/api/claw/sessions/${activeSessionId}`, {
+				cache: "no-store",
+			});
+			if (!response.ok || cancelled) return;
+			const payload = await response.json();
+			replaySession(
+				activeSessionId,
+				toReplayRepository(payload, activeSessionId, runStatus),
+			);
+			window.dispatchEvent(new Event("claw-sessions-changed"));
+		}
+
+		refreshIfFinished();
+		const interval = window.setInterval(refreshIfFinished, 3000);
+		return () => {
+			cancelled = true;
+			window.clearInterval(interval);
+		};
+	}, [hasReplayedRunMessage, isRunning, latestSessionId, replaySession]);
+}
+
 const ChatTopActions: FC = () => {
 	const { openFiles } = useActivityPanel();
 	const latestSessionId = useAuiState((s) =>