From 9a42b00bfea4016a6af02dd06e2c21d526e59a05 Mon Sep 17 00:00:00 2001
From: wuyang6 <wuyang6@xiaomi.com>
Date: Wed, 13 May 2026 16:56:14 +0800
Subject: [PATCH] Fix Jupyter workspace session scoping

---
 .../app/components/assistant-ui/thread.tsx    | 57 ++++++++++++++-----
 frontend/app/lib/claw-active-session.ts       | 14 ++++-
 2 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx
index 32ba6db..61332cf 100644
--- a/frontend/app/components/assistant-ui/thread.tsx
+++ b/frontend/app/components/assistant-ui/thread.tsx
@@ -81,7 +81,9 @@ import {
 } from "@/components/ui/tooltip";
 import {
 	clearActiveSessionId,
+	PENDING_WORKSPACE_SESSION_CHANGED_EVENT,
 	readActiveSessionId,
+	readPendingWorkspaceSessionId,
 	writeActiveSessionId,
 	writePendingWorkspaceSessionId,
 } from "@/lib/claw-active-session";
@@ -273,17 +275,46 @@ function useReplayedRunState() {
 			readMetadataValue(replayRunMessage?.metadata, "runId"),
 		);
 	});
-	const localSessionId =
-		typeof window !== "undefined"
-			? normalizeSessionId(readActiveSessionId())
-			: null;
-	const sessionId = replaySessionId ?? latestSessionId ?? localSessionId;
+	const sessionId = replaySessionId ?? latestSessionId;
 	return useMemo(
 		() => ({ active, sessionId, runId }),
 		[active, sessionId, runId],
 	);
 }
 
+function usePendingWorkspaceSessionId() {
+	const [sessionId, setSessionId] = useState<string | null>(() =>
+		typeof window !== "undefined"
+			? normalizeSessionId(readPendingWorkspaceSessionId())
+			: null,
+	);
+
+	useEffect(() => {
+		const refresh = () => {
+			setSessionId(normalizeSessionId(readPendingWorkspaceSessionId()));
+		};
+		window.addEventListener(PENDING_WORKSPACE_SESSION_CHANGED_EVENT, refresh);
+		window.addEventListener("storage", refresh);
+		window.addEventListener("focus", refresh);
+		return () => {
+			window.removeEventListener(
+				PENDING_WORKSPACE_SESSION_CHANGED_EVENT,
+				refresh,
+			);
+			window.removeEventListener("storage", refresh);
+			window.removeEventListener("focus", refresh);
+		};
+	}, []);
+
+	return sessionId;
+}
+
+function useCurrentThreadSessionId({ includePending = false } = {}) {
+	const currentRun = useReplayedRunState();
+	const pendingSessionId = usePendingWorkspaceSessionId();
+	return currentRun.sessionId ?? (includePending ? pendingSessionId : null);
+}
+
 async function cancelLatestRun(sessionId: string, runId?: string | null) {
 	await fetch("/api/claw/runs/cancel", {
 		method: "POST",
@@ -298,8 +329,9 @@ async function cancelLatestRun(sessionId: string, runId?: string | null) {
 const ChatTopActions: FC = () => {
 	const { openFiles } = useActivityPanel();
 	const [workspaceOpen, setWorkspaceOpen] = useState(false);
-	const currentRun = useReplayedRunState();
-	const normalizedSessionId = currentRun.sessionId;
+	const normalizedSessionId = useCurrentThreadSessionId({
+		includePending: true,
+	});
 	const workspaceStatus = useJupyterWorkspaceStatus(normalizedSessionId);
 
 	return (
@@ -444,11 +476,7 @@ function WorkspaceSwitchDialog({
 		null,
 	);
 	const fieldId = useId();
-	const effectiveSessionId =
-		sessionId ??
-		(typeof window !== "undefined"
-			? normalizeSessionId(readActiveSessionId())
-			: null);
+	const effectiveSessionId = sessionId;
 
 	useEffect(() => {
 		if (!open || !effectiveSessionId) return;
@@ -726,6 +754,9 @@ const ThreadSuggestionItem: FC = () => {
 
 const Composer: FC = () => {
 	const replayedRun = useReplayedRunState();
+	const workspaceSessionId = useCurrentThreadSessionId({
+		includePending: true,
+	});
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
 			<ComposerPrimitive.AttachmentDropzone asChild>
@@ -734,7 +765,7 @@ const Composer: FC = () => {
 					className="flex w-full flex-col gap-2 rounded-(--composer-radius) border bg-background p-(--composer-padding) transition-shadow focus-within:border-ring/75 focus-within:ring-2 focus-within:ring-ring/20 data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50"
 				>
 					<ComposerAttachments />
-					<ComposerWorkspaceStatus sessionId={replayedRun.sessionId} />
+					<ComposerWorkspaceStatus sessionId={workspaceSessionId} />
 					<ImeComposerInput
 						placeholder="Send a message..."
 						className="aui-composer-input max-h-32 min-h-10 w-full resize-none bg-transparent px-1.75 py-1 text-sm outline-none placeholder:text-muted-foreground/80"
diff --git a/frontend/app/lib/claw-active-session.ts b/frontend/app/lib/claw-active-session.ts
index 252f9e1..b88cc97 100644
--- a/frontend/app/lib/claw-active-session.ts
+++ b/frontend/app/lib/claw-active-session.ts
@@ -3,6 +3,8 @@
 const ACTIVE_SESSION_ID_KEY = "claw.activeSessionId";
 const PENDING_WORKSPACE_SESSION_ID_KEY = "claw.pendingWorkspaceSessionId";
 export const ACTIVE_SESSION_CHANGED_EVENT = "claw-active-session-changed";
+export const PENDING_WORKSPACE_SESSION_CHANGED_EVENT =
+	"claw-pending-workspace-session-changed";
 
 export function readActiveSessionId() {
 	if (typeof window === "undefined") return null;
@@ -55,6 +57,11 @@ export function writePendingWorkspaceSessionId(
 	}
 	window.sessionStorage.setItem(PENDING_WORKSPACE_SESSION_ID_KEY, normalized);
 	window.localStorage.setItem(PENDING_WORKSPACE_SESSION_ID_KEY, normalized);
+	window.dispatchEvent(
+		new CustomEvent(PENDING_WORKSPACE_SESSION_CHANGED_EVENT, {
+			detail: { sessionId: normalized },
+		}),
+	);
 }
 
 export function consumePendingWorkspaceSessionId() {
@@ -63,7 +70,7 @@ export function consumePendingWorkspaceSessionId() {
 	return sessionId;
 }
 
-function readPendingWorkspaceSessionId() {
+export function readPendingWorkspaceSessionId() {
 	if (typeof window === "undefined") return null;
 	return (
 		normalizeSessionId(
@@ -79,6 +86,11 @@ function clearPendingWorkspaceSessionId() {
 	if (typeof window === "undefined") return;
 	window.sessionStorage.removeItem(PENDING_WORKSPACE_SESSION_ID_KEY);
 	window.localStorage.removeItem(PENDING_WORKSPACE_SESSION_ID_KEY);
+	window.dispatchEvent(
+		new CustomEvent(PENDING_WORKSPACE_SESSION_CHANGED_EVENT, {
+			detail: { sessionId: null },
+		}),
+	);
 }
 
 function normalizeSessionId(value: unknown) {