From bccd3a9de608b113c17913036e1392abead66b79 Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:14:35 -0700 Subject: [PATCH] fix(ssh): collapse duplicate PTY leases per pane and reap the orphans (3/5) STA-3077 part 3 of 5. One pane may hold at most one live SSH lease, and a retirement that proves the pane is gone hands its remote PTY to a reaper. - D1a: an incoming lease write supersedes any other live lease for the same pane, flagging the loser for reaping. Transactional with the insert. - D1b: lease states advance monotonically by finality, replacing the ad-hoc directional guards some callers remembered and others forgot. `expired` -> `terminated` stays legal so a relay-expired lease the user then closes is not stranded visible to the 30s recovery consumers. - M1: collapse each pane to one live lease at load, so profiles already poisoned by past fan-outs heal without waiting for a reconnect. - D3: drain `pendingRemoteShutdown` rows on each successful connect, oldest first, capped at 32 per connect and dropped after 7 days. A flagged id this connection is live on is spared, not killed: a reset relay restarts `pty-N` numbering, so a flag minted in an earlier generation can name another pane's shell. - F8: dropping a lease whose spawn failed to persist also undoes the supersede it caused, so a failed spawn cannot retire a working pane's lease. --- src/main/persistence.test.ts | 266 +++++++++++++++++- src/main/persistence.ts | 248 +++++++++++++++- ...ay-session-agent-hooks.integration.test.ts | 4 +- .../ssh/ssh-relay-session-orphan-reap.test.ts | 170 +++++++++++ .../ssh-relay-session-terminal-error.test.ts | 4 +- .../ssh/ssh-relay-session-test-fixtures.ts | 2 + src/main/ssh/ssh-relay-session.ts | 45 +++ 7 files changed, 721 insertions(+), 18 deletions(-) create mode 100644 src/main/ssh/ssh-relay-session-orphan-reap.test.ts diff --git a/src/main/persistence.test.ts b/src/main/persistence.test.ts index df434466aaa..f3dd7ea5b9a 100644 --- a/src/main/persistence.test.ts +++ b/src/main/persistence.test.ts @@ -38,6 +38,7 @@ import { import { folderWorkspaceKey, worktreeWorkspaceKey } from '../shared/workspace-scope' import { toRuntimeExecutionHostId, toSshExecutionHostId } from '../shared/execution-host' import { SshConnectionStore } from './ssh/ssh-connection-store' +import { SSH_LEASE_REAP_BATCH_LIMIT, SSH_LEASE_REAP_MAX_AGE_MS } from './persistence' import { setSourceControlActionDefault } from '../shared/source-control-ai-actions' import { LEGACY_DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS } from '../shared/ssh-types' import { closeTerminalTabInWorkspaceSession } from '../shared/workspace-session-terminal-tab-close' @@ -9425,7 +9426,7 @@ describe('Store', () => { expect(lease.pendingRemoteShutdown).toBeUndefined() }) - it('refuses to re-retire, so a spared lease can never later acquire a reap flag', async () => { + it('lets a spared lease advance to terminated but never acquire a reap flag', async () => { const store = await createStore() store.upsertSshRemotePtyLease({ targetId: 'target-1', @@ -9440,9 +9441,201 @@ describe('Store', () => { store.retireLeaseAndReap('target-1', 'pty-1', 'pane closed') const lease = store.getSshRemotePtyLeases('target-1')[0] - expect(lease.state).toBe('expired') - expect(lease.retiredReason).toBe('relay generation replaced') + // Advancing is required — otherwise a relay-expired lease whose pane is then closed stays visible + // to the 30 s recovery consumers — but the spared row still names a PTY we have no claim on. + expect(lease.state).toBe('terminated') expect(lease.pendingRemoteShutdown).toBeUndefined() + + store.retireLeaseAndReap('target-1', 'pty-1', 'closed again') + expect(store.getSshRemotePtyLeases('target-1')[0].pendingRemoteShutdown).toBeUndefined() + }) + + it('blocks the reproduced resurrection of a retired lease back to detached', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-1', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + store.retireLeaseAndReap('target-1', 'pty-1', 'pane closed') + + // `abandonPtySourceRecovery` writes this unconditionally; it made retired leases reattach-eligible. + store.markSshRemotePtyLease('target-1', 'pty-1', 'detached') + expect(store.getSshRemotePtyLeases('target-1')[0].state).toBe('terminated') + + store.markSshRemotePtyLease('target-1', 'pty-1', 'attached') + expect(store.getSshRemotePtyLeases('target-1')[0].state).toBe('terminated') + }) + + it('supersedes an older lease for the same pane, flagging it for reaping', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-1', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-2', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + + const leases = store.getSshRemotePtyLeases('target-1') + const older = leases.find((l) => l.ptyId === 'pty-1') + const newer = leases.find((l) => l.ptyId === 'pty-2') + expect(older?.state).toBe('terminated') + expect(older?.pendingRemoteShutdown).toBe(true) + expect(newer?.state).toBe('attached') + expect(newer?.pendingRemoteShutdown).toBeUndefined() + }) + + it('leaves leases for other panes and legacy identity-less rows alone when superseding', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-other-pane', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_2, + state: 'attached' + }) + // Legacy pre-leafId row: cannot be grouped, so destroying it would strand a real session. + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-legacy', + state: 'attached' + }) + store.upsertSshRemotePtyLease({ + targetId: 'target-other', + ptyId: 'pty-other-target', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-1', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + + for (const ptyId of ['pty-other-pane', 'pty-legacy', 'pty-other-target', 'pty-1']) { + const lease = store.getSshRemotePtyLeases().find((l) => l.ptyId === ptyId) + expect(lease?.state, ptyId).toBe('attached') + } + }) + + it('undoes the supersede when the spawn that caused it fails to persist', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-1', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + const before = structuredClone(store.getSshRemotePtyLeases('target-1')[0]) + + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-2', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + expect(store.getSshRemotePtyLeases('target-1').find((l) => l.ptyId === 'pty-1')?.state).toBe( + 'terminated' + ) + + store.removeSshRemotePtyLease('target-1', 'pty-2') + + // Restoring must survive D1b, which would refuse this transition through the public setter. + const leases = store.getSshRemotePtyLeases('target-1') + expect(leases).toHaveLength(1) + expect(leases[0]).toEqual(before) + }) + + it('collapses each pane to one live lease on load and keeps the newest', async () => { + const store = await createStore() + const pane = { worktreeId: 'wt1', tabId: 'tab-1', leafId: TEST_LEAF_1 } + // Bypass the supersede by writing rows that are already retired-free but land via distinct panes, + // then rewrite them to share a pane so the file on disk looks like a pre-fix poisoned profile. + for (const ptyId of ['pty-1', 'pty-2', 'pty-3']) { + store.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId, state: 'attached' }) + } + const raw = store.getSshRemotePtyLeases('target-1') + raw[0].updatedAt = 100 + raw[0].createdAt = 100 + raw[1].updatedAt = 300 + raw[1].createdAt = 300 + raw[2].updatedAt = 200 + raw[2].createdAt = 200 + for (const lease of raw) { + Object.assign(lease, pane) + } + store.flush() + + const reloaded = await createStore() + const healed = reloaded.getSshRemotePtyLeases('target-1') + expect(healed.find((l) => l.ptyId === 'pty-2')?.state).toBe('attached') + for (const ptyId of ['pty-1', 'pty-3']) { + const lease = healed.find((l) => l.ptyId === ptyId) + expect(lease?.state, ptyId).toBe('terminated') + expect(lease?.pendingRemoteShutdown, ptyId).toBe(true) + } + }) + + it('leaves an already-clean profile untouched on load', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: 'pty-1', + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + state: 'attached' + }) + store.flush() + const before = structuredClone(store.getSshRemotePtyLeases('target-1')) + + const reloaded = await createStore() + expect(reloaded.getSshRemotePtyLeases('target-1')).toEqual(before) + }) + + it('breaks load-time ties deterministically on createdAt then ptyId', async () => { + const store = await createStore() + for (const ptyId of ['pty-b', 'pty-a', 'pty-c']) { + store.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId, state: 'attached' }) + } + for (const lease of store.getSshRemotePtyLeases('target-1')) { + Object.assign(lease, { + worktreeId: 'wt1', + tabId: 'tab-1', + leafId: TEST_LEAF_1, + updatedAt: 500, + createdAt: 500 + }) + } + store.flush() + + const healed = (await createStore()).getSshRemotePtyLeases('target-1') + expect(healed.find((l) => l.ptyId === 'pty-c')?.state).toBe('attached') + expect(healed.find((l) => l.ptyId === 'pty-a')?.state).toBe('terminated') + expect(healed.find((l) => l.ptyId === 'pty-b')?.state).toBe('terminated') }) it('round-trips retirement fields through normalize on restart', async () => { @@ -9494,6 +9687,73 @@ describe('Store', () => { expect(store.getSshRemotePtyLeases()).toEqual(before) }) + it('claims only flagged leases for the target, oldest first', async () => { + const store = await createStore() + for (const [ptyId, targetId] of [ + ['pty-1', 'target-1'], + ['pty-2', 'target-1'], + ['pty-3', 'target-1'], + ['pty-4', 'target-other'] + ] as const) { + store.upsertSshRemotePtyLease({ targetId, ptyId, state: 'attached' }) + } + store.retireLeaseAndReap('target-1', 'pty-1', 'pane closed') + store.retireLeaseAndReap('target-1', 'pty-2', 'pane closed') + store.retireLeaseSparingPty('target-1', 'pty-3', 'relay pty missing') + store.retireLeaseAndReap('target-other', 'pty-4', 'pane closed') + const leases = store.getSshRemotePtyLeases('target-1') + leases.find((l) => l.ptyId === 'pty-1')!.updatedAt = Date.now() - 1_000 + leases.find((l) => l.ptyId === 'pty-2')!.updatedAt = Date.now() - 9_000 + + expect(store.claimSshRemotePtyLeasesToReap('target-1')).toEqual(['pty-2', 'pty-1']) + }) + + it('keeps the reap flag until the drain reports success, then drops it', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId: 'pty-1', state: 'attached' }) + store.retireLeaseAndReap('target-1', 'pty-1', 'pane closed') + + // Why: a transient shutdown failure must be retried on the next connect, so claiming cannot consume it. + expect(store.claimSshRemotePtyLeasesToReap('target-1')).toEqual(['pty-1']) + expect(store.claimSshRemotePtyLeasesToReap('target-1')).toEqual(['pty-1']) + + store.clearSshRemotePtyLeaseReapFlag('target-1', 'pty-1') + expect(store.claimSshRemotePtyLeasesToReap('target-1')).toEqual([]) + expect(store.getSshRemotePtyLeases('target-1')[0]?.state).toBe('terminated') + }) + + it('drops reap flags older than the bounded age instead of retrying them forever', async () => { + const store = await createStore() + store.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId: 'pty-old', state: 'attached' }) + store.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId: 'pty-new', state: 'attached' }) + store.retireLeaseAndReap('target-1', 'pty-old', 'pane closed') + store.retireLeaseAndReap('target-1', 'pty-new', 'pane closed') + const stale = store.getSshRemotePtyLeases('target-1').find((l) => l.ptyId === 'pty-old')! + stale.updatedAt = Date.now() - SSH_LEASE_REAP_MAX_AGE_MS - 1 + + expect(store.claimSshRemotePtyLeasesToReap('target-1')).toEqual(['pty-new']) + expect(stale.pendingRemoteShutdown).toBeUndefined() + expect(stale.state).toBe('terminated') + }) + + it('caps one connect drain so a poisoned profile cannot stall the reconnect', async () => { + const store = await createStore() + const total = SSH_LEASE_REAP_BATCH_LIMIT + 5 + for (let index = 0; index < total; index++) { + const ptyId = `pty-${String(index).padStart(3, '0')}` + store.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId, state: 'attached' }) + store.retireLeaseAndReap('target-1', ptyId, 'pane closed') + } + const sameInstant = Date.now() + for (const lease of store.getSshRemotePtyLeases('target-1')) { + lease.updatedAt = sameInstant + } + + const claimed = store.claimSshRemotePtyLeasesToReap('target-1') + expect(claimed).toHaveLength(SSH_LEASE_REAP_BATCH_LIMIT) + expect(claimed[0]).toBe('pty-000') + }) + it('keeps worktree deletion authoritative against stale writes and later same-path reuse', async () => { const store = await createStore() store.setWorktreeMeta('wt1', { displayName: 'Worktree' }) diff --git a/src/main/persistence.ts b/src/main/persistence.ts index 0a9d08c8c23..60540090d0d 100644 --- a/src/main/persistence.ts +++ b/src/main/persistence.ts @@ -104,6 +104,7 @@ import { type RemovedSshTargetTombstone, type SshPtyConsumerRecovery, type SshRemotePtyLease, + type SshRemotePtyLeaseState, type SshTarget } from '../shared/ssh-types' import { isFolderRepo } from '../shared/repo-kind' @@ -1615,6 +1616,95 @@ function normalizeSshRemotePtyLease(value: unknown): SshRemotePtyLease | null { } } +/** D3 — a reap flag no relay ever accepts must not be retried until the profile is deleted. */ +export const SSH_LEASE_REAP_MAX_AGE_MS = 7 * 24 * 60 * 60 * 1000 +/** D3 — a profile poisoned by a past fan-out must not stall every reconnect draining it. */ +export const SSH_LEASE_REAP_BATCH_LIMIT = 32 + +// Why: retirement is one-way. Ordering states by finality lets a single predicate replace the ad-hoc +// current-state guards that some lease callers remembered and others forgot. +const SSH_LEASE_FINALITY: Record = { + attached: 0, + detached: 0, + expired: 1, + terminated: 2 +} + +/** + * Monotonic by finality, not absorbing: `expired` → `terminated` must stay legal, or a lease the relay + * expired and the user then closed is stranded `expired` and visible to the 30 s recovery consumers. + */ +function canAdvanceSshLeaseState( + from: SshRemotePtyLeaseState, + to: SshRemotePtyLeaseState +): boolean { + return SSH_LEASE_FINALITY[to] >= SSH_LEASE_FINALITY[from] +} + +function isRetiredSshLeaseState(state: SshRemotePtyLeaseState): boolean { + return state === 'terminated' || state === 'expired' +} + +/** Prior state of a lease D1a superseded, kept only long enough for the spawn-failure rollback (F8). */ +type SupersededSshLeaseSnapshot = { + ptyId: string + state: SshRemotePtyLeaseState + updatedAt: number + retiredReason?: string + pendingRemoteShutdown?: boolean +} + +/** Legacy rows without a complete pane identity cannot be grouped, so they are exempt from supersede. */ +function sshLeasePaneGroupKey(lease: SshRemotePtyLease): string | null { + if (!lease.worktreeId || !lease.tabId || !lease.leafId) { + return null + } + return [lease.targetId, lease.worktreeId, lease.tabId, lease.leafId].join('\u0000') +} + +/** Newest wins; `createdAt` then `ptyId` break ties, since a fan-out lands rows inside one second. */ +function isNewerSshLease(candidate: SshRemotePtyLease, incumbent: SshRemotePtyLease): boolean { + if (candidate.updatedAt !== incumbent.updatedAt) { + return candidate.updatedAt > incumbent.updatedAt + } + if (candidate.createdAt !== incumbent.createdAt) { + return candidate.createdAt > incumbent.createdAt + } + return candidate.ptyId > incumbent.ptyId +} + +/** + * M1 — collapse each pane to one live lease at load, so profiles already poisoned by past fan-outs heal + * without waiting for a reconnect. Idempotent and unconditional: a no-op on clean profiles. + */ +function reconcileSshLeasesToOnePerPane(leases: SshRemotePtyLease[]): boolean { + const winners = new Map() + for (const lease of leases) { + const key = isRetiredSshLeaseState(lease.state) ? null : sshLeasePaneGroupKey(lease) + if (!key) { + continue + } + const incumbent = winners.get(key) + if (!incumbent || isNewerSshLease(lease, incumbent)) { + winners.set(key, lease) + } + } + const now = Date.now() + let changed = false + for (const lease of leases) { + const key = isRetiredSshLeaseState(lease.state) ? null : sshLeasePaneGroupKey(lease) + if (!key || winners.get(key) === lease) { + continue + } + lease.state = 'terminated' + lease.retiredReason = 'superseded on load: another lease holds this pane' + lease.pendingRemoteShutdown = true + lease.updatedAt = now + changed = true + } + return changed +} + const SSH_PTY_OWNER_LEASE_MAX_LENGTH = 512 const ENCRYPTED_SSH_PTY_OWNER_LEASE_MAX_LENGTH = 4096 @@ -2800,6 +2890,8 @@ export class Store { private readonly staleGithubCacheTempCleanup: Promise private gitUsernameCache = new Map() private loadNeedsSave = false + // In-memory only: if the app dies before the rollback runs, M1 re-collapses the pane on next load. + private supersededByLeaseInsert = new Map>() private settingsChangeListeners = new Set< ( updates: Partial, @@ -3635,9 +3727,16 @@ export class Store { (alias): alias is string => typeof alias === 'string' ) : [], - sshRemotePtyLeases: (parsed.sshRemotePtyLeases ?? []) - .map(normalizeSshRemotePtyLease) - .filter((lease): lease is SshRemotePtyLease => lease !== null), + sshRemotePtyLeases: (() => { + const leases = (parsed.sshRemotePtyLeases ?? []) + .map(normalizeSshRemotePtyLease) + .filter((lease): lease is SshRemotePtyLease => lease !== null) + // Why: nothing else marks dirty, so a healed profile would otherwise wait for an unrelated write. + if (reconcileSshLeasesToOnePerPane(leases)) { + this.loadNeedsSave = true + } + return leases + })(), sshPtyConsumerRecoveries: parsed.sshPtyConsumerRecoveries, claudeLivePtySessionIds: normalizeClaudeLivePtySessionIds(parsed.claudeLivePtySessionIds), migrationUnsupportedPtyEntries: normalizeMigrationUnsupportedPtyEntries( @@ -7089,9 +7188,84 @@ export class Store { } else { this.state.sshRemotePtyLeases.push(next) } + // Why: D1a — the incoming write is the current binding for that pane, so any other live lease for it + // is stale by construction. Transactional with the insert so the spawn-failure rollback can undo both. + const superseded = this.supersedeOtherLeasesForPane(next) + const byPtyId = this.supersededByLeaseInsert.get(next.targetId) + if (superseded.length === 0) { + byPtyId?.delete(next.ptyId) + } else if (byPtyId) { + byPtyId.set(next.ptyId, superseded) + } else { + this.supersededByLeaseInsert.set(next.targetId, new Map([[next.ptyId, superseded]])) + } this.flush() } + private supersedeOtherLeasesForPane(inserted: SshRemotePtyLease): SupersededSshLeaseSnapshot[] { + const key = isRetiredSshLeaseState(inserted.state) ? null : sshLeasePaneGroupKey(inserted) + if (!key) { + return [] + } + const now = Date.now() + const snapshots: SupersededSshLeaseSnapshot[] = [] + for (const lease of this.state.sshRemotePtyLeases ?? []) { + if (lease === inserted || isRetiredSshLeaseState(lease.state)) { + continue + } + if (sshLeasePaneGroupKey(lease) !== key) { + continue + } + snapshots.push({ + ptyId: lease.ptyId, + state: lease.state, + updatedAt: lease.updatedAt, + retiredReason: lease.retiredReason, + pendingRemoteShutdown: lease.pendingRemoteShutdown + }) + lease.state = 'terminated' + lease.retiredReason = `superseded by ${inserted.ptyId}` + lease.pendingRemoteShutdown = true + lease.updatedAt = now + } + return snapshots + } + + /** + * F8 — undo a supersede when the spawn that caused it fails to persist. Assigns directly rather than + * going through the public setter, which D1b would correctly refuse: this restores, it does not retire. + */ + private restoreLeasesSupersededByInsert(targetId: string, relayPtyId: string): boolean { + const snapshots = this.supersededByLeaseInsert.get(targetId)?.get(relayPtyId) + if (!snapshots) { + return false + } + this.supersededByLeaseInsert.get(targetId)?.delete(relayPtyId) + let restored = false + for (const snapshot of snapshots) { + const lease = this.state.sshRemotePtyLeases?.find( + (entry) => entry.targetId === targetId && entry.ptyId === snapshot.ptyId + ) + if (!lease) { + continue + } + lease.state = snapshot.state + lease.updatedAt = snapshot.updatedAt + if (snapshot.retiredReason === undefined) { + delete lease.retiredReason + } else { + lease.retiredReason = snapshot.retiredReason + } + if (snapshot.pendingRemoteShutdown === undefined) { + delete lease.pendingRemoteShutdown + } else { + lease.pendingRemoteShutdown = snapshot.pendingRemoteShutdown + } + restored = true + } + return restored + } + markSshRemotePtyLeases(targetId: string, state: SshRemotePtyLease['state']): void { if (this.updateSshRemotePtyLeaseStates(targetId, state)) { this.flush() @@ -7133,10 +7307,8 @@ export class Store { if (lease.targetId !== targetId || (ptyIds && !ptyIds.has(lease.ptyId))) { continue } - if (state === 'attached' && (lease.state === 'terminated' || lease.state === 'expired')) { - continue - } - if (state === 'detached' && lease.state !== 'attached') { + // D1b: subsumes the two directional special cases this guard replaced. + if (!canAdvanceSshLeaseState(lease.state, state)) { continue } if (lease.state !== state) { @@ -7168,7 +7340,9 @@ export class Store { return } const shouldClearBindings = state === 'terminated' || state === 'expired' - if (lease.state === state) { + // Why: D1b — retirement must be terminal in fact. `abandonPtySourceRecovery` writes `detached` here + // unconditionally, which resurrected leases that had already been retired (reproduced under SIGSTOP). + if (lease.state === state || !canAdvanceSshLeaseState(lease.state, state)) { if (shouldClearBindings && this.clearSshRemotePtyBindingsForLeases(targetId, [lease])) { this.flush() } @@ -7215,24 +7389,68 @@ export class Store { if (!lease) { return } - // Why: retirement is terminal. Re-retiring would let a spared row later acquire a reap flag for a PTY - // we already decided we have no claim on; keep binding cleanup idempotent as the singular setter does. - if (lease.state === 'terminated' || lease.state === 'expired') { + // D1b's rule, so `expired` → `terminated` stays legal and nothing walks back toward live. + if (lease.state === state || !canAdvanceSshLeaseState(lease.state, state)) { if (this.clearSshRemotePtyBindingsForLeases(targetId, [lease])) { this.flush() } return } + // Why: a row already retired by the sparing primitive named a PTY we have no claim on. Advancing it + // to `terminated` must not retroactively authorize killing that id — a reset relay reuses `pty-N`. + const wasSparedEarlier = lease.retiredReason !== undefined && !lease.pendingRemoteShutdown lease.state = state lease.updatedAt = Date.now() lease.retiredReason = reason - if (pendingRemoteShutdown) { + if (pendingRemoteShutdown && !wasSparedEarlier) { lease.pendingRemoteShutdown = true } this.clearSshRemotePtyBindingsForLeases(targetId, [lease]) this.flush() } + /** + * D3 — the reap queue for one target, oldest first, in relay id space. Over-age rows are dropped here + * rather than returned: a flag no relay ever accepts must not be retried until the profile is deleted. + * The flag survives a claim, so a transient shutdown failure is retried on the next connect. + */ + claimSshRemotePtyLeasesToReap(targetId: string): string[] { + const now = Date.now() + const claimed: SshRemotePtyLease[] = [] + let expired = false + for (const lease of this.state.sshRemotePtyLeases ?? []) { + if (lease.targetId !== targetId || !lease.pendingRemoteShutdown) { + continue + } + if (now - lease.updatedAt > SSH_LEASE_REAP_MAX_AGE_MS) { + delete lease.pendingRemoteShutdown + expired = true + continue + } + claimed.push(lease) + } + if (expired) { + this.flush() + } + return claimed + .sort((a, b) => a.updatedAt - b.updatedAt || (a.ptyId < b.ptyId ? -1 : 1)) + .slice(0, SSH_LEASE_REAP_BATCH_LIMIT) + .map((lease) => lease.ptyId) + } + + /** D3 — the relay accepted the shutdown (or reports the PTY already gone), so the flag has done its job. */ + clearSshRemotePtyLeaseReapFlag(targetId: string, ptyId: string): void { + const relayPtyId = this.getRelayPtyIdForSshLeaseStorage(targetId, ptyId) + const lease = this.state.sshRemotePtyLeases?.find( + (entry) => entry.targetId === targetId && entry.ptyId === relayPtyId + ) + if (!lease?.pendingRemoteShutdown) { + return + } + delete lease.pendingRemoteShutdown + this.flush() + } + removeSshRemotePtyLease(targetId: string, ptyId: string): void { const relayPtyId = this.getRelayPtyIdForSshLeaseStorage(targetId, ptyId) const leases = (this.state.sshRemotePtyLeases ?? []).filter( @@ -7243,13 +7461,17 @@ export class Store { this.state.sshRemotePtyLeases = (this.state.sshRemotePtyLeases ?? []).filter( (lease) => lease.targetId !== targetId || lease.ptyId !== relayPtyId ) - if (this.state.sshRemotePtyLeases.length !== before) { + // Why: this method's only production caller is the spawn-persist rollback, so dropping the new lease + // must also undo the supersede it caused — otherwise a failed spawn retires a working pane's lease. + const restored = this.restoreLeasesSupersededByInsert(targetId, relayPtyId) + if (this.state.sshRemotePtyLeases.length !== before || restored) { this.flush() } } removeSshRemotePtyLeases(targetId: string): void { this.state.sshRemotePtyLeases ??= [] + this.supersededByLeaseInsert.delete(targetId) this.clearSshRemotePtyBindingsForTarget(targetId) const before = this.state.sshRemotePtyLeases.length this.state.sshRemotePtyLeases = this.state.sshRemotePtyLeases.filter( diff --git a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts index 1a384a32f72..629b59800f6 100644 --- a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts +++ b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts @@ -158,7 +158,9 @@ function createSession(targetId: string): InstanceType { markSshRemotePtyLease: vi.fn(), markSshRemotePtyLeases: vi.fn(), markSshRemotePtyLeasesAsync: vi.fn(), - markSshRemotePtyLeasesAttachedAsync: vi.fn() + markSshRemotePtyLeasesAttachedAsync: vi.fn(), + claimSshRemotePtyLeasesToReap: vi.fn().mockReturnValue([]), + clearSshRemotePtyLeaseReapFlag: vi.fn() } as unknown as Store const portForwardManager = { removeAllForwards: vi.fn().mockResolvedValue(undefined) diff --git a/src/main/ssh/ssh-relay-session-orphan-reap.test.ts b/src/main/ssh/ssh-relay-session-orphan-reap.test.ts new file mode 100644 index 00000000000..65934798825 --- /dev/null +++ b/src/main/ssh/ssh-relay-session-orphan-reap.test.ts @@ -0,0 +1,170 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { SshRelaySession } from './ssh-relay-session' +import { createMockDeps, mockDeploySuccess } from './ssh-relay-session-test-fixtures' + +const { acceptOutputExitMock, muxRequestMock } = vi.hoisted(() => ({ + acceptOutputExitMock: vi.fn().mockResolvedValue(undefined), + muxRequestMock: vi.fn() +})) + +vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: vi.fn(async (_mux, options) => ({ + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) +})) +vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ + acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined), + acceptSshPtyOutputExit: acceptOutputExitMock, + allocateSshPtyProviderGeneration: vi.fn(() => 31), + beginSshPtyOutputGenerationMigration: vi.fn(() => ({ + byPty: new Map(), + completion: Promise.resolve() + })), + closeSshPtyOutputGeneration: vi.fn(), + getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []), + applySshPtySourceCancellationProof: vi.fn(() => true), + applySshPtySourceRecoveryCancellationProof: vi.fn(() => true), + installSshPtySourceAckPublisher: vi.fn(() => () => {}), + installSshPtySourceCancellationPublisher: vi.fn(() => () => {}) +})) +vi.mock('./ssh-relay-deploy-helpers', () => ({ execCommand: vi.fn().mockResolvedValue('') })) +vi.mock('./ssh-channel-multiplexer', () => ({ + SshChannelMultiplexer: class MockSshChannelMultiplexer { + notify = vi.fn() + notifyWithSettlement = vi.fn() + request = muxRequestMock + onNotification = vi.fn().mockReturnValue(() => {}) + onNotificationByMethod = vi.fn().mockReturnValue(() => {}) + onRequest = vi.fn().mockReturnValue(() => {}) + onDispose = vi.fn().mockReturnValue(() => {}) + dispose = vi.fn() + isDisposed = vi.fn().mockReturnValue(false) + } +})) +vi.mock('../agent-hooks/remote-managed-hook-installers', () => ({ + installRemoteManagedAgentHooks: vi.fn().mockResolvedValue([]) +})) +vi.mock('../providers/ssh-pty-provider', () => ({ + SshPtyProvider: class MockSshPtyProvider { + onData = vi.fn().mockReturnValue(() => {}) + onReplay = vi.fn().mockReturnValue(() => {}) + onExit = vi.fn().mockReturnValue(() => {}) + attach = vi.fn().mockResolvedValue(undefined) + attachForReconnect = vi.fn().mockResolvedValue({}) + dispose = vi.fn() + } +})) +vi.mock('../providers/ssh-filesystem-provider', () => ({ + SshFilesystemProvider: class MockSshFilesystemProvider { + dispose = vi.fn() + } +})) +vi.mock('../providers/ssh-git-provider', () => ({ + SshGitProvider: class MockSshGitProvider {} +})) +vi.mock('../ipc/pty', () => ({ + registerSshPtyProvider: vi.fn(), + unregisterSshPtyProvider: vi.fn(), + getSshPtyProvider: vi.fn().mockReturnValue({ dispose: vi.fn() }), + getPtyIdsForConnection: vi.fn().mockReturnValue([]), + clearPtyOwnershipForConnection: vi.fn(), + clearProviderPtyState: vi.fn(), + deletePtyOwnership: vi.fn(), + setPtyOwnership: vi.fn(), + restorePtyIncarnation: vi.fn(), + isCurrentPtyExit: vi.fn(() => true), + answerStartupTerminalColorQueriesForPty: vi.fn((_id: string, data: string) => data) +})) +vi.mock('../providers/ssh-filesystem-dispatch', () => ({ + registerSshFilesystemProvider: vi.fn(), + unregisterSshFilesystemProvider: vi.fn(), + getSshFilesystemProvider: vi.fn().mockReturnValue({ dispose: vi.fn() }) +})) +vi.mock('../providers/ssh-git-dispatch', () => ({ + registerSshGitProvider: vi.fn(), + unregisterSshGitProvider: vi.fn() +})) + +const { getPtyIdsForConnection } = await import('../ipc/pty') + +describe('SSH relay orphaned PTY reaping', () => { + beforeEach(() => { + vi.clearAllMocks() + muxRequestMock.mockResolvedValue([]) + mockDeploySuccess() + }) + + it('reaps orphaned remote PTYs on connect and clears the flag once the relay accepts', async () => { + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + const { getSshPtyProvider } = await import('../ipc/pty') + const mockShutdown = vi + .fn() + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error('relay unreachable')) + vi.mocked(getSshPtyProvider).mockReturnValue({ + attachForReconnect: vi.fn().mockResolvedValue(undefined), + shutdown: mockShutdown, + hasPty: vi.fn().mockReturnValue(false), + dispose: vi.fn() + } as unknown as ReturnType) + vi.mocked(getPtyIdsForConnection).mockReturnValue([]) + vi.mocked(mockStore.claimSshRemotePtyLeasesToReap).mockReturnValue(['pty-dead', 'pty-stuck']) + + const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + + expect(mockShutdown).toHaveBeenCalledWith('ssh:target-1@@pty-dead', { + immediate: true, + keepHistory: false + }) + expect(mockStore.clearSshRemotePtyLeaseReapFlag).toHaveBeenCalledWith('target-1', 'pty-dead') + // Why: a failed kill keeps its flag, so the next connect retries it. + expect(mockStore.clearSshRemotePtyLeaseReapFlag).not.toHaveBeenCalledWith( + 'target-1', + 'pty-stuck' + ) + }) + + it('spares a flagged id this connection is live on, since a reset relay reuses pty-N', async () => { + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + const { getSshPtyProvider } = await import('../ipc/pty') + const mockShutdown = vi.fn().mockResolvedValue(undefined) + vi.mocked(getSshPtyProvider).mockReturnValue({ + attachForReconnect: vi.fn().mockResolvedValue(undefined), + shutdown: mockShutdown, + hasPty: vi.fn().mockImplementation((id: string) => id === 'ssh:target-1@@pty-1'), + dispose: vi.fn() + } as unknown as ReturnType) + vi.mocked(getPtyIdsForConnection).mockReturnValue([]) + vi.mocked(mockStore.claimSshRemotePtyLeasesToReap).mockReturnValue(['pty-1']) + + const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + + expect(mockShutdown).not.toHaveBeenCalled() + expect(mockStore.clearSshRemotePtyLeaseReapFlag).toHaveBeenCalledWith('target-1', 'pty-1') + }) + + it('treats an already-gone remote PTY as a successful reap', async () => { + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + const { getSshPtyProvider } = await import('../ipc/pty') + const notFound = new Error('PTY "pty-gone" not found') + vi.mocked(getSshPtyProvider).mockReturnValue({ + attachForReconnect: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn().mockRejectedValue(notFound), + hasPty: vi.fn().mockReturnValue(false), + dispose: vi.fn() + } as unknown as ReturnType) + vi.mocked(getPtyIdsForConnection).mockReturnValue([]) + vi.mocked(mockStore.claimSshRemotePtyLeasesToReap).mockReturnValue(['pty-gone']) + + const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + + expect(mockStore.clearSshRemotePtyLeaseReapFlag).toHaveBeenCalledWith('target-1', 'pty-gone') + }) +}) diff --git a/src/main/ssh/ssh-relay-session-terminal-error.test.ts b/src/main/ssh/ssh-relay-session-terminal-error.test.ts index a0f855543ef..175393f9931 100644 --- a/src/main/ssh/ssh-relay-session-terminal-error.test.ts +++ b/src/main/ssh/ssh-relay-session-terminal-error.test.ts @@ -107,7 +107,9 @@ function createMockDeps(): { markSshRemotePtyLease: vi.fn(), markSshRemotePtyLeases: vi.fn(), markSshRemotePtyLeasesAsync: vi.fn(), - markSshRemotePtyLeasesAttachedAsync: vi.fn() + markSshRemotePtyLeasesAttachedAsync: vi.fn(), + claimSshRemotePtyLeasesToReap: vi.fn().mockReturnValue([]), + clearSshRemotePtyLeaseReapFlag: vi.fn() } as unknown as Store const mockPortForward = { removeAllForwards: vi.fn() diff --git a/src/main/ssh/ssh-relay-session-test-fixtures.ts b/src/main/ssh/ssh-relay-session-test-fixtures.ts index b0ed98eaf2e..1533c96368b 100644 --- a/src/main/ssh/ssh-relay-session-test-fixtures.ts +++ b/src/main/ssh/ssh-relay-session-test-fixtures.ts @@ -26,6 +26,8 @@ export function createMockDeps(): SshRelaySessionTestDeps { markSshRemotePtyLeases: vi.fn(), markSshRemotePtyLeasesAsync: vi.fn(), markSshRemotePtyLeasesAttachedAsync: vi.fn(), + claimSshRemotePtyLeasesToReap: vi.fn().mockReturnValue([]), + clearSshRemotePtyLeaseReapFlag: vi.fn(), persistPtyBinding: vi.fn() } as unknown as Store const mockPortForward = { diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index f5a7302237d..0f0605b4a8f 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1999,6 +1999,51 @@ export class SshRelaySession { Array.from(attachedLeaseIds) ) } + if (shouldContinue()) { + await this.reapOrphanedRemotePtys(ptyProvider) + } + } + + /** + * D3 — kill the relay PTYs whose panes I3 proved gone. Deferred to connect time because that is the only + * moment a relay exists to accept the kill; best-effort, since a failure just leaves the flag for next time. + */ + private async reapOrphanedRemotePtys(ptyProvider: SshPtyProvider): Promise { + const claimed = this.store.claimSshRemotePtyLeasesToReap(this.targetId) + const relayPtyIds: string[] = [] + for (const relayPtyId of claimed) { + // Why: a reset relay restarts `pty-N` numbering, so a flag minted in an earlier generation can name + // an id this connection just bound. Anything live here belongs to a pane; drop the stale flag instead. + if (ptyProvider.hasPty(toAppSshPtyId(this.targetId, relayPtyId))) { + this.store.clearSshRemotePtyLeaseReapFlag(this.targetId, relayPtyId) + continue + } + relayPtyIds.push(relayPtyId) + } + if (relayPtyIds.length === 0) { + return + } + const results = await Promise.allSettled( + relayPtyIds.map((relayPtyId) => + ptyProvider.shutdown(toAppSshPtyId(this.targetId, relayPtyId), { + immediate: true, + keepHistory: false + }) + ) + ) + for (const [index, result] of results.entries()) { + const relayPtyId = relayPtyIds[index] + // Already gone is the outcome we wanted, so it retires the flag just like a successful kill. + if (result.status === 'fulfilled' || isSshPtyNotFoundError(result.reason)) { + this.store.clearSshRemotePtyLeaseReapFlag(this.targetId, relayPtyId) + continue + } + console.warn( + `[ssh-relay-session] orphaned PTY ${relayPtyId} reap failed for ${this.targetId}: ${ + result.reason instanceof Error ? result.reason.message : String(result.reason) + }` + ) + } } private async reattachKnownPty(args: {