Fix active session replay state

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

@@ -62,6 +62,11 @@ import {
 	TooltipContent,
 	TooltipTrigger,
 } from "@/components/ui/tooltip";
+import {
+	clearActiveSessionId,
+	readActiveSessionId,
+	writeActiveSessionId,
+} from "@/lib/claw-active-session";
 import { useClawSessionReplay } from "@/lib/claw-session-replay";
 import { cn } from "@/lib/utils";
 
@@ -193,6 +198,15 @@ function useRefreshReplayedRun() {
 			const runStatus = await fetchLatestRunStatus(activeSessionId);
 			if (cancelled) return;
 			if (runStatus?.status === "queued" || runStatus?.status === "running") {
+				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),
+				);
 				return;
 			}
 			const response = await fetch(`/api/claw/sessions/${activeSessionId}`, {
@@ -243,7 +257,7 @@ function useReplayedRunState() {
 	});
 	const localSessionId =
 		typeof window !== "undefined"
-			? normalizeSessionId(window.localStorage.getItem("claw.activeSessionId"))
+			? normalizeSessionId(readActiveSessionId())
 			: null;
 	const sessionId = replaySessionId ?? latestSessionId ?? localSessionId;
 	return useMemo(
@@ -722,7 +736,7 @@ function useComposerContextStatus() {
 	useEffect(() => {
 		const cleanLatestSessionId = normalizeSessionId(latestSessionId);
 		if (cleanLatestSessionId) {
-			window.localStorage.setItem("claw.activeSessionId", cleanLatestSessionId);
+			writeActiveSessionId(cleanLatestSessionId);
 			window.setTimeout(() => {
 				window.dispatchEvent(new Event("claw-sessions-changed"));
 			}, 0);
@@ -752,13 +766,10 @@ function useComposerContextStatus() {
 					: {};
 				let sessionId = normalizeSessionId(
 					normalizeSessionId(latestSessionId) ||
-						normalizeSessionId(
-							window.localStorage.getItem("claw.activeSessionId"),
-						) ||
+						normalizeSessionId(readActiveSessionId()) ||
 						normalizeSessionId(statePayload.active_session_id),
 				);
-				if (sessionId)
-					window.localStorage.setItem("claw.activeSessionId", sessionId);
+				if (sessionId) writeActiveSessionId(sessionId);
 				let sessionPayload: ClawStoredSession | null = null;
 				let contextBudget: ContextBudget | undefined;
 				if (sessionId) {
@@ -772,7 +783,7 @@ function useComposerContextStatus() {
 						sessionPayload =
 							(await sessionResponse.json()) as ClawStoredSession;
 					} else if (sessionResponse.status === 404) {
-						window.localStorage.removeItem("claw.activeSessionId");
+						clearActiveSessionId();
 						sessionId = null;
 					}
 				}