fix(ssh): recover bounded PTY output across reconnect

This commit is contained in:
OrcaWin
2026-07-28 20:13:36 -07:00
parent efa043d088
commit 022856af3b
62 changed files with 4333 additions and 339 deletions
+31 -7
View File
@@ -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.",
@@ -21,7 +21,7 @@ function fakeChild(): FakeChild {
resume: ReturnType<typeof vi.fn>
}
stderr: EventEmitter
stdin: { write: ReturnType<typeof vi.fn> }
stdin: EventEmitter & { write: ReturnType<typeof vi.fn> }
kill: ReturnType<typeof vi.fn>
}
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<typeof vi.fn>
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)
})
})
@@ -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)
+68
View File
@@ -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 {
+26 -3
View File
@@ -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<string>()
const rendererCreditBeforeExitByPty = new Map<string, boolean>()
const rendererDeliveryRestoreNeededPtys = new Set<string>()
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<string>()
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()
}
@@ -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
@@ -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({
+4
View File
@@ -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(),
@@ -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()
}
@@ -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)
}
@@ -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()
@@ -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 {
@@ -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<string, TokenRecord>,
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<string, TokenRecord>,
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<string, TokenRecord>,
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
}
+3 -1
View File
@@ -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)
@@ -3,11 +3,13 @@ import { subscribeSshPtyNotifications } from './ssh-pty-notification-routing'
type MockMux = {
onNotification: ReturnType<typeof vi.fn>
request: ReturnType<typeof vi.fn>
}
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<string>()
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()
})
})
@@ -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<string, SourceDeliveryState>()
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<string, SourceDeliveryState>,
relayPtyId: string,
params: Record<string, unknown>,
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<string, unknown>
): 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
}
+6 -7
View File
@@ -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<PtySpawnResult> {
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<SshPtyAttachResult> {
// 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 } : {})
})
+59 -2
View File
@@ -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<string, unknown>
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
+3 -1
View File
@@ -1,6 +1,6 @@
import type { SshPtyDataCallback } from './ssh-pty-provider-contract'
type SshPtySourceFrame = NonNullable<Parameters<SshPtyDataCallback>[0]['source']>
export type SshPtySourceFrame = NonNullable<Parameters<SshPtyDataCallback>[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)
) {
+3 -18
View File
@@ -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
@@ -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' })
})
})
@@ -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()
})
})
+3 -1
View File
@@ -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,
+21 -41
View File
@@ -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'
}
@@ -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<typeof vi.fn>
}
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()
})
})
@@ -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<WriterEntry>()
private readonly settleOnDrain = new Set<WriterEntry>()
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)
}
}
@@ -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)
})
})
@@ -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'
}
+36 -3
View File
@@ -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<typeof vi.fn> }
}
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', () => {
+6 -1
View File
@@ -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
+5 -2
View File
@@ -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)')
+18 -5
View File
@@ -115,7 +115,8 @@ export async function deployAndLaunchRelay(
conn: SshConnection,
onProgress?: (status: string) => void,
graceTimeSeconds?: number,
relayInstanceId?: string
relayInstanceId?: string,
enablePtySourceCreditV1 = false
): Promise<RelayDeployResult> {
let timeoutHandle: ReturnType<typeof setTimeout>
// 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<RelayDeployResult> {
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<RelayDeployResult> {
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 </dev/null &`
const sourceCreditFlag = enablePtySourceCreditV1 ? ' --pty-source-credit-v1' : ''
const launchCmd = `cd ${escapedDir} && chmod 600 ${shellEscape(credentialFile)} && nohup ${escapedNode} relay.js --detached${sourceCreditFlag} --grace-time ${graceTime} --sock-path ${shellEscape(sockFile)} --credential-file ${shellEscape(credentialFile)} --log-file ${shellEscape(logFile)} > ${shellEscape(logFile)} 2>&1 </dev/null &`
const launchChannel = await conn.exec(launchCmd, { signal })
launchChannel.on('data', () => {})
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',
@@ -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<string, unknown>) => 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<string, unknown>) => 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<string, unknown>) => void) | undefined
canceled?.({
id: 'pty-1',
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-1',
deliveryToken: 'old-token'
})
expect(closeSshPtyOutputGeneration).toHaveBeenCalledTimes(closeCount)
})
})
@@ -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()
@@ -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()
+3 -1
View File
@@ -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 () => {
+556 -41
View File
@@ -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<SshPtyExitCallback>[0]
type PendingPtyReattach = { exits: SshPtyExitPayload[] }
type SshPtyDataPayload = Parameters<SshPtyDataCallback>[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<unknown>[]
queuedBytes: number
queuedFrames: number
livePassthrough: boolean
activated: boolean
}
type RemoteCliBridgeEnv = {
remoteHome: string
@@ -115,6 +140,40 @@ function expectedIdentityForLease(lease: {
}
}
function parseRecoveryComplete(params: Record<string, unknown>): 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<string, SshPtyAcceptedSourceCheckpoint>
}
const ptyConsumerRecoveryByTarget = new Map<string, PtyConsumerRecovery>()
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<string, SshPtyAcceptedSourceCheckpoint>()
}
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<string, PendingPtyReattach>()
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<string, SshPtyAcceptedSourceCheckpoint>()
})
}
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<unknown> {
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<void> {
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<boolean> {
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<void> {
const deadline = Date.now() + 30_000
while (
shouldContinue() &&
!pending.recoveryComplete &&
!pending.restoreRequired &&
Date.now() < deadline
) {
await new Promise<void>((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<void> {
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 })
}
}
}
@@ -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()
+9 -4
View File
@@ -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()) {
+6 -1
View File
@@ -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
)
}
+74 -1
View File
@@ -234,6 +234,43 @@ export class RelayDispatcher {
)
}
tryNotifyPtyDataToMatchingClients(
matchesClient: (clientId: number) => boolean,
params: Record<string, unknown>,
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<string, unknown>,
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<string, unknown>): boolean {
if (this.disposed) {
return false
@@ -249,6 +286,42 @@ export class RelayDispatcher {
)
}
tryNotifyPtyExitToMatchingClients(
matchesClient: (clientId: number) => boolean,
params: Record<string, unknown>
): 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<string, unknown>,
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<string, unknown>,
@@ -806,7 +879,7 @@ export class RelayDispatcher {
private publishBulkWhenAvailable(client: RelayClient, msg: JsonRpcNotification): Promise<void> {
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<void>((resolve, reject) => {
+1 -8
View File
@@ -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'
+72
View File
@@ -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({
+125 -14
View File
@@ -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<string, unknown>
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<string>()
private consumerPausedOutputPtys = new Set<string>()
private removeLegacyCapacityListener: (() => void) | null = null
private sourcePublication: RelayPtySourcePublication | null = null
private lastInputAtByPty = new Map<string, number>()
private interactiveOutputCharsByPty = new Map<string, number>()
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<typeof NodePty | null> {
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<string, unknown>
): Promise<{ incarnationId: string; replay?: string }> {
params: Record<string, unknown>,
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)
@@ -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) =>
+62 -60
View File
@@ -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<string, DeliveryRecord>()
private readonly upstreamOwnerByPty = new Map<string, string>()
private readonly closedSnapshots = new Map<string, PtySourceDeliverySnapshot>()
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) {
+39
View File
@@ -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<PtySourceCreditLedgerOptions>
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)
}
+81 -3
View File
@@ -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<number>
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<DeliveryRecord>): number {
let total = 0
for (const record of records) {
@@ -171,6 +222,22 @@ export function retainedSourceTotal(records: Iterable<DeliveryRecord>): number {
return total
}
export function retainedDataBytesTotal(records: Iterable<DeliveryRecord>): number {
let total = 0
for (const record of records) {
total += record.retainedDataBytes
}
return total
}
export function retainedSpanTotal(records: Iterable<DeliveryRecord>): 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
}
@@ -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<PtyConsumerSessionGrant>,
id: string
): PtySourceDeliveryIdentity | null {
if (
!identity ||
identity.id !== id ||
identity.clientGeneration !== grant.clientGeneration ||
identity.ownerGeneration !== grant.ownerGeneration
) {
return null
}
return identity
}
+3 -18
View File
@@ -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
+8
View File
@@ -0,0 +1,8 @@
export type RelayPtySourceOutput = {
data: string
rawLength?: number
transformed?: boolean
seq?: number
sourceAccepted?: boolean
sourceSpanId?: string
}
@@ -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<string, unknown>): Buffer {
return encodeJsonRpcFrame({ jsonrpc: '2.0', id, method, params }, id, 0)
}
function notification(buffer: Buffer): { method: string; params: Record<string, unknown> } | 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<string, unknown> | 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<void> {
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<typeof frame> => 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<typeof frame> => frame !== null)
.map((frame) => frame.method)
).toEqual(['pty.deliveryCanceled', 'pty.data', 'pty.recoveryComplete', 'pty.data'])
await flushRequests()
})
})
+308
View File
@@ -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<string, RelayPtySourceDeliveryRecord>()
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<boolean> =>
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
}
}
@@ -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<string, RelayPtySourceDeliveryRecord>,
private readonly counters: RelayPtySourcePublicationCounters,
private readonly onCapacity: (id: string) => void
) {}
async waitForPendingSend(id: string, timeoutMs = 5_000): Promise<boolean> {
const record = this.deliveries.get(id)
if (!record) {
return true
}
record.rotationPending = true
if (!record.sending) {
return true
}
return new Promise<boolean>((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)
}
}
+23 -4
View File
@@ -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<void> {
sockPath,
endpointDir,
logFile,
credentialFile
credentialFile,
ptySourceCreditV1
} = parseArgs(process.argv)
const endpointCredential = readEndpointCredential(credentialFile)
@@ -556,9 +563,16 @@ async function main(): Promise<void> {
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<void> {
ptys: {
active: ptyHandler.activePtyCount
},
ptySourceCredit: {
enabled: ptySourceCreditV1,
session: ptyConsumerSessionAdapter.getDebugSnapshot(),
publication: ptySourcePublication.getDebugSnapshot()
},
socket: {
path: sockPath,
owned: ownsSocketPath,
@@ -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(
+83 -6
View File
@@ -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<string, PtySourceDeliveryIdentity>()
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<string, unknown>)
this.sourceCredit = new SshPtySourceCreditAdapter(
(proof) =>
dispatcher.notifyControl(
'pty.deliveryCanceled',
proof as unknown as Record<string, unknown>
),
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<string, unknown>,
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<PtyConsumerSessionGrant>): 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)
}
}
}
@@ -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')
+31 -12
View File
@@ -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<PtyConsumerSessionGrant>,
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)) {
@@ -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<PtySourceRecoveryPending, 'status'> &
Readonly<{ id: string }>
@@ -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
}
+33 -5
View File
@@ -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(