diff --git a/src/main/ipc/ssh.test.ts b/src/main/ipc/ssh.test.ts index 09b741076cc..c7f850452c5 100644 --- a/src/main/ipc/ssh.test.ts +++ b/src/main/ipc/ssh.test.ts @@ -259,6 +259,7 @@ import { getPtyIdsForConnection } from './pty' import { assertSshMutationExpectation } from '../ssh/ssh-connection-generation' +import { getSshPtyConsumerRecovery } from '../ssh/ssh-pty-consumer-recovery' describe('SSH IPC handlers', () => { const relayBuildId = '0.1.0+ipc-test' @@ -280,6 +281,8 @@ describe('SSH IPC handlers', () => { getSshRemotePtyLeases: vi.fn().mockReturnValue([]), markSshRemotePtyLease: vi.fn(), markSshRemotePtyLeases: vi.fn(), + markSshRemotePtyLeasesAsync: vi.fn(), + markSshRemotePtyLeasesAttachedAsync: vi.fn(), removeSshRemotePtyLeases: vi.fn() } const mockWindow = { @@ -355,7 +358,10 @@ describe('SSH IPC handlers', () => { mockStore.getSshRemotePtyLeases.mockReset().mockReturnValue([]) mockStore.markSshRemotePtyLease.mockReset() mockStore.markSshRemotePtyLeases.mockReset() + mockStore.markSshRemotePtyLeasesAsync.mockReset() + mockStore.markSshRemotePtyLeasesAttachedAsync.mockReset() mockStore.removeSshRemotePtyLeases.mockReset() + mockStore.upsertSshPtyConsumerRecovery.mockReset() mockConnectionManager.connect.mockReset() mockConnectionManager.disconnect.mockReset() @@ -527,7 +533,7 @@ describe('SSH IPC handlers', () => { expect(mockPortForwardManager.removeAllForwards).toHaveBeenCalledWith('ssh-1') expect(mockMux.dispose).toHaveBeenCalledWith('shutdown') - expect(mockStore.markSshRemotePtyLeases).toHaveBeenCalledWith('ssh-1', 'terminated') + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'terminated') expect(mockConnectionManager.disconnect).toHaveBeenCalledWith('ssh-1') expect(mockStore.removeSshRemotePtyLeases).toHaveBeenCalledWith('ssh-1') expect(mockSshStore.removeTarget).toHaveBeenCalledWith('ssh-1') @@ -2620,7 +2626,7 @@ describe('SSH IPC handlers', () => { }) await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' }) - mockStore.markSshRemotePtyLeases.mockClear() + mockStore.markSshRemotePtyLeasesAsync.mockClear() mockStore.markSshRemotePtyLease.mockClear() mockStore.getSshRemotePtyLeases.mockReturnValue([ { targetId: 'ssh-1', ptyId: 'pty-1', state: 'attached' } @@ -2628,12 +2634,220 @@ describe('SSH IPC handlers', () => { await handlers.get('ssh:resetRelay')!(null, { targetId: 'ssh-1' }) - expect(mockStore.markSshRemotePtyLeases).not.toHaveBeenCalledWith('ssh-1', 'terminated') - expect(mockStore.markSshRemotePtyLeases).toHaveBeenCalledWith('ssh-1', 'detached') + expect(mockStore.markSshRemotePtyLeasesAsync).not.toHaveBeenCalledWith('ssh-1', 'terminated') + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'detached') expect(mockStore.markSshRemotePtyLease).toHaveBeenCalledWith('ssh-1', 'pty-1', 'expired') expect(mockForceStopRelayForTarget).toHaveBeenCalledWith(conn, 'ssh-1') }) + describe('SSH PTY consumer identity across failed connects', () => { + function makeTarget(id: string): SshTarget { + return { id, label: 'Server', host: 'example.com', port: 22, username: 'deploy' } + } + + function markConnected(targetId: string): void { + mockConnectionManager.getState.mockReturnValue({ + targetId, + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + } + + it('reclaims the consumer identity after a failed transport connect', async () => { + const targetId = 'ssh-consumer-identity-connect-failure' + let settleLeasePersistence!: () => void + mockStore.markSshRemotePtyLeasesAsync.mockImplementationOnce( + () => + new Promise((resolve) => { + settleLeasePersistence = resolve + }) + ) + mockSshStore.getTarget.mockReturnValue(makeTarget(targetId)) + mockConnectionManager.connect.mockRejectedValueOnce(new Error('transport refused')) + + const failedConnect = handlers.get('ssh:connect')!(null, { targetId }) as Promise + const failure = expect(failedConnect).rejects.toThrow('transport refused') + const settled = vi.fn() + void failedConnect.then(settled, settled) + + await vi.waitFor(() => + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith(targetId, 'detached') + ) + expect(mockStore.markSshRemotePtyLeases).not.toHaveBeenCalled() + // Why: the connect rejection is gated on the durable 'detached' write, so the retry it + // triggers cannot re-mark leases 'attached' ahead of the abandoned session's release. + expect(settled).not.toHaveBeenCalled() + + settleLeasePersistence() + await failure + const claimedId = getSshPtyConsumerRecovery(targetId)?.clientInstanceId + expect(claimedId).toEqual(expect.any(String)) + + mockConnectionManager.connect.mockResolvedValue({}) + markConnected(targetId) + await handlers.get('ssh:connect')!(null, { targetId }) + + expect(getSshPtyConsumerRecovery(targetId)?.clientInstanceId).toBe(claimedId) + expect(mockStore.upsertSshPtyConsumerRecovery).toHaveBeenCalledWith( + expect.objectContaining({ targetId, clientInstanceId: claimedId }) + ) + }) + + it('releases the abandoned leases before a fast reconnect re-owns them', async () => { + const targetId = 'ssh-consumer-identity-fast-reconnect' + const order: string[] = [] + let settleLeaseRelease!: () => void + mockStore.markSshRemotePtyLeasesAsync.mockImplementationOnce((_id: string, state: string) => { + order.push(`leases:${state}`) + return new Promise((resolve) => { + settleLeaseRelease = () => { + order.push(`leases:${state}:persisted`) + resolve() + } + }) + }) + mockStore.upsertSshPtyConsumerRecovery.mockImplementation(async () => { + order.push('recovery:upsert') + }) + mockSshStore.getTarget.mockReturnValue(makeTarget(targetId)) + mockConnectionManager.connect.mockRejectedValueOnce(new Error('transport refused')) + + const failedConnect = handlers.get('ssh:connect')!(null, { targetId }) + const failure = expect(failedConnect).rejects.toThrow('transport refused') + await vi.waitFor(() => expect(order).toContain('leases:detached')) + expect(order).not.toContain('leases:detached:persisted') + + settleLeaseRelease() + await failure + + mockConnectionManager.connect.mockResolvedValue({}) + markConnected(targetId) + await handlers.get('ssh:connect')!(null, { targetId }) + + // Why: the reclaimed owner is only re-persisted after the abandoned 'detached' write landed, + // so no late release can strand the reconnected leases in 'detached'. + expect(order).toEqual(['leases:detached', 'leases:detached:persisted', 'recovery:upsert']) + }) + + it('holds a retry that starts while the detach write is still pending', async () => { + const targetId = 'ssh-consumer-identity-pending-retry' + const order: string[] = [] + let settleLeaseRelease!: () => void + mockStore.markSshRemotePtyLeasesAsync.mockImplementationOnce((_id: string, state: string) => { + order.push(`leases:${state}`) + return new Promise((resolve) => { + settleLeaseRelease = () => { + order.push(`leases:${state}:persisted`) + resolve() + } + }) + }) + mockStore.upsertSshPtyConsumerRecovery.mockImplementation(async () => { + order.push('recovery:upsert') + }) + mockSshStore.getTarget.mockReturnValue(makeTarget(targetId)) + mockConnectionManager.connect.mockRejectedValueOnce(new Error('transport refused')) + + const failedConnect = handlers.get('ssh:connect')!(null, { targetId }) + const failure = expect(failedConnect).rejects.toThrow('transport refused') + await vi.waitFor(() => expect(order).toContain('leases:detached')) + + // Retry mid-write: it must not mint a session or re-own leases while the release is pending. + mockConnectionManager.connect.mockResolvedValue({}) + markConnected(targetId) + const retry = handlers.get('ssh:connect')!(null, { targetId }) as Promise + const retrySettled = vi.fn() + void retry.then(retrySettled, retrySettled) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(retrySettled).not.toHaveBeenCalled() + expect(order).not.toContain('leases:detached:persisted') + + settleLeaseRelease() + await failure + // Why: the retry folds onto the still-latched attempt rather than starting a second connect, + // so it inherits that attempt's failure instead of racing its teardown. + await expect(retry).rejects.toThrow('transport refused') + expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1) + expect(order).toEqual(['leases:detached', 'leases:detached:persisted']) + + await handlers.get('ssh:connect')!(null, { targetId }) + + expect(order).toEqual(['leases:detached', 'leases:detached:persisted', 'recovery:upsert']) + }) + + it('replaces a live session whose detach write keeps failing', async () => { + const targetId = 'ssh-consumer-identity-detach-write-failure' + mockSshStore.getTarget.mockReturnValue(makeTarget(targetId)) + mockConnectionManager.connect.mockResolvedValue({}) + markConnected(targetId) + await handlers.get('ssh:connect')!(null, { targetId }) + const claimedId = getSshPtyConsumerRecovery(targetId)?.clientInstanceId + expect(claimedId).toEqual(expect.any(String)) + + // Why a permanent reject, not once: it proves the failed session is gone rather than merely + // retried — a second connect that still holds it would fail on the same write again. + mockStore.markSshRemotePtyLeasesAsync.mockImplementation((_id: string, state: string) => + state === 'detached' + ? Promise.reject(new Error('lease write failed')) + : Promise.resolve(undefined) + ) + await expect(handlers.get('ssh:connect')!(null, { targetId })).rejects.toThrow( + 'lease write failed' + ) + + mockConnectionManager.connect.mockClear() + await handlers.get('ssh:connect')!(null, { targetId }) + + expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1) + // Why: the abandoned session still released its identity synchronously, so the replacement + // reclaims the owner instead of minting a new one. + expect(getSshPtyConsumerRecovery(targetId)?.clientInstanceId).toBe(claimedId) + }) + + it('resumes the remembered owner lease after a failed establish', async () => { + const targetId = 'ssh-consumer-identity-establish-failure' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + mockSshStore.getTarget.mockReturnValue(makeTarget(targetId)) + mockConnectionManager.connect.mockResolvedValue({}) + markConnected(targetId) + // Why: fail the first request after the consumer session opens, so establish() rejects with an + // owner lease already remembered — the state a retry must be able to resume from. + const openClientResponse = await mockMux.request('pty.openClient') + mockMux.request.mockImplementationOnce(() => Promise.resolve(openClientResponse)) + mockMux.request.mockImplementationOnce(() => + Promise.reject(new Error('relay handshake aborted')) + ) + mockStore.markSshRemotePtyLeasesAsync.mockRejectedValueOnce( + new Error('lease persistence failed') + ) + + try { + await expect(handlers.get('ssh:connect')!(null, { targetId })).rejects.toThrow( + 'relay handshake aborted' + ) + await vi.waitFor(() => expect(warn).toHaveBeenCalled()) + const claimedId = getSshPtyConsumerRecovery(targetId)?.clientInstanceId + expect(claimedId).toEqual(expect.any(String)) + + mockMux.request.mockClear() + await handlers.get('ssh:connect')!(null, { targetId }) + + expect(getSshPtyConsumerRecovery(targetId)?.clientInstanceId).toBe(claimedId) + expect(mockMux.request).toHaveBeenCalledWith( + 'pty.openClient', + expect.objectContaining({ + clientInstanceId: claimedId, + resume: { ownerGeneration: 1, ownerLease: 'ipc-test-owner' } + }), + expect.anything() + ) + } finally { + warn.mockRestore() + } + }) + }) + it('ssh:getState returns connection state', async () => { const state = { targetId: 'ssh-1', diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index 91234031aa3..995d3e88b72 100644 --- a/src/main/ipc/ssh.ts +++ b/src/main/ipc/ssh.ts @@ -115,7 +115,7 @@ export function listRegisteredRemovedSshTargetLabels(): Record { export async function disconnectRegisteredSshTarget(targetId: string): Promise { invalidateConnectAttempt(targetId) await runTargetLifecycle(targetId, () => - teardownSshTargetTransport(targetId, (session) => session.detach()) + teardownSshTargetTransport(targetId, (session) => session.detachAndPersist()) ) } @@ -128,7 +128,7 @@ export async function removeRegisteredSshTarget(targetId: string): Promise await runTargetLifecycle(targetId, async () => { try { // Why: removal is destructive; dispose so remote PTYs cannot reattach to a deleted target. - await teardownSshTargetTransport(targetId, (session) => session.dispose()) + await teardownSshTargetTransport(targetId, (session) => session.disposeAndPersist()) } catch (err) { // Why: a failed disconnect must not block metadata removal, else the target lingers in the store with uncleaned leases. console.warn( @@ -191,7 +191,7 @@ async function awaitTargetLifecycle(targetId: string): Promise { async function teardownSshTargetTransport( targetId: string, - teardown: (session: SshRelaySession) => void + teardown: (session: SshRelaySession) => void | Promise ): Promise { let transportDisconnect: Promise<{ ok: true } | { ok: false; error: unknown }> try { @@ -220,7 +220,7 @@ async function teardownSshTargetTransport( async function teardownActiveSshSession( targetId: string, - teardown: (session: SshRelaySession) => void + teardown: (session: SshRelaySession) => void | Promise ): Promise { const session = activeSessions.get(targetId) if (!session) { @@ -234,7 +234,7 @@ async function teardownActiveSshSession( teardownError = { error } } try { - teardown(session) + await teardown(session) } catch (error) { teardownError ??= { error } } @@ -248,6 +248,27 @@ async function teardownActiveSshSession( } } +// Why: a dropped session must detach, not just leave activeSessions — detach releases the SSH PTY +// consumer identity so the next connect reclaims its owner lease instead of minting a new one. +// Why awaited, and why the map entry outlives the await: a retry can start the moment this returns, +// so it must either find the session (and await this same latched teardown at the existing-session +// path) or find nothing because the 'detached' lease write is already durable. Deleting first lets a +// fast reconnect mark leases 'attached' and then have this session's late 'detached' write clobber it. +async function abandonFailedSshSession(targetId: string, session: SshRelaySession): Promise { + // Why: detachAndPersist transitions recovery ownership synchronously; only durability is awaited. + try { + await session.detachAndPersist() + } catch (error) { + // Why: a teardown throw must not mask the connect error the caller is about to rethrow. + console.warn( + `[ssh] Failed to detach abandoned session for ${targetId}: ${error instanceof Error ? error.message : String(error)}` + ) + } + if (activeSessions.get(targetId) === session) { + activeSessions.delete(targetId) + } +} + function relayGracePeriodForTarget(target: SshTarget | null | undefined): number | undefined { return target?.relayGracePeriodSeconds } @@ -1046,11 +1067,18 @@ export function registerSshHandlers( if (!isCurrentConnectAttempt(targetId, authority)) { throw createCancelledConnectAttemptError() } - existingSession.detach() - if (activeSessions.get(targetId) === existingSession) { - activeSessions.delete(targetId) - clearRelayLostBackoff(targetId) - clearRelayStateOverride(targetId) + try { + await existingSession.detachAndPersist() + } finally { + // Why finally: detachAndPersist runs its in-memory half synchronously, so the session is + // dead even when the lease write rejects — keeping it in activeSessions would strand every + // later connect on the same dead session. Why still after the await, not before it: the + // write has settled by now, so it can no longer clobber the replacement's 'attached' write. + if (activeSessions.get(targetId) === existingSession) { + activeSessions.delete(targetId) + clearRelayLostBackoff(targetId) + clearRelayStateOverride(targetId) + } } } @@ -1092,7 +1120,7 @@ export function registerSshHandlers( } // Why: clear this failed connect's flag so a later non-prompting connect isn't deferred. credentialRequestedForTarget.delete(targetId) - activeSessions.delete(targetId) + await abandonFailedSshSession(targetId, session) clearRelayLostBackoff(targetId) clearRelayStateOverride(targetId) broadcastSshState(getCurrentMainWindow, targetId, { @@ -1130,9 +1158,16 @@ export function registerSshHandlers( if (!ownsSession()) { throw createCancelledConnectAttemptError() } - activeSessions.delete(targetId) + await abandonFailedSshSession(targetId, session) clearRelayLostBackoff(targetId) - await connectionManager!.disconnect(targetId) + try { + await connectionManager!.disconnect(targetId) + } catch (disconnectError) { + // Why: the establish failure is the actionable error; a teardown throw must not replace it. + console.warn( + `[ssh] Failed to disconnect transport after failed establish for ${targetId}: ${disconnectError instanceof Error ? disconnectError.message : String(disconnectError)}` + ) + } throw err } @@ -1202,7 +1237,7 @@ export function registerSshHandlers( // Why: a failed relay shutdown can leave the remote process alive in the grace window; keep the lease/session so the user can retry. throw new Error(`Failed to terminate SSH host sessions: ${shutdownFailures.join('; ')}`) } - await teardownSshTargetTransport(args.targetId, (session) => session.dispose()) + await teardownSshTargetTransport(args.targetId, (session) => session.disposeAndPersist()) }) }) @@ -1221,7 +1256,9 @@ export function registerSshHandlers( const session = activeSessions.get(targetId) if (session) { // Why: detach() not dispose() — reset has its own stale-lease semantics below that dispose()'s clean-termination recording would hide. - await teardownActiveSshSession(targetId, (capturedSession) => capturedSession.detach()) + await teardownActiveSshSession(targetId, (capturedSession) => + capturedSession.detachAndPersist() + ) } const existingConn = connectionManager!.getConnection(targetId) @@ -1445,9 +1482,10 @@ export async function resetSshHandlerStateForTests(): Promise { } ipcMain.removeHandler('ssh:submitCredential') - for (const session of activeSessions.values()) { - session.dispose() - } + // Why: allSettled — a rejected disposal write must not abort the rest of the reset and leak state into the next test. + await Promise.allSettled( + [...activeSessions.values()].map((session) => session.disposeAndPersist()) + ) activeSessions.clear() for (const targetId of relayLostBackoff.keys()) { clearRelayLostBackoff(targetId) diff --git a/src/main/persistence-async-write-syscalls.test.ts b/src/main/persistence-async-write-syscalls.test.ts index 5ca2cbe5e74..2b3a6ea68b1 100644 --- a/src/main/persistence-async-write-syscalls.test.ts +++ b/src/main/persistence-async-write-syscalls.test.ts @@ -12,6 +12,7 @@ import type * as NodeFs from 'node:fs' import type * as NodeFsPromises from 'node:fs/promises' import { join } from 'node:path' import { tmpdir } from 'node:os' +import type { SshRemotePtyLeaseState } from '../shared/ssh-types' const testState = { dir: '' } @@ -143,6 +144,33 @@ type TestStore = { flushOrThrow(): void flushPendingAsync(): Promise flushPendingOrThrowAsync(): Promise + upsertSshPtyConsumerRecovery(record: { + targetId: string + clientInstanceId: string + serverBuildId: string + clientGeneration: number + ownerGeneration: number + ownerLease: string + }): Promise + removeSshPtyConsumerRecovery(targetId: string): Promise + upsertSshRemotePtyLease(lease: { + targetId: string + ptyId: string + state: SshRemotePtyLeaseState + }): void + markSshRemotePtyLeasesAsync(targetId: string, state: SshRemotePtyLeaseState): Promise + markSshRemotePtyLeasesAttachedAsync(targetId: string, ptyIds: readonly string[]): Promise +} + +function consumerRecovery(clientInstanceId: string) { + return { + targetId: 'ssh-1', + clientInstanceId, + serverBuildId: 'relay-build-1', + clientGeneration: 3, + ownerGeneration: 5, + ownerLease: 'secret-owner-lease' + } } async function createStore(dir: string): Promise { @@ -526,6 +554,8 @@ describe('async persistence write path avoids synchronous fs syscalls', () => { vi.advanceTimersByTime(SAVE_DEBOUNCE_MS) allWrites = store.waitForPendingWrite() expect(fsCalls.asyncCalls.filter((call) => call === statCall)).toHaveLength(1) + releaseRotation() + await Promise.all([firstWrite, allWrites]) } finally { releaseRotation() await allWrites @@ -654,6 +684,201 @@ describe('async persistence write path avoids synchronous fs syscalls', () => { expect(Object.keys(ring)).toContain(`orca-data.json.bak.${BACKUP_COUNT - 1}`) }) + it('persists SSH PTY consumer recovery without a sync syscall, durable once awaited', async () => { + const dir = makeDir() + const store = await createStore(dir) + + fsCalls.dirPrefix = dir + fsCalls.recording = true + try { + await store.upsertSshPtyConsumerRecovery(consumerRecovery('client-1')) + } finally { + fsCalls.recording = false + } + + expect(fsCalls.syncCalls).toEqual([]) + // Durability is awaited, not merely debounced: the record is on disk when the promise resolves. + const persisted = JSON.parse(readFileSync(dataFile(dir), 'utf-8')) as { + sshPtyConsumerRecoveries: { clientInstanceId: string }[] + } + expect(persisted.sshPtyConsumerRecoveries).toHaveLength(1) + expect(persisted.sshPtyConsumerRecoveries[0]?.clientInstanceId).toBe('client-1') + }) + + it('rejects the consumer-recovery durability barrier when the primary write fails', async () => { + const dir = makeDir() + const store = await createStore(dir) + const writeError = Object.assign(new Error('profile mount rejected write'), { code: 'EIO' }) + const errors = vi.spyOn(console, 'error').mockImplementation(() => {}) + fsCalls.dirPrefix = dir + fsCalls.recording = true + fsCalls.failAsync = (fn, target) => + fn === 'open' && target.startsWith(`${dataFile(dir)}.`) ? writeError : null + + try { + await expect(store.upsertSshPtyConsumerRecovery(consumerRecovery('client-1'))).rejects.toBe( + writeError + ) + } finally { + fsCalls.recording = false + errors.mockRestore() + } + }) + + it('removes SSH PTY consumer recovery without a sync syscall, durable once awaited', async () => { + const dir = makeDir() + const store = await createStore(dir) + await store.upsertSshPtyConsumerRecovery(consumerRecovery('client-1')) + + fsCalls.dirPrefix = dir + fsCalls.recording = true + try { + await store.removeSshPtyConsumerRecovery('ssh-1') + } finally { + fsCalls.recording = false + } + + expect(fsCalls.syncCalls).toEqual([]) + const persisted = JSON.parse(readFileSync(dataFile(dir), 'utf-8')) as { + sshPtyConsumerRecoveries: unknown[] + } + expect(persisted.sshPtyConsumerRecoveries).toEqual([]) + }) + + it('persists failed-session lease detachment without a sync syscall', async () => { + const dir = makeDir() + const store = await createStore(dir) + store.upsertSshRemotePtyLease({ targetId: 'ssh-1', ptyId: 'pty-1', state: 'attached' }) + + fsCalls.dirPrefix = dir + fsCalls.recording = true + try { + await store.markSshRemotePtyLeasesAsync('ssh-1', 'detached') + } finally { + fsCalls.recording = false + } + + expect(fsCalls.syncCalls).toEqual([]) + const persisted = JSON.parse(readFileSync(dataFile(dir), 'utf-8')) as { + sshRemotePtyLeases: { state: string }[] + } + expect(persisted.sshRemotePtyLeases[0]?.state).toBe('detached') + }) + + it('persists selected reattach leases in one async write', async () => { + const dir = makeDir() + const store = await createStore(dir) + store.upsertSshRemotePtyLease({ targetId: 'ssh-1', ptyId: 'pty-1', state: 'detached' }) + store.upsertSshRemotePtyLease({ targetId: 'ssh-1', ptyId: 'pty-2', state: 'expired' }) + store.upsertSshRemotePtyLease({ targetId: 'ssh-1', ptyId: 'pty-3', state: 'detached' }) + // Why: a PTY that exits mid-reattach is terminated before the batch write lands; it must stay dead. + store.upsertSshRemotePtyLease({ targetId: 'ssh-1', ptyId: 'pty-4', state: 'terminated' }) + + fsCalls.dirPrefix = dir + fsCalls.recording = true + try { + await store.markSshRemotePtyLeasesAttachedAsync('ssh-1', ['pty-1', 'pty-2', 'pty-4']) + } finally { + fsCalls.recording = false + } + + expect(fsCalls.syncCalls).toEqual([]) + const persisted = JSON.parse(readFileSync(dataFile(dir), 'utf-8')) as { + sshRemotePtyLeases: { ptyId: string; state: string }[] + } + expect(persisted.sshRemotePtyLeases).toEqual( + expect.arrayContaining([ + expect.objectContaining({ ptyId: 'pty-1', state: 'attached' }), + expect.objectContaining({ ptyId: 'pty-2', state: 'expired' }), + expect.objectContaining({ ptyId: 'pty-3', state: 'detached' }), + expect.objectContaining({ ptyId: 'pty-4', state: 'terminated' }) + ]) + ) + }) + + it('keeps async writers serialized across a synchronous shutdown flush', async () => { + const dir = makeDir() + const store = await createStore(dir) + let signalFirstOpen!: () => void + const firstOpen = new Promise((resolve) => { + signalFirstOpen = resolve + }) + let releaseFirstOpen!: () => void + const firstOpenRelease = new Promise((resolve) => { + releaseFirstOpen = resolve + }) + let held = false + fsCalls.waitAsync = (fn, target) => { + if (held || fn !== 'open' || !target.endsWith('.tmp')) { + return null + } + held = true + signalFirstOpen() + return firstOpenRelease + } + + fsCalls.dirPrefix = dir + fsCalls.recording = true + try { + const firstWrite = store.upsertSshPtyConsumerRecovery(consumerRecovery('client-1')) + await firstOpen + store.flushOrThrow() + const secondWrite = store.upsertSshPtyConsumerRecovery(consumerRecovery('client-2')) + await Promise.resolve() + await Promise.resolve() + + expect(fsCalls.asyncCalls.filter((call) => call.startsWith('open:'))).toHaveLength(1) + releaseFirstOpen() + await Promise.all([firstWrite, secondWrite]) + } finally { + releaseFirstOpen() + fsCalls.recording = false + fsCalls.waitAsync = null + } + + const persisted = JSON.parse(readFileSync(dataFile(dir), 'utf-8')) as { + sshPtyConsumerRecoveries: { clientInstanceId: string }[] + } + expect(persisted.sshPtyConsumerRecoveries[0]?.clientInstanceId).toBe('client-2') + }) + + it('lets a main-thread timer keep firing while a consumer-recovery write is in flight', async () => { + // The P1-A freeze itself: a stalled profile mount must not park the main thread on establish. + vi.useRealTimers() + const dir = makeDir() + const store = await createStore(dir) + const stallMs = 1_000 + // Why half: a sync write parks the loop for at least stallMs while the async path ticks every + // ~10ms, so this leaves room for scheduler jitter on a loaded runner without going vacuous. + const maxAcceptableGapMs = stallMs / 2 + + let lastTick = Date.now() + let worstGapMs = 0 + const heartbeat = setInterval(() => { + const now = Date.now() + worstGapMs = Math.max(worstGapMs, now - lastTick) + lastTick = now + }, 10) + + try { + fsCalls.dirPrefix = dir + fsCalls.stallMs = stallMs + fsCalls.recording = true + lastTick = Date.now() + const write = store.upsertSshPtyConsumerRecovery(consumerRecovery('client-1')) + // Why not await first: a fully synchronous write finishes before the interval can fire, so the + // heartbeat would never observe the stall it exists to detect. + await new Promise((resolve) => setTimeout(resolve, stallMs + 200)) + await write + } finally { + fsCalls.recording = false + clearInterval(heartbeat) + } + + expect(worstGapMs).toBeLessThan(maxAcceptableGapMs) + expect(readFileSync(dataFile(dir), 'utf-8')).toContain('client-1') + }, 20_000) + it('keeps the sync quit/crash fallback on synchronous syscalls', async () => { const dir = makeDir() const store = await createStore(dir) diff --git a/src/main/persistence.test.ts b/src/main/persistence.test.ts index 6cddc809c29..8fb3ad85904 100644 --- a/src/main/persistence.test.ts +++ b/src/main/persistence.test.ts @@ -6800,7 +6800,7 @@ describe('Store', () => { it('durably encrypts SSH PTY consumer ownership for process restart recovery', async () => { const store = await createStore() store.setGitHubCache({ pr: { 'o/r#1': { fetchedAt: 1 } as never }, issue: {} }) - store.upsertSshPtyConsumerRecovery({ + await store.upsertSshPtyConsumerRecovery({ targetId: 'ssh-1', clientInstanceId: 'client-1', serverBuildId: 'relay-build-1', @@ -6858,7 +6858,7 @@ describe('Store', () => { port: 22, username: 'orca' }) - store.upsertSshPtyConsumerRecovery({ + await store.upsertSshPtyConsumerRecovery({ targetId: 'ssh-1', clientInstanceId: 'client-1', serverBuildId: 'relay-build-1', diff --git a/src/main/persistence.ts b/src/main/persistence.ts index e7245d27e00..966686bfbdf 100644 --- a/src/main/persistence.ts +++ b/src/main/persistence.ts @@ -4044,6 +4044,7 @@ export class Store { } } } + // Why: later async flushes must remain serialized behind the invalidated writer. this.writeToDiskSync({ force: asyncWriteWasInFlight, skipBackupRotation: this.backupRotationInFlight @@ -6888,7 +6889,7 @@ export class Store { return record ? structuredClone(record) : null } - upsertSshPtyConsumerRecovery(record: SshPtyConsumerRecovery): void { + async upsertSshPtyConsumerRecovery(record: SshPtyConsumerRecovery): Promise { const normalized = normalizeSshPtyConsumerRecovery(record) if (!normalized) { throw new Error('Invalid SSH PTY consumer recovery record') @@ -6898,26 +6899,24 @@ export class Store { ...recoveries.filter((candidate) => candidate.targetId !== normalized.targetId), normalized ] - this.flushSshPtyConsumerRecovery() + await this.flushSshPtyConsumerRecovery() } - removeSshPtyConsumerRecovery(targetId: string): void { + async removeSshPtyConsumerRecovery(targetId: string): Promise { const recoveries = this.state.sshPtyConsumerRecoveries ?? [] const next = recoveries.filter((record) => record.targetId !== targetId) if (next.length === recoveries.length) { return } this.state.sshPtyConsumerRecoveries = next - this.flushSshPtyConsumerRecovery() + await this.flushSshPtyConsumerRecovery() } - private flushSshPtyConsumerRecovery(): void { - try { - // Why: ownership must be durable before relay setup continues, but active-view and GitHub sidecars are unrelated startup work. - this.flushOrThrow() - } catch (err) { - console.error('[persistence] Failed to flush SSH PTY consumer recovery:', err) - } + private async flushSshPtyConsumerRecovery(): Promise { + // Why: ownership must be durable before relay setup continues, but this runs on the live + // establish/reconnect path — a sync flush would park the main thread on a stalled profile mount. + // Why not caught here: the failure must reach the awaiting caller. + await this.flushDurableStateOrThrowAsync() } // ── SSH Remote PTY Leases ────────────────────────────────────────── @@ -6962,13 +6961,47 @@ export class Store { } markSshRemotePtyLeases(targetId: string, state: SshRemotePtyLease['state']): void { + if (this.updateSshRemotePtyLeaseStates(targetId, state)) { + this.flush() + } + } + + async markSshRemotePtyLeasesAsync( + targetId: string, + state: SshRemotePtyLease['state'] + ): Promise { + if (this.updateSshRemotePtyLeaseStates(targetId, state)) { + await this.flushDurableStateOrThrowAsync() + } + } + + async markSshRemotePtyLeasesAttachedAsync( + targetId: string, + ptyIds: readonly string[] + ): Promise { + const relayPtyIds = new Set( + ptyIds.map((ptyId) => this.getRelayPtyIdForSshLeaseStorage(targetId, ptyId)) + ) + if (this.updateSshRemotePtyLeaseStates(targetId, 'attached', relayPtyIds)) { + await this.flushDurableStateOrThrowAsync() + } + } + + private updateSshRemotePtyLeaseStates( + targetId: string, + state: SshRemotePtyLease['state'], + ptyIds?: ReadonlySet + ): boolean { const now = Date.now() let changed = false const shouldClearBindings = state === 'terminated' || state === 'expired' const leasesToClear: SshRemotePtyLease[] = [] this.state.sshRemotePtyLeases ??= [] for (const lease of this.state.sshRemotePtyLeases) { - if (lease.targetId !== targetId) { + 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') { @@ -6991,9 +7024,7 @@ export class Store { const bindingsChanged = shouldClearBindings ? this.clearSshRemotePtyBindingsForLeases(targetId, leasesToClear) : false - if (changed || bindingsChanged) { - this.flush() - } + return changed || bindingsChanged } markSshRemotePtyLease(targetId: string, ptyId: string, state: SshRemotePtyLease['state']): void { @@ -7174,6 +7205,26 @@ export class Store { return this.flushCurrentStateAsync(false, options.signal) } + // Async twin of flushOrThrow: durable state only. Active-view and GitHub sidecars are + // quit/startup work and must not be snapshotted on the live SSH establish/reconnect path. + private async flushDurableStateOrThrowAsync(): Promise { + if (this.writesFrozen || this.quitFlushStarted) { + throw new Error('Cannot flush while persistence is finalized') + } + for (;;) { + if (this.writeTimer) { + clearTimeout(this.writeTimer) + this.writeTimer = null + } + this.firstPendingSaveAt = null + const generation = this.writeGeneration + await this.enqueueWrite() + if (generation === this.writeGeneration) { + break + } + } + } + private async flushCurrentStateAsync( final: boolean, signal?: AbortSignal, diff --git a/src/main/ssh/ssh-pty-consumer-recovery.test.ts b/src/main/ssh/ssh-pty-consumer-recovery.test.ts new file mode 100644 index 00000000000..09405f9ee0b --- /dev/null +++ b/src/main/ssh/ssh-pty-consumer-recovery.test.ts @@ -0,0 +1,45 @@ +import { describe, expect, it, vi } from 'vitest' +import type { Store } from '../persistence' +import type { SshPtyConsumerOwnerState } from './ssh-pty-consumer-session' +import { + claimSshPtyConsumerRecovery, + detachSshPtyConsumerRecovery, + rememberSshPtyConsumerRecovery +} from './ssh-pty-consumer-recovery' + +describe('SSH PTY consumer recovery', () => { + it('keeps a detached identity when a concurrent open finishes late', async () => { + const targetId = 'remember-detach-race' + const store = { + getSshPtyConsumerRecovery: vi.fn().mockReturnValue(null), + upsertSshPtyConsumerRecovery: vi.fn() + } as unknown as Store + const claimed = claimSshPtyConsumerRecovery(targetId, store) + let finishOpen!: (owner: SshPtyConsumerOwnerState) => void + const opened = new Promise((resolve) => { + finishOpen = resolve + }) + const remembering = opened.then((owner) => + rememberSshPtyConsumerRecovery({ + targetId, + clientInstanceId: claimed.clientInstanceId, + serverBuildId: 'relay-build', + owner, + store + }) + ) + + detachSshPtyConsumerRecovery(targetId, claimed.clientInstanceId) + finishOpen({ + mode: 'negotiated', + clientInstanceId: claimed.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'late-owner' + }) + await remembering + + expect(store.upsertSshPtyConsumerRecovery).not.toHaveBeenCalled() + expect(claimSshPtyConsumerRecovery(targetId, store)).toBe(claimed) + }) +}) diff --git a/src/main/ssh/ssh-pty-consumer-recovery.ts b/src/main/ssh/ssh-pty-consumer-recovery.ts index 7b3fe215b76..1c357d32a12 100644 --- a/src/main/ssh/ssh-pty-consumer-recovery.ts +++ b/src/main/ssh/ssh-pty-consumer-recovery.ts @@ -56,21 +56,22 @@ export function getSshPtyConsumerRecovery( return recoveryByTarget.get(targetId) } -export function rememberSshPtyConsumerRecovery(args: { +// Why async: the in-memory update lands synchronously (before the first await) so no caller can +// observe a torn state; only the awaited durability barrier is deferred. +export async function rememberSshPtyConsumerRecovery(args: { targetId: string clientInstanceId: string serverBuildId: string owner: SshPtyConsumerOwnerState store: Store -}): void { +}): Promise { const current = recoveryByTarget.get(args.targetId) - if (current?.clientInstanceId !== args.clientInstanceId) { + if (current?.clientInstanceId !== args.clientInstanceId || current.detached) { return } - current.detached = false current.serverBuildId = args.serverBuildId current.owner = args.owner - args.store.upsertSshPtyConsumerRecovery({ + await args.store.upsertSshPtyConsumerRecovery({ targetId: args.targetId, clientInstanceId: args.clientInstanceId, serverBuildId: args.serverBuildId, @@ -81,14 +82,14 @@ export function rememberSshPtyConsumerRecovery(args: { }) } -export function removeSshPtyConsumerOwnerRecovery( +export async function removeSshPtyConsumerOwnerRecovery( targetId: string, clientInstanceId: string, store: Store -): void { +): Promise { const persisted = store.getSshPtyConsumerRecovery(targetId) if (persisted?.clientInstanceId === clientInstanceId) { - store.removeSshPtyConsumerRecovery(targetId) + await store.removeSshPtyConsumerRecovery(targetId) } } @@ -99,14 +100,14 @@ export function detachSshPtyConsumerRecovery(targetId: string, clientInstanceId: } } -export function forgetSshPtyConsumerRecovery( +export async function forgetSshPtyConsumerRecovery( targetId: string, clientInstanceId: string, store: Store -): void { +): Promise { const current = recoveryByTarget.get(targetId) if (current?.clientInstanceId === clientInstanceId) { recoveryByTarget.delete(targetId) } - removeSshPtyConsumerOwnerRecovery(targetId, clientInstanceId, store) + await removeSshPtyConsumerOwnerRecovery(targetId, clientInstanceId, store) } 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 01e4ec1dfb5..f92b72b09a8 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 @@ -154,7 +154,9 @@ function createSession(targetId: string): InstanceType { removeSshPtyConsumerRecovery: vi.fn(), getSshRemotePtyLeases: vi.fn().mockReturnValue([]), markSshRemotePtyLease: vi.fn(), - markSshRemotePtyLeases: vi.fn() + markSshRemotePtyLeases: vi.fn(), + markSshRemotePtyLeasesAsync: vi.fn(), + markSshRemotePtyLeasesAttachedAsync: vi.fn() } as unknown as Store const portForwardManager = { removeAllForwards: vi.fn().mockResolvedValue(undefined) diff --git a/src/main/ssh/ssh-relay-session-data-delivery.test.ts b/src/main/ssh/ssh-relay-session-data-delivery.test.ts index 19939ae3857..56f6c6eeff2 100644 --- a/src/main/ssh/ssh-relay-session-data-delivery.test.ts +++ b/src/main/ssh/ssh-relay-session-data-delivery.test.ts @@ -819,7 +819,7 @@ describe('SshRelaySession data delivery', () => { ) expect(mockStore.markSshRemotePtyLease).toHaveBeenCalledWith(targetId, 'pty-1', 'detached') expect(mockStore.markSshRemotePtyLease).not.toHaveBeenCalledWith(targetId, 'pty-1', 'expired') - expect(mockStore.markSshRemotePtyLeases).not.toHaveBeenCalled() + expect(mockStore.markSshRemotePtyLeasesAsync).not.toHaveBeenCalled() expect(clearProviderPtyState).not.toHaveBeenCalled() expect(clearPtyOwnershipForConnection).not.toHaveBeenCalled() expect(deletePtyOwnership).not.toHaveBeenCalled() @@ -836,10 +836,8 @@ describe('SshRelaySession data delivery', () => { }) expect(acceptOutputDataMock.mock.calls.map(([payload]) => payload.data)).toEqual(['live']) - expect(mockStore.markSshRemotePtyLease).toHaveBeenCalledWith( - 'empty-recovery', - 'pty-1', - 'attached' - ) + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith('empty-recovery', [ + 'pty-1' + ]) }) }) 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 c624e72564d..c34c4298dd8 100644 --- a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts @@ -220,12 +220,8 @@ describe('SshRelaySession reconnect incarnation ordering', () => { const reconnect = session.reconnect(mockConn) await vi.advanceTimersByTimeAsync(750) - expect( - vi - .mocked(mockStore.markSshRemotePtyLease) - .mock.calls.filter(([, , state]) => state === 'attached') - .map(([, id]) => id) - ).toHaveLength(48) + expect(setPtyOwnership).toHaveBeenCalledTimes(48) + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).not.toHaveBeenCalled() expect(peakActive).toBeLessThanOrEqual(8) await vi.advanceTimersByTimeAsync(20_000) @@ -235,6 +231,11 @@ describe('SshRelaySession reconnect incarnation ordering', () => { expect(attempts.get('pty-0')).toBe(2) expect(attempts.get('pty-1')).toBe(2) expect(session.getState()).toBe('ready') + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledOnce() + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith( + 'target-1', + expect.arrayContaining(ptyIds.slice(2)) + ) expect(mockStore.markSshRemotePtyLease).not.toHaveBeenCalledWith( 'target-1', expect.any(String), @@ -420,7 +421,7 @@ describe('SshRelaySession reconnect incarnation ordering', () => { incarnationId }) expect(vi.mocked(mockStore.persistPtyBinding).mock.invocationCallOrder[0]).toBeLessThan( - vi.mocked(mockStore.markSshRemotePtyLease).mock.invocationCallOrder[0]! + vi.mocked(mockStore.markSshRemotePtyLeasesAttachedAsync).mock.invocationCallOrder[0]! ) }) @@ -564,7 +565,9 @@ describe('SshRelaySession reconnect incarnation ordering', () => { leafId: INCARNATION_LEAF_ID, incarnationId }) - expect(mockStore.markSshRemotePtyLease).toHaveBeenCalledWith('target-1', 'pty-live', 'attached') + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith('target-1', [ + 'pty-live' + ]) expect(consoleError).toHaveBeenCalledWith( '[ssh-relay-session] Failed to persist reconnect incarnation:', expect.any(Error) diff --git a/src/main/ssh/ssh-relay-session-recovery-durability.test.ts b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts new file mode 100644 index 00000000000..df7bce635ec --- /dev/null +++ b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts @@ -0,0 +1,355 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { SshRelaySession } from './ssh-relay-session' +import { + createMismatchedOwnerRecoveryError, + createMockDeps +} from './ssh-relay-session-test-fixtures' +import { getSshPtyConsumerRecovery } from './ssh-pty-consumer-recovery' + +const { muxRequestMock, openConsumerSessionMock } = vi.hoisted(() => ({ + muxRequestMock: vi.fn(), + openConsumerSessionMock: vi.fn() +})) + +vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: openConsumerSessionMock +})) +vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ + acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined), + acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined), + allocateSshPtyProviderGeneration: vi.fn(() => 23), + beginSshPtyOutputGenerationMigration: vi.fn(() => ({ + byPty: new Map(), + completion: Promise.resolve() + })), + closeSshPtyOutputGeneration: vi.fn(), + getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []), + installSshPtySourceAckPublisher: vi.fn(() => () => {}), + installSshPtySourceCancellationPublisher: vi.fn(() => () => {}), + applySshPtySourceCancellationProof: vi.fn(), + applySshPtySourceRecoveryCancellationProof: vi.fn() +})) + +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', () => ({ + isSshPtyNotFoundError: vi.fn().mockReturnValue(false), + isSshPtyIdentityMismatchError: vi.fn().mockReturnValue(false), + SshPtyProvider: class MockSshPtyProvider { + onData = vi.fn().mockReturnValue(() => {}) + onReplay = vi.fn().mockReturnValue(() => {}) + onExit = vi.fn().mockReturnValue(() => {}) + attach = vi.fn().mockResolvedValue(undefined) + attachForReconnect = vi.fn().mockResolvedValue({}) + setPtyDeliveryPauseAdapter = vi.fn() + 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(), + restorePtyIncarnation: vi.fn(), + setPtyOwnership: vi.fn() +})) + +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 { deployAndLaunchRelay } = await import('./ssh-relay-deploy') +const { clearPtyOwnershipForConnection, unregisterSshPtyProvider } = await import('../ipc/pty') + +describe('SshRelaySession consumer recovery durability', () => { + beforeEach(() => { + vi.clearAllMocks() + muxRequestMock.mockResolvedValue([]) + openConsumerSessionMock.mockImplementation(async (_mux, options) => ({ + mode: 'negotiated', + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) + vi.mocked(deployAndLaunchRelay).mockResolvedValue({ + transport: { write: vi.fn(), onData: vi.fn(), onClose: vi.fn() }, + platform: 'linux-x64', + serverBuildId: 'test-relay-build' + }) + }) + + it('holds establish open until the consumer recovery write is durable', async () => { + const deps = createMockDeps() + let settleWrite!: () => void + let signalWriteStarted!: () => void + const writeStarted = new Promise((resolve) => { + signalWriteStarted = resolve + }) + vi.mocked(deps.mockStore.upsertSshPtyConsumerRecovery).mockImplementation(() => { + signalWriteStarted() + return new Promise((resolve) => { + settleWrite = resolve + }) + }) + + const session = new SshRelaySession( + 'durability-barrier-target', + deps.getMainWindow, + deps.mockStore, + deps.mockPortForward + ) + let established = false + const establishing = session.establish(deps.mockConn).then(() => { + established = true + }) + + await writeStarted + // Why a macrotask, not a microtask count: establish() has several awaits after the write starts, + // so only yielding past the whole microtask queue proves the write is the thing blocking it. + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(established).toBe(false) + + settleWrite() + await establishing + expect(established).toBe(true) + session.dispose() + }) + + it('keeps destructive disposal pending until consumer recovery is removed', async () => { + const { mockStore, mockPortForward, getMainWindow } = createMockDeps() + vi.mocked(mockStore.getSshPtyConsumerRecovery).mockReturnValue({ + targetId: 'target-disposal-durability', + clientInstanceId: 'client-disposal-durability', + serverBuildId: 'test-relay-build', + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'owner-lease' + }) + let settleRemoval!: () => void + vi.mocked(mockStore.removeSshPtyConsumerRecovery).mockImplementation( + () => + new Promise((resolve) => { + settleRemoval = resolve + }) + ) + const session = new SshRelaySession( + 'target-disposal-durability', + getMainWindow, + mockStore, + mockPortForward + ) + let completed = false + + const disposal = session.disposeAndPersist().then(() => { + completed = true + }) + await Promise.resolve() + + expect(session.getState()).toBe('disposed') + expect(completed).toBe(false) + expect(mockStore.markSshRemotePtyLeases).not.toHaveBeenCalled() + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith( + 'target-disposal-durability', + 'terminated' + ) + + settleRemoval() + await disposal + expect(completed).toBe(true) + }) + + it('does not retry a stale owner after disposal wins the recovery-removal race', async () => { + const targetId = 'target-stale-owner-disposal' + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + vi.mocked(mockStore.getSshPtyConsumerRecovery).mockReturnValue({ + targetId, + clientInstanceId: 'persisted-client', + serverBuildId: 'test-relay-build', + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'stale-owner' + }) + let signalRemovalStarted!: () => void + const removalStarted = new Promise((resolve) => { + signalRemovalStarted = resolve + }) + let settleRemoval!: () => void + vi.mocked(mockStore.removeSshPtyConsumerRecovery).mockImplementationOnce(() => { + signalRemovalStarted() + return new Promise((resolve) => { + settleRemoval = resolve + }) + }) + openConsumerSessionMock.mockRejectedValueOnce(createMismatchedOwnerRecoveryError()) + const session = new SshRelaySession(targetId, getMainWindow, mockStore, mockPortForward) + + const establishing = session.establish(mockConn) + const failed = expect(establishing).rejects.toThrow('Session disposed during establish') + await removalStarted + const disposal = session.disposeAndPersist() + settleRemoval() + await Promise.all([failed, disposal]) + + expect(openConsumerSessionMock).toHaveBeenCalledTimes(1) + }) + + it('does not remember a consumer opened after establish was disposed', async () => { + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + let signalOpenStarted!: () => void + const openStarted = new Promise((resolve) => { + signalOpenStarted = resolve + }) + let finishOpen!: (value: unknown) => void + openConsumerSessionMock.mockImplementationOnce(() => { + signalOpenStarted() + return new Promise((resolve) => { + finishOpen = resolve + }) + }) + const session = new SshRelaySession( + 'target-open-disposal', + getMainWindow, + mockStore, + mockPortForward + ) + + const establishing = session.establish(mockConn) + const failed = expect(establishing).rejects.toThrow('Session disposed during establish') + await openStarted + await session.disposeAndPersist() + finishOpen({ + mode: 'negotiated', + clientInstanceId: 'late-client', + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'late-owner' + }) + await failed + + expect(mockStore.upsertSshPtyConsumerRecovery).not.toHaveBeenCalled() + }) + + it('upgrades a pending detach to a full disposal', async () => { + const targetId = 'target-teardown-upgrade' + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + let settleDetachPersistence!: () => void + let settleDisposalPersistence!: () => void + vi.mocked(mockStore.markSshRemotePtyLeasesAsync) + .mockImplementationOnce( + () => + new Promise((resolve) => { + settleDetachPersistence = resolve + }) + ) + .mockImplementationOnce( + () => + new Promise((resolve) => { + settleDisposalPersistence = resolve + }) + ) + const session = new SshRelaySession(targetId, getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + + let detachCompleted = false + const detach = session.detachAndPersist().then(() => { + detachCompleted = true + }) + const disposal = session.disposeAndPersist() + + // Why: dispose supersedes the in-flight detach, so the destructive half must still run. + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith(targetId, 'terminated') + expect(getSshPtyConsumerRecovery(targetId)).toBeUndefined() + // Why: 'shutdown' teardown, not detach's 'connection_lost' — PTY ownership is released for good. + expect(clearPtyOwnershipForConnection).toHaveBeenCalledWith(targetId) + + settleDetachPersistence() + await Promise.resolve() + await Promise.resolve() + expect(detachCompleted).toBe(false) + settleDisposalPersistence() + await Promise.all([detach, disposal]) + + // Why: the reverse order is not an upgrade — a detach after disposal must not re-open ownership. + vi.mocked(mockStore.markSshRemotePtyLeasesAsync).mockClear() + await session.detachAndPersist() + expect(mockStore.markSshRemotePtyLeasesAsync).not.toHaveBeenCalled() + }) + + it('re-issues only the lease write after a rejected detach persistence', async () => { + const targetId = 'target-detach-write-retry' + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + vi.mocked(mockStore.markSshRemotePtyLeasesAsync).mockRejectedValueOnce( + new Error('lease write failed') + ) + const session = new SshRelaySession(targetId, getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + vi.mocked(clearPtyOwnershipForConnection).mockClear() + + await expect(session.detachAndPersist()).rejects.toThrow('lease write failed') + + const teardownCalls = vi.mocked(unregisterSshPtyProvider).mock.calls.length + vi.mocked(mockStore.markSshRemotePtyLeasesAsync).mockClear() + await session.detachAndPersist() + + // Why: the retry re-issues the write only — re-running provider teardown would tear down + // whatever a replacement session has already registered for this target. + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith(targetId, 'detached') + expect(unregisterSshPtyProvider).toHaveBeenCalledTimes(teardownCalls) + expect(clearPtyOwnershipForConnection).not.toHaveBeenCalled() + }) + + it('still upgrades to disposal after a rejected detach persistence', async () => { + const targetId = 'target-detach-write-failure-disposal' + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + vi.mocked(mockStore.markSshRemotePtyLeasesAsync).mockRejectedValueOnce( + new Error('lease write failed') + ) + const session = new SshRelaySession(targetId, getMainWindow, mockStore, mockPortForward) + await session.establish(mockConn) + + await expect(session.detachAndPersist()).rejects.toThrow('lease write failed') + await session.disposeAndPersist() + + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith(targetId, 'terminated') + expect(getSshPtyConsumerRecovery(targetId)).toBeUndefined() + }) +}) diff --git a/src/main/ssh/ssh-relay-session-recovery-races.test.ts b/src/main/ssh/ssh-relay-session-recovery-races.test.ts index 43354ebb90c..81039d901ec 100644 --- a/src/main/ssh/ssh-relay-session-recovery-races.test.ts +++ b/src/main/ssh/ssh-relay-session-recovery-races.test.ts @@ -279,11 +279,7 @@ describe('SshRelaySession recovery race fencing', () => { expect(recoveryActivationLease.retire).not.toHaveBeenCalled() expect(muxRequestMock).not.toHaveBeenCalledWith('pty.cancelDelivery', expect.anything()) expect(setPtyOwnership).not.toHaveBeenCalled() - expect(deps.mockStore.markSshRemotePtyLease).not.toHaveBeenCalledWith( - targetId, - 'pty-1', - 'attached' - ) + expect(deps.mockStore.markSshRemotePtyLeasesAttachedAsync).not.toHaveBeenCalled() }) it('settles exact cancellation before publishing an exit with incomplete recovery data', async () => { @@ -775,8 +771,10 @@ describe('SshRelaySession recovery race fencing', () => { expect(muxRequestMock.mock.calls.filter(([method]) => method === 'pty.cancelDelivery')).toEqual( [] ) - expect(deps.mockStore.markSshRemotePtyLease).toHaveBeenCalledTimes(1) - expect(deps.mockStore.markSshRemotePtyLease).toHaveBeenCalledWith(targetId, 'pty-1', 'attached') + expect(deps.mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledOnce() + expect(deps.mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith(targetId, [ + 'pty-1' + ]) expect(setPtyOwnership).toHaveBeenCalledTimes(1) expect(staleLease.transferToRecovery).toHaveBeenCalledOnce() expect(staleLease.commit).not.toHaveBeenCalled() 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 24d95fe18a5..a0f855543ef 100644 --- a/src/main/ssh/ssh-relay-session-terminal-error.test.ts +++ b/src/main/ssh/ssh-relay-session-terminal-error.test.ts @@ -105,7 +105,9 @@ function createMockDeps(): { removeSshPtyConsumerRecovery: vi.fn(), getSshRemotePtyLeases: vi.fn().mockReturnValue([]), markSshRemotePtyLease: vi.fn(), - markSshRemotePtyLeases: vi.fn() + markSshRemotePtyLeases: vi.fn(), + markSshRemotePtyLeasesAsync: vi.fn(), + markSshRemotePtyLeasesAttachedAsync: 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 41c4bca32d4..b0ed98eaf2e 100644 --- a/src/main/ssh/ssh-relay-session-test-fixtures.ts +++ b/src/main/ssh/ssh-relay-session-test-fixtures.ts @@ -24,6 +24,8 @@ export function createMockDeps(): SshRelaySessionTestDeps { getSshRemotePtyLeases: vi.fn().mockReturnValue([]), markSshRemotePtyLease: vi.fn(), markSshRemotePtyLeases: vi.fn(), + markSshRemotePtyLeasesAsync: vi.fn(), + markSshRemotePtyLeasesAttachedAsync: vi.fn(), persistPtyBinding: vi.fn() } as unknown as Store const mockPortForward = { diff --git a/src/main/ssh/ssh-relay-session.test.ts b/src/main/ssh/ssh-relay-session.test.ts index ad563d002b5..eb81cda4cb1 100644 --- a/src/main/ssh/ssh-relay-session.test.ts +++ b/src/main/ssh/ssh-relay-session.test.ts @@ -497,7 +497,9 @@ describe('SshRelaySession', () => { expect(mockAttach).toHaveBeenCalledWith('pty-1') expect(setPtyOwnership).toHaveBeenCalledWith('ssh:target-1@@pty-1', 'target-1') - expect(mockStore.markSshRemotePtyLease).toHaveBeenCalledWith('target-1', 'pty-1', 'attached') + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith('target-1', [ + 'pty-1' + ]) }) it('establish re-attaches durable leases after app restart', async () => { @@ -511,6 +513,7 @@ describe('SshRelaySession', () => { vi.mocked(getPtyIdsForConnection).mockReturnValue([]) vi.mocked(mockStore.getSshRemotePtyLeases).mockReturnValue([ { targetId: 'target-1', ptyId: 'pty-live', state: 'detached' }, + { targetId: 'target-1', ptyId: 'pty-live-2', state: 'detached' }, { targetId: 'target-1', ptyId: 'pty-expired', state: 'expired' } ] as ReturnType) @@ -519,9 +522,14 @@ describe('SshRelaySession', () => { await session.establish(mockConn) expect(mockAttach).toHaveBeenCalledWith('pty-live') + expect(mockAttach).toHaveBeenCalledWith('pty-live-2') expect(mockAttach).not.toHaveBeenCalledWith('pty-expired') expect(setPtyOwnership).toHaveBeenCalledWith('ssh:target-1@@pty-live', 'target-1') - expect(mockStore.markSshRemotePtyLease).toHaveBeenCalledWith('target-1', 'pty-live', 'attached') + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledOnce() + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith( + 'target-1', + expect.arrayContaining(['pty-live', 'pty-live-2']) + ) }) it('forwards a lease tab identity to reattach so a reset relay cannot cross-wire it', async () => { @@ -631,11 +639,7 @@ describe('SshRelaySession', () => { await expect(establish).rejects.toThrow('Session disposed during establish') expect(setPtyOwnership).not.toHaveBeenCalledWith('pty-1', 'target-1') - expect(mockStore.markSshRemotePtyLease).not.toHaveBeenCalledWith( - 'target-1', - 'pty-1', - 'attached' - ) + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).not.toHaveBeenCalled() }) it('does not mark PTYs attached if detach wins while reattach is in flight', async () => { @@ -665,11 +669,7 @@ describe('SshRelaySession', () => { await reconnect expect(setPtyOwnership).not.toHaveBeenCalledWith('pty-1', 'target-1') - expect(mockStore.markSshRemotePtyLease).not.toHaveBeenCalledWith( - 'target-1', - 'pty-1', - 'attached' - ) + expect(mockStore.markSshRemotePtyLeasesAttachedAsync).not.toHaveBeenCalled() }) it('invalidates and broadcasts remote PTYs that cannot reattach after relay reconnect', async () => { diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index 129ef53960d..d413f45f044 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -233,6 +233,27 @@ function normalizeRelayGracePeriodSeconds(graceTimeSeconds: number | undefined): ) } +// Why: teardown barriers are independent, so one failing store write must not hide the others — +// settle them all and aggregate, rather than rethrowing only whichever rejected first. +async function settleSshSessionTeardown( + barriers: (Promise | null | undefined)[] +): Promise { + const results = await Promise.allSettled(barriers.map((barrier) => barrier ?? Promise.resolve())) + const errors = results.flatMap((result) => + result.status === 'rejected' ? [result.reason as unknown] : [] + ) + if (errors.length === 1) { + throw errors[0] + } + if (errors.length > 1) { + throw new AggregateError(errors, 'SSH relay session teardown failed') + } +} + +// Why: dispose is strictly more destructive than detach, so the mode records which teardown a +// session has already committed to and lets dispose supersede an in-flight detach. +type SshRelaySessionTeardownMode = 'detach' | 'dispose' + export class SshRelaySession { private _state: RelaySessionState = 'idle' private mux: SshChannelMultiplexer | null = null @@ -256,6 +277,12 @@ export class SshRelaySession { private activePtyProviderGeneration: number | null = null private sourceAckPublisherCleanup: (() => void) | null = null private sourceCancellationPublisherCleanup: (() => void) | null = null + private teardownMode: SshRelaySessionTeardownMode | null = null + private teardownCompletion: Promise | null = null + // Why: detach's in-memory half is one-shot but its lease write is retryable, so they are tracked + // apart — a rejected write can be re-issued without re-running provider teardown. + private detachedInMemory = false + private detachFlushRejected = false private ptyRecoveryNotificationCleanups: (() => void)[] = [] private readonly sourceIdentityByRelayPtyId = new Map< string, @@ -405,10 +432,27 @@ export class SshRelaySession { const mux = new SshChannelMultiplexer(transport) this.mux = mux - const ownsAttempt = (): boolean => this.mux === mux && !this.isDisposed() + const ownsAttempt = (): boolean => this.mux === mux && !mux.isDisposed() && !this.isDisposed() - this.ptyConsumerSessionState = await this.openPtyConsumerSession(mux, serverBuildId) - this.rememberPtyConsumerRecovery(serverBuildId) + const ptyConsumerSessionState = await this.openPtyConsumerSession( + mux, + serverBuildId, + ownsAttempt + ) + if (!ownsAttempt()) { + if (!mux.isDisposed()) { + mux.dispose() + } + throw new Error('Session disposed during establish') + } + this.ptyConsumerSessionState = ptyConsumerSessionState + await this.rememberPtyConsumerRecovery(serverBuildId) + if (!ownsAttempt()) { + if (!mux.isDisposed()) { + mux.dispose() + } + throw new Error('Session disposed during establish') + } await mux.request('session.resolveHome', { path: '~' }) if (!ownsAttempt()) { @@ -529,12 +573,25 @@ export class SshRelaySession { this.mux = mux const ownsAttempt = (): boolean => + this.mux === mux && + !mux.isDisposed() && this.abortController === abortController && !abortController.signal.aborted && !this.isDisposed() - this.ptyConsumerSessionState = await this.openPtyConsumerSession(mux, serverBuildId) - this.rememberPtyConsumerRecovery(serverBuildId) + const ptyConsumerSessionState = await this.openPtyConsumerSession( + mux, + serverBuildId, + ownsAttempt + ) + if (!ownsAttempt()) { + if (!mux.isDisposed()) { + mux.dispose() + } + return + } + this.ptyConsumerSessionState = ptyConsumerSessionState + await this.rememberPtyConsumerRecovery(serverBuildId) if (!ownsAttempt()) { if (!mux.isDisposed()) { mux.dispose() @@ -620,35 +677,126 @@ export class SshRelaySession { } } + /** Fire-and-forget disposal; prefer {@link disposeAndPersist} when the caller can await durability. */ dispose(): void { - if (this._state === 'disposed') { - return - } - this.abortController?.abort() - this.stopPortScanning() - // Why: fire-and-forget — nothing rebinds after dispose, so no need to await port release. - void this.portForwardManager.removeAllForwards(this.targetId) - this.broadcastEmptyLists() - this.teardownProviders('shutdown') - this.store.markSshRemotePtyLeases(this.targetId, 'terminated') - this.currentConnection = null - this._state = 'disposed' - forgetSshPtyConsumerRecovery(this.targetId, this.ptyConsumerClientInstanceId, this.store) + void this.disposeAndPersist().catch((error) => { + console.warn( + `[ssh-relay-session] Failed to persist disposal for ${this.targetId}: ${error instanceof Error ? error.message : String(error)}` + ) + }) } - detach(): void { - if (this._state === 'disposed') { - return + /** + * Destructive teardown: forgets the consumer recovery record and terminates the leases. + * Repeat calls share one completion, and a dispose requested after a detach supersedes it — + * the destructive half re-runs so detach-then-dispose still forgets recovery and terminates + * leases. (The reverse, detach after dispose, is a no-op.) + */ + disposeAndPersist(): Promise { + if (this.teardownMode === 'dispose') { + return this.teardownCompletion ?? Promise.resolve() } + const pendingDetach = this.teardownCompletion + this.teardownMode = 'dispose' + try { + this.teardownCompletion = this.runDisposal(pendingDetach) + } catch (error) { + // Why: a synchronous failure in the in-memory half must ride the completion promise, or a later + // disposeAndPersist reports success on a null completion and detachAndPersist re-runs runDetach. + this.teardownCompletion = Promise.reject(error) + } + return this.teardownCompletion + } + + private runDisposal(pendingDetach: Promise | null): Promise { + // Why: the whole in-memory half runs before any await so a concurrent connect can never observe + // a half-torn session; only the durability barriers below are deferred onto the returned promise. this.abortController?.abort() this.stopPortScanning() this.broadcastEmptyLists() - // Why: window disconnect is non-destructive — unregister local providers but keep PTY ownership so reattach works (relay owns the grace timer). - this.teardownProviders('connection_lost') - this.store.markSshRemotePtyLeases(this.targetId, 'detached') + this.teardownProviders('shutdown') this.currentConnection = null this._state = 'disposed' - detachSshPtyConsumerRecovery(this.targetId, this.ptyConsumerClientInstanceId) + const recoveryRemoval = forgetSshPtyConsumerRecovery( + this.targetId, + this.ptyConsumerClientInstanceId, + this.store + ) + const leaseTermination = this.store.markSshRemotePtyLeasesAsync(this.targetId, 'terminated') + return settleSshSessionTeardown([ + // Why: a superseded detach keeps its own rejection for its own caller; swallow it here so a + // failed detach write cannot fail the disposal that replaced it. + pendingDetach?.catch(() => undefined), + // Why: nothing rebinds after dispose, but direct callers (no IPC-side teardown) still need + // this session's local listeners released. + this.portForwardManager.removeAllForwards(this.targetId), + recoveryRemoval, + leaseTermination + ]) + } + + /** Fire-and-forget detach; prefer {@link detachAndPersist} when the caller can await durability. */ + detach(): void { + void this.detachAndPersist().catch((error) => { + console.warn( + `[ssh-relay-session] Failed to persist detach for ${this.targetId}: ${error instanceof Error ? error.message : String(error)}` + ) + }) + } + + /** + * Non-destructive teardown: keeps PTY ownership so a later connect reclaims this consumer + * identity. Once dispose has been requested this is a no-op — see {@link disposeAndPersist}. + */ + async detachAndPersist(): Promise { + // Why: a rejected lease write must not latch forever — re-issue just the write so the caller can + // retry. Never under 'dispose': that supersedes detach for good and re-issuing would resurrect + // the 'detached' state the disposal already replaced with 'terminated'. + if (!this.teardownCompletion || (this.teardownMode === 'detach' && this.detachFlushRejected)) { + this.teardownCompletion = this.runDetach() + } + this.teardownMode ??= 'detach' + let completion = this.teardownCompletion + while (completion) { + try { + await completion + } catch (error) { + if (completion === this.teardownCompletion) { + throw error + } + } + if (completion === this.teardownCompletion) { + return + } + completion = this.teardownCompletion + } + } + + private runDetach(): Promise { + if (!this.detachedInMemory) { + if (this._state === 'disposed') { + return Promise.resolve() + } + // Why first: same synchronous-half-first rule as runDisposal, and this is the highest-value + // step — a fast reconnect must reclaim this identity instead of minting one, even if a + // teardown call below throws unexpectedly. + detachSshPtyConsumerRecovery(this.targetId, this.ptyConsumerClientInstanceId) + this.abortController?.abort() + this.stopPortScanning() + this.broadcastEmptyLists() + // Why: disconnect keeps PTY ownership so a later manual connect can reattach. + this.teardownProviders('connection_lost') + this.currentConnection = null + this._state = 'disposed' + this.detachedInMemory = true + } + this.detachFlushRejected = false + return settleSshSessionTeardown([ + this.store.markSshRemotePtyLeasesAsync(this.targetId, 'detached') + ]).catch((error: unknown) => { + this.detachFlushRejected = true + throw error + }) } // ── Private ─────────────────────────────────────────────────────── @@ -832,7 +980,8 @@ export class SshRelaySession { private async openPtyConsumerSession( mux: SshChannelMultiplexer, - serverBuildId: string | undefined + serverBuildId: string | undefined, + ownsAttempt: () => boolean ): Promise { const previousOwner = this.recoverablePtyConsumerOwner(serverBuildId) const options = { @@ -876,18 +1025,25 @@ export class SshRelaySession { ) } } - removeSshPtyConsumerOwnerRecovery(this.targetId, this.ptyConsumerClientInstanceId, this.store) + await removeSshPtyConsumerOwnerRecovery( + this.targetId, + this.ptyConsumerClientInstanceId, + this.store + ) + if (!ownsAttempt()) { + throw new Error('Session disposed during establish') + } this.ptyConsumerSessionState = null return openSshPtyConsumerSession(mux, options) } } - private rememberPtyConsumerRecovery(serverBuildId: string | undefined): void { + private async rememberPtyConsumerRecovery(serverBuildId: string | undefined): Promise { const owner = this.activePtyConsumerOwner() if (!owner || !serverBuildId) { return } - rememberSshPtyConsumerRecovery({ + await rememberSshPtyConsumerRecovery({ targetId: this.targetId, clientInstanceId: this.ptyConsumerClientInstanceId, serverBuildId, @@ -1707,6 +1863,7 @@ export class SshRelaySession { }) .filter((entry): entry is [string, ExpectedPtyIdentity] => entry !== null) ) + const attachedLeaseIds = new Set() // Why: after app restart ptyOwnership is empty, but durable SSH leases still describe grace-window survivors. const ptyIds = Array.from( new Set([ @@ -1734,6 +1891,7 @@ export class SshRelaySession { ptyId, activeLeaseByPtyId, expectedIdentityByPtyId, + attachedLeaseIds, mux, providerGeneration, shouldContinue @@ -1753,6 +1911,12 @@ export class SshRelaySession { await Promise.all( Array.from({ length: Math.min(SSH_PTY_REATTACH_MAX_CONCURRENCY, ptyIds.length) }, worker) ) + if (attachedLeaseIds.size > 0 && shouldContinue()) { + await this.store.markSshRemotePtyLeasesAttachedAsync( + this.targetId, + Array.from(attachedLeaseIds) + ) + } } private async reattachKnownPty(args: { @@ -1760,6 +1924,7 @@ export class SshRelaySession { ptyId: string activeLeaseByPtyId: Map expectedIdentityByPtyId: Map + attachedLeaseIds: Set mux: SshChannelMultiplexer providerGeneration: number shouldContinue: () => boolean @@ -1769,6 +1934,7 @@ export class SshRelaySession { ptyId, activeLeaseByPtyId, expectedIdentityByPtyId, + attachedLeaseIds, mux, providerGeneration, shouldContinue @@ -1884,7 +2050,7 @@ export class SshRelaySession { activeLeaseByPtyId.get(ptyId) ) } - this.store.markSshRemotePtyLease(this.targetId, ptyId, 'attached') + attachedLeaseIds.add(ptyId) pendingReattach.activated = true recoveryActivationLease?.commit() recoveryActivationLease = undefined