perf: index selected-host SSH leases by PTY (#19467)

* perf: index selected-host SSH leases by PTY

* perf(ssh): build lease index lazily and cover pty id normalization

* style(ssh): restore repository oxfmt formatting

---------

Co-authored-by: m4air <m4air@m4airs-MacBook-Air.local>
Co-authored-by: Neil <4138956+nwparker@users.noreply.github.com>

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:<target>@@<id>" 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()
+  })
+})

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<string, SshRemotePtyLease[]> | undefined
+  const referencesBinding = (
+    binding: Parameters<typeof sshRemotePtyLeaseMayReferenceBinding>[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) {