diff --git a/src/main/ssh/ssh-owner-recovery-retry.test.ts b/src/main/ssh/ssh-owner-recovery-retry.test.ts index 9b6775a5f48..9c8f4bd6a14 100644 --- a/src/main/ssh/ssh-owner-recovery-retry.test.ts +++ b/src/main/ssh/ssh-owner-recovery-retry.test.ts @@ -1,6 +1,9 @@ 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' +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'), { @@ -8,6 +11,12 @@ function publicationPendingError(): Error & { code: number } { }) } +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, @@ -28,7 +37,7 @@ describe('SSH owner recovery retry', () => { .mockRejectedValueOnce(publicationPendingError()) .mockResolvedValue('recovered') - const recovery = retrySshOwnerRecoveryWhilePublicationPending(attempt, openGate()) + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate()) await vi.advanceTimersByTimeAsync(75) await expect(recovery).resolves.toBe('recovered') @@ -40,7 +49,7 @@ describe('SSH owner recovery retry', () => { const error = publicationPendingError() const attempt = vi.fn<() => Promise>().mockRejectedValue(error) - const recovery = retrySshOwnerRecoveryWhilePublicationPending(attempt, openGate(), 60) + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate(), 60) const rejection = expect(recovery).rejects.toBe(error) await vi.advanceTimersByTimeAsync(60) @@ -52,19 +61,31 @@ describe('SSH owner recovery retry', () => { 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 - ) + 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 = retrySshOwnerRecoveryWhilePublicationPending(attempt, { + const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, { isCurrent: () => current, onClosed: (listener) => { close = listener diff --git a/src/main/ssh/ssh-owner-recovery-retry.ts b/src/main/ssh/ssh-owner-recovery-retry.ts index 9449853acdb..e5e78d9a4f2 100644 --- a/src/main/ssh/ssh-owner-recovery-retry.ts +++ b/src/main/ssh/ssh-owner-recovery-retry.ts @@ -1,7 +1,9 @@ -import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../../shared/pty-consumer-session' +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_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. +// 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 @@ -34,7 +36,7 @@ function waitForRetry(delayMs: number, gate: SshOwnerRecoveryRetryGate): Promise }) } -export async function retrySshOwnerRecoveryWhilePublicationPending( +export async function retrySshOwnerRecoveryWhileBlocked( attempt: () => Promise, gate: SshOwnerRecoveryRetryGate, waitMs: number = SSH_OWNER_RECOVERY_WAIT_MS @@ -45,9 +47,10 @@ export async function retrySshOwnerRecoveryWhilePublicationPending( try { return await attempt() } catch (error) { + const code = (error as { code?: unknown } | null | undefined)?.code if ( - (error as { code?: unknown } | null | undefined)?.code !== - PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR || + (code !== PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR && + code !== PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR) || !gate.isCurrent() ) { throw error diff --git a/src/main/ssh/ssh-relay-session-recovery-durability.test.ts b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts index be5cc07f80a..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,5 +1,8 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../../shared/pty-consumer-session' +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, @@ -195,6 +198,42 @@ describe('SshRelaySession consumer recovery durability', () => { } }) + 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({ diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index ea293e5b8ca..3f1c5105d29 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -81,7 +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 { retrySshOwnerRecoveryWhileBlocked } from './ssh-owner-recovery-retry' import { runRemoteOrcaCli } from './ssh-remote-orca-cli' import { acknowledgeRemoteOrcaCliPostOutput, @@ -1028,7 +1028,7 @@ export class SshRelaySession { outputFlowControl: { requestedWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU } } try { - return await retrySshOwnerRecoveryWhilePublicationPending( + return await retrySshOwnerRecoveryWhileBlocked( () => openSshPtyConsumerSession(mux, { ...options, diff --git a/src/relay/relay-pty-consumer-owner-displacement.test.ts b/src/relay/relay-pty-consumer-owner-displacement.test.ts index 1c7bd076b20..89059119f5f 100644 --- a/src/relay/relay-pty-consumer-owner-displacement.test.ts +++ b/src/relay/relay-pty-consumer-owner-displacement.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import { RelayDispatcher, type RelayClientSessionIdentity, @@ -6,7 +6,10 @@ import { } from './dispatcher' import { encodeJsonRpcFrame, MessageType } from './protocol' import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' -import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../shared/pty-consumer-session' +import { + PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR, + PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR +} from '../shared/pty-consumer-session' const endpointIdentity: RelayClientSessionIdentity = { principal: 'endpoint-principal', @@ -59,13 +62,14 @@ describe('relay PTY consumer owner displacement', () => { // 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 }, + { supportsWriteCallback: true, close: closeStaleTransport }, endpointIdentity ) const detached: number[] = [] @@ -120,6 +124,7 @@ describe('relay PTY consumer owner displacement', () => { 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 }) @@ -274,7 +279,7 @@ describe('relay PTY consumer owner displacement', () => { expect(detached).toContain(1) }) - it('serializes overlapping reconnect grants without invalidating the original proof', async () => { + it('fences an old proof while its replacement is live', async () => { const incumbentWrites: Buffer[] = [] dispatcher = new RelayDispatcher( (data, settle) => { @@ -340,12 +345,24 @@ describe('relay PTY consumer owner displacement', () => { ) await flushRequests() - expect(response(retryWrites, 4)!.result).toMatchObject({ + 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(firstReconnectClientId)).toBe('subscriber') expect(adapter.deliveryMode(retryClientId)).toBe('source-owner') }) }) diff --git a/src/relay/relay.ts b/src/relay/relay.ts index d5003bdf649..568d6f30f3c 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -10,7 +10,7 @@ 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 } from 'node:fs' import { RELAY_SENTINEL, FrameDecoder, @@ -597,6 +597,17 @@ async function main(): Promise { } stdoutDrainWaiters.add(cb) return () => stdoutDrainWaiters.delete(cb) + }, + close: () => { + stdoutAlive = false + flushStdoutDrainWaiters() + for (const fd of [process.stdin.fd, process.stdout.fd]) { + try { + closeSync(fd) + } catch { + // Already closed by the peer. + } + } } }, undefined, 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 6d299a7658b..1e359a11f53 100644 --- a/src/shared/pty-consumer-session-contract.ts +++ b/src/shared/pty-consumer-session-contract.ts @@ -4,6 +4,7 @@ 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' diff --git a/src/shared/pty-consumer-session.test.ts b/src/shared/pty-consumer-session.test.ts index d1999639e6f..ce6282e8664 100644 --- a/src/shared/pty-consumer-session.test.ts +++ b/src/shared/pty-consumer-session.test.ts @@ -1,6 +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, @@ -251,7 +252,7 @@ describe('PtyConsumerSession', () => { ) }) - it('retries overlapping recovery against the stable lease after publication commits', () => { + it('fences an old recovery generation after its replacement commits', () => { const session = createSession() const first = session.admit(ownerHello(), auth('connection-1')) first.commitPublication() @@ -266,7 +267,19 @@ describe('PtyConsumerSession', () => { ) replacement.commitPublication() - const retry = session.admit(ownerHello({ resume }), auth('connection-3')) + 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', @@ -278,6 +291,27 @@ describe('PtyConsumerSession', () => { 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')) diff --git a/src/shared/pty-consumer-session.ts b/src/shared/pty-consumer-session.ts index 4afad1dd822..3baf502acd8 100644 --- a/src/shared/pty-consumer-session.ts +++ b/src/shared/pty-consumer-session.ts @@ -1,7 +1,6 @@ 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, @@ -17,6 +16,7 @@ import { MAX_CAPABILITY_VERSIONS, validateHello } from './pty-consumer-session-hello' +import { assertPtyConsumerOwnerRecovery } from './pty-consumer-owner-recovery' export * from './pty-consumer-session-contract' @@ -228,23 +228,7 @@ export class PtyConsumerSession { if (!hello.resume) { return null } - const recoveryMatches = - 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') { - // 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 - }) - } + 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 —