Merge pull request #74526 from NousResearch/bb/tab-restore-homing

fix(desktop): ⌘R lands on main — cold-start resume clobbers the persisted active tab

apps/desktop/src/app/session/hooks/use-route-resume.test.tsx

@@ -3,9 +3,15 @@ import type { MutableRefObject } from 'react'
 import { afterEach, describe, expect, it, vi } from 'vitest'
 
 import { $resumeExhaustedSessionId, setResumeExhaustedSessionId } from '@/store/session'
+import { markSelectionRestore } from '@/store/session-states'
 
 import { useRouteResume } from './use-route-resume'
 
+// The hook only arms the boot-restore one-shot; the listener consuming it lives
+// in the real store (covered by session-states.test.ts). Mock the module so the
+// store's side effects (persistence listeners) stay out of this harness.
+vi.mock('@/store/session-states', () => ({ markSelectionRestore: vi.fn() }))
+
 interface HarnessProps {
   activeSessionId: null | string
   activeSessionIdRef: MutableRefObject<null | string>
@@ -192,6 +198,45 @@ describe('useRouteResume', () => {
     expect(resumeSession).toHaveBeenCalledWith('session-2', true)
   })
 
+  it('arms the boot-restore one-shot for the FIRST resume only (⌘R tab persistence)', () => {
+    // Factory mocks survive restoreAllMocks — drop calls earlier tests made.
+    vi.mocked(markSelectionRestore).mockClear()
+    const resumeSession = vi.fn(async () => undefined)
+    const startFreshSessionDraft = vi.fn()
+    const activeSessionIdRef: MutableRefObject<null | string> = { current: null }
+    const creatingSessionRef = { current: false }
+    const runtimeIdByStoredSessionIdRef = { current: new Map() }
+    const selectedStoredSessionIdRef: MutableRefObject<null | string> = { current: null }
+
+    const props = {
+      activeSessionId: null,
+      activeSessionIdRef,
+      creatingSessionRef,
+      currentView: 'chat',
+      freshDraftReady: false,
+      gatewayState: 'open',
+      resumeSession,
+      runtimeIdByStoredSessionIdRef,
+      selectedStoredSessionId: null,
+      selectedStoredSessionIdRef,
+      startFreshSessionDraft
+    }
+
+    // Cold start: the window mounts already routed at /session-1 (the reload).
+    const { rerender } = render(
+      <RouteResumeHarness {...props} locationPathname="/session-1" routedSessionId="session-1" />
+    )
+
+    expect(resumeSession).toHaveBeenCalledWith('session-1', true)
+    expect(markSelectionRestore).toHaveBeenCalledTimes(1)
+
+    // A later route change is a real navigation: resume fires, one-shot doesn't.
+    rerender(<RouteResumeHarness {...props} locationPathname="/session-2" routedSessionId="session-2" />)
+
+    expect(resumeSession).toHaveBeenCalledWith('session-2', true)
+    expect(markSelectionRestore).toHaveBeenCalledTimes(1)
+  })
+
   it('resumes the selected route again when the gateway reconnects', () => {
     const resumeSession = vi.fn(async () => undefined)
     const startFreshSessionDraft = vi.fn()

apps/desktop/src/app/session/hooks/use-route-resume.ts

@@ -2,6 +2,7 @@ import { type MutableRefObject, useEffect, useRef } from 'react'
 
 import { isNewChatRoute } from '@/app/routes'
 import { setResumeExhaustedSessionId } from '@/store/session'
+import { markSelectionRestore } from '@/store/session-states'
 
 interface RouteResumeOptions {
   activeSessionId: string | null
@@ -85,6 +86,12 @@ export function useRouteResume({
   const lastPathnameRef = useRef<string | null>(null)
   const seenGatewayStateRef = useRef(false)
   const wasGatewayOpenRef = useRef(false)
+  // True until the FIRST resume this window dispatches. That resume is the
+  // cold-start restore of a route that was already loaded before the reload —
+  // a re-attachment, not a navigation — so it must not home focus/tabs to the
+  // workspace (which would clobber the persisted active tab; ⌘R always landed
+  // on main). Every later resume is a real navigation and homes as usual.
+  const bootResumeRef = useRef(true)
   // Per-session retry bookkeeping for the bounded auto-retry effect below. Keyed
   // by the session id we're retrying so switching chats resets the counter.
   const retrySessionIdRef = useRef<string | null>(null)
@@ -148,6 +155,16 @@ export function useRouteResume({
       // rebinds/reaps the session on its side, and trusting it strands Desktop on
       // a dead id ("session not found"). Otherwise keep skipping when already active.
       if ((gatewayBecameOpen || !alreadyActive) && shouldResume && !creatingSessionRef.current) {
+        // The window's FIRST resume re-attaches the pre-reload route rather
+        // than navigating anywhere, so the selection listener must not home
+        // focus/tabs to the workspace over the persisted layout (see
+        // markSelectionRestore). One-shot: consumed by the selection change
+        // resumeSession makes synchronously at entry.
+        if (bootResumeRef.current) {
+          markSelectionRestore()
+        }
+
+        bootResumeRef.current = false
         void resumeSession(routedSessionId, true)
       }
 
@@ -160,6 +177,8 @@ export function useRouteResume({
       (selectedStoredSessionId || activeSessionId || !freshDraftReady) &&
       !rawHashLooksLikeSession()
     ) {
+      // A fresh draft is a real navigation — any later resume homes normally.
+      bootResumeRef.current = false
       startFreshSessionDraft(true)
     }
   }, [

apps/desktop/src/store/session-states.test.ts

@@ -2,10 +2,13 @@ import { describe, expect, it } from 'vitest'
 
 import type { ClientSessionState } from '@/app/types'
 import { group, split } from '@/components/pane-shell/tree/model'
+import { $layoutTree } from '@/components/pane-shell/tree/store'
+import { $selectedStoredSessionId } from '@/store/session'
 import type { SessionTile } from '@/store/session-states'
 import {
   blankDraftTile,
   focusedSessionNeedsRoute,
+  markSelectionRestore,
   orderTilesByTree,
   selectionHomesToWorkspace
 } from '@/store/session-states'
@@ -51,6 +54,36 @@ describe('selectionHomesToWorkspace', () => {
   })
 })
 
+describe('boot-restore selection homing (⌘R tab persistence)', () => {
+  const mainGroup = () => group(['workspace', tilePane('t')], { active: tilePane('t'), id: 'main' })
+
+  const activePane = () => {
+    const tree = $layoutTree.get()
+
+    return tree?.type === 'group' ? tree.active : null
+  }
+
+  it('a normal selection change fronts the workspace tab over an active tile', () => {
+    $layoutTree.set(mainGroup())
+    $selectedStoredSessionId.set('nav-1')
+
+    expect(activePane()).toBe('workspace')
+  })
+
+  it('markSelectionRestore skips homing exactly once, so the persisted active tab survives a reload', () => {
+    $layoutTree.set(mainGroup())
+    markSelectionRestore()
+    $selectedStoredSessionId.set('boot-1')
+
+    // Boot restore: the tile tab the user reloaded on stays fronted.
+    expect(activePane()).toBe(tilePane('t'))
+
+    // One-shot consumed: the next selection change is a real navigation.
+    $selectedStoredSessionId.set('nav-2')
+    expect(activePane()).toBe('workspace')
+  })
+})
+
 describe('focusedSessionNeedsRoute', () => {
   it('routes when the session is not on screen', () => {
     expect(focusedSessionNeedsRoute(null, false)).toBe(true)

apps/desktop/src/store/session-states.ts

@@ -768,10 +768,26 @@ export const $focusedSessionState = computed([$focusedRuntimeId, $sessionStates]
 export const selectionHomesToWorkspace = (selected: null | string, tiles: readonly SessionTile[]): boolean =>
   !(selected && tiles.some(t => t.storedSessionId === selected))
 
+// Cold-start restore is the one selection change that is NOT a navigation: the
+// route already pointed at the primary session before the window loaded, and
+// homing on it would front the workspace tab over the PERSISTED active tab —
+// then persist that clobber, so the tab you reloaded on never comes back
+// (⌘R always landing on main). use-route-resume arms this one-shot right
+// before dispatching the boot resume; the very next selection change skips
+// homing and the restored layout tree keeps its say.
+let selectionRestoreInFlight = false
+
+export function markSelectionRestore() {
+  selectionRestoreInFlight = true
+}
+
 // Homing also FRONTS the workspace tab: the resumed chat loads in the workspace
 // pane, so a zone parked on a tile tab must switch back or the click looks dead.
 $selectedStoredSessionId.listen(selected => {
-  if (!selectionHomesToWorkspace(selected, $sessionTiles.get())) {
+  const restoring = selectionRestoreInFlight
+  selectionRestoreInFlight = false
+
+  if (restoring || !selectionHomesToWorkspace(selected, $sessionTiles.get())) {
     return
   }