diff --git a/src/main/index.ts b/src/main/index.ts index 3521070b55c..7ffd94df50e 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -251,6 +251,7 @@ 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 { detachAllSshSessionsForShutdown } from './ipc/ssh' import { PluginService } from './plugins/plugin-service' import { PluginKillListService } from './plugins/plugin-kill-list-service' import { getPluginsDataDir } from './plugins/plugin-discovery' @@ -2994,6 +2995,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() @@ -3036,6 +3038,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 7a7b01563e1..62f1ef26985 100644 --- a/src/main/ipc/ssh.test.ts +++ b/src/main/ipc/ssh.test.ts @@ -257,7 +257,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, @@ -557,6 +563,89 @@ 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('drains the replacement session a paused connect publishes 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 + }) + mockStore.markSshRemotePtyLeasesAsync.mockImplementationOnce(() => detachFlush) + + const replacement = handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' }) + // Why: let the replacement connect reach the flush before shutdown snapshots what to drain. + for (let tick = 0; tick < 10; tick++) { + await Promise.resolve() + } + expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'detached') + + mockConnectionManager.disconnectAll.mockClear() + const shutdown = detachAllSshSessionsForShutdown() + releaseDetach() + + await expect(replacement).rejects.toThrow() + await shutdown + + // Why two detaches: the old session's, plus the replacement the paused connect published after the + // first drain had already snapshotted activeSessions. + const detaches = mockStore.markSshRemotePtyLeasesAsync.mock.calls.filter( + (call) => call[1] === 'detached' + ) + expect(detaches).toHaveLength(2) + 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('ssh:importConfig returns imported targets', async () => { const imported: SshTarget[] = [ { id: 'ssh-imp', label: 'staging', host: 'staging.com', port: 22, username: '' } diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index 0dd32ce534e..75dfa1f7b9f 100644 --- a/src/main/ipc/ssh.ts +++ b/src/main/ipc/ssh.ts @@ -306,6 +306,16 @@ type ConnectAttempt = { const connectInFlight = new Map() const pendingTransportReconnects = new Set() +// Why: once the committed quit path starts draining SSH, a connect that has not yet published its +// session must not start one — the drain has already snapshotted what it will tear down. +let sshShutdownFenced = false + +function assertSshConnectsNotFenced(): void { + if (sshShutdownFenced) { + throw new Error('SSH connects are closed for app shutdown') + } +} + function invalidateConnectAttempt(targetId: string): void { rotateSshProviderAuthority(targetId) pendingTransportReconnects.delete(targetId) @@ -1042,6 +1052,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) @@ -1333,6 +1346,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, () => @@ -1393,6 +1408,8 @@ export function registerSshHandlers( testingTargets.add(args.targetId) try { + // 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) @@ -1510,6 +1527,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?.() + }) + ]) +} + +export async function detachAllSshSessionsForShutdown(): Promise { + // Why synchronous, before the first await: connectTarget registers its attempt and checks the fence + // in one uninterrupted block, so fencing and snapshotting here leaves no gap for an unjoined connect. + sshShutdownFenced = true + const inFlightWork: Promise[] = [ + ...[...connectInFlight.values()].map((attempt) => attempt.promise), + ...resetRelayInFlight.values() + ] + 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 @@ -1537,6 +1609,7 @@ export async function resetSshHandlerStateForTests(): Promise { resetRelayInFlight.clear() testingTargets.clear() credentialRequestedForTarget.clear() + sshShutdownFenced = false await connectionManager?.disconnectAll() portForwardManager?.dispose() 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..9b6775a5f48 --- /dev/null +++ b/src/main/ssh/ssh-owner-recovery-retry.test.ts @@ -0,0 +1,85 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../../shared/pty-consumer-session' +import { retrySshOwnerRecoveryWhilePublicationPending } 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 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 = retrySshOwnerRecoveryWhilePublicationPending(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 = retrySshOwnerRecoveryWhilePublicationPending(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(retrySshOwnerRecoveryWhilePublicationPending(attempt, openGate())).rejects.toBe( + error + ) + expect(attempt).toHaveBeenCalledOnce() + }) + + 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 = retrySshOwnerRecoveryWhilePublicationPending(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..ea460f63d13 --- /dev/null +++ b/src/main/ssh/ssh-owner-recovery-retry.ts @@ -0,0 +1,65 @@ +import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../../shared/pty-consumer-session' + +// Why this budget is safe: the relay displaces a still-attached owner outright, so the only refusal +// left is an incumbent grant whose publication has not settled — bounded by one relay response write. +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 retrySshOwnerRecoveryWhilePublicationPending( + 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) { + if ( + (error as { code?: unknown }).code !== PTY_CONSUMER_OWNER_RECOVERY_PENDING_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..9a28d491512 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,5 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' +import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../../shared/pty-consumer-session' import { SshRelaySession } from './ssh-relay-session' import { createMismatchedOwnerRecoveryError, @@ -154,6 +155,46 @@ 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('keeps destructive disposal pending until consumer recovery is removed', async () => { const { mockStore, mockPortForward, getMainWindow } = createMockDeps() vi.mocked(mockStore.getSshPtyConsumerRecovery).mockReturnValue({ diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index c8440f6c239..e6a821f74a8 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -81,6 +81,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 { retrySshOwnerRecoveryWhilePublicationPending } from './ssh-owner-recovery-retry' import { runRemoteOrcaCli } from './ssh-remote-orca-cli' import { acknowledgeRemoteOrcaCliPostOutput, @@ -1027,17 +1028,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 retrySshOwnerRecoveryWhilePublicationPending( + () => + openSshPtyConsumerSession(mux, { + ...options, + ...(previousOwner + ? { + resume: { + ownerGeneration: previousOwner.ownerGeneration, + ownerLease: previousOwner.ownerLease + } + } + : {}) + }), + { + isCurrent: () => ownsAttempt() && !mux.isDisposed(), + onClosed: (listener) => mux.onDispose(listener) + } + ) } catch (error) { if ( !previousOwner || 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..ede575333d5 --- /dev/null +++ b/src/relay/relay-pty-consumer-owner-displacement.test.ts @@ -0,0 +1,213 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { + RelayDispatcher, + type RelayClientSessionIdentity, + type SinkWriteSettlement +} from './dispatcher' +import { encodeJsonRpcFrame, MessageType } from './protocol' +import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' + +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[] = [] + dispatcher = new RelayDispatcher( + (data, settle) => { + staleWrites.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 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) + + // 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('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) + }) +}) 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-session-contract.ts b/src/shared/pty-consumer-session-contract.ts index 661138896cd..6d299a7658b 100644 --- a/src/shared/pty-consumer-session-contract.ts +++ b/src/shared/pty-consumer-session-contract.ts @@ -1,6 +1,9 @@ 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 type PtyConsumerRole = 'session-owner' | 'subscriber' @@ -41,8 +44,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..bbfed831e1a 100644 --- a/src/shared/pty-consumer-session.test.ts +++ b/src/shared/pty-consumer-session.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from 'vitest' import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR, PtyConsumerSession, type PtyConsumerAuthentication, @@ -145,11 +146,93 @@ describe('PtyConsumerSession', () => { }) }) - 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 +243,12 @@ describe('PtyConsumerSession', () => { }), auth('connection-2') ) - ).toThrow('Active owner') + ).toThrow( + expect.objectContaining({ + code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + message: expect.stringContaining('still pending') + }) + ) }) it('types mismatched recovery without disturbing principal or lease ownership', () => { diff --git a/src/shared/pty-consumer-session.ts b/src/shared/pty-consumer-session.ts index 953b77c3a8d..084c6638ec5 100644 --- a/src/shared/pty-consumer-session.ts +++ b/src/shared/pty-consumer-session.ts @@ -1,25 +1,31 @@ import { randomUUID } from 'node:crypto' import { PTY_CONSUMER_OWNER_GRACE_MS, + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, 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' 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 } } }, @@ -266,14 +240,45 @@ export class PtyConsumerSession { ) } 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') + // Why coded: unlike a still-attached owner this cannot be fenced — the incumbent grant is mid + // publication, so there is nothing to displace yet. Bounded by one response write, so retryable. + throw Object.assign(new Error('Owner grant publication is still pending'), { + code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR + }) } + // 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,