Merge pull request #397 from nwparker/lc2/p1-wake-during-suspend

fix(sleep): tabs survive a wake pressed while the suspend is still committing (LC2-P1)

src/renderer/src/components/browser-pane/stream-remote/use-remote-browser-page-lifecycle.ts

@@ -1,4 +1,8 @@
-import { consumeReopenAfterSleep } from '@/components/sidebar/workspace-actions/sleep/browser-tab-sleep'
+import {
+  noteSleptRemotePageAnswered,
+  sleptRemotePageMissingAction,
+  useReopenSleptRemotePage
+} from './use-slept-remote-page'
 import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react'
 import { useAppStore } from '@/store'
 import { runtimeEnvironmentIdOrLocalUnlessTwin } from '@/lib/workspace/ownership/worktree-runtime-owner'
@@ -60,6 +64,8 @@ export function useRemoteBrowserPageLifecycle({
   const mountedRef = useRef(true)
   const isActiveRef = useRef(isActive)
   const currentBrowserTabIdRef = useRef(browserTab.id)
+  const currentWorktreeIdRef = useRef(browserTab.worktreeId)
+  const currentWorkspaceIdRef = useRef(browserTab.workspaceId)
   const currentBrowserTabUrlRef = useRef(browserTab.url)
   const runtimeWorktree = useMemo(() => toRuntimeWorktreeSelector(worktreeId), [worktreeId])
   const runtimeWorktreeRef = useRef(runtimeWorktree)
@@ -72,11 +78,21 @@ export function useRemoteBrowserPageLifecycle({
 
   useLayoutEffect(() => {
     currentBrowserTabIdRef.current = browserTab.id
+    currentWorktreeIdRef.current = browserTab.worktreeId
+    currentWorkspaceIdRef.current = browserTab.workspaceId
     currentBrowserTabUrlRef.current = browserTab.url
     activeRuntimeEnvironmentIdRef.current = activeRuntimeEnvironmentId
     isActiveRef.current = isActive
     runtimeWorktreeRef.current = runtimeWorktree
-  }, [activeRuntimeEnvironmentId, browserTab.id, browserTab.url, isActive, runtimeWorktree])
+  }, [
+    activeRuntimeEnvironmentId,
+    browserTab.id,
+    browserTab.url,
+    browserTab.workspaceId,
+    browserTab.worktreeId,
+    isActive,
+    runtimeWorktree
+  ])
 
   if (!lifecycleRef.current) {
     lifecycleRef.current = new RemoteBrowserStreamLifecycle({
@@ -107,7 +123,14 @@ export function useRemoteBrowserPageLifecycle({
         // Every transition invalidates the pane notice; stale direct feedback would outrank the stream’s failure beside Reconnect.
         setPaneNotice(null)
       },
-      applyTabInfo: (tab) => streamBridgeRef.current.applyTabInfo(tab),
+      applyTabInfo: (tab) => {
+        noteSleptRemotePageAnswered({
+          id: currentBrowserTabIdRef.current,
+          workspaceId: currentWorkspaceIdRef.current,
+          worktreeId: currentWorktreeIdRef.current
+        })
+        streamBridgeRef.current.applyTabInfo(tab)
+      },
       clearFrame: () => streamBridgeRef.current.clearFrame(),
       handleFrameBytes: (token, bytes) => streamBridgeRef.current.handleFrameBytes(token, bytes),
       closeMissingRemotePage: (remotePageId) =>
@@ -119,6 +142,13 @@ export function useRemoteBrowserPageLifecycle({
   }
   const lifecycle = lifecycleRef.current
 
+  const isMounted = useCallback(() => mountedRef.current, [])
+  useReopenSleptRemotePage(
+    { id: browserTab.id, worktreeId: browserTab.worktreeId },
+    lifecycle,
+    isMounted
+  )
+
   const runtimeTarget = useCallback(() => {
     return activeRuntimeEnvironmentId
       ? ({
@@ -148,15 +178,16 @@ export function useRemoteBrowserPageLifecycle({
         state.removeRemoteBrowserPageHandle(browserTab.id, remotePageId)
       }
       lifecycle.abandonRemotePage()
-      // Slept with its VM (NEW-01): the tab stays. While asleep it waits; once the workspace is
-      // open again it opens a new host page at its URL.
-      if (state.asleepBrowserWorktreeIds[browserTab.worktreeId]) {
+      // Slept with its VM (NEW-01, LC2-P1): wait, or reopen at its URL, instead of closing.
+      const action = sleptRemotePageMissingAction({
+        id: browserTab.id,
+        worktreeId: browserTab.worktreeId
+      })
+      if (action !== 'close') {
         clearStreamFrame()
-        return
-      }
-      if (consumeReopenAfterSleep(browserTab.id)) {
-        clearStreamFrame()
-        lifecycle.open()
+        if (action === 'reopen') {
+          lifecycle.open()
+        }
         return
       }
       if (remoteViewportTimerRef.current !== null) {

src/renderer/src/components/browser-pane/stream-remote/use-slept-remote-page.ts

@@ -0,0 +1,62 @@
+// A host-run browser page and its workspace's sleep (NEW-01, LC2-P1). Pages die with the VM; the
+// tab stays and gets a new host page at its URL once the sandbox is up for good, including after
+// a wake pressed while the suspend was still being committed.
+import { useEffect, useRef } from 'react'
+import { useAppStore } from '@/store'
+import { useIsWorktreeSuspendCommitted } from '@/lib/worktree-sleep-view'
+import {
+  isAwaitingReopenAfterSleep,
+  isBrowserSandboxUp,
+  missingRemotePageAction,
+  readSleepInFlight,
+  settleReopenAfterSleep
+} from '@/components/sidebar/workspace-actions/sleep/browser-tab-sleep'
+
+type Page = { id: string; worktreeId: string }
+
+/** What to do with a page the host reports missing: wait, reopen it, or close the tab. */
+export function sleptRemotePageMissingAction(page: Page): 'wait' | 'reopen' | 'close' {
+  const state = useAppStore.getState()
+  return missingRemotePageAction({
+    awaitingReopen: isAwaitingReopenAfterSleep(page.id),
+    asleep: Boolean(state.asleepBrowserWorktreeIds[page.worktreeId]),
+    sandboxUp: isBrowserSandboxUp(state, page.worktreeId, readSleepInFlight(page.worktreeId))
+  })
+}
+
+/** The host answered for this page with the sandbox up for good: it (and its tab) is no longer owed. */
+export function noteSleptRemotePageAnswered(page: Page & { workspaceId: string }): void {
+  if (
+    isBrowserSandboxUp(useAppStore.getState(), page.worktreeId, readSleepInFlight(page.worktreeId))
+  ) {
+    settleReopenAfterSleep(page.id)
+    settleReopenAfterSleep(page.workspaceId)
+  }
+}
+
+/** Opens a page owed since sleep the moment the sandbox comes up (nothing else may ask again). */
+export function useReopenSleptRemotePage(
+  page: Page,
+  lifecycle: { open: () => unknown; tokens: { remotePage: string | null } },
+  isMounted: () => boolean
+): void {
+  const suspendCommitted = useIsWorktreeSuspendCommitted(page.worktreeId)
+  const sandboxUp = useAppStore(
+    (state) =>
+      !suspendCommitted &&
+      isBrowserSandboxUp(state, page.worktreeId, readSleepInFlight(page.worktreeId))
+  )
+  const wasUpRef = useRef(sandboxUp)
+  useEffect(() => {
+    const cameUp = sandboxUp && !wasUpRef.current
+    wasUpRef.current = sandboxUp
+    if (
+      cameUp &&
+      isMounted() &&
+      isAwaitingReopenAfterSleep(page.id) &&
+      lifecycle.tokens.remotePage === null
+    ) {
+      lifecycle.open()
+    }
+  }, [isMounted, lifecycle, page.id, sandboxUp])
+}

src/renderer/src/components/sidebar/workspace-actions/sleep/browser-tab-sleep-wake-race.test.ts

@@ -0,0 +1,132 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const store = vi.hoisted(() => ({
+  state: {
+    browserTabsByWorktree: {
+      wt: [
+        { id: 'tab-a', worktreeId: 'wt' },
+        { id: 'tab-b', worktreeId: 'wt' }
+      ]
+    },
+    browserPagesByWorkspace: { 'tab-a': [{ id: 'page-a' }], 'tab-b': [{ id: 'page-b' }] },
+    asleepBrowserWorktreeIds: {}
+  }
+}))
+vi.mock('@/store', () => ({
+  useAppStore: {
+    getState: () => store.state,
+    setState: () => {}
+  }
+}))
+vi.mock('@/store/slices/browser/cleanup/browser-webview-cleanup', () => ({
+  destroyWorktreeBrowserGuests: () => {}
+}))
+
+const {
+  isAwaitingReopenAfterSleep,
+  isBrowserSandboxUp,
+  isSleptBrowserTabAwaitingReopen,
+  missingRemotePageAction,
+  resetReopenAfterSleepForTest,
+  settleReopenAfterSleep,
+  sleepWorktreeBrowserTabs
+} = await import('./browser-tab-sleep')
+
+type Vm = 'running' | 'suspended'
+const WT = 'wt'
+
+/** One moment of the race: the VM record, a wake in flight, and the window's sleep flags. */
+function moment(args: {
+  vm: Vm
+  waking?: boolean
+  asleepMark?: boolean
+  sleepIntent?: boolean
+  goingToSleep?: boolean
+  suspendCommitted?: boolean
+}) {
+  const waking: Record<string, boolean> = {}
+  const asleep: Record<string, number> = {}
+  if (args.waking) {
+    waking[WT] = true
+  }
+  if (args.asleepMark) {
+    asleep[WT] = 1
+  }
+  const state = {
+    ephemeralVmRuntimesByWorkspaceId: { [WT]: { status: args.vm } },
+    wakingEphemeralVmWorkspaceIds: waking,
+    asleepBrowserWorktreeIds: asleep
+  }
+  const sleep = {
+    sleepIntent: Boolean(args.sleepIntent),
+    goingToSleep: Boolean(args.goingToSleep),
+    suspendCommitted: Boolean(args.suspendCommitted)
+  }
+  return {
+    up: isBrowserSandboxUp(state, WT, sleep),
+    onMissing: (pageId: string) =>
+      missingRemotePageAction({
+        awaitingReopen: isAwaitingReopenAfterSleep(pageId),
+        asleep: Boolean(args.asleepMark),
+        sandboxUp: isBrowserSandboxUp(state, WT, sleep)
+      })
+  }
+}
+
+beforeEach(() => {
+  resetReopenAfterSleepForTest()
+})
+
+describe('wake pressed while the suspend is still being committed (LC2-P1)', () => {
+  it('keeps both tabs: they wait through the race and reopen once the sandbox is up', () => {
+    sleepWorktreeBrowserTabs(WT, 1)
+    expect(isAwaitingReopenAfterSleep('page-a')).toBe(true)
+
+    // +8 s: "Wake to continue" pressed. The asleep mark is cleared, but main is still
+    // committing the suspend and the VM record still reads "running": hosts pages vanish.
+    const pressed = moment({ vm: 'running', sleepIntent: true, suspendCommitted: true })
+    expect(pressed.up).toBe(false)
+    expect(pressed.onMissing('page-a')).toBe('wait')
+    expect(pressed.onMissing('page-b')).toBe('wait')
+
+    // ~4 s later the VM is suspended, then waking: still nothing opens or closes.
+    expect(moment({ vm: 'suspended' }).onMissing('page-a')).toBe('wait')
+    expect(moment({ vm: 'suspended', waking: true }).onMissing('page-b')).toBe('wait')
+    expect(isAwaitingReopenAfterSleep('page-a')).toBe(true)
+    expect(isAwaitingReopenAfterSleep('page-b')).toBe(true)
+
+    // Live: the missing pages reopen in place instead of closing their tabs.
+    const live = moment({ vm: 'running' })
+    expect(live.up).toBe(true)
+    expect(live.onMissing('page-a')).toBe('reopen')
+    expect(live.onMissing('page-b')).toBe('reopen')
+    // A second "missing" before the new page answers still reopens, never closes.
+    expect(live.onMissing('page-a')).toBe('reopen')
+
+    // The new host page answered: from now on a missing page is a real close.
+    settleReopenAfterSleep('page-a')
+    expect(live.onMissing('page-a')).toBe('close')
+    expect(live.onMissing('page-b')).toBe('reopen')
+  })
+
+  it('an ordinary sleep and wake still waits while asleep and reopens after', () => {
+    sleepWorktreeBrowserTabs(WT, 1)
+    expect(moment({ vm: 'suspended', asleepMark: true }).onMissing('page-a')).toBe('wait')
+    expect(moment({ vm: 'running' }).onMissing('page-a')).toBe('reopen')
+  })
+
+  it('a tab that never slept closes when its host page is gone, as before', () => {
+    expect(moment({ vm: 'running' }).onMissing('page-x')).toBe('close')
+  })
+
+  it("the host's session snapshot can't drop a slept tab before its page is reopened", () => {
+    const pages = { browserPagesByWorkspace: store.state.browserPagesByWorkspace }
+    expect(isSleptBrowserTabAwaitingReopen(pages, 'tab-a')).toBe(false)
+    sleepWorktreeBrowserTabs(WT, 1)
+    expect(isSleptBrowserTabAwaitingReopen(pages, 'tab-a')).toBe(true)
+    settleReopenAfterSleep('tab-a')
+    settleReopenAfterSleep('page-a')
+    expect(isSleptBrowserTabAwaitingReopen(pages, 'tab-a')).toBe(false)
+    expect(isSleptBrowserTabAwaitingReopen(pages, 'tab-b')).toBe(true)
+  })
+})

src/renderer/src/components/sidebar/workspace-actions/sleep/browser-tab-sleep.ts

@@ -6,6 +6,8 @@
 // wake. The mark lives outside the layout (asleepBrowserWorktreeIds), so sleep writes no layout.
 import { useAppStore } from '@/store'
 import { destroyWorktreeBrowserGuests } from '@/store/slices/browser/cleanup/browser-webview-cleanup'
+import { hasWorktreeSleepIntent } from '@/lib/worktree-sleep-intent'
+import { isWorktreeGoingToSleep, isWorktreeSuspendCommitted } from '@/lib/worktree-sleep-view'
 import type { BrowserWorkspace } from '../../../../../../shared/browser-workspace-types'
 
 type BrowserTabsByWorktree = Record<string, BrowserWorkspace[]>
@@ -46,9 +48,92 @@ export function selectAwakeBrowserTabsByWorktree(state: {
 // reported missing after the wake, its tab opens a new page at the same URL instead of closing.
 const reopenAfterSleep = new Set<string>()
 
-/** True once per slept page: the pane recreates the host page instead of closing the tab. */
-export function consumeReopenAfterSleep(browserPageId: string): boolean {
-  return reopenAfterSleep.delete(browserPageId)
+/**
+ * A slept page still owed a new host page. It stays owed until a host page answers for it again
+ * (a new one opened, or the old one still there), however many times "missing" is reported on
+ * the way: a wake that races the suspend reports it while the VM is still going down (LC2-P1).
+ */
+export function isAwaitingReopenAfterSleep(browserPageId: string): boolean {
+  return reopenAfterSleep.has(browserPageId)
+}
+
+/** A tab (or any of its pages) still owed a host page since sleep: host snapshots can't drop it. */
+export function isSleptBrowserTabAwaitingReopen(
+  state: { browserPagesByWorkspace: Record<string, readonly { id: string }[]> },
+  tabId: string
+): boolean {
+  return (
+    reopenAfterSleep.has(tabId) ||
+    (state.browserPagesByWorkspace[tabId] ?? []).some((page) => reopenAfterSleep.has(page.id))
+  )
+}
+
+/** A host page answered for this tab again: from now on a missing page is a real close. */
+export function settleReopenAfterSleep(browserPageId: string): void {
+  reopenAfterSleep.delete(browserPageId)
+}
+
+export function resetReopenAfterSleepForTest(): void {
+  reopenAfterSleep.clear()
+}
+
+type SandboxState = {
+  ephemeralVmRuntimesByWorkspaceId?: Record<string, { status: string } | undefined>
+  wakingEphemeralVmWorkspaceIds?: Record<string, boolean | undefined>
+  asleepBrowserWorktreeIds?: Record<string, number>
+}
+
+/** This window's view of a sleep in flight for one workspace (workspace-sleep-follower). */
+export type SleepInFlight = {
+  /** Sleep started and has not ended (also while the user presses Wake to continue). */
+  sleepIntent: boolean
+  goingToSleep: boolean
+  /** Main is committing the suspend: the VM still reads "running" but is going down. */
+  suspendCommitted: boolean
+}
+
+export function readSleepInFlight(worktreeId: string): SleepInFlight {
+  return {
+    sleepIntent: hasWorktreeSleepIntent(worktreeId),
+    goingToSleep: isWorktreeGoingToSleep(worktreeId),
+    suspendCommitted: isWorktreeSuspendCommitted(worktreeId)
+  }
+}
+
+/**
+ * The workspace's host can run pages that will last: no sleep in flight (a wake pressed while
+ * the suspend is still being committed included), not marked asleep, not waking, and its VM (if
+ * it has one) running. A page opened before all of this holds would die with the suspend.
+ */
+export function isBrowserSandboxUp(
+  state: SandboxState,
+  worktreeId: string,
+  sleep: SleepInFlight
+): boolean {
+  const runtime = state.ephemeralVmRuntimesByWorkspaceId?.[worktreeId]
+  return (
+    !sleep.sleepIntent &&
+    !sleep.goingToSleep &&
+    !sleep.suspendCommitted &&
+    !state.asleepBrowserWorktreeIds?.[worktreeId] &&
+    !state.wakingEphemeralVmWorkspaceIds?.[worktreeId] &&
+    (runtime === undefined || runtime.status === 'running')
+  )
+}
+
+/** What a pane does when its host page is reported missing (NEW-01, LC2-P1). */
+export function missingRemotePageAction(args: {
+  awaitingReopen: boolean
+  asleep: boolean
+  sandboxUp: boolean
+}): 'wait' | 'reopen' | 'close' {
+  if (args.asleep) {
+    return 'wait'
+  }
+  if (args.awaitingReopen) {
+    return args.sandboxUp ? 'reopen' : 'wait'
+  }
+  return 'close'
 }
 
 /** Sleep: drop the pages, keep the tabs, mark the workspace's tabs asleep. */

src/renderer/src/runtime/web-session-tabs-sync/apply-preparation-browser.ts

@@ -1,4 +1,5 @@
 import { buildWebSessionExistingTabIndex } from './placement/web-session-existing-tab-index'
+import { isSleptBrowserTabAwaitingReopen } from '@/components/sidebar/workspace-actions/sleep/browser-tab-sleep'
 import { resolveWebSessionVisibleTabId } from './navigation/web-session-focus-intent'
 import { chooseTargetGroupId, buildHostGroupIdByTabId } from './tab-group-layout-tree'
 import {
@@ -78,6 +79,11 @@ export function prepareWebSessionTabsSnapshotBrowser(
         if (mirroredBrowserWorkspaceIds.has(tab.id)) {
           return true
         }
+        // Slept with its VM (NEW-01, LC2-P1): the host's pages died with it, so its snapshot no
+        // longer lists them. The tab stays and reopens its page after the wake.
+        if (isSleptBrowserTabAwaitingReopen(state, tab.id)) {
+          return false
+        }
         if (!browserWorkspaceHasRemoteEnvironmentPage(state, tab, environmentId)) {
           return false
         }

src/renderer/src/runtime/web-session-tabs-sync/placement/web-session-tabs-sync-slept-browser-tabs.test.ts

@@ -0,0 +1,132 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import type { BrowserPage, BrowserWorkspace } from '../../../../../shared/browser-workspace-types'
+import type { Tab } from '../../../../../shared/tab-types'
+import { applyWebSessionTabsSnapshot } from '../snapshot-api'
+import {
+  ENV,
+  NOW,
+  WT,
+  makeSnapshot,
+  makeState,
+  resetWebSessionTabsSyncTestState
+} from '../testing/web-session-tabs-sync.test-fixture'
+import {
+  resetReopenAfterSleepForTest,
+  settleReopenAfterSleep,
+  sleepWorktreeBrowserTabs
+} from '@/components/sidebar/workspace-actions/sleep/browser-tab-sleep'
+
+const sleepStore = vi.hoisted(() => {
+  const holder: { state: object } = { state: {} }
+  return holder
+})
+vi.mock('../../../store/index', () => ({
+  useAppStore: {
+    getState: () => sleepStore.state,
+    setState: vi.fn()
+  }
+}))
+vi.mock('@/store', () => ({
+  useAppStore: {
+    getState: () => sleepStore.state,
+    setState: vi.fn()
+  }
+}))
+vi.mock('@/store/slices/browser/cleanup/browser-webview-cleanup', () => ({
+  destroyWorktreeBrowserGuests: () => {}
+}))
+
+function hostRunBrowserTab(id: string) {
+  const workspace: BrowserWorkspace = {
+    id: `ws-${id}`,
+    worktreeId: WT,
+    activePageId: `page-${id}`,
+    pageIds: [`page-${id}`],
+    url: `https://${id}.example/`,
+    title: id,
+    loading: false,
+    faviconUrl: null,
+    canGoBack: false,
+    canGoForward: false,
+    loadError: null,
+    createdAt: NOW - 10
+  }
+  const page: BrowserPage = {
+    id: `page-${id}`,
+    workspaceId: workspace.id,
+    worktreeId: WT,
+    url: workspace.url,
+    title: id,
+    loading: false,
+    faviconUrl: null,
+    canGoBack: false,
+    canGoForward: false,
+    loadError: null,
+    createdAt: NOW - 10
+  }
+  const unified: Tab = {
+    id: `unified-${id}`,
+    entityId: workspace.id,
+    groupId: 'group',
+    worktreeId: WT,
+    contentType: 'browser',
+    label: id,
+    customLabel: null,
+    color: null,
+    sortOrder: 0,
+    createdAt: NOW - 10,
+    isPreview: false,
+    isPinned: false
+  }
+  return { workspace, page, unified }
+}
+
+describe('host snapshot after Sleep (LC2-P1)', () => {
+  beforeEach(() => {
+    resetWebSessionTabsSyncTestState()
+    resetReopenAfterSleepForTest()
+  })
+
+  it("keeps slept tabs the host no longer lists (their pages died with the VM), then culls once they're answered", () => {
+    const a = hostRunBrowserTab('a')
+    const b = hostRunBrowserTab('b')
+    const state = makeState({
+      browserTabsByWorktree: { [WT]: [a.workspace, b.workspace] },
+      browserPagesByWorkspace: { [a.workspace.id]: [a.page], [b.workspace.id]: [b.page] },
+      remoteBrowserPageHandlesByPageId: {
+        [a.page.id]: { environmentId: ENV, remotePageId: 'host-a' },
+        [b.page.id]: { environmentId: ENV, remotePageId: 'host-b' }
+      },
+      unifiedTabsByWorktree: { [WT]: [a.unified, b.unified] },
+      groupsByWorktree: {
+        [WT]: [
+          {
+            id: 'group',
+            worktreeId: WT,
+            activeTabId: b.unified.id,
+            tabOrder: [a.unified.id, b.unified.id],
+            recentTabIds: [b.unified.id]
+          }
+        ]
+      }
+    })
+    sleepStore.state = state
+    sleepWorktreeBrowserTabs(WT, NOW)
+
+    // The host's snapshot after the suspend lists no browser pages.
+    const asleep = applyWebSessionTabsSnapshot(state, makeSnapshot([]), ENV, NOW)
+    const keptIds = (asleep.browserTabsByWorktree ?? state.browserTabsByWorktree)[WT]?.map(
+      (tab) => tab.id
+    )
+    expect(keptIds).toEqual([a.workspace.id, b.workspace.id])
+    expect(
+      (asleep.unifiedTabsByWorktree ?? state.unifiedTabsByWorktree)[WT]?.map((tab) => tab.id)
+    ).toEqual([a.unified.id, b.unified.id])
+
+    // Once a host page answered for "a" after the wake, a snapshot without it is a real close.
+    settleReopenAfterSleep(a.page.id)
+    settleReopenAfterSleep(a.workspace.id)
+    const later = applyWebSessionTabsSnapshot(state, makeSnapshot([]), ENV, NOW)
+    expect(later.browserTabsByWorktree?.[WT]?.map((tab) => tab.id)).toEqual([b.workspace.id])
+  })
+})