diff --git a/src/main/providers/ssh-pty-liveness-state.ts b/src/main/providers/ssh-pty-liveness-state.ts index afd9c9199ae..8f4786cabae 100644 --- a/src/main/providers/ssh-pty-liveness-state.ts +++ b/src/main/providers/ssh-pty-liveness-state.ts @@ -2,21 +2,12 @@ export const MAX_SSH_PTY_EXIT_TOMBSTONES = 1000 export class SshPtyLivenessState { readonly livePtyIds = new Set() - private readonly unverifiablePtyIds = new Set() private readonly exitedPtyIds = new Set() - constructor( - private readonly toAppPtyId: (id: string) => string, - initialUnverifiablePtyIds: Iterable - ) { - for (const id of initialUnverifiablePtyIds) { - this.acceptUnverifiable(id) - } - } + constructor(private readonly toAppPtyId: (id: string) => string) {} clear(): void { this.livePtyIds.clear() - this.unverifiablePtyIds.clear() this.exitedPtyIds.clear() } @@ -26,7 +17,6 @@ export class SshPtyLivenessState { acceptLive(id: string): void { const appPtyId = this.toAppPtyId(id) - this.unverifiablePtyIds.delete(appPtyId) this.exitedPtyIds.delete(appPtyId) this.livePtyIds.add(appPtyId) } @@ -35,13 +25,11 @@ export class SshPtyLivenessState { const appPtyId = this.toAppPtyId(id) this.livePtyIds.delete(appPtyId) this.exitedPtyIds.delete(appPtyId) - this.unverifiablePtyIds.add(appPtyId) } acceptExited(id: string): void { const appPtyId = this.toAppPtyId(id) this.livePtyIds.delete(appPtyId) - this.unverifiablePtyIds.delete(appPtyId) this.exitedPtyIds.delete(appPtyId) this.exitedPtyIds.add(appPtyId) if (this.exitedPtyIds.size > MAX_SSH_PTY_EXIT_TOMBSTONES) { diff --git a/src/main/providers/ssh-pty-provider.ts b/src/main/providers/ssh-pty-provider.ts index 5aeed6c18ce..f0316129300 100644 --- a/src/main/providers/ssh-pty-provider.ts +++ b/src/main/providers/ssh-pty-provider.ts @@ -47,14 +47,13 @@ export class SshPtyProvider implements IPtyProvider { connectionId: string, mux: SshChannelMultiplexer, private readonly remoteCliBridgeEnv?: RemoteCliBridgeEnv, - readonly providerGeneration = 1, - initialUnverifiablePtyIds: Iterable = [] + readonly providerGeneration = 1 ) { this.connectionId = connectionId this.mux = mux this.agentSessionCapabilities = new SshAgentSessionCapabilities(mux) this.getAppliedSize = createSshPtyAppliedSizeReader(mux, connectionId) - this.livenessState = new SshPtyLivenessState(this.toAppPtyId, initialUnverifiablePtyIds) + this.livenessState = new SshPtyLivenessState(this.toAppPtyId) this.outputState = new SshPtyProviderOutputState(providerGeneration, { mux, 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 bc6755961cf..c0349b090a6 100644 --- a/src/main/providers/ssh-pty-restore-required-reattach.test.ts +++ b/src/main/providers/ssh-pty-restore-required-reattach.test.ts @@ -79,7 +79,7 @@ describe('SSH PTY reattach when the relay requires source restoration', () => { const id = 'ssh:conn-1@@pty-reconnecting' const mux = createMockMux() mux.request.mockResolvedValueOnce([]) - const provider = new SshPtyProvider('conn-1', mux as never, undefined, 1, [id]) + const provider = new SshPtyProvider('conn-1', mux as never) await provider.listProcesses() diff --git a/src/main/ssh/ssh-relay-session-incarnation.test.ts b/src/main/ssh/ssh-relay-session-incarnation.test.ts index d983ecaec21..b3a52b0798d 100644 --- a/src/main/ssh/ssh-relay-session-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-incarnation.test.ts @@ -56,6 +56,7 @@ vi.mock('../providers/ssh-pty-provider', () => ({ onExit = vi.fn().mockReturnValue(() => {}) attach = vi.fn().mockResolvedValue(undefined) attachForReconnect = vi.fn().mockResolvedValue({}) + acceptUnverifiablePty = vi.fn() dispose = vi.fn() } })) @@ -120,6 +121,7 @@ describe('SSH relay PTY incarnation exits', () => { await session.establish(mockConn) const provider = vi.mocked(registerSshPtyProvider).mock.calls[0]?.[1] as unknown as { onExit: ReturnType + acceptUnverifiablePty: ReturnType } const onExit = provider.onExit.mock.calls[0]?.[0] as (payload: { id: string @@ -209,4 +211,33 @@ describe('SSH relay PTY incarnation exits', () => { expect(acceptLivePty).toHaveBeenCalledExactlyOnceWith('ssh:target-1@@pty-reused') ) }) + + it('marks a legacy exit unverifiable after current incarnation ownership is known', async () => { + 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 { + onExit: ReturnType + acceptUnverifiablePty: ReturnType + } + const onExit = provider.onExit.mock.calls[0]?.[0] as (payload: { + id: string + code: number + providerGeneration: number + ptyIncarnation: string + }) => void + vi.mocked(isCurrentPtyExit).mockReturnValueOnce(false) + + onExit({ + id: 'ssh:target-1@@pty-current', + code: 0, + providerGeneration: 31, + ptyIncarnation: 'incarnation-current' + }) + + expect(provider.acceptUnverifiablePty).toHaveBeenCalledExactlyOnceWith( + 'ssh:target-1@@pty-current' + ) + expect(acceptOutputExitMock).not.toHaveBeenCalled() + }) }) diff --git a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts index b4733471c4f..ead85a31414 100644 --- a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts @@ -84,7 +84,6 @@ vi.mock('../providers/ssh-pty-provider', () => ({ isSshPtyNotFoundError: (error: unknown) => String(error).includes('not found'), isSshPtyIdentityMismatchError: (error: unknown) => String(error).includes('identity mismatch'), SshPtyProvider: class MockSshPtyProvider { - readonly initialUnverifiablePtyIds: Set onData = vi.fn().mockReturnValue(() => {}) onReplay = vi.fn().mockReturnValue(() => {}) onExit = vi.fn().mockReturnValue(() => {}) @@ -93,16 +92,6 @@ vi.mock('../providers/ssh-pty-provider', () => ({ acceptUnverifiablePty = vi.fn() acceptExitedPty = vi.fn() dispose = vi.fn() - - constructor( - _connectionId: string, - _mux: unknown, - _remoteCliBridgeEnv: unknown, - _providerGeneration: number, - initialUnverifiablePtyIds: Iterable = [] - ) { - this.initialUnverifiablePtyIds = new Set(initialUnverifiablePtyIds) - } } })) vi.mock('../providers/ssh-filesystem-provider', () => ({ @@ -210,14 +199,12 @@ describe('SshRelaySession reconnect incarnation ordering', () => { try { await session.establish(mockConn) const provider = vi.mocked(registerSshPtyProvider).mock.calls.at(-1)?.[1] as unknown as { - initialUnverifiablePtyIds: Set acceptUnverifiablePty: ReturnType acceptExitedPty: ReturnType } const appPtyId = 'ssh:target-1@@pty-unverifiable' expect(session.getState()).toBe('ready') - expect(provider.initialUnverifiablePtyIds).toContain(appPtyId) expect(provider.acceptUnverifiablePty).toHaveBeenCalledWith(appPtyId) expect(provider.acceptExitedPty).not.toHaveBeenCalled() expect(mockStore.markSshRemotePtyLease).not.toHaveBeenCalledWith( diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index d1bb7aae3b0..d03d7c040e9 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1009,21 +1009,11 @@ export class SshRelaySession { this.wireUpRemoteOrcaCli(mux, connectionIncarnation) const providerGeneration = allocateSshPtyProviderGeneration() - const initialUnverifiablePtyIds = new Set( - [ - ...getPtyIdsForConnection(this.targetId), - ...this.store - .getSshRemotePtyLeases(this.targetId) - .filter((lease) => lease.state !== 'terminated' && lease.state !== 'expired') - .map((lease) => lease.ptyId) - ].map((ptyId) => toAppSshPtyId(this.targetId, ptyId)) - ) const ptyProvider = new SshPtyProvider( this.targetId, mux, this.remoteCliBridgeEnv ?? undefined, - providerGeneration, - initialUnverifiablePtyIds + providerGeneration ) const consumerOwnerState = this.activePtyConsumerOwner() if (consumerOwnerState) { @@ -1788,6 +1778,9 @@ export class SshRelaySession { return } if (!isCurrentPtyExit(payload)) { + if (!payload.incarnationId) { + ptyProvider.acceptUnverifiablePty(payload.id) + } return } void this.acceptPtyExit(payload).catch(() => {})