diff --git a/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.test.ts b/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.test.ts new file mode 100644 index 00000000000..03badd28f94 --- /dev/null +++ b/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.test.ts @@ -0,0 +1,136 @@ +import { describe, expect, it, vi } from 'vitest' +import { getDefaultPersistedState, getDefaultWorkspaceSession } from '../../../shared/constants' +import type { SshRemotePtyLease } from '../../../shared/ssh-types' +import { toComparableRelaySshPtyId } from '../../../shared/ssh-pty-id' +import { clearSshRemotePtyBindingsForLeases } from './ssh-pty-binding-cleanup' + +function fixture(count: number) { + const state = getDefaultPersistedState('/home/test') + state.workspaceSession = getDefaultWorkspaceSession() + state.workspaceSession.tabsByWorktree.wt = Array.from({ length: count }, (_, i) => ({ + id: `tab-${i}`, + worktreeId: 'wt', + ptyId: `pty-${i}`, + title: '', + customTitle: null, + color: null, + sortOrder: i, + createdAt: 1 + })) + const leases: SshRemotePtyLease[] = Array.from({ length: count }, (_, i) => ({ + targetId: 'ssh-one', + ptyId: `pty-${i}`, + tabId: `tab-${i}`, + worktreeId: 'wt', + state: 'detached', + createdAt: 1, + updatedAt: 1 + })) + return { + state, + leases, + toComparablePtyId: vi.fn((_target: string, ptyId: string) => ptyId), + scheduleSave: vi.fn() + } +} + +describe('SSH binding cleanup indexing', () => { + it('normalizes each binding once across a large lease inventory', () => { + const operations = fixture(1000) + expect(clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases)).toBe(true) + expect(operations.toComparablePtyId).toHaveBeenCalledTimes(1000) + expect( + operations.state.workspaceSession!.tabsByWorktree.wt.every((tab) => tab.ptyId === null) + ).toBe(true) + expect(operations.scheduleSave).toHaveBeenCalledTimes(1) + }) + + it('retains foreign hosts and conflicting tab/workspace leases', () => { + const operations = fixture(4) + operations.leases[0].targetId = 'ssh-two' + operations.leases[1].tabId = 'other-tab' + operations.leases[2].worktreeId = 'other-workspace' + delete operations.leases[3].tabId + clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases) + expect(operations.state.workspaceSession!.tabsByWorktree.wt.map((tab) => tab.ptyId)).toEqual([ + 'pty-0', + 'pty-1', + 'pty-2', + null + ]) + }) + + it('matches layout leaves against every lease for a PTY while preserving leaf conflicts', () => { + const operations = fixture(1) + const session = operations.state.workspaceSession! + session.tabsByWorktree.wt[0].ptyId = null + session.terminalLayoutsByTabId['tab-0'] = { + root: null, + activeLeafId: null, + expandedLeafId: null, + ptyIdsByLeafId: { + matched: 'pty-0', + protected: 'pty-0', + wildcard: 'pty-1' + } + } + const lease = operations.leases[0] + operations.leases = [ + { ...lease, leafId: 'wrong' }, + { ...lease, leafId: 'matched' }, + { ...lease, ptyId: 'pty-1' } + ] + expect(clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases)).toBe(true) + expect(session.terminalLayoutsByTabId['tab-0'].ptyIdsByLeafId).toEqual({ + protected: 'pty-0' + }) + expect(operations.toComparablePtyId).toHaveBeenCalledTimes(3) + }) + + it('normalizes app-form binding ids onto the relay-form lease key', () => { + // Leases store the relay-local id; sessions may hold the app-wide "ssh:@@" form. + // The index key is the normalized form, so both spellings still name the same PTY. + const operations = fixture(1) + operations.toComparablePtyId = vi.fn(toComparableRelaySshPtyId) + operations.state.workspaceSession!.tabsByWorktree.wt[0].ptyId = 'ssh:ssh-one@@pty-0' + + expect(clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases)).toBe(true) + expect(operations.state.workspaceSession!.tabsByWorktree.wt[0].ptyId).toBeNull() + }) + + it('keeps a binding whose app-form id names a different SSH target', () => { + // Relay-local ids collide across targets ("pty-0" exists on every host). Clearing ssh-one must + // never scrub a pane still bound to a live ssh-two shell. + const operations = fixture(1) + operations.toComparablePtyId = vi.fn(toComparableRelaySshPtyId) + operations.state.workspaceSession!.tabsByWorktree.wt[0].ptyId = 'ssh:ssh-two@@pty-0' + + expect(clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases)).toBe(false) + expect(operations.state.workspaceSession!.tabsByWorktree.wt[0].ptyId).toBe('ssh:ssh-two@@pty-0') + expect(operations.scheduleSave).not.toHaveBeenCalled() + }) + + it('keeps every binding when no lease names its PTY', () => { + // A bucket miss must fail closed: leak a stale id rather than unbind a live pane. + const operations = fixture(2) + for (const lease of operations.leases) { + lease.ptyId = `unrelated-${lease.ptyId}` + } + + expect(clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases)).toBe(false) + expect(operations.state.workspaceSession!.tabsByWorktree.wt.map((tab) => tab.ptyId)).toEqual([ + 'pty-0', + 'pty-1' + ]) + expect(operations.scheduleSave).not.toHaveBeenCalled() + }) + + it('does not index leases when the session holds no bindings to check', () => { + const operations = fixture(500) + operations.state.workspaceSession!.tabsByWorktree = {} + operations.state.workspaceSession!.terminalLayoutsByTabId = {} + + expect(clearSshRemotePtyBindingsForLeases(operations, 'ssh-one', operations.leases)).toBe(false) + expect(operations.toComparablePtyId).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.ts b/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.ts index 65b879795dc..c467e0fc3ab 100644 --- a/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.ts +++ b/src/main/persistence/leasing-ssh-ptys/ssh-pty-binding-cleanup.ts @@ -9,8 +9,8 @@ export type SshPtyBindingCleanupOperations = { scheduleSave: () => void } +/** `binding.ptyId` must already be in lease-comparable (relay) form; callers normalize it. */ function sshRemotePtyLeaseMayReferenceBinding( - operations: SshPtyBindingCleanupOperations, lease: SshRemotePtyLease, binding: { ptyId: string @@ -20,8 +20,7 @@ function sshRemotePtyLeaseMayReferenceBinding( leafId?: string } ): boolean { - const bindingPtyId = operations.toComparablePtyId(binding.targetId, binding.ptyId) - if (lease.targetId !== binding.targetId || lease.ptyId !== bindingPtyId) { + if (lease.targetId !== binding.targetId || lease.ptyId !== binding.ptyId) { return false } // Why: target removal is destructive; scrub matching bindings before deleting the lease, else removing the tombstone can revive stale PTY ids. @@ -50,6 +49,33 @@ export function clearSshRemotePtyBindingsForLeases( if (!leases?.length) { return false } + // Keyed by the stored (relay) pty id, which is the only form a lease holds; every lookup below + // normalizes the binding id to that form first, so a bucket miss means "no lease names this pty" + // and the binding is KEPT. Failing closed here leaves a stale id to be retired on reattach, + // where clearing on a bad match would strand a live remote shell behind a respawned pane. + let leasesByPtyId: Map | undefined + const referencesBinding = ( + binding: Parameters[1] + ): boolean => { + if (!leasesByPtyId) { + leasesByPtyId = new Map() + for (const lease of leases) { + if (lease.targetId !== targetId) { + continue + } + const entries = leasesByPtyId.get(lease.ptyId) + if (entries) { + entries.push(lease) + } else { + leasesByPtyId.set(lease.ptyId, [lease]) + } + } + } + const ptyId = operations.toComparablePtyId(binding.targetId, binding.ptyId) + return (leasesByPtyId.get(ptyId) ?? []).some((lease) => + sshRemotePtyLeaseMayReferenceBinding(lease, { ...binding, ptyId }) + ) + } let changed = false const sessions = new Set( [ @@ -62,14 +88,7 @@ export function clearSshRemotePtyBindingsForLeases( for (const tab of tabs) { if ( tab.ptyId && - leases.some((lease) => - sshRemotePtyLeaseMayReferenceBinding(operations, lease, { - ptyId: tab.ptyId!, - worktreeId, - targetId, - tabId: tab.id - }) - ) + referencesBinding({ ptyId: tab.ptyId, worktreeId, targetId, tabId: tab.id }) ) { tab.ptyId = null changed = true @@ -92,16 +111,7 @@ export function clearSshRemotePtyBindingsForLeases( const worktreeId = worktreeIdByTabId.get(tabId) const nextBindings = Object.fromEntries( Object.entries(bindings).filter( - ([leafId, ptyId]) => - !leases.some((lease) => - sshRemotePtyLeaseMayReferenceBinding(operations, lease, { - ptyId, - targetId, - worktreeId, - tabId, - leafId - }) - ) + ([leafId, ptyId]) => !referencesBinding({ ptyId, targetId, worktreeId, tabId, leafId }) ) ) if (Object.keys(nextBindings).length !== Object.keys(bindings).length) {