diff --git a/src/main/runtime/expired-ssh-lease-pane-candidacy.test.ts b/src/main/runtime/expired-ssh-lease-pane-candidacy.test.ts index 82986a4b3fd..3d8157c9a28 100644 --- a/src/main/runtime/expired-ssh-lease-pane-candidacy.test.ts +++ b/src/main/runtime/expired-ssh-lease-pane-candidacy.test.ts @@ -12,6 +12,12 @@ const TARGET = 'ssh-target' const TAB_ID = 'tab-candidacy' type LeaseReader = { + workspaceSessionWorktreeHasRuntimeOwnedPtyCandidate: ( + session: { terminalLayoutsByTabId?: Record }, + worktreeId: string, + tabs: { id: string; ptyId: string | null }[] + ) => boolean + collectRecentExpiredSshLeaseTabIds: (worktreeId: string) => ReadonlySet getRecentExpiredSshLease: ( worktreeId: string, tabId: string, @@ -87,4 +93,45 @@ describe('recent expired SSH lease candidacy', () => { ) expect(reader.hasRecentExpiredSshLeasePane(TEST_WORKTREE_ID, pane)).toBe(true) }) + + it('collects the same tabs the per-tab reader reports, in one sweep of the leases', () => { + const leases = [leaseFor('pty-1', { supersededBy: 'pty-2' }), leaseFor('pty-2')] + let sweeps = 0 + const reader = new OrcaRuntimeService({ + ...store, + getSshRemotePtyLeases: () => { + sweeps += 1 + return leases + } + }) as unknown as LeaseReader + const tabs = Array.from({ length: 8 }, (_, index) => ({ + id: index === 7 ? TAB_ID : `tab-${index}`, + ptyId: null + })) + + expect( + reader.workspaceSessionWorktreeHasRuntimeOwnedPtyCandidate( + { terminalLayoutsByTabId: {} }, + TEST_WORKTREE_ID, + tabs + ) + ).toBe(true) + // One sweep answers all eight tabs; the per-tab reader used to sweep once per tab. + expect(sweeps).toBe(1) + expect([...reader.collectRecentExpiredSshLeaseTabIds(TEST_WORKTREE_ID)]).toEqual([TAB_ID]) + }) + + it('reports no candidate when no lease names any of the worktree tabs', () => { + const reader = readerWithLeases([ + leaseFor('pty-1', { tabId: 'somewhere-else', leafId: undefined }) + ]) + + expect( + reader.workspaceSessionWorktreeHasRuntimeOwnedPtyCandidate( + { terminalLayoutsByTabId: {} }, + TEST_WORKTREE_ID, + [{ id: TAB_ID, ptyId: null }] + ) + ).toBe(false) + }) }) diff --git a/src/main/runtime/orca-runtime-reconcile-headless-mobile-session-browser-tabs.ts b/src/main/runtime/orca-runtime-reconcile-headless-mobile-session-browser-tabs.ts index c4687c96d14..a9f377e5a7b 100644 --- a/src/main/runtime/orca-runtime-reconcile-headless-mobile-session-browser-tabs.ts +++ b/src/main/runtime/orca-runtime-reconcile-headless-mobile-session-browser-tabs.ts @@ -82,17 +82,24 @@ export class OrcaRuntimeWithReconcileHeadlessMobileSessionBrowserTabs extends Or worktreeId: string, tabs: WorkspaceSessionState['tabsByWorktree'][string] ): boolean { + // Why resolved lazily and reused: the per-tab question is the same lease sweep with a + // different tabId, so asking it once per worktree answers every tab. Kept lazy so a + // worktree whose first tab already owns a serve/SSH pty never sweeps at all. + let recoverableTabIds: ReadonlySet | undefined return tabs.some((tab) => { if (this.isServeOrSshOwnedPtyId(tab.ptyId)) { return true } const leafPtyIds = session.terminalLayoutsByTabId?.[tab.id]?.ptyIdsByLeafId - return ( - (leafPtyIds && - Object.values(leafPtyIds).some((ptyId) => this.isServeOrSshOwnedPtyId(ptyId))) || - // Why: expiry keeps pane coordinates so paired viewers can request a fresh shell. - this.getRecentExpiredSshLease(worktreeId, tab.id, undefined) !== null - ) + if ( + leafPtyIds && + Object.values(leafPtyIds).some((ptyId) => this.isServeOrSshOwnedPtyId(ptyId)) + ) { + return true + } + // Why: expiry keeps pane coordinates so paired viewers can request a fresh shell. + recoverableTabIds ??= this.collectRecentExpiredSshLeaseTabIds(worktreeId) + return recoverableTabIds.has(tab.id) }) } @@ -128,6 +135,54 @@ export class OrcaRuntimeWithReconcileHeadlessMobileSessionBrowserTabs extends Or ) } + /** + * Why eligibility belongs in the selection, not after it: a pane accumulates leases as it + * re-leases under new relay ids, so `(worktreeId, tabId, leafId)` names several. A superseded or + * relay-id-recycled predecessor is `expired` for a reason that already names its successor, and + * the unqualified callers use this answer to decide a pane is still recoverable — reporting one + * would offer paired viewers a recovery `recoverTerminalPane` then refuses. Picking the first + * ELIGIBLE orphan also keeps a predecessor from shadowing the successor that is genuinely + * reattachable. + */ + private isRecentExpiredSshLeaseForWorktree( + lease: ReturnType>[number], + worktreeId: string, + now: number + ): boolean { + return ( + lease.state === 'expired' && + lease.worktreeId === worktreeId && + sshRemotePtyLeaseAllowsReattach(lease) && + lease.updatedAt <= now && + now - lease.updatedAt <= SSH_PANE_RECOVERY_GRACE_MS + ) + } + + /** + * Leaf is the pane's identity; the frozen tabId is only trustworthy while nothing else can say + * where the leaf actually lives. + */ + private resolveExpiredSshLeaseTabId( + lease: ReturnType>[number] + ): string { + const currentTabId = lease.leafId + ? this.findCurrentTerminalTabIdForLeaf(lease.targetId, lease.leafId) + : undefined + return currentTabId ?? lease.tabId + } + + /** The tabs a recent eligible expired lease still names, resolved in one sweep of the leases. */ + protected collectRecentExpiredSshLeaseTabIds(worktreeId: string): ReadonlySet { + const now = Date.now() + const tabIds = new Set() + for (const lease of this.store?.getSshRemotePtyLeases?.() ?? []) { + if (this.isRecentExpiredSshLeaseForWorktree(lease, worktreeId, now)) { + tabIds.add(this.resolveExpiredSshLeaseTabId(lease)) + } + } + return tabIds + } + protected getRecentExpiredSshLease( worktreeId: string, tabId: string, @@ -137,34 +192,17 @@ export class OrcaRuntimeWithReconcileHeadlessMobileSessionBrowserTabs extends Or const now = Date.now() return ( this.store?.getSshRemotePtyLeases?.().find((lease) => { - if (lease.state !== 'expired' || lease.worktreeId !== worktreeId) { + if (!this.isRecentExpiredSshLeaseForWorktree(lease, worktreeId, now)) { return false } - // Why eligibility belongs in the selection, not after it: a pane accumulates leases as it - // re-leases under new relay ids, so `(worktreeId, tabId, leafId)` names several. A - // superseded or relay-id-recycled predecessor is `expired` for a reason that already names - // its successor, and the unqualified callers use this answer to decide a pane is still - // recoverable — reporting one would offer paired viewers a recovery `recoverTerminalPane` - // then refuses. Picking the first ELIGIBLE orphan also keeps a predecessor from shadowing - // the successor that is genuinely reattachable. - if (!sshRemotePtyLeaseAllowsReattach(lease)) { - return false - } - // Leaf is the pane's identity; the frozen tabId is only trustworthy while nothing else can - // say where the leaf actually lives. - const currentTabId = lease.leafId - ? this.findCurrentTerminalTabIdForLeaf(lease.targetId, lease.leafId) - : undefined return ( - (currentTabId ?? lease.tabId) === tabId && + this.resolveExpiredSshLeaseTabId(lease) === tabId && // Leases store RELAY form (`toStoredPtyId` -> `toRelaySshPtyId`); the runtime hands us // the APP form (`ssh:@@pty-3`). A raw `===` therefore never held for an SSH // pane, which is what kept this reader's only ptyId-qualified caller inert. (ptyId === undefined || lease.ptyId === toComparableRelaySshPtyId(lease.targetId, ptyId)) && - (leafId === undefined || lease.leafId === undefined || lease.leafId === leafId) && - lease.updatedAt <= now && - now - lease.updatedAt <= SSH_PANE_RECOVERY_GRACE_MS + (leafId === undefined || lease.leafId === undefined || lease.leafId === leafId) ) }) ?? null ) diff --git a/src/main/runtime/terminal-leaf-membership-index.test.ts b/src/main/runtime/terminal-leaf-membership-index.test.ts new file mode 100644 index 00000000000..fcaf51f1937 --- /dev/null +++ b/src/main/runtime/terminal-leaf-membership-index.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it } from 'vitest' + +import type { TerminalLayoutSnapshot } from '../../shared/terminal-tab-types' +import type { WorkspaceSessionState } from '../../shared/workspace-session-state-types' +import { findTerminalTabIdForLeaf } from './workspace-session-terminal-membership-authority' +import { getTerminalLeafMembershipIndex } from './terminal-leaf-membership-index' + +function layout(...leafIds: string[]): TerminalLayoutSnapshot { + let root = { type: 'leaf' as const, leafId: leafIds[0] } + for (const leafId of leafIds.slice(1)) { + root = { + type: 'split', + direction: 'row', + first: root, + second: { type: 'leaf' as const, leafId } + } as never + } + return { root, activeLeafId: leafIds[0], ptyIdsByLeafId: {} } as TerminalLayoutSnapshot +} + +function session(layouts: Record): WorkspaceSessionState { + return { terminalLayoutsByTabId: layouts } as WorkspaceSessionState +} + +describe('getTerminalLeafMembershipIndex', () => { + it('maps every leaf in a split tree to its tab', () => { + const index = getTerminalLeafMembershipIndex({ + 'tab-a': layout('leaf-1', 'leaf-2', 'leaf-3'), + 'tab-b': layout('leaf-4') + }) + expect([...index]).toEqual([ + ['leaf-1', 'tab-a'], + ['leaf-2', 'tab-a'], + ['leaf-3', 'tab-a'], + ['leaf-4', 'tab-b'] + ]) + }) + + it('keeps the first tab in record order when two layouts claim one leaf', () => { + const layouts = { 'tab-a': layout('shared'), 'tab-b': layout('shared') } + expect(getTerminalLeafMembershipIndex(layouts).get('shared')).toBe('tab-a') + expect(findTerminalTabIdForLeaf(session(layouts), 'shared')).toBe('tab-a') + }) + + it('returns the same map instance for an unchanged record', () => { + const layouts = { 'tab-a': layout('leaf-1') } + expect(getTerminalLeafMembershipIndex(layouts)).toBe(getTerminalLeafMembershipIndex(layouts)) + }) + + it('rebuilds when a layout object inside the same record is replaced', () => { + const layouts: Record = { 'tab-a': layout('leaf-1') } + expect(getTerminalLeafMembershipIndex(layouts).get('leaf-1')).toBe('tab-a') + layouts['tab-a'] = layout('leaf-9') + expect(getTerminalLeafMembershipIndex(layouts).get('leaf-1')).toBeUndefined() + expect(getTerminalLeafMembershipIndex(layouts).get('leaf-9')).toBe('tab-a') + }) + + it('rebuilds when a tab is added to or removed from the same record', () => { + const layouts: Record = { 'tab-a': layout('leaf-1') } + expect(getTerminalLeafMembershipIndex(layouts).get('leaf-1')).toBe('tab-a') + layouts['tab-b'] = layout('leaf-2') + expect(getTerminalLeafMembershipIndex(layouts).get('leaf-2')).toBe('tab-b') + delete layouts['tab-a'] + expect(getTerminalLeafMembershipIndex(layouts).get('leaf-1')).toBeUndefined() + }) + + it('answers misses and empty sessions the same way the linear scan did', () => { + expect(findTerminalTabIdForLeaf(undefined, 'leaf-1')).toBeUndefined() + expect(findTerminalTabIdForLeaf(session({}), 'leaf-1')).toBeUndefined() + expect(findTerminalTabIdForLeaf(session({ 'tab-a': layout('leaf-1') }), 'nope')).toBeUndefined() + }) + + it('does not walk a layout tree again once the record is indexed', () => { + let rootReads = 0 + const tracked = { + get root() { + rootReads += 1 + return layout('leaf-1').root + }, + activeLeafId: 'leaf-1', + ptyIdsByLeafId: {} + } as unknown as TerminalLayoutSnapshot + const layouts = { 'tab-a': tracked } + for (let i = 0; i < 50; i += 1) { + findTerminalTabIdForLeaf(session(layouts), `leaf-${i}`) + } + expect(rootReads).toBe(1) + }) +}) diff --git a/src/main/runtime/terminal-leaf-membership-index.ts b/src/main/runtime/terminal-leaf-membership-index.ts new file mode 100644 index 00000000000..faab6db7920 --- /dev/null +++ b/src/main/runtime/terminal-leaf-membership-index.ts @@ -0,0 +1,83 @@ +import type { + TerminalLayoutSnapshot, + TerminalPaneLayoutNode +} from '../../shared/terminal-tab-types' + +type LayoutsByTabId = Record + +type LeafMembershipIndex = { + /** First tab (in record order) whose layout tree holds the leaf — matches the linear scan. */ + readonly tabIdByLeafId: ReadonlyMap + /** Layout objects the index was built from, used to detect in-place record edits. */ + readonly builtFrom: ReadonlyMap +} + +const EMPTY_INDEX: LeafMembershipIndex = { + tabIdByLeafId: new Map(), + builtFrom: new Map() +} + +// Sessions hand back the same layouts record across reads, so one index serves every lookup. +const indexByLayoutsRecord = new WeakMap() + +function addLeafIds( + node: TerminalPaneLayoutNode | null | undefined, + tabId: string, + tabIdByLeafId: Map +): void { + if (!node) { + return + } + if (node.type === 'leaf') { + if (!tabIdByLeafId.has(node.leafId)) { + tabIdByLeafId.set(node.leafId, tabId) + } + return + } + addLeafIds(node.first, tabId, tabIdByLeafId) + addLeafIds(node.second, tabId, tabIdByLeafId) +} + +/** + * Why an identity check and not just the WeakMap: a few writers copy the layouts record and then + * assign into the copy, so the record can be new while most layout objects are shared — and a + * cached index must not outlive a replaced layout. Comparing layout references is O(tabs) pointer + * loads with no allocation, versus the rebuild's set-per-tab tree walk. + */ +function isIndexCurrent(index: LeafMembershipIndex, layouts: LayoutsByTabId): boolean { + let seen = 0 + for (const tabId of Object.keys(layouts)) { + if (index.builtFrom.get(tabId) !== layouts[tabId]) { + return false + } + seen += 1 + } + return seen === index.builtFrom.size +} + +function buildIndex(layouts: LayoutsByTabId): LeafMembershipIndex { + const tabIdByLeafId = new Map() + const builtFrom = new Map() + for (const tabId of Object.keys(layouts)) { + const layout = layouts[tabId] + builtFrom.set(tabId, layout) + addLeafIds(layout?.root, tabId, tabIdByLeafId) + } + return { tabIdByLeafId, builtFrom } +} + +/** leafId -> owning tabId for one session's layouts, rebuilt only when a layout object changes. */ +export function getTerminalLeafMembershipIndex( + layouts: LayoutsByTabId | undefined +): ReadonlyMap { + if (!layouts) { + return EMPTY_INDEX.tabIdByLeafId + } + const cached = indexByLayoutsRecord.get(layouts) + if (cached && isIndexCurrent(cached, layouts)) { + return cached.tabIdByLeafId + } + const built = buildIndex(layouts) + indexByLayoutsRecord.set(layouts, built) + return built.tabIdByLeafId +} diff --git a/src/main/runtime/workspace-session-terminal-membership-authority.ts b/src/main/runtime/workspace-session-terminal-membership-authority.ts index 4c85f46b0c8..b787b1547c8 100644 --- a/src/main/runtime/workspace-session-terminal-membership-authority.ts +++ b/src/main/runtime/workspace-session-terminal-membership-authority.ts @@ -6,6 +6,7 @@ import type { import type { WorkspaceSessionState } from '../../shared/workspace-session-state-types' import { getRepoIdFromWorktreeId } from '../../shared/worktree/id' import { pruneTabGroupLayoutAfterRetirement } from './mobile-session-terminal-retirement' +import { getTerminalLeafMembershipIndex } from './terminal-leaf-membership-index' function collectLeafIds(node: TerminalPaneLayoutNode | null, ids: Set): void { if (!node) { @@ -169,14 +170,7 @@ export function findTerminalTabIdForLeaf( session: WorkspaceSessionState | undefined, leafId: string ): string | undefined { - for (const [tabId, layout] of Object.entries(session?.terminalLayoutsByTabId ?? {})) { - const leafIds = new Set() - collectLeafIds(layout.root, leafIds) - if (leafIds.has(leafId)) { - return tabId - } - } - return undefined + return getTerminalLeafMembershipIndex(session?.terminalLayoutsByTabId).get(leafId) } export function hasHostAuthoritativeTerminalMembership(