From cdbc7d9fb7d1e05ea24880a2705ce74dc488b70d Mon Sep 17 00:00:00 2001
From: wuyang6 <wuyang6@xiaomi.com>
Date: Wed, 13 May 2026 17:22:58 +0800
Subject: [PATCH] Fix pending Jupyter workspace handoff

---
 frontend/app/app/assistant.tsx                | 12 +++------
 .../app/components/assistant-ui/thread.tsx    |  6 +++++
 frontend/app/lib/claw-active-session.ts       | 26 +++++--------------
 3 files changed, 15 insertions(+), 29 deletions(-)

diff --git a/frontend/app/app/assistant.tsx b/frontend/app/app/assistant.tsx
index acf177d..b9d0a23 100644
--- a/frontend/app/app/assistant.tsx
+++ b/frontend/app/app/assistant.tsx
@@ -20,8 +20,8 @@ import {
 	SidebarTrigger,
 } from "@/components/ui/sidebar";
 import {
-	consumePendingWorkspaceSessionId,
 	readActiveSessionId,
+	readPendingWorkspaceSessionId,
 	writeActiveSessionId,
 } from "@/lib/claw-active-session";
 import { ClawSessionReplayProvider } from "@/lib/claw-session-replay";
@@ -39,19 +39,13 @@ export const Assistant = () => {
 					const selectedSessionId = readActiveSessionId();
 					const lastSessionId = getLastSessionId(messages);
 					const pendingWorkspaceSessionId =
-						messages.length <= 1 ? consumePendingWorkspaceSessionId() : null;
-					const usablePendingWorkspaceSessionId =
-						pendingWorkspaceSessionId &&
-						(!selectedSessionId ||
-							selectedSessionId === pendingWorkspaceSessionId)
-							? pendingWorkspaceSessionId
-							: null;
+						messages.length <= 1 ? readPendingWorkspaceSessionId() : null;
 					const selectedResumeSessionId =
 						messages.length > 1 ? selectedSessionId : null;
 					const outgoingSessionId =
 						lastSessionId ??
 						selectedResumeSessionId ??
-						usablePendingWorkspaceSessionId ??
+						pendingWorkspaceSessionId ??
 						options.id;
 					const lastUserMessage = [...messages]
 						.reverse()
diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx
index 61332cf..7c3b8dc 100644
--- a/frontend/app/components/assistant-ui/thread.tsx
+++ b/frontend/app/components/assistant-ui/thread.tsx
@@ -81,6 +81,7 @@ import {
 } from "@/components/ui/tooltip";
 import {
 	clearActiveSessionId,
+	clearPendingWorkspaceSessionId,
 	PENDING_WORKSPACE_SESSION_CHANGED_EVENT,
 	readActiveSessionId,
 	readPendingWorkspaceSessionId,
@@ -312,6 +313,11 @@ function usePendingWorkspaceSessionId() {
 function useCurrentThreadSessionId({ includePending = false } = {}) {
 	const currentRun = useReplayedRunState();
 	const pendingSessionId = usePendingWorkspaceSessionId();
+	useEffect(() => {
+		if (currentRun.sessionId && currentRun.sessionId === pendingSessionId) {
+			clearPendingWorkspaceSessionId();
+		}
+	}, [currentRun.sessionId, pendingSessionId]);
 	return currentRun.sessionId ?? (includePending ? pendingSessionId : null);
 }
 
diff --git a/frontend/app/lib/claw-active-session.ts b/frontend/app/lib/claw-active-session.ts
index b88cc97..9b8ebe6 100644
--- a/frontend/app/lib/claw-active-session.ts
+++ b/frontend/app/lib/claw-active-session.ts
@@ -1,11 +1,12 @@
 "use client";
 
 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";
 
+let pendingWorkspaceSessionId: string | null = null;
+
 export function readActiveSessionId() {
 	if (typeof window === "undefined") return null;
 	return (
@@ -55,8 +56,7 @@ export function writePendingWorkspaceSessionId(
 		clearPendingWorkspaceSessionId();
 		return;
 	}
-	window.sessionStorage.setItem(PENDING_WORKSPACE_SESSION_ID_KEY, normalized);
-	window.localStorage.setItem(PENDING_WORKSPACE_SESSION_ID_KEY, normalized);
+	pendingWorkspaceSessionId = normalized;
 	window.dispatchEvent(
 		new CustomEvent(PENDING_WORKSPACE_SESSION_CHANGED_EVENT, {
 			detail: { sessionId: normalized },
@@ -64,28 +64,14 @@ export function writePendingWorkspaceSessionId(
 	);
 }
 
-export function consumePendingWorkspaceSessionId() {
-	const sessionId = readPendingWorkspaceSessionId();
-	clearPendingWorkspaceSessionId();
-	return sessionId;
-}
-
 export function readPendingWorkspaceSessionId() {
 	if (typeof window === "undefined") return null;
-	return (
-		normalizeSessionId(
-			window.sessionStorage.getItem(PENDING_WORKSPACE_SESSION_ID_KEY),
-		) ??
-		normalizeSessionId(
-			window.localStorage.getItem(PENDING_WORKSPACE_SESSION_ID_KEY),
-		)
-	);
+	return pendingWorkspaceSessionId;
 }
 
-function clearPendingWorkspaceSessionId() {
+export function clearPendingWorkspaceSessionId() {
 	if (typeof window === "undefined") return;
-	window.sessionStorage.removeItem(PENDING_WORKSPACE_SESSION_ID_KEY);
-	window.localStorage.removeItem(PENDING_WORKSPACE_SESSION_ID_KEY);
+	pendingWorkspaceSessionId = null;
 	window.dispatchEvent(
 		new CustomEvent(PENDING_WORKSPACE_SESSION_CHANGED_EVENT, {
 			detail: { sessionId: null },