Fix queued workspace wake after sleep (#4406)

Co-authored-by: Orca <help@stably.ai>

src/renderer/src/components/sidebar/sleep-worktree-activation-race.test.ts

@@ -0,0 +1,118 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => {
+  const state = {
+    activeWorktreeId: null as string | null,
+    setActiveWorktree: vi.fn((worktreeId: string | null) => {
+      state.activeWorktreeId = worktreeId
+    }),
+    shutdownWorktreeBrowsers: vi.fn().mockResolvedValue(undefined),
+    shutdownWorktreeTerminals: vi.fn(async (worktreeId: string) => {
+      for (const tab of state.tabsByWorktree[worktreeId] ?? []) {
+        state.ptyIdsByTabId[tab.id] = []
+      }
+    }),
+    tabsByWorktree: {} as Record<string, { id: string }[]>,
+    ptyIdsByTabId: {} as Record<string, string[]>,
+    browserTabsByWorktree: {} as Record<string, { id: string }[]>,
+    openFiles: [] as { worktreeId: string }[]
+  }
+  const activateAndRevealWorktree = vi.fn()
+  const markInputQuietSchedulerInput = vi.fn()
+  const pendingCallbacks: (() => void)[] = []
+  const pendingCancels: ReturnType<typeof vi.fn>[] = []
+  const scheduleAfterInputQuiet = vi.fn((callback: () => void) => {
+    let cancelled = false
+    const cancel = vi.fn(() => {
+      cancelled = true
+    })
+    pendingCallbacks.push(() => {
+      if (!cancelled) {
+        callback()
+      }
+    })
+    pendingCancels.push(cancel)
+    return cancel
+  })
+  return {
+    activateAndRevealWorktree,
+    markInputQuietSchedulerInput,
+    pendingCallbacks,
+    pendingCancels,
+    scheduleAfterInputQuiet,
+    state,
+    toastError: vi.fn()
+  }
+})
+
+vi.mock('@/store', () => ({
+  useAppStore: {
+    getState: () => mocks.state
+  }
+}))
+
+vi.mock('@/lib/worktree-activation', () => ({
+  activateAndRevealWorktree: mocks.activateAndRevealWorktree
+}))
+
+vi.mock('@/lib/input-quiet-scheduler', () => ({
+  markInputQuietSchedulerInput: mocks.markInputQuietSchedulerInput,
+  scheduleAfterInputQuiet: mocks.scheduleAfterInputQuiet
+}))
+
+vi.mock('sonner', () => ({ toast: { error: mocks.toastError } }))
+
+import { activateWorktreeFromSidebar } from '@/lib/sidebar-worktree-activation'
+import { runSleepWorktrees } from './sleep-worktree-flow'
+
+describe('sleep flow vs queued slept-workspace activation', () => {
+  beforeEach(() => {
+    mocks.activateAndRevealWorktree.mockClear()
+    mocks.markInputQuietSchedulerInput.mockClear()
+    mocks.scheduleAfterInputQuiet.mockClear()
+    mocks.pendingCallbacks.length = 0
+    mocks.pendingCancels.length = 0
+    mocks.toastError.mockClear()
+    mocks.state.activeWorktreeId = 'wt-parent'
+    mocks.state.setActiveWorktree.mockClear()
+    mocks.state.shutdownWorktreeBrowsers.mockClear().mockResolvedValue(undefined)
+    mocks.state.shutdownWorktreeTerminals.mockClear().mockImplementation(async (worktreeId) => {
+      for (const tab of mocks.state.tabsByWorktree[worktreeId] ?? []) {
+        mocks.state.ptyIdsByTabId[tab.id] = []
+      }
+    })
+    mocks.state.tabsByWorktree = {
+      'wt-parent': [{ id: 'tab-parent' }],
+      'wt-child-1': [{ id: 'tab-child-1' }],
+      'wt-child-2': [{ id: 'tab-child-2' }],
+      'wt-child-3': [{ id: 'tab-child-3' }]
+    }
+    mocks.state.ptyIdsByTabId = {
+      'tab-parent': ['pty-parent'],
+      'tab-child-1': ['pty-child-1'],
+      'tab-child-2': ['pty-child-2'],
+      'tab-child-3': ['pty-child-3']
+    }
+    mocks.state.browserTabsByWorktree = {}
+    mocks.state.openFiles = []
+  })
+
+  it('does not let an old slept-parent activation fire after sleeping children', async () => {
+    await runSleepWorktrees(['wt-parent'])
+
+    expect(mocks.state.activeWorktreeId).toBeNull()
+    expect(mocks.state.ptyIdsByTabId['tab-parent']).toEqual([])
+
+    // A normal click on the slept parent row during selection setup queues a
+    // wake internally, even though the user is just trying to select children.
+    activateWorktreeFromSidebar('wt-parent')
+    expect(mocks.activateAndRevealWorktree).not.toHaveBeenCalled()
+    expect(mocks.pendingCallbacks).toHaveLength(1)
+
+    await runSleepWorktrees(['wt-child-1', 'wt-child-2', 'wt-child-3'])
+    mocks.pendingCallbacks[0]?.()
+
+    expect(mocks.activateAndRevealWorktree).not.toHaveBeenCalled()
+    expect(mocks.pendingCancels[0]).toHaveBeenCalledTimes(1)
+  })
+})

