From cc84b69865219facfba51546aca177916c000ff4 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:19:04 -0700 Subject: [PATCH] fix(ssh): apply exact exit evidence immediately --- src/main/providers/ssh-pty-liveness-state.ts | 20 +++++++- .../ssh-pty-provider-output-state.ts | 2 +- .../ssh-pty-provider-process-events.test.ts | 2 +- src/main/providers/ssh-pty-provider.ts | 30 ++++++------ .../ssh-pty-restore-required-reattach.test.ts | 20 +++++++- src/main/providers/ssh-pty-shutdown.ts | 19 ++++++++ .../ssh/ssh-relay-session-incarnation.test.ts | 47 ++++++++++++++++++- src/main/ssh/ssh-relay-session.ts | 6 ++- 8 files changed, 124 insertions(+), 22 deletions(-) create mode 100644 src/main/providers/ssh-pty-shutdown.ts diff --git a/src/main/providers/ssh-pty-liveness-state.ts b/src/main/providers/ssh-pty-liveness-state.ts index eaf15389a97..6d3066321e1 100644 --- a/src/main/providers/ssh-pty-liveness-state.ts +++ b/src/main/providers/ssh-pty-liveness-state.ts @@ -1,4 +1,5 @@ export const MAX_SSH_PTY_EXIT_TOMBSTONES = 1000 +export const MAX_SSH_PTY_AMBIGUOUS_EXIT_STATES = 1000 export type SshPtyLiveEvidence = { valid: boolean } export type SshPtyLiveEvidenceWindow = { @@ -9,10 +10,14 @@ export type SshPtyLiveEvidenceWindow = { export class SshPtyLivenessState { readonly livePtyIds = new Set() private readonly exitedPtyIds = new Set() + private readonly ambiguousExitPtyIds = new Set() private readonly pendingLiveEvidenceByPtyId = new Map>() private readonly liveEvidenceWindows = new Set() - constructor(private readonly toAppPtyId: (id: string) => string) {} + constructor( + private readonly toAppPtyId: (id: string) => string, + private readonly onAmbiguousExitCapacityExceeded?: () => void + ) {} clear(): void { for (const evidence of this.pendingLiveEvidenceByPtyId.values()) { @@ -28,6 +33,7 @@ export class SshPtyLivenessState { this.liveEvidenceWindows.clear() this.livePtyIds.clear() this.exitedPtyIds.clear() + this.ambiguousExitPtyIds.clear() } probe(id: string): boolean | null { @@ -36,6 +42,7 @@ export class SshPtyLivenessState { acceptLive(id: string): void { const appPtyId = this.toAppPtyId(id) + this.ambiguousExitPtyIds.delete(appPtyId) this.exitedPtyIds.delete(appPtyId) this.livePtyIds.add(appPtyId) } @@ -92,6 +99,7 @@ export class SshPtyLivenessState { acceptExited(id: string): void { const appPtyId = this.toAppPtyId(id) + this.ambiguousExitPtyIds.delete(appPtyId) this.invalidateLiveEvidence(appPtyId) this.livePtyIds.delete(appPtyId) this.exitedPtyIds.delete(appPtyId) @@ -104,6 +112,16 @@ export class SshPtyLivenessState { } } + acceptUnverifiableExit(id: string): void { + const appPtyId = this.toAppPtyId(id) + this.acceptUnverifiable(appPtyId) + this.ambiguousExitPtyIds.add(appPtyId) + if (this.ambiguousExitPtyIds.size > MAX_SSH_PTY_AMBIGUOUS_EXIT_STATES) { + this.ambiguousExitPtyIds.clear() + this.onAmbiguousExitCapacityExceeded?.() + } + } + private invalidateLiveEvidence(appPtyId: string): void { for (const window of this.liveEvidenceWindows) { window.invalidatedPtyIds.add(appPtyId) diff --git a/src/main/providers/ssh-pty-provider-output-state.ts b/src/main/providers/ssh-pty-provider-output-state.ts index bdb5d7813c7..a62057c8643 100644 --- a/src/main/providers/ssh-pty-provider-output-state.ts +++ b/src/main/providers/ssh-pty-provider-output-state.ts @@ -153,7 +153,7 @@ export class SshPtyProviderOutputState { private resolvePtyExitIncarnation(relayPtyId: string, incarnationId: unknown): string { if (typeof incarnationId === 'string' && incarnationId.length > 0) { - return this.resolvePtyIncarnation(relayPtyId, incarnationId) + return incarnationId } return ( this.incarnationByRelayPtyId.get(relayPtyId) ?? diff --git a/src/main/providers/ssh-pty-provider-process-events.test.ts b/src/main/providers/ssh-pty-provider-process-events.test.ts index 2b2f07f4689..8934c7039e3 100644 --- a/src/main/providers/ssh-pty-provider-process-events.test.ts +++ b/src/main/providers/ssh-pty-provider-process-events.test.ts @@ -158,7 +158,7 @@ describe('SshPtyProvider process listings and events', () => { id: scopedPty1, code: 0, providerGeneration: 1, - ptyIncarnation: 'legacy:1:1:pty-1', + ptyIncarnation: 'incarnation-1', incarnationId: 'incarnation-1' }) }) diff --git a/src/main/providers/ssh-pty-provider.ts b/src/main/providers/ssh-pty-provider.ts index e5f0f7564ea..da25cf00016 100644 --- a/src/main/providers/ssh-pty-provider.ts +++ b/src/main/providers/ssh-pty-provider.ts @@ -26,11 +26,7 @@ import { SshPtyLivenessState } from './ssh-pty-liveness-state' import type { SshPtyLiveEvidence } from './ssh-pty-liveness-state' import { listSshPtyProcessesWithLiveEvidence } from './ssh-pty-process-list-liveness' import { attachSshPtyWithLiveEvidence } from './ssh-pty-attach-liveness' - -// Why: sequential relay teardown calls share one absolute budget; convert to the mux-relative timeout only at dispatch. -function relayTimeoutOptions(deadlineMs: number | undefined): { timeoutMs: number } | undefined { - return deadlineMs === undefined ? undefined : { timeoutMs: Math.max(1, deadlineMs - Date.now()) } -} +import { shutdownSshPty } from './ssh-pty-shutdown' /** Remote PTY provider that proxies IPtyProvider operations through the relay. */ export class SshPtyProvider implements IPtyProvider { @@ -55,7 +51,9 @@ export class SshPtyProvider implements IPtyProvider { this.mux = mux this.agentSessionCapabilities = new SshAgentSessionCapabilities(mux) this.getAppliedSize = createSshPtyAppliedSizeReader(mux, connectionId) - this.livenessState = new SshPtyLivenessState(this.toAppPtyId) + this.livenessState = new SshPtyLivenessState(this.toAppPtyId, () => + this.closeOutputIntake('ambiguous-exit-state-cap') + ) this.outputState = new SshPtyProviderOutputState(providerGeneration, { mux, @@ -211,15 +209,7 @@ export class SshPtyProvider implements IPtyProvider { id: string, opts: { immediate?: boolean; keepHistory?: boolean; deadlineMs?: number } ): Promise { - await this.mux.request( - 'pty.shutdown', - { - id: this.toRelayPtyId(id), - immediate: opts.immediate ?? false, - keepHistory: opts.keepHistory ?? false - }, - relayTimeoutOptions(opts.deadlineMs) - ) + await shutdownSshPty(this.mux, this.toRelayPtyId(id), opts) this.acceptExitedPty(id) } @@ -312,9 +302,17 @@ export class SshPtyProvider implements IPtyProvider { this.livenessState.acceptUnverifiable(id) } + acceptAmbiguousExitPty(id: string): void { + this.livenessState.acceptUnverifiableExit(id) + } + + acceptExitedPtyLiveness(id: string): void { + this.livenessState.acceptExited(id) + } + acceptExitedPty(id: string): void { this.outputState.acceptExit(this.toRelayPtyId(id)) - this.livenessState.acceptExited(id) + this.acceptExitedPtyLiveness(id) } async getDefaultShell(): Promise { diff --git a/src/main/providers/ssh-pty-restore-required-reattach.test.ts b/src/main/providers/ssh-pty-restore-required-reattach.test.ts index 2e0a7fbb556..c98ed114d81 100644 --- a/src/main/providers/ssh-pty-restore-required-reattach.test.ts +++ b/src/main/providers/ssh-pty-restore-required-reattach.test.ts @@ -4,7 +4,10 @@ import { SSH_PTY_RESTORE_REQUIRED_ERROR, SSH_SESSION_EXPIRED_ERROR } from './ssh-pty-errors' -import { MAX_SSH_PTY_EXIT_TOMBSTONES } from './ssh-pty-liveness-state' +import { + MAX_SSH_PTY_AMBIGUOUS_EXIT_STATES, + MAX_SSH_PTY_EXIT_TOMBSTONES +} from './ssh-pty-liveness-state' import { SshPtyProvider } from './ssh-pty-provider' // The relay answers `restoreRequired` from its DELIVERY layer only: `requireRestore` @@ -154,6 +157,21 @@ describe('SSH PTY reattach when the relay requires source restoration', () => { ).resolves.toBe(false) }) + it('bounds ambiguous exit retention by resetting the output intake', () => { + const mux = createMockMux() + const provider = new SshPtyProvider('conn-1', mux as never) + const log = vi.spyOn(console, 'error').mockImplementation(() => {}) + try { + for (let index = 0; index <= MAX_SSH_PTY_AMBIGUOUS_EXIT_STATES; index += 1) { + provider.acceptAmbiguousExitPty(`pty-${index}`) + } + } finally { + log.mockRestore() + } + + expect(mux.dispose).toHaveBeenCalledExactlyOnceWith('connection_lost') + }) + it('keeps reconnect liveness unverifiable until relay activation', async () => { const id = 'ssh:conn-1@@pty-reconnecting' const mux = createMockMux() diff --git a/src/main/providers/ssh-pty-shutdown.ts b/src/main/providers/ssh-pty-shutdown.ts new file mode 100644 index 00000000000..98930815db1 --- /dev/null +++ b/src/main/providers/ssh-pty-shutdown.ts @@ -0,0 +1,19 @@ +import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer' + +export async function shutdownSshPty( + mux: SshChannelMultiplexer, + relayPtyId: string, + opts: { immediate?: boolean; keepHistory?: boolean; deadlineMs?: number } +): Promise { + await mux.request( + 'pty.shutdown', + { + id: relayPtyId, + immediate: opts.immediate ?? false, + keepHistory: opts.keepHistory ?? false + }, + opts.deadlineMs === undefined + ? undefined + : { timeoutMs: Math.max(1, opts.deadlineMs - Date.now()) } + ) +} diff --git a/src/main/ssh/ssh-relay-session-incarnation.test.ts b/src/main/ssh/ssh-relay-session-incarnation.test.ts index 0684d2adac9..83102e6e7ee 100644 --- a/src/main/ssh/ssh-relay-session-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-incarnation.test.ts @@ -72,12 +72,15 @@ vi.mock('../providers/ssh-pty-provider', () => ({ } } ) - acceptUnverifiablePty = vi.fn(() => { + acceptUnverifiablePty = vi.fn((_id: string) => { for (const evidence of this.pendingLiveEvidence) { evidence.valid = false } this.pendingLiveEvidence.clear() }) + acceptAmbiguousExitPty = vi.fn((id: string) => this.acceptUnverifiablePty(id)) + acceptExitedPtyLiveness = vi.fn() + acceptExitedPty = vi.fn() dispose = vi.fn() constructor(_connectionId: string, _mux: unknown, _env: unknown, providerGeneration: number) { @@ -156,8 +159,10 @@ describe('SSH relay PTY incarnation exits', () => { ptyIncarnation: string }) => void const acceptExitedPty = vi.fn() + const acceptExitedPtyLiveness = vi.fn() vi.mocked(getSshPtyProvider).mockReturnValue({ providerGeneration: 31, + acceptExitedPtyLiveness, acceptExitedPty } as never) vi.mocked(isCurrentPtyExit).mockReturnValueOnce(true).mockReturnValueOnce(false) @@ -197,6 +202,46 @@ describe('SSH relay PTY incarnation exits', () => { expect(runtime.onPtyExit).not.toHaveBeenCalled() }) + it('records exact exited liveness before output-exit settlement', async () => { + let settleExit!: () => void + acceptOutputExitMock.mockReturnValueOnce( + new Promise((resolve) => { + settleExit = resolve + }) + ) + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + const provider = vi.mocked(registerSshPtyProvider).mock.calls[0]?.[1] as unknown as { + providerGeneration: number + onExit: ReturnType + acceptExitedPtyLiveness: ReturnType + acceptExitedPty: ReturnType + } + vi.mocked(getSshPtyProvider).mockReturnValue(provider as never) + const onExit = provider.onExit.mock.calls[0]?.[0] as (payload: { + id: string + code: number + incarnationId: string + providerGeneration: number + ptyIncarnation: string + }) => void + const id = 'ssh:target-1@@pty-current' + + onExit({ + id, + code: 0, + incarnationId: 'incarnation-current', + providerGeneration: provider.providerGeneration, + ptyIncarnation: 'incarnation-current' + }) + + expect(provider.acceptExitedPtyLiveness).toHaveBeenCalledExactlyOnceWith(id) + expect(provider.acceptExitedPty).not.toHaveBeenCalled() + settleExit() + await vi.waitFor(() => expect(provider.acceptExitedPty).toHaveBeenCalledExactlyOnceWith(id)) + }) + it('promotes liveness only after current output passes intake fencing', async () => { const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward) diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index 126bcdad35c..aaeab5bb336 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1783,7 +1783,7 @@ export class SshRelaySession { return } if (!payload.incarnationId) { - ptyProvider.acceptUnverifiablePty(payload.id) + ptyProvider.acceptAmbiguousExitPty(payload.id) return } if (!isCurrentPtyExit(payload)) { @@ -2235,6 +2235,10 @@ export class SshRelaySession { if (!isCurrentPtyExit(payload)) { return } + const evidenceProvider = getSshPtyProvider(this.targetId) as SshPtyProvider | undefined + if (evidenceProvider?.providerGeneration === payload.providerGeneration) { + evidenceProvider.acceptExitedPtyLiveness(payload.id) + } await acceptSshPtyOutputExit({ id: payload.id, code: payload.code,