fix(flow-chat): preserve explicit session selection during restore

limityan · limityan · commit 650d77880100 · 2026-06-10T11:10:29.000+08:00
diff --git a/src/web-ui/src/app/layout/AppLayout.tsx b/src/web-ui/src/app/layout/AppLayout.tsx
@@ -276,6 +276,7 @@ const AppLayout: React.FC<AppLayoutProps> = ({ className = '' }) => {
 
   // Initialize FlowChatManager
   React.useEffect(() => {
+    let cancelled = false;
     const initializeFlowChat = async () => {
       if (!currentWorkspace?.rootPath) return;
 
@@ -305,15 +306,24 @@ const AppLayout: React.FC<AppLayoutProps> = ({ className = '' }) => {
             ? currentWorkspace.sshHost
             : undefined
         );
+        if (cancelled) {
+          return;
+        }
 
         let sessionId: string | undefined;
         const { flowChatStore } = await import('@/flow_chat/store/FlowChatStore');
+        if (cancelled) {
+          return;
+        }
         if (!hasHistoricalSessions) {
           const initialSessionMode =
             currentWorkspace.workspaceKind === WorkspaceKind.Assistant
               ? 'Claw'
               : explicitPreferredMode || 'agentic';
           sessionId = await flowChatManager.createChatSession({}, initialSessionMode);
+          if (cancelled) {
+            return;
+          }
         }
 
         const activeSessionId = sessionId || flowChatStore.getState().activeSessionId;
@@ -326,6 +336,9 @@ const AppLayout: React.FC<AppLayoutProps> = ({ className = '' }) => {
           sessionStorage.removeItem('pendingProjectDescription');
 
           setTimeout(async () => {
+            if (cancelled) {
+              return;
+            }
             try {
               const targetSessionId = sessionId || flowChatStore.getState().activeSessionId;
 
@@ -353,6 +366,9 @@ const AppLayout: React.FC<AppLayoutProps> = ({ className = '' }) => {
         if (pendingSettings) {
           sessionStorage.removeItem('pendingOpenSettings');
           setTimeout(async () => {
+            if (cancelled) {
+              return;
+            }
             try {
               const { quickActions } = await import('@/shared/services/ide-control');
               await quickActions.openSettings(pendingSettings);
@@ -362,6 +378,9 @@ const AppLayout: React.FC<AppLayoutProps> = ({ className = '' }) => {
           }, 500);
         }
       } catch (error) {
+        if (cancelled) {
+          return;
+        }
         log.error('FlowChatManager initialization failed', error);
         import('@/shared/notification-system').then(({ notificationService }) => {
           notificationService.error(t('appLayout.flowChatInitFailed'), { duration: 5000 });
@@ -370,6 +389,9 @@ const AppLayout: React.FC<AppLayoutProps> = ({ className = '' }) => {
     };
 
     initializeFlowChat();
+    return () => {
+      cancelled = true;
+    };
   }, [
     currentWorkspace,
     currentWorkspace?.id,
diff --git a/src/web-ui/src/flow_chat/services/FlowChatManager.test.ts b/src/web-ui/src/flow_chat/services/FlowChatManager.test.ts
@@ -153,4 +153,162 @@ describe('FlowChatManager initialization', () => {
     expect(storeMocks.store.switchSession).toHaveBeenCalledTimes(1);
     expect(storeMocks.store.switchSession).toHaveBeenCalledWith('history-1');
   });
+
+  it('does not overwrite a user-selected workspace session after initial history restore completes', async () => {
+    const historyRestore = createDeferred<void>();
+    const sessions = new Map<string, any>([
+      ['history-1', createHistoricalSession({
+        sessionId: 'history-1',
+        title: 'Newest history',
+        lastActiveAt: 30,
+        lastFinishedAt: 30,
+      })],
+      ['history-2', createHistoricalSession({
+        sessionId: 'history-2',
+        title: 'User selected history',
+        lastActiveAt: 10,
+        lastFinishedAt: 10,
+      })],
+    ]);
+    let activeSessionId: string | null = null;
+
+    storeMocks.store = {
+      registerPersistUnreadCompletionCallback: vi.fn(),
+      loadSessionMetadataPage: vi.fn(async () => ({
+        sessions: [],
+        totalTopLevelCount: 2,
+        hasMore: false,
+      })),
+      getState: vi.fn(() => ({
+        sessions,
+        activeSessionId,
+      })),
+      loadSessionHistory: vi.fn(() => historyRestore.promise),
+      switchSession: vi.fn((sessionId: string) => {
+        activeSessionId = sessionId;
+      }),
+    };
+
+    const manager = FlowChatManager.getInstance();
+    const initialize = manager.initialize('D:/workspace/BitFun');
+
+    await flushAsyncWork();
+    expect(storeMocks.store.loadSessionHistory).toHaveBeenCalledWith(
+      'history-1',
+      'D:/workspace/BitFun',
+      undefined,
+      undefined,
+      undefined,
+      { deferFullHistoryUntilActive: true },
+    );
+
+    activeSessionId = 'history-2';
+    historyRestore.resolve();
+
+    await expect(initialize).resolves.toBe(true);
+
+    expect(storeMocks.store.switchSession).not.toHaveBeenCalled();
+    expect(activeSessionId).toBe('history-2');
+  });
+
+  it('does not let a stale workspace initialization overwrite a newer active workspace', async () => {
+    const historyRestore = createDeferred<void>();
+    const sessions = new Map<string, any>([
+      ['history-1', createHistoricalSession()],
+      ['other-1', createHistoricalSession({
+        sessionId: 'other-1',
+        title: 'Other workspace',
+        workspacePath: 'D:/workspace/Other',
+        lastActiveAt: 40,
+        lastFinishedAt: 40,
+      })],
+    ]);
+    let activeSessionId: string | null = null;
+
+    storeMocks.store = {
+      registerPersistUnreadCompletionCallback: vi.fn(),
+      loadSessionMetadataPage: vi.fn(async () => ({
+        sessions: [],
+        totalTopLevelCount: 1,
+        hasMore: false,
+      })),
+      getState: vi.fn(() => ({
+        sessions,
+        activeSessionId,
+      })),
+      loadSessionHistory: vi.fn(() => historyRestore.promise),
+      switchSession: vi.fn((sessionId: string) => {
+        activeSessionId = sessionId;
+      }),
+    };
+
+    const manager = FlowChatManager.getInstance();
+    const initialize = manager.initialize('D:/workspace/BitFun');
+
+    await flushAsyncWork();
+    activeSessionId = 'other-1';
+    (manager as unknown as { context: { currentWorkspacePath: string | null } })
+      .context.currentWorkspacePath = 'D:/workspace/Other';
+    historyRestore.resolve();
+
+    await expect(initialize).resolves.toBe(true);
+
+    expect(storeMocks.store.switchSession).not.toHaveBeenCalled();
+    expect(activeSessionId).toBe('other-1');
+    expect((manager as unknown as { context: { currentWorkspacePath: string | null } })
+      .context.currentWorkspacePath).toBe('D:/workspace/Other');
+  });
+
+  it('does not let an older workspace initialization switch after a newer workspace initialize starts', async () => {
+    const bitfunHistoryRestore = createDeferred<void>();
+    const sessions = new Map<string, any>([
+      ['history-1', createHistoricalSession()],
+      ['other-1', createHistoricalSession({
+        sessionId: 'other-1',
+        title: 'Other workspace',
+        workspacePath: 'D:/workspace/Other',
+        lastActiveAt: 40,
+        lastFinishedAt: 40,
+      })],
+    ]);
+    let activeSessionId: string | null = null;
+
+    storeMocks.store = {
+      registerPersistUnreadCompletionCallback: vi.fn(),
+      loadSessionMetadataPage: vi.fn(async (
+        workspacePath: string,
+      ) => ({
+        sessions: [],
+        totalTopLevelCount: workspacePath === 'D:/workspace/BitFun' ? 1 : 0,
+        hasMore: false,
+      })),
+      getState: vi.fn(() => ({
+        sessions,
+        activeSessionId,
+      })),
+      loadSessionHistory: vi.fn((sessionId: string) => {
+        if (sessionId === 'history-1') {
+          return bitfunHistoryRestore.promise;
+        }
+        return Promise.resolve();
+      }),
+      switchSession: vi.fn((sessionId: string) => {
+        activeSessionId = sessionId;
+      }),
+    };
+
+    const manager = FlowChatManager.getInstance();
+    const bitfunInitialize = manager.initialize('D:/workspace/BitFun');
+
+    await flushAsyncWork();
+    await expect(manager.initialize('D:/workspace/Other')).resolves.toBe(true);
+
+    bitfunHistoryRestore.resolve();
+    await expect(bitfunInitialize).resolves.toBe(true);
+
+    expect(storeMocks.store.switchSession).toHaveBeenCalledWith('other-1');
+    expect(storeMocks.store.switchSession).not.toHaveBeenCalledWith('history-1');
+    expect((manager as unknown as { context: { currentWorkspacePath: string | null } })
+      .context.currentWorkspacePath).toBe('D:/workspace/Other');
+  });
 });
diff --git a/src/web-ui/src/flow_chat/services/FlowChatManager.ts b/src/web-ui/src/flow_chat/services/FlowChatManager.ts
@@ -54,6 +54,7 @@ export class FlowChatManager {
   private eventListenerInitializationPromise: Promise<void> | null = null;
   private eventListenerCleanup: (() => void) | null = null;
   private initializationRequests = new Map<string, Promise<boolean>>();
+  private latestInitializationRequestKey: string | null = null;
 
   private constructor() {
     this.context = {
@@ -107,12 +108,14 @@ export class FlowChatManager {
       remoteSshHost,
     );
     const existingRequest = this.initializationRequests.get(requestKey);
+    this.latestInitializationRequestKey = requestKey;
     if (existingRequest) {
       return existingRequest;
     }
 
     let request: Promise<boolean>;
     request = this.initializeWorkspace(
+      requestKey,
       workspacePath,
       preferredMode,
       remoteConnectionId,
@@ -141,6 +144,7 @@ export class FlowChatManager {
   }
 
   private async initializeWorkspace(
+    requestKey: string,
     workspacePath: string,
     preferredMode?: string,
     remoteConnectionId?: string,
@@ -214,13 +218,19 @@ export class FlowChatManager {
         workspaceSessions.length > 0 ||
         initialMetadataPage.totalTopLevelCount > 0 ||
         initialMetadataPage.sessions.length > 0;
+      const isCurrentInitializationRequest = () =>
+        this.latestInitializationRequestKey === requestKey;
       const activeSession = state.activeSessionId
         ? state.sessions.get(state.activeSessionId) ?? null
         : null;
       const activeSessionBelongsToWorkspace =
         !!activeSession && sessionMatchesWorkspace(activeSession);
+      const activeSessionIdAtAutoSelectStart = state.activeSessionId;
 
       if (hasHistoricalSessions && !activeSessionBelongsToWorkspace) {
+        if (!isCurrentInitializationRequest()) {
+          return hasHistoricalSessions;
+        }
         const sortedWorkspaceSessions = [...workspaceSessions].sort(compareSessionsForDisplay);
         const latestSession = (preferredMode
           ? sortedWorkspaceSessions.find(session => session.mode === preferredMode)
@@ -242,10 +252,33 @@ export class FlowChatManager {
           );
         }
 
+        if (!isCurrentInitializationRequest()) {
+          return hasHistoricalSessions;
+        }
+
+        const currentState = this.context.flowChatStore.getState();
+        const currentActiveSession = currentState.activeSessionId
+          ? currentState.sessions.get(currentState.activeSessionId) ?? null
+          : null;
+        const currentActiveSessionBelongsToWorkspace =
+          !!currentActiveSession && sessionMatchesWorkspace(currentActiveSession);
+        const activeSessionChangedDuringAutoSelect =
+          currentState.activeSessionId !== activeSessionIdAtAutoSelectStart &&
+          currentState.activeSessionId !== null;
+        if (currentActiveSessionBelongsToWorkspace) {
+          this.context.currentWorkspacePath = workspacePath;
+          return hasHistoricalSessions;
+        }
+        if (activeSessionChangedDuringAutoSelect) {
+          return hasHistoricalSessions;
+        }
+
         this.context.flowChatStore.switchSession(latestSession.sessionId);
       }
 
-      this.context.currentWorkspacePath = workspacePath;
+      if (isCurrentInitializationRequest()) {
+        this.context.currentWorkspacePath = workspacePath;
+      }
 
       return hasHistoricalSessions;
     } catch (error) {
diff --git a/tests/e2e/scripts/generate-long-session-fixture.mjs b/tests/e2e/scripts/generate-long-session-fixture.mjs
@@ -100,8 +100,8 @@ function parseArgs(argv) {
   if (!options.sessionPrefix.trim()) {
     throw new Error('--session-prefix cannot be empty');
   }
-  if (!['explore-only', 'mixed-visible', 'dense-visible'].includes(options.scenario)) {
-    throw new Error('--scenario must be one of: explore-only, mixed-visible, dense-visible');
+  if (!['explore-only', 'mixed-visible', 'dense-visible', 'user-only-latest'].includes(options.scenario)) {
+    throw new Error('--scenario must be one of: explore-only, mixed-visible, dense-visible, user-only-latest');
   }
   return options;
 }
@@ -122,7 +122,7 @@ Options:
   --tool-items <n>          Tool item count per turn. Default: 2
   --dense-groups <n>        Groups in the latest dense-visible turn. Default: 160
   --session-prefix <text>   Session id prefix. Default: perf-long-session
-  --scenario <name>         Fixture shape: mixed-visible, dense-visible, or explore-only. Default: mixed-visible
+  --scenario <name>         Fixture shape: mixed-visible, dense-visible, explore-only, or user-only-latest. Default: mixed-visible
   --bitfun-home <path>      BitFun home root. Default: BITFUN_HOME or ~/.bitfun
   --cleanup                 Remove generated sessions for the prefix.
 `);
@@ -196,7 +196,7 @@ function makeMetadata({ sessionId, sessionName, workspacePath, createdAt, lastAc
     createdAt,
     lastActiveAt,
     turnCount,
-    messageCount: turnCount * 2,
+    messageCount: scenario === 'user-only-latest' ? (turnCount * 2) - 1 : turnCount * 2,
     toolCallCount: turnCount * toolItems,
     status: 'active',
     terminalSessionId: null,
@@ -415,7 +415,9 @@ function makeTurn({
   const turnId = `${sessionId}-turn-${String(turnIndex).padStart(4, '0')}`;
   const toolItemsData = makeToolItems({ turnId, turnIndex, timestamp, toolResultChars, toolItems });
   const isLatestTurn = turnIndex === totalTurns - 1;
-  const modelRounds = scenario === 'dense-visible' && isLatestTurn
+  const modelRounds = scenario === 'user-only-latest' && isLatestTurn
+    ? []
+    : scenario === 'dense-visible' && isLatestTurn
     ? makeDenseLatestModelRounds({
       turnId,
       turnIndex,
diff --git a/tests/e2e/specs/performance/startup-session-perf.spec.ts b/tests/e2e/specs/performance/startup-session-perf.spec.ts
