fix: guard workspace runtime reattach races

Author: pallyoung

apps/web/src/features/app/WorkbenchRuntimeCoordinator.tsx

@@ -52,6 +52,10 @@ import {
 } from "../workspace/workspace-controller";
 import { attachWorkspaceRuntimeWithRetry } from "../workspace/runtime-attach";
 import { createWorkspaceSessionActions } from "../workspace/session-actions";
+import {
+  advanceWorkspaceSyncVersion,
+  isWorkspaceSyncVersionCurrent,
+} from "../workspace/workspace-sync-version.ts";
 import { useWorkspaceTransportSync } from "../workspace/workspace-sync-hooks";
 
 const withServiceFallback = async <T,>(
@@ -400,17 +404,25 @@ export const WorkbenchRuntimeCoordinator = ({
     }
 
     const task = (async () => {
+      const syncVersion = advanceWorkspaceSyncVersion(workspaceId);
       const runtimeSnapshot = await attachWorkspaceRuntimeWithRetry(
         workspaceId,
         deviceId,
         clientId,
         withServiceFallback,
       );
-      if (!runtimeSnapshot || !hasLiveWorkspaceTab(workspaceId)) {
+      if (
+        !runtimeSnapshot
+        || !hasLiveWorkspaceTab(workspaceId)
+        || !isWorkspaceSyncVersionCurrent(workspaceId, syncVersion)
+      ) {
         return;
       }
       updateState((current) => {
-        if (!current.tabs.some((tab) => tab.id === workspaceId)) {
+        if (
+          !current.tabs.some((tab) => tab.id === workspaceId)
+          || !isWorkspaceSyncVersionCurrent(workspaceId, syncVersion)
+        ) {
           return current;
         }
         return applyWorkspaceRuntimeSnapshot(

apps/web/src/features/workspace/WorkspaceScreen.tsx

@@ -120,6 +120,10 @@ import {
   updateWorkspaceOverlayInput,
   updateWorkspaceOverlayTarget
 } from "./workspace-overlay-actions";
+import {
+  advanceWorkspaceSyncVersion,
+  isWorkspaceSyncVersionCurrent,
+} from "./workspace-sync-version.ts";
 import {
   buildWorkspaceFileSearchResults,
   closeWorkspaceFileSearch,
@@ -461,7 +465,6 @@ export default function WorkspaceScreen({ locale, appSettings, onOpenSettings }:
   const validatedRuntimeTargetsRef = useRef(new Set<string>());
   const runtimeValidationRequestIdRef = useRef(0);
   const overlayBrowseRequestIdRef = useRef(0);
-  const workspaceSyncVersionRef = useRef(new Map<string, number>());
   const persistedLayoutRef = useRef<string>("");
   const agentStartupStateRef = useRef(new Map<string, {
     token: number;
@@ -579,16 +582,6 @@ export default function WorkspaceScreen({ locale, appSettings, onOpenSettings }:
     setState(next);
   };
 
-  const advanceWorkspaceSyncVersion = useCallback((workspaceId: string) => {
-    const nextVersion = (workspaceSyncVersionRef.current.get(workspaceId) ?? 0) + 1;
-    workspaceSyncVersionRef.current.set(workspaceId, nextVersion);
-    return nextVersion;
-  }, []);
-
-  const isWorkspaceSyncVersionCurrent = useCallback((workspaceId: string, version: number) => (
-    (workspaceSyncVersionRef.current.get(workspaceId) ?? 0) === version
-  ), []);
-
   const agentRuntimeRefs = useMemo(() => ({
     draftPromptInputRefs,
     agentTerminalRefs,

apps/web/src/features/workspace/session-actions.ts

@@ -43,6 +43,7 @@ import { createTabFromWorkspaceSnapshot } from "../../shared/utils/workspace.ts"
 import type { AppSettings, BackendArchiveEntry, BackendSession, SessionPatch, Toast, WorkspaceSnapshot } from "../../types/app.ts";
 
 import type { CompletionReminderTarget } from "./completion-reminders.ts";
+import { advanceWorkspaceSyncVersion } from "./workspace-sync-version.ts";
 
 type UpdateTab = (tabId: string, updater: (tab: Tab) => Tab) => void;
 type WithServiceFallback = <T>(operation: () => Promise<T>, fallback: T) => Promise<T>;
@@ -495,6 +496,7 @@ export const createWorkspaceSessionActions = ({
     const target = ensureRestorePane(tabId, preferredPaneId);
     if (!target) return null;
 
+    advanceWorkspaceSyncVersion(tabId);
     const restored = await withServiceFallback(
       () => restoreSessionRequest(tabId, numericSessionId, controllerForTab(tabId)),
       null,

apps/web/src/features/workspace/workspace-sync-version.ts

@@ -0,0 +1,35 @@
+export type WorkspaceSyncVersionTracker = {
+  advance: (workspaceId: string) => number;
+  read: (workspaceId: string) => number;
+  isCurrent: (workspaceId: string, version: number) => boolean;
+};
+
+export const createWorkspaceSyncVersionTracker = (): WorkspaceSyncVersionTracker => {
+  const versions = new Map<string, number>();
+
+  const read = (workspaceId: string) => versions.get(workspaceId) ?? 0;
+
+  return {
+    advance: (workspaceId: string) => {
+      const nextVersion = read(workspaceId) + 1;
+      versions.set(workspaceId, nextVersion);
+      return nextVersion;
+    },
+    read,
+    isCurrent: (workspaceId: string, version: number) => read(workspaceId) === version,
+  };
+};
+
+const workspaceSyncVersionTracker = createWorkspaceSyncVersionTracker();
+
+export const advanceWorkspaceSyncVersion = (workspaceId: string) => (
+  workspaceSyncVersionTracker.advance(workspaceId)
+);
+
+export const readWorkspaceSyncVersion = (workspaceId: string) => (
+  workspaceSyncVersionTracker.read(workspaceId)
+);
+
+export const isWorkspaceSyncVersionCurrent = (workspaceId: string, version: number) => (
+  workspaceSyncVersionTracker.isCurrent(workspaceId, version)
+);

tests/session-actions.test.ts

@@ -2,6 +2,7 @@ import test from 'node:test';
 import assert from 'node:assert/strict';
 import { createTranslator } from '../apps/web/src/i18n.ts';
 import { createWorkspaceSessionActions } from '../apps/web/src/features/workspace/session-actions.ts';
+import { readWorkspaceSyncVersion } from '../apps/web/src/features/workspace/workspace-sync-version.ts';
 import type { AppSettings, Toast } from '../apps/web/src/types/app.ts';
 import type { WorkbenchState } from '../apps/web/src/state/workbench.ts';
 
@@ -200,3 +201,134 @@ test('markSessionIdle does not trigger completion reminder for agent exit notes'
 
   assert.deepEqual(reminders, []);
 });
+
+test('restoreSessionIntoPane bumps the workspace sync version before applying the restored session', async () => {
+  const locale = 'en';
+  const t = createTranslator(locale);
+  const workspaceId = 'ws-restore-sync';
+  const beforeVersion = readWorkspaceSyncVersion(workspaceId);
+  const stateRef = {
+    current: {
+      activeTabId: workspaceId,
+      layout: {
+        leftWidth: 320,
+        rightWidth: 320,
+        rightSplit: 64,
+        showCodePanel: false,
+        showTerminalPanel: false,
+      },
+      overlay: {
+        visible: false,
+        mode: 'local' as const,
+        input: '',
+        target: { type: 'native' as const },
+      },
+      tabs: [
+        {
+          id: workspaceId,
+          title: 'Workspace Restore',
+          status: 'ready' as const,
+          controller: {
+            role: 'controller' as const,
+            deviceId: 'device-a',
+            clientId: 'client-a',
+            fencingToken: 1,
+            takeoverPending: false,
+            takeoverRequestedBySelf: false,
+          },
+          agent: {
+            provider: 'claude' as const,
+            command: 'claude',
+            useWsl: false,
+          },
+          git: { branch: 'main', changes: 0, lastCommit: 'abc123' },
+          gitChanges: [],
+          worktrees: [],
+          sessions: [
+            {
+              id: 'draft-restore',
+              title: 'Session 1',
+              status: 'idle' as const,
+              mode: 'branch' as const,
+              autoFeed: true,
+              queue: [],
+              messages: [],
+              stream: '',
+              unread: 0,
+              lastActiveAt: 1,
+              isDraft: true,
+            },
+          ],
+          activeSessionId: 'draft-restore',
+          archive: [],
+          terminals: [],
+          activeTerminalId: '',
+          fileTree: [],
+          changesTree: [],
+          filePreview: {
+            path: '',
+            content: '',
+            mode: 'preview' as const,
+            originalContent: '',
+            modifiedContent: '',
+            dirty: false,
+          },
+          paneLayout: {
+            type: 'leaf' as const,
+            id: 'pane-draft',
+            sessionId: 'draft-restore',
+          },
+          activePaneId: 'pane-draft',
+          idlePolicy: {
+            enabled: true,
+            idleMinutes: 10,
+            maxActive: 3,
+            pressure: true,
+          },
+        },
+      ],
+    } satisfies WorkbenchState,
+  };
+
+  const actions = createWorkspaceSessionActions({
+    appSettings: defaultAppSettings(),
+    locale,
+    t,
+    stateRef,
+    updateTab: (tabId, updater) => {
+      stateRef.current = {
+        ...stateRef.current,
+        tabs: stateRef.current.tabs.map((tab) => (tab.id === tabId ? updater(tab) : tab)),
+      };
+    },
+    withServiceFallback: async (_operation, fallback) => {
+      if (fallback === null) {
+        return {
+          session: {
+            id: 7,
+            title: 'History Restore Session',
+            status: 'idle' as const,
+            mode: 'branch' as const,
+            auto_feed: true,
+            queue: [],
+            messages: [],
+            stream: '',
+            unread: 0,
+            last_active_at: 10,
+            claude_session_id: 'claude-restore-sync',
+          },
+          alreadyActive: false,
+        };
+      }
+      return fallback;
+    },
+    addToast: () => {},
+  });
+
+  const restored = await actions.restoreSessionIntoPane(workspaceId, '7');
+
+  assert.equal(restored?.id, 7);
+  assert.equal(readWorkspaceSyncVersion(workspaceId), beforeVersion + 1);
+  assert.equal(stateRef.current.tabs[0]?.activeSessionId, '7');
+  assert.equal(stateRef.current.tabs[0]?.sessions[0]?.title, 'History Restore Session');
+});

tests/workspace-sync-version.test.ts

@@ -0,0 +1,38 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+import { createWorkspaceSyncVersionTracker } from "../apps/web/src/features/workspace/workspace-sync-version.ts";
+
+test("later sync versions invalidate earlier in-flight syncs for the same workspace", () => {
+  const tracker = createWorkspaceSyncVersionTracker();
+
+  const firstVersion = tracker.advance("ws-sync-race");
+  const secondVersion = tracker.advance("ws-sync-race");
+
+  assert.equal(firstVersion, 1);
+  assert.equal(secondVersion, 2);
+  assert.equal(tracker.isCurrent("ws-sync-race", firstVersion), false);
+  assert.equal(tracker.isCurrent("ws-sync-race", secondVersion), true);
+});
+
+test("sync versions stay isolated per workspace", () => {
+  const tracker = createWorkspaceSyncVersionTracker();
+
+  const leftVersion = tracker.advance("ws-sync-left");
+  const rightVersion = tracker.advance("ws-sync-right");
+
+  assert.equal(leftVersion, 1);
+  assert.equal(rightVersion, 1);
+  assert.equal(tracker.isCurrent("ws-sync-left", leftVersion), true);
+  assert.equal(tracker.isCurrent("ws-sync-right", rightVersion), true);
+});
+
+test("a local mutation bump invalidates an older attach version", () => {
+  const tracker = createWorkspaceSyncVersionTracker();
+
+  const attachVersion = tracker.advance("ws-sync-restore");
+  const restoreVersion = tracker.advance("ws-sync-restore");
+
+  assert.equal(restoreVersion, attachVersion + 1);
+  assert.equal(tracker.isCurrent("ws-sync-restore", attachVersion), false);
+  assert.equal(tracker.isCurrent("ws-sync-restore", restoreVersion), true);
+});