src/renderer/src/components/sidebar/sleep-worktree-flow.test.ts

@@ -14,7 +14,14 @@ const mocks = vi.hoisted(() => {
   const toastError = vi.fn()
   const markWorktreeSleepIntent = vi.fn()
   const clearWorktreeSleepIntent = vi.fn()
-  return { clearWorktreeSleepIntent, markWorktreeSleepIntent, state, toastError }
+  const cancelPendingSidebarWorktreeActivation = vi.fn()
+  return {
+    cancelPendingSidebarWorktreeActivation,
+    clearWorktreeSleepIntent,
+    markWorktreeSleepIntent,
+    state,
+    toastError
+  }
 })
 
 vi.mock('@/store', () => ({
@@ -28,6 +35,9 @@ vi.mock('@/lib/worktree-sleep-intent', () => ({
   clearWorktreeSleepIntent: mocks.clearWorktreeSleepIntent,
   markWorktreeSleepIntent: mocks.markWorktreeSleepIntent
 }))
+vi.mock('@/lib/sidebar-worktree-activation', () => ({
+  cancelPendingSidebarWorktreeActivation: mocks.cancelPendingSidebarWorktreeActivation
+}))
 
 import { runSleepWorktree, runSleepWorktrees } from './sleep-worktree-flow'
 
@@ -40,6 +50,7 @@ describe('runSleepWorktree', () => {
     mocks.state.consumeSuppressedPtyExit.mockClear()
     mocks.markWorktreeSleepIntent.mockClear()
     mocks.clearWorktreeSleepIntent.mockClear()
+    mocks.cancelPendingSidebarWorktreeActivation.mockClear()
     mocks.toastError.mockClear()
     mocks.state.activeWorktreeId = null
     mocks.state.tabsByWorktree = {}
@@ -64,6 +75,15 @@ describe('runSleepWorktree', () => {
     expect(browsersCallOrder).toBeLessThan(terminalsCallOrder)
   })
 
+  it('cancels a queued sidebar wake before sleeping worktrees', async () => {
+    await runSleepWorktrees(['wt-1', 'wt-2'])
+
+    expect(mocks.cancelPendingSidebarWorktreeActivation).toHaveBeenCalledTimes(1)
+    const cancelCall = mocks.cancelPendingSidebarWorktreeActivation.mock.invocationCallOrder[0]
+    const browserCall = mocks.state.shutdownWorktreeBrowsers.mock.invocationCallOrder[0]
+    expect(cancelCall).toBeLessThan(browserCall)
+  })
+
   it('clears activeWorktreeId before teardown when the slept worktree is active', async () => {
     mocks.state.activeWorktreeId = 'wt-1'
 

src/renderer/src/components/sidebar/sleep-worktree-flow.ts

@@ -1,6 +1,7 @@
 import { toast } from 'sonner'
 import { useAppStore } from '@/store'
 import { clearWorktreeSleepIntent, markWorktreeSleepIntent } from '@/lib/worktree-sleep-intent'
+import { cancelPendingSidebarWorktreeActivation } from '@/lib/sidebar-worktree-activation'
 import { VIRTUALIZED_SCROLL_ANCHOR_RECORD_EVENT } from '@/hooks/useVirtualizedScrollAnchor'
 
 /**
@@ -84,6 +85,9 @@ export async function runSleepWorktrees(worktreeIds: readonly string[]): Promise
   if (worktreeIds.length === 0) {
     return
   }
+  // Why: clicking a slept sidebar row queues a delayed wake. A later Sleep
+  // command is the newer intent, so do not let that stale wake respawn PTYs.
+  cancelPendingSidebarWorktreeActivation()
   const {
     activeWorktreeId,
     setActiveWorktree,

src/renderer/src/lib/sidebar-worktree-activation.test.ts

@@ -0,0 +1,92 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => {
+  const state = {
+    tabsByWorktree: {} as Record<string, { id: string }[]>,
+    ptyIdsByTabId: {} as Record<string, string[]>,
+    browserTabsByWorktree: {} as Record<string, { id: string }[]>,
+    openFiles: [] as { worktreeId: string }[]
+  }
+  const activateAndRevealWorktree = vi.fn()
+  const markInputQuietSchedulerInput = vi.fn()
+  const pendingCallbacks: (() => void)[] = []
+  const pendingCancels: ReturnType<typeof vi.fn>[] = []
+  const scheduleAfterInputQuiet = vi.fn((callback: () => void) => {
+    let cancelled = false
+    const cancel = vi.fn(() => {
+      cancelled = true
+    })
+    pendingCallbacks.push(() => {
+      if (!cancelled) {
+        callback()
+      }
+    })
+    pendingCancels.push(cancel)
+    return cancel
+  })
+  return {
+    activateAndRevealWorktree,
+    markInputQuietSchedulerInput,
+    pendingCallbacks,
+    pendingCancels,
+    scheduleAfterInputQuiet,
+    state
+  }
+})
+
+vi.mock('@/store', () => ({
+  useAppStore: {
+    getState: () => mocks.state
+  }
+}))
+
+vi.mock('@/lib/worktree-activation', () => ({
+  activateAndRevealWorktree: mocks.activateAndRevealWorktree
+}))
+
+vi.mock('@/lib/input-quiet-scheduler', () => ({
+  markInputQuietSchedulerInput: mocks.markInputQuietSchedulerInput,
+  scheduleAfterInputQuiet: mocks.scheduleAfterInputQuiet
+}))
+
+import {
+  activateWorktreeFromSidebar,
+  cancelPendingSidebarWorktreeActivation
+} from './sidebar-worktree-activation'
+
+describe('sidebar worktree activation', () => {
+  beforeEach(() => {
+    cancelPendingSidebarWorktreeActivation()
+    mocks.activateAndRevealWorktree.mockClear()
+    mocks.markInputQuietSchedulerInput.mockClear()
+    mocks.scheduleAfterInputQuiet.mockClear()
+    mocks.pendingCallbacks.length = 0
+    mocks.pendingCancels.length = 0
+    mocks.state.tabsByWorktree = {}
+    mocks.state.ptyIdsByTabId = {}
+    mocks.state.browserTabsByWorktree = {}
+    mocks.state.openFiles = []
+  })
+
+  it('cancels a queued slept-workspace activation', () => {
+    mocks.state.tabsByWorktree = { 'wt-parent': [{ id: 'tab-1' }] }
+    mocks.state.ptyIdsByTabId = { 'tab-1': [] }
+
+    activateWorktreeFromSidebar('wt-parent')
+    cancelPendingSidebarWorktreeActivation()
+
+    expect(mocks.pendingCancels[0]).toHaveBeenCalledTimes(1)
+    mocks.pendingCallbacks[0]?.()
+    expect(mocks.activateAndRevealWorktree).not.toHaveBeenCalled()
+  })
+
+  it('does not defer a workspace with a live PTY', () => {
+    mocks.state.tabsByWorktree = { 'wt-live': [{ id: 'tab-1' }] }
+    mocks.state.ptyIdsByTabId = { 'tab-1': ['pty-1'] }
+
+    activateWorktreeFromSidebar('wt-live')
+
+    expect(mocks.scheduleAfterInputQuiet).not.toHaveBeenCalled()
+    expect(mocks.activateAndRevealWorktree).toHaveBeenCalledWith('wt-live')
+  })
+})

src/renderer/src/lib/sidebar-worktree-activation.ts

@@ -11,6 +11,11 @@ let pendingSidebarWorktreeActivation: {
   cancel: () => void
 } | null = null
 
+export function cancelPendingSidebarWorktreeActivation(): void {
+  pendingSidebarWorktreeActivation?.cancel()
+  pendingSidebarWorktreeActivation = null
+}
+
 function shouldDeferSidebarWorktreeActivation(worktreeId: string): boolean {
   const state = useAppStore.getState()
   const tabs = state.tabsByWorktree[worktreeId] ?? []
@@ -27,8 +32,7 @@ function shouldDeferSidebarWorktreeActivation(worktreeId: string): boolean {
 }
 
 export function activateWorktreeFromSidebar(worktreeId: string): void {
-  pendingSidebarWorktreeActivation?.cancel()
-  pendingSidebarWorktreeActivation = null
+  cancelPendingSidebarWorktreeActivation()
 
   const activate = (): void => {
     if (pendingSidebarWorktreeActivation?.worktreeId === worktreeId) {