diff --git a/config/reliability-gates.jsonc b/config/reliability-gates.jsonc index f5aa19d93e7..f89c1b11c86 100644 --- a/config/reliability-gates.jsonc +++ b/config/reliability-gates.jsonc @@ -3373,8 +3373,10 @@ "coveredPlatforms": [ "macos" ], - "coveredProviders": [], - "coverageNotes": "Local macOS evidence covers the existing main-process pending-output caps plus deterministic runtime path-provenance history reuse and saturated-tail wait detection. Daemon stream write(false)/drain contracts, cross-session drain priority, and bounded queued tails arrive with the pending perf slice; live flood/latency artifacts remain gaps.", + "coveredProviders": [ + "ssh" + ], + "coverageNotes": "Local macOS evidence covers main-process pending-output caps, runtime path-provenance history reuse, and saturated-tail wait detection. A macOS-hosted Docker OpenSSH run exercises the deployed Linux relay and SSH provider with an exact 256 KiB source-credit plateau, concurrent PTY typing, fixed-size filesystem frames, Git churn, and owner-lease reconnect. It does not cover local/daemon, paired remote-runtime, headed desktop, headless orca serve, prior-version daemons, WSL, Windows named pipes/ConPTY, folder workspaces, or mixed-version clients.", "motivatingLinks": [ "https://github.com/stablyai/orca/pull/6836", "https://github.com/stablyai/orca/pull/6858", @@ -3386,13 +3388,15 @@ "commands": [ "pnpm exec vitest run --config config/vitest.config.ts src/main/ipc/pty.test.ts", "pnpm exec vitest run --config config/vitest.config.ts src/main/runtime/orca-runtime-path-candidate-history.test.ts", - "pnpm exec vitest run --config config/vitest.config.ts src/main/runtime/orca-runtime.test.ts src/main/runtime/orca-runtime-tail-wait-memo.test.ts" + "pnpm exec vitest run --config config/vitest.config.ts src/main/runtime/orca-runtime.test.ts src/main/runtime/orca-runtime-tail-wait-memo.test.ts", + "ORCA_E2E_SSH_DOCKER=1 ORCA_SSH_PTY_SOURCE_CREDIT_V1=1 pnpm exec playwright test tests/e2e/ssh-docker-relay-perf.spec.ts --config tests/playwright.config.ts --project electron-headless --workers=1" ], "testFiles": [ "src/main/ipc/pty.test.ts", "src/main/runtime/orca-runtime-path-candidate-history.test.ts", "src/main/runtime/orca-runtime.test.ts", - "src/main/runtime/orca-runtime-tail-wait-memo.test.ts" + "src/main/runtime/orca-runtime-tail-wait-memo.test.ts", + "tests/e2e/ssh-docker-relay-perf.spec.ts" ], "assertionRefs": [ { @@ -3426,6 +3430,14 @@ "returns a blocked wait result for generic Codex interactive prompts", "resolves tui-idle when a stale Codex prompt is followed by the ready header" ] + }, + { + "file": "tests/e2e/ssh-docker-relay-perf.spec.ts", + "assertions": [ + "a stalled renderer ACK produces an exact 256 KiB negotiated source-credit plateau while a second SSH PTY remains responsive", + "fixed-size filesystem frames and Git churn complete without stream corruption while active typing remains within budget", + "the negotiated owner lease reconnects and the existing SSH workspace terminal remains usable" + ] } ], "evidenceRuns": [ @@ -3455,6 +3467,15 @@ "result": "passed", "durationSeconds": 21.33, "summary": "2 test files passed, 609 tests passed on pushed commit 4fb14eac3897." + }, + { + "date": "2026-07-27", + "runner": "local", + "platform": "macos", + "command": "ORCA_E2E_SSH_DOCKER=1 ORCA_SSH_PTY_SOURCE_CREDIT_V1=1 pnpm exec playwright test tests/e2e/ssh-docker-relay-perf.spec.ts --config tests/playwright.config.ts --project electron-headless --workers=1", + "result": "passed", + "durationSeconds": 57.4, + "summary": "Four Docker OpenSSH/deployed-relay tests passed: direct typing median/worst 106.2/114.0 ms, stalled-ACK typing 3.4/6.1 ms at an exact 262,144-source-unit plateau, fixed-size filesystem/Git churn 141.8/153.6 ms with 85 completed reads, and terminal owner-lease reconnect in 15.7 seconds." } ], "runtimeBudget": { @@ -3463,7 +3484,7 @@ }, "flakeHistory": { "status": "unknown", - "evidence": "Focused main-process backlog tests are deterministic unit slices. The combined perf soak still needs runtime history before promotion." + "evidence": "Focused main-process backlog tests are deterministic unit slices. On 2026-07-27 the source-credit reconnect case passed once in the full four-test Docker suite (15.7 seconds) and three consecutive isolated repeat-each runs (21.1, 19.8, and 15.3 seconds); the combined perf soak still needs longer runtime history before promotion." }, "redGreenEvidence": { "status": "partial", @@ -3471,7 +3492,7 @@ }, "performanceBudget": { "required": true, - "evidence": "Suggested ceilings: renderer in-flight <=8MB total, <=512KB per PTY plus active reserve, renderer queued chars <=2MB, dropped backlogs 0, hidden restore <=1000ms, active key median/worst <=75ms/300ms in perf scenarios. With 1,024 retained provenance candidates, 4,096 pathless chunks dropped from 78.84ms and 4,096 array replacements to 1.74ms and zero replacements. Brace-free 1 KiB output dropped from 20.32/79.12/318.23ms to 0.31/0.64/2.35ms across 4,096/16,384/65,536 chunks. Repeated ordinary 252,000-character tail checks dropped from 96.34/464.20/941.79/3,796.70ms to 3.21/15.68/31.34/129.72ms across 100/500/1,000/4,096 updates." + "evidence": "Suggested ceilings: renderer in-flight <=8MB total, <=512KB per PTY plus active reserve, renderer queued chars <=2MB, dropped backlogs 0, hidden restore <=1000ms, active key median/worst <=75ms/300ms in perf scenarios. The Docker SSH gate enforces median/worst typing below 500/2,000 ms and observed 3.4/6.1 ms during an exact 256 KiB stalled-credit plateau plus 141.8/153.6 ms during 85 completed filesystem/Git operations. With 1,024 retained provenance candidates, 4,096 pathless chunks dropped from 78.84ms and 4,096 array replacements to 1.74ms and zero replacements. Brace-free 1 KiB output dropped from 20.32/79.12/318.23ms to 0.31/0.64/2.35ms across 4,096/16,384/65,536 chunks. Repeated ordinary 252,000-character tail checks dropped from 96.34/464.20/941.79/3,796.70ms to 3.21/15.68/31.34/129.72ms across 100/500/1,000/4,096 updates." }, "promotionCriteria": [ "Split daemon stream backpressure into a deterministic provider/IPC contract if full E2E is flaky.", @@ -3479,7 +3500,10 @@ "Keep stress variants non-blocking until stable runtime history exists." ], "knownGaps": [ - "Daemon stream write(false)/drain/cleanup and bounded queued-tail contracts are not registered in the current command.", + "Local daemon/provider and prior-version daemon topologies were not executed.", + "Headed paired desktop, headless orca serve, folder workspace, and mixed-version client topologies were not executed.", + "Physical WSL and Windows named-pipe/ConPTY topologies were unavailable and were not executed.", + "Ubuntu 20.04/glibc 2.31 packaging was not physically executed; cross-target relay bundling is build evidence only.", "Runtime provenance coverage is deterministic and does not include a live high-throughput provider artifact.", "Terminal-wait scale evidence is deterministic and does not yet include a live saturated-tail event-loop artifact.", "Current command does not prove renderer parse pressure, scheduler queue depth, event-loop delay, or active key latency.", diff --git a/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts b/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts index b806f43da74..b57be2544b6 100644 --- a/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts +++ b/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts @@ -21,7 +21,7 @@ function fakeChild(): FakeChild { resume: ReturnType } stderr: EventEmitter - stdin: { write: ReturnType } + stdin: EventEmitter & { write: ReturnType } kill: ReturnType } child.stdout = Object.assign(new EventEmitter(), { @@ -29,7 +29,7 @@ function fakeChild(): FakeChild { resume: vi.fn() }) child.stderr = new EventEmitter() - child.stdin = { write: vi.fn(() => true) } + child.stdin = Object.assign(new EventEmitter(), { write: vi.fn(() => true) }) child.kill = vi.fn() return child as unknown as FakeChild } @@ -159,4 +159,25 @@ describe('waitForWslRelaySentinel', () => { expect(err.message).toContain('E_FAIL') expect(err.message).not.toContain(String.fromCharCode(0)) }) + + it('forwards write(false), callback settlement, and drain from WSL stdin', async () => { + const child = fakeChild() + const callback = vi.fn() + const drain = vi.fn() + const writeMock = child.stdin.write as ReturnType + writeMock.mockImplementation((_data, onWritten) => { + onWritten(null) + return false + }) + const promise = waitForWslRelaySentinel(child) + emitStdout(child, RELAY_SENTINEL) + const transport = await promise + + transport.onDrain?.(drain) + expect(transport.write(Buffer.from('frame'), callback)).toBe(false) + expect(callback).toHaveBeenCalledWith({ ok: true }) + child.stdin.emit('drain') + expect(drain).toHaveBeenCalledOnce() + expect(transport.supportsWriteSettlement).toBe(true) + }) }) diff --git a/src/main/agent-hooks/wsl-hook-relay-sentinel.ts b/src/main/agent-hooks/wsl-hook-relay-sentinel.ts index db4c48257a7..5e1b5b5a872 100644 --- a/src/main/agent-hooks/wsl-hook-relay-sentinel.ts +++ b/src/main/agent-hooks/wsl-hook-relay-sentinel.ts @@ -124,17 +124,14 @@ export function waitForWslRelaySentinel( } const transport: MultiplexerTransport = { write: (data, onSettled) => { - try { - child.stdin.write(data, (error?: Error | null) => { - onSettled?.(error ? { ok: false, error } : { ok: true }) - }) - } catch (error) { - onSettled?.({ - ok: false, - error: error instanceof Error ? error : new Error(String(error)) - }) - // Channel already closing — mux close handling takes over. - } + return child.stdin.write(data, (error?: Error | null) => { + onSettled?.(error ? { ok: false, error } : { ok: true }) + }) + }, + supportsWriteSettlement: true, + onDrain: (cb) => { + child.stdin.on('drain', cb) + return () => child.stdin.off('drain', cb) }, onData: (cb) => { dataCallbacks.push(cb) diff --git a/src/main/ipc/pty.test.ts b/src/main/ipc/pty.test.ts index 58b841508bd..4ff670577e9 100644 --- a/src/main/ipc/pty.test.ts +++ b/src/main/ipc/pty.test.ts @@ -247,6 +247,7 @@ import { resolveWindowsShellLaunchArgs } from '../providers/windows-shell-args' import { _resetWslCachesForTests, _setWslCachesForTests } from '../wsl' import { wslHookRelayManager } from '../agent-hooks/wsl-hook-relay-manager' import { acquireWatcherRemovalGate } from './watcher-removal-gate' +import { acceptSshPtyOutputData } from './ssh-pty-output-intake-registry' // Why: Windows resolves a bare PowerShell name to an absolute exe before ConPTY, else CreateProcessW fails with error 5 (PR #6537 / #5161). const RESOLVED_WINDOWS_POWERSHELL = 'C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe' @@ -11467,6 +11468,73 @@ describe('registerPtyHandlers', () => { } }) + it('keeps negotiated source-credit overflow off the legacy PTY-global pause path', async () => { + vi.useFakeTimers() + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + try { + const provider = installObservableDaemonTestProvider() + let modelSequence = 0 + const runtime = { + setPtyController: vi.fn(), + setRemoteTerminalSourceRangeConsumerHooks: vi.fn(), + getPtyOutputSequence: vi.fn(() => modelSequence), + onPtyData: vi.fn( + (_id: string, data: string, _at: number, rawLength = data.length) => + (modelSequence += rawLength) + ), + acceptPtyDataBounded: vi.fn( + (_id: string, _data: string, _at: number, rawLength: number) => { + modelSequence += rawLength + return { sequence: modelSequence, completion: Promise.resolve() } + } + ) + } + registerPtyHandlers(mainWindow as never, runtime as never) + mainWindow.webContents.send.mockClear() + + const sourceChunk = 's'.repeat(128 * 1024) + for (let index = 0; index < 17; index++) { + const sourceStartSu = index * sourceChunk.length + await acceptSshPtyOutputData({ + id: 'source-credit-pty', + data: sourceChunk, + providerGeneration: 41, + ptyIncarnation: 'source-incarnation', + rawLength: sourceChunk.length, + transformed: false, + source: { + relayPtyId: 'relay-source-pty', + spanId: `source-token:${sourceStartSu}:${sourceStartSu + sourceChunk.length}`, + clientGeneration: 2, + ownerGeneration: 3, + deliveryToken: 'source-token', + sourceStartSu, + sourceEndSu: sourceStartSu + sourceChunk.length + } + }) + } + + expect(getPtyRendererDeliveryDebugSnapshot()).toMatchObject({ + pendingPtyCount: 1, + pendingChars: 0 + }) + expect(provider.pauseProducer).not.toHaveBeenCalledWith('source-credit-pty') + expect(provider.resumeProducer).not.toHaveBeenCalledWith('source-credit-pty') + + provider.emitData('legacy-pty', 'l'.repeat(320 * 1024)) + expect(provider.pauseProducer).toHaveBeenCalledTimes(1) + expect(provider.pauseProducer).toHaveBeenCalledWith('legacy-pty') + expect(provider.pauseProducer).not.toHaveBeenCalledWith('unrelated-pty') + + vi.runAllTimers() + expect(provider.resumeProducer).toHaveBeenCalledTimes(1) + expect(provider.resumeProducer).toHaveBeenCalledWith('legacy-pty') + } finally { + errorSpy.mockRestore() + vi.useRealTimers() + } + }) + it('resumes a paused producer when the PTY exits before draining', async () => { vi.useFakeTimers() try { diff --git a/src/main/ipc/pty.ts b/src/main/ipc/pty.ts index 11b1dfca7c0..daf15c65a03 100644 --- a/src/main/ipc/pty.ts +++ b/src/main/ipc/pty.ts @@ -1842,6 +1842,7 @@ export function registerPtyHandlers( let sshOutputIntake: SshPtyOutputIntake | null = null // Why: resuming a paused producer during exit can synchronously emit; those bytes must not queue behind pty:exit. const rendererExitingPtyIds = new Set() + const rendererCreditBeforeExitByPty = new Map() const rendererDeliveryRestoreNeededPtys = new Set() function transitionHiddenRendererPtyDeliveryState(id: string, hidden: boolean) { @@ -1927,11 +1928,18 @@ export function registerPtyHandlers( pauseProducer: (id) => tryGetProviderForPty(id)?.pauseProducer?.(id), resumeProducer: (id) => tryGetProviderForPty(id)?.resumeProducer?.(id) }) + const sourceCreditPendingPtys = new Set() function updateProducerFlowControl(id: string): void { if (!PRODUCER_FLOW_CONTROL_ENABLED) { return } + if (sourceCreditPendingPtys.has(id)) { + if (pendingData.get(id)) { + return + } + sourceCreditPendingPtys.delete(id) + } producerFlowControl.update(id, pendingData.get(id)?.data.length ?? 0) } @@ -2009,6 +2017,7 @@ export function registerPtyHandlers( } } pendingData.clear() + sourceCreditPendingPtys.clear() } function readCurrentPtyRendererDeliveryDebugSnapshot(): PtyRendererDeliveryDebugSnapshot { @@ -2856,6 +2865,12 @@ export function registerPtyHandlers( } rendererExitingPtyIds.add(payload.id) try { + if (!rendererCreditBeforeExitByPty.has(payload.id)) { + rendererCreditBeforeExitByPty.set( + payload.id, + getRendererInFlightCharsForPty(payload.id) > 0 + ) + } // Why flush before exit: the renderer tears down the terminal on pty:exit, so any batched output not yet flushed would be silently lost. const remaining = pendingData.delete(payload.id) clearFlushTimerIfIdle() @@ -2893,6 +2908,7 @@ export function registerPtyHandlers( function finalizePtyExitForRenderer(payload: { id: string; code: number }): void { if (mainWindow.isDestroyed()) { + rendererCreditBeforeExitByPty.delete(payload.id) return } if (rendererExitingPtyIds.has(payload.id)) { @@ -2900,9 +2916,13 @@ export function registerPtyHandlers( } rendererExitingPtyIds.add(payload.id) try { - const hadReleasableRendererCredit = getRendererInFlightCharsForPty(payload.id) > 0 + const hadReleasableRendererCredit = + rendererCreditBeforeExitByPty.get(payload.id) ?? + getRendererInFlightCharsForPty(payload.id) > 0 + rendererCreditBeforeExitByPty.delete(payload.id) // Why resume a dead PTY (no-op): avoid leaving a stale paused mark behind for a reused id. producerFlowControl.release(payload.id) + sourceCreditPendingPtys.delete(payload.id) pendingOverflowMarkedPtys.delete(payload.id) rendererDeliveryRestoreNeededPtys.delete(payload.id) lastInputAtByPty.delete(payload.id) @@ -3002,6 +3022,9 @@ export function registerPtyHandlers( markHiddenRendererResizeOutputDelivered(payload.id) } const overflowMarkedBeforeAppend = pendingOverflowMarkedPtys.has(payload.id) + if (projection?.desktopSpan) { + sourceCreditPendingPtys.add(payload.id) + } const pending = appendPendingPtyData( payload.id, pendingData.get(payload.id), @@ -3148,12 +3171,12 @@ export function registerPtyHandlers( }, publishSourceAck: publishSshPtySourceAck }) - runtime?.setRemoteTerminalSourceRangeConsumerHooks( + runtime?.setRemoteTerminalSourceRangeConsumerHooks?.( sshOutputIntake.getRemoteSourceRangeConsumerHooks() ) const cleanupSshOutputIntakeRegistry = installSshPtyOutputIntake(sshOutputIntake) sshOutputIntakeCleanup = () => { - runtime?.setRemoteTerminalSourceRangeConsumerHooks(null) + runtime?.setRemoteTerminalSourceRangeConsumerHooks?.(null) cleanupSshOutputIntakeRegistry() } diff --git a/src/main/ipc/ssh-pty-output-intake-registry.ts b/src/main/ipc/ssh-pty-output-intake-registry.ts index e55f5f253e9..103171f9557 100644 --- a/src/main/ipc/ssh-pty-output-intake-registry.ts +++ b/src/main/ipc/ssh-pty-output-intake-registry.ts @@ -4,6 +4,7 @@ import type { SshPtyOutputExitEvent, SshPtyOutputReceipt } from './ssh-pty-output-intake-contract' +import type { SshPtyAcceptedSourceCheckpoint } from './ssh-pty-output-source-obligations' import type { PtySourceCreditAckBatch } from '../../shared/pty-source-credit-contract' let installedIntake: SshPtyOutputIntake | null = null @@ -47,6 +48,12 @@ export function closeSshPtyOutputGeneration(providerGeneration: number, reason: installedIntake?.closeGeneration(providerGeneration, reason) } +export function getSshPtyAcceptedSourceCheckpoints( + providerGeneration: number +): readonly SshPtyAcceptedSourceCheckpoint[] { + return installedIntake?.getAcceptedSourceCheckpoints(providerGeneration) ?? [] +} + export function installSshPtySourceAckPublisher( providerGeneration: number, publish: SshPtySourceAckPublisher diff --git a/src/main/ipc/ssh-pty-output-intake.test.ts b/src/main/ipc/ssh-pty-output-intake.test.ts index 5d61e39c556..370894a34aa 100644 --- a/src/main/ipc/ssh-pty-output-intake.test.ts +++ b/src/main/ipc/ssh-pty-output-intake.test.ts @@ -239,6 +239,37 @@ describe('SshPtyOutputIntake', () => { expect(Object.isFrozen(projections[0]?.desktopSpan)).toBe(true) }) + it('exports only the model-settled source boundary before generation close', async () => { + const harness = createHarness() + const receipt = harness.intake.acceptData( + event({ + source: { + spanId: 'recovery-span-1', + clientGeneration: 3, + ownerGeneration: 4, + deliveryToken: 'delivery-token-1', + sourceStartSu: 10, + sourceEndSu: 14 + } + }) + ) + + expect(harness.intake.getAcceptedSourceCheckpoints(1)[0]?.acceptedSourceEndSu).toBe(10) + harness.completions[0]!.resolve() + await receipt + expect(harness.intake.getAcceptedSourceCheckpoints(1)).toEqual([ + { + id: 'pty-1', + providerGeneration: 1, + clientGeneration: 3, + ownerGeneration: 4, + ptyIncarnation: 'incarnation-1', + deliveryToken: 'delivery-token-1', + acceptedSourceEndSu: 14 + } + ]) + }) + it('rolls back projection staging when model capture throws synchronously', async () => { const project = vi.fn() const harness = createHarness({ diff --git a/src/main/ipc/ssh-pty-output-intake.ts b/src/main/ipc/ssh-pty-output-intake.ts index b496f2bfa7c..4cbaddb7088 100644 --- a/src/main/ipc/ssh-pty-output-intake.ts +++ b/src/main/ipc/ssh-pty-output-intake.ts @@ -222,6 +222,10 @@ export class SshPtyOutputIntake { return this.sourceObligations.remoteHooks } + getAcceptedSourceCheckpoints(providerGeneration: number) { + return this.sourceObligations.acceptedCheckpoints(providerGeneration) + } + getDebugSnapshot() { return { model: this.admission.getDebugSnapshot(), diff --git a/src/main/ipc/ssh-pty-output-source-obligations.ts b/src/main/ipc/ssh-pty-output-source-obligations.ts index 92306f012eb..804cc7d0fb6 100644 --- a/src/main/ipc/ssh-pty-output-source-obligations.ts +++ b/src/main/ipc/ssh-pty-output-source-obligations.ts @@ -20,6 +20,16 @@ export type SshPtyOutputSourceReservation = Readonly<{ span: PtySourceSpan }> +export type SshPtyAcceptedSourceCheckpoint = Readonly<{ + id: string + providerGeneration: number + clientGeneration: number + ownerGeneration: number + ptyIncarnation: string + deliveryToken: string + acceptedSourceEndSu: number +}> + export class SshPtyOutputSourceObligations { private readonly coordinator: SshPtySourceObligationCoordinator private readonly remoteConsumers: SshPtyRemoteSourceRangeConsumers @@ -139,6 +149,27 @@ export class SshPtyOutputSourceObligations { } } + acceptedCheckpoints(providerGeneration: number): readonly SshPtyAcceptedSourceCheckpoint[] { + const checkpoints: SshPtyAcceptedSourceCheckpoint[] = [] + for (const identity of this.identityByPty.values()) { + if (identity.providerGeneration !== providerGeneration) { + continue + } + checkpoints.push( + Object.freeze({ + id: identity.id, + providerGeneration: identity.providerGeneration, + clientGeneration: identity.clientGeneration, + ownerGeneration: identity.ownerGeneration, + ptyIncarnation: identity.ptyIncarnation, + deliveryToken: identity.deliveryToken, + acceptedSourceEndSu: this.coordinator.modelAcceptedEnd(identity) + }) + ) + } + return Object.freeze(checkpoints) + } + dispose(): void { this.coordinator.dispose() } diff --git a/src/main/ipc/ssh-pty-source-obligation-coordinator.ts b/src/main/ipc/ssh-pty-source-obligation-coordinator.ts index a52a034a79c..9d9d0ef4222 100644 --- a/src/main/ipc/ssh-pty-source-obligation-coordinator.ts +++ b/src/main/ipc/ssh-pty-source-obligation-coordinator.ts @@ -119,6 +119,10 @@ export class SshPtySourceObligationCoordinator { return this.ledger.snapshot(identity) } + modelAcceptedEnd(identity: PtySourceDeliveryIdentity): number { + return this.ledger.modelAcceptedEnd(identity) + } + obligation(spanId: string, consumer: SshPtySourceConsumerId): SshPtySourceObligationState { return this.ledger.obligation(spanId, consumer) } diff --git a/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts b/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts index 42c93610736..a329474b5bb 100644 --- a/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts +++ b/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts @@ -111,6 +111,36 @@ describe('SshPtySourceObligationLedger', () => { expect(ledger.snapshot(owner)).toMatchObject({ ackPublishedEndSu: 4, openSpans: 0 }) }) + it('checkpoints only the contiguous model-settled prefix', () => { + const ledger = new SshPtySourceObligationLedger() + const owner = identity() + ledger.open(owner) + commitSpan(ledger, owner, span(owner, 'span-1', 0, 'abcd')) + commitSpan(ledger, owner, span(owner, 'span-2', 4, 'efgh')) + + expect(ledger.modelAcceptedEnd(owner)).toBe(0) + ledger.settle('span-2', 'model', 'out-of-order') + expect(ledger.modelAcceptedEnd(owner)).toBe(0) + ledger.settle('span-1', 'model', 'emulator-receipt') + expect(ledger.modelAcceptedEnd(owner)).toBe(8) + }) + + it('continues checkpoints from the reclaimed ACK-published prefix', () => { + const ledger = new SshPtySourceObligationLedger() + const owner = identity() + ledger.open(owner) + commitSpan(ledger, owner, span(owner, 'span-1', 0, 'abcd')) + commitSpan(ledger, owner, span(owner, 'span-2', 4, 'efgh')) + ledger.settle('span-1', 'model', 'emulator-receipt') + ledger.settle('span-1', 'desktop', 'renderer-parse') + ledger.queueAck(owner)!.onSettled({ ok: true }) + + expect(ledger.snapshot(owner)).toMatchObject({ ackPublishedEndSu: 4, openSpans: 1 }) + expect(ledger.modelAcceptedEnd(owner)).toBe(4) + ledger.settle('span-2', 'model', 'emulator-receipt') + expect(ledger.modelAcceptedEnd(owner)).toBe(8) + }) + it('requires an exact transfer fence before a desktop obligation becomes terminal', () => { const ledger = new SshPtySourceObligationLedger() const owner = identity() diff --git a/src/main/ipc/ssh-pty-source-obligation-ledger.ts b/src/main/ipc/ssh-pty-source-obligation-ledger.ts index 3e3501fc013..f45e6296c51 100644 --- a/src/main/ipc/ssh-pty-source-obligation-ledger.ts +++ b/src/main/ipc/ssh-pty-source-obligation-ledger.ts @@ -18,7 +18,6 @@ import type { } from './ssh-pty-source-obligation-contract' import { createSshPtySourceAckPublication } from './ssh-pty-source-ack-publication' import { - advanceSourceTerminalEnd, beginSourceExitTimeout, cancelOpenSourceObligations, closeAllSourceTokens, @@ -34,6 +33,12 @@ import { type ReservationRecord, type TokenRecord } from './ssh-pty-source-obligation-state' +import { + cancelSourceObligationTransfer, + commitSourceObligationTransfer, + modelAcceptedSourceEnd, + transitionOpenSourceObligation +} from './ssh-pty-source-obligation-transitions' export type { SshPtySourceAckPublication, @@ -148,32 +153,11 @@ export class SshPtySourceObligationLedger { } commitTransfer(spanId: string, consumer: SshPtySourceConsumerId): boolean { - const { token, span } = requireSourceSpan(this.spanOwners, spanId) - const current = span.obligations.get(consumer) - if (current?.state !== 'transferring') { - return false - } - span.obligations.set( - consumer, - Object.freeze({ - state: 'transferred', - to: current.to, - reason: current.reason - }) - ) - advanceSourceTerminalEnd(token) - return true + return commitSourceObligationTransfer(this.spanOwners, spanId, consumer) } cancelTransfer(spanId: string, consumer: SshPtySourceConsumerId, reason: string): boolean { - const { token, span } = requireSourceSpan(this.spanOwners, spanId) - const current = span.obligations.get(consumer) - if (current?.state !== 'transferring') { - return false - } - span.obligations.set(consumer, Object.freeze({ state: 'canceled', reason })) - advanceSourceTerminalEnd(token) - return true + return cancelSourceObligationTransfer(this.spanOwners, spanId, consumer, reason) } queueAck(identity: PtySourceDeliveryIdentity): SshPtySourceAckPublication | null { @@ -263,6 +247,10 @@ export class SshPtySourceObligationLedger { throw new Error('Unknown or stale SSH PTY source token') } + modelAcceptedEnd(identity: PtySourceDeliveryIdentity): number { + return modelAcceptedSourceEnd(this.requireToken(identity)) + } + obligation(spanId: string, consumer: SshPtySourceConsumerId): SshPtySourceObligationState { const obligation = requireSourceSpan(this.spanOwners, spanId).span.obligations.get(consumer) if (!obligation) { @@ -280,13 +268,7 @@ export class SshPtySourceObligationLedger { consumer: SshPtySourceConsumerId, next: SshPtySourceObligationState ): boolean { - const { token, span } = requireSourceSpan(this.spanOwners, spanId) - if (span.obligations.get(consumer)?.state !== 'open') { - return false - } - span.obligations.set(consumer, next) - advanceSourceTerminalEnd(token) - return true + return transitionOpenSourceObligation(this.spanOwners, spanId, consumer, next) } private maybeClose(token: TokenRecord): void { diff --git a/src/main/ipc/ssh-pty-source-obligation-transitions.ts b/src/main/ipc/ssh-pty-source-obligation-transitions.ts new file mode 100644 index 00000000000..f136540cbb4 --- /dev/null +++ b/src/main/ipc/ssh-pty-source-obligation-transitions.ts @@ -0,0 +1,71 @@ +import type { + SshPtySourceConsumerId, + SshPtySourceObligationState +} from './ssh-pty-source-obligation-contract' +import { + advanceSourceTerminalEnd, + requireSourceSpan, + type TokenRecord +} from './ssh-pty-source-obligation-state' + +export function transitionOpenSourceObligation( + spanOwners: ReadonlyMap, + spanId: string, + consumer: SshPtySourceConsumerId, + next: SshPtySourceObligationState +): boolean { + const { token, span } = requireSourceSpan(spanOwners, spanId) + if (span.obligations.get(consumer)?.state !== 'open') { + return false + } + span.obligations.set(consumer, next) + advanceSourceTerminalEnd(token) + return true +} + +export function commitSourceObligationTransfer( + spanOwners: ReadonlyMap, + spanId: string, + consumer: SshPtySourceConsumerId +): boolean { + const { token, span } = requireSourceSpan(spanOwners, spanId) + const current = span.obligations.get(consumer) + if (current?.state !== 'transferring') { + return false + } + span.obligations.set( + consumer, + Object.freeze({ state: 'transferred', to: current.to, reason: current.reason }) + ) + advanceSourceTerminalEnd(token) + return true +} + +export function cancelSourceObligationTransfer( + spanOwners: ReadonlyMap, + spanId: string, + consumer: SshPtySourceConsumerId, + reason: string +): boolean { + const { token, span } = requireSourceSpan(spanOwners, spanId) + if (span.obligations.get(consumer)?.state !== 'transferring') { + return false + } + span.obligations.set(consumer, Object.freeze({ state: 'canceled', reason })) + advanceSourceTerminalEnd(token) + return true +} + +export function modelAcceptedSourceEnd(token: TokenRecord): number { + let acceptedEndSu = token.ackPublishedEndSu + for (const record of token.spans) { + if ( + record.span.sourceStartSu !== acceptedEndSu || + record.obligations.get('model')?.state !== 'settled' + ) { + break + } + acceptedEndSu = record.span.sourceEndSu + } + return acceptedEndSu +} diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index 7e1a54026df..ad8abb87398 100644 --- a/src/main/ipc/ssh.ts +++ b/src/main/ipc/ssh.ts @@ -7,6 +7,7 @@ import type { 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' +import { isSshPtySourceCreditV1Enabled } from '../ssh/ssh-pty-source-credit-rollout' import { SshPortForwardManager } from '../ssh/ssh-port-forward' import type { DetectedPort, @@ -1021,7 +1022,8 @@ export function registerSshHandlers( persistedStore!, portForwardManager!, currentRuntime, - broadcastDetectedPortsFromCurrentWindow + broadcastDetectedPortsFromCurrentWindow, + isSshPtySourceCreditV1Enabled ) configureRelaySessionCallbacks(session) activeSessions.set(targetId, session) diff --git a/src/main/providers/ssh-pty-notification-routing.test.ts b/src/main/providers/ssh-pty-notification-routing.test.ts index 572cffda4d2..845247faa7c 100644 --- a/src/main/providers/ssh-pty-notification-routing.test.ts +++ b/src/main/providers/ssh-pty-notification-routing.test.ts @@ -3,11 +3,13 @@ import { subscribeSshPtyNotifications } from './ssh-pty-notification-routing' type MockMux = { onNotification: ReturnType + request: ReturnType } function createSubscription() { const mux: MockMux = { - onNotification: vi.fn() + onNotification: vi.fn(), + request: vi.fn(async () => ({ canceled: true })) } const dataListeners = new Set<(payload: { id: string; data: string }) => void>() const replayListeners = new Set<(payload: { id: string; data: string }) => void>() @@ -15,6 +17,7 @@ function createSubscription() { const livePtyIds = new Set() const recordExit = vi.fn() const toAppPtyId = vi.fn((id: string) => `ssh:conn@@${id}`) + const resolvePtyIncarnation = vi.fn((id: string) => `incarnation:${id}`) subscribeSshPtyNotifications({ mux: mux as never, @@ -25,7 +28,7 @@ function createSubscription() { livePtyIds, recordExit, providerGeneration: 7, - resolvePtyIncarnation: (id) => `incarnation:${id}` + resolvePtyIncarnation }) const handler = mux.onNotification.mock.calls[0]?.[0] as ( @@ -38,12 +41,14 @@ function createSubscription() { return { handler, + mux, toAppPtyId, dataListeners, replayListeners, exitListeners, livePtyIds, - recordExit + recordExit, + resolvePtyIncarnation } } @@ -102,10 +107,12 @@ describe('subscribeSshPtyNotifications', () => { }) }) - it('derives exact source ranges and marks malformed negotiated frames', () => { - const { handler, dataListeners } = createSubscription() + it('derives exact immutable source ranges and cancels malformed frames without side effects', () => { + const { handler, mux, dataListeners, livePtyIds, toAppPtyId, resolvePtyIncarnation } = + createSubscription() const onData = vi.fn() dataListeners.add(onData) + livePtyIds.add('ssh:conn@@unrelated') handler('pty.data', { id: 'pty-1', @@ -117,6 +124,11 @@ describe('subscribeSshPtyNotifications', () => { sourceEndSu: 14, sourceLengthSu: 4 }) + const acceptedSource = onData.mock.calls[0]?.[0].source + expect(Object.isFrozen(acceptedSource)).toBe(true) + const liveBeforeMalformed = new Set(livePtyIds) + toAppPtyId.mockClear() + resolvePtyIncarnation.mockClear() handler('pty.data', { id: 'pty-1', data: 'bad', @@ -139,7 +151,256 @@ describe('subscribeSshPtyNotifications', () => { sourceEndSu: 14 } }) - expect(onData.mock.calls[1]?.[0]).toMatchObject({ sourceMalformed: true }) + expect(onData).toHaveBeenCalledTimes(1) + expect(livePtyIds).toEqual(liveBeforeMalformed) + expect(toAppPtyId).not.toHaveBeenCalled() + expect(resolvePtyIncarnation).not.toHaveBeenCalled() + expect(mux.request).toHaveBeenCalledWith('pty.cancelDelivery', { + id: 'pty-1', + clientGeneration: 2, + ownerGeneration: 3, + deliveryToken: 'token-1' + }) + }) + + it('keeps exact source incarnation independent from prior legacy delivery state', () => { + const { handler, dataListeners, resolvePtyIncarnation } = createSubscription() + const onData = vi.fn() + dataListeners.add(onData) + handler('pty.data', { id: 'pty-1', data: 'legacy' }) + handler('pty.data', { + id: 'pty-1', + data: 'data', + ptyIncarnation: 'incarnation-1', + deliveryToken: 'token-1', + clientGeneration: 2, + ownerGeneration: 3, + sourceEndSu: 4, + sourceLengthSu: 4 + }) + + expect(onData.mock.calls.map(([payload]) => payload.ptyIncarnation)).toEqual([ + 'incarnation:pty-1', + 'incarnation-1' + ]) + expect(resolvePtyIncarnation).toHaveBeenCalledTimes(2) + }) + + it('drops stale delivery generations without touching their PTY or unrelated PTYs', () => { + const { handler, mux, dataListeners, livePtyIds, toAppPtyId, resolvePtyIncarnation } = + createSubscription() + const onData = vi.fn() + dataListeners.add(onData) + livePtyIds.add('ssh:conn@@unrelated') + + handler('pty.data', { + id: 'pty-1', + data: 'new', + ptyIncarnation: 'incarnation-1', + deliveryToken: 'token-new', + clientGeneration: 4, + ownerGeneration: 5, + sourceEndSu: 3, + sourceLengthSu: 3 + }) + const liveBeforeStale = new Set(livePtyIds) + toAppPtyId.mockClear() + resolvePtyIncarnation.mockClear() + + handler('pty.data', { + id: 'pty-1', + data: 'old', + ptyIncarnation: 'incarnation-1', + deliveryToken: 'token-old', + clientGeneration: 3, + ownerGeneration: 4, + sourceEndSu: 6, + sourceLengthSu: 3 + }) + + expect(onData).toHaveBeenCalledTimes(1) + expect(livePtyIds).toEqual(liveBeforeStale) + expect(toAppPtyId).not.toHaveBeenCalled() + expect(resolvePtyIncarnation).not.toHaveBeenCalled() + expect(mux.request).toHaveBeenCalledWith('pty.cancelDelivery', { + id: 'pty-1', + clientGeneration: 3, + ownerGeneration: 4, + deliveryToken: 'token-old' + }) + }) + + it('rejects same-generation token changes and source discontinuities', () => { + const { handler, mux, dataListeners } = createSubscription() + const onData = vi.fn() + dataListeners.add(onData) + const sourceParams = { + id: 'pty-1', + ptyIncarnation: 'incarnation-1', + clientGeneration: 2, + ownerGeneration: 3, + sourceLengthSu: 3 + } + + handler('pty.data', { + ...sourceParams, + data: 'one', + deliveryToken: 'token-1', + sourceEndSu: 3 + }) + handler('pty.data', { + ...sourceParams, + data: 'two', + deliveryToken: 'token-2', + sourceEndSu: 6 + }) + handler('pty.data', { + ...sourceParams, + data: 'gap', + deliveryToken: 'token-1', + sourceEndSu: 9 + }) + handler('pty.data', { + ...sourceParams, + data: 'two', + deliveryToken: 'token-1', + sourceEndSu: 6 + }) + + expect(onData.mock.calls.map((call) => call[0].data)).toEqual(['one', 'two']) + expect(mux.request).toHaveBeenCalledTimes(2) + expect(mux.request).toHaveBeenCalledWith('pty.cancelDelivery', { + id: 'pty-1', + clientGeneration: 2, + ownerGeneration: 3, + deliveryToken: 'token-2' + }) + expect(mux.request).toHaveBeenCalledWith('pty.cancelDelivery', { + id: 'pty-1', + clientGeneration: 2, + ownerGeneration: 3, + deliveryToken: 'token-1' + }) + }) + + it('accepts a strictly newer rotation, rejects late old data, and preserves new continuity', () => { + const { handler, mux, dataListeners } = createSubscription() + const onData = vi.fn() + dataListeners.add(onData) + const frame = ( + data: string, + deliveryToken: string, + clientGeneration: number, + ownerGeneration: number, + sourceEndSu: number + ) => ({ + id: 'pty-1', + data, + ptyIncarnation: 'incarnation-1', + deliveryToken, + clientGeneration, + ownerGeneration, + sourceEndSu, + sourceLengthSu: data.length + }) + + handler('pty.data', frame('old', 'token-old', 2, 3, 3)) + handler('pty.data', frame('new', 'token-new', 3, 4, 13)) + handler('pty.data', frame('old', 'token-old', 2, 3, 6)) + handler('pty.data', frame('next', 'token-new', 3, 4, 17)) + + expect(onData.mock.calls.map((call) => call[0].data)).toEqual(['old', 'new', 'next']) + expect(mux.request).toHaveBeenCalledTimes(1) + expect(mux.request).toHaveBeenCalledWith('pty.cancelDelivery', { + id: 'pty-1', + clientGeneration: 2, + ownerGeneration: 3, + deliveryToken: 'token-old' + }) + }) + + it.each([ + ['client-only advance', 3, 3, 'token-client'], + ['owner-only advance', 2, 4, 'token-owner'], + ['crossed generations', 3, 2, 'token-crossed'], + ['replayed client generation', 1, 4, 'token-replayed'], + ['reused token on newer generations', 3, 4, 'token-current'] + ])( + 'rejects a %s without replacing the accepted continuity record', + (_case, clientGeneration, ownerGeneration, deliveryToken) => { + const { handler, mux, dataListeners, livePtyIds, toAppPtyId, resolvePtyIncarnation } = + createSubscription() + const onData = vi.fn() + dataListeners.add(onData) + const base = { + id: 'pty-1', + ptyIncarnation: 'incarnation-1', + sourceLengthSu: 3 + } + handler('pty.data', { + ...base, + data: 'one', + deliveryToken: 'token-current', + clientGeneration: 2, + ownerGeneration: 3, + sourceEndSu: 3 + }) + const liveBeforeInvalid = new Set(livePtyIds) + toAppPtyId.mockClear() + resolvePtyIncarnation.mockClear() + + handler('pty.data', { + ...base, + data: 'bad', + deliveryToken, + clientGeneration, + ownerGeneration, + sourceEndSu: 6 + }) + handler('pty.data', { + ...base, + data: 'two', + deliveryToken: 'token-current', + clientGeneration: 2, + ownerGeneration: 3, + sourceEndSu: 6 + }) + + expect(onData.mock.calls.map((call) => call[0].data)).toEqual(['one', 'two']) + expect(livePtyIds).toEqual(liveBeforeInvalid) + expect(toAppPtyId).toHaveBeenCalledTimes(1) + expect(resolvePtyIncarnation).toHaveBeenCalledTimes(1) + expect(mux.request).toHaveBeenCalledWith('pty.cancelDelivery', { + id: 'pty-1', + clientGeneration, + ownerGeneration, + deliveryToken + }) + } + ) + + it('does not cancel an incomplete malformed identity or mutate provider state', () => { + const { handler, mux, dataListeners, livePtyIds, toAppPtyId, resolvePtyIncarnation } = + createSubscription() + const onData = vi.fn() + dataListeners.add(onData) + livePtyIds.add('ssh:conn@@unrelated') + + handler('pty.data', { + id: 'pty-1', + data: 'bad', + ptyIncarnation: 'incarnation-1', + clientGeneration: 2, + ownerGeneration: 3, + sourceEndSu: 3, + sourceLengthSu: 3 + }) + + expect(onData).not.toHaveBeenCalled() + expect(livePtyIds).toEqual(new Set(['ssh:conn@@unrelated'])) + expect(toAppPtyId).not.toHaveBeenCalled() + expect(resolvePtyIncarnation).not.toHaveBeenCalled() + expect(mux.request).not.toHaveBeenCalled() }) it('ignores PTY methods with missing ids', () => { @@ -151,4 +412,49 @@ describe('subscribeSshPtyNotifications', () => { expect(toAppPtyId).not.toHaveBeenCalled() expect(onData).not.toHaveBeenCalled() }) + + it('leaves recovery and cancellation control methods to their dedicated handlers', () => { + const { + handler, + mux, + toAppPtyId, + dataListeners, + replayListeners, + exitListeners, + livePtyIds, + recordExit, + resolvePtyIncarnation + } = createSubscription() + const onData = vi.fn() + const onReplay = vi.fn() + const onExit = vi.fn() + dataListeners.add(onData) + replayListeners.add(onReplay) + exitListeners.add(onExit) + livePtyIds.add('ssh:conn@@unrelated') + + for (const method of [ + 'pty.recoveryData', + 'pty.recoveryComplete', + 'pty.restoreRequired', + 'pty.deliveryCanceled' + ]) { + handler(method, { + id: 'pty-1', + data: 'control', + deliveryToken: 'token-1', + clientGeneration: 2, + ownerGeneration: 3 + }) + } + + expect(toAppPtyId).not.toHaveBeenCalled() + expect(resolvePtyIncarnation).not.toHaveBeenCalled() + expect(recordExit).not.toHaveBeenCalled() + expect(onData).not.toHaveBeenCalled() + expect(onReplay).not.toHaveBeenCalled() + expect(onExit).not.toHaveBeenCalled() + expect(livePtyIds).toEqual(new Set(['ssh:conn@@unrelated'])) + expect(mux.request).not.toHaveBeenCalled() + }) }) diff --git a/src/main/providers/ssh-pty-notification-routing.ts b/src/main/providers/ssh-pty-notification-routing.ts index 4f0f36ff475..a75ad6a6a4b 100644 --- a/src/main/providers/ssh-pty-notification-routing.ts +++ b/src/main/providers/ssh-pty-notification-routing.ts @@ -5,10 +5,20 @@ import type { SshPtyExitCallback, SshPtyReplayCallback } from './ssh-pty-provider-contract' -import { parseSshPtySourceFrame } from './ssh-pty-source-frame' +import { parseSshPtySourceFrame, type SshPtySourceFrame } from './ssh-pty-source-frame' export type { SshPtyDataCallback, SshPtyExitCallback, SshPtyReplayCallback } +type SourceDeliveryState = Readonly<{ + clientGeneration: number + ownerGeneration: number + ptyIncarnation: string + deliveryToken: string + sourceEndSu: number +}> + +type SourceDeliveryOrder = 'continuation' | 'rotation' | 'invalid' + export function subscribeSshPtyNotifications(args: { mux: SshChannelMultiplexer toAppPtyId: (id: string) => string @@ -20,6 +30,7 @@ export function subscribeSshPtyNotifications(args: { providerGeneration: number resolvePtyIncarnation: (relayPtyId: string, incarnationId?: unknown) => string }): () => void { + const sourceDeliveryByPty = new Map() return args.mux.onNotification((method, params) => { // Why: mux delivers every method to generic handlers; non-PTY payloads // (workspace.changed, fs.changed, …) have no `id` and must not reach @@ -31,11 +42,12 @@ export function subscribeSshPtyNotifications(args: { return } const relayPtyId = params.id - const id = args.toAppPtyId(relayPtyId) if (method === 'pty.exit') { + const id = args.toAppPtyId(relayPtyId) const ptyIncarnation = args.resolvePtyIncarnation(relayPtyId, params.incarnationId) args.recordExit(relayPtyId, params.incarnationId) args.livePtyIds.delete(id) + sourceDeliveryByPty.delete(relayPtyId) for (const listener of args.exitListeners) { listener({ id, @@ -49,8 +61,9 @@ export function subscribeSshPtyNotifications(args: { } return } - args.livePtyIds.add(id) if (method === 'pty.replay') { + const id = args.toAppPtyId(relayPtyId) + args.livePtyIds.add(id) for (const listener of args.replayListeners) { listener({ id, data: params.data as string }) } @@ -58,21 +71,117 @@ export function subscribeSshPtyNotifications(args: { } const data = typeof params.data === 'string' ? params.data : '' const sourceFrame = parseSshPtySourceFrame(params, data, relayPtyId) + if ( + sourceFrame.malformed || + (sourceFrame.source && + !acceptSourceDelivery(sourceDeliveryByPty, relayPtyId, params, sourceFrame.source)) + ) { + cancelExactSourceDelivery(args.mux, relayPtyId, params) + return + } + const id = args.toAppPtyId(relayPtyId) + const ptyIncarnation = args.resolvePtyIncarnation( + relayPtyId, + params.ptyIncarnation ?? params.incarnationId + ) + args.livePtyIds.add(id) for (const listener of args.dataListeners) { listener({ id, data, providerGeneration: args.providerGeneration, - ptyIncarnation: args.resolvePtyIncarnation( - relayPtyId, - params.ptyIncarnation ?? params.incarnationId - ), + ptyIncarnation: sourceFrame.source ? (params.ptyIncarnation as string) : ptyIncarnation, ...(typeof params.rawLength === 'number' ? { sequenceChars: params.rawLength } : {}), ...(params.transformed === true ? { transformed: true } : {}), ...(typeof params.seq === 'number' ? { seq: params.seq } : {}), - ...(sourceFrame.source ? { source: sourceFrame.source } : {}), - ...(sourceFrame.malformed ? { sourceMalformed: true } : {}) + ...(sourceFrame.source ? { source: sourceFrame.source } : {}) }) } }) } + +function acceptSourceDelivery( + sourceDeliveryByPty: Map, + relayPtyId: string, + params: Record, + source: SshPtySourceFrame +): boolean { + const ptyIncarnation = params.ptyIncarnation as string + const current = sourceDeliveryByPty.get(relayPtyId) + if (current) { + const order = sourceDeliveryOrder(current, source) + const incarnationMatches = current.ptyIncarnation === ptyIncarnation + if (!incarnationMatches || order === 'invalid') { + return false + } + if ( + order === 'continuation' && + (current.deliveryToken !== source.deliveryToken || + current.sourceEndSu !== source.sourceStartSu) + ) { + return false + } + if (order === 'rotation' && current.deliveryToken === source.deliveryToken) { + return false + } + } + sourceDeliveryByPty.set( + relayPtyId, + Object.freeze({ + clientGeneration: source.clientGeneration, + ownerGeneration: source.ownerGeneration, + ptyIncarnation, + deliveryToken: source.deliveryToken, + sourceEndSu: source.sourceEndSu + }) + ) + return true +} + +function sourceDeliveryOrder( + current: SourceDeliveryState, + source: SshPtySourceFrame +): SourceDeliveryOrder { + if ( + source.clientGeneration === current.clientGeneration && + source.ownerGeneration === current.ownerGeneration + ) { + return 'continuation' + } + if ( + source.clientGeneration > current.clientGeneration && + source.ownerGeneration > current.ownerGeneration + ) { + return 'rotation' + } + return 'invalid' +} + +function cancelExactSourceDelivery( + mux: SshChannelMultiplexer, + relayPtyId: string, + params: Record +): void { + if ( + typeof params.deliveryToken !== 'string' || + params.deliveryToken.length === 0 || + !positiveSafeInteger(params.clientGeneration) || + !positiveSafeInteger(params.ownerGeneration) + ) { + return + } + try { + void mux + .request('pty.cancelDelivery', { + id: relayPtyId, + clientGeneration: params.clientGeneration, + ownerGeneration: params.ownerGeneration, + deliveryToken: params.deliveryToken + }) + .catch(() => {}) + } catch {} +} + +function positiveSafeInteger(value: unknown): value is number { + return Number.isSafeInteger(value) && (value as number) > 0 +} diff --git a/src/main/providers/ssh-pty-provider.ts b/src/main/providers/ssh-pty-provider.ts index 7c8c3438c75..a7845bd0a78 100644 --- a/src/main/providers/ssh-pty-provider.ts +++ b/src/main/providers/ssh-pty-provider.ts @@ -19,6 +19,7 @@ import { mapSshPtyProcessList } from './ssh-agent-session-process-list' import { parseSshPtyAttachResult, reattachSshPtySessionWithExitFence, + type PtySourceRecoveryRequest, type SshPtyAttachResult } from './ssh-pty-session-reattach' import { buildSshPtySpawnRequest } from './ssh-pty-spawn-request' @@ -69,13 +70,9 @@ export class SshPtyProvider implements IPtyProvider { getConnectionId = (): string => this.connectionId - private toRelayPtyId(id: string): string { - return toRelaySshPtyId(this.connectionId, id) - } + private toRelayPtyId = (id: string): string => toRelaySshPtyId(this.connectionId, id) - private toAppPtyId(id: string): string { - return toAppSshPtyId(this.connectionId, id) - } + private toAppPtyId = (id: string): string => toAppSshPtyId(this.connectionId, id) async spawn(opts: PtySpawnOptions): Promise { if (opts.agentSessionEnsure && opts.sessionId) { @@ -194,7 +191,8 @@ export class SshPtyProvider implements IPtyProvider { async attachForReconnect( id: string, - expected?: { paneKey?: string; tabId?: string } + expected?: { paneKey?: string; tabId?: string }, + sourceRecovery?: PtySourceRecoveryRequest ): Promise { // Why: reconnect owns replay delivery so stale/duplicate attach results can // be filtered before they reach the renderer. The expected identity lets the @@ -204,6 +202,7 @@ export class SshPtyProvider implements IPtyProvider { await this.mux.request('pty.attach', { id: this.toRelayPtyId(id), suppressReplayNotification: true, + ...(sourceRecovery ? { sourceRecovery } : {}), ...(expected?.paneKey ? { expectedPaneKey: expected.paneKey } : {}), ...(expected?.tabId ? { expectedTabId: expected.tabId } : {}) }) diff --git a/src/main/providers/ssh-pty-session-reattach.ts b/src/main/providers/ssh-pty-session-reattach.ts index 5970f85da0a..ffc23468f2e 100644 --- a/src/main/providers/ssh-pty-session-reattach.ts +++ b/src/main/providers/ssh-pty-session-reattach.ts @@ -9,10 +9,15 @@ import { import { toAppSshPtyId, toRelaySshPtyId } from './ssh-pty-id' import type { PtySpawnOptions, PtySpawnResult } from './types' import type { SshPtySpawnExitRaceTracker } from './ssh-pty-spawn-exit-race' +import type { + PtySourceRecoveryRequest, + PtySourceRecoveryResult +} from '../../shared/pty-source-recovery-contract' export type SshPtyAttachResult = { replay?: string incarnationId?: PtyIncarnationId + sourceRecovery?: PtySourceRecoveryResult } export function parseSshPtyAttachResult(value: unknown): SshPtyAttachResult { @@ -22,7 +27,11 @@ export function parseSshPtyAttachResult(value: unknown): SshPtyAttachResult { if (typeof value !== 'object' || Array.isArray(value)) { throw new Error('Invalid SSH PTY attach response') } - const result = value as { replay?: unknown; incarnationId?: unknown } + const result = value as { + replay?: unknown + incarnationId?: unknown + sourceRecovery?: unknown + } if (result.replay !== undefined && typeof result.replay !== 'string') { throw new Error('Invalid SSH PTY attach replay') } @@ -30,12 +39,60 @@ export function parseSshPtyAttachResult(value: unknown): SshPtyAttachResult { // Why: a present-but-invalid identity cannot safely fence delayed exits from a reused relay id. throw new Error('Invalid SSH PTY attach incarnation') } + const sourceRecovery = parseSourceRecoveryResult(result.sourceRecovery) return { ...(typeof result.replay === 'string' ? { replay: result.replay } : {}), - ...(isPtyIncarnationId(result.incarnationId) ? { incarnationId: result.incarnationId } : {}) + ...(isPtyIncarnationId(result.incarnationId) ? { incarnationId: result.incarnationId } : {}), + ...(sourceRecovery ? { sourceRecovery } : {}) } } +function parseSourceRecoveryResult(value: unknown): PtySourceRecoveryResult | undefined { + if (value === undefined) { + return undefined + } + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + throw new Error('Invalid SSH PTY source recovery response') + } + const input = value as Record + if (input.status === 'restoreRequired' && typeof input.reason === 'string') { + return Object.freeze({ status: 'restoreRequired', reason: input.reason }) + } + if ( + input.status !== 'pending' || + typeof input.deliveryToken !== 'string' || + input.deliveryToken.length === 0 || + typeof input.ptyIncarnation !== 'string' || + input.ptyIncarnation.length === 0 || + !positiveInteger(input.clientGeneration) || + !positiveInteger(input.ownerGeneration) || + !nonNegativeInteger(input.checkpointSourceEndSu) || + !nonNegativeInteger(input.recoveryEndSu) || + Number(input.recoveryEndSu) < Number(input.checkpointSourceEndSu) + ) { + throw new Error('Invalid SSH PTY source recovery response') + } + return Object.freeze({ + status: 'pending', + deliveryToken: input.deliveryToken, + ptyIncarnation: input.ptyIncarnation, + clientGeneration: Number(input.clientGeneration), + ownerGeneration: Number(input.ownerGeneration), + checkpointSourceEndSu: Number(input.checkpointSourceEndSu), + recoveryEndSu: Number(input.recoveryEndSu) + }) +} + +function positiveInteger(value: unknown): boolean { + return Number.isSafeInteger(value) && Number(value) > 0 +} + +function nonNegativeInteger(value: unknown): boolean { + return Number.isSafeInteger(value) && Number(value) >= 0 +} + +export type { PtySourceRecoveryRequest } + export async function reattachSshPtySession(args: { mux: SshChannelMultiplexer connectionId: string diff --git a/src/main/providers/ssh-pty-source-frame.ts b/src/main/providers/ssh-pty-source-frame.ts index 41d36dd0e3f..2165b4ba7c4 100644 --- a/src/main/providers/ssh-pty-source-frame.ts +++ b/src/main/providers/ssh-pty-source-frame.ts @@ -1,6 +1,6 @@ import type { SshPtyDataCallback } from './ssh-pty-provider-contract' -type SshPtySourceFrame = NonNullable[0]['source']> +export type SshPtySourceFrame = NonNullable[0]['source']> const SOURCE_KEYS = [ 'deliveryToken', @@ -24,6 +24,7 @@ export function parseSshPtySourceFrame( const rawLength = params.rawLength const transformed = params.transformed === true if ( + typeof params.data !== 'string' || typeof params.deliveryToken !== 'string' || params.deliveryToken.length === 0 || !positiveSafeInteger(params.clientGeneration) || @@ -33,6 +34,7 @@ export function parseSshPtySourceFrame( sourceEndSu < sourceLengthSu || typeof params.ptyIncarnation !== 'string' || params.ptyIncarnation.length === 0 || + (params.seq !== undefined && !nonNegativeSafeInteger(params.seq)) || (transformed ? rawLength !== sourceLengthSu : data.length !== sourceLengthSu) || (rawLength !== undefined && rawLength !== sourceLengthSu) ) { diff --git a/src/main/ssh/relay-frame-decoder.ts b/src/main/ssh/relay-frame-decoder.ts index 009f799c4fb..3ec15d6e53e 100644 --- a/src/main/ssh/relay-frame-decoder.ts +++ b/src/main/ssh/relay-frame-decoder.ts @@ -1,3 +1,6 @@ +import type { DecodedFrame, FrameDecoderOptions } from '../../shared/relay-frame-decoder-contract' +export type { DecodedFrame, FrameDecoderOptions } from '../../shared/relay-frame-decoder-contract' + export const HEADER_LENGTH = 13 export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 export const FRAME_DECODER_MAX_FRAMES_PER_TURN = 64 @@ -5,24 +8,6 @@ export const FRAME_DECODER_MAX_BYTES_PER_TURN = MAX_MESSAGE_SIZE + HEADER_LENGTH export const FRAME_DECODER_MAX_TURN_MS = 4, FRAME_DECODER_MAX_RETAINED_BYTES = 2 * (MAX_MESSAGE_SIZE + HEADER_LENGTH) -export type DecodedFrame = { - type: number - id: number - ack: number - payload: Buffer -} - -export type FrameDecoderOptions = { - maxFramesPerTurn?: number - maxBytesPerTurn?: number - maxTurnMs?: number - now?: () => number - schedule?: (callback: () => void) => unknown - cancelScheduled?: (handle: unknown) => void - pause?: () => void - resume?: () => void -} - export class FrameDecoder { private chunks: Buffer[] = [] private bufferedLength = 0 diff --git a/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts b/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts index 1e3d415dbfe..582a3d66c35 100644 --- a/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts +++ b/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { SshChannelMultiplexer, type MultiplexerTransport } from './ssh-channel-multiplexer' -import { encodeFrame, encodeKeepAliveFrame, MessageType } from './relay-protocol' +import { encodeFrame, encodeKeepAliveFrame, HEADER_LENGTH, MessageType } from './relay-protocol' type MockTransport = MultiplexerTransport & { data: (chunk: Buffer) => void @@ -135,4 +135,59 @@ describe('SshChannelMultiplexer backpressure hardening', () => { vi.advanceTimersByTime(0) expect(seen).toEqual(Array.from({ length: 65 }, (_, index) => index)) }) + + it('keeps acknowledgements, cancellation, exit, requests, and responses ordered after drain', async () => { + mux.dispose() + const written: Buffer[] = [] + let drain = (): void => {} + let deliver = (_data: Buffer): void => {} + transport = { + write: (data) => { + written.push(data) + return written.length !== 1 + }, + onDrain: (callback) => { + drain = callback + }, + onData: (callback) => { + deliver = callback + }, + onClose: vi.fn(), + pauseReads: vi.fn<() => void>(), + resumeReads: vi.fn<() => void>(), + data: vi.fn<(chunk: Buffer) => void>(), + written, + close: vi.fn() + } + mux = new SshChannelMultiplexer(transport) + mux.onRequest('client.control', () => ({ accepted: true })) + const controller = new AbortController() + + mux.notify('pty.data', { id: 'pty-1', data: 'ordinary' }) + mux.notify('pty.ackData', { acknowledgements: [] }) + const request = mux.request('fs.scan', {}, { signal: controller.signal }) + controller.abort() + mux.notify('pty.exit', { id: 'pty-1', code: 0 }) + const remoteRequest = Buffer.from( + JSON.stringify({ jsonrpc: '2.0', id: 91, method: 'client.control' }) + ) + deliver(encodeFrame(MessageType.Regular, 1, 0, remoteRequest)) + await Promise.resolve() + await Promise.resolve() + + expect(written).toHaveLength(1) + drain() + const payloads = written.map((frame) => + JSON.parse(frame.subarray(HEADER_LENGTH, HEADER_LENGTH + frame.readUInt32BE(9)).toString()) + ) + expect(payloads.map((payload) => payload.method ?? `response:${payload.id}`)).toEqual([ + 'pty.data', + 'pty.ackData', + 'fs.scan', + 'rpc.cancel', + 'pty.exit', + 'response:91' + ]) + await expect(request).rejects.toMatchObject({ name: 'AbortError' }) + }) }) diff --git a/src/main/ssh/ssh-channel-multiplexer-settlement.test.ts b/src/main/ssh/ssh-channel-multiplexer-settlement.test.ts index 75b7a7cb226..272f84399b3 100644 --- a/src/main/ssh/ssh-channel-multiplexer-settlement.test.ts +++ b/src/main/ssh/ssh-channel-multiplexer-settlement.test.ts @@ -13,6 +13,7 @@ function transportHarness(): { settlements.push(onSettled) } }, + supportsWriteSettlement: true, onData: vi.fn(), onClose: vi.fn() }, @@ -39,6 +40,7 @@ describe('SshChannelMultiplexer notification settlement', () => { write: () => { throw error }, + supportsWriteSettlement: true, onData: vi.fn(), onClose: vi.fn() }) @@ -55,6 +57,7 @@ describe('SshChannelMultiplexer notification settlement', () => { onSettled?.({ ok: true }) throw new Error('late throw') }, + supportsWriteSettlement: true, onData: vi.fn(), onClose: vi.fn() }) @@ -64,4 +67,27 @@ describe('SshChannelMultiplexer notification settlement', () => { expect(settled).toHaveBeenCalledOnce() expect(settled).toHaveBeenCalledWith({ ok: true }) }) + + it('fails an unsettled publication when the multiplexer is disposed', () => { + const close = vi.fn() + const mux = new SshChannelMultiplexer({ + write: () => false, + supportsWriteSettlement: true, + onDrain: vi.fn(), + onData: vi.fn(), + onClose: vi.fn(), + close + }) + const settled = vi.fn() + + mux.notifyWithSettlement('pty.ackData', { acknowledgements: [] }, settled) + expect(settled).not.toHaveBeenCalled() + mux.dispose() + + expect(settled).toHaveBeenCalledWith({ + ok: false, + error: expect.objectContaining({ code: 'DISPOSED' }) + }) + expect(close).toHaveBeenCalledOnce() + }) }) diff --git a/src/main/ssh/ssh-channel-multiplexer.test.ts b/src/main/ssh/ssh-channel-multiplexer.test.ts index d17e756f75b..d46f269fd34 100644 --- a/src/main/ssh/ssh-channel-multiplexer.test.ts +++ b/src/main/ssh/ssh-channel-multiplexer.test.ts @@ -12,7 +12,9 @@ function createMockTransport(): MultiplexerTransport & { const written: Buffer[] = [] return { - write: (data: Buffer) => written.push(data), + write: (data: Buffer) => { + written.push(data) + }, onData: (cb) => dataCallbacks.push(cb), onClose: (cb) => closeCallbacks.push(cb), dataCallbacks, diff --git a/src/main/ssh/ssh-channel-multiplexer.ts b/src/main/ssh/ssh-channel-multiplexer.ts index 91c2cb8422b..686659c3de5 100644 --- a/src/main/ssh/ssh-channel-multiplexer.ts +++ b/src/main/ssh/ssh-channel-multiplexer.ts @@ -14,18 +14,14 @@ import { type JsonRpcResponse, type JsonRpcNotification } from './relay-protocol' +import { + SshMultiplexerTransportWriter, + type MultiplexerTransport, + type MultiplexerWriteSettlement, + type MultiplexerWriterLane +} from './ssh-multiplexer-transport-writer' -export type MultiplexerTransport = { - write: ( - data: Buffer, - onSettled?: (result: { ok: true } | { ok: false; error: Error }) => void - ) => void - onData: (cb: (data: Buffer) => void) => void - onClose: (cb: () => void) => void - pauseReads?: () => void - resumeReads?: () => void - close?: () => void -} +export type { MultiplexerTransport, MultiplexerWriteSettlement } type PendingRequest = { resolve: (result: unknown) => void @@ -48,6 +44,7 @@ const WAKE_GAP_MS = KEEPALIVE_SEND_MS * 3 export class SshChannelMultiplexer { private decoder: FrameDecoder private transport: MultiplexerTransport + private writer: SshMultiplexerTransportWriter private nextRequestId = 1 private nextOutgoingSeq = 1 private highestReceivedSeq = 0 @@ -74,6 +71,9 @@ export class SshChannelMultiplexer { constructor(transport: MultiplexerTransport) { this.transport = transport + this.writer = new SshMultiplexerTransportWriter(transport, (error) => + this.handleProtocolError(error) + ) this.decoder = new FrameDecoder( (frame) => this.handleFrame(frame), @@ -328,6 +328,9 @@ export class SshChannelMultiplexer { this.pendingRequests.delete(id) } + const writerError = new Error(errorMessage) as Error & { code: string } + writerError.code = errorCode + this.writer.dispose(writerError) this.unackedTimestamps.clear() // Why: relay teardown can race with late provider registration; disposed // muxes must not retain provider/session closures through subscribers. @@ -354,33 +357,12 @@ export class SshChannelMultiplexer { private sendMessage( msg: JsonRpcMessage, - onSettled?: (result: { ok: true } | { ok: false; error: Error }) => void + onSettled?: (result: MultiplexerWriteSettlement) => void ): void { const seq = this.nextOutgoingSeq++ const frame = encodeJsonRpcFrame(msg, seq, this.highestReceivedSeq) this.trackOutgoingTimestamp(seq, false) - let writeSettled = false - const settleWrite = onSettled - ? (result: { ok: true } | { ok: false; error: Error }): void => { - if (writeSettled) { - return - } - writeSettled = true - onSettled(result) - } - : undefined - try { - this.transport.write(frame, settleWrite) - } catch (err) { - settleWrite?.({ - ok: false, - error: err instanceof Error ? err : new Error(String(err)) - }) - // Why: a remote reboot can make the SSH channel's stdin throw EPIPE - // from a timer/request path. Scope it to this mux instead of letting - // the Electron main process treat it as an uncaught exception. - this.handleProtocolError(err) - } + this.writer.enqueue(frame, messageLane(msg), onSettled) } private sendKeepAlive(): void { @@ -390,13 +372,7 @@ export class SshChannelMultiplexer { const seq = this.nextOutgoingSeq++ const frame = encodeKeepAliveFrame(seq, this.highestReceivedSeq) this.trackOutgoingTimestamp(seq, true) - try { - this.transport.write(frame) - } catch (err) { - // Why: keepalive runs on an interval; without catching transport - // write failures here, a dead SSH host can terminate the whole app. - this.handleProtocolError(err) - } + this.writer.enqueue(frame, 'control') } private handleFrame(frame: DecodedFrame): void { @@ -621,3 +597,7 @@ export class SshChannelMultiplexer { } } } + +function messageLane(msg: JsonRpcMessage): MultiplexerWriterLane { + return 'method' in msg && msg.method === 'pty.data' ? 'ordinary' : 'control' +} diff --git a/src/main/ssh/ssh-multiplexer-transport-writer.test.ts b/src/main/ssh/ssh-multiplexer-transport-writer.test.ts new file mode 100644 index 00000000000..7087fbff5ad --- /dev/null +++ b/src/main/ssh/ssh-multiplexer-transport-writer.test.ts @@ -0,0 +1,202 @@ +import { EventEmitter } from 'node:events' +import { describe, expect, it, vi } from 'vitest' +import { + MULTIPLEXER_CONTROL_RESERVE_BYTES, + MULTIPLEXER_ORDINARY_QUEUE_MAX_BYTES, + SshMultiplexerTransportWriter, + type MultiplexerTransport, + type MultiplexerWriteSettlement +} from './ssh-multiplexer-transport-writer' + +type WriterHarness = { + transport: MultiplexerTransport + drain: () => void + writes: Buffer[] + callbacks: ((result: MultiplexerWriteSettlement) => void)[] + removeDrain: ReturnType +} + +function transportHarness(writeResults: (boolean | void)[]): WriterHarness { + const emitter = new EventEmitter() + const writes: Buffer[] = [] + const callbacks: ((result: MultiplexerWriteSettlement) => void)[] = [] + const removeDrain = vi.fn() + return { + transport: { + write: (data, onSettled) => { + writes.push(data) + callbacks.push(onSettled!) + return writeResults.shift() + }, + supportsWriteSettlement: true, + onDrain: (callback) => { + emitter.on('drain', callback) + return () => { + removeDrain() + emitter.off('drain', callback) + } + }, + onData: vi.fn(), + onClose: vi.fn() + }, + drain: () => emitter.emit('drain'), + writes, + callbacks, + removeDrain + } +} + +describe('SshMultiplexerTransportWriter', () => { + it('preserves FIFO order and waits for drain after write(false)', () => { + const harness = transportHarness([false, true, true]) + const failed = vi.fn() + const writer = new SshMultiplexerTransportWriter(harness.transport, failed) + const settlements = [vi.fn(), vi.fn(), vi.fn()] + + writer.enqueue(Buffer.from('ordinary-1'), 'ordinary', settlements[0]) + writer.enqueue(Buffer.from('control'), 'control', settlements[1]) + writer.enqueue(Buffer.from('ordinary-2'), 'ordinary', settlements[2]) + + expect(harness.writes.map(String)).toEqual(['ordinary-1']) + harness.callbacks[0]({ ok: true }) + expect(settlements[0]).toHaveBeenCalledWith({ ok: true }) + expect(harness.writes.map(String)).toEqual(['ordinary-1']) + + harness.drain() + expect(harness.writes.map(String)).toEqual(['ordinary-1', 'control', 'ordinary-2']) + harness.callbacks[1]({ ok: true }) + harness.callbacks[2]({ ok: true }) + expect(settlements.every((settle) => settle.mock.calls.length === 1)).toBe(true) + expect(failed).not.toHaveBeenCalled() + }) + + it('keeps a full ordinary lane from consuming the control reserve', () => { + const harness = transportHarness([false, true, true]) + const writer = new SshMultiplexerTransportWriter(harness.transport, vi.fn()) + + expect(writer.enqueue(Buffer.alloc(1), 'ordinary')).toBe(true) + expect(writer.enqueue(Buffer.alloc(MULTIPLEXER_ORDINARY_QUEUE_MAX_BYTES - 1), 'ordinary')).toBe( + true + ) + expect(writer.enqueue(Buffer.alloc(MULTIPLEXER_CONTROL_RESERVE_BYTES), 'control')).toBe(true) + + expect(harness.writes).toHaveLength(1) + harness.drain() + expect(harness.writes).toHaveLength(3) + writer.dispose() + }) + + it('fails all retained writes once when a callback reports an error', () => { + const harness = transportHarness([false]) + const failed = vi.fn() + const writer = new SshMultiplexerTransportWriter(harness.transport, failed) + const first = vi.fn() + const queued = vi.fn() + const error = new Error('broken pipe') + + writer.enqueue(Buffer.from('first'), 'ordinary', first) + writer.enqueue(Buffer.from('queued'), 'control', queued) + harness.callbacks[0]({ ok: false, error }) + harness.callbacks[0]({ ok: true }) + + expect(first).toHaveBeenCalledOnce() + expect(first).toHaveBeenCalledWith({ ok: false, error }) + expect(queued).toHaveBeenCalledWith({ ok: false, error }) + expect(failed).toHaveBeenCalledWith(error) + expect(harness.removeDrain).toHaveBeenCalledOnce() + }) + + it('closes on queue overflow and rejects the overflowing write', () => { + const harness = transportHarness([false]) + const failed = vi.fn() + const writer = new SshMultiplexerTransportWriter(harness.transport, failed) + const retained = vi.fn() + const overflow = vi.fn() + + writer.enqueue(Buffer.alloc(MULTIPLEXER_ORDINARY_QUEUE_MAX_BYTES), 'ordinary', retained) + expect(writer.enqueue(Buffer.alloc(1), 'ordinary', overflow)).toBe(false) + + expect(overflow).toHaveBeenCalledWith({ + ok: false, + error: expect.objectContaining({ message: expect.stringContaining('bounded capacity') }) + }) + expect(retained).toHaveBeenCalledWith({ + ok: false, + error: expect.objectContaining({ message: expect.stringContaining('bounded capacity') }) + }) + expect(failed).toHaveBeenCalledOnce() + }) + + it('settles legacy callback-less transports on acceptance and drain', () => { + const emitter = new EventEmitter() + const write = vi.fn().mockReturnValueOnce(false).mockReturnValueOnce(undefined) + const transport: MultiplexerTransport = { + write, + onDrain: (callback) => { + emitter.on('drain', callback) + }, + onData: vi.fn(), + onClose: vi.fn() + } + const writer = new SshMultiplexerTransportWriter(transport, vi.fn()) + const first = vi.fn() + const second = vi.fn() + + writer.enqueue(Buffer.from('first'), 'ordinary', first) + writer.enqueue(Buffer.from('second'), 'control', second) + expect(first).not.toHaveBeenCalled() + expect(second).not.toHaveBeenCalled() + + emitter.emit('drain') + expect(first).toHaveBeenCalledWith({ ok: true }) + expect(second).toHaveBeenCalledWith({ ok: true }) + }) + + it('does not miss a drain emitted synchronously by a hostile transport', () => { + let drain = (): void => {} + const write = vi.fn(() => { + drain() + return false + }) + const writer = new SshMultiplexerTransportWriter( + { + write, + onDrain: (callback) => { + drain = callback + }, + onData: vi.fn(), + onClose: vi.fn() + }, + vi.fn() + ) + const first = vi.fn() + const second = vi.fn() + + writer.enqueue(Buffer.from('first'), 'ordinary', first) + writer.enqueue(Buffer.from('second'), 'control', second) + + expect(write).toHaveBeenCalledTimes(2) + expect(first).toHaveBeenCalledWith({ ok: true }) + expect(second).toHaveBeenCalledWith({ ok: true }) + }) + + it('fails deterministically when write(false) has no drain source', () => { + const failed = vi.fn() + const writer = new SshMultiplexerTransportWriter( + { + write: () => false, + onData: vi.fn(), + onClose: vi.fn() + }, + failed + ) + const settled = vi.fn() + + expect(writer.enqueue(Buffer.from('data'), 'ordinary', settled)).toBe(true) + expect(settled).toHaveBeenCalledWith({ + ok: false, + error: expect.objectContaining({ message: expect.stringContaining('without drain support') }) + }) + expect(failed).toHaveBeenCalledOnce() + }) +}) diff --git a/src/main/ssh/ssh-multiplexer-transport-writer.ts b/src/main/ssh/ssh-multiplexer-transport-writer.ts new file mode 100644 index 00000000000..bb9d45b1068 --- /dev/null +++ b/src/main/ssh/ssh-multiplexer-transport-writer.ts @@ -0,0 +1,246 @@ +import { HEADER_LENGTH, MAX_MESSAGE_SIZE } from './relay-protocol' + +export type MultiplexerWriteSettlement = { ok: true } | { ok: false; error: Error } + +export type MultiplexerTransport = { + write: (data: Buffer, onSettled?: (result: MultiplexerWriteSettlement) => void) => boolean | void + onData: (cb: (data: Buffer) => void) => void + onClose: (cb: () => void) => void + onDrain?: (cb: () => void) => void | (() => void) + supportsWriteSettlement?: boolean + pauseReads?: () => void + resumeReads?: () => void + close?: () => void +} + +export type MultiplexerWriterLane = 'ordinary' | 'control' + +type WriterEntry = { + data: Buffer + lane: MultiplexerWriterLane + onSettled: (result: MultiplexerWriteSettlement) => void + settled: boolean +} + +export const MULTIPLEXER_ORDINARY_QUEUE_MAX_BYTES = 2 * 1024 * 1024 +export const MULTIPLEXER_CONTROL_RESERVE_BYTES = MAX_MESSAGE_SIZE + HEADER_LENGTH +const ORDINARY_QUEUE_MAX_FRAMES = 2048 +const CONTROL_QUEUE_MAX_FRAMES = 512 + +function onceSettlement( + callback: (result: MultiplexerWriteSettlement) => void +): (result: MultiplexerWriteSettlement) => void { + let settled = false + return (result) => { + if (settled) { + return + } + settled = true + callback(result) + } +} + +export class SshMultiplexerTransportWriter { + private readonly queue: WriterEntry[] = [] + private readonly inFlight = new Set() + private readonly settleOnDrain = new Set() + private ordinaryBytes = 0 + private controlBytes = 0 + private ordinaryFrames = 0 + private controlFrames = 0 + private saturated = false + private writing = false + private drainObservedDuringWrite = false + private pumping = false + private closed = false + private removeDrainListener: (() => void) | null = null + + constructor( + private readonly transport: MultiplexerTransport, + private readonly onFailure: (error: Error) => void + ) { + if (transport.onDrain) { + const remove = transport.onDrain(() => this.handleDrain()) + this.removeDrainListener = typeof remove === 'function' ? remove : null + } + } + + enqueue( + data: Buffer, + lane: MultiplexerWriterLane, + onSettled: (result: MultiplexerWriteSettlement) => void = () => {} + ): boolean { + const settle = onceSettlement(onSettled) + if (this.closed) { + settle({ ok: false, error: new Error('Multiplexer writer is closed') }) + return false + } + const admissionError = this.admissionError(data.length, lane) + if (admissionError) { + settle({ ok: false, error: admissionError }) + this.fail(admissionError) + return false + } + const entry = { data, lane, onSettled: settle, settled: false } + this.retain(entry) + this.queue.push(entry) + this.pump() + return true + } + + dispose(error = new Error('Multiplexer writer disposed')): void { + if (this.closed) { + return + } + this.closed = true + this.saturated = false + this.removeDrainListener?.() + this.removeDrainListener = null + for (const entry of this.queue.splice(0)) { + this.release(entry, { ok: false, error }) + } + for (const entry of Array.from(this.inFlight)) { + this.release(entry, { ok: false, error }) + } + this.settleOnDrain.clear() + } + + private admissionError(bytes: number, lane: MultiplexerWriterLane): Error | null { + const byteLimit = + lane === 'ordinary' ? MULTIPLEXER_ORDINARY_QUEUE_MAX_BYTES : MULTIPLEXER_CONTROL_RESERVE_BYTES + const retainedBytes = lane === 'ordinary' ? this.ordinaryBytes : this.controlBytes + const frameLimit = lane === 'ordinary' ? ORDINARY_QUEUE_MAX_FRAMES : CONTROL_QUEUE_MAX_FRAMES + const retainedFrames = lane === 'ordinary' ? this.ordinaryFrames : this.controlFrames + if (retainedBytes + bytes <= byteLimit && retainedFrames < frameLimit) { + return null + } + return new Error(`Multiplexer ${lane} write queue exceeded its bounded capacity`) + } + + private pump(): void { + if (this.pumping || this.closed || this.saturated) { + return + } + this.pumping = true + try { + while (!this.closed && !this.saturated) { + const entry = this.queue.shift() + if (!entry) { + return + } + this.writeEntry(entry) + } + } finally { + this.pumping = false + } + } + + private writeEntry(entry: WriterEntry): void { + this.inFlight.add(entry) + let callbackResult: MultiplexerWriteSettlement | undefined + let writeReturned = false + const onWriteSettled = (result: MultiplexerWriteSettlement): void => { + if (!writeReturned) { + callbackResult = result + return + } + this.handleWriteSettlement(entry, result) + } + try { + this.writing = true + this.drainObservedDuringWrite = false + const accepted = this.transport.write(entry.data, onWriteSettled) + this.writing = false + writeReturned = true + if (accepted === false) { + if (!this.transport.onDrain) { + throw new Error('Multiplexer transport returned write(false) without drain support') + } + this.saturated = !this.drainObservedDuringWrite + if (this.transport.supportsWriteSettlement !== true && this.saturated) { + this.settleOnDrain.add(entry) + } else if (this.transport.supportsWriteSettlement !== true) { + this.handleWriteSettlement(entry, { ok: true }) + } + } else if (this.transport.supportsWriteSettlement !== true) { + this.handleWriteSettlement(entry, { ok: true }) + } + if (callbackResult) { + this.handleWriteSettlement(entry, callbackResult) + } + } catch (error) { + this.writing = false + writeReturned = true + if (callbackResult) { + this.handleWriteSettlement(entry, callbackResult) + } + this.fail(error instanceof Error ? error : new Error(String(error))) + } + } + + private handleWriteSettlement(entry: WriterEntry, result: MultiplexerWriteSettlement): void { + if (entry.settled) { + return + } + this.release(entry, result) + if (!result.ok) { + this.fail(result.error) + return + } + this.pump() + } + + private handleDrain(): void { + if (this.closed) { + return + } + if (this.writing) { + this.drainObservedDuringWrite = true + return + } + if (!this.saturated) { + return + } + this.saturated = false + for (const entry of Array.from(this.settleOnDrain)) { + this.release(entry, { ok: true }) + } + this.settleOnDrain.clear() + this.pump() + } + + private retain(entry: WriterEntry): void { + if (entry.lane === 'ordinary') { + this.ordinaryBytes += entry.data.length + this.ordinaryFrames++ + } else { + this.controlBytes += entry.data.length + this.controlFrames++ + } + } + + private release(entry: WriterEntry, result: MultiplexerWriteSettlement): void { + if (entry.settled) { + return + } + entry.settled = true + this.inFlight.delete(entry) + this.settleOnDrain.delete(entry) + if (entry.lane === 'ordinary') { + this.ordinaryBytes -= entry.data.length + this.ordinaryFrames-- + } else { + this.controlBytes -= entry.data.length + this.controlFrames-- + } + entry.onSettled(result) + } + + private fail(error: Error): void { + if (this.closed) { + return + } + this.dispose(error) + this.onFailure(error) + } +} diff --git a/src/main/ssh/ssh-pty-source-credit-rollout.test.ts b/src/main/ssh/ssh-pty-source-credit-rollout.test.ts new file mode 100644 index 00000000000..7a92eb0684e --- /dev/null +++ b/src/main/ssh/ssh-pty-source-credit-rollout.test.ts @@ -0,0 +1,14 @@ +import { describe, expect, it } from 'vitest' +import { + isSshPtySourceCreditV1Enabled, + SSH_PTY_SOURCE_CREDIT_V1_ENV +} from './ssh-pty-source-credit-rollout' + +describe('SSH PTY source-credit rollout', () => { + it('defaults off and requires the exact opt-in value', () => { + expect(isSshPtySourceCreditV1Enabled({})).toBe(false) + expect(isSshPtySourceCreditV1Enabled({ [SSH_PTY_SOURCE_CREDIT_V1_ENV]: '0' })).toBe(false) + expect(isSshPtySourceCreditV1Enabled({ [SSH_PTY_SOURCE_CREDIT_V1_ENV]: 'true' })).toBe(false) + expect(isSshPtySourceCreditV1Enabled({ [SSH_PTY_SOURCE_CREDIT_V1_ENV]: '1' })).toBe(true) + }) +}) diff --git a/src/main/ssh/ssh-pty-source-credit-rollout.ts b/src/main/ssh/ssh-pty-source-credit-rollout.ts new file mode 100644 index 00000000000..9355ac3be52 --- /dev/null +++ b/src/main/ssh/ssh-pty-source-credit-rollout.ts @@ -0,0 +1,5 @@ +export const SSH_PTY_SOURCE_CREDIT_V1_ENV = 'ORCA_SSH_PTY_SOURCE_CREDIT_V1' + +export function isSshPtySourceCreditV1Enabled(env: NodeJS.ProcessEnv = process.env): boolean { + return env[SSH_PTY_SOURCE_CREDIT_V1_ENV] === '1' +} diff --git a/src/main/ssh/ssh-relay-deploy-helpers.test.ts b/src/main/ssh/ssh-relay-deploy-helpers.test.ts index 30e436b979f..acda4594ca2 100644 --- a/src/main/ssh/ssh-relay-deploy-helpers.test.ts +++ b/src/main/ssh/ssh-relay-deploy-helpers.test.ts @@ -9,13 +9,17 @@ import { RELAY_EXIT_CODE_VERSION_MISMATCH } from './ssh-relay-version-mismatch-error' -function createMockChannel(): ClientChannel { +type MockChannel = ClientChannel & { + stdin: EventEmitter & { write: ReturnType } +} + +function createMockChannel(): MockChannel { return Object.assign(new EventEmitter(), { stderr: Object.assign(new EventEmitter(), { resume: vi.fn() }), - stdin: { write: vi.fn() }, + stdin: Object.assign(new EventEmitter(), { write: vi.fn(() => true) }), close: vi.fn(), resume: vi.fn() - }) as unknown as ClientChannel + }) as unknown as MockChannel } // execCommand only rejects with Error; narrow the caught reason (its resolve @@ -204,6 +208,35 @@ describe('waitForSentinel', () => { expect(Buffer.concat(chunks)).toEqual(postSentinelPayload) expect(channel.close).not.toHaveBeenCalled() }) + + it.each(['ssh2 channel', 'system-SSH child stdio'])( + 'forwards write(false), callback settlement, and drain for a %s', + async (shape) => { + const channel = createMockChannel() + if (shape.startsWith('system')) { + Object.assign(channel, { _process: new EventEmitter() }) + } + const callback = vi.fn() + const drain = vi.fn() + channel.stdin.write.mockImplementation((...args: unknown[]) => { + const onWritten = args.find((arg) => typeof arg === 'function') as ( + error?: Error | null + ) => void + onWritten(null) + return false + }) + const transportPromise = waitForSentinel(channel) + channel.emit('data', Buffer.from(RELAY_SENTINEL)) + const transport = await transportPromise + + transport.onDrain?.(drain) + expect(transport.write(Buffer.from('frame'), callback)).toBe(false) + expect(callback).toHaveBeenCalledWith({ ok: true }) + channel.stdin.emit('drain') + expect(drain).toHaveBeenCalledOnce() + expect(transport.supportsWriteSettlement).toBe(true) + } + ) }) describe('execCommand', () => { diff --git a/src/main/ssh/ssh-relay-deploy-helpers.ts b/src/main/ssh/ssh-relay-deploy-helpers.ts index 88b0e2befa7..a035133ddc5 100644 --- a/src/main/ssh/ssh-relay-deploy-helpers.ts +++ b/src/main/ssh/ssh-relay-deploy-helpers.ts @@ -215,10 +215,15 @@ export function waitForSentinel( } const transport: MultiplexerTransport = { write: (buf: Buffer, onSettled) => { - channel.stdin.write(buf, (error?: Error | null) => { + return channel.stdin.write(buf, (error?: Error | null) => { onSettled?.(error ? { ok: false, error } : { ok: true }) }) }, + supportsWriteSettlement: true, + onDrain: (cb) => { + channel.stdin.on('drain', cb) + return () => channel.stdin.off('drain', cb) + }, onData: (cb) => { dataCallbacks.push(cb) // Why: deliver buffered post-sentinel data to the first diff --git a/src/main/ssh/ssh-relay-deploy.test.ts b/src/main/ssh/ssh-relay-deploy.test.ts index 244d0c43f09..bbf8d0a6e0f 100644 --- a/src/main/ssh/ssh-relay-deploy.test.ts +++ b/src/main/ssh/ssh-relay-deploy.test.ts @@ -483,6 +483,7 @@ describe('deployAndLaunchRelay', () => { .find((cmd) => cmd.includes('--detached')) expect(launchCommand).toContain(`--grace-time ${DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS}`) + expect(launchCommand).not.toContain('--pty-source-credit-v1') }) it('allows an unlimited SSH disconnect grace window', async () => { @@ -494,7 +495,7 @@ describe('deployAndLaunchRelay', () => { mockExecCommand.mockResolvedValueOnce('DEAD') mockExecCommand.mockResolvedValueOnce('READY') - await deployAndLaunchRelay(conn, undefined, 0, 'target-a') + await deployAndLaunchRelay(conn, undefined, 0, 'target-a', true) const launchCommand = vi .mocked(conn.exec) @@ -502,6 +503,7 @@ describe('deployAndLaunchRelay', () => { .find((cmd) => cmd.includes('--detached')) expect(launchCommand).toContain('--grace-time 0') + expect(launchCommand).toContain('--pty-source-credit-v1') }) it('clamps configured SSH disconnect grace to the seven-day maximum', async () => { @@ -748,7 +750,7 @@ describe('deployAndLaunchRelay', () => { .mockResolvedValueOnce('READY') // named pipe poll .mockResolvedValueOnce('') // persist active pipe marker - const result = await deployAndLaunchRelay(conn, undefined, 300, 'target-a') + const result = await deployAndLaunchRelay(conn, undefined, 300, 'target-a', true) expect(result.platform).toBe('win32-x64') expect(result.remoteHome).toBe('C:/Users/me user') @@ -767,6 +769,7 @@ describe('deployAndLaunchRelay', () => { '"C:/Users/me user/.orca-remote/relay-0.1.0+abcdef012345/agent-hooks/orca-relay-' ) expect(launchScript).toContain('--endpoint-dir') + expect(launchScript).toContain('--pty-source-credit-v1') expect(launchScript).not.toContain('\\\\.\\pipe\\agent-hooks') const waitScript = decodedScripts.find((script) => script.includes('deadline=Date.now()')) ?? '' expect(waitScript).toContain('setTimeout(attempt,intervalMs)') diff --git a/src/main/ssh/ssh-relay-deploy.ts b/src/main/ssh/ssh-relay-deploy.ts index ac9bdcaffc9..fd28f5badb0 100644 --- a/src/main/ssh/ssh-relay-deploy.ts +++ b/src/main/ssh/ssh-relay-deploy.ts @@ -115,7 +115,8 @@ export async function deployAndLaunchRelay( conn: SshConnection, onProgress?: (status: string) => void, graceTimeSeconds?: number, - relayInstanceId?: string + relayInstanceId?: string, + enablePtySourceCreditV1 = false ): Promise { let timeoutHandle: ReturnType // Why: Promise.race doesn't cancel its loser; abort so a contended install-lock waiter can't mutate the relay after this call times out. @@ -134,6 +135,7 @@ export async function deployAndLaunchRelay( onProgress, graceTimeSeconds, relayInstanceId, + enablePtySourceCreditV1, deployAbortController.signal ), timeoutPromise @@ -265,6 +267,7 @@ async function deployAndLaunchRelayInner( onProgress?: (status: string) => void, graceTimeSeconds?: number, relayInstanceId?: string, + enablePtySourceCreditV1 = false, deploySignal?: AbortSignal ): Promise { while (true) { @@ -275,6 +278,7 @@ async function deployAndLaunchRelayInner( onProgress, graceTimeSeconds, relayInstanceId, + enablePtySourceCreditV1, deploySignal ) } catch (err) { @@ -292,6 +296,7 @@ async function deployAndLaunchRelayAttempt( onProgress?: (status: string) => void, graceTimeSeconds?: number, relayInstanceId?: string, + enablePtySourceCreditV1 = false, deploySignal?: AbortSignal ): Promise { onProgress?.('Detecting remote platform...') @@ -414,6 +419,7 @@ async function deployAndLaunchRelayAttempt( nodePath, graceTimeSeconds, relayInstanceId, + enablePtySourceCreditV1, deploySignal ) launchLivenessObserved = true @@ -1161,6 +1167,7 @@ async function launchRelay( nodePath: string, graceTimeSeconds?: number, relayInstanceId?: string, + enablePtySourceCreditV1 = false, signal?: AbortSignal ): Promise<{ transport: MultiplexerTransport @@ -1209,7 +1216,8 @@ async function launchRelay( graceTime, activePipeMarkerPath, reconnectFallback: fallbackEndpoint, - credentialFile + credentialFile, + enablePtySourceCreditV1 }, signal ) @@ -1264,7 +1272,8 @@ async function launchRelay( const logFile = `${remoteDir}/relay.log` await writeRelayEndpointCredential(conn, hostPlatform, credentialFile, signal) // Why: --log-file lets the relay rotate relay.log in-process; the shell redirect stays to capture pre-JS boot/crash output. - const launchCmd = `cd ${escapedDir} && chmod 600 ${shellEscape(credentialFile)} && nohup ${escapedNode} relay.js --detached --grace-time ${graceTime} --sock-path ${shellEscape(sockFile)} --credential-file ${shellEscape(credentialFile)} --log-file ${shellEscape(logFile)} > ${shellEscape(logFile)} 2>&1 ${shellEscape(logFile)} 2>&1 {}) launchChannel.on('error', () => {}) @@ -1428,6 +1437,7 @@ type WindowsRelayLaunchOptions = { graceTime: number activePipeMarkerPath: string credentialFile: string + enablePtySourceCreditV1: boolean } & WindowsRelayEndpoint & { reconnectFallback?: WindowsRelayEndpoint } @@ -1510,7 +1520,8 @@ async function launchWindowsRelay( launchOpts.graceTime, logFile, errFile, - launchOpts.credentialFile + launchOpts.credentialFile, + launchOpts.enablePtySourceCreditV1 ), { signal } ) @@ -1602,7 +1613,8 @@ function windowsRelayLaunchCommand( graceTime: number, logFile: string, errFile: string, - credentialFile: string + credentialFile: string, + enablePtySourceCreditV1 = false ): string { const relayScript = joinRemotePath(hostPlatform, remoteDir, 'relay.js') // Why: Windows sshd kills the exec channel's process tree on close; WMI re-parents the detached relay to survive. @@ -1611,6 +1623,7 @@ function windowsRelayLaunchCommand( quoted(nodePath), quoted(relayScript), '--detached', + ...(enablePtySourceCreditV1 ? ['--pty-source-credit-v1'] : []), '--grace-time', String(graceTime), '--sock-path', diff --git a/src/main/ssh/ssh-relay-session-data-delivery.test.ts b/src/main/ssh/ssh-relay-session-data-delivery.test.ts index f036ac24d7a..1f7e811fb98 100644 --- a/src/main/ssh/ssh-relay-session-data-delivery.test.ts +++ b/src/main/ssh/ssh-relay-session-data-delivery.test.ts @@ -5,15 +5,23 @@ import { createMockDeps, mockDeploySuccess } from './ssh-relay-session-test-fixt const { acceptOutputDataMock, muxRequestMock, + onNotificationByMethodMock, notifyWithSettlementMock, openConsumerSessionMock, - pauseAdapterMock + pauseAdapterMock, + muxDisposeMock, + attachForReconnectMock, + ptyDataHandlerRef } = vi.hoisted(() => ({ acceptOutputDataMock: vi.fn().mockResolvedValue(undefined), muxRequestMock: vi.fn(), + onNotificationByMethodMock: vi.fn(), notifyWithSettlementMock: vi.fn(), openConsumerSessionMock: vi.fn(), - pauseAdapterMock: vi.fn() + pauseAdapterMock: vi.fn(), + muxDisposeMock: vi.fn(), + attachForReconnectMock: vi.fn().mockResolvedValue({}), + ptyDataHandlerRef: { current: undefined as undefined | ((payload: unknown) => void) } })) vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) @@ -25,6 +33,7 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined), allocateSshPtyProviderGeneration: vi.fn(() => 23), closeSshPtyOutputGeneration: vi.fn(), + getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []), installSshPtySourceAckPublisher: vi.fn(() => () => {}) })) @@ -34,9 +43,10 @@ vi.mock('./ssh-channel-multiplexer', () => ({ notifyWithSettlement = notifyWithSettlementMock request = muxRequestMock onNotification = vi.fn().mockReturnValue(() => {}) + onNotificationByMethod = onNotificationByMethodMock.mockImplementation(() => () => {}) onRequest = vi.fn().mockReturnValue(() => {}) onDispose = vi.fn().mockReturnValue(() => {}) - dispose = vi.fn() + dispose = muxDisposeMock isDisposed = vi.fn().mockReturnValue(false) } })) @@ -49,11 +59,14 @@ vi.mock('../providers/ssh-pty-provider', () => ({ isSshPtyNotFoundError: vi.fn().mockReturnValue(false), isSshPtyIdentityMismatchError: vi.fn().mockReturnValue(false), SshPtyProvider: class MockSshPtyProvider { - onData = vi.fn().mockReturnValue(() => {}) + onData = vi.fn().mockImplementation((handler) => { + ptyDataHandlerRef.current = handler + return () => {} + }) onReplay = vi.fn().mockReturnValue(() => {}) onExit = vi.fn().mockReturnValue(() => {}) attach = vi.fn().mockResolvedValue(undefined) - attachForReconnect = vi.fn().mockResolvedValue({}) + attachForReconnect = attachForReconnectMock setPtyDeliveryPauseAdapter = pauseAdapterMock dispose = vi.fn() } @@ -77,6 +90,7 @@ vi.mock('../ipc/pty', () => ({ clearPtyOwnershipForConnection: vi.fn(), clearProviderPtyState: vi.fn(), deletePtyOwnership: vi.fn(), + restorePtyIncarnation: vi.fn(), setPtyOwnership: vi.fn() })) @@ -91,13 +105,18 @@ vi.mock('../providers/ssh-git-dispatch', () => ({ unregisterSshGitProvider: vi.fn() })) -const { registerSshPtyProvider } = await import('../ipc/pty') +const { getSshPtyProvider, getPtyIdsForConnection, registerSshPtyProvider } = + await import('../ipc/pty') const { closeSshPtyOutputGeneration } = await import('../ipc/ssh-pty-output-intake-registry') +const { getSshPtyAcceptedSourceCheckpoints } = await import('../ipc/ssh-pty-output-intake-registry') const { installSshPtySourceAckPublisher } = await import('../ipc/ssh-pty-output-intake-registry') +const { deployAndLaunchRelay } = await import('./ssh-relay-deploy') describe('SshRelaySession data delivery', () => { beforeEach(() => { vi.clearAllMocks() + ptyDataHandlerRef.current = undefined + attachForReconnectMock.mockResolvedValue({}) openConsumerSessionMock.mockImplementation(async (_mux, options) => ({ mode: 'negotiated', clientInstanceId: options.clientInstanceId, @@ -114,6 +133,36 @@ describe('SshRelaySession data delivery', () => { mockDeploySuccess() }) + it('transfers negotiated owner recovery exactly once across an explicit detach', async () => { + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + vi.mocked(deployAndLaunchRelay).mockResolvedValue({ + transport: { write: vi.fn(), onData: vi.fn(), onClose: vi.fn() }, + platform: 'linux-x64', + serverBuildId: 'test-relay-build' + }) + const first = new SshRelaySession('recovery-target', getMainWindow, mockStore, mockPortForward) + + await first.establish(mockConn) + first.detach() + const second = new SshRelaySession('recovery-target', getMainWindow, mockStore, mockPortForward) + await second.establish(mockConn) + + const firstOpen = openConsumerSessionMock.mock.calls[0]?.[1] + const recoveredOpen = openConsumerSessionMock.mock.calls[1]?.[1] + expect(recoveredOpen).toMatchObject({ + clientInstanceId: firstOpen.clientInstanceId, + resume: { ownerGeneration: 1, ownerLease: 'test-owner-lease' } + }) + + second.dispose() + const fresh = new SshRelaySession('recovery-target', getMainWindow, mockStore, mockPortForward) + await fresh.establish(mockConn) + const freshOpen = openConsumerSessionMock.mock.calls[2]?.[1] + expect(freshOpen.clientInstanceId).not.toBe(firstOpen.clientInstanceId) + expect(freshOpen).not.toHaveProperty('resume') + fresh.dispose() + }) + it('delivers empty transformed relay spans with raw sequence metadata', async () => { const { mockConn, mockStore, mockPortForward, getMainWindow, mockWindow } = createMockDeps() const runtime = { onPtyData: vi.fn(() => 17), onPtyExit: vi.fn() } @@ -321,6 +370,13 @@ describe('SshRelaySession data delivery', () => { expect(openConsumerSessionMock.mock.calls[0][1]).toMatchObject({ outputFlowControl: { requestedWindowSu: 256 * 1024 } }) + expect(deployAndLaunchRelay).toHaveBeenCalledWith( + mockConn, + undefined, + undefined, + 'target-1', + true + ) expect(notifyWithSettlementMock).toHaveBeenCalledWith('pty.ackData', batch, settled) }) @@ -343,5 +399,167 @@ describe('SshRelaySession data delivery', () => { expect(openConsumerSessionMock.mock.calls[0][1]).toHaveProperty('outputFlowControl') expect(openConsumerSessionMock.mock.calls[1][1]).not.toHaveProperty('outputFlowControl') + expect(deployAndLaunchRelay).toHaveBeenNthCalledWith( + 1, + mockConn, + undefined, + undefined, + 'target-1', + true + ) + expect(deployAndLaunchRelay).toHaveBeenNthCalledWith( + 2, + mockConn, + undefined, + undefined, + 'target-1', + false + ) + }) + + it('reattaches V1 from an exact checkpoint and quarantines live data until recoveryComplete', async () => { + let generation = 0 + openConsumerSessionMock.mockImplementation(async (_mux, options) => { + generation++ + return { + mode: 'negotiated', + clientInstanceId: options.clientInstanceId, + clientGeneration: generation, + ownerGeneration: generation, + ownerLease: `owner-lease-${generation}`, + outputFlowControl: { version: 1, windowSu: 256 * 1024 } + } + }) + vi.mocked(getSshPtyAcceptedSourceCheckpoints).mockReturnValue([ + { + id: 'ssh:target-1@@pty-1', + providerGeneration: 23, + clientGeneration: 1, + ownerGeneration: 1, + ptyIncarnation: 'incarnation-1', + deliveryToken: 'old-token', + acceptedSourceEndSu: 4 + } + ]) + const { mockConn, mockStore, mockPortForward, getMainWindow, mockWindow } = createMockDeps() + const session = new SshRelaySession( + 'target-1', + getMainWindow, + mockStore, + mockPortForward, + undefined, + undefined, + () => true + ) + await session.establish(mockConn) + vi.mocked(getPtyIdsForConnection).mockReturnValue(['ssh:target-1@@pty-1']) + vi.mocked(getSshPtyProvider).mockImplementation( + () => vi.mocked(registerSshPtyProvider).mock.calls.at(-1)?.[1] + ) + let transferDisposedMux = false + attachForReconnectMock.mockImplementation(async () => { + const canceled = onNotificationByMethodMock.mock.calls.findLast( + ([method]) => method === 'pty.deliveryCanceled' + )?.[1] as ((params: Record) => void) | undefined + const disposeCount = muxDisposeMock.mock.calls.length + canceled?.({ + id: 'pty-1', + clientGeneration: 1, + ownerGeneration: 1, + ptyIncarnation: 'incarnation-1', + deliveryToken: 'old-token', + replacementDeliveryToken: 'new-token' + }) + transferDisposedMux = muxDisposeMock.mock.calls.length !== disposeCount + queueMicrotask(() => { + ptyDataHandlerRef.current?.({ + id: 'ssh:target-1@@pty-1', + data: 'reco', + providerGeneration: 23, + ptyIncarnation: 'incarnation-1', + sequenceChars: 4, + source: { + relayPtyId: 'pty-1', + spanId: 'new-token:4:8', + clientGeneration: 2, + ownerGeneration: 2, + deliveryToken: 'new-token', + sourceStartSu: 4, + sourceEndSu: 8 + } + }) + const complete = onNotificationByMethodMock.mock.calls.findLast( + ([method]) => method === 'pty.recoveryComplete' + )?.[1] as ((params: Record) => void) | undefined + complete?.({ + id: 'pty-1', + clientGeneration: 2, + ownerGeneration: 2, + ptyIncarnation: 'incarnation-1', + deliveryToken: 'new-token', + checkpointSourceEndSu: 4, + recoveryEndSu: 8 + }) + ptyDataHandlerRef.current?.({ + id: 'ssh:target-1@@pty-1', + data: 'live', + providerGeneration: 23, + ptyIncarnation: 'incarnation-1', + sequenceChars: 4, + source: { + relayPtyId: 'pty-1', + spanId: 'new-token:8:12', + clientGeneration: 2, + ownerGeneration: 2, + deliveryToken: 'new-token', + sourceStartSu: 8, + sourceEndSu: 12 + } + }) + }) + return { + incarnationId: 'incarnation-1', + sourceRecovery: { + status: 'pending', + clientGeneration: 2, + ownerGeneration: 2, + ptyIncarnation: 'incarnation-1', + deliveryToken: 'new-token', + checkpointSourceEndSu: 4, + recoveryEndSu: 8 + } + } + }) + + await session.reconnect(mockConn) + + expect(attachForReconnectMock).toHaveBeenCalledWith( + 'pty-1', + undefined, + expect.objectContaining({ + status: 'checkpoint', + deliveryToken: 'old-token', + acceptedSourceEndSu: 4 + }) + ) + expect(acceptOutputDataMock.mock.calls.map(([payload]) => payload.data)).toEqual([ + 'reco', + 'live' + ]) + expect(transferDisposedMux).toBe(false) + expect(mockWindow.webContents.send).not.toHaveBeenCalledWith('pty:replay', expect.anything()) + + const closeCount = vi.mocked(closeSshPtyOutputGeneration).mock.calls.length + const canceled = onNotificationByMethodMock.mock.calls.findLast( + ([method]) => method === 'pty.deliveryCanceled' + )?.[1] as ((params: Record) => void) | undefined + canceled?.({ + id: 'pty-1', + clientGeneration: 1, + ownerGeneration: 1, + ptyIncarnation: 'incarnation-1', + deliveryToken: 'old-token' + }) + expect(closeSshPtyOutputGeneration).toHaveBeenCalledTimes(closeCount) }) }) diff --git a/src/main/ssh/ssh-relay-session-incarnation.test.ts b/src/main/ssh/ssh-relay-session-incarnation.test.ts index 2e9312d7d4f..e782348eda7 100644 --- a/src/main/ssh/ssh-relay-session-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-incarnation.test.ts @@ -21,6 +21,7 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputExit: acceptOutputExitMock, allocateSshPtyProviderGeneration: vi.fn(() => 31), closeSshPtyOutputGeneration: vi.fn(), + getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []), installSshPtySourceAckPublisher: vi.fn(() => () => {}) })) vi.mock('./ssh-relay-deploy-helpers', () => ({ execCommand: vi.fn().mockResolvedValue('') })) @@ -30,6 +31,7 @@ vi.mock('./ssh-channel-multiplexer', () => ({ notifyWithSettlement = vi.fn() request = muxRequestMock onNotification = vi.fn().mockReturnValue(() => {}) + onNotificationByMethod = vi.fn().mockReturnValue(() => {}) onRequest = vi.fn().mockReturnValue(() => {}) onDispose = vi.fn().mockReturnValue(() => {}) dispose = vi.fn() diff --git a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts index cd48b9c3c3d..1a7159da4e3 100644 --- a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts @@ -21,6 +21,7 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputExit: acceptOutputExitMock, allocateSshPtyProviderGeneration: vi.fn(() => 17), closeSshPtyOutputGeneration: vi.fn(), + getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []), installSshPtySourceAckPublisher: vi.fn(() => () => {}) })) vi.mock('./ssh-relay-deploy-helpers', () => ({ execCommand: vi.fn().mockResolvedValue('') })) @@ -30,6 +31,7 @@ vi.mock('./ssh-channel-multiplexer', () => ({ notifyWithSettlement = vi.fn() request = muxRequestMock onNotification = vi.fn().mockReturnValue(() => {}) + onNotificationByMethod = vi.fn().mockReturnValue(() => {}) onRequest = vi.fn().mockReturnValue(() => {}) onDispose = vi.fn().mockReturnValue(() => {}) dispose = vi.fn() diff --git a/src/main/ssh/ssh-relay-session.test.ts b/src/main/ssh/ssh-relay-session.test.ts index 48322a31d34..f6c9c80fc03 100644 --- a/src/main/ssh/ssh-relay-session.test.ts +++ b/src/main/ssh/ssh-relay-session.test.ts @@ -30,6 +30,7 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined), allocateSshPtyProviderGeneration: vi.fn(() => 41), closeSshPtyOutputGeneration: vi.fn(), + getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []), installSshPtySourceAckPublisher: vi.fn(() => () => {}) })) @@ -44,6 +45,7 @@ vi.mock('./ssh-channel-multiplexer', () => { notifyWithSettlement = vi.fn() request = muxRequestMock onNotification = vi.fn().mockReturnValue(() => {}) + onNotificationByMethod = vi.fn().mockReturnValue(() => {}) onRequest = vi.fn().mockReturnValue(() => {}) onDispose = vi.fn().mockReturnValue(() => {}) dispose = vi.fn() @@ -820,7 +822,7 @@ describe('SshRelaySession', () => { await session.establish(mockConn, 600) - expect(deployAndLaunchRelay).toHaveBeenCalledWith(mockConn, undefined, 600, 'target-1') + expect(deployAndLaunchRelay).toHaveBeenCalledWith(mockConn, undefined, 600, 'target-1', false) }) it('restores the configured relay grace after establish', async () => { diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index 9d42ef8dd81..cea22a190de 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -9,7 +9,8 @@ import { isRelayVersionMismatchError } from './ssh-relay-version-mismatch-error' import type { RelayVersionMismatchError } from './ssh-relay-version-mismatch-error' import { SshChannelMultiplexer } from './ssh-channel-multiplexer' import { SshPtyProvider } from '../providers/ssh-pty-provider' -import type { SshPtyExitCallback } from '../providers/ssh-pty-provider-contract' +import type { SshPtyAttachResult } from '../providers/ssh-pty-session-reattach' +import type { SshPtyDataCallback, SshPtyExitCallback } from '../providers/ssh-pty-provider-contract' import { isSshPtyIdentityMismatchError, isSshPtyNotFoundError } from '../providers/ssh-pty-errors' import { toAppSshPtyId, toRelaySshPtyId } from '../providers/ssh-pty-id' import { SshFilesystemProvider } from '../providers/ssh-filesystem-provider' @@ -42,8 +43,10 @@ import { acceptSshPtyOutputExit, allocateSshPtyProviderGeneration, closeSshPtyOutputGeneration, + getSshPtyAcceptedSourceCheckpoints, installSshPtySourceAckPublisher } from '../ipc/ssh-pty-output-intake-registry' +import type { SshPtyAcceptedSourceCheckpoint } from '../ipc/ssh-pty-output-source-obligations' import { registerSshFilesystemProvider, unregisterSshFilesystemProvider, @@ -77,11 +80,33 @@ import { type SshPtyConsumerOwnerState, type SshPtyConsumerSessionState } from './ssh-pty-consumer-session' +import type { + PtySourceRecoveryComplete, + PtySourceRecoveryPending, + PtySourceRecoveryRequest +} from '../../shared/pty-source-recovery-contract' export type RelaySessionState = 'idle' | 'deploying' | 'ready' | 'reconnecting' | 'disposed' type SshPtyExitPayload = Parameters[0] -type PendingPtyReattach = { exits: SshPtyExitPayload[] } +type SshPtyDataPayload = Parameters[0] +const SSH_PTY_RECOVERY_QUARANTINE_MAX_BYTES = DEFAULT_PTY_SOURCE_WINDOW_SU * 8 +const SSH_PTY_RECOVERY_QUARANTINE_MAX_FRAMES = 1_024 +type PendingPtyReattach = { + exits: SshPtyExitPayload[] + queuedData: SshPtyDataPayload[] + liveData: SshPtyDataPayload[] + recovery?: PtySourceRecoveryPending + recoveryComplete?: PtySourceRecoveryComplete + replacementDeliveryToken?: string + restoreRequired?: string + recoveryWaiters: Set<() => void> + recoveryAdmissions: Promise[] + queuedBytes: number + queuedFrames: number + livePassthrough: boolean + activated: boolean +} type RemoteCliBridgeEnv = { remoteHome: string @@ -115,6 +140,40 @@ function expectedIdentityForLease(lease: { } } +function parseRecoveryComplete(params: Record): PtySourceRecoveryComplete | null { + if ( + typeof params.id !== 'string' || + typeof params.deliveryToken !== 'string' || + params.deliveryToken.length === 0 || + typeof params.ptyIncarnation !== 'string' || + params.ptyIncarnation.length === 0 || + !positiveSafeInteger(params.clientGeneration) || + !positiveSafeInteger(params.ownerGeneration) || + !nonNegativeSafeInteger(params.checkpointSourceEndSu) || + !nonNegativeSafeInteger(params.recoveryEndSu) || + Number(params.recoveryEndSu) < Number(params.checkpointSourceEndSu) + ) { + return null + } + return Object.freeze({ + id: params.id, + deliveryToken: params.deliveryToken, + ptyIncarnation: params.ptyIncarnation, + clientGeneration: Number(params.clientGeneration), + ownerGeneration: Number(params.ownerGeneration), + checkpointSourceEndSu: Number(params.checkpointSourceEndSu), + recoveryEndSu: Number(params.recoveryEndSu) + }) +} + +function positiveSafeInteger(value: unknown): boolean { + return Number.isSafeInteger(value) && Number(value) > 0 +} + +function nonNegativeSafeInteger(value: unknown): boolean { + return Number.isSafeInteger(value) && Number(value) >= 0 +} + export type SshRelayAiVaultHostInfo = { targetId: string executionHostId: ExecutionHostId @@ -133,6 +192,31 @@ function normalizeRelayGracePeriodSeconds(graceTimeSeconds: number | undefined): ) } +type PtyConsumerRecovery = { + clientInstanceId: string + detached: boolean + serverBuildId?: string + owner?: SshPtyConsumerOwnerState + checkpointsByAppPtyId: Map +} + +const ptyConsumerRecoveryByTarget = new Map() + +function ptyConsumerRecoveryForTarget(targetId: string): PtyConsumerRecovery { + const current = ptyConsumerRecoveryByTarget.get(targetId) + if (current?.detached) { + current.detached = false + return current + } + const created = { + clientInstanceId: randomUUID(), + detached: false, + checkpointsByAppPtyId: new Map() + } + ptyConsumerRecoveryByTarget.set(targetId, created) + return created +} + export class SshRelaySession { private _state: RelaySessionState = 'idle' private mux: SshChannelMultiplexer | null = null @@ -154,7 +238,17 @@ export class SshRelaySession { private pendingPtyReattaches = new Map() private activePtyProviderGeneration: number | null = null private sourceAckPublisherCleanup: (() => void) | null = null - private readonly ptyConsumerClientInstanceId = randomUUID() + private ptyRecoveryNotificationCleanups: (() => void)[] = [] + private readonly sourceIdentityByRelayPtyId = new Map< + string, + Readonly<{ + deliveryToken: string + clientGeneration: number + ownerGeneration: number + ptyIncarnation: string + }> + >() + private readonly ptyConsumerClientInstanceId: string private ptyConsumerSessionState: SshPtyConsumerSessionState | null = null constructor( @@ -169,7 +263,9 @@ export class SshRelaySession { platform: string ) => void, private readonly isPtySourceCreditEnabled: () => boolean = () => false - ) {} + ) { + this.ptyConsumerClientInstanceId = ptyConsumerRecoveryForTarget(targetId).clientInstanceId + } refreshEnvironment( getMainWindow: () => BrowserWindow | null, @@ -264,7 +360,13 @@ export class SshRelaySession { sockPath, credentialFile, hostPlatform - } = await deployAndLaunchRelay(conn, undefined, graceTimeSeconds, this.targetId) + } = await deployAndLaunchRelay( + conn, + undefined, + graceTimeSeconds, + this.targetId, + this.isPtySourceCreditEnabled() + ) this.hostPlatform = hostPlatform ?? null this.remoteCliBridgeEnv = remoteHome && remoteRelayDir && nodePath && sockPath && hostPlatform @@ -291,7 +393,7 @@ export class SshRelaySession { this.mux = mux const ownsAttempt = (): boolean => this.mux === mux && !this.isDisposed() - const previousOwner = this.negotiatedPtyConsumerOwner() + const previousOwner = this.negotiatedPtyConsumerOwner(serverBuildId) this.ptyConsumerSessionState = await openSshPtyConsumerSession(mux, { clientInstanceId: this.ptyConsumerClientInstanceId, expectedServerBuildId: serverBuildId, @@ -308,6 +410,7 @@ export class SshRelaySession { } : {}) }) + this.rememberPtyConsumerRecovery(serverBuildId) await mux.request('session.resolveHome', { path: '~' }) @@ -389,7 +492,13 @@ export class SshRelaySession { sockPath, credentialFile, hostPlatform - } = await deployAndLaunchRelay(conn, undefined, graceTimeSeconds, this.targetId) + } = await deployAndLaunchRelay( + conn, + undefined, + graceTimeSeconds, + this.targetId, + this.isPtySourceCreditEnabled() + ) this.hostPlatform = hostPlatform ?? null this.remoteCliBridgeEnv = remoteHome && remoteRelayDir && nodePath && sockPath && hostPlatform @@ -420,7 +529,7 @@ export class SshRelaySession { !abortController.signal.aborted && !this.isDisposed() - const previousOwner = this.negotiatedPtyConsumerOwner() + const previousOwner = this.negotiatedPtyConsumerOwner(serverBuildId) this.ptyConsumerSessionState = await openSshPtyConsumerSession(mux, { clientInstanceId: this.ptyConsumerClientInstanceId, expectedServerBuildId: serverBuildId, @@ -437,6 +546,7 @@ export class SshRelaySession { } : {}) }) + this.rememberPtyConsumerRecovery(serverBuildId) if (!ownsAttempt()) { if (!mux.isDisposed()) { mux.dispose() @@ -529,6 +639,7 @@ export class SshRelaySession { this.store.markSshRemotePtyLeases(this.targetId, 'terminated') this.currentConnection = null this._state = 'disposed' + ptyConsumerRecoveryByTarget.delete(this.targetId) } detach(): void { @@ -543,6 +654,10 @@ export class SshRelaySession { this.store.markSshRemotePtyLeases(this.targetId, 'detached') this.currentConnection = null this._state = 'disposed' + const recovery = ptyConsumerRecoveryByTarget.get(this.targetId) + if (recovery?.clientInstanceId === this.ptyConsumerClientInstanceId) { + recovery.detached = true + } } // ── Private ─────────────────────────────────────────────────────── @@ -613,11 +728,16 @@ export class SshRelaySession { ) { return } + const sourceIdentity = this.sourceIdentityByRelayPtyId.get(id) + if (consumerOwnerState.outputFlowControl && !sourceIdentity) { + return + } mux.notify('pty.setDeliveryPaused', { id, paused, clientGeneration: consumerOwnerState.clientGeneration, - ownerGeneration: consumerOwnerState.ownerGeneration + ownerGeneration: consumerOwnerState.ownerGeneration, + ...(sourceIdentity ? { deliveryToken: sourceIdentity.deliveryToken } : {}) }) }) } @@ -636,6 +756,7 @@ export class SshRelaySession { } this.activePtyProviderGeneration = providerGeneration registerSshPtyProvider(this.targetId, ptyProvider) + this.installPtyRecoveryNotifications(mux) const connection = this.requireReadyConnection() const createSftp = @@ -681,9 +802,31 @@ export class SshRelaySession { return true } - private negotiatedPtyConsumerOwner(): SshPtyConsumerOwnerState | null { + private negotiatedPtyConsumerOwner(serverBuildId?: string): SshPtyConsumerOwnerState | null { const state = this.ptyConsumerSessionState - return state && state.mode !== 'legacy-fallback' ? (state as SshPtyConsumerOwnerState) : null + if (state && state.mode !== 'legacy-fallback') { + return state as SshPtyConsumerOwnerState + } + const recovery = ptyConsumerRecoveryByTarget.get(this.targetId) + return !serverBuildId || recovery?.serverBuildId === serverBuildId + ? (recovery?.owner ?? null) + : null + } + + private rememberPtyConsumerRecovery(serverBuildId: string | undefined): void { + const owner = this.negotiatedPtyConsumerOwner() + if (!owner || !serverBuildId) { + return + } + const previous = ptyConsumerRecoveryByTarget.get(this.targetId) + ptyConsumerRecoveryByTarget.set(this.targetId, { + clientInstanceId: this.ptyConsumerClientInstanceId, + detached: false, + serverBuildId, + owner, + checkpointsByAppPtyId: + previous?.checkpointsByAppPtyId ?? new Map() + }) } private configureRelayGraceTime( @@ -929,7 +1072,21 @@ export class SshRelaySession { this.muxDisposeCleanup = null this.muxNotificationCleanup?.() this.muxNotificationCleanup = null + for (const cleanup of this.ptyRecoveryNotificationCleanups) { + cleanup() + } + this.ptyRecoveryNotificationCleanups = [] if (this.activePtyProviderGeneration !== null) { + if (reason === 'connection_lost' && this.negotiatedPtyConsumerOwner()?.outputFlowControl) { + const recovery = ptyConsumerRecoveryByTarget.get(this.targetId) + if (recovery) { + for (const checkpoint of getSshPtyAcceptedSourceCheckpoints( + this.activePtyProviderGeneration + )) { + recovery.checkpointsByAppPtyId.set(checkpoint.id, checkpoint) + } + } + } closeSshPtyOutputGeneration(this.activePtyProviderGeneration, reason) this.activePtyProviderGeneration = null } @@ -959,6 +1116,13 @@ export class SshRelaySession { unregisterSshPtyProvider(this.targetId) unregisterSshFilesystemProvider(this.targetId) unregisterSshGitProvider(this.targetId) + this.sourceIdentityByRelayPtyId.clear() + for (const pending of this.pendingPtyReattaches.values()) { + for (const resolve of pending.recoveryWaiters) { + resolve() + } + } + this.pendingPtyReattaches.clear() } // Why: back-compat for old relays that gate FS ops on registered roots; removable post-cutover (docs/relay-fs-allowlist-removal.md). @@ -1032,33 +1196,16 @@ export class SshRelaySession { private wireUpPtyEvents(ptyProvider: SshPtyProvider): void { ptyProvider.onData((payload) => { - const consumerOwner = this.negotiatedPtyConsumerOwner() - if ( - consumerOwner?.outputFlowControl && - (!payload.source || - payload.sourceMalformed || - payload.source.clientGeneration !== consumerOwner.clientGeneration || - payload.source.ownerGeneration !== consumerOwner.ownerGeneration) - ) { - closeSshPtyOutputGeneration( - payload.providerGeneration, - 'ssh_source_frame_malformed_or_missing' - ) - this.mux?.dispose('connection_lost') + const pending = this.pendingPtyReattaches.get(payload.id) + if (pending && this.negotiatedPtyConsumerOwner()?.outputFlowControl) { + if (pending.livePassthrough) { + void this.acceptPtyData(payload).catch(() => {}) + return + } + this.quarantineReattachData(pending, payload) return } - const source = consumerOwner?.outputFlowControl ? payload.source : undefined - const rawLength = payload.sequenceChars ?? payload.data.length - void acceptSshPtyOutputData({ - id: payload.id, - data: payload.data, - providerGeneration: payload.providerGeneration, - ptyIncarnation: payload.ptyIncarnation, - rawLength, - transformed: payload.transformed === true, - ...(typeof payload.seq === 'number' ? { sequence: payload.seq } : {}), - ...(source ? { source } : {}) - }).catch(() => {}) + void this.acceptPtyData(payload).catch(() => {}) }) ptyProvider.onReplay((payload) => { const win = this.getMainWindow() @@ -1068,7 +1215,7 @@ export class SshRelaySession { }) ptyProvider.onExit((payload) => { const pendingReattach = this.pendingPtyReattaches.get(payload.id) - if (pendingReattach) { + if (pendingReattach && !pendingReattach.activated) { // Why: attach response and exit can share one transport batch, before incarnation restoration runs. pendingReattach.exits.push(payload) return @@ -1080,6 +1227,173 @@ export class SshRelaySession { }) } + private acceptPtyData(payload: SshPtyDataPayload): Promise { + const consumerOwner = this.negotiatedPtyConsumerOwner() + if ( + consumerOwner?.outputFlowControl && + (!payload.source || + payload.sourceMalformed || + payload.source.clientGeneration !== consumerOwner.clientGeneration || + payload.source.ownerGeneration !== consumerOwner.ownerGeneration) + ) { + closeSshPtyOutputGeneration( + payload.providerGeneration, + 'ssh_source_frame_malformed_or_missing' + ) + this.mux?.dispose('connection_lost') + return Promise.reject(new Error('ssh_source_frame_malformed_or_missing')) + } + const source = consumerOwner?.outputFlowControl ? payload.source : undefined + if (source) { + this.sourceIdentityByRelayPtyId.set(source.relayPtyId, { + deliveryToken: source.deliveryToken, + clientGeneration: source.clientGeneration, + ownerGeneration: source.ownerGeneration, + ptyIncarnation: payload.ptyIncarnation + }) + } + const rawLength = payload.sequenceChars ?? payload.data.length + return acceptSshPtyOutputData({ + id: payload.id, + data: payload.data, + providerGeneration: payload.providerGeneration, + ptyIncarnation: payload.ptyIncarnation, + rawLength, + transformed: payload.transformed === true, + ...(typeof payload.seq === 'number' ? { sequence: payload.seq } : {}), + ...(source ? { source } : {}) + }) + } + + private quarantineReattachData(pending: PendingPtyReattach, payload: SshPtyDataPayload): void { + pending.queuedBytes += Buffer.byteLength(payload.data, 'utf8') + pending.queuedFrames++ + if ( + pending.queuedBytes > SSH_PTY_RECOVERY_QUARANTINE_MAX_BYTES || + pending.queuedFrames > SSH_PTY_RECOVERY_QUARANTINE_MAX_FRAMES + ) { + pending.restoreRequired = 'recoveryQuarantineCapacityExceeded' + this.wakeRecovery(pending) + return + } + if (pending.recoveryComplete) { + pending.liveData.push(payload) + return + } + if (!pending.recovery) { + pending.queuedData.push(payload) + return + } + this.admitRecoveryData(pending, payload) + } + + private admitRecoveryData(pending: PendingPtyReattach, payload: SshPtyDataPayload): void { + const recovery = pending.recovery + if ( + !recovery || + !payload.source || + payload.source.deliveryToken !== recovery.deliveryToken || + payload.source.clientGeneration !== recovery.clientGeneration || + payload.source.ownerGeneration !== recovery.ownerGeneration || + payload.source.sourceEndSu > recovery.recoveryEndSu || + payload.ptyIncarnation !== recovery.ptyIncarnation + ) { + pending.restoreRequired = 'recoveryFrameIdentityMismatch' + this.wakeRecovery(pending) + return + } + const admission = this.acceptPtyData(payload).catch((error) => { + pending.restoreRequired = + error instanceof Error ? error.message : 'recoveryIntakeAdmissionFailed' + this.wakeRecovery(pending) + throw error + }) + pending.recoveryAdmissions.push(admission) + } + + private installPtyRecoveryNotifications(mux: SshChannelMultiplexer): void { + for (const cleanup of this.ptyRecoveryNotificationCleanups) { + cleanup() + } + this.ptyRecoveryNotificationCleanups = [ + mux.onNotificationByMethod('pty.recoveryComplete', (params) => { + const id = typeof params.id === 'string' ? toAppSshPtyId(this.targetId, params.id) : '' + const pending = this.pendingPtyReattaches.get(id) + if (!pending) { + return + } + const complete = parseRecoveryComplete(params) + if (!complete) { + pending.restoreRequired = 'invalidRecoveryComplete' + } else { + pending.recoveryComplete = complete + } + this.wakeRecovery(pending) + }), + mux.onNotificationByMethod('pty.restoreRequired', (params) => { + const id = typeof params.id === 'string' ? toAppSshPtyId(this.targetId, params.id) : '' + const pending = this.pendingPtyReattaches.get(id) + if (!pending) { + return + } + pending.restoreRequired = + typeof params.reason === 'string' ? params.reason : 'relayRestoreRequired' + this.wakeRecovery(pending) + }), + mux.onNotificationByMethod('pty.deliveryCanceled', (params) => { + const id = typeof params.id === 'string' ? params.id : '' + const identity = this.sourceIdentityByRelayPtyId.get(id) + if ( + !identity || + params.deliveryToken !== identity.deliveryToken || + params.clientGeneration !== identity.clientGeneration || + params.ownerGeneration !== identity.ownerGeneration || + params.ptyIncarnation !== identity.ptyIncarnation + ) { + return + } + const replacementDeliveryToken = + typeof params.replacementDeliveryToken === 'string' ? params.replacementDeliveryToken : '' + const pending = this.pendingPtyReattaches.get(toAppSshPtyId(this.targetId, id)) + if (pending) { + if ( + replacementDeliveryToken.length === 0 || + replacementDeliveryToken === identity.deliveryToken + ) { + pending.restoreRequired = + typeof params.reason === 'string' + ? `relayDeliveryCanceled:${params.reason}` + : 'relayDeliveryCanceled' + this.wakeRecovery(pending) + return + } + if ( + pending.replacementDeliveryToken && + pending.replacementDeliveryToken !== replacementDeliveryToken + ) { + pending.restoreRequired = 'recoveryReplacementTokenMismatch' + this.wakeRecovery(pending) + return + } + pending.replacementDeliveryToken = replacementDeliveryToken + return + } + const generation = this.activePtyProviderGeneration + if (generation !== null) { + closeSshPtyOutputGeneration(generation, 'ssh_source_delivery_canceled') + } + mux.dispose('connection_lost') + }) + ] + } + + private wakeRecovery(pending: PendingPtyReattach): void { + for (const resolve of pending.recoveryWaiters) { + resolve() + } + pending.recoveryWaiters.clear() + } + private async acceptPtyExit(payload: SshPtyExitPayload): Promise { await acceptSshPtyOutputExit({ id: payload.id, @@ -1096,6 +1410,10 @@ export class SshRelaySession { const relayPtyId = toRelaySshPtyId(this.targetId, payload.id) clearProviderPtyState(payload.id) deletePtyOwnership(payload.id) + ptyConsumerRecoveryByTarget.get(this.targetId)?.checkpointsByAppPtyId.delete(payload.id) + ptyConsumerRecoveryByTarget + .get(this.targetId) + ?.checkpointsByAppPtyId.delete(toRelaySshPtyId(this.targetId, payload.id)) this.store.markSshRemotePtyLease(this.targetId, relayPtyId, 'terminated') if (deliveryHandled) { return @@ -1150,14 +1468,29 @@ export class SshRelaySession { return } const appPtyId = toAppSshPtyId(this.targetId, ptyId) - const pendingReattach: PendingPtyReattach = { exits: [] } + const pendingReattach: PendingPtyReattach = { + exits: [], + queuedData: [], + liveData: [], + recoveryWaiters: new Set(), + recoveryAdmissions: [], + queuedBytes: 0, + queuedFrames: 0, + livePassthrough: false, + activated: false + } this.pendingPtyReattaches.set(appPtyId, pendingReattach) try { const expectedIdentity = expectedIdentityByPtyId.get(ptyId) + const recoveryRequest = this.sourceRecoveryRequest(appPtyId) const attachResult = (expectedIdentity - ? await ptyProvider.attachForReconnect(ptyId, expectedIdentity) - : await ptyProvider.attachForReconnect(ptyId)) ?? {} + ? recoveryRequest + ? await ptyProvider.attachForReconnect(ptyId, expectedIdentity, recoveryRequest) + : await ptyProvider.attachForReconnect(ptyId, expectedIdentity) + : recoveryRequest + ? await ptyProvider.attachForReconnect(ptyId, undefined, recoveryRequest) + : await ptyProvider.attachForReconnect(ptyId)) ?? {} if (!shouldContinue()) { return } @@ -1175,6 +1508,20 @@ export class SshRelaySession { await this.acceptPtyExit(exitDuringAttach) continue } + if ( + recoveryRequest && + !(await this.finishSourceRecovery( + ptyProvider, + ptyId, + appPtyId, + attachResult, + recoveryRequest, + pendingReattach, + shouldContinue + )) + ) { + continue + } setPtyOwnership(appPtyId, this.targetId) if (attachResult.incarnationId) { restorePtyIncarnation(appPtyId, attachResult.incarnationId) @@ -1205,7 +1552,20 @@ export class SshRelaySession { } } this.store.markSshRemotePtyLease(this.targetId, ptyId, 'attached') - this.forwardReattachReplay(appPtyId, attachResult.replay ?? '') + pendingReattach.activated = true + const exitAfterActivation = pendingReattach.exits.find( + (exit) => + !exit.incarnationId || + !attachResult.incarnationId || + exit.incarnationId === attachResult.incarnationId + ) + if (exitAfterActivation) { + await this.acceptPtyExit(exitAfterActivation) + continue + } + if (!recoveryRequest) { + this.forwardReattachReplay(appPtyId, attachResult.replay ?? '') + } } catch (err) { if (!isSshPtyNotFoundError(err)) { throw err @@ -1238,4 +1598,159 @@ export class SshRelaySession { } } } + + private sourceRecoveryRequest(appPtyId: string): PtySourceRecoveryRequest | undefined { + if (!this.negotiatedPtyConsumerOwner()?.outputFlowControl) { + return undefined + } + const checkpoints = ptyConsumerRecoveryByTarget.get(this.targetId)?.checkpointsByAppPtyId + const relayPtyId = toRelaySshPtyId(this.targetId, appPtyId) + const checkpoint = checkpoints?.get(appPtyId) ?? checkpoints?.get(relayPtyId) + if (!checkpoint) { + return Object.freeze({ status: 'checkpointUnavailable' }) + } + return Object.freeze({ + status: 'checkpoint', + clientGeneration: checkpoint.clientGeneration, + ownerGeneration: checkpoint.ownerGeneration, + ptyIncarnation: checkpoint.ptyIncarnation, + deliveryToken: checkpoint.deliveryToken, + acceptedSourceEndSu: checkpoint.acceptedSourceEndSu + }) + } + + private async finishSourceRecovery( + ptyProvider: SshPtyProvider, + relayPtyId: string, + appPtyId: string, + attachResult: SshPtyAttachResult, + request: PtySourceRecoveryRequest, + pending: PendingPtyReattach, + shouldContinue: () => boolean + ): Promise { + const recovery = attachResult.sourceRecovery + const owner = this.negotiatedPtyConsumerOwner() + if ( + !owner?.outputFlowControl || + recovery?.status !== 'pending' || + request.status !== 'checkpoint' || + recovery.clientGeneration !== owner.clientGeneration || + recovery.ownerGeneration !== owner.ownerGeneration || + recovery.ptyIncarnation !== attachResult.incarnationId || + recovery.ptyIncarnation !== request.ptyIncarnation || + recovery.checkpointSourceEndSu !== request.acceptedSourceEndSu || + (pending.replacementDeliveryToken !== undefined && + pending.replacementDeliveryToken !== recovery.deliveryToken) + ) { + await this.resetPtyAfterRecoveryFailure(ptyProvider, relayPtyId, appPtyId) + return false + } + pending.recovery = recovery + this.sourceIdentityByRelayPtyId.set(relayPtyId, { + deliveryToken: recovery.deliveryToken, + clientGeneration: recovery.clientGeneration, + ownerGeneration: recovery.ownerGeneration, + ptyIncarnation: recovery.ptyIncarnation + }) + for (const payload of pending.queuedData.splice(0)) { + this.admitRecoveryData(pending, payload) + } + await this.waitForRecoveryFence(pending, shouldContinue) + const complete = pending.recoveryComplete + if ( + !shouldContinue() || + pending.restoreRequired || + !complete || + complete.deliveryToken !== recovery.deliveryToken || + complete.clientGeneration !== recovery.clientGeneration || + complete.ownerGeneration !== recovery.ownerGeneration || + complete.ptyIncarnation !== recovery.ptyIncarnation || + complete.checkpointSourceEndSu !== recovery.checkpointSourceEndSu || + complete.recoveryEndSu !== recovery.recoveryEndSu + ) { + await this.resetPtyAfterRecoveryFailure(ptyProvider, relayPtyId, appPtyId) + return false + } + try { + await Promise.all(pending.recoveryAdmissions) + for (const payload of pending.liveData) { + await this.acceptPtyData(payload) + } + pending.livePassthrough = true + } catch { + await this.resetPtyAfterRecoveryFailure(ptyProvider, relayPtyId, appPtyId) + return false + } + const acceptedSourceEndSu = pending.liveData.reduce( + (endSu, payload) => Math.max(endSu, payload.source?.sourceEndSu ?? endSu), + recovery.recoveryEndSu + ) + ptyConsumerRecoveryByTarget.get(this.targetId)?.checkpointsByAppPtyId.set( + relayPtyId, + Object.freeze({ + id: relayPtyId, + providerGeneration: this.activePtyProviderGeneration!, + clientGeneration: recovery.clientGeneration, + ownerGeneration: recovery.ownerGeneration, + ptyIncarnation: recovery.ptyIncarnation, + deliveryToken: recovery.deliveryToken, + acceptedSourceEndSu + }) + ) + return true + } + + private async waitForRecoveryFence( + pending: PendingPtyReattach, + shouldContinue: () => boolean + ): Promise { + const deadline = Date.now() + 30_000 + while ( + shouldContinue() && + !pending.recoveryComplete && + !pending.restoreRequired && + Date.now() < deadline + ) { + await new Promise((resolve) => { + const timer = setTimeout( + () => { + pending.recoveryWaiters.delete(settle) + resolve() + }, + Math.max(1, deadline - Date.now()) + ) + timer.unref?.() + const settle = (): void => { + clearTimeout(timer) + resolve() + } + pending.recoveryWaiters.add(settle) + }) + } + if (!pending.recoveryComplete && !pending.restoreRequired) { + pending.restoreRequired = 'recoveryFenceTimeout' + } + } + + private async resetPtyAfterRecoveryFailure( + ptyProvider: SshPtyProvider, + relayPtyId: string, + appPtyId: string + ): Promise { + try { + await ptyProvider.shutdown(appPtyId, { immediate: true, deadlineMs: Date.now() + 5_000 }) + } catch { + /* The recovery fence already prevents stale output activation. */ + } + clearProviderPtyState(appPtyId) + deletePtyOwnership(appPtyId) + this.sourceIdentityByRelayPtyId.delete(relayPtyId) + ptyConsumerRecoveryByTarget.get(this.targetId)?.checkpointsByAppPtyId.delete(appPtyId) + ptyConsumerRecoveryByTarget.get(this.targetId)?.checkpointsByAppPtyId.delete(relayPtyId) + this.store.markSshRemotePtyLease(this.targetId, relayPtyId, 'expired') + const win = this.getMainWindow() + if (win && !win.isDestroyed()) { + win.webContents.send('pty:exit', { id: appPtyId, code: -1 }) + } + } } diff --git a/src/relay/dispatcher-client-writer.test.ts b/src/relay/dispatcher-client-writer.test.ts index 4fcf758492c..9218844bc1e 100644 --- a/src/relay/dispatcher-client-writer.test.ts +++ b/src/relay/dispatcher-client-writer.test.ts @@ -300,6 +300,26 @@ describe('DispatcherClientWriter', () => { expect(sink.accepted[1].data).toBe('cancel') }) + it('admits one bounded producer frame larger than an empty native sink window', () => { + const sink = new FakeSink() + sink.highWaterMark = 4096 + sink.saturateNext = true + const writer = createWriter(sink, vi.fn(), 8192) + const oversized = 'x'.repeat(5000) + + expect(writer.producerFrameCapacity).toBe(3072) + const first = enqueue(writer, 'bulk', oversized) + enqueue(writer, 'ordinary', 'after-drain') + expect(writer.enqueue('bulk', () => Buffer.alloc(8193), 8193)).toBe(false) + expect(sink.accepted.map((write) => write.data)).toEqual([oversized]) + + sink.accepted[0].settle({ ok: true }) + expect(first).toHaveBeenCalledWith({ ok: true }) + expect(sink.accepted).toHaveLength(1) + sink.drain() + expect(sink.accepted.map((write) => write.data)).toEqual([oversized, 'after-drain']) + }) + it('gives independent writers progress when one client stays saturated', () => { const slowSink = new FakeSink() const fastSink = new FakeSink() diff --git a/src/relay/dispatcher-client-writer.ts b/src/relay/dispatcher-client-writer.ts index 7a27e670ceb..d5c4898bb89 100644 --- a/src/relay/dispatcher-client-writer.ts +++ b/src/relay/dispatcher-client-writer.ts @@ -59,6 +59,10 @@ export class DispatcherClientWriter { return this.sink.producerFrameCapacity } + get producerQueueCapacity(): number { + return this.admission.producerQueueCapacity + } + canEnqueueProducer(bytes: number): boolean { return !this.closed && this.admission.canAdmitProducer(bytes, this.producerFrameCapacity) } @@ -191,7 +195,10 @@ export class DispatcherClientWriter { if (!Number.isFinite(highWaterMark)) { return true } - return this.sink.writableLength + entry.estimatedBytes <= this.sink.producerFrameCapacity + return ( + this.sink.writableLength + entry.estimatedBytes <= this.sink.producerFrameCapacity || + this.sink.writableLength === 0 + ) } private writeEntry(entry: DispatcherWriterEntry): void { @@ -303,9 +310,7 @@ export class DispatcherClientWriter { remove?.() } - private isIdle(): boolean { - return this.inFlight.size === 0 && this.admission.queuedEntries === 0 - } + private isIdle = (): boolean => this.inFlight.size === 0 && this.admission.queuedEntries === 0 private notifyIdle(): void { if (!this.isIdle()) { diff --git a/src/relay/dispatcher-writer-admission.ts b/src/relay/dispatcher-writer-admission.ts index 32a3c9cee75..01634cbf91b 100644 --- a/src/relay/dispatcher-writer-admission.ts +++ b/src/relay/dispatcher-writer-admission.ts @@ -66,13 +66,18 @@ export class DispatcherWriterAdmission { return this.producerBytes } + get producerQueueCapacity(): number { + return this.producerQueueMaxBytes + } + get queuedEntries(): number { return Object.values(this.queues).reduce((total, queue) => total + queue.length, 0) } canAdmitProducer(bytes: number, producerFrameCapacity: number): boolean { return ( - bytes <= producerFrameCapacity && this.producerBytes + bytes <= this.producerQueueMaxBytes + bytes <= Math.max(producerFrameCapacity, this.producerQueueMaxBytes) && + this.producerBytes + bytes <= this.producerQueueMaxBytes ) } diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 4e719bd3ae5..f723899a63f 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -234,6 +234,43 @@ export class RelayDispatcher { ) } + tryNotifyPtyDataToMatchingClients( + matchesClient: (clientId: number) => boolean, + params: Record, + options: { interactive?: boolean } = {} + ): boolean { + if (this.disposed) { + return false + } + return this.tryPublishToClients( + this.activeClients().filter((client) => matchesClient(client.id)), + { jsonrpc: '2.0', method: 'pty.data', params }, + options.interactive ? 'interactive' : 'ordinary' + ) + } + + tryNotifyPtyDataToClient( + clientId: number, + params: Record, + onSettled: (result: SinkWriteSettlement) => void + ): boolean { + if (this.disposed) { + onSettled({ ok: false, error: new Error('Relay dispatcher is disposed') }) + return false + } + const client = this.clients.get(clientId) + if (!client || client.closed) { + onSettled({ ok: false, error: new Error('Relay client is not connected') }) + return false + } + return this.publishToClient( + client, + { jsonrpc: '2.0', method: 'pty.data', params }, + 'ordinary', + onSettled + ) + } + tryNotifyPtyExit(params: Record): boolean { if (this.disposed) { return false @@ -249,6 +286,42 @@ export class RelayDispatcher { ) } + tryNotifyPtyExitToMatchingClients( + matchesClient: (clientId: number) => boolean, + params: Record + ): boolean { + if (this.disposed) { + return false + } + return this.tryPublishToClients( + this.activeClients().filter((client) => matchesClient(client.id)), + { jsonrpc: '2.0', method: 'pty.exit', params }, + 'ordinary' + ) + } + + tryNotifyPtyExitToClient( + clientId: number, + params: Record, + onSettled: (result: SinkWriteSettlement) => void + ): boolean { + if (this.disposed) { + onSettled({ ok: false, error: new Error('Relay dispatcher is disposed') }) + return false + } + const client = this.clients.get(clientId) + if (!client || client.closed) { + onSettled({ ok: false, error: new Error('Relay client is not connected') }) + return false + } + return this.publishToClient( + client, + { jsonrpc: '2.0', method: 'pty.exit', params }, + 'ordinary', + onSettled + ) + } + producerDataBudget( method: string, paramsWithoutData: Record, @@ -806,7 +879,7 @@ export class RelayDispatcher { private publishBulkWhenAvailable(client: RelayClient, msg: JsonRpcNotification): Promise { const bytes = this.estimateFrameBytes(msg) - if (bytes > client.writer.producerFrameCapacity) { + if (bytes > client.writer.producerQueueCapacity) { return Promise.reject(new Error('Relay bulk frame exceeds sink producer capacity')) } return new Promise((resolve, reject) => { diff --git a/src/relay/fs-handler-file-read.ts b/src/relay/fs-handler-file-read.ts index cb8630331e4..003172c08fd 100644 --- a/src/relay/fs-handler-file-read.ts +++ b/src/relay/fs-handler-file-read.ts @@ -183,14 +183,7 @@ async function pumpChunks( break } } - const base64Budget = - dispatcher.producerDataBudget?.( - 'fs.streamChunk', - { streamId, seq }, - pumpOptions.clientId - ) ?? Number.MAX_SAFE_INTEGER - const sinkChunkBytes = Math.max(1, Math.floor(Math.max(0, base64Budget - 32) / 4) * 3) - const want = Math.min(STREAM_CHUNK_SIZE, sinkChunkBytes, totalSize - offset) + const want = Math.min(STREAM_CHUNK_SIZE, totalSize - offset) const bytesRead = await readFullStreamChunk(entry.handle, buffer, want, offset) if (bytesRead !== want) { endReason = 'error' diff --git a/src/relay/pty-handler.test.ts b/src/relay/pty-handler.test.ts index f461baf5a54..ba2a776660c 100644 --- a/src/relay/pty-handler.test.ts +++ b/src/relay/pty-handler.test.ts @@ -1515,6 +1515,78 @@ describe('PtyHandler', () => { expect(dispatcher.notify).not.toHaveBeenCalledWith('pty.data', expect.anything()) }) + it('suppresses legacy replay after the V1 owner is already active', async () => { + let dataCallback: ((data: string) => void) | undefined + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + onData: vi.fn((cb: (data: string) => void) => { + dataCallback = cb + }), + onExit: vi.fn() + }) + const spawn = await spawnPty() + dataCallback!('buffered output') + handler.setSourcePublication({ + activate: vi.fn(() => 'existing'), + accepts: vi.fn(() => true), + publish: vi.fn(() => true), + dispose: vi.fn() + } as never) + + const result = await attachPty({ + id: 'pty-1', + suppressReplayNotification: true + }) + + expect(result).toEqual({ incarnationId: spawn.incarnationId }) + }) + + it('requires restore when the V1 pending-send recovery fence expires', async () => { + const spawn = await spawnPty() + const waitForPendingSend = vi.fn().mockResolvedValue(false) + const activate = vi + .fn() + .mockReturnValue({ status: 'restoreRequired', reason: 'checkpointUnavailable' }) + handler.setSourcePublication({ + activate, + accepts: vi.fn(() => true), + waitForPendingSend, + dispose: vi.fn() + } as never) + + const result = await dispatcher.callRequest( + 'pty.attach', + { + id: spawn.id, + sourceRecovery: { + status: 'checkpoint', + deliveryToken: 'old-token', + ptyIncarnation: spawn.incarnationId, + clientGeneration: 1, + ownerGeneration: 1, + acceptedSourceEndSu: 0 + } + }, + { + clientId: 2, + isStale: () => false, + onResponseSettled: vi.fn() + } as never + ) + + expect(waitForPendingSend).toHaveBeenCalledWith(spawn.id) + expect(activate).toHaveBeenCalledWith( + spawn.id, + spawn.incarnationId, + expect.anything(), + Object.freeze({ status: 'checkpointUnavailable' }) + ) + expect(result).toEqual({ + incarnationId: spawn.incarnationId, + sourceRecovery: { status: 'restoreRequired', reason: 'checkpointUnavailable' } + }) + }) + it('notifies replay on normal attach', async () => { let dataCallback: ((data: string) => void) | undefined mockPtySpawn.mockReturnValue({ diff --git a/src/relay/pty-handler.ts b/src/relay/pty-handler.ts index 0117907eda2..32d34210032 100644 --- a/src/relay/pty-handler.ts +++ b/src/relay/pty-handler.ts @@ -53,6 +53,12 @@ import { agentSessionOwnerBindingsEqual, ClaimedAgentPtyOwnerRegistry } from '../shared/claimed-agent-pty-owner' +import type { RelayPtySourceOutput } from './relay-pty-source-output' +import type { RelayPtySourcePublication } from './relay-pty-source-publication' +import type { + PtySourceRecoveryRequest, + PtySourceRecoveryResult +} from '../shared/pty-source-recovery-contract' import { AGENT_SESSION_CREATE_OPERATION_PROTOCOL_VERSION, AGENT_SESSION_EXECUTION_OWNER_PROTOCOL_VERSION, @@ -79,6 +85,39 @@ function isMissingNodePtyNativeBinding(error: unknown): boolean { ) } +function parseSourceRecoveryRequest(value: unknown): PtySourceRecoveryRequest | undefined { + if (typeof value !== 'object' || value === null) { + return undefined + } + const input = value as Record + if (input.status === 'checkpointUnavailable') { + return Object.freeze({ status: 'checkpointUnavailable' }) + } + if ( + input.status !== 'checkpoint' || + typeof input.deliveryToken !== 'string' || + input.deliveryToken.length === 0 || + typeof input.ptyIncarnation !== 'string' || + input.ptyIncarnation.length === 0 || + !Number.isSafeInteger(input.clientGeneration) || + Number(input.clientGeneration) <= 0 || + !Number.isSafeInteger(input.ownerGeneration) || + Number(input.ownerGeneration) <= 0 || + !Number.isSafeInteger(input.acceptedSourceEndSu) || + Number(input.acceptedSourceEndSu) < 0 + ) { + return Object.freeze({ status: 'checkpointUnavailable' }) + } + return Object.freeze({ + status: 'checkpoint', + deliveryToken: input.deliveryToken, + ptyIncarnation: input.ptyIncarnation, + clientGeneration: Number(input.clientGeneration), + ownerGeneration: Number(input.ownerGeneration), + acceptedSourceEndSu: Number(input.acceptedSourceEndSu) + }) +} + type ManagedPty = { id: string incarnationId: string @@ -124,12 +163,10 @@ const AGENT_SESSION_CREATE_OPERATION_ID_PATTERN = /^[A-Za-z0-9_-]{43}$/ const AGENT_SESSION_CREATE_OPERATION_RETENTION_MS = 24 * 60 * 60 * 1000 const AGENT_SESSION_CREATE_OPERATION_LIMIT = 4_096 -type PendingPtyOutput = { +type PendingPtyOutput = RelayPtySourceOutput & { data: string - rawLength?: number - transformed?: boolean - seq?: number interactive?: boolean + sourceChunk?: RelayPtySourceOutput } type ManagedStartupCommand = { @@ -317,6 +354,7 @@ export class PtyHandler { private pausedOutputPtys = new Set() private consumerPausedOutputPtys = new Set() private removeLegacyCapacityListener: (() => void) | null = null + private sourcePublication: RelayPtySourcePublication | null = null private lastInputAtByPty = new Map() private interactiveOutputCharsByPty = new Map() private pendingSpawnCount = 0 @@ -356,6 +394,22 @@ export class PtyHandler { this.maybeResumePtyOutput(id) } + setSourcePublication(publication: RelayPtySourcePublication): void { + this.sourcePublication = publication + } + + handleSourceCreditAvailable(id: string): void { + this.sourcePublication?.onCreditAvailable(id) + } + + handleSourcePublicationCapacity(id: string): void { + if (this.pendingOutputByPty.has(id)) { + this.scheduleOutputFlush(0) + } + this.maybeResumePtyOutput(id) + this.publishPendingExit(id) + } + private async loadPty(): Promise { if (this.ptyModule) { return this.ptyModule @@ -645,7 +699,7 @@ export class PtyHandler { private registerHandlers(): void { this.dispatcher.onRequest('pty.spawn', (p, context) => this.spawn(p, context)) - this.dispatcher.onRequest('pty.attach', (p) => this.attach(p)) + this.dispatcher.onRequest('pty.attach', (p, context) => this.attach(p, context)) this.dispatcher.onRequest('pty.shutdown', (p) => this.shutdown(p)) this.dispatcher.onRequest('pty.sendSignal', (p) => this.sendSignal(p)) this.dispatcher.onRequest('pty.getCwd', (p) => this.getCwd(p)) @@ -720,6 +774,21 @@ export class PtyHandler { meta: { rawLength?: number; transformed?: boolean; seq?: number } = {} ): void { const queue = this.pendingOutputByPty.get(id) ?? [] + if (this.sourcePublication?.accepts(id)) { + queue.push({ data, ...meta }) + this.pendingOutputByPty.set(id, queue) + if (queue.length === 1 && this.shouldSendInteractiveOutputNow(id, data)) { + queue[0].interactive = true + if (this.flushPtyOutput(id)) { + return + } + } + if (this.pendingProducerBytes(id) >= PTY_OUTPUT_PRODUCER_HIGH_BYTES) { + this.pausePtyOutput(id) + } + this.scheduleOutputFlush(PTY_OUTPUT_BATCH_INTERVAL_MS) + return + } const existing = queue.at(-1) if (meta.transformed === true) { if (queue.length === 0) { @@ -839,16 +908,16 @@ export class PtyHandler { : chunk.length const chunkSeq = pending.seq === undefined ? undefined : pending.seq - (pending.data.length - chunk.length) - const published = this.publishPtyOutput( - id, - { + const sourceChunk = + pending.sourceChunk ?? + ({ data: chunk, ...(chunkSeq === undefined ? {} : { seq: chunkSeq }), ...(chunkRawLength === undefined ? {} : { rawLength: chunkRawLength }), ...(pending.transformed ? { transformed: true } : {}) - }, - pending.interactive === true - ) + } satisfies RelayPtySourceOutput) + pending.sourceChunk = sourceChunk + const published = this.publishPtyOutput(id, sourceChunk, pending.interactive === true) if (!published) { this.pendingOutputByPty.set(id, queue) this.pausePtyOutput(id) @@ -896,7 +965,14 @@ export class PtyHandler { this.interactiveOutputCharsByPty.delete(id) } - private publishPtyOutput(id: string, output: PendingPtyOutput, interactive: boolean): boolean { + private publishPtyOutput( + id: string, + output: RelayPtySourceOutput, + interactive: boolean + ): boolean { + if (this.sourcePublication?.accepts(id)) { + return this.sourcePublication.publish(id, output, interactive) + } if (this.dispatcher.tryNotifyPtyData) { return this.dispatcher.tryNotifyPtyData( { @@ -927,6 +1003,13 @@ export class PtyHandler { if (!exit) { return } + if (this.sourcePublication?.accepts(id)) { + if (!this.sourcePublication.sealAndPublishExit(exit)) { + return + } + this.pendingExitByPty.delete(id) + return + } const published = this.dispatcher.tryNotifyPtyExit ? this.dispatcher.tryNotifyPtyExit(exit) : (this.dispatcher.notify('pty.exit', exit), true) @@ -1323,6 +1406,7 @@ export class PtyHandler { } : {}) } + this.sourcePublication?.activate(id, managed.incarnationId, context) this.wireAndStore(managed) if (context?.isStale() && !params.agentSessionEnsure && !params.agentSessionCreateOperationId) { // Why: if the client reconnected while pty.spawn was in flight, the @@ -1342,8 +1426,13 @@ export class PtyHandler { } private async attach( - params: Record - ): Promise<{ incarnationId: string; replay?: string }> { + params: Record, + context?: RequestContext + ): Promise<{ + incarnationId: string + replay?: string + sourceRecovery?: PtySourceRecoveryResult + }> { const id = params.id as string const managed = this.ptys.get(id) // Why: after dispose, pty.kill is a POSIX no-op; treat disposed as not-found so failures aren't silent. @@ -1377,6 +1466,26 @@ export class PtyHandler { } managed.startupIngress?.snapshotBarrier() + let sourceRecovery = parseSourceRecoveryRequest(params.sourceRecovery) + if ( + sourceRecovery?.status === 'checkpoint' && + this.sourcePublication && + !(await this.sourcePublication.waitForPendingSend(id)) + ) { + sourceRecovery = Object.freeze({ status: 'checkpointUnavailable' }) + } + const activation = this.sourcePublication?.activate( + id, + managed.incarnationId, + context, + sourceRecovery + ) + if (typeof activation === 'object') { + return { incarnationId: managed.incarnationId, sourceRecovery: activation } + } + if (activation === 'existing' && this.sourcePublication?.accepts(id)) { + return { incarnationId: managed.incarnationId } + } // Why: renderer hasn't registered replay handlers yet during spawn, so return to the caller instead of notifying too early. // Why: buffer intentionally NOT cleared after replay (client clears xterm first) so later restarts still replay full history. @@ -1815,6 +1924,8 @@ export class PtyHandler { this.consumerPausedOutputPtys.clear() this.lastInputAtByPty.clear() this.interactiveOutputCharsByPty.clear() + this.sourcePublication?.dispose() + this.sourcePublication = null const results = await Promise.allSettled( [...this.ptys.values()].map((managed) => this.disposePtyForRelayShutdown(managed, waitForPhysicalExit) diff --git a/src/relay/pty-source-credit-ledger.test.ts b/src/relay/pty-source-credit-ledger.test.ts index 0e40349bbdf..384722210b2 100644 --- a/src/relay/pty-source-credit-ledger.test.ts +++ b/src/relay/pty-source-credit-ledger.test.ts @@ -281,6 +281,29 @@ describe('RelayPtySourceCreditLedger', () => { }) }) + it('accepts an exact receiver checkpoint for a rolled-back send attempt', () => { + const ledger = new RelayPtySourceCreditLedger() + const oldOwner = identity() + const replacement = identity('token-2', { + clientGeneration: 4, + ownerGeneration: 5 + }) + ledger.open(oldOwner, 16) + append(ledger, oldOwner, 'abcdefgh') + const attempted = ledger.reserveNextSend(oldOwner, 4)! + ledger.rollbackSend(attempted) + + const rotation = ledger.rotate(oldOwner, replacement, 4, 16) + + expect(rotation.recovery.map((span) => span.data).join('')).toBe('efgh') + expect(rotation.cancellation).toMatchObject({ sentEndSu: 4, creditedEndSu: 0 }) + expect(ledger.snapshot(replacement)).toMatchObject({ + receivedEndSu: 8, + sentEndSu: 4, + creditedEndSu: 4 + }) + }) + it('generation-closes every retained token exactly once', () => { const ledger = new RelayPtySourceCreditLedger() const owner = identity() @@ -311,6 +334,41 @@ describe('RelayPtySourceCreditLedger', () => { expect(ledger.retainedSourceSu()).toBe(4) }) + it('bounds exact UTF-8 retained bytes for transformed spans', () => { + const ledger = new RelayPtySourceCreditLedger({ + maxRetainedDataBytes: 7, + maxAggregateRetainedDataBytes: 7 + }) + const owner = identity() + ledger.open(owner, 8) + + expect(() => + ledger.append(owner, { + spanId: 'expanded', + data: '😀😀', + displayStart: 0, + displayEnd: 4, + splittable: false, + transform: { transformed: true, rawLengthSu: 1, scalarSafe: false } + }) + ).toThrow('encoded-data') + expect(ledger.retentionSnapshot()).toEqual({ sourceSu: 0, dataBytes: 0, spans: 0 }) + }) + + it('bounds retained frame metadata independently of source units', () => { + const ledger = new RelayPtySourceCreditLedger({ + maxRetainedSpans: 2, + maxAggregateRetainedSpans: 2 + }) + const owner = identity() + ledger.open(owner, 8) + append(ledger, owner, 'a', 'span-a') + append(ledger, owner, 'b', 'span-b') + + expect(() => append(ledger, owner, 'c', 'span-c')).toThrow('retained-span') + expect(ledger.retentionSnapshot()).toEqual({ sourceSu: 2, dataBytes: 2, spans: 2 }) + }) + it('bounds closed token tombstones', () => { const ledger = new RelayPtySourceCreditLedger() const owners = Array.from({ length: 300 }, (_, index) => diff --git a/src/relay/pty-source-credit-ledger.ts b/src/relay/pty-source-credit-ledger.ts index f0ea92aab4f..0360392cfb9 100644 --- a/src/relay/pty-source-credit-ledger.ts +++ b/src/relay/pty-source-credit-ledger.ts @@ -11,53 +11,46 @@ import { import { assertNonNegativeSafeInteger, assertPositiveSafeInteger, - assertPtySourceAck, assertPtySourceIdentity } from '../shared/pty-source-credit-validation' +import { + resolvePtySourceCreditLimits, + type PtySourceCreditLedgerOptions, + type PtySourceCreditLimits +} from './pty-source-credit-limits' import { CLOSED_DELIVERY_TOMBSTONE_LIMIT, + acknowledgeDeliveryRecord, closeDeliveryGeneration, createAppendedSourceSpan, createDeliveryCancellation, createDeliveryRecord, createReplacementDeliveryRecord, - DEFAULT_AGGREGATE_RETAINED_SOURCE_SU, - DEFAULT_RETAINED_SOURCE_SU, + MAX_SOURCE_SPAN_DATA_BYTES, ptyOwnerKey, - reclaimCreditedSpans, + retainedDataBytesTotal, + retainedSpanTotal, retainedSourceTotal, sliceForSend, snapshotDeliveryRecord, type DeliveryRecord, + type PtySourceAckResult, type PtySourceAppendInput, type PtySourceSendReservation } from './pty-source-credit-record' export type { PtySourceSendReservation } from './pty-source-credit-record' - -export type PtySourceAckResult = 'advanced' | 'duplicate' | 'regression' - -export type PtySourceCreditLedgerOptions = { - maxRetainedSourceSu?: number - maxAggregateRetainedSourceSu?: number -} +export type { PtySourceCreditLedgerOptions } from './pty-source-credit-limits' export class RelayPtySourceCreditLedger { private readonly deliveries = new Map() private readonly upstreamOwnerByPty = new Map() private readonly closedSnapshots = new Map() - private readonly maxRetainedSourceSu: number - private readonly maxAggregateRetainedSourceSu: number + private readonly limits: PtySourceCreditLimits private nextReservationId = 1 constructor(options: PtySourceCreditLedgerOptions = {}) { - const maxRetainedSourceSu = options.maxRetainedSourceSu ?? DEFAULT_RETAINED_SOURCE_SU - assertPositiveSafeInteger(maxRetainedSourceSu, 'maxRetainedSourceSu') - const maxAggregateRetainedSourceSu = - options.maxAggregateRetainedSourceSu ?? DEFAULT_AGGREGATE_RETAINED_SOURCE_SU - assertPositiveSafeInteger(maxAggregateRetainedSourceSu, 'maxAggregateRetainedSourceSu') - this.maxRetainedSourceSu = maxRetainedSourceSu - this.maxAggregateRetainedSourceSu = maxAggregateRetainedSourceSu + this.limits = resolvePtySourceCreditLimits(options) } open( @@ -85,14 +78,34 @@ export class RelayPtySourceCreditLedger { append(identity: PtySourceDeliveryIdentity, input: PtySourceAppendInput): PtySourceSpan { const record = this.requireActive(identity) const sourceEndSu = record.receivedEndSu + input.transform.rawLengthSu - if (sourceEndSu - record.creditedEndSu > this.maxRetainedSourceSu) { + const dataBytes = Buffer.byteLength(input.data, 'utf8') + if (dataBytes > MAX_SOURCE_SPAN_DATA_BYTES) { + throw new Error('PTY source span encoded-data budget exceeded') + } + if (sourceEndSu - record.creditedEndSu > this.limits.maxRetainedSourceSu) { throw new Error('PTY source retained-range budget exceeded') } - if (this.retainedSourceSu() + input.transform.rawLengthSu > this.maxAggregateRetainedSourceSu) { + if ( + this.retainedSourceSu() + input.transform.rawLengthSu > + this.limits.maxAggregateRetainedSourceSu + ) { throw new Error('Aggregate PTY source retained-range budget exceeded') } + if (record.retainedDataBytes + dataBytes > this.limits.maxRetainedDataBytes) { + throw new Error('PTY source retained encoded-data budget exceeded') + } + if (this.retainedDataBytes() + dataBytes > this.limits.maxAggregateRetainedDataBytes) { + throw new Error('Aggregate PTY source retained encoded-data budget exceeded') + } + if (record.spans.length + 1 > this.limits.maxRetainedSpans) { + throw new Error('PTY source retained-span budget exceeded') + } + if (this.retainedSpans() + 1 > this.limits.maxAggregateRetainedSpans) { + throw new Error('Aggregate PTY source retained-span budget exceeded') + } const span = createAppendedSourceSpan(record, input) record.spans.push(span) + record.retainedDataBytes += dataBytes record.receivedEndSu = sourceEndSu return span } @@ -133,6 +146,7 @@ export class RelayPtySourceCreditLedger { span }) record.pendingSend = reservation + record.attemptedEndSu = span.sourceEndSu return reservation } @@ -144,6 +158,7 @@ export class RelayPtySourceCreditLedger { record.pendingSend = null record.sentEndSu = reservation.span.sourceEndSu record.sentBoundaries.add(record.sentEndSu) + record.attemptedEndSu = null } rollbackSend(reservation: PtySourceSendReservation): void { @@ -154,40 +169,12 @@ export class RelayPtySourceCreditLedger { } acknowledge(identity: PtySourceDeliveryIdentity, ack: PtySourceCreditAck): PtySourceAckResult { - assertPtySourceAck(ack) const record = this.requireDelivery(identity) - if (record.state === 'closed' || record.state === 'closing') { - throw new Error('PTY source ACK targets a closed delivery') + const result = acknowledgeDeliveryRecord(record, ack) + if (result === 'advanced') { + this.maybeClose(record) } - if ( - ack.id !== record.identity.id || - ack.clientGeneration !== record.identity.clientGeneration || - ack.ownerGeneration !== record.identity.ownerGeneration || - ack.deliveryToken !== record.identity.deliveryToken - ) { - throw new Error('PTY source ACK does not own this delivery') - } - if (ack.creditedEndSu > record.sentEndSu) { - throw new Error('PTY source ACK exceeds sent source credit') - } - if (ack.creditedEndSu === record.creditedEndSu) { - return 'duplicate' - } - if (ack.creditedEndSu < record.creditedEndSu) { - return 'regression' - } - if (!record.sentBoundaries.has(ack.creditedEndSu)) { - throw new Error('PTY source ACK does not match a committed send boundary') - } - record.creditedEndSu = ack.creditedEndSu - for (const boundary of record.sentBoundaries) { - if (boundary < record.creditedEndSu) { - record.sentBoundaries.delete(boundary) - } - } - reclaimCreditedSpans(record) - this.maybeClose(record) - return 'advanced' + return result } seal(identity: PtySourceDeliveryIdentity): void { @@ -240,16 +227,16 @@ export class RelayPtySourceCreditLedger { windowSu: number ): Readonly<{ cancellation: PtySourceDeliveryCancellation; recovery: readonly PtySourceSpan[] }> { const old = this.requireDelivery(oldIdentity) + const replacementKey = ptySourceDeliveryKey(newIdentity) + if (this.deliveries.has(replacementKey) || this.closedSnapshots.has(replacementKey)) { + throw new Error('PTY source replacement token was already used') + } const replacement = createReplacementDeliveryRecord( old, newIdentity, acceptedSourceEndSu, windowSu ) - const replacementKey = ptySourceDeliveryKey(newIdentity) - if (this.deliveries.has(replacementKey) || this.closedSnapshots.has(replacementKey)) { - throw new Error('PTY source replacement token was already used') - } this.upstreamOwnerByPty.delete(ptyOwnerKey(old.identity)) this.deliveries.set(replacementKey, replacement) this.upstreamOwnerByPty.set(ptyOwnerKey(replacement.identity), replacementKey) @@ -269,8 +256,22 @@ export class RelayPtySourceCreditLedger { throw new Error('Unknown or stale PTY source delivery') } - retainedSourceSu(): number { - return retainedSourceTotal(this.deliveries.values()) + retainedSourceSu = (): number => retainedSourceTotal(this.deliveries.values()) + + retainedDataBytes = (): number => retainedDataBytesTotal(this.deliveries.values()) + + retainedSpans = (): number => retainedSpanTotal(this.deliveries.values()) + + retentionSnapshot(): Readonly<{ + sourceSu: number + dataBytes: number + spans: number + }> { + return Object.freeze({ + sourceSu: this.retainedSourceSu(), + dataBytes: this.retainedDataBytes(), + spans: this.retainedSpans() + }) } private requireActive(identity: PtySourceDeliveryIdentity): DeliveryRecord { @@ -303,6 +304,7 @@ export class RelayPtySourceCreditLedger { record.state = 'closed' record.pendingSend = null record.spans = [] + record.retainedDataBytes = 0 const key = ptySourceDeliveryKey(record.identity) const ownerKey = ptyOwnerKey(record.identity) if (this.upstreamOwnerByPty.get(ownerKey) === key) { diff --git a/src/relay/pty-source-credit-limits.ts b/src/relay/pty-source-credit-limits.ts new file mode 100644 index 00000000000..4a271c3d2fa --- /dev/null +++ b/src/relay/pty-source-credit-limits.ts @@ -0,0 +1,39 @@ +import { assertPositiveSafeInteger } from '../shared/pty-source-credit-validation' +import { + DEFAULT_AGGREGATE_RETAINED_DATA_BYTES, + DEFAULT_AGGREGATE_RETAINED_SPANS, + DEFAULT_AGGREGATE_RETAINED_SOURCE_SU, + DEFAULT_RETAINED_DATA_BYTES, + DEFAULT_RETAINED_SPANS, + DEFAULT_RETAINED_SOURCE_SU +} from './pty-source-credit-record' + +export type PtySourceCreditLedgerOptions = { + maxRetainedSourceSu?: number + maxAggregateRetainedSourceSu?: number + maxRetainedDataBytes?: number + maxAggregateRetainedDataBytes?: number + maxRetainedSpans?: number + maxAggregateRetainedSpans?: number +} + +export type PtySourceCreditLimits = Required + +export function resolvePtySourceCreditLimits( + options: PtySourceCreditLedgerOptions +): PtySourceCreditLimits { + const limits = { + maxRetainedSourceSu: options.maxRetainedSourceSu ?? DEFAULT_RETAINED_SOURCE_SU, + maxAggregateRetainedSourceSu: + options.maxAggregateRetainedSourceSu ?? DEFAULT_AGGREGATE_RETAINED_SOURCE_SU, + maxRetainedDataBytes: options.maxRetainedDataBytes ?? DEFAULT_RETAINED_DATA_BYTES, + maxAggregateRetainedDataBytes: + options.maxAggregateRetainedDataBytes ?? DEFAULT_AGGREGATE_RETAINED_DATA_BYTES, + maxRetainedSpans: options.maxRetainedSpans ?? DEFAULT_RETAINED_SPANS, + maxAggregateRetainedSpans: options.maxAggregateRetainedSpans ?? DEFAULT_AGGREGATE_RETAINED_SPANS + } + for (const [name, value] of Object.entries(limits)) { + assertPositiveSafeInteger(value, name) + } + return Object.freeze(limits) +} diff --git a/src/relay/pty-source-credit-record.ts b/src/relay/pty-source-credit-record.ts index 5f0bb971727..985369bbc4c 100644 --- a/src/relay/pty-source-credit-record.ts +++ b/src/relay/pty-source-credit-record.ts @@ -1,4 +1,5 @@ import type { + PtySourceCreditAck, PtySourceDeliveryIdentity, PtySourceDeliveryCancellation, PtySourceDeliverySnapshot, @@ -9,12 +10,18 @@ import { ptySourceSpanIsSplittable } from '../shared/pty-source-credit-contract' import { assertNonNegativeSafeInteger, assertPositiveSafeInteger, + assertPtySourceAck, assertPtySourceIdentity, assertPtySourceSpan } from '../shared/pty-source-credit-validation' export const DEFAULT_RETAINED_SOURCE_SU = 512 * 1024 export const DEFAULT_AGGREGATE_RETAINED_SOURCE_SU = 48 * 1024 * 1024 +export const DEFAULT_RETAINED_DATA_BYTES = 2 * 1024 * 1024 +export const DEFAULT_AGGREGATE_RETAINED_DATA_BYTES = 64 * 1024 * 1024 +export const DEFAULT_RETAINED_SPANS = 1_024 +export const DEFAULT_AGGREGATE_RETAINED_SPANS = 64 * 1_024 +export const MAX_SOURCE_SPAN_DATA_BYTES = 1024 * 1024 export const CLOSED_DELIVERY_TOMBSTONE_LIMIT = 256 export type PtySourceSendReservation = Readonly<{ @@ -23,6 +30,8 @@ export type PtySourceSendReservation = Readonly<{ span: PtySourceSpan }> +export type PtySourceAckResult = 'advanced' | 'duplicate' | 'regression' + export type DeliveryRecord = { identity: PtySourceDeliveryIdentity state: 'active' | 'sealed-unsettled' | 'closing' | 'closed' @@ -30,9 +39,11 @@ export type DeliveryRecord = { receivedEndSu: number sentEndSu: number creditedEndSu: number + retainedDataBytes: number spans: PtySourceSpan[] sentBoundaries: Set pendingSend: PtySourceSendReservation | null + attemptedEndSu: number | null exitPublished: boolean generationClosed: boolean } @@ -53,9 +64,11 @@ export function createDeliveryRecord( receivedEndSu: checkpointSourceEndSu, sentEndSu: checkpointSourceEndSu, creditedEndSu: checkpointSourceEndSu, + retainedDataBytes: 0, spans: [], sentBoundaries: new Set([checkpointSourceEndSu]), pendingSend: null, + attemptedEndSu: null, exitPublished: false, generationClosed: false } @@ -159,10 +172,48 @@ export function snapshotDeliveryRecord(record: DeliveryRecord): PtySourceDeliver export function reclaimCreditedSpans(record: DeliveryRecord): void { while (record.spans[0]?.sourceEndSu <= record.creditedEndSu) { - record.spans.shift() + record.retainedDataBytes -= Buffer.byteLength(record.spans.shift()!.data, 'utf8') } } +export function acknowledgeDeliveryRecord( + record: DeliveryRecord, + ack: PtySourceCreditAck +): PtySourceAckResult { + assertPtySourceAck(ack) + if (record.state === 'closed' || record.state === 'closing') { + throw new Error('PTY source ACK targets a closed delivery') + } + if ( + ack.id !== record.identity.id || + ack.clientGeneration !== record.identity.clientGeneration || + ack.ownerGeneration !== record.identity.ownerGeneration || + ack.deliveryToken !== record.identity.deliveryToken + ) { + throw new Error('PTY source ACK does not own this delivery') + } + if (ack.creditedEndSu > record.sentEndSu) { + throw new Error('PTY source ACK exceeds sent source credit') + } + if (ack.creditedEndSu === record.creditedEndSu) { + return 'duplicate' + } + if (ack.creditedEndSu < record.creditedEndSu) { + return 'regression' + } + if (!record.sentBoundaries.has(ack.creditedEndSu)) { + throw new Error('PTY source ACK does not match a committed send boundary') + } + record.creditedEndSu = ack.creditedEndSu + for (const boundary of record.sentBoundaries) { + if (boundary < record.creditedEndSu) { + record.sentBoundaries.delete(boundary) + } + } + reclaimCreditedSpans(record) + return 'advanced' +} + export function retainedSourceTotal(records: Iterable): number { let total = 0 for (const record of records) { @@ -171,6 +222,22 @@ export function retainedSourceTotal(records: Iterable): number { return total } +export function retainedDataBytesTotal(records: Iterable): number { + let total = 0 + for (const record of records) { + total += record.retainedDataBytes + } + return total +} + +export function retainedSpanTotal(records: Iterable): number { + let total = 0 + for (const record of records) { + total += record.spans.length + } + return total +} + export function createReplacementDeliveryRecord( old: DeliveryRecord, newIdentity: PtySourceDeliveryIdentity, @@ -180,10 +247,12 @@ export function createReplacementDeliveryRecord( assertPtySourceIdentity(newIdentity) assertPositiveSafeInteger(windowSu, 'windowSu') assertNonNegativeSafeInteger(acceptedSourceEndSu, 'acceptedSourceEndSu') + const committedCheckpoint = + acceptedSourceEndSu <= old.sentEndSu && old.sentBoundaries.has(acceptedSourceEndSu) + const attemptedCheckpoint = acceptedSourceEndSu === old.attemptedEndSu if ( acceptedSourceEndSu < old.creditedEndSu || - acceptedSourceEndSu > old.sentEndSu || - !old.sentBoundaries.has(acceptedSourceEndSu) || + (!committedCheckpoint && !attemptedCheckpoint) || old.pendingSend || newIdentity.id !== old.identity.id || newIdentity.ptyIncarnation !== old.identity.ptyIncarnation || @@ -191,12 +260,21 @@ export function createReplacementDeliveryRecord( ) { throw new Error('PTY source recovery checkpoint does not exactly cover the retained delivery') } + if (acceptedSourceEndSu > old.sentEndSu) { + old.sentEndSu = acceptedSourceEndSu + old.sentBoundaries.add(acceptedSourceEndSu) + } + old.attemptedEndSu = null const replacement = createDeliveryRecord(newIdentity, windowSu, acceptedSourceEndSu) replacement.receivedEndSu = old.receivedEndSu replacement.spans = old.spans .filter((span) => span.sourceEndSu > acceptedSourceEndSu) .map((span) => sliceAtSourceStart(span, Math.max(span.sourceStartSu, acceptedSourceEndSu))) .map((span) => Object.freeze({ ...span, ...replacement.identity })) + replacement.retainedDataBytes = replacement.spans.reduce( + (bytes, span) => bytes + Buffer.byteLength(span.data, 'utf8'), + 0 + ) return replacement } diff --git a/src/relay/pty-source-delivery-ownership.ts b/src/relay/pty-source-delivery-ownership.ts new file mode 100644 index 00000000000..b7f1306e72a --- /dev/null +++ b/src/relay/pty-source-delivery-ownership.ts @@ -0,0 +1,18 @@ +import type { PtySourceDeliveryIdentity } from '../shared/pty-source-credit-contract' +import type { PtyConsumerSessionGrant } from '../shared/pty-consumer-session' + +export function ownedPtySourceDelivery( + identity: PtySourceDeliveryIdentity | undefined, + grant: Readonly, + id: string +): PtySourceDeliveryIdentity | null { + if ( + !identity || + identity.id !== id || + identity.clientGeneration !== grant.clientGeneration || + identity.ownerGeneration !== grant.ownerGeneration + ) { + return null + } + return identity +} diff --git a/src/relay/relay-frame-decoder.ts b/src/relay/relay-frame-decoder.ts index 2751d984bfd..f6673d8b6ac 100644 --- a/src/relay/relay-frame-decoder.ts +++ b/src/relay/relay-frame-decoder.ts @@ -1,3 +1,6 @@ +import type { DecodedFrame, FrameDecoderOptions } from '../shared/relay-frame-decoder-contract' +export type { DecodedFrame, FrameDecoderOptions } from '../shared/relay-frame-decoder-contract' + export const HEADER_LENGTH = 13 export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 export const FRAME_DECODER_MAX_FRAMES_PER_TURN = 64 @@ -5,24 +8,6 @@ export const FRAME_DECODER_MAX_BYTES_PER_TURN = MAX_MESSAGE_SIZE + HEADER_LENGTH export const FRAME_DECODER_MAX_TURN_MS = 4, FRAME_DECODER_MAX_RETAINED_BYTES = 2 * (MAX_MESSAGE_SIZE + HEADER_LENGTH) -export type DecodedFrame = { - type: number - id: number - ack: number - payload: Buffer -} - -export type FrameDecoderOptions = { - maxFramesPerTurn?: number - maxBytesPerTurn?: number - maxTurnMs?: number - now?: () => number - schedule?: (callback: () => void) => unknown - cancelScheduled?: (handle: unknown) => void - pause?: () => void - resume?: () => void -} - export class FrameDecoder { private chunks: Buffer[] = [] private bufferedLength = 0 diff --git a/src/relay/relay-pty-source-output.ts b/src/relay/relay-pty-source-output.ts new file mode 100644 index 00000000000..6beefefa6cd --- /dev/null +++ b/src/relay/relay-pty-source-output.ts @@ -0,0 +1,8 @@ +export type RelayPtySourceOutput = { + data: string + rawLength?: number + transformed?: boolean + seq?: number + sourceAccepted?: boolean + sourceSpanId?: string +} diff --git a/src/relay/relay-pty-source-publication.test.ts b/src/relay/relay-pty-source-publication.test.ts new file mode 100644 index 00000000000..65248d53111 --- /dev/null +++ b/src/relay/relay-pty-source-publication.test.ts @@ -0,0 +1,580 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { + RelayDispatcher, + type RelayClientSessionIdentity, + type SinkWriteSettlement +} from './dispatcher' +import { encodeJsonRpcFrame, MessageType } from './protocol' +import { RelayPtySourcePublication } from './relay-pty-source-publication' +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 notification(buffer: Buffer): { method: string; params: Record } | null { + if (buffer[0] !== MessageType.Regular) { + return null + } + const length = buffer.readUInt32BE(9) + const message = JSON.parse(buffer.subarray(13, 13 + length).toString('utf8')) + return typeof message.method === 'string' && message.id === undefined ? message : null +} + +function responseResult(buffer: Buffer): Record | null { + if (buffer[0] !== MessageType.Regular) { + return null + } + const length = buffer.readUInt32BE(9) + const message = JSON.parse(buffer.subarray(13, 13 + length).toString('utf8')) + return message.id === undefined ? null : (message.result ?? null) +} + +async function flushRequests(): Promise { + await new Promise((resolve) => setImmediate(resolve)) +} + +describe('RelayPtySourcePublication', () => { + let dispatcher: RelayDispatcher | null = null + + afterEach(() => { + dispatcher?.dispose() + dispatcher = null + }) + + async function createHarness(windowSu = 4, settleSourceImmediately = true) { + const writes: Buffer[] = [] + const sourceSettlements: ((result: SinkWriteSettlement) => void)[] = [] + dispatcher = new RelayDispatcher( + (data, onSettled) => { + writes.push(Buffer.from(data)) + const frame = notification(data) + if (frame?.method === 'pty.data' || frame?.method === 'pty.exit') { + sourceSettlements.push(onSettled) + if (settleSourceImmediately) { + onSettled({ ok: true }) + } + } else { + onSettled({ ok: true }) + } + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + let publication: RelayPtySourcePublication + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a', undefined, (id) => + publication.onCreditAvailable(id) + ) + publication = new RelayPtySourcePublication(dispatcher, adapter, () => {}) + dispatcher.feed( + requestFrame(1, 'pty.openClient', { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: windowSu } } + }) + ) + await flushRequests() + const activationSettlements: ((result: SinkWriteSettlement) => void)[] = [] + expect( + publication.activate('pty-1', 'incarnation-1', { + clientId: 1, + isStale: () => false, + sessionIdentity: endpointIdentity, + onResponseSettled: (callback) => activationSettlements.push(callback) + }) + ).toBe('opened') + activationSettlements[0]({ ok: true }) + return { adapter, publication, sourceSettlements, writes } + } + + it('commits only from writer settlement and resumes exactly after cumulative ACK', async () => { + const harness = await createHarness(4, false) + expect(harness.publication.publish('pty-1', { data: 'abcdefgh' }, false)).toBe(true) + expect(harness.publication.getDebugSnapshot()).toMatchObject({ + outstandingSourceUnits: 0, + sendCommitted: 0 + }) + + harness.sourceSettlements[0]({ ok: true }) + expect(harness.publication.getDebugSnapshot()).toMatchObject({ + outstandingSourceUnits: 4, + sendCommitted: 1 + }) + const first = harness.writes.map(notification).find((frame) => frame?.method === 'pty.data')! + dispatcher!.feed( + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.ackData', + params: { + acknowledgements: [ + { + id: 'pty-1', + clientGeneration: first.params.clientGeneration, + ownerGeneration: first.params.ownerGeneration, + deliveryToken: first.params.deliveryToken, + creditedEndSu: 4 + } + ] + } + }, + 2, + 0 + ) + ) + expect( + harness.writes.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(2) + await flushRequests() + }) + + it('keeps mixed legacy and V1 clients on distinct frame authority', async () => { + const harness = await createHarness(8) + const legacyWrites: Buffer[] = [] + dispatcher!.attachClient( + (data, onSettled) => { + legacyWrites.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true } + ) + + harness.publication.publish('pty-1', { data: 'data' }, false) + + const sourceFrame = harness.writes + .map(notification) + .find((frame) => frame?.method === 'pty.data')! + const oldGrant = harness.writes.map(responseResult).find((result) => result?.ownerLease)! + const legacyFrame = legacyWrites + .map(notification) + .find((frame) => frame?.method === 'pty.data')! + expect(sourceFrame.params).toMatchObject({ + data: 'data', + sourceEndSu: 4, + sourceLengthSu: 4 + }) + expect(legacyFrame.params).toEqual({ id: 'pty-1', data: 'data' }) + + dispatcher!.invalidateClient() + const replacementWrites: Buffer[] = [] + const replacementClientId = dispatcher!.attachClient( + (data, onSettled) => { + replacementWrites.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher!.feedClient( + replacementClientId, + requestFrame(2, 'pty.openClient', { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + resume: { + ownerGeneration: oldGrant.ownerGeneration, + ownerLease: oldGrant.ownerLease + }, + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 8 } } + }) + ) + await flushRequests() + harness.publication.publish('pty-1', { data: 'next' }, false) + + expect( + replacementWrites.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(0) + expect( + legacyWrites.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(2) + }) + + it('rolls back a failed source write without advancing the sent window', async () => { + const harness = await createHarness(4, false) + harness.publication.publish('pty-1', { data: 'data' }, false) + harness.sourceSettlements[0]({ ok: false, error: new Error('socket write failed') }) + + expect(harness.publication.getDebugSnapshot()).toMatchObject({ + outstandingSourceUnits: 0, + sendCommitted: 0, + sendRolledBack: 1 + }) + }) + + it('fences idle publication and pumping before the wait continuation', async () => { + const harness = await createHarness(4) + harness.publication.publish('pty-1', { data: 'abcdefgh' }, false) + const firstData = harness.writes + .map(notification) + .find((frame) => frame?.method === 'pty.data')! + const fence = harness.publication.waitForPendingSend('pty-1') + + expect(harness.publication.publish('pty-1', { data: 'ijkl' }, false)).toBe(false) + dispatcher!.feed( + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.ackData', + params: { + acknowledgements: [ + { + id: 'pty-1', + clientGeneration: firstData.params.clientGeneration, + ownerGeneration: firstData.params.ownerGeneration, + deliveryToken: firstData.params.deliveryToken, + creditedEndSu: 4 + } + ] + } + }, + 2, + 0 + ) + ) + expect( + harness.writes.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(1) + expect(await fence).toBe(true) + + expect( + harness.publication.activate('pty-1', 'incarnation-1', { + clientId: 1, + isStale: () => false, + sessionIdentity: endpointIdentity, + onResponseSettled: () => {} + }) + ).toBe('existing') + expect( + harness.writes.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(2) + await flushRequests() + }) + + it('does not pump an old-token suffix after a successful fenced settlement', async () => { + const harness = await createHarness(4, false) + harness.publication.publish('pty-1', { data: 'abcdefgh' }, false) + const oldData = harness.writes.map(notification).find((frame) => frame?.method === 'pty.data')! + const oldGrant = harness.writes.map(responseResult).find((result) => result?.ownerLease)! + const fence = harness.publication.waitForPendingSend('pty-1') + + harness.sourceSettlements[0]({ ok: true }) + expect(await fence).toBe(true) + expect( + harness.writes.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(1) + dispatcher!.invalidateClient() + + const recoveredWrites: Buffer[] = [] + const recoveredClientId = dispatcher!.attachClient( + (data, onSettled) => { + recoveredWrites.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher!.feedClient( + recoveredClientId, + requestFrame(2, 'pty.openClient', { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + resume: { + ownerGeneration: oldGrant.ownerGeneration, + ownerLease: oldGrant.ownerLease + }, + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 4 } } + }) + ) + await flushRequests() + const activationSettlements: ((result: SinkWriteSettlement) => void)[] = [] + expect( + harness.publication.activate( + 'pty-1', + 'incarnation-1', + { + clientId: recoveredClientId, + isStale: () => false, + sessionIdentity: endpointIdentity, + onResponseSettled: (callback) => activationSettlements.push(callback) + }, + { + status: 'checkpoint', + deliveryToken: String(oldData.params.deliveryToken), + clientGeneration: Number(oldData.params.clientGeneration), + ownerGeneration: Number(oldData.params.ownerGeneration), + ptyIncarnation: 'incarnation-1', + acceptedSourceEndSu: 4 + } + ) + ).toMatchObject({ status: 'pending', checkpointSourceEndSu: 4, recoveryEndSu: 8 }) + activationSettlements[0]({ ok: true }) + + const replacementData = recoveredWrites + .map(notification) + .find((frame) => frame?.method === 'pty.data')! + expect(replacementData.params).toMatchObject({ data: 'efgh', sourceEndSu: 8 }) + expect(replacementData.params.deliveryToken).not.toBe(oldData.params.deliveryToken) + expect( + harness.writes.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(1) + }) + + it('recovers a receiver-accepted frame after detach rolls back its send', async () => { + const harness = await createHarness(4, false) + harness.publication.publish('pty-1', { data: 'data' }, false) + const oldData = harness.writes.map(notification).find((frame) => frame?.method === 'pty.data')! + const oldGrant = harness.writes.map(responseResult).find((result) => result?.ownerLease)! + let fenceSettled = false + const fence = harness.publication.waitForPendingSend('pty-1').then((result) => { + fenceSettled = true + return result + }) + await Promise.resolve() + expect(fenceSettled).toBe(false) + + dispatcher!.invalidateClient() + expect(await fence).toBe(true) + expect(harness.publication.getDebugSnapshot()).toMatchObject({ + sendCommitted: 0, + sendRolledBack: 1 + }) + + const recoveredWrites: Buffer[] = [] + const recoveredClientId = dispatcher!.attachClient( + (data, onSettled) => { + recoveredWrites.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher!.feedClient( + recoveredClientId, + requestFrame(2, 'pty.openClient', { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + resume: { + ownerGeneration: oldGrant.ownerGeneration, + ownerLease: oldGrant.ownerLease + }, + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 4 } } + }) + ) + await flushRequests() + const activationSettlements: ((result: SinkWriteSettlement) => void)[] = [] + const activation = harness.publication.activate( + 'pty-1', + 'incarnation-1', + { + clientId: recoveredClientId, + isStale: () => false, + sessionIdentity: endpointIdentity, + onResponseSettled: (callback) => activationSettlements.push(callback) + }, + { + status: 'checkpoint', + deliveryToken: String(oldData.params.deliveryToken), + clientGeneration: Number(oldData.params.clientGeneration), + ownerGeneration: Number(oldData.params.ownerGeneration), + ptyIncarnation: 'incarnation-1', + acceptedSourceEndSu: 4 + } + ) + expect(activation).toMatchObject({ + status: 'pending', + checkpointSourceEndSu: 4, + recoveryEndSu: 4 + }) + activationSettlements[0]({ ok: true }) + expect( + recoveredWrites.map(notification).filter((frame) => frame?.method === 'pty.data') + ).toHaveLength(0) + expect( + recoveredWrites.map(notification).find((frame) => frame?.method === 'pty.recoveryComplete') + ?.params + ).toMatchObject({ checkpointSourceEndSu: 4, recoveryEndSu: 4 }) + const beforeLateSettlement = harness.publication.getDebugSnapshot() + + harness.sourceSettlements[0]({ ok: true }) + + expect(harness.publication.getDebugSnapshot()).toEqual(beforeLateSettlement) + }) + + it('publishes exit only after preceding source data settles', async () => { + const harness = await createHarness(4, false) + harness.publication.publish('pty-1', { data: 'data' }, false) + expect( + harness.publication.sealAndPublishExit({ + id: 'pty-1', + code: 0, + incarnationId: 'incarnation-1' + }) + ).toBe(false) + + harness.sourceSettlements[0]({ ok: true }) + expect( + harness.publication.sealAndPublishExit({ + id: 'pty-1', + code: 0, + incarnationId: 'incarnation-1' + }) + ).toBe(true) + expect( + harness.writes + .map(notification) + .filter((frame): frame is NonNullable => frame !== null) + .map((frame) => frame.method) + ).toEqual(['pty.data', 'pty.exit']) + }) + + it('keeps gate-off same-build sessions on legacy capability omission', async () => { + const writes: Buffer[] = [] + dispatcher = new RelayDispatcher( + (data, onSettled) => { + writes.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + new SshPtyConsumerSessionAdapter(dispatcher, 'build-a', undefined, undefined, false) + dispatcher.feed( + requestFrame(1, 'pty.openClient', { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 4 } } + }) + ) + await flushRequests() + + expect( + writes + .map((buffer) => { + const length = buffer.readUInt32BE(9) + return JSON.parse(buffer.subarray(13, 13 + length).toString('utf8')).result + }) + .find(Boolean) + ).not.toHaveProperty('capabilities') + }) + + it('settles retained recovery before the fence and publishes later live output after it', async () => { + const harness = await createHarness(4) + harness.publication.publish('pty-1', { data: 'abcdefgh' }, false) + const firstData = harness.writes + .map(notification) + .find((frame) => frame?.method === 'pty.data')! + const firstGrant = harness.writes.map(responseResult).find((result) => result?.ownerLease)! + dispatcher!.invalidateClient() + + const recoveredWrites: Buffer[] = [] + const recoverySettlements: ((result: SinkWriteSettlement) => void)[] = [] + const recoveredClientId = dispatcher!.attachClient( + (data, onSettled) => { + recoveredWrites.push(Buffer.from(data)) + if (notification(data)?.method === 'pty.data') { + recoverySettlements.push(onSettled) + } else { + onSettled({ ok: true }) + } + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher!.feedClient( + recoveredClientId, + requestFrame(2, 'pty.openClient', { + protocolVersion: 1, + clientInstanceId: 'client-1', + requestedRole: 'session-owner', + resume: { + ownerGeneration: firstGrant.ownerGeneration, + ownerLease: firstGrant.ownerLease + }, + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 4 } } + }) + ) + await flushRequests() + const activationSettlements: ((result: SinkWriteSettlement) => void)[] = [] + const activation = harness.publication.activate( + 'pty-1', + 'incarnation-1', + { + clientId: recoveredClientId, + isStale: () => false, + sessionIdentity: endpointIdentity, + onResponseSettled: (callback) => activationSettlements.push(callback) + }, + { + status: 'checkpoint', + deliveryToken: String(firstData.params.deliveryToken), + clientGeneration: Number(firstData.params.clientGeneration), + ownerGeneration: Number(firstData.params.ownerGeneration), + ptyIncarnation: 'incarnation-1', + acceptedSourceEndSu: 4 + } + ) + expect(activation).toMatchObject({ + status: 'pending', + checkpointSourceEndSu: 4, + recoveryEndSu: 8 + }) + activationSettlements[0]({ ok: true }) + harness.publication.publish('pty-1', { data: 'ijkl' }, false) + + expect( + recoveredWrites.map(notification).filter((frame) => frame?.method === 'pty.recoveryComplete') + ).toHaveLength(0) + recoverySettlements[0]({ ok: true }) + const recoveredData = recoveredWrites + .map(notification) + .find((frame) => frame?.method === 'pty.data')! + dispatcher!.feedClient( + recoveredClientId, + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.ackData', + params: { + acknowledgements: [ + { + id: 'pty-1', + clientGeneration: recoveredData.params.clientGeneration, + ownerGeneration: recoveredData.params.ownerGeneration, + deliveryToken: recoveredData.params.deliveryToken, + creditedEndSu: 8 + } + ] + } + }, + 3, + 0 + ) + ) + + expect( + recoveredWrites + .map(notification) + .filter((frame): frame is NonNullable => frame !== null) + .map((frame) => frame.method) + ).toEqual(['pty.deliveryCanceled', 'pty.data', 'pty.recoveryComplete', 'pty.data']) + await flushRequests() + }) +}) diff --git a/src/relay/relay-pty-source-publication.ts b/src/relay/relay-pty-source-publication.ts new file mode 100644 index 00000000000..670b72a32e6 --- /dev/null +++ b/src/relay/relay-pty-source-publication.ts @@ -0,0 +1,308 @@ +import { randomUUID } from 'node:crypto' +import type { PtySourceDeliveryIdentity } from '../shared/pty-source-credit-contract' +import type { + PtySourceRecoveryRequest, + PtySourceRecoveryResult +} from '../shared/pty-source-recovery-contract' +import { + onceSinkSettlement, + RelayPtySourceSendScheduler, + type RelayPtySourceDeliveryRecord, + type RelayPtySourcePublicationCounters +} from './relay-pty-source-send-scheduler' +import type { RelayDispatcher, RequestContext } from './dispatcher' +import type { RelayPtySourceOutput } from './relay-pty-source-output' +import type { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' + +export class RelayPtySourcePublication { + private readonly deliveries = new Map() + private readonly counters: RelayPtySourcePublicationCounters = { + opened: 0, + rotated: 0, + appendDenied: 0, + sendCommitted: 0, + sendRolledBack: 0, + exitCommitted: 0, + exitRolledBack: 0 + } + + constructor( + private readonly dispatcher: RelayDispatcher, + private readonly session: SshPtyConsumerSessionAdapter, + private readonly onCapacity: (id: string) => void + ) { + this.sender = new RelayPtySourceSendScheduler( + dispatcher, + session, + this.deliveries, + this.counters, + onCapacity + ) + } + + private readonly sender: RelayPtySourceSendScheduler + + activate( + id: string, + ptyIncarnation: string, + context: RequestContext | undefined, + recovery?: PtySourceRecoveryRequest + ): false | 'opened' | 'rotated' | 'existing' | PtySourceRecoveryResult { + if (!context?.onResponseSettled) { + this.sender.releaseRotationFence(this.deliveries.get(id)) + return false + } + const mode = this.session.deliveryMode(context.clientId) + const current = this.deliveries.get(id) + if (mode === 'subscriber') { + this.sender.releaseRotationFence(current) + return false + } + if (mode === 'legacy-owner') { + if (current) { + this.session.cancelDelivery(current.identity, 'source-credit-disabled') + this.deliveries.delete(id) + this.onCapacity(id) + } + return false + } + if (current?.clientId === context.clientId) { + this.sender.releaseRotationFence(current) + return 'existing' + } + let identity: PtySourceDeliveryIdentity | null = null + let displayEnd = 0 + let recoveryCheckpointSourceEndSu: number | null = null + let recoveryEndSu: number | null = null + if (!current && recovery) { + return this.publishRestoreRequired(id, context, 'deliveryUnavailable') + } + if (current) { + try { + const snapshot = this.session.sourceDeliverySnapshot(current.identity) + if ( + snapshot.state === 'closed' || + snapshot.state === 'closing' || + recovery?.status !== 'checkpoint' || + recovery.deliveryToken !== current.identity.deliveryToken || + recovery.clientGeneration !== current.identity.clientGeneration || + recovery.ownerGeneration !== current.identity.ownerGeneration || + recovery.ptyIncarnation !== current.identity.ptyIncarnation + ) { + return this.requireRestore(id, current, context, 'checkpointUnavailable') + } + const rotation = this.session.rotateDelivery( + current.identity, + context.clientId, + recovery.acceptedSourceEndSu + ) + identity = rotation.identity + displayEnd = current.displayEnd + recoveryCheckpointSourceEndSu = recovery.acceptedSourceEndSu + recoveryEndSu = snapshot.receivedEndSu + this.counters.rotated++ + } catch (error) { + return this.requireRestore( + id, + current, + context, + error instanceof Error ? error.message : 'invalidCheckpoint' + ) + } + } + identity ??= this.session.openDelivery(context.clientId, id, ptyIncarnation) + if (!identity) { + return false + } + if (!current || identity !== current.identity) { + this.counters.opened++ + } + const record: RelayPtySourceDeliveryRecord = { + clientId: context.clientId, + identity, + displayEnd, + activating: true, + sealed: false, + legacyExitAccepted: false, + sourceExitAccepted: false, + sending: false, + turnFrames: 0, + turnSourceSu: 0, + turnScheduled: false, + sendWaiters: new Set(), + recoveryCheckpointSourceEndSu, + recoveryEndSu, + restoreRequired: false, + rotationPending: false + } + this.deliveries.set(id, record) + context.onResponseSettled((result) => { + if (this.deliveries.get(id) !== record) { + return + } + if (!result.ok) { + this.session.cancelDelivery(record.identity, 'activation-publication-failed') + this.deliveries.delete(id) + return + } + record.activating = false + this.sender.completeRecoveryIfReady(record) + this.sender.pump(record) + this.onCapacity(id) + }) + if (recoveryEndSu !== null && recoveryCheckpointSourceEndSu !== null) { + return Object.freeze({ + status: 'pending', + clientGeneration: identity.clientGeneration, + ownerGeneration: identity.ownerGeneration, + ptyIncarnation: identity.ptyIncarnation, + deliveryToken: identity.deliveryToken, + checkpointSourceEndSu: recoveryCheckpointSourceEndSu, + recoveryEndSu + }) + } + return current ? 'rotated' : 'opened' + } + + accepts = (id: string): boolean => this.deliveries.has(id) + + waitForPendingSend = (id: string, timeoutMs = 5_000): Promise => + this.sender.waitForPendingSend(id, timeoutMs) + + publish(id: string, output: RelayPtySourceOutput, interactive: boolean): boolean { + const record = this.deliveries.get(id) + if (!record || record.sealed || record.restoreRequired || record.rotationPending) { + return false + } + if (!output.sourceAccepted) { + const rawLength = output.rawLength ?? output.data.length + output.sourceSpanId ??= randomUUID() + try { + this.session.appendSource(record.identity, { + spanId: output.sourceSpanId, + data: output.data, + displayStart: record.displayEnd, + displayEnd: record.displayEnd + output.data.length, + splittable: output.transformed !== true, + transform: { + transformed: output.transformed === true, + rawLengthSu: rawLength, + scalarSafe: output.transformed !== true + } + }) + } catch { + this.counters.appendDenied++ + return false + } + output.sourceAccepted = true + record.displayEnd += output.data.length + } + if ( + !this.dispatcher.tryNotifyPtyDataToMatchingClients( + (clientId) => this.session.deliveryMode(clientId) !== 'source-owner', + { + id, + data: output.data, + ...(output.seq === undefined ? {} : { seq: output.seq }), + ...(output.rawLength === undefined ? {} : { rawLength: output.rawLength }), + ...(output.transformed ? { transformed: true } : {}) + }, + { interactive } + ) + ) { + return false + } + this.sender.pump(record) + return true + } + + sealAndPublishExit(params: { id: string; code: number; incarnationId: string }): boolean { + const record = this.deliveries.get(params.id) + if (!record) { + return false + } + if (record.restoreRequired) { + const published = this.dispatcher.tryNotifyPtyExit(params) + if (published && this.deliveries.get(params.id) === record) { + this.deliveries.delete(params.id) + } + return published + } + if (!record.sealed) { + this.session.sealDelivery(record.identity) + record.sealed = true + } + this.sender.pump(record) + const snapshot = this.session.sourceDeliverySnapshot(record.identity) + if (snapshot.sentEndSu !== snapshot.receivedEndSu) { + return false + } + if (!record.legacyExitAccepted) { + record.legacyExitAccepted = this.dispatcher.tryNotifyPtyExitToMatchingClients( + (clientId) => this.session.deliveryMode(clientId) !== 'source-owner', + params + ) + if (!record.legacyExitAccepted) { + return false + } + } + if (record.sourceExitAccepted) { + return true + } + const settle = onceSinkSettlement((result) => { + try { + this.session.settleExitPublication(record.identity, result) + if (result.ok) { + this.counters.exitCommitted++ + } else { + this.counters.exitRolledBack++ + } + this.sender.pruneClosed(params.id, record) + } finally { + this.onCapacity(params.id) + } + }) + record.sourceExitAccepted = this.dispatcher.tryNotifyPtyExitToClient( + record.clientId, + params, + settle + ) + return record.sourceExitAccepted + } + + onCreditAvailable = (id: string): void => this.sender.onCreditAvailable(id) + + getDebugSnapshot = () => this.sender.getDebugSnapshot() + + dispose = (): void => this.sender.dispose() + + private requireRestore( + id: string, + current: RelayPtySourceDeliveryRecord, + context: RequestContext, + reason: string + ): Readonly<{ status: 'restoreRequired'; reason: string }> { + this.session.cancelDelivery(current.identity, `recovery-${reason}`) + current.restoreRequired = true + current.activating = false + this.sender.wakeSendWaiters(current) + current.recoveryCheckpointSourceEndSu = null + current.recoveryEndSu = null + return this.publishRestoreRequired(id, context, reason) + } + + private publishRestoreRequired( + id: string, + context: RequestContext, + reason: string + ): Readonly<{ status: 'restoreRequired'; reason: string }> { + const result = Object.freeze({ status: 'restoreRequired' as const, reason }) + context.onResponseSettled?.((settlement) => { + if (settlement.ok) { + this.dispatcher.notifyClient(context.clientId, 'pty.restoreRequired', { id, reason }) + } + }) + this.onCapacity(id) + return result + } +} diff --git a/src/relay/relay-pty-source-send-scheduler.ts b/src/relay/relay-pty-source-send-scheduler.ts new file mode 100644 index 00000000000..e8c0676a062 --- /dev/null +++ b/src/relay/relay-pty-source-send-scheduler.ts @@ -0,0 +1,259 @@ +import type { PtySourceDeliveryIdentity } from '../shared/pty-source-credit-contract' +import type { RelayDispatcher, SinkWriteSettlement } from './dispatcher' +import { + PTY_SOURCE_SCHEDULER_MAX_FRAMES, + PTY_SOURCE_SCHEDULER_MAX_SU +} from './pty-source-credit-scheduler' +import type { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' + +export type RelayPtySourceDeliveryRecord = { + clientId: number + identity: PtySourceDeliveryIdentity + displayEnd: number + activating: boolean + sealed: boolean + legacyExitAccepted: boolean + sourceExitAccepted: boolean + sending: boolean + turnFrames: number + turnSourceSu: number + turnScheduled: boolean + sendWaiters: Set<() => void> + recoveryCheckpointSourceEndSu: number | null + recoveryEndSu: number | null + restoreRequired: boolean + rotationPending: boolean +} + +export type RelayPtySourcePublicationCounters = { + opened: number + rotated: number + appendDenied: number + sendCommitted: number + sendRolledBack: number + exitCommitted: number + exitRolledBack: number +} + +export function onceSinkSettlement( + callback: (result: SinkWriteSettlement) => void +): (result: SinkWriteSettlement) => void { + let settled = false + return (result) => { + if (settled) { + return + } + settled = true + callback(result) + } +} + +export class RelayPtySourceSendScheduler { + constructor( + private readonly dispatcher: RelayDispatcher, + private readonly session: SshPtyConsumerSessionAdapter, + private readonly deliveries: Map, + private readonly counters: RelayPtySourcePublicationCounters, + private readonly onCapacity: (id: string) => void + ) {} + + async waitForPendingSend(id: string, timeoutMs = 5_000): Promise { + const record = this.deliveries.get(id) + if (!record) { + return true + } + record.rotationPending = true + if (!record.sending) { + return true + } + return new Promise((resolve) => { + const settle = (settled: boolean): void => { + clearTimeout(timer) + record.sendWaiters.delete(onSettled) + resolve(settled) + } + const onSettled = (): void => settle(true) + const timer = setTimeout(() => settle(false), timeoutMs) + timer.unref?.() + record.sendWaiters.add(onSettled) + }) + } + + onCreditAvailable(id: string): void { + const record = this.deliveries.get(id) + if (!record || record.restoreRequired) { + return + } + if (this.pruneClosed(id, record)) { + this.onCapacity(id) + return + } + this.pump(record) + this.onCapacity(id) + } + + getDebugSnapshot() { + let active = 0 + let activating = 0 + let sealedUnsettled = 0 + let outstandingSourceUnits = 0 + for (const record of this.deliveries.values()) { + const snapshot = this.session.sourceDeliverySnapshot(record.identity) + outstandingSourceUnits += snapshot.sentEndSu - snapshot.creditedEndSu + if (record.activating) { + activating++ + } else if (record.sealed && snapshot.state !== 'closed') { + sealedUnsettled++ + } else if (snapshot.state === 'active') { + active++ + } + } + return Object.freeze({ + active, + activating, + sealedUnsettled, + outstandingSourceUnits, + ...this.counters + }) + } + + dispose(): void { + for (const record of this.deliveries.values()) { + this.session.cancelDelivery(record.identity, 'source-publication-disposed') + this.wakeSendWaiters(record) + } + this.deliveries.clear() + } + + pump(record: RelayPtySourceDeliveryRecord): void { + if ( + record.activating || + record.sending || + record.turnScheduled || + record.restoreRequired || + record.rotationPending || + this.deliveries.get(record.identity.id) !== record + ) { + return + } + if ( + record.turnFrames >= PTY_SOURCE_SCHEDULER_MAX_FRAMES || + record.turnSourceSu >= PTY_SOURCE_SCHEDULER_MAX_SU + ) { + record.turnScheduled = true + setImmediate(() => { + record.turnScheduled = false + record.turnFrames = 0 + record.turnSourceSu = 0 + if (this.deliveries.get(record.identity.id) !== record || record.restoreRequired) { + return + } + this.pump(record) + this.onCapacity(record.identity.id) + }) + return + } + const reservation = this.session.reserveSourceSend(record.identity) + if (!reservation) { + return + } + record.sending = true + const sourceLengthSu = reservation.span.sourceEndSu - reservation.span.sourceStartSu + const settle = onceSinkSettlement((result) => { + record.sending = false + this.wakeSendWaiters(record) + if (this.deliveries.get(record.identity.id) !== record || record.restoreRequired) { + this.onCapacity(record.identity.id) + return + } + if (result.ok) { + this.session.commitSourceSend(reservation) + record.turnFrames++ + record.turnSourceSu += sourceLengthSu + this.counters.sendCommitted++ + this.completeRecoveryIfReady(record) + } else { + this.session.rollbackSourceSend(reservation) + this.counters.sendRolledBack++ + } + this.pruneClosed(record.identity.id, record) + this.onCapacity(record.identity.id) + if (result.ok && this.deliveries.get(record.identity.id) === record) { + this.pump(record) + } + }) + const accepted = this.dispatcher.tryNotifyPtyDataToClient( + record.clientId, + { + id: reservation.identity.id, + data: reservation.span.data, + rawLength: sourceLengthSu, + transformed: reservation.span.transform.transformed, + deliveryToken: reservation.identity.deliveryToken, + clientGeneration: reservation.identity.clientGeneration, + ownerGeneration: reservation.identity.ownerGeneration, + ptyIncarnation: reservation.identity.ptyIncarnation, + sourceEndSu: reservation.span.sourceEndSu, + sourceLengthSu + }, + settle + ) + if (!accepted) { + settle({ ok: false, error: new Error('PTY source publication was not admitted') }) + } + } + + completeRecoveryIfReady(record: RelayPtySourceDeliveryRecord): void { + const recoveryEndSu = record.recoveryEndSu + const checkpointSourceEndSu = record.recoveryCheckpointSourceEndSu + if ( + record.activating || + record.sending || + recoveryEndSu === null || + checkpointSourceEndSu === null + ) { + return + } + const snapshot = this.session.sourceDeliverySnapshot(record.identity) + if (snapshot.sentEndSu < recoveryEndSu) { + return + } + record.recoveryEndSu = null + record.recoveryCheckpointSourceEndSu = null + this.dispatcher.notifyClient(record.clientId, 'pty.recoveryComplete', { + id: record.identity.id, + clientGeneration: record.identity.clientGeneration, + ownerGeneration: record.identity.ownerGeneration, + ptyIncarnation: record.identity.ptyIncarnation, + deliveryToken: record.identity.deliveryToken, + checkpointSourceEndSu, + recoveryEndSu + }) + } + + pruneClosed(id: string, record: RelayPtySourceDeliveryRecord): boolean { + if (this.session.sourceDeliverySnapshot(record.identity).state !== 'closed') { + return false + } + if (this.deliveries.get(id) === record) { + this.deliveries.delete(id) + } + return true + } + + wakeSendWaiters(record: RelayPtySourceDeliveryRecord): void { + for (const resolve of record.sendWaiters) { + resolve() + } + record.sendWaiters.clear() + } + + releaseRotationFence(record: RelayPtySourceDeliveryRecord | undefined): void { + if (!record?.rotationPending) { + return + } + record.rotationPending = false + this.pump(record) + this.onCapacity(record.identity.id) + } +} diff --git a/src/relay/relay.ts b/src/relay/relay.ts index 31e7fab109a..7cbb48763a8 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -55,6 +55,7 @@ import { registerManagedHookInstaller } from './managed-hook-installer' import { registerRelayPluginHostCallHandlers } from './plugin-host-call-handler' import { DispatcherClientWriter } from './dispatcher-client-writer' import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' +import { RelayPtySourcePublication } from './relay-pty-source-publication' const DEFAULT_GRACE_MS = DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 const SOCK_NAME = 'relay.sock' @@ -109,6 +110,7 @@ function parseArgs(argv: string[]): { endpointDir?: string logFile?: string credentialFile?: string + ptySourceCreditV1: boolean } { let graceTimeMs = DEFAULT_GRACE_MS let connectMode = false @@ -118,6 +120,7 @@ function parseArgs(argv: string[]): { let endpointDir: string | undefined let logFile: string | undefined let credentialFile: string | undefined + let ptySourceCreditV1 = false for (let i = 2; i < argv.length; i++) { if (argv[i] === '--grace-time' && argv[i + 1]) { const parsed = Number.parseInt(argv[i + 1], 10) @@ -144,6 +147,8 @@ function parseArgs(argv: string[]): { } else if (argv[i] === '--credential-file' && argv[i + 1]) { credentialFile = argv[i + 1] i++ + } else if (argv[i] === '--pty-source-credit-v1') { + ptySourceCreditV1 = true } } if (!sockPath) { @@ -157,7 +162,8 @@ function parseArgs(argv: string[]): { sockPath, endpointDir, logFile, - credentialFile + credentialFile, + ptySourceCreditV1 } } @@ -420,7 +426,8 @@ async function main(): Promise { sockPath, endpointDir, logFile, - credentialFile + credentialFile, + ptySourceCreditV1 } = parseArgs(process.argv) const endpointCredential = readEndpointCredential(credentialFile) @@ -556,9 +563,16 @@ async function main(): Promise { const ptyConsumerSessionAdapter = new SshPtyConsumerSessionAdapter( dispatcher, launchVersion, - (id, paused) => ptyHandler.setConsumerDeliveryPaused(id, paused) + (id, paused) => ptyHandler.setConsumerDeliveryPaused(id, paused), + (id) => ptyHandler.handleSourceCreditAvailable(id), + ptySourceCreditV1 ) - void ptyConsumerSessionAdapter + const ptySourcePublication = new RelayPtySourcePublication( + dispatcher, + ptyConsumerSessionAdapter, + (id) => ptyHandler.handleSourcePublicationCapacity(id) + ) + ptyHandler.setSourcePublication(ptySourcePublication) const fsHandler = new FsHandler(dispatcher, context) const watchRegistry = fsHandler.getWatchRegistry() ptyHandler.setWorktreeRemovalCoordinator(watchRegistry) @@ -742,6 +756,11 @@ async function main(): Promise { ptys: { active: ptyHandler.activePtyCount }, + ptySourceCredit: { + enabled: ptySourceCreditV1, + session: ptyConsumerSessionAdapter.getDebugSnapshot(), + publication: ptySourcePublication.getDebugSnapshot() + }, socket: { path: sockPath, owned: ownsSocketPath, diff --git a/src/relay/ssh-pty-consumer-session-adapter.test.ts b/src/relay/ssh-pty-consumer-session-adapter.test.ts index bf54c20cf97..a0f040814ed 100644 --- a/src/relay/ssh-pty-consumer-session-adapter.test.ts +++ b/src/relay/ssh-pty-consumer-session-adapter.test.ts @@ -210,6 +210,103 @@ describe('SshPtyConsumerSessionAdapter', () => { expect(setPaused).toEqual([{ id: 'pty-1', paused: true }]) }) + it('token-fences V1 pause and clears an owned pause before detach retention', async () => { + const setPaused: { id: string; paused: boolean }[] = [] + dispatcher = new RelayDispatcher( + (_data, onSettled) => { + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a', (id, paused) => { + setPaused.push({ id, paused }) + }) + dispatcher.feed( + openFrame(1, { + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 8 } } + }) + ) + await flushRequests() + const identity = adapter.openDelivery(1, 'pty-1', 'incarnation-1')! + + for (const [sequence, deliveryToken] of [ + [2, 'stale-token'], + [3, identity.deliveryToken] + ] as const) { + dispatcher.feed( + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.setDeliveryPaused', + params: { + id: identity.id, + paused: true, + clientGeneration: identity.clientGeneration, + ownerGeneration: identity.ownerGeneration, + deliveryToken + } + }, + sequence, + 0 + ) + ) + } + dispatcher.invalidateClient() + + expect(setPaused).toEqual([ + { id: 'pty-1', paused: true }, + { id: 'pty-1', paused: false } + ]) + }) + + it('clears the exact token pause before cancellation cleanup', async () => { + const setPaused: { id: string; paused: boolean }[] = [] + dispatcher = new RelayDispatcher( + (_data, onSettled) => { + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a', (id, paused) => { + setPaused.push({ id, paused }) + }) + dispatcher.feed( + openFrame(1, { + capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 8 } } + }) + ) + await flushRequests() + const identity = adapter.openDelivery(1, 'pty-1', 'incarnation-1')! + dispatcher.feed( + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.setDeliveryPaused', + params: { + id: identity.id, + paused: true, + clientGeneration: identity.clientGeneration, + ownerGeneration: identity.ownerGeneration, + deliveryToken: identity.deliveryToken + } + }, + 2, + 0 + ) + ) + + adapter.cancelDelivery(identity, 'restore-required') + + expect(setPaused).toEqual([ + { id: 'pty-1', paused: true }, + { id: 'pty-1', paused: false } + ]) + }) + it('intersects the negotiated V1 source window without changing legacy omission', async () => { const writes: Buffer[] = [] dispatcher = new RelayDispatcher( diff --git a/src/relay/ssh-pty-consumer-session-adapter.ts b/src/relay/ssh-pty-consumer-session-adapter.ts index 93c401c0ed3..ecdcfc5c109 100644 --- a/src/relay/ssh-pty-consumer-session-adapter.ts +++ b/src/relay/ssh-pty-consumer-session-adapter.ts @@ -71,24 +71,37 @@ function requireIdentity(context: RequestContext): RelayClientSessionIdentity { export class SshPtyConsumerSessionAdapter { private readonly session: PtyConsumerSession private readonly sourceCredit: SshPtySourceCreditAdapter + private readonly pausedDeliveryByPty = new Map() constructor( dispatcher: RelayDispatcher, serverBuildId: string, - setDeliveryPaused?: (id: string, paused: boolean) => void + private readonly setDeliveryPaused?: (id: string, paused: boolean) => void, + onSourceCreditAvailable?: (id: string) => void, + sourceCreditEnabled = true ) { - this.sourceCredit = new SshPtySourceCreditAdapter((proof) => - dispatcher.notifyControl('pty.deliveryCanceled', proof as unknown as Record) + this.sourceCredit = new SshPtySourceCreditAdapter( + (proof) => + dispatcher.notifyControl( + 'pty.deliveryCanceled', + proof as unknown as Record + ), + onSourceCreditAvailable ) this.session = new PtyConsumerSession({ serverBuildId, - outputFlowControl: { versions: [1], maxWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU } + ...(sourceCreditEnabled + ? { outputFlowControl: { versions: [1], maxWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU } } + : {}) }) dispatcher.onRequest(SSH_PTY_OPEN_CLIENT_METHOD, (params, context) => this.openClient(params, context) ) dispatcher.onClientDetached((clientId) => { const grant = this.session.activeGrant(String(clientId)) + if (grant) { + this.clearPausedForGrant(grant) + } this.session.close(String(clientId)) if (grant) { this.sourceCredit.retainOrCloseOnDetach(grant) @@ -105,7 +118,21 @@ export class SshPtyConsumerSessionAdapter { ) { return } - setDeliveryPaused?.(params.id, params.paused) + if (grant.capabilities?.outputFlowControl) { + const token = typeof params.deliveryToken === 'string' ? params.deliveryToken : '' + const identity = this.sourceCredit.ownsDelivery(token, grant, params.id) + if (!identity) { + return + } + if (params.paused) { + this.pausedDeliveryByPty.set(params.id, identity) + } else if (this.pausedDeliveryByPty.get(params.id) !== identity) { + return + } else { + this.pausedDeliveryByPty.delete(params.id) + } + } + this.setDeliveryPaused?.(params.id, params.paused) }) dispatcher.onNotification('pty.ackData', (params, context) => { this.sourceCredit.acknowledge(params, this.session.activeGrant(String(context.clientId))) @@ -113,7 +140,13 @@ export class SshPtyConsumerSessionAdapter { dispatcher.onRequest('pty.cancelDelivery', async (params, context) => this.sourceCredit.cancel(params, this.session.activeGrant(String(context.clientId))) ) - dispatcher.onDisposed(() => this.sourceCredit.dispose()) + dispatcher.onDisposed(() => { + for (const id of this.pausedDeliveryByPty.keys()) { + this.setDeliveryPaused?.(id, false) + } + this.pausedDeliveryByPty.clear() + this.sourceCredit.dispose() + }) } openDelivery( @@ -135,6 +168,7 @@ export class SshPtyConsumerSessionAdapter { newClientId: number, acceptedSourceEndSu: number ) { + this.clearPausedIdentity(oldIdentity) return this.sourceCredit.rotate( oldIdentity, this.session.activeGrant(String(newClientId)), @@ -186,6 +220,29 @@ export class SshPtyConsumerSessionAdapter { return this.sourceCredit.snapshot(identity) } + cancelDelivery(identity: PtySourceDeliveryIdentity, reason: string): void { + this.clearPausedIdentity(identity) + this.sourceCredit.cancelIdentity(identity, reason) + } + + getDebugSnapshot(): Readonly<{ + deliveryTokens: number + graceTimers: number + sourceSu: number + dataBytes: number + spans: number + }> { + return this.sourceCredit.retentionSnapshot() + } + + deliveryMode(clientId: number): 'source-owner' | 'legacy-owner' | 'subscriber' { + const grant = this.session.activeGrant(String(clientId)) + if (grant?.role !== 'session-owner') { + return 'subscriber' + } + return grant.capabilities?.outputFlowControl ? 'source-owner' : 'legacy-owner' + } + private async openClient( rawParams: Record, context: RequestContext @@ -216,4 +273,24 @@ export class SshPtyConsumerSessionAdapter { }) return admission.grant } + + private clearPausedIdentity(identity: PtySourceDeliveryIdentity): void { + if (this.pausedDeliveryByPty.get(identity.id) === identity) { + this.pausedDeliveryByPty.delete(identity.id) + this.setDeliveryPaused?.(identity.id, false) + } + } + + private clearPausedForGrant(grant: Readonly): void { + for (const [id, identity] of this.pausedDeliveryByPty) { + if ( + identity.clientGeneration !== grant.clientGeneration || + identity.ownerGeneration !== grant.ownerGeneration + ) { + continue + } + this.pausedDeliveryByPty.delete(id) + this.setDeliveryPaused?.(id, false) + } + } } diff --git a/src/relay/ssh-pty-source-credit-adapter.test.ts b/src/relay/ssh-pty-source-credit-adapter.test.ts index da9b4514868..f738395a154 100644 --- a/src/relay/ssh-pty-source-credit-adapter.test.ts +++ b/src/relay/ssh-pty-source-credit-adapter.test.ts @@ -76,7 +76,13 @@ describe('SshPtySourceCreditAdapter cleanup', () => { adapter.settleExit(identity, { ok: true }) } - expect(adapter.retentionSnapshot()).toEqual({ deliveryTokens: 0, graceTimers: 0 }) + expect(adapter.retentionSnapshot()).toEqual({ + deliveryTokens: 0, + graceTimers: 0, + sourceSu: 0, + dataBytes: 0, + spans: 0 + }) }) it('disposes grace timers and exact retained tokens', () => { @@ -86,11 +92,23 @@ describe('SshPtySourceCreditAdapter cleanup', () => { const first = adapter.open(grant, 'pty-1', 'incarnation-1')! const second = adapter.open(grant, 'pty-2', 'incarnation-2')! adapter.retainOrCloseOnDetach(grant) - expect(adapter.retentionSnapshot()).toEqual({ deliveryTokens: 2, graceTimers: 1 }) + expect(adapter.retentionSnapshot()).toEqual({ + deliveryTokens: 2, + graceTimers: 1, + sourceSu: 0, + dataBytes: 0, + spans: 0 + }) adapter.dispose() - expect(adapter.retentionSnapshot()).toEqual({ deliveryTokens: 0, graceTimers: 0 }) + expect(adapter.retentionSnapshot()).toEqual({ + deliveryTokens: 0, + graceTimers: 0, + sourceSu: 0, + dataBytes: 0, + spans: 0 + }) expect(adapter.snapshot(first)).toMatchObject({ state: 'closed', generationClosed: true }) expect(adapter.snapshot(second)).toMatchObject({ state: 'closed', generationClosed: true }) expect(() => adapter.open(grant, 'pty-late', 'incarnation-late')).toThrow('disposed') diff --git a/src/relay/ssh-pty-source-credit-adapter.ts b/src/relay/ssh-pty-source-credit-adapter.ts index 938fce0422b..d551dec8142 100644 --- a/src/relay/ssh-pty-source-credit-adapter.ts +++ b/src/relay/ssh-pty-source-credit-adapter.ts @@ -15,6 +15,7 @@ import { ptySourceCancellationResult, RecentPtySourceCancellationIndex } from './pty-source-cancellation-index' +import { ownedPtySourceDelivery } from './pty-source-delivery-ownership' import { RelayPtySourceCreditLedger, type PtySourceSendReservation @@ -29,7 +30,8 @@ export class SshPtySourceCreditAdapter { private disposed = false constructor( - private readonly publishCancellation?: (proof: PtySourceDeliveryCancellation) => void + private readonly publishCancellation?: (proof: PtySourceDeliveryCancellation) => void, + private readonly onCreditAvailable?: (id: string) => void ) {} open( @@ -88,6 +90,7 @@ export class SshPtySourceCreditAdapter { ) this.identityByToken.delete(oldIdentity.deliveryToken) this.recentCancellations.remember(rotation.cancellation) + this.publishCancellation?.(rotation.cancellation) this.identityByToken.set(replacement.deliveryToken, replacement) this.clearGraceWhenSettled(oldIdentity.ownerGeneration) return Object.freeze({ identity: replacement, ...rotation }) @@ -161,13 +164,16 @@ export class SshPtySourceCreditAdapter { continue } try { - this.sourceCredit.acknowledge(identity, { + const result = this.sourceCredit.acknowledge(identity, { id: String(candidate.id ?? ''), clientGeneration: Number(candidate.clientGeneration), ownerGeneration: Number(candidate.ownerGeneration), deliveryToken: String(candidate.deliveryToken), creditedEndSu: Number(candidate.creditedEndSu) }) + if (result === 'advanced') { + this.onCreditAvailable?.(identity.id) + } this.pruneClosed(token, identity) } catch { /* Invalid and stale cumulative ACKs never mutate credit. */ @@ -212,13 +218,30 @@ export class SshPtySourceCreditAdapter { this.closeClientDeliveries(grant.clientGeneration) } - retentionSnapshot(): Readonly<{ deliveryTokens: number; graceTimers: number }> { + retentionSnapshot() { return Object.freeze({ deliveryTokens: this.identityByToken.size, - graceTimers: this.graceTimers.size + graceTimers: this.graceTimers.size, + ...this.sourceCredit.retentionSnapshot() }) } + ownsDelivery( + token: string, + grant: Readonly, + id: string + ): PtySourceDeliveryIdentity | null { + return ownedPtySourceDelivery(this.identityByToken.get(token), grant, id) + } + + cancelIdentity(identity: PtySourceDeliveryIdentity, reason: string): void { + const token = identity.deliveryToken + if (this.identityByToken.get(token) !== identity) { + return + } + this.cancelExact(token, identity, reason) + } + dispose(): void { if (this.disposed) { return @@ -279,14 +302,10 @@ export class SshPtySourceCreditAdapter { } } - private hasOwnerDeliveries(ownerGeneration: number): boolean { - for (const identity of this.identityByToken.values()) { - if (identity.ownerGeneration === ownerGeneration) { - return true - } - } - return false - } + private hasOwnerDeliveries = (ownerGeneration: number): boolean => + Array.from(this.identityByToken.values()).some( + (identity) => identity.ownerGeneration === ownerGeneration + ) private clearGraceWhenSettled(ownerGeneration: number): void { if (!this.hasOwnerDeliveries(ownerGeneration)) { diff --git a/src/shared/pty-source-recovery-contract.ts b/src/shared/pty-source-recovery-contract.ts new file mode 100644 index 00000000000..2634072396c --- /dev/null +++ b/src/shared/pty-source-recovery-contract.ts @@ -0,0 +1,29 @@ +export type PtySourceRecoveryCheckpoint = Readonly<{ + status: 'checkpoint' + clientGeneration: number + ownerGeneration: number + ptyIncarnation: string + deliveryToken: string + acceptedSourceEndSu: number +}> + +export type PtySourceRecoveryRequest = + | PtySourceRecoveryCheckpoint + | Readonly<{ status: 'checkpointUnavailable' }> + +export type PtySourceRecoveryPending = Readonly<{ + status: 'pending' + clientGeneration: number + ownerGeneration: number + ptyIncarnation: string + deliveryToken: string + checkpointSourceEndSu: number + recoveryEndSu: number +}> + +export type PtySourceRecoveryResult = + | PtySourceRecoveryPending + | Readonly<{ status: 'restoreRequired'; reason: string }> + +export type PtySourceRecoveryComplete = Omit & + Readonly<{ id: string }> diff --git a/src/shared/relay-frame-decoder-contract.ts b/src/shared/relay-frame-decoder-contract.ts new file mode 100644 index 00000000000..344df44c93d --- /dev/null +++ b/src/shared/relay-frame-decoder-contract.ts @@ -0,0 +1,17 @@ +export type DecodedFrame = { + type: number + id: number + ack: number + payload: Buffer +} + +export type FrameDecoderOptions = { + maxFramesPerTurn?: number + maxBytesPerTurn?: number + maxTurnMs?: number + now?: () => number + schedule?: (callback: () => void) => unknown + cancelScheduled?: (handle: unknown) => void + pause?: () => void + resume?: () => void +} diff --git a/tests/e2e/ssh-docker-relay-perf.spec.ts b/tests/e2e/ssh-docker-relay-perf.spec.ts index 1d7606fcbe3..c5f0836870b 100644 --- a/tests/e2e/ssh-docker-relay-perf.spec.ts +++ b/tests/e2e/ssh-docker-relay-perf.spec.ts @@ -21,6 +21,7 @@ import { } from './helpers/docker-ssh-relay-connection' const RUN_DOCKER_SSH = process.env.ORCA_E2E_SSH_DOCKER === '1' +const RUN_SOURCE_CREDIT_V1 = process.env.ORCA_SSH_PTY_SOURCE_CREDIT_V1 === '1' const KEY_LATENCY_SAMPLES = 'abcdefghij' const MAX_MEDIAN_KEY_LATENCY_MS = 500 const MAX_WORST_KEY_LATENCY_MS = 2_000 @@ -217,12 +218,16 @@ test.describe('Docker SSH relay perf', () => { `node -e ${shellQuote(remoteTypingLoadScript(activeRunId))}` ) await waitForTerminalOutput(orcaPage, `REMOTE_TUI_READY_${activeRunId}`, 30_000, 80_000) - await expect - .poll(async () => (await readSshPtyAckGate(orcaPage))?.heldAckChars ?? 0, { + const heldAckPressure = expect.poll( + async () => (await readSshPtyAckGate(orcaPage))?.heldAckChars ?? 0, + { timeout: 30_000, message: 'remote background SSH PTY stream did not build held ACK pressure' - }) - .toBeGreaterThan(MIN_HELD_SSH_ACK_CHARS) + } + ) + await (RUN_SOURCE_CREDIT_V1 + ? heldAckPressure.toBe(MIN_HELD_SSH_ACK_CHARS) + : heldAckPressure.toBeGreaterThan(MIN_HELD_SSH_ACK_CHARS)) const measurement = await measureRemoteTyping(orcaPage, activePtyId, activeRunId) const ackGate = await readSshPtyAckGate(orcaPage) @@ -238,7 +243,11 @@ test.describe('Docker SSH relay perf', () => { type: 'docker-ssh-relay-pty-ack-pressure', description: summary }) - expect(ackGate?.heldAckChars ?? 0).toBeGreaterThan(MIN_HELD_SSH_ACK_CHARS) + if (RUN_SOURCE_CREDIT_V1) { + expect(ackGate?.heldAckChars ?? 0).toBe(MIN_HELD_SSH_ACK_CHARS) + } else { + expect(ackGate?.heldAckChars ?? 0).toBeGreaterThan(MIN_HELD_SSH_ACK_CHARS) + } expect(measurement.medianLatencyMs).toBeLessThan(MAX_MEDIAN_KEY_LATENCY_MS) expect(measurement.worstLatencyMs).toBeLessThan(MAX_WORST_KEY_LATENCY_MS) @@ -368,11 +377,30 @@ test.describe('Docker SSH relay perf', () => { const beforeMarker = `SSH_RECONNECT_BEFORE_${Date.now()}` await execInTerminal(orcaPage, beforePtyId, `printf ${shellQuote(beforeMarker)}`) await waitForTerminalOutput(orcaPage, beforeMarker, 20_000, 60_000) + const recoveryStartedMarker = `SSH_RECONNECT_RECOVERY_STARTED_${Date.now()}` + const recoveryMarker = `SSH_RECONNECT_RECOVERY_${Date.now()}` + if (RUN_SOURCE_CREDIT_V1) { + const recoveryScript = [ + 'let frame = 0', + "const chunk = 'Q'.repeat(4096)", + `process.stdout.write('${recoveryStartedMarker}\\n')`, + 'const timer = setInterval(() => {', + 'frame += 1', + "process.stdout.write('RECOVERY_FRAME_' + frame + '_' + chunk + '\\n')", + `if (frame === 256) { clearInterval(timer); process.stdout.write('${recoveryMarker}\\n') }`, + '}, 10)' + ].join(';') + await execInTerminal(orcaPage, beforePtyId, `node -e ${shellQuote(recoveryScript)}`) + await waitForTerminalOutput(orcaPage, recoveryStartedMarker, 30_000, 80_000) + } await reconnectDockerSshRelayTarget(orcaPage, remote.targetId) await ensureTerminalVisible(orcaPage, 45_000) await waitForActiveTerminalManager(orcaPage, 60_000) const afterPtyId = await waitForActivePanePtyId(orcaPage, 60_000) + if (RUN_SOURCE_CREDIT_V1) { + await waitForTerminalOutput(orcaPage, recoveryMarker, 30_000, 80_000) + } const afterMarker = `SSH_RECONNECT_AFTER_${Date.now()}` const remoteProofPath = `/tmp/${afterMarker}` await execInTerminal(