diff --git a/src/main/index.ts b/src/main/index.ts index aa3c2b7d493..a40b5ee508b 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -250,7 +250,8 @@ import { buildHeadlessAutomationWorktreeCreateArgs } from './automations/headles import { AgentAwakeService } from './agent-awake-service' import { registerSystemResumeBroadcast } from './system-resume-broadcast' import { settleTeardownWithinDeadline } from './quit-teardown-deadline' -import { QuitTeardownStartGate } from './quit-teardown-start-gate' +import { quitTeardownStartGate } from './quit-teardown-start-gate' +import { detachAllSshSessionsForShutdown } from './ipc/ssh' import { PluginService } from './plugins/plugin-service' import { PluginKillListService } from './plugins/plugin-kill-list-service' import { getPluginsDataDir } from './plugins/plugin-discovery' @@ -2940,7 +2941,6 @@ app.on('before-quit', () => { // Why: will-quit fires twice — first pass preventDefaults and runs teardown; second pass exits. let daemonDisconnectDone = false -const quitTeardownStartGate = new QuitTeardownStartGate() app.on('will-quit', (e) => { // Why return instead of re-running teardown: the second pass is Electron re-firing after // our own app.quit(), so every step below already ran and every durable write already @@ -2992,6 +2992,7 @@ app.on('will-quit', (e) => { runtime?.getOffscreenBrowserBackend()?.destroyAll?.() browserManager.setBrowserGuestStateChangedListener(null) const emulatorShutdown = runtime?.getEmulatorBridge()?.destroyAllSessions() ?? Promise.resolve() + const sshShutdown = detachAllSshSessionsForShutdown() killAllPty() const watcherShutdown = shutdownWatchersOnce() const storeFlush = store?.flushAsync() ?? Promise.resolve() @@ -3034,6 +3035,7 @@ app.on('will-quit', (e) => { { name: 'runtime-rpc', promise: rpcStopAndClear }, { name: 'watchers', promise: watcherShutdown }, { name: 'emulator', promise: emulatorShutdown }, + { name: 'ssh', promise: sshShutdown }, { name: 'plugin-hosts', promise: pluginHostShutdown }, { name: 'usage-cache', promise: usageCacheFlush }, { name: 'stats', promise: statsFlush }, diff --git a/src/main/ipc/ssh.test.ts b/src/main/ipc/ssh.test.ts index c7f850452c5..3191960409f 100644 --- a/src/main/ipc/ssh.test.ts +++ b/src/main/ipc/ssh.test.ts @@ -42,6 +42,7 @@ const { mockConnectionManager: { connect: vi.fn(), disconnect: vi.fn(), + disconnectConnection: vi.fn(), reconnect: vi.fn(), getConnection: vi.fn(), getState: vi.fn(), @@ -241,7 +242,13 @@ vi.mock('../ssh/ssh-port-scanner', () => ({ } })) -import { getSshConnectionManager, registerSshHandlers, resetSshHandlerStateForTests } from './ssh' +import { + detachAllSshSessionsForShutdown, + getActiveMultiplexer, + getSshConnectionManager, + registerSshHandlers, + resetSshHandlerStateForTests +} from './ssh' import { RelayVersionMismatchError } from '../ssh/ssh-relay-version-mismatch-error' import { SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, @@ -260,6 +267,7 @@ import { } from './pty' import { assertSshMutationExpectation } from '../ssh/ssh-connection-generation' import { getSshPtyConsumerRecovery } from '../ssh/ssh-pty-consumer-recovery' +import { quitTeardownStartGate } from '../quit-teardown-start-gate' describe('SSH IPC handlers', () => { const relayBuildId = '0.1.0+ipc-test' @@ -296,6 +304,7 @@ describe('SSH IPC handlers', () => { const createConnectionManagerMock = () => ({ connect: vi.fn(), disconnect: vi.fn(), + disconnectConnection: vi.fn(), reconnect: vi.fn(), getConnection: vi.fn(), getState: vi.fn(), @@ -365,6 +374,7 @@ describe('SSH IPC handlers', () => { mockConnectionManager.connect.mockReset() mockConnectionManager.disconnect.mockReset() + mockConnectionManager.disconnectConnection.mockReset().mockResolvedValue(undefined) mockConnectionManager.reconnect.mockReset() mockConnectionManager.getConnection.mockReset() mockConnectionManager.getState.mockReset() @@ -539,6 +549,145 @@ describe('SSH IPC handlers', () => { expect(mockSshStore.removeTarget).toHaveBeenCalledWith('ssh-1') }) + it('detaches active SSH sessions during app shutdown without terminating recovery', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.connect.mockResolvedValue({}) + mockConnectionManager.getState.mockReturnValue({ + targetId: 'ssh-1', + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' }) + mockMux.dispose.mockClear() + mockConnectionManager.disconnectAll.mockClear().mockResolvedValue(undefined) + + await detachAllSshSessionsForShutdown() + + expect(mockPortForwardManager.removeAllForwards).toHaveBeenCalledWith('ssh-1') + expect(mockMux.dispose).toHaveBeenCalledWith('connection_lost') + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'detached') + expect(mockStore.removeSshPtyConsumerRecovery).not.toHaveBeenCalled() + expect(mockConnectionManager.disconnectAll).toHaveBeenCalled() + }) + + it('refuses the replacement session a paused connect would publish after shutdown began', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.connect.mockResolvedValue({}) + mockConnectionManager.getState.mockReturnValue({ + targetId: 'ssh-1', + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + mockConnectionManager.disconnectAll.mockResolvedValue(undefined) + await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' }) + + // Why here: the old session's lease flush is the one window where doConnect resumes straight into + // publishing its replacement session without re-checking authority. + let releaseDetach = (): void => {} + const detachFlush = new Promise((resolve) => { + releaseDetach = resolve + }) + let signalEnteredDetach = (): void => {} + const enteredDetach = new Promise((resolve) => { + signalEnteredDetach = resolve + }) + mockStore.markSshRemotePtyLeasesAsync.mockImplementationOnce(() => { + signalEnteredDetach() + return detachFlush + }) + + const replacement = handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' }) + // Why await the flush entry and not ticks: shutdown must not snapshot what to drain until the + // replacement connect is parked in the lease flush. + await enteredDetach + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'detached') + + mockConnectionManager.disconnectAll.mockClear() + // Why latch the gate here: the committed quit path owns it, so the drain alone must not fence. + quitTeardownStartGate.tryStart({ preventDefault() {} }) + const shutdown = detachAllSshSessionsForShutdown() + releaseDetach() + + await expect(replacement).rejects.toThrow('closed for app shutdown') + await shutdown + + // Why one detach: the fence sits at the publication point, so the resumed connect never registers a + // replacement session — only the old session it had already torn down was detached. + const detaches = mockStore.markSshRemotePtyLeasesAsync.mock.calls.filter( + (call) => call[1] === 'detached' + ) + expect(detaches).toHaveLength(1) + // Why twice: once for the drain's snapshot, once after joining the connect that was still in flight. + expect(mockConnectionManager.disconnectAll).toHaveBeenCalledTimes(2) + expect(getActiveMultiplexer('ssh-1')).toBeUndefined() + await expect(handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })).rejects.toThrow( + 'closed for app shutdown' + ) + }) + + it('joins an in-flight test-connection probe before the final shutdown disconnect', async () => { + const target: SshTarget = { + id: 'ssh-probe', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getState.mockReturnValue({ + targetId: 'ssh-probe', + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + mockConnectionManager.disconnectAll.mockResolvedValue(undefined) + mockConnectionManager.disconnect.mockClear().mockResolvedValue(undefined) + + // Why: a probe holds a transport no session owns, so shutdown has to wait for it to hand it back. + const probeState = { + targetId: 'ssh-probe', + status: 'connected', + error: null, + reconnectAttempt: 0 + } + let openProbeTransport = (): void => {} + mockConnectionManager.connect.mockImplementationOnce( + () => + new Promise((resolve) => { + openProbeTransport = () => resolve({ getState: () => probeState }) + }) + ) + const probe = handlers.get('ssh:testConnection')!(null, { targetId: 'ssh-probe' }) + for (let tick = 0; tick < 5; tick++) { + await Promise.resolve() + } + expect(mockConnectionManager.connect).toHaveBeenCalledWith(target) + + quitTeardownStartGate.tryStart({ preventDefault() {} }) + const shutdown = detachAllSshSessionsForShutdown() + openProbeTransport() + await shutdown + + expect(await probe).toMatchObject({ success: true }) + expect(mockConnectionManager.disconnect).toHaveBeenCalledWith('ssh-probe') + }) + it('ssh:importConfig returns imported targets', async () => { const imported: SshTarget[] = [ { id: 'ssh-imp', label: 'staging', host: 'staging.com', port: 22, username: '' } @@ -1839,6 +1988,106 @@ describe('SSH IPC handlers', () => { expect(mockDeployAndLaunchRelay).toHaveBeenCalledTimes(1) }) + it('closes the transport a cancelled connect opened after the disconnect finished', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + const lateConn = { id: 'late-transport' } + let resolveStaleConnect!: (connection: unknown) => void + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getConnection.mockReturnValue(undefined) + mockConnectionManager.connect.mockReturnValueOnce( + new Promise((resolve) => { + resolveStaleConnect = resolve + }) + ) + mockConnectionManager.disconnect.mockResolvedValue(undefined) + + const staleConnect = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1)) + // Why await the whole disconnect: the leak only exists once its teardown has already run, so + // nothing else is left to close the transport this attempt opens afterwards. + await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' }) + + resolveStaleConnect(lateConn) + + await expect(staleConnect).rejects.toThrow('SSH connection attempt was cancelled') + expect(mockConnectionManager.disconnectConnection).toHaveBeenCalledWith('ssh-1', lateConn) + expect(getActiveMultiplexer('ssh-1')).toBeUndefined() + }) + + it('closes the transport when establish resumes after the connect was invalidated', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + const conn = { id: 'establishing-transport' } + let releaseRelayLaunch = (): void => {} + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getConnection.mockReturnValue(undefined) + mockConnectionManager.connect.mockResolvedValue(conn) + mockConnectionManager.disconnect.mockResolvedValue(undefined) + mockDeployAndLaunchRelay.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseRelayLaunch = () => resolve(createRelayLaunchResult()) + }) + ) + + const connect = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockDeployAndLaunchRelay).toHaveBeenCalledTimes(1)) + await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' }) + + releaseRelayLaunch() + + await expect(connect).rejects.toThrow('SSH connection attempt was cancelled') + expect(mockConnectionManager.disconnectConnection).toHaveBeenCalledWith('ssh-1', conn) + expect(getActiveMultiplexer('ssh-1')).toBeUndefined() + }) + + it('leaves a reused transport to its replacement when the connect is cancelled', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + // Why: connect() hands back the already-open transport, so this attempt never owned it. + const sharedConn = { id: 'shared-transport' } + let resolveStaleConnect!: (connection: unknown) => void + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getConnection.mockReturnValue(sharedConn) + mockConnectionManager.connect.mockReturnValueOnce( + new Promise((resolve) => { + resolveStaleConnect = resolve + }) + ) + mockConnectionManager.disconnect.mockResolvedValue(undefined) + + const staleConnect = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1)) + await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' }) + + resolveStaleConnect(sharedConn) + + await expect(staleConnect).rejects.toThrow('SSH connection attempt was cancelled') + expect(mockConnectionManager.disconnectConnection).not.toHaveBeenCalled() + }) + it('keeps reconnect behind transport disconnect when forward teardown fails', async () => { const target: SshTarget = { id: 'ssh-1', @@ -2327,6 +2576,51 @@ describe('SSH IPC handlers', () => { expect(mockConnectionManager.disconnect).toHaveBeenCalledWith('ssh-1') }) + it('ssh:resetRelay does not open a transport when shutdown starts while it waits for a connect', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + let failConnect!: (error: unknown) => void + const connectResult = new Promise((_resolve, reject) => { + failConnect = reject + }) + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.connect.mockReturnValue(connectResult) + // Why undefined: reset must fall through to opening its own transport, which is the call under test. + mockConnectionManager.getConnection.mockReturnValue(undefined) + mockConnectionManager.getState.mockReturnValue({ + targetId: 'ssh-1', + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + + const connectPromise = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1)) + + // Why admitted first: the gate latches only after reset is already parked behind the connect, so + // the entry fence cannot be what stops it. + const resetPromise = handlers.get('ssh:resetRelay')!(null, { + targetId: 'ssh-1' + }) as Promise + await Promise.resolve() + + quitTeardownStartGate.tryStart({ preventDefault() {} }) + failConnect(new Error('transport lost')) + await expect(connectPromise).rejects.toThrow('transport lost') + + await expect(resetPromise).rejects.toThrow('closed for app shutdown') + // Why once: the resumed reset must not open a second transport that would outlive the drain. + expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1) + expect(mockForceStopRelayForTarget).not.toHaveBeenCalled() + }) + it('ssh:connect waits for an in-flight reset before starting a new connection', async () => { const target: SshTarget = { id: 'ssh-1', diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index 995d3e88b72..3fc1d732bb3 100644 --- a/src/main/ipc/ssh.ts +++ b/src/main/ipc/ssh.ts @@ -3,7 +3,7 @@ import { ipcMain, powerMonitor, type BrowserWindow } from 'electron' import { appendFileSync } from 'node:fs' import type { Store } from '../persistence' import { SshConnectionStore } from '../ssh/ssh-connection-store' -import type { SshConnectionCallbacks } from '../ssh/ssh-connection' +import type { SshConnection, SshConnectionCallbacks } from '../ssh/ssh-connection' import { SshConnectionManager } from '../ssh/ssh-connection-manager' import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer' import { SshRelaySession, type SshRelayAiVaultHostInfo } from '../ssh/ssh-relay-session' @@ -19,6 +19,7 @@ import type { DirectSshAuthority } from '../../shared/ssh-types' import { SSH_TERMINATE_RECONNECT_REQUIRED } from '../../shared/constants' +import { quitTeardownStartGate } from '../quit-teardown-start-gate' import { isRuntimeOwnedSshTargetId } from '../../shared/execution-host' import { isAuthError } from '../ssh/ssh-connection-utils' import { createCancelledConnectAttemptError } from '../ssh/ssh-connect-attempt-cancellation' @@ -269,6 +270,33 @@ async function abandonFailedSshSession(targetId: string, session: SshRelaySessio } } +// Why: a connect cancelled after its transport already opened still owns that transport and its +// unpublished session — nothing else will reach them, so it has to close them itself. Why only the +// connection it minted, and by identity: disconnecting by target id (or closing a transport it merely +// reused) would tear down the replacement connect's live transport. +async function abandonCancelledConnectAttempt( + targetId: string, + session: SshRelaySession, + mintedConnection: SshConnection | null +): Promise { + // Why the identity guard: every path that removes a session from activeSessions detaches it first, + // so re-detaching a superseded session would only clobber the replacement's lease state. + if (activeSessions.get(targetId) === session) { + await abandonFailedSshSession(targetId, session) + } + if (!mintedConnection) { + return + } + try { + await connectionManager!.disconnectConnection(targetId, mintedConnection) + } catch (error) { + // Why: the caller is about to throw the cancellation; a teardown throw must not replace it. + console.warn( + `[ssh] Failed to disconnect cancelled connect transport for ${targetId}: ${error instanceof Error ? error.message : String(error)}` + ) + } +} + function relayGracePeriodForTarget(target: SshTarget | null | undefined): number | undefined { return target?.relayGracePeriodSeconds } @@ -283,6 +311,15 @@ type ConnectAttempt = { const connectInFlight = new Map() const pendingTransportReconnects = new Set() +// Why the quit gate rather than a local latch: "the committed quit has begun" already has an owner, +// and a private copy could be set by something that is not actually quitting — leaving SSH connects +// refused for the rest of the process lifetime. +function assertSshConnectsNotFenced(): void { + if (quitTeardownStartGate.hasStarted()) { + throw new Error('SSH connects are closed for app shutdown') + } +} + function invalidateConnectAttempt(targetId: string): void { rotateSshProviderAuthority(targetId) pendingTransportReconnects.delete(targetId) @@ -299,6 +336,7 @@ const resetRelayInFlight = new Map>() // Why: ssh:testConnection connects then disconnects; suppressing broadcasts during the test avoids worktree cards flashing connected → disconnected. const testingTargets = new Set() +const testConnectionProbes = new Set>() // Why: without backoff, a relay channel that keeps dying reconnects as fast as the network allows, hammering local + remote sshd; track attempts and back off to end the loop recoverably. type RelayLostBackoffState = { @@ -1004,6 +1042,9 @@ export function registerSshHandlers( if (!isCurrentSshProviderAuthority(observedAuthority)) { throw createCancelledConnectAttemptError() } + // Why: the shutdown drain fences and snapshots synchronously, so a connect either registers in + // connectInFlight below (and gets joined) or fails here — it can never slip between the two. + assertSshConnectsNotFenced() pendingTransportReconnects.delete(targetId) const promise = doConnect(targetId, replacePendingTransport) @@ -1092,6 +1133,11 @@ export function registerSshHandlers( } } + // Why here and not only at entry: this is the publication point, and it is the last statement + // before the transport opens. Checking it in the same synchronous block as activeSessions.set + // means a connect either registers before the shutdown drain snapshots, or registers never and + // owns nothing to clean up. + assertSshConnectsNotFenced() // Why: create the session early so onStateChange sees it in 'deploying' and skips reconnect logic. const session = new SshRelaySession( targetId, @@ -1106,6 +1152,12 @@ export function registerSshHandlers( const ownsSession = (): boolean => isCurrentConnectAttempt(targetId, authority) && activeSessions.get(targetId) === session + // Why captured here and not with existingState: connect() reuses an already-connected transport, + // and only a transport this attempt opened is this attempt's to close when it loses the race. + const priorConnection = connectionManager!.getConnection(targetId) + const mintedConnection = (): SshConnection | null => + conn && conn !== priorConnection ? conn : null + try { conn = await connectionManager!.connect(target) if (!ownsSession()) { @@ -1116,6 +1168,7 @@ export function registerSshHandlers( const errObj = err instanceof Error ? err : new Error(String(err)) const status: SshConnectionStatus = isAuthError(errObj) ? 'auth-failed' : 'error' if (!ownsSession()) { + await abandonCancelledConnectAttempt(targetId, session, mintedConnection()) throw createCancelledConnectAttemptError() } // Why: clear this failed connect's flag so a later non-prompting connect isn't deferred. @@ -1156,6 +1209,7 @@ export function registerSshHandlers( }) } catch (err) { if (!ownsSession()) { + await abandonCancelledConnectAttempt(targetId, session, mintedConnection()) throw createCancelledConnectAttemptError() } await abandonFailedSshSession(targetId, session) @@ -1262,7 +1316,13 @@ export function registerSshHandlers( } const existingConn = connectionManager!.getConnection(targetId) - const conn = existingConn ?? (await connectionManager!.connect(target)) + let conn = existingConn + if (!conn) { + // Why re-check: admission fenced this reset before it parked on the in-flight connect, so shutdown + // may have started (and drained) while we waited — opening a transport now would outlive the drain. + assertSshConnectsNotFenced() + conn = await connectionManager!.connect(target) + } try { await forceStopRelayForTarget(conn, targetId) } finally { @@ -1295,6 +1355,8 @@ export function registerSshHandlers( if (!target) { throw new Error(`SSH target "${args.targetId}" not found`) } + // Why: reset opens its own transport, so it must be fenced by shutdown the same way connect is. + assertSshConnectsNotFenced() let resetPromise: Promise resetPromise = runTargetLifecycle(args.targetId, () => @@ -1354,17 +1416,26 @@ export function registerSshHandlers( } testingTargets.add(args.targetId) - try { + // Why a tracked promise and not just the id: a probe holds a real transport that no session owns, + // so shutdown has to be able to join it before the final drain disconnects what is left. + const probe = (async () => { + // Why: a probe transport opened after the shutdown drain would outlive orderly teardown. + assertSshConnectsNotFenced() const conn = await connectionManager!.connect(target) const state = conn.getState() await connectionManager!.disconnect(args.targetId) - return { success: true, state } + return state + })() + testConnectionProbes.add(probe) + try { + return { success: true, state: await probe } } catch (err) { return { success: false, error: err instanceof Error ? err.message : String(err) } } finally { + testConnectionProbes.delete(probe) testingTargets.delete(args.targetId) // Why: clear so a test's credential prompt doesn't leave lastRequiredPassphrase=true and defer this target at startup. credentialRequestedForTarget.delete(args.targetId) @@ -1472,6 +1543,61 @@ export function getSshConnectionManager(): SshConnectionManager | null { return connectionManager } +// Why: an invalidated connect only observes its cancellation at the next checkpoint, and one blocked +// in the transport handshake can sit there for the whole SSH timeout. Bound the join so the final +// drain always runs well inside the quit deadline instead of racing it. +const SSH_SHUTDOWN_INFLIGHT_JOIN_MS = 2_000 + +async function settleWithinMs(work: readonly Promise[], timeoutMs: number): Promise { + if (work.length === 0) { + return + } + await Promise.race([ + Promise.allSettled(work), + new Promise((resolve) => { + const timer = setTimeout(resolve, timeoutMs) + timer.unref?.() + }) + ]) +} + +// Why no fence latch here: the committed quit path sets it before calling this, so it is already on +// for the snapshot below. Called without that gate (tests), this degrades to a plain drain. +export async function detachAllSshSessionsForShutdown(): Promise { + const inFlightWork: Promise[] = [ + ...[...connectInFlight.values()].map((attempt) => attempt.promise), + ...resetRelayInFlight.values(), + ...testConnectionProbes + ] + for (const targetId of Array.from(connectInFlight.keys())) { + invalidateConnectAttempt(targetId) + } + + const errors: unknown[] = [] + const drain = async (): Promise => { + const results = await Promise.allSettled([ + ...[...activeSessions.keys()].map((targetId) => + teardownActiveSshSession(targetId, (session) => session.detachAndPersist()) + ), + connectionManager?.disconnectAll() ?? Promise.resolve() + ]) + for (const result of results) { + if (result.status === 'rejected') { + errors.push(result.reason as unknown) + } + } + } + + await drain() + // Why: a connect paused in old-session teardown still publishes its replacement session and opens a + // transport before it reaches the cancellation checkpoint, so the first drain can miss both. + await settleWithinMs(inFlightWork, SSH_SHUTDOWN_INFLIGHT_JOIN_MS) + await drain() + if (errors.length > 0) { + throw new AggregateError(errors, 'SSH shutdown detach failed') + } +} + export async function resetSshHandlerStateForTests(): Promise { advertisedUrlWatcherUnsubscribe?.() advertisedUrlWatcherUnsubscribe = null @@ -1498,7 +1624,9 @@ export async function resetSshHandlerStateForTests(): Promise { resetSshProviderAuthorities() resetRelayInFlight.clear() testingTargets.clear() + testConnectionProbes.clear() credentialRequestedForTarget.clear() + quitTeardownStartGate.resetForTests() await connectionManager?.disconnectAll() portForwardManager?.dispose() diff --git a/src/main/quit-teardown-start-gate.ts b/src/main/quit-teardown-start-gate.ts index be970495a4b..a96b7b3f38b 100644 --- a/src/main/quit-teardown-start-gate.ts +++ b/src/main/quit-teardown-start-gate.ts @@ -9,4 +9,16 @@ export class QuitTeardownStartGate { this.started = true return true } + + hasStarted(): boolean { + return this.started + } + + resetForTests(): void { + this.started = false + } } + +// Why shared, not per-caller: subsystems fence new work on "the committed quit has begun", and a +// second instance would answer no while the real teardown is already draining what it snapshotted. +export const quitTeardownStartGate = new QuitTeardownStartGate() diff --git a/src/main/ssh/ssh-connection-manager.test.ts b/src/main/ssh/ssh-connection-manager.test.ts index 3a7a73cbcbf..b2de497a006 100644 --- a/src/main/ssh/ssh-connection-manager.test.ts +++ b/src/main/ssh/ssh-connection-manager.test.ts @@ -94,4 +94,36 @@ describe('SshConnectionManager', () => { expect(await manager.connect(target)).toBe(replacement) expect(manager.getConnection(target.id)).toBe(replacement) }) + + it('closes a late-resolved connection without evicting the replacement', async () => { + let resolveFirst!: () => void + mockState.connectResults.push( + new Promise((resolve) => { + resolveFirst = resolve + }), + Promise.resolve() + ) + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + + const firstConnect = manager.connect(target) + await manager.disconnect(target.id) + const replacement = await manager.connect(target) + resolveFirst() + const late = await firstConnect + + await manager.disconnectConnection(target.id, late) + + expect(mockState.instances[0].disconnect).toHaveBeenCalledTimes(2) + expect(manager.getConnection(target.id)).toBe(replacement) + }) + + it('clears the pool entry when the closed connection is still the registered one', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + const conn = await manager.connect(target) + + await manager.disconnectConnection(target.id, conn) + + expect(mockState.instances[0].disconnect).toHaveBeenCalledOnce() + expect(manager.getConnection(target.id)).toBeUndefined() + }) }) diff --git a/src/main/ssh/ssh-connection-manager.ts b/src/main/ssh/ssh-connection-manager.ts index 8bc17d06b70..e553bf906b5 100644 --- a/src/main/ssh/ssh-connection-manager.ts +++ b/src/main/ssh/ssh-connection-manager.ts @@ -76,6 +76,18 @@ export class SshConnectionManager { } } + /** + * Close one specific connection, clearing the pool entry only when it is still the registered one. + * Why: a cancelled connect whose transport opened late owns that exact connection — disconnecting + * by target id would tear down the replacement's live transport instead. + */ + async disconnectConnection(targetId: string, conn: SshConnection): Promise { + await conn.disconnect() + if (this.connections.get(targetId) === conn) { + this.connections.delete(targetId) + } + } + async reconnect(targetId: string): Promise { const conn = this.connections.get(targetId) if (!conn) { diff --git a/src/main/ssh/ssh-owner-recovery-retry.test.ts b/src/main/ssh/ssh-owner-recovery-retry.test.ts new file mode 100644 index 00000000000..9c8f4bd6a14 --- /dev/null +++ b/src/main/ssh/ssh-owner-recovery-retry.test.ts @@ -0,0 +1,106 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR +} from '../../shared/pty-consumer-session' +import { retrySshOwnerRecoveryWhileBlocked } from './ssh-owner-recovery-retry' + +function publicationPendingError(): Error & { code: number } { + return Object.assign(new Error('Owner grant publication is still pending'), { + code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR + }) +} + +function supersededError(): Error & { code: number } { + return Object.assign(new Error('Owner recovery generation was superseded'), { + code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR + }) +} + +function openGate() { + return { + isCurrent: () => true, + onClosed: () => () => {} + } +} + +describe('SSH owner recovery retry', () => { + afterEach(() => { + vi.useRealTimers() + }) + + it('recovers once the incumbent grant publication settles', async () => { + vi.useFakeTimers() + const attempt = vi + .fn<() => Promise>() + .mockRejectedValueOnce(publicationPendingError()) + .mockRejectedValueOnce(publicationPendingError()) + .mockResolvedValue('recovered') + + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate()) + await vi.advanceTimersByTimeAsync(75) + + await expect(recovery).resolves.toBe('recovered') + expect(attempt).toHaveBeenCalledTimes(3) + }) + + it('stops at the publication-settlement deadline', async () => { + vi.useFakeTimers() + const error = publicationPendingError() + const attempt = vi.fn<() => Promise>().mockRejectedValue(error) + + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate(), 60) + const rejection = expect(recovery).rejects.toBe(error) + await vi.advanceTimersByTimeAsync(60) + + await rejection + expect(attempt).toHaveBeenCalledTimes(3) + }) + + it('does not retry unrelated failures', async () => { + const error = Object.assign(new Error('stale owner'), { code: -32041 }) + const attempt = vi.fn<() => Promise>().mockRejectedValue(error) + + await expect(retrySshOwnerRecoveryWhileBlocked(attempt, openGate())).rejects.toBe(error) + expect(attempt).toHaveBeenCalledOnce() + }) + + it('retries while a superseded transport is closing', async () => { + vi.useFakeTimers() + const attempt = vi + .fn<() => Promise>() + .mockRejectedValueOnce(supersededError()) + .mockResolvedValue('recovered') + + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate()) + await vi.advanceTimersByTimeAsync(25) + + await expect(recovery).resolves.toBe('recovered') + expect(attempt).toHaveBeenCalledTimes(2) + }) + + it('stops waiting when the relay channel closes', async () => { + vi.useFakeTimers() + let current = true + let close: (() => void) | undefined + const error = publicationPendingError() + const attempt = vi.fn<() => Promise>().mockRejectedValue(error) + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, { + isCurrent: () => current, + onClosed: (listener) => { + close = listener + return () => { + close = undefined + } + } + }) + await vi.advanceTimersByTimeAsync(0) + expect(close).toBeTypeOf('function') + + current = false + close?.() + + await expect(recovery).rejects.toBe(error) + expect(attempt).toHaveBeenCalledOnce() + }) +}) diff --git a/src/main/ssh/ssh-owner-recovery-retry.ts b/src/main/ssh/ssh-owner-recovery-retry.ts new file mode 100644 index 00000000000..e5e78d9a4f2 --- /dev/null +++ b/src/main/ssh/ssh-owner-recovery-retry.ts @@ -0,0 +1,69 @@ +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR +} from '../../shared/pty-consumer-session' + +// Why: bound polling when publication is settling or a superseded attempt is closing its transport. +export const SSH_OWNER_RECOVERY_WAIT_MS = 3_000 +const SSH_OWNER_RECOVERY_INITIAL_DELAY_MS = 25 +const SSH_OWNER_RECOVERY_MAX_DELAY_MS = 250 + +type SshOwnerRecoveryRetryGate = { + isCurrent: () => boolean + onClosed: (listener: () => void) => () => void +} + +function waitForRetry(delayMs: number, gate: SshOwnerRecoveryRetryGate): Promise { + return new Promise((resolve) => { + let settled = false + let timer: ReturnType | undefined + let unsubscribe = (): void => {} + const finish = (): void => { + if (settled) { + return + } + settled = true + clearTimeout(timer) + unsubscribe() + resolve() + } + unsubscribe = gate.onClosed(finish) + timer = setTimeout(finish, delayMs) + timer.unref?.() + if (!gate.isCurrent()) { + finish() + } + }) +} + +export async function retrySshOwnerRecoveryWhileBlocked( + attempt: () => Promise, + gate: SshOwnerRecoveryRetryGate, + waitMs: number = SSH_OWNER_RECOVERY_WAIT_MS +): Promise { + const deadline = Date.now() + waitMs + let delayMs = SSH_OWNER_RECOVERY_INITIAL_DELAY_MS + while (true) { + try { + return await attempt() + } catch (error) { + const code = (error as { code?: unknown } | null | undefined)?.code + if ( + (code !== PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR && + code !== PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR) || + !gate.isCurrent() + ) { + throw error + } + const remainingMs = deadline - Date.now() + if (remainingMs <= 0) { + throw error + } + await waitForRetry(Math.min(delayMs, remainingMs), gate) + if (!gate.isCurrent()) { + throw error + } + delayMs = Math.min(delayMs * 2, SSH_OWNER_RECOVERY_MAX_DELAY_MS) + } + } +} diff --git a/src/main/ssh/ssh-relay-session-recovery-durability.test.ts b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts index df7bce635ec..6e1fe81c5a5 100644 --- a/src/main/ssh/ssh-relay-session-recovery-durability.test.ts +++ b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts @@ -1,4 +1,8 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR +} from '../../shared/pty-consumer-session' import { SshRelaySession } from './ssh-relay-session' import { createMismatchedOwnerRecoveryError, @@ -154,6 +158,82 @@ describe('SshRelaySession consumer recovery durability', () => { session.dispose() }) + it('retries a pending incumbent publication before recovering its persisted lease', async () => { + vi.useFakeTimers() + try { + const targetId = 'target-owner-publication-pending' + const deps = createMockDeps() + vi.mocked(deps.mockStore.getSshPtyConsumerRecovery).mockReturnValue({ + targetId, + clientInstanceId: 'persisted-client', + serverBuildId: 'test-relay-build', + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'persisted-owner' + }) + openConsumerSessionMock.mockRejectedValueOnce( + Object.assign(new Error('Owner grant publication is still pending'), { + code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR + }) + ) + const session = new SshRelaySession( + targetId, + deps.getMainWindow, + deps.mockStore, + deps.mockPortForward + ) + + const establishing = session.establish(deps.mockConn) + await vi.advanceTimersByTimeAsync(25) + await establishing + + expect(openConsumerSessionMock).toHaveBeenCalledTimes(2) + expect(openConsumerSessionMock.mock.calls[0]?.[1]).toMatchObject({ + clientInstanceId: 'persisted-client', + resume: { ownerGeneration: 1, ownerLease: 'persisted-owner' } + }) + session.dispose() + } finally { + vi.useRealTimers() + } + }) + + it('preserves persisted recovery while a superseding transport is still live', async () => { + vi.useFakeTimers() + try { + const targetId = 'target-owner-generation-superseded' + const deps = createMockDeps() + vi.mocked(deps.mockStore.getSshPtyConsumerRecovery).mockReturnValue({ + targetId, + clientInstanceId: 'persisted-client', + serverBuildId: 'test-relay-build', + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'persisted-owner' + }) + const superseded = Object.assign(new Error('Owner recovery generation was superseded'), { + code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR + }) + openConsumerSessionMock.mockRejectedValue(superseded) + const session = new SshRelaySession( + targetId, + deps.getMainWindow, + deps.mockStore, + deps.mockPortForward + ) + + const failed = expect(session.establish(deps.mockConn)).rejects.toBe(superseded) + await vi.advanceTimersByTimeAsync(3_000) + await failed + + expect(openConsumerSessionMock.mock.calls.length).toBeGreaterThan(1) + expect(deps.mockStore.removeSshPtyConsumerRecovery).not.toHaveBeenCalled() + session.dispose() + } finally { + vi.useRealTimers() + } + }) + it('keeps destructive disposal pending until consumer recovery is removed', async () => { const { mockStore, mockPortForward, getMainWindow } = createMockDeps() vi.mocked(mockStore.getSshPtyConsumerRecovery).mockReturnValue({ @@ -232,6 +312,46 @@ describe('SshRelaySession consumer recovery durability', () => { expect(openConsumerSessionMock).toHaveBeenCalledTimes(1) }) + it('leaves recovery state alone when a newer owner already claimed the target record', async () => { + const targetId = 'target-stale-owner-loser' + 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' + }) + const winner = { + mode: 'negotiated' as const, + clientInstanceId: 'persisted-client', + clientGeneration: 2, + ownerGeneration: 5, + ownerLease: 'winner-owner' + } + openConsumerSessionMock.mockImplementationOnce(() => { + // Why inside the rejection: the record is target-scoped, so the winner can land while this + // attempt is still unwinding its own resume. + const record = getSshPtyConsumerRecovery(targetId)! + record.owner = winner + record.checkpointsByAppPtyId.set('pty-1', { + id: 'pty-1' + } as unknown as never) + return Promise.reject(createMismatchedOwnerRecoveryError()) + }) + openConsumerSessionMock.mockRejectedValueOnce(new Error('fresh open failed')) + const session = new SshRelaySession(targetId, getMainWindow, mockStore, mockPortForward) + + await expect(session.establish(mockConn)).rejects.toThrow('fresh open failed') + + const record = getSshPtyConsumerRecovery(targetId)! + expect(record.owner).toBe(winner) + expect(record.checkpointsByAppPtyId.has('pty-1')).toBe(true) + expect(mockStore.removeSshPtyConsumerRecovery).not.toHaveBeenCalled() + session.dispose() + }) + it('does not remember a consumer opened after establish was disposed', async () => { const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() let signalOpenStarted!: () => void diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index d413f45f044..56d2ab6da8a 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -80,6 +80,7 @@ import type { Store } from '../persistence' import type { OrcaRuntimeService } from '../runtime/orca-runtime' import { DEFAULT_PTY_SOURCE_WINDOW_SU } from '../../shared/pty-source-credit-contract' import { PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR } from '../../shared/pty-consumer-session' +import { retrySshOwnerRecoveryWhileBlocked } from './ssh-owner-recovery-retry' import { runRemoteOrcaCli } from './ssh-remote-orca-cli' import { acknowledgeRemoteOrcaCliPostOutput, @@ -991,17 +992,24 @@ export class SshRelaySession { outputFlowControl: { requestedWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU } } try { - return await openSshPtyConsumerSession(mux, { - ...options, - ...(previousOwner - ? { - resume: { - ownerGeneration: previousOwner.ownerGeneration, - ownerLease: previousOwner.ownerLease - } - } - : {}) - }) + return await retrySshOwnerRecoveryWhileBlocked( + () => + openSshPtyConsumerSession(mux, { + ...options, + ...(previousOwner + ? { + resume: { + ownerGeneration: previousOwner.ownerGeneration, + ownerLease: previousOwner.ownerLease + } + } + : {}) + }), + { + isCurrent: () => ownsAttempt() && !mux.isDisposed(), + onClosed: (listener) => mux.onDispose(listener) + } + ) } catch (error) { if ( !previousOwner || @@ -1010,26 +1018,36 @@ export class SshRelaySession { throw error } const recovery = getSshPtyConsumerRecovery(this.targetId) - if (recovery) { - delete recovery.owner - recovery.checkpointsByAppPtyId.clear() - for (const [ptyId, migration] of recovery.modelMigrationsByAppPtyId) { - recovery.modelMigrationsByAppPtyId.set( - ptyId, - migration.then(() => - Object.freeze({ - status: 'checkpoint-unavailable' as const, - reason: 'completion-failed' as const - }) + // Why identity-guarded: the record is target-scoped and its clientInstanceId is shared by every + // session for that target, so only a record still describing the owner this attempt tried to + // resume is ours to drop — otherwise a loser wipes the winner's checkpoints. + const ownsRecoveryRecord = + ownsAttempt() && + (!recovery?.owner || + (recovery.owner.ownerGeneration === previousOwner.ownerGeneration && + recovery.owner.ownerLease === previousOwner.ownerLease)) + if (ownsRecoveryRecord) { + if (recovery) { + delete recovery.owner + recovery.checkpointsByAppPtyId.clear() + for (const [ptyId, migration] of recovery.modelMigrationsByAppPtyId) { + recovery.modelMigrationsByAppPtyId.set( + ptyId, + migration.then(() => + Object.freeze({ + status: 'checkpoint-unavailable' as const, + reason: 'completion-failed' as const + }) + ) ) - ) + } } + await removeSshPtyConsumerOwnerRecovery( + this.targetId, + this.ptyConsumerClientInstanceId, + this.store + ) } - await removeSshPtyConsumerOwnerRecovery( - this.targetId, - this.ptyConsumerClientInstanceId, - this.store - ) if (!ownsAttempt()) { throw new Error('Session disposed during establish') } diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 98e8959547f..d89d0795ed9 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -161,6 +161,16 @@ export class RelayDispatcher { this.closeClient(client, new Error('Relay client detached'), true) } + // Why: a displaced owner must lose its transport whichever client holds it, and the launch channel is + // often the half-open one. The primary keeps its id for setWrite() revival, so invalidate it instead. + releaseDisplacedClient(clientId: number): void { + if (clientId === this.primaryClient.id) { + this.invalidateClient() + return + } + this.detachClient(clientId) + } + feedClient(clientId: number, data: Buffer): void { const client = this.clients.get(clientId) if (!client) { diff --git a/src/relay/relay-pty-consumer-owner-displacement.test.ts b/src/relay/relay-pty-consumer-owner-displacement.test.ts new file mode 100644 index 00000000000..89059119f5f --- /dev/null +++ b/src/relay/relay-pty-consumer-owner-displacement.test.ts @@ -0,0 +1,368 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { + RelayDispatcher, + type RelayClientSessionIdentity, + type SinkWriteSettlement +} from './dispatcher' +import { encodeJsonRpcFrame, MessageType } from './protocol' +import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR +} from '../shared/pty-consumer-session' + +const endpointIdentity: RelayClientSessionIdentity = { + principal: 'endpoint-principal', + authenticated: true, + allowSessionOwner: true, + authenticationKind: 'endpoint-credential' +} + +function requestFrame(id: number, method: string, params: Record): Buffer { + return encodeJsonRpcFrame({ jsonrpc: '2.0', id, method, params }, id, 0) +} + +function message(buffer: Buffer): Record | null { + if (buffer[0] !== MessageType.Regular) { + return null + } + const length = buffer.readUInt32BE(9) + return JSON.parse(buffer.subarray(13, 13 + length).toString('utf8')) +} + +function response(writes: readonly Buffer[], id: number): Record | undefined { + return writes.map(message).find((entry) => entry?.id === id) as + | Record + | undefined +} + +function ownerHelloParams(resume?: Record): Record { + return { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 8 } }, + ...(resume ? { resume } : {}) + } +} + +async function flushRequests(): Promise { + await new Promise((resolve) => setImmediate(resolve)) +} + +describe('relay PTY consumer owner displacement', () => { + let dispatcher: RelayDispatcher | null = null + + afterEach(() => { + dispatcher?.dispose() + dispatcher = null + }) + + it('displaces a half-open owner and retains its deliveries for the reconnected client', async () => { + // Why no settle deferral: a half-open peer is indistinguishable from a live one at the sink — writes + // still "succeed" locally, so the relay never learns the owner is gone. + const staleWrites: Buffer[] = [] + const closeStaleTransport = vi.fn() + dispatcher = new RelayDispatcher( + (data, settle) => { + staleWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true, close: closeStaleTransport }, + endpointIdentity + ) + const detached: number[] = [] + dispatcher.onClientDetached((clientId) => detached.push(clientId)) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams())) + await flushRequests() + const staleGrant = response(staleWrites, 1)!.result as Record + expect(adapter.deliveryMode(1)).toBe('source-owner') + + const staleDelivery = adapter.openDelivery(1, 'pty-1', 'incarnation-1')! + adapter.appendSource(staleDelivery, { + spanId: 'span-1', + data: 'tail', + displayStart: 0, + displayEnd: 4, + splittable: true, + transform: { transformed: false, rawLengthSu: 4, scalarSafe: true } + }) + + const reconnectWrites: Buffer[] = [] + const reconnectClientId = dispatcher.attachClient( + (data, settle) => { + reconnectWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + reconnectClientId, + requestFrame( + 2, + 'pty.openClient', + ownerHelloParams({ + ownerGeneration: staleGrant.ownerGeneration, + ownerLease: staleGrant.ownerLease + }) + ) + ) + await flushRequests() + + expect(response(reconnectWrites, 2)!.result).toMatchObject({ + role: 'session-owner', + ownerGeneration: 2 + }) + expect(adapter.deliveryMode(reconnectClientId)).toBe('source-owner') + // Why: revocation is what bounds the takeover — the stale owner cannot open or drive deliveries again + // even though its socket never closed. + expect(adapter.deliveryMode(1)).toBe('subscriber') + expect(adapter.openDelivery(1, 'pty-2', 'incarnation-1')).toBeNull() + expect(detached).toContain(1) + expect(closeStaleTransport).toHaveBeenCalledOnce() + + // Why: retained, not closed — the reconnected owner still has to rotate the stale delivery forward. + expect(adapter.getDebugSnapshot()).toMatchObject({ deliveryTokens: 1, graceTimers: 1 }) + const rotation = adapter.rotateDelivery(staleDelivery, reconnectClientId, 0) + expect(rotation.identity).toMatchObject({ ownerGeneration: 2, clientGeneration: 2 }) + expect(rotation.recovery.map((span) => span.data)).toEqual(['tail']) + }) + + it('retains the incumbent delivery once when its socket closes mid-publication', async () => { + const incumbentWrites: Buffer[] = [] + dispatcher = new RelayDispatcher( + (data, settle) => { + incumbentWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams())) + await flushRequests() + const incumbentGrant = response(incumbentWrites, 1)!.result as Record + const incumbentDelivery = adapter.openDelivery(1, 'pty-1', 'incarnation-1')! + adapter.appendSource(incumbentDelivery, { + spanId: 'span-1', + data: 'tail', + displayStart: 0, + displayEnd: 4, + splittable: true, + transform: { transformed: false, rawLengthSu: 4, scalarSafe: true } + }) + + const reconnectWrites: Buffer[] = [] + let grantSettlement: ((result: SinkWriteSettlement) => void) | undefined + const reconnectClientId = dispatcher.attachClient( + (data, settle) => { + reconnectWrites.push(Buffer.from(data)) + grantSettlement = settle + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + reconnectClientId, + requestFrame( + 2, + 'pty.openClient', + ownerHelloParams({ + ownerGeneration: incumbentGrant.ownerGeneration, + ownerLease: incumbentGrant.ownerLease + }) + ) + ) + await flushRequests() + + // Why this ordering: the incumbent's socket close and the replacement's grant write are independent + // events, so detach-then-commit has to leave exactly one retention, not two competing ones. + dispatcher.invalidateClient() + grantSettlement!({ ok: true }) + + expect(adapter.deliveryMode(reconnectClientId)).toBe('source-owner') + expect(adapter.getDebugSnapshot()).toMatchObject({ deliveryTokens: 1, graceTimers: 1 }) + const rotation = adapter.rotateDelivery(incumbentDelivery, reconnectClientId, 0) + expect(rotation.identity).toMatchObject({ ownerGeneration: 2, clientGeneration: 2 }) + expect(adapter.getDebugSnapshot()).toMatchObject({ graceTimers: 0 }) + }) + + it('keeps the incumbent owner when the replacement grant fails to publish', async () => { + const incumbentWrites: Buffer[] = [] + dispatcher = new RelayDispatcher( + (data, settle) => { + incumbentWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + const detached: number[] = [] + dispatcher.onClientDetached((clientId) => detached.push(clientId)) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams())) + await flushRequests() + const incumbentGrant = response(incumbentWrites, 1)!.result as Record + + const reconnectWrites: Buffer[] = [] + let grantSettlement: ((result: SinkWriteSettlement) => void) | undefined + const reconnectClientId = dispatcher.attachClient( + (data, settle) => { + reconnectWrites.push(Buffer.from(data)) + // Why: the takeover must not commit until the replacement grant is on the wire. + grantSettlement = settle + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + reconnectClientId, + requestFrame( + 2, + 'pty.openClient', + ownerHelloParams({ + ownerGeneration: incumbentGrant.ownerGeneration, + ownerLease: incumbentGrant.ownerLease + }) + ) + ) + await flushRequests() + + expect(adapter.deliveryMode(1)).toBe('source-owner') + expect(detached).not.toContain(1) + + grantSettlement!({ ok: false, error: new Error('reconnect socket closed mid-response') }) + expect(adapter.deliveryMode(1)).toBe('source-owner') + expect(adapter.deliveryMode(reconnectClientId)).toBe('subscriber') + + // Why: the rolled-back attempt must leave the incumbent's lease reclaimable by the next reconnect. + const retryWrites: Buffer[] = [] + const retryClientId = dispatcher.attachClient( + (data, settle) => { + retryWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + retryClientId, + requestFrame( + 3, + 'pty.openClient', + ownerHelloParams({ + ownerGeneration: incumbentGrant.ownerGeneration, + ownerLease: incumbentGrant.ownerLease + }) + ) + ) + await flushRequests() + + // Why generation 3: the rolled-back attempt consumed generation 2, but not the incumbent's lease. + expect(response(retryWrites, 3)!.result).toMatchObject({ + role: 'session-owner', + ownerGeneration: 3 + }) + expect(adapter.deliveryMode(1)).toBe('subscriber') + expect(detached).toContain(1) + }) + + it('fences an old proof while its replacement is live', async () => { + const incumbentWrites: Buffer[] = [] + dispatcher = new RelayDispatcher( + (data, settle) => { + incumbentWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams())) + await flushRequests() + const incumbentGrant = response(incumbentWrites, 1)!.result as Record + const resume = { + ownerGeneration: incumbentGrant.ownerGeneration, + ownerLease: incumbentGrant.ownerLease + } + + let firstReconnectSettlement: ((result: SinkWriteSettlement) => void) | undefined + const firstReconnectWrites: Buffer[] = [] + const firstReconnectClientId = dispatcher.attachClient( + (data, settle) => { + firstReconnectWrites.push(Buffer.from(data)) + // Why ??=: the test settles the grant response, not whatever frame the dispatcher wrote last. + firstReconnectSettlement ??= settle + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + firstReconnectClientId, + requestFrame(2, 'pty.openClient', ownerHelloParams(resume)) + ) + await flushRequests() + + const retryWrites: Buffer[] = [] + const retryClientId = dispatcher.attachClient( + (data, settle) => { + retryWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + retryClientId, + requestFrame(3, 'pty.openClient', ownerHelloParams(resume)) + ) + await flushRequests() + + expect(response(retryWrites, 3)!.error).toMatchObject({ + code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR + }) + + firstReconnectSettlement!({ ok: true }) + dispatcher.feedClient( + retryClientId, + requestFrame(4, 'pty.openClient', ownerHelloParams(resume)) + ) + await flushRequests() + + expect(response(retryWrites, 4)!.error).toMatchObject({ + code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR + }) + expect(adapter.deliveryMode(firstReconnectClientId)).toBe('source-owner') + expect(adapter.deliveryMode(retryClientId)).toBe('subscriber') + + dispatcher.detachClient(firstReconnectClientId) + dispatcher.feedClient( + retryClientId, + requestFrame(5, 'pty.openClient', ownerHelloParams(resume)) + ) + await flushRequests() + + expect(response(retryWrites, 5)!.result).toMatchObject({ + role: 'session-owner', + ownerGeneration: 3, + ownerLease: incumbentGrant.ownerLease + }) + expect(adapter.deliveryMode(retryClientId)).toBe('source-owner') + }) +}) diff --git a/src/relay/relay.ts b/src/relay/relay.ts index 4bfc8bf3e92..f90126977ab 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -10,7 +10,15 @@ import { createServer, createConnection, type Socket, type Server } from 'node:net' import { join } from 'node:path' -import { unlinkSync, existsSync, statSync, readFileSync, chmodSync } from 'node:fs' +import { + unlinkSync, + existsSync, + statSync, + readFileSync, + chmodSync, + closeSync, + openSync +} from 'node:fs' import { RELAY_SENTINEL, FrameDecoder, @@ -596,6 +604,31 @@ async function main(): Promise { } stdoutDrainWaiters.add(cb) return () => stdoutDrainWaiters.delete(cb) + }, + close: () => { + stdoutAlive = false + flushStdoutDrainWaiters() + // Why close then re-pin: the SSH peer must see EOF, but a long-lived daemon that + // frees fds 0/1 lets accept()/open() recycle them while Node still treats + // process.stdin/stdout as those numbers — corrupting socket clients and shutdown. + for (const fd of [process.stdin.fd, process.stdout.fd]) { + try { + closeSync(fd) + } catch { + // Already closed by the peer. + } + } + const devNull = process.platform === 'win32' ? 'NUL' : '/dev/null' + try { + openSync(devNull, 'r') + } catch { + /* best-effort pin of the lowest free fd (normally 0) */ + } + try { + openSync(devNull, 'w') + } catch { + /* best-effort pin of the next free fd (normally 1) */ + } } }, undefined, diff --git a/src/relay/ssh-pty-consumer-session-adapter.ts b/src/relay/ssh-pty-consumer-session-adapter.ts index 1617fdb09dd..b28082aa7bc 100644 --- a/src/relay/ssh-pty-consumer-session-adapter.ts +++ b/src/relay/ssh-pty-consumer-session-adapter.ts @@ -1,6 +1,7 @@ import { PTY_CONSUMER_SESSION_PROTOCOL_VERSION, PtyConsumerSession, + type PtyConsumerSessionAdmission, type PtyConsumerSessionGrant, type PtyConsumerSessionHello } from '../shared/pty-consumer-session' @@ -74,7 +75,7 @@ export class SshPtyConsumerSessionAdapter { private readonly pausedDeliveryByPty = new Map() constructor( - dispatcher: RelayDispatcher, + private readonly dispatcher: RelayDispatcher, serverBuildId: string, private readonly setDeliveryPaused?: (id: string, paused: boolean) => void, onSourceCreditAvailable?: (id: string) => void @@ -268,15 +269,30 @@ export class SshPtyConsumerSessionAdapter { throw new Error('SSH PTY consumer response publication fence is unavailable') } context.onResponseSettled((result) => { - if (result.ok) { - admission.commitPublication() - } else { + if (!result.ok) { admission.rollbackPublication() + return } + admission.commitPublication() + this.closeDisplacedOwner(admission.displacedOwner) }) return admission.grant } + // Why: only after the replacement grant is published — until then the admission can still roll back + // onto the incumbent. Its deliveries are retained for the new owner to rotate, exactly as on detach. + private closeDisplacedOwner(displaced: PtyConsumerSessionAdmission['displacedOwner']): void { + if (!displaced) { + return + } + this.clearPausedForGrant(displaced.grant) + this.sourceCredit.retainOrCloseOnDetach(displaced.grant) + const clientId = Number(displaced.connectionId) + if (Number.isSafeInteger(clientId)) { + this.dispatcher.releaseDisplacedClient(clientId) + } + } + private clearPausedIdentity(identity: PtySourceDeliveryIdentity): void { if (this.pausedDeliveryByPty.get(identity.id) === identity) { this.pausedDeliveryByPty.delete(identity.id) diff --git a/src/shared/pty-consumer-owner-recovery.ts b/src/shared/pty-consumer-owner-recovery.ts new file mode 100644 index 00000000000..dfac94d7e9c --- /dev/null +++ b/src/shared/pty-consumer-owner-recovery.ts @@ -0,0 +1,63 @@ +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR, + PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR, + type PtyConsumerAuthentication, + type PtyConsumerSessionHello +} from './pty-consumer-session-contract' + +type IncumbentOwner = { + principal: string + clientInstanceId: string + generation: number + lease: string + state: 'pending' | 'active' | 'disconnected' + replaces?: { generation: number } +} + +function throwRecoveryError(message: string, code: number): never { + throw Object.assign(new Error(message), { code }) +} + +export function assertPtyConsumerOwnerRecovery( + hello: PtyConsumerSessionHello, + authentication: PtyConsumerAuthentication, + current: IncumbentOwner +): void { + const resume = hello.resume + if (!resume) { + throw new Error('Owner recovery proof is required') + } + if ( + resume.ownerLease !== current.lease || + hello.clientInstanceId !== current.clientInstanceId || + authentication.principal !== current.principal + ) { + throwRecoveryError( + 'Owner recovery lease is stale or belongs to another principal', + PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR + ) + } + if (current.state === 'active' && resume.ownerGeneration < current.generation) { + throwRecoveryError( + 'Owner recovery generation was superseded', + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR + ) + } + const generationMatches = + resume.ownerGeneration === current.generation || + (current.state === 'pending' && resume.ownerGeneration === current.replaces?.generation) || + (current.state === 'disconnected' && resume.ownerGeneration < current.generation) + if (!generationMatches) { + throwRecoveryError( + 'Owner recovery generation is stale', + PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR + ) + } + if (current.state === 'pending') { + throwRecoveryError( + 'Owner grant publication is still pending', + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR + ) + } +} diff --git a/src/shared/pty-consumer-session-contract.ts b/src/shared/pty-consumer-session-contract.ts index 661138896cd..1e359a11f53 100644 --- a/src/shared/pty-consumer-session-contract.ts +++ b/src/shared/pty-consumer-session-contract.ts @@ -1,6 +1,10 @@ export const PTY_CONSUMER_SESSION_PROTOCOL_VERSION = 1 export const PTY_CONSUMER_OWNER_GRACE_MS = 30_000 export const PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR = -32041 +// Why: recovery is blocked only while the incumbent owner's grant publication is still settling — a +// window bounded by one response write, so the client may retry within a short budget. +export const PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR = -32042 +export const PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR = -32043 export type PtyConsumerRole = 'session-owner' | 'subscriber' @@ -41,8 +45,16 @@ export type PtyConsumerAuthentication = { allowSessionOwner: boolean } +export type PtyConsumerDisplacedOwner = { + connectionId: string + grant: Readonly +} + export type PtyConsumerSessionAdmission = { grant: Readonly + // Why: set when this admission takes over a still-attached owner. The transport layer owns closing + // that connection and releasing its deliveries — do it only once the new grant has been published. + displacedOwner?: Readonly commitPublication: () => void rollbackPublication: () => void } diff --git a/src/shared/pty-consumer-session-hello.ts b/src/shared/pty-consumer-session-hello.ts new file mode 100644 index 00000000000..c8119da88ee --- /dev/null +++ b/src/shared/pty-consumer-session-hello.ts @@ -0,0 +1,52 @@ +import type { PtyConsumerSessionHello } from './pty-consumer-session-contract' + +export const MAX_CAPABILITY_VERSIONS = 8 + +export function assertNonEmptyString(value: unknown, name: string): asserts value is string { + if (typeof value !== 'string' || value.length === 0 || value.length > 512) { + throw new Error(`${name} must be a non-empty string of at most 512 characters`) + } +} + +export function validateHello(hello: PtyConsumerSessionHello): void { + assertNonEmptyString(hello.clientInstanceId, 'clientInstanceId') + if (hello.requestedRole !== 'session-owner' && hello.requestedRole !== 'subscriber') { + throw new Error('requestedRole must be session-owner or subscriber') + } + if (hello.resume) { + if (!Number.isSafeInteger(hello.resume.ownerGeneration) || hello.resume.ownerGeneration <= 0) { + throw new Error('resume.ownerGeneration must be a positive safe integer') + } + assertNonEmptyString(hello.resume.ownerLease, 'resume.ownerLease') + } + const flow = hello.capabilities?.outputFlowControl + if (!flow) { + return + } + if ( + !Array.isArray(flow.versions) || + flow.versions.length > MAX_CAPABILITY_VERSIONS || + flow.versions.some((version) => !Number.isSafeInteger(version) || version <= 0) + ) { + throw new Error('outputFlowControl.versions must contain positive safe integers') + } + if (!Number.isSafeInteger(flow.requestedWindowSu) || flow.requestedWindowSu <= 0) { + throw new Error('outputFlowControl.requestedWindowSu must be a positive safe integer') + } +} + +// Why: a duplicate open on one connection is only idempotent if it asks for exactly the same thing. +export function helloFingerprint(hello: PtyConsumerSessionHello): string { + const flow = hello.capabilities?.outputFlowControl + return JSON.stringify({ + clientInstanceId: hello.clientInstanceId, + requestedRole: hello.requestedRole, + resume: hello.resume, + outputFlowControl: flow + ? { + versions: [...flow.versions].sort((a, b) => a - b), + requestedWindowSu: flow.requestedWindowSu + } + : undefined + }) +} diff --git a/src/shared/pty-consumer-session.test.ts b/src/shared/pty-consumer-session.test.ts index 07f9f7f9bf9..ce6282e8664 100644 --- a/src/shared/pty-consumer-session.test.ts +++ b/src/shared/pty-consumer-session.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it } from 'vitest' import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR, PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR, PtyConsumerSession, type PtyConsumerAuthentication, @@ -121,7 +123,7 @@ describe('PtyConsumerSession', () => { ).toThrow('authentication required') }) - it('rotates the lease and increments owner generation on valid recovery', () => { + it('keeps the lease stable and increments owner generation on valid recovery', () => { const session = createSession() const first = session.admit(ownerHello(), auth('connection-1')) first.commitPublication() @@ -141,15 +143,97 @@ describe('PtyConsumerSession', () => { expect(recovered.grant).toMatchObject({ role: 'session-owner', ownerGeneration: 2, - ownerLease: 'lease-2' + ownerLease: 'lease-1' }) }) - it('rejects recovery while the current owner is still active', () => { + it('displaces a still-attached owner that a matching resume proof reclaims', () => { const session = createSession() const first = session.admit(ownerHello(), auth('connection-1')) first.commitPublication() + const recovered = session.admit( + ownerHello({ + resume: { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + }), + auth('connection-2') + ) + + expect(recovered.displacedOwner).toEqual({ + connectionId: 'connection-1', + grant: first.grant + }) + // Why: the incumbent keeps authority until the replacement grant is actually published. + expect(session.activeGrant('connection-1')).toBe(first.grant) + + recovered.commitPublication() + expect(recovered.grant).toMatchObject({ role: 'session-owner', ownerGeneration: 2 }) + expect(session.activeGrant('connection-1')).toBeNull() + expect(session.activeGrant('connection-2')).toBe(recovered.grant) + }) + + it('restores the displaced owner when the replacement publication rolls back', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + + const recovered = session.admit( + ownerHello({ + resume: { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + }), + auth('connection-2') + ) + recovered.rollbackPublication() + + expect(session.activeGrant('connection-1')).toBe(first.grant) + expect(session.activeGrant('connection-2')).toBeNull() + // Why: the restored incumbent must still hold the lease it was admitted with. + const reclaimed = session.admit( + ownerHello({ + resume: { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + }), + auth('connection-3') + ) + expect(reclaimed.displacedOwner?.connectionId).toBe('connection-1') + }) + + it('expires a displaced owner restored after its connection already closed', () => { + let now = 10 + const session = createSession({ now: () => now }) + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + + const recovered = session.admit( + ownerHello({ + resume: { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + }), + auth('connection-2') + ) + session.close('connection-1') + recovered.rollbackPublication() + + now += 30_000 + session.sweepExpired() + const fresh = session.admit(ownerHello(), auth('connection-4')) + expect(fresh.grant.role).toBe('session-owner') + }) + + it('refuses recovery while the incumbent grant publication is still settling', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + expect(() => session.admit( ownerHello({ @@ -160,7 +244,99 @@ describe('PtyConsumerSession', () => { }), auth('connection-2') ) - ).toThrow('Active owner') + ).toThrow( + expect.objectContaining({ + code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + message: expect.stringContaining('still pending') + }) + ) + }) + + it('fences an old recovery generation after its replacement commits', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + const resume = { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + const replacement = session.admit(ownerHello({ resume }), auth('connection-2')) + + expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow( + expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR }) + ) + + replacement.commitPublication() + expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow( + expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR }) + ) + + const retry = session.admit( + ownerHello({ + resume: { + ownerGeneration: replacement.grant.ownerGeneration!, + ownerLease: replacement.grant.ownerLease! + } + }), + auth('connection-3') + ) + + expect(retry.grant).toMatchObject({ + role: 'session-owner', + ownerGeneration: 3, + ownerLease: first.grant.ownerLease + }) + expect(retry.displacedOwner?.connectionId).toBe('connection-2') + // Why: the committed replacement already retired the incumbent, so only connection-2 is left to displace. + expect(session.activeGrant('connection-1')).toBeNull() + }) + + it('accepts an older stable proof after its replacement disconnects', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + const resume = { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + const replacement = session.admit(ownerHello({ resume }), auth('connection-2')) + replacement.commitPublication() + session.close('connection-2') + + const retry = session.admit(ownerHello({ resume }), auth('connection-3')) + + expect(retry.grant).toMatchObject({ + role: 'session-owner', + ownerGeneration: 3, + ownerLease: first.grant.ownerLease + }) + }) + + it('retries overlapping recovery against the incumbent after publication rolls back', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + const resume = { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + const replacement = session.admit(ownerHello({ resume }), auth('connection-2')) + + expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow( + expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR }) + ) + + replacement.rollbackPublication() + const retry = session.admit(ownerHello({ resume }), auth('connection-3')) + + expect(retry.grant).toMatchObject({ + role: 'session-owner', + ownerGeneration: 3, + ownerLease: first.grant.ownerLease + }) + expect(retry.displacedOwner?.connectionId).toBe('connection-1') + // Why: the rolled-back replacement never published, so it holds no grant the retry could displace. + expect(session.activeGrant('connection-2')).toBeNull() }) it('types mismatched recovery without disturbing principal or lease ownership', () => { @@ -201,7 +377,7 @@ describe('PtyConsumerSession', () => { expect(recovered.grant).toMatchObject({ role: 'session-owner', ownerGeneration: 2, - ownerLease: 'lease-2' + ownerLease: 'lease-1' }) }) diff --git a/src/shared/pty-consumer-session.ts b/src/shared/pty-consumer-session.ts index 953b77c3a8d..3baf502acd8 100644 --- a/src/shared/pty-consumer-session.ts +++ b/src/shared/pty-consumer-session.ts @@ -4,22 +4,28 @@ import { PTY_CONSUMER_SESSION_PROTOCOL_VERSION, PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR, type PtyConsumerAuthentication, + type PtyConsumerDisplacedOwner, type PtyConsumerSessionAdmission, type PtyConsumerSessionGrant, type PtyConsumerSessionHello, type PtyConsumerSessionOptions } from './pty-consumer-session-contract' +import { + assertNonEmptyString, + helloFingerprint, + MAX_CAPABILITY_VERSIONS, + validateHello +} from './pty-consumer-session-hello' +import { assertPtyConsumerOwnerRecovery } from './pty-consumer-owner-recovery' export * from './pty-consumer-session-contract' -const MAX_CAPABILITY_VERSIONS = 8 - type ClientRecord = { fingerprint: string principal: string clientInstanceId: string grant: Readonly - state: 'pending' | 'active' + state: 'pending' | 'active' | 'displaced' publicationState: 'pending' | 'committed' | 'rolled-back' } @@ -34,54 +40,6 @@ type OwnerRecord = { replaces?: OwnerRecord } -function assertNonEmptyString(value: unknown, name: string): asserts value is string { - if (typeof value !== 'string' || value.length === 0 || value.length > 512) { - throw new Error(`${name} must be a non-empty string of at most 512 characters`) - } -} - -function validateHello(hello: PtyConsumerSessionHello): void { - assertNonEmptyString(hello.clientInstanceId, 'clientInstanceId') - if (hello.requestedRole !== 'session-owner' && hello.requestedRole !== 'subscriber') { - throw new Error('requestedRole must be session-owner or subscriber') - } - if (hello.resume) { - if (!Number.isSafeInteger(hello.resume.ownerGeneration) || hello.resume.ownerGeneration <= 0) { - throw new Error('resume.ownerGeneration must be a positive safe integer') - } - assertNonEmptyString(hello.resume.ownerLease, 'resume.ownerLease') - } - const flow = hello.capabilities?.outputFlowControl - if (!flow) { - return - } - if ( - !Array.isArray(flow.versions) || - flow.versions.length > MAX_CAPABILITY_VERSIONS || - flow.versions.some((version) => !Number.isSafeInteger(version) || version <= 0) - ) { - throw new Error('outputFlowControl.versions must contain positive safe integers') - } - if (!Number.isSafeInteger(flow.requestedWindowSu) || flow.requestedWindowSu <= 0) { - throw new Error('outputFlowControl.requestedWindowSu must be a positive safe integer') - } -} - -function helloFingerprint(hello: PtyConsumerSessionHello): string { - const flow = hello.capabilities?.outputFlowControl - return JSON.stringify({ - clientInstanceId: hello.clientInstanceId, - requestedRole: hello.requestedRole, - resume: hello.resume, - outputFlowControl: flow - ? { - versions: [...flow.versions].sort((a, b) => a - b), - requestedWindowSu: flow.requestedWindowSu - } - : undefined - }) -} - export class PtyConsumerSession { private readonly clients = new Map() private readonly now: () => number @@ -160,7 +118,7 @@ export class PtyConsumerSession { if (owner) { this.owner = owner } - return this.admissionFor(client) + return this.admissionFor(client, this.displacedOwnerFor(owner)) } close(connectionId: string): void { @@ -170,6 +128,17 @@ export class PtyConsumerSession { } this.clients.delete(connectionId) if (this.owner?.connectionId !== connectionId) { + // Why: a pending replacement can still roll back onto the owner it is displacing; restoring an + // 'active' record whose connection has since closed would wedge an owner that can never expire. + if ( + this.owner?.replaces?.connectionId === connectionId && + this.owner.replaces.state === 'active' + ) { + this.owner = { + ...this.owner, + replaces: { ...this.owner.replaces, state: 'disconnected', disconnectedAt: this.now() } + } + } return } if (this.owner.state === 'pending') { @@ -192,9 +161,13 @@ export class PtyConsumerSession { return client?.state === 'active' ? client.grant : null } - private admissionFor(client: ClientRecord): PtyConsumerSessionAdmission { + private admissionFor( + client: ClientRecord, + displacedOwner?: Readonly + ): PtyConsumerSessionAdmission { return { grant: client.grant, + ...(displacedOwner ? { displacedOwner } : {}), commitPublication: () => { if (client.publicationState !== 'pending') { return @@ -206,6 +179,7 @@ export class PtyConsumerSession { client.state = 'active' const owner = this.owner if (owner?.connectionId === this.connectionIdFor(client) && owner.state === 'pending') { + this.retireDisplacedOwner(owner.replaces) this.owner = { ...owner, state: 'active', replaces: undefined } } }, @@ -254,32 +228,46 @@ export class PtyConsumerSession { if (!hello.resume) { return null } - const recoveryMatches = - hello.resume.ownerGeneration === current.generation && - hello.resume.ownerLease === current.lease && - hello.clientInstanceId === current.clientInstanceId && - authentication.principal === current.principal - if (!recoveryMatches) { - throw Object.assign( - new Error('Owner recovery lease is stale or belongs to another principal'), - { code: PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR } - ) - } - if (current.state === 'pending') { - throw new Error('Owner grant publication is still pending') - } - if (current.state !== 'disconnected') { - throw new Error('Active owner cannot be replaced by recovery') - } + assertPtyConsumerOwnerRecovery(hello, authentication, current) + // Why an active owner is displaced rather than refused: the resume proof matched this owner's + // generation, lease, client instance, and principal on a *different* transport, so the requester is + // the same logical owner reconnecting. Waiting for the incumbent's socket to close is unbounded — + // a half-open connection after sleep/resume or NAT loss never gets there. return this.newOwner(hello, authentication, current) } + private displacedOwnerFor( + owner: OwnerRecord | null + ): Readonly | undefined { + const replaced = owner?.replaces + if (replaced?.state !== 'active') { + return undefined + } + const client = this.clients.get(replaced.connectionId) + if (client?.state !== 'active') { + return undefined + } + return Object.freeze({ connectionId: replaced.connectionId, grant: client.grant }) + } + + // Why: the displaced connection may still be writable (half-open), so revoke its grant the moment the + // replacement is published — a stale owner must not keep driving deliveries under the old generation. + private retireDisplacedOwner(replaced: OwnerRecord | undefined): void { + if (replaced?.state !== 'active') { + return + } + const client = this.clients.get(replaced.connectionId) + if (client?.state === 'active') { + client.state = 'displaced' + } + } + private newOwner( hello: PtyConsumerSessionHello, authentication: PtyConsumerAuthentication, replaces: OwnerRecord | null ): OwnerRecord { - const lease = this.createLease() + const lease = replaces?.lease ?? this.createLease() assertNonEmptyString(lease, 'ownerLease') return { connectionId: authentication.connectionId,