fix(ssh): handle rejected PTY deliveries with targeted recovery (#12746)

Add targeted recovery for rejected PTY source frames instead of
terminating the relay channel. Classify rejection reasons (malformed,
generation mismatch, range invalid) and attempt recovery based on the
rejection type. Implement admission control at publication time to ensure
frames aren't delivered after ownership changes. Bound recovery attempts
and retry with backoff to prevent exhaustion. Diagnose and log rejection
reasons to aid debugging.
This commit is contained in:
Jinjing 2026-08-05 10:30:04 -07:00 committed by GitHub
parent eea0bb64db
commit 9fb4dbe8eb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
31 changed files with 2181 additions and 359 deletions

View File

@ -0,0 +1,183 @@
import { expect, it, vi } from 'vitest'
import { subscribeSshPtyNotifications } from './ssh-pty-notification-routing'
it('routes malformed and unadmitted source frames only to rejection diagnostics', async () => {
const mux = {
onNotification: vi.fn(),
request: vi.fn(async (_method: string, params: Record<string, unknown>) => {
if (params.deliveryToken !== 'token-1') {
throw new Error('Unknown or stale PTY source delivery cancellation')
}
return { canceled: true, sentEndSu: 0, creditedEndSu: 0 }
})
}
const dataListeners = new Set()
const rejectedDataListeners = new Set()
const subscription = subscribeSshPtyNotifications({
mux: mux as never,
toAppPtyId: (id) => `ssh:conn@@${id}`,
dataListeners: dataListeners as never,
rejectedDataListeners: rejectedDataListeners as never,
replayListeners: new Set() as never,
exitListeners: new Set() as never,
livePtyIds: new Set(),
recordExit: vi.fn(),
providerGeneration: 7,
resolvePtyIncarnation: (id) => `incarnation:${id}`,
peekPtyIncarnation: () => undefined
})
const handler = mux.onNotification.mock.calls[0]?.[0] as (
method: string,
params: Record<string, unknown>
) => void
const onData = vi.fn()
const onRejectedData = vi.fn()
dataListeners.add(onData)
rejectedDataListeners.add(onRejectedData)
subscription
.installReceivingActivation('pty-1', {
status: 'pending',
clientGeneration: 2,
ownerGeneration: 3,
ptyIncarnation: 'incarnation-1',
deliveryToken: 'token-1',
checkpointSourceEndSu: 0,
recoveryEndSu: 4
})
.commit()
handler('pty.data', {
id: 'pty-1',
data: 'bad',
deliveryToken: 'token-1',
clientGeneration: 2,
ownerGeneration: 3,
ptyIncarnation: 'incarnation-1',
sourceEndSu: 4,
sourceLengthSu: 4
})
handler('pty.data', {
id: 'pty-2',
data: 'old',
deliveryToken: 'token-old',
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-1',
sourceEndSu: 3,
sourceLengthSu: 3
})
await vi.waitFor(() => expect(onRejectedData).toHaveBeenCalledTimes(2))
expect(onData).not.toHaveBeenCalled()
expect(onRejectedData).toHaveBeenNthCalledWith(
1,
expect.objectContaining({ sourceMalformed: true })
)
expect(onRejectedData).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
sourceRejected: true,
source: expect.objectContaining({
deliveryToken: 'token-old',
clientGeneration: 1,
ownerGeneration: 1
})
})
)
})
it('never resolves an incarnation for a rejected frame', async () => {
const mux = {
onNotification: vi.fn(),
request: vi.fn(async () => ({ canceled: true, sentEndSu: 0, creditedEndSu: 0 }))
}
const rejectedDataListeners = new Set()
const resolvePtyIncarnation = vi.fn((id: string) => `incarnation:${id}`)
subscribeSshPtyNotifications({
mux: mux as never,
toAppPtyId: (id) => `ssh:conn@@${id}`,
dataListeners: new Set() as never,
rejectedDataListeners: rejectedDataListeners as never,
replayListeners: new Set() as never,
exitListeners: new Set() as never,
livePtyIds: new Set(),
recordExit: vi.fn(),
providerGeneration: 7,
resolvePtyIncarnation,
peekPtyIncarnation: () => undefined
})
const handler = mux.onNotification.mock.calls[0]?.[0] as (
method: string,
params: Record<string, unknown>
) => void
const onRejectedData = vi.fn()
rejectedDataListeners.add(onRejectedData)
// Why this matters: resolvePtyIncarnation caches what it mints and rememberPtyIncarnation is
// first-write-wins, so a synthetic id minted here would fence the PTY's real incarnation off for
// the rest of the generation.
handler('pty.data', { id: 'pty-1', data: 'bad', deliveryToken: 'token-1', sourceEndSu: 4 })
await vi.waitFor(() => expect(onRejectedData).toHaveBeenCalledOnce())
expect(onRejectedData).toHaveBeenCalledOnce()
expect(onRejectedData).toHaveBeenCalledWith(expect.objectContaining({ sourceMalformed: true }))
expect(resolvePtyIncarnation).not.toHaveBeenCalled()
})
it('coalesces repeated exact rejections into one fresh-activation recovery', async () => {
const mux = {
onNotification: vi.fn(),
request: vi.fn(async () => ({ canceled: true, sentEndSu: 3, creditedEndSu: 0 }))
}
const rejectedDataListeners = new Set()
const subscription = subscribeSshPtyNotifications({
mux: mux as never,
toAppPtyId: (id) => `ssh:conn@@${id}`,
dataListeners: new Set() as never,
rejectedDataListeners: rejectedDataListeners as never,
replayListeners: new Set() as never,
exitListeners: new Set() as never,
livePtyIds: new Set(),
recordExit: vi.fn(),
providerGeneration: 7,
resolvePtyIncarnation: (id) => `incarnation:${id}`,
peekPtyIncarnation: () => 'incarnation-1'
})
subscription
.installReceivingActivation('pty-1', {
status: 'pending',
clientGeneration: 2,
ownerGeneration: 3,
ptyIncarnation: 'incarnation-1',
deliveryToken: 'token-1',
checkpointSourceEndSu: 0,
recoveryEndSu: 0
})
.commit()
const handler = mux.onNotification.mock.calls[0]?.[0] as (
method: string,
params: Record<string, unknown>
) => void
const onRejectedData = vi.fn()
rejectedDataListeners.add(onRejectedData)
const rejected = {
id: 'pty-1',
data: 'gap',
deliveryToken: 'token-1',
clientGeneration: 2,
ownerGeneration: 3,
ptyIncarnation: 'incarnation-1',
sourceEndSu: 6,
sourceLengthSu: 3
}
handler('pty.data', rejected)
handler('pty.data', rejected)
await vi.waitFor(() => expect(onRejectedData).toHaveBeenCalledOnce())
expect(onRejectedData).toHaveBeenCalledWith(
expect.objectContaining({ rejectedSourceRecovery: 'fresh-activation' })
)
})

View File

@ -29,7 +29,8 @@ function createSubscription() {
livePtyIds,
recordExit,
providerGeneration: 7,
resolvePtyIncarnation
resolvePtyIncarnation,
peekPtyIncarnation: () => undefined
})
const handler = mux.onNotification.mock.calls[0]?.[0] as (

View File

@ -7,10 +7,11 @@ import type {
SshPtyReplayCallback
} from './ssh-pty-provider-contract'
import { parseSshPtySourceFrame } from './ssh-pty-source-frame'
import {
SshPtySourceDeliveryLedger,
type PendingSshPtySourceData
} from './ssh-pty-source-delivery-ledger'
import { SshPtySourceDeliveryLedger } from './ssh-pty-source-delivery-ledger'
import type {
PendingSshPtySourceData,
SshPtyRejectedSourceRecovery
} from './ssh-pty-source-delivery-state'
export type { SshPtyDataCallback, SshPtyExitCallback, SshPtyReplayCallback }
export type SshPtyRecoveryActivationLease = Readonly<{
@ -35,18 +36,25 @@ export function subscribeSshPtyNotifications(args: {
mux: SshChannelMultiplexer
toAppPtyId: (id: string) => string
dataListeners: Set<SshPtyDataCallback>
rejectedDataListeners?: Set<SshPtyDataCallback>
replayListeners: Set<SshPtyReplayCallback>
exitListeners: Set<SshPtyExitCallback>
livePtyIds: Set<string>
recordExit: (relayPtyId: string, incarnationId: unknown) => void
providerGeneration: number
resolvePtyIncarnation: (relayPtyId: string, incarnationId?: unknown) => string
peekPtyIncarnation: (relayPtyId: string) => string | undefined
}): SshPtyNotificationSubscription {
const toDataPayload = (pending: PendingSshPtySourceData): Parameters<SshPtyDataCallback>[0] => {
const toDataPayload = (
pending: PendingSshPtySourceData,
incarnationOverride?: string
): Parameters<SshPtyDataCallback>[0] => {
const id = args.toAppPtyId(pending.relayPtyId)
const ptyIncarnation = pending.source
? (pending.params.ptyIncarnation as string)
: args.resolvePtyIncarnation(pending.relayPtyId, pending.params.incarnationId)
const ptyIncarnation =
incarnationOverride ??
(pending.source
? (pending.params.ptyIncarnation as string)
: args.resolvePtyIncarnation(pending.relayPtyId, pending.params.incarnationId))
return {
id,
data: pending.data,
@ -67,7 +75,80 @@ export function subscribeSshPtyNotifications(args: {
listener(payload)
}
}
// Why: a rejected frame is diagnostic and must never mint an incarnation. resolvePtyIncarnation
// caches what it synthesizes and rememberPtyIncarnation is first-write-wins, so a malformed frame
// that lands before the PTY's first good one would pin a `legacy:` id the real attach can never
// displace — fencing every later frame of that generation off as a mismatch.
const rejectedPtyIncarnation = (pending: PendingSshPtySourceData): string => {
const offered = pending.params.ptyIncarnation
if (typeof offered === 'string' && offered.length > 0) {
return offered
}
return args.peekPtyIncarnation(pending.relayPtyId) ?? ''
}
const publishRejectedData = (
pending: PendingSshPtySourceData,
rejection: 'malformed' | 'unadmitted',
recovery: SshPtyRejectedSourceRecovery
): void => {
const listeners = args.rejectedDataListeners
if (!listeners || listeners.size === 0) {
return
}
const payload = {
...toDataPayload(pending, rejectedPtyIncarnation(pending)),
...(rejection === 'malformed' ? { sourceMalformed: true } : {}),
...(rejection === 'unadmitted' ? { sourceRejected: true } : {}),
rejectedSourceRecovery: recovery
}
for (const listener of listeners) {
listener(payload)
}
}
const sourceDeliveries = new SshPtySourceDeliveryLedger(args.mux, publishData)
const rejectedPublications = new Map<
string,
{
pending: number
payload: PendingSshPtySourceData
rejection: 'malformed' | 'unadmitted'
recovery?: SshPtyRejectedSourceRecovery
}
>()
const rejectSourceData = (
pending: PendingSshPtySourceData,
rejection: 'malformed' | 'unadmitted'
): void => {
const batch = rejectedPublications.get(pending.relayPtyId) ?? {
pending: 0,
payload: pending,
rejection
}
batch.pending++
rejectedPublications.set(pending.relayPtyId, batch)
void sourceDeliveries
.reject(pending.relayPtyId, rejectedSourceIdentity(pending.params))
.then((recovery) => {
if (
!batch.recovery ||
rejectedRecoveryPriority(recovery) > rejectedRecoveryPriority(batch.recovery)
) {
batch.payload = pending
batch.rejection = rejection
batch.recovery = recovery
}
})
.finally(() => {
batch.pending--
queueMicrotask(() => {
if (batch.pending > 0 || rejectedPublications.get(pending.relayPtyId) !== batch) {
return
}
rejectedPublications.delete(pending.relayPtyId)
publishRejectedData(batch.payload, batch.rejection, batch.recovery ?? 'reconnect-channel')
})
})
}
const dispose = 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
@ -82,6 +163,7 @@ export function subscribeSshPtyNotifications(args: {
if (method === 'pty.exit') {
const id = args.toAppPtyId(relayPtyId)
const ptyIncarnation = args.resolvePtyIncarnation(relayPtyId, params.incarnationId)
rejectedPublications.delete(relayPtyId)
args.recordExit(relayPtyId, params.incarnationId)
args.livePtyIds.delete(id)
sourceDeliveries.recordExit(relayPtyId)
@ -109,7 +191,8 @@ export function subscribeSshPtyNotifications(args: {
const data = typeof params.data === 'string' ? params.data : ''
const sourceFrame = parseSshPtySourceFrame(params, data, relayPtyId)
if (sourceFrame.malformed) {
cancelExactSourceDelivery(args.mux, relayPtyId, params)
const pending = Object.freeze({ relayPtyId, params, data })
rejectSourceData(pending, 'malformed')
return
}
const pending = Object.freeze({
@ -120,14 +203,17 @@ export function subscribeSshPtyNotifications(args: {
})
if (sourceFrame.source) {
if (!sourceDeliveries.admit({ ...pending, source: sourceFrame.source })) {
cancelExactSourceDelivery(args.mux, relayPtyId, params)
rejectSourceData(pending, 'unadmitted')
}
return
}
publishData(pending)
})
return Object.freeze({
dispose,
dispose: () => {
rejectedPublications.clear()
dispose()
},
installReceivingActivation: (relayPtyId, activation) => {
const lease = sourceDeliveries.install(relayPtyId, activation)
return Object.freeze({
@ -140,33 +226,42 @@ export function subscribeSshPtyNotifications(args: {
})
}
function cancelExactSourceDelivery(
mux: SshChannelMultiplexer,
relayPtyId: string,
params: {
deliveryToken?: unknown
clientGeneration?: unknown
ownerGeneration?: unknown
function rejectedRecoveryPriority(recovery: SshPtyRejectedSourceRecovery): number {
if (recovery === 'reconnect-channel') {
return 3
}
): void {
return recovery === 'fresh-activation' ? 2 : 1
}
function rejectedSourceIdentity(params: {
deliveryToken?: unknown
clientGeneration?: unknown
ownerGeneration?: unknown
ptyIncarnation?: unknown
}):
| Readonly<{
deliveryToken: string
clientGeneration: number
ownerGeneration: number
ptyIncarnation: string
}>
| undefined {
if (
typeof params.deliveryToken !== 'string' ||
params.deliveryToken.length === 0 ||
!positiveSafeInteger(params.clientGeneration) ||
!positiveSafeInteger(params.ownerGeneration)
!positiveSafeInteger(params.ownerGeneration) ||
typeof params.ptyIncarnation !== 'string' ||
params.ptyIncarnation.length === 0
) {
return
return undefined
}
try {
void mux
.request('pty.cancelDelivery', {
id: relayPtyId,
clientGeneration: params.clientGeneration,
ownerGeneration: params.ownerGeneration,
deliveryToken: params.deliveryToken
})
.catch(() => {})
} catch {}
return Object.freeze({
deliveryToken: params.deliveryToken,
clientGeneration: params.clientGeneration,
ownerGeneration: params.ownerGeneration,
ptyIncarnation: params.ptyIncarnation
})
}
function positiveSafeInteger(value: unknown): value is number {

View File

@ -27,6 +27,8 @@ export type SshPtyDataCallback = (payload: {
sourceEndSu: number
}>
sourceMalformed?: boolean
sourceRejected?: boolean
rejectedSourceRecovery?: 'confirm-existing' | 'fresh-activation' | 'reconnect-channel'
}) => void
export type SshPtyReplayCallback = (payload: { id: string; data: string }) => void
export type SshPtyExitCallback = (payload: {

View File

@ -14,6 +14,7 @@ import type { PtySourceReceivingActivation } from '../../shared/pty-source-recei
export class SshPtyProviderOutputState {
private readonly dataListeners = new Set<SshPtyDataCallback>()
private readonly rejectedDataListeners = new Set<SshPtyDataCallback>()
private readonly replayListeners = new Set<SshPtyReplayCallback>()
private readonly exitListeners = new Set<SshPtyExitCallback>()
private readonly incarnationByRelayPtyId = new Map<string, string>()
@ -34,11 +35,13 @@ export class SshPtyProviderOutputState {
this.subscription = subscribeSshPtyNotifications({
...args,
dataListeners: this.dataListeners,
rejectedDataListeners: this.rejectedDataListeners,
replayListeners: this.replayListeners,
exitListeners: this.exitListeners,
providerGeneration,
resolvePtyIncarnation: (relayPtyId, incarnationId) =>
this.resolvePtyIncarnation(relayPtyId, incarnationId),
peekPtyIncarnation: (relayPtyId) => this.incarnationByRelayPtyId.get(relayPtyId),
recordExit: (relayPtyId, incarnationId) => {
args.recordExit(relayPtyId, incarnationId)
this.incarnationByRelayPtyId.delete(relayPtyId)
@ -52,6 +55,7 @@ export class SshPtyProviderOutputState {
this.subscription?.dispose()
this.subscription = null
this.dataListeners.clear()
this.rejectedDataListeners.clear()
this.replayListeners.clear()
this.exitListeners.clear()
this.incarnationByRelayPtyId.clear()
@ -63,6 +67,11 @@ export class SshPtyProviderOutputState {
return () => this.dataListeners.delete(callback)
}
onRejectedData(callback: SshPtyDataCallback): () => void {
this.rejectedDataListeners.add(callback)
return () => this.rejectedDataListeners.delete(callback)
}
onReplay(callback: SshPtyReplayCallback): () => void {
this.replayListeners.add(callback)
return () => this.replayListeners.delete(callback)

View File

@ -314,6 +314,8 @@ export class SshPtyProvider implements IPtyProvider {
}
onData = (callback: SshPtyDataCallback): (() => void) => this.outputState.onData(callback)
onRejectedData = (callback: SshPtyDataCallback): (() => void) =>
this.outputState.onRejectedData(callback)
onReplay = (callback: SshPtyReplayCallback): (() => void) => this.outputState.onReplay(callback)
onExit = (callback: SshPtyExitCallback): (() => void) => this.outputState.onExit(callback)

View File

@ -2,6 +2,17 @@ import { describe, expect, it, vi } from 'vitest'
import { SshPtySourceDeliveryLedger } from './ssh-pty-source-delivery-ledger'
describe('SshPtySourceDeliveryLedger', () => {
const activation = (deliveryToken: string) =>
Object.freeze({
status: 'pending' as const,
clientGeneration: 2,
ownerGeneration: 3,
ptyIncarnation: 'incarnation-1',
deliveryToken,
checkpointSourceEndSu: 0,
recoveryEndSu: 0
})
it('retains cancellation ownership when recovery transfer is superseded', async () => {
const request = vi.fn(async () => ({ canceled: true, sentEndSu: 0, creditedEndSu: 0 }))
const ledger = new SshPtySourceDeliveryLedger({ request } as never, vi.fn())
@ -41,4 +52,44 @@ describe('SshPtySourceDeliveryLedger', () => {
deliveryToken: 'token-old'
})
})
it('retires an exact rejected activation only after cancellation is proven', async () => {
const request = vi.fn(async () => ({ canceled: true, sentEndSu: 4, creditedEndSu: 0 }))
const ledger = new SshPtySourceDeliveryLedger({ request } as never, vi.fn())
ledger.install('pty-1', activation('token-old')).commit()
await expect(ledger.reject('pty-1', activation('token-old'))).resolves.toBe('fresh-activation')
expect(() => ledger.install('pty-1', activation('token-new')).commit()).not.toThrow()
})
it('keeps the current activation when a stale rejected token cannot cancel it', async () => {
const publish = vi.fn()
const request = vi.fn(async () => {
throw new Error('Unknown or stale PTY source delivery cancellation')
})
const ledger = new SshPtySourceDeliveryLedger({ request } as never, publish)
ledger.install('pty-1', activation('token-current')).commit()
await expect(ledger.reject('pty-1', activation('token-stale'))).resolves.toBe(
'confirm-existing'
)
expect(
ledger.admit({
relayPtyId: 'pty-1',
params: { ptyIncarnation: 'incarnation-1' },
data: 'healthy',
source: {
relayPtyId: 'pty-1',
spanId: 'span-1',
clientGeneration: 2,
ownerGeneration: 3,
deliveryToken: 'token-current',
sourceStartSu: 0,
sourceEndSu: 7
}
})
).toBe(true)
expect(publish).toHaveBeenCalledWith(expect.objectContaining({ data: 'healthy' }))
})
})

View File

@ -1,40 +1,21 @@
import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer'
import type { PtySourceReceivingActivation } from '../../shared/pty-source-receiving-activation'
import type { SshPtySourceFrame } from './ssh-pty-source-frame'
export type PendingSshPtySourceData = Readonly<{
relayPtyId: string
params: Record<string, unknown>
data: string
source?: SshPtySourceFrame
}>
type SourceDeliveryLeaseState = {
phase: 'provisional' | 'recovery' | 'committing' | 'committed' | 'retired'
pendingData: PendingSshPtySourceData[]
recoverySink?: (pending: PendingSshPtySourceData) => void
exited: boolean
}
type SourceDeliveryState = Readonly<{
activation: PtySourceReceivingActivation
sourceEndSu: number
lease: SourceDeliveryLeaseState
previous?: SourceDeliveryState
}>
export type SshPtyRecoveryActivationLease = Readonly<{
commit: () => void
retire: () => void
}>
export type SshPtySourceDeliveryLease = Readonly<{
commit: () => void
rollback: () => Promise<boolean>
transferToRecovery: (
sink: (pending: PendingSshPtySourceData) => void
) => SshPtyRecoveryActivationLease
}>
import {
acceptsSourceFrame,
activePredecessor,
sameReceivingActivation,
sameRejectedSourceIdentity,
settleExactSourceDeliveryCancellation,
settledReceivingActivationLease,
type PendingSshPtySourceData,
type RejectedSourceIdentity,
type SourceDeliveryLeaseState,
type SourceDeliveryState,
type SshPtyRecoveryActivationLease,
type SshPtyRejectedSourceRecovery,
type SshPtySourceDeliveryLease
} from './ssh-pty-source-delivery-state'
export class SshPtySourceDeliveryLedger {
private readonly deliveryByPty = new Map<string, SourceDeliveryState>()
@ -95,6 +76,32 @@ export class SshPtySourceDeliveryLedger {
}
}
async reject(
relayPtyId: string,
identity: RejectedSourceIdentity | undefined
): Promise<SshPtyRejectedSourceRecovery> {
if (!identity) {
return 'reconnect-channel'
}
const offeredCurrent = this.deliveryByPty.get(relayPtyId)
const matched = Boolean(
offeredCurrent && sameRejectedSourceIdentity(offeredCurrent.activation, identity)
)
const canceled = await settleExactSourceDeliveryCancellation(this.mux, relayPtyId, identity)
if (matched && canceled) {
const current = this.deliveryByPty.get(relayPtyId)
if (current && sameRejectedSourceIdentity(current.activation, identity)) {
this.retire(relayPtyId, current.previous, current.lease)
return 'fresh-activation'
}
return current ? 'reconnect-channel' : 'fresh-activation'
}
if (!matched && !canceled) {
return 'confirm-existing'
}
return 'reconnect-channel'
}
private installProvisional(
relayPtyId: string,
activation: PtySourceReceivingActivation,
@ -245,76 +252,3 @@ export class SshPtySourceDeliveryLedger {
}
}
}
function settledReceivingActivationLease(): SshPtySourceDeliveryLease {
return Object.freeze({
commit: () => {},
rollback: async () => true,
transferToRecovery: () => Object.freeze({ commit: () => {}, retire: () => {} })
})
}
function activePredecessor(previous?: SourceDeliveryState): SourceDeliveryState | undefined {
while (previous?.lease.phase === 'retired') {
previous = previous.previous
}
return previous
}
function sameReceivingActivation(
left: PtySourceReceivingActivation,
right: PtySourceReceivingActivation
): boolean {
return (
left.clientGeneration === right.clientGeneration &&
left.ownerGeneration === right.ownerGeneration &&
left.ptyIncarnation === right.ptyIncarnation &&
left.deliveryToken === right.deliveryToken &&
left.checkpointSourceEndSu === right.checkpointSourceEndSu &&
left.recoveryEndSu === right.recoveryEndSu
)
}
function acceptsSourceFrame(
current: SourceDeliveryState | undefined,
params: Record<string, unknown>,
source: SshPtySourceFrame
): current is SourceDeliveryState {
return Boolean(
current &&
current.lease.phase !== 'retired' &&
!current.lease.exited &&
current.activation.ptyIncarnation === params.ptyIncarnation &&
current.activation.deliveryToken === source.deliveryToken &&
current.activation.clientGeneration === source.clientGeneration &&
current.activation.ownerGeneration === source.ownerGeneration &&
current.sourceEndSu === source.sourceStartSu
)
}
async function settleExactSourceDeliveryCancellation(
mux: SshChannelMultiplexer,
relayPtyId: string,
activation: PtySourceReceivingActivation
): Promise<boolean> {
try {
const result = (await mux.request('pty.cancelDelivery', {
id: relayPtyId,
clientGeneration: activation.clientGeneration,
ownerGeneration: activation.ownerGeneration,
deliveryToken: activation.deliveryToken
})) as Record<string, unknown>
return (
result.canceled === true &&
nonNegativeSafeInteger(result.sentEndSu) &&
nonNegativeSafeInteger(result.creditedEndSu) &&
result.creditedEndSu <= result.sentEndSu
)
} catch {
return false
}
}
function nonNegativeSafeInteger(value: unknown): value is number {
return Number.isSafeInteger(value) && (value as number) >= 0
}

View File

@ -0,0 +1,134 @@
import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer'
import type { PtySourceReceivingActivation } from '../../shared/pty-source-receiving-activation'
import type { SshPtySourceFrame } from './ssh-pty-source-frame'
export type PendingSshPtySourceData = Readonly<{
relayPtyId: string
params: Record<string, unknown>
data: string
source?: SshPtySourceFrame
}>
export type SourceDeliveryLeaseState = {
phase: 'provisional' | 'recovery' | 'committing' | 'committed' | 'retired'
pendingData: PendingSshPtySourceData[]
recoverySink?: (pending: PendingSshPtySourceData) => void
exited: boolean
}
export type SourceDeliveryState = Readonly<{
activation: PtySourceReceivingActivation
sourceEndSu: number
lease: SourceDeliveryLeaseState
previous?: SourceDeliveryState
}>
export type SshPtyRecoveryActivationLease = Readonly<{
commit: () => void
retire: () => void
}>
export type SshPtySourceDeliveryLease = Readonly<{
commit: () => void
rollback: () => Promise<boolean>
transferToRecovery: (
sink: (pending: PendingSshPtySourceData) => void
) => SshPtyRecoveryActivationLease
}>
export type SshPtyRejectedSourceRecovery =
| 'confirm-existing'
| 'fresh-activation'
| 'reconnect-channel'
export type RejectedSourceIdentity = Readonly<{
clientGeneration: number
ownerGeneration: number
ptyIncarnation: string
deliveryToken: string
}>
export function settledReceivingActivationLease(): SshPtySourceDeliveryLease {
return Object.freeze({
commit: () => {},
rollback: async () => true,
transferToRecovery: () => Object.freeze({ commit: () => {}, retire: () => {} })
})
}
export function activePredecessor(previous?: SourceDeliveryState): SourceDeliveryState | undefined {
while (previous?.lease.phase === 'retired') {
previous = previous.previous
}
return previous
}
export function sameReceivingActivation(
left: PtySourceReceivingActivation,
right: PtySourceReceivingActivation
): boolean {
return (
left.clientGeneration === right.clientGeneration &&
left.ownerGeneration === right.ownerGeneration &&
left.ptyIncarnation === right.ptyIncarnation &&
left.deliveryToken === right.deliveryToken &&
left.checkpointSourceEndSu === right.checkpointSourceEndSu &&
left.recoveryEndSu === right.recoveryEndSu
)
}
export function sameRejectedSourceIdentity(
left: PtySourceReceivingActivation,
right: RejectedSourceIdentity
): boolean {
return (
left.clientGeneration === right.clientGeneration &&
left.ownerGeneration === right.ownerGeneration &&
left.ptyIncarnation === right.ptyIncarnation &&
left.deliveryToken === right.deliveryToken
)
}
export function acceptsSourceFrame(
current: SourceDeliveryState | undefined,
params: Record<string, unknown>,
source: SshPtySourceFrame
): current is SourceDeliveryState {
return Boolean(
current &&
current.lease.phase !== 'retired' &&
!current.lease.exited &&
current.activation.ptyIncarnation === params.ptyIncarnation &&
current.activation.deliveryToken === source.deliveryToken &&
current.activation.clientGeneration === source.clientGeneration &&
current.activation.ownerGeneration === source.ownerGeneration &&
current.sourceEndSu === source.sourceStartSu
)
}
export async function settleExactSourceDeliveryCancellation(
mux: SshChannelMultiplexer,
relayPtyId: string,
activation: RejectedSourceIdentity
): Promise<boolean> {
try {
const result = (await mux.request('pty.cancelDelivery', {
id: relayPtyId,
clientGeneration: activation.clientGeneration,
ownerGeneration: activation.ownerGeneration,
deliveryToken: activation.deliveryToken
})) as Record<string, unknown>
return (
result.canceled === true &&
nonNegativeSafeInteger(result.sentEndSu) &&
nonNegativeSafeInteger(result.creditedEndSu) &&
result.creditedEndSu <= result.sentEndSu
)
} catch {
return false
}
}
function nonNegativeSafeInteger(value: unknown): value is number {
return Number.isSafeInteger(value) && (value as number) >= 0
}

View File

@ -0,0 +1,115 @@
import type { SshPtyDataCallback } from '../providers/ssh-pty-provider-contract'
import type { SshPtyConsumerOwnerState } from './ssh-pty-consumer-session'
type SshPtyDataPayload = Parameters<SshPtyDataCallback>[0]
const SSH_PTY_FRAME_REJECTION_LOG_KEY_LIMIT = 12
export type SshPtyFrameRejectionReason =
| 'source-missing'
| 'source-malformed'
| 'client-generation-stale'
| 'owner-generation-stale'
| 'client-and-owner-generation-stale'
| 'source-range-invalid'
export type SshPtyFrameRejection = Readonly<{
reason: SshPtyFrameRejectionReason
action: 'quarantine-legacy-frame' | 'retire-and-reattach-delivery'
}>
// Why a source-less frame is quarantined rather than reattached: it is the signature of legacy
// output already in flight when the flow-control grant landed, not of a broken delivery. The relay's
// own admission stops publishing those once the grant is live, and source recovery replays from
// sourceStartSu, so dropping them loses nothing. Reattaching instead would restart the PTY on every
// openClient handshake.
export function classifySshPtyFrameRejection(
payload: SshPtyDataPayload,
owner: SshPtyConsumerOwnerState | null
): SshPtyFrameRejection | null {
if (payload.sourceMalformed) {
return Object.freeze({
reason: 'source-malformed',
action: 'retire-and-reattach-delivery'
})
}
const source = payload.source
if (!source) {
return owner?.outputFlowControl
? Object.freeze({
reason: 'source-missing',
action: 'quarantine-legacy-frame'
})
: null
}
if (!owner?.outputFlowControl) {
return payload.sourceRejected
? Object.freeze({
reason: 'source-range-invalid',
action: 'retire-and-reattach-delivery'
})
: null
}
const staleClient = source.clientGeneration !== owner.clientGeneration
const staleOwner = source.ownerGeneration !== owner.ownerGeneration
if (staleClient || staleOwner) {
return Object.freeze({
reason:
staleClient && staleOwner
? 'client-and-owner-generation-stale'
: staleClient
? 'client-generation-stale'
: 'owner-generation-stale',
action: 'retire-and-reattach-delivery'
})
}
if (payload.sourceRejected) {
return Object.freeze({
reason: 'source-range-invalid',
action: 'retire-and-reattach-delivery'
})
}
return null
}
/**
* Why bounded: a quarantined frame carries no delivery token to retire, so nothing dedupes it a
* relay that keeps publishing legacy output to a flow-control owner would log once per frame and
* bury every other diagnostic. One line per PTY and reason per generation names the fault; the
* generation is itself the counter for how often it recurs.
*/
export class SshPtyFrameRejectionLog {
private generation: number | null = null
private readonly loggedKeys = new Set<string>()
record(
payload: SshPtyDataPayload,
owner: SshPtyConsumerOwnerState | null,
rejection: SshPtyFrameRejection
): void {
if (this.generation !== payload.providerGeneration) {
this.generation = payload.providerGeneration
this.loggedKeys.clear()
}
const key = `${payload.id}\0${rejection.reason}`
if (this.loggedKeys.has(key) || this.loggedKeys.size >= SSH_PTY_FRAME_REJECTION_LOG_KEY_LIMIT) {
return
}
this.loggedKeys.add(key)
console.warn('[ssh-relay-session] Rejected PTY delivery', {
ptyId: payload.id,
providerGeneration: payload.providerGeneration,
expectedClientGeneration: owner?.clientGeneration ?? null,
expectedOwnerGeneration: owner?.ownerGeneration ?? null,
offeredClientGeneration: payload.source?.clientGeneration ?? null,
offeredOwnerGeneration: payload.source?.ownerGeneration ?? null,
sourceState: rejection.reason,
recoveryAction: rejection.action
})
}
clear(): void {
this.generation = null
this.loggedKeys.clear()
}
}

View File

@ -0,0 +1,57 @@
type TargetedReattach = Readonly<{ start: () => void }>
// Why bounded: every rejected frame asks for its own attach round trip, so a relay that starts
// rejecting across many PTYs at once fans out one concurrent reattach per PTY. The bulk reconnect
// path caps the identical work, and each attach also costs a lease read and a lease write.
export class SshPtyTargetedReattachQueue {
private readonly running = new Map<string, TargetedReattach>()
private readonly waiting: TargetedReattach[] = []
private active = 0
constructor(private readonly maxConcurrency: number) {}
has(key: string): boolean {
return this.running.has(key)
}
run(key: string, task: () => Promise<boolean>): Promise<boolean> {
return new Promise<boolean>((resolve, reject) => {
const entry: TargetedReattach = {
start: () => {
this.active++
task().then(
(recovered) => {
this.settle(key, entry)
resolve(recovered)
},
(error: unknown) => {
this.settle(key, entry)
reject(error instanceof Error ? error : new Error(String(error)))
}
)
}
}
this.running.set(key, entry)
if (this.active < this.maxConcurrency) {
entry.start()
return
}
this.waiting.push(entry)
})
}
// Why queued entries are dropped rather than started: each one is keyed to the provider generation
// the teardown just ended, so running it would attach onto a mux that is already gone.
clear(): void {
this.waiting.length = 0
this.running.clear()
}
private settle(key: string, entry: TargetedReattach): void {
if (this.running.get(key) === entry) {
this.running.delete(key)
}
this.active--
this.waiting.shift()?.start()
}
}

View File

@ -519,7 +519,7 @@ describe('SshRelaySession data delivery', () => {
)
})
it('rejects missing negotiated source identity before main admission', async () => {
it('quarantines missing negotiated source identity before main admission', async () => {
const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps()
const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward)
await session.establish(mockConn)
@ -536,10 +536,8 @@ describe('SshRelaySession data delivery', () => {
})
expect(acceptOutputDataMock).not.toHaveBeenCalled()
expect(closeSshPtyOutputGeneration).toHaveBeenCalledWith(
23,
'ssh_source_frame_malformed_or_missing'
)
expect(closeSshPtyOutputGeneration).not.toHaveBeenCalled()
expect(muxDisposeMock).not.toHaveBeenCalled()
})
it('keeps unoffered source metadata out of legacy intake', async () => {

View File

@ -420,6 +420,7 @@ describe('SshRelaySession recovery race fencing', () => {
})
await session.reconnect(deps.mockConn)
const closeCount = vi.mocked(closeSshPtyOutputGeneration).mock.calls.length
const muxDisposeCount = muxDisposeMock.mock.calls.length
emitSourceFrame({
targetId,
@ -431,11 +432,9 @@ describe('SshRelaySession recovery race fencing', () => {
})
expect(acceptOutputDataMock).not.toHaveBeenCalled()
expect(closeSshPtyOutputGeneration).toHaveBeenCalledTimes(closeCount + 1)
expect(closeSshPtyOutputGeneration).toHaveBeenLastCalledWith(
23,
'ssh_source_frame_stale_or_non_contiguous'
)
expect(closeSshPtyOutputGeneration).toHaveBeenCalledTimes(closeCount)
expect(muxDisposeMock).toHaveBeenCalledTimes(muxDisposeCount)
await vi.waitFor(() => expect(attachForReconnectMock).toHaveBeenCalledTimes(2))
})
it('drops late frames from a token after its cancellation proof is validated', async () => {

View File

@ -0,0 +1,496 @@
import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { SshPtyDataCallback } from '../providers/ssh-pty-provider-contract'
import type { SshPtyProvider } from '../providers/ssh-pty-provider'
import type { SshChannelMultiplexer } from './ssh-channel-multiplexer'
import type { SshPtyConsumerSessionState } from './ssh-pty-consumer-session'
import { SshRelaySession } from './ssh-relay-session'
import { createMockDeps } from './ssh-relay-session-test-fixtures'
const { acceptOutputDataMock, getSshPtyProviderMock } = vi.hoisted(() => ({
acceptOutputDataMock: vi.fn().mockResolvedValue(undefined),
getSshPtyProviderMock: vi.fn()
}))
vi.mock('../ipc/ssh-pty-output-intake-registry', async (importOriginal) => {
const original = (await importOriginal()) as object
return { ...original, acceptSshPtyOutputData: acceptOutputDataMock }
})
vi.mock('../ipc/pty', async (importOriginal) => {
const original = (await importOriginal()) as object
return { ...original, getSshPtyProvider: getSshPtyProviderMock }
})
type SshPtyDataPayload = Parameters<SshPtyDataCallback>[0]
type RejectedDeliverySession = {
mux: SshChannelMultiplexer | null
activePtyProviderGeneration: number | null
ptyConsumerSessionState: SshPtyConsumerSessionState | null
acceptPtyData: (payload: SshPtyDataPayload) => Promise<unknown>
reattachKnownPty: (args: {
ptyId: string
activeLeaseByPtyId: Map<string, unknown>
expectedIdentityByPtyId: Map<string, unknown>
attachedLeaseIds: Set<string>
targetedDeliveryRecovery?: 'confirm-existing' | 'fresh-activation'
}) => Promise<void>
reattachRejectedPty: (
relayPtyId: string,
mux: SshChannelMultiplexer,
providerGeneration: number,
targetedDeliveryRecovery: 'confirm-existing' | 'fresh-activation'
) => Promise<boolean>
sourceRecoveryRequest: (appPtyId: string) => Promise<
| {
status: 'checkpoint'
clientGeneration: number
ownerGeneration: number
ptyIncarnation: string
deliveryToken: string
acceptedSourceEndSu: number
}
| undefined
>
rejectedPtyRecoveryAttempts: Map<string, unknown>
sourceIdentityByRelayPtyId: Map<string, unknown>
retireExitedPty: (payload: {
id: string
code: number
providerGeneration: number
ptyIncarnation: string
}) => void
}
function source(overrides: Partial<NonNullable<SshPtyDataPayload['source']>> = {}) {
return {
relayPtyId: 'pty-bad',
spanId: 'token-bad:0:4',
clientGeneration: 1,
ownerGeneration: 1,
deliveryToken: 'token-bad',
sourceStartSu: 0,
sourceEndSu: 4,
...overrides
}
}
function rejectedPayload(overrides: Partial<SshPtyDataPayload> = {}): SshPtyDataPayload {
return {
id: 'ssh:target-1@@pty-bad',
data: 'rejected',
providerGeneration: 23,
ptyIncarnation: 'incarnation-bad',
source: source(),
sourceRejected: true,
...overrides
}
}
function prepareSession() {
const deps = createMockDeps()
const mux = {
isDisposed: vi.fn(() => false),
dispose: vi.fn()
} as unknown as SshChannelMultiplexer
const session = new SshRelaySession(
'target-1',
deps.getMainWindow,
deps.mockStore,
deps.mockPortForward
)
const internals = session as unknown as RejectedDeliverySession
internals.mux = mux
internals.activePtyProviderGeneration = 23
internals.ptyConsumerSessionState = {
mode: 'negotiated',
clientInstanceId: 'client-1',
clientGeneration: 1,
ownerGeneration: 1,
ownerLease: 'owner-lease',
outputFlowControl: { version: 1, windowSu: 64 }
}
return { deps, internals, mux, session }
}
describe('SshRelaySession rejected PTY delivery recovery', () => {
beforeEach(() => {
vi.clearAllMocks()
})
it('quarantines source-less output queued before a flow-control grant', async () => {
const { internals, mux } = prepareSession()
const reattach = vi.fn().mockResolvedValue(true)
internals.reattachRejectedPty = reattach
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
await internals.acceptPtyData(rejectedPayload({ source: undefined, sourceRejected: undefined }))
expect(warn).toHaveBeenCalledWith(
'[ssh-relay-session] Rejected PTY delivery',
expect.objectContaining({
sourceState: 'source-missing',
recoveryAction: 'quarantine-legacy-frame'
})
)
expect(reattach).not.toHaveBeenCalled()
expect(acceptOutputDataMock).not.toHaveBeenCalled()
expect(mux.dispose).not.toHaveBeenCalled()
})
it.each([
['client', 9, 1],
['owner', 1, 9]
] as const)(
'retires a superseded %s generation without disposing the mux',
async (_generation, clientGeneration, ownerGeneration) => {
const { internals, mux } = prepareSession()
const reattach = vi.fn().mockResolvedValue(true)
internals.reattachRejectedPty = reattach
vi.spyOn(console, 'warn').mockImplementation(() => {})
await internals.acceptPtyData(
rejectedPayload({ source: source({ clientGeneration, ownerGeneration }) })
)
expect(reattach).toHaveBeenCalledWith('pty-bad', mux, 23, 'confirm-existing')
expect(acceptOutputDataMock).not.toHaveBeenCalled()
expect(mux.dispose).not.toHaveBeenCalled()
}
)
it('logs malformed source metadata and reattaches from the exact outer PTY identity', async () => {
const { internals } = prepareSession()
const reattach = vi.fn().mockResolvedValue(true)
internals.reattachRejectedPty = reattach
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
await internals.acceptPtyData(
rejectedPayload({ source: undefined, sourceMalformed: true, sourceRejected: undefined })
)
expect(warn).toHaveBeenCalledWith(
'[ssh-relay-session] Rejected PTY delivery',
expect.objectContaining({
ptyId: 'ssh:target-1@@pty-bad',
providerGeneration: 23,
expectedClientGeneration: 1,
expectedOwnerGeneration: 1,
sourceState: 'source-malformed',
recoveryAction: 'retire-and-reattach-delivery'
})
)
expect(reattach).toHaveBeenCalledWith('pty-bad', expect.anything(), 23, 'confirm-existing')
expect(acceptOutputDataMock).not.toHaveBeenCalled()
})
it('reattaches only the rejected PTY while a healthy sibling keeps delivering', async () => {
const { deps, internals, mux } = prepareSession()
const provider = {} as SshPtyProvider
getSshPtyProviderMock.mockReturnValue(provider)
vi.mocked(deps.mockStore.getSshRemotePtyLeases).mockReturnValue([
{
targetId: 'target-1',
ptyId: 'pty-bad',
state: 'detached',
createdAt: 1,
updatedAt: 1
},
{
targetId: 'target-1',
ptyId: 'pty-healthy',
state: 'attached',
createdAt: 1,
updatedAt: 1
}
])
const reattachKnownPty = vi.fn(
async (args: Parameters<RejectedDeliverySession['reattachKnownPty']>[0]) => {
args.attachedLeaseIds.add(args.ptyId)
}
)
internals.reattachKnownPty = reattachKnownPty
await internals.acceptPtyData(rejectedPayload({ rejectedSourceRecovery: 'fresh-activation' }))
await vi.waitFor(() => expect(reattachKnownPty).toHaveBeenCalledOnce())
await internals.acceptPtyData(
rejectedPayload({
id: 'ssh:target-1@@pty-healthy',
data: 'healthy',
ptyIncarnation: 'incarnation-healthy',
source: source({
relayPtyId: 'pty-healthy',
spanId: 'token-healthy:0:7',
deliveryToken: 'token-healthy',
sourceEndSu: 7
}),
sourceRejected: undefined
})
)
const recovery = reattachKnownPty.mock.calls[0]?.[0]
expect(recovery?.ptyId).toBe('pty-bad')
expect(Array.from(recovery?.activeLeaseByPtyId.keys() ?? [])).toEqual(['pty-bad'])
expect(Array.from(recovery?.expectedIdentityByPtyId.keys() ?? [])).toEqual([])
expect(recovery?.targetedDeliveryRecovery).toBe('fresh-activation')
expect(deps.mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith('target-1', [
'pty-bad'
])
expect(acceptOutputDataMock).toHaveBeenCalledOnce()
expect(acceptOutputDataMock).toHaveBeenCalledWith(
expect.objectContaining({ id: 'ssh:target-1@@pty-healthy', data: 'healthy' })
)
expect(mux.dispose).not.toHaveBeenCalled()
})
it('reconnects instead of canceling an unprovable malformed delivery', async () => {
const { internals, mux } = prepareSession()
const reattach = vi.fn().mockResolvedValue(true)
internals.reattachRejectedPty = reattach
vi.spyOn(console, 'warn').mockImplementation(() => {})
await internals.acceptPtyData(
rejectedPayload({
source: undefined,
sourceMalformed: true,
sourceRejected: undefined,
rejectedSourceRecovery: 'reconnect-channel'
})
)
expect(mux.dispose).toHaveBeenCalledWith('connection_lost')
expect(reattach).not.toHaveBeenCalled()
})
it('opens a fresh activation after exact rejected-delivery cancellation', async () => {
const { deps, internals, mux } = prepareSession()
internals.sourceIdentityByRelayPtyId.set('pty-bad', {
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-bad',
deliveryToken: 'token-old',
nextSourceSu: 4
})
const commit = vi.fn()
const attachForReconnect = vi.fn(async () => ({
incarnationId: 'incarnation-bad',
sourceActivation: {
status: 'pending' as const,
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-bad',
deliveryToken: 'token-fresh',
checkpointSourceEndSu: 0,
recoveryEndSu: 0
},
sourceActivationLease: { commit, rollback: vi.fn(async () => true) }
}))
getSshPtyProviderMock.mockReturnValue({ attachForReconnect } as unknown as SshPtyProvider)
await expect(
internals.reattachRejectedPty('pty-bad', mux, 23, 'fresh-activation')
).resolves.toBe(true)
expect(attachForReconnect).toHaveBeenCalledWith('pty-bad')
expect(commit).toHaveBeenCalledOnce()
expect(deps.mockStore.markSshRemotePtyLeasesAttachedAsync).toHaveBeenCalledWith('target-1', [
'pty-bad'
])
await internals.acceptPtyData(
rejectedPayload({
data: 'fresh',
sourceRejected: undefined,
source: source({ deliveryToken: 'token-fresh' })
})
)
expect(acceptOutputDataMock).toHaveBeenCalledWith(expect.objectContaining({ data: 'fresh' }))
expect(mux.dispose).not.toHaveBeenCalled()
})
it('accepts an exact existing activation as stale-frame confirmation', async () => {
const { internals, mux } = prepareSession()
const checkpoint = {
status: 'checkpoint' as const,
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-bad',
deliveryToken: 'token-current',
acceptedSourceEndSu: 7
}
internals.sourceRecoveryRequest = vi.fn(async () => checkpoint)
const commit = vi.fn()
const attachForReconnect = vi.fn(async () => ({
incarnationId: 'incarnation-bad',
sourceActivation: {
status: 'pending' as const,
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-bad',
deliveryToken: 'token-current',
checkpointSourceEndSu: 0,
recoveryEndSu: 0
},
sourceActivationLease: { commit, rollback: vi.fn(async () => true) }
}))
getSshPtyProviderMock.mockReturnValue({ attachForReconnect } as unknown as SshPtyProvider)
await expect(
internals.reattachRejectedPty('pty-bad', mux, 23, 'confirm-existing')
).resolves.toBe(true)
expect(attachForReconnect).toHaveBeenCalledWith('pty-bad', undefined, checkpoint)
expect(commit).toHaveBeenCalledOnce()
expect(mux.dispose).not.toHaveBeenCalled()
})
it('forgets rejected-delivery recovery history when the PTY exits', () => {
const { internals } = prepareSession()
internals.rejectedPtyRecoveryAttempts.set('ssh:target-1@@pty-bad', {})
internals.retireExitedPty({
id: 'ssh:target-1@@pty-bad',
code: 0,
providerGeneration: 23,
ptyIncarnation: 'incarnation-bad'
})
expect(internals.rejectedPtyRecoveryAttempts).toHaveLength(0)
})
// Why a channel drop rather than a terminal relay error: a terminal error clears the reconnect
// backoff and rotates provider authority, aborting every fs and git request on the target, so one
// PTY's undeliverable output would strand the whole connection in manual recovery.
it('bounds failed targeted recovery and escalates to a recoverable relay reconnect', async () => {
const { internals, mux, session } = prepareSession()
const reattach = vi.fn().mockResolvedValue(false)
internals.reattachRejectedPty = reattach
getSshPtyProviderMock.mockReturnValue({ hasPty: () => true } as unknown as SshPtyProvider)
const onTerminalError = vi.fn()
session.setOnTerminalRelayError(onTerminalError)
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
await internals.acceptPtyData(rejectedPayload())
await vi.waitFor(() => expect(mux.dispose).toHaveBeenCalledOnce(), { timeout: 2000 })
expect(reattach).toHaveBeenCalledTimes(2)
expect(reattach.mock.calls).toEqual([
['pty-bad', mux, 23, 'confirm-existing'],
['pty-bad', mux, 23, 'confirm-existing']
])
expect(mux.dispose).toHaveBeenCalledWith('connection_lost')
expect(warn).toHaveBeenCalledWith(
expect.stringContaining('PTY pty-bad delivery recovery exhausted')
)
expect(onTerminalError).not.toHaveBeenCalled()
expect(acceptOutputDataMock).not.toHaveBeenCalled()
})
it('stops without an error when the rejected PTY exited during recovery', async () => {
const { internals, mux, session } = prepareSession()
const reattach = vi.fn().mockResolvedValue(false)
internals.reattachRejectedPty = reattach
// Why: reattachKnownPty resolves without claiming the lease when the PTY exits mid-attach, so a
// plain "not recovered" is indistinguishable from a failure until liveness is checked.
getSshPtyProviderMock.mockReturnValue({ hasPty: () => false } as unknown as SshPtyProvider)
const onTerminalError = vi.fn()
session.setOnTerminalRelayError(onTerminalError)
await internals.acceptPtyData(rejectedPayload())
await vi.waitFor(() => expect(reattach).toHaveBeenCalledOnce())
await new Promise((resolve) => setTimeout(resolve, 400))
expect(reattach).toHaveBeenCalledOnce()
expect(onTerminalError).not.toHaveBeenCalled()
expect(mux.dispose).not.toHaveBeenCalled()
})
// Why bounded: the bulk reconnect path caps the identical attach work at 8, and a relay that
// starts rejecting across every PTY at once would otherwise open one attach round trip per PTY.
it('caps concurrent targeted reattaches and coalesces repeats for one PTY', async () => {
const { internals } = prepareSession()
getSshPtyProviderMock.mockReturnValue({ hasPty: () => true } as unknown as SshPtyProvider)
vi.spyOn(console, 'warn').mockImplementation(() => {})
const release: (() => void)[] = []
let active = 0
let peak = 0
const reattach = vi.fn(async () => {
active++
peak = Math.max(peak, active)
await new Promise<void>((resolve) => release.push(resolve))
active--
return true
})
internals.reattachRejectedPty = reattach
for (let index = 0; index < 20; index++) {
const relayPtyId = `pty-${index}`
await internals.acceptPtyData(
rejectedPayload({
id: `ssh:target-1@@${relayPtyId}`,
source: source({ relayPtyId, deliveryToken: `token-${index}` })
})
)
// Why twice: a repeat for a PTY already recovering must not consume a second slot.
await internals.acceptPtyData(
rejectedPayload({
id: `ssh:target-1@@${relayPtyId}`,
source: source({ relayPtyId, deliveryToken: `token-${index}-repeat` })
})
)
}
expect(peak).toBe(8)
expect(reattach).toHaveBeenCalledTimes(8)
for (const resolve of release.splice(0)) {
resolve()
}
await vi.waitFor(() => expect(reattach).toHaveBeenCalledTimes(16))
expect(peak).toBe(8)
for (const resolve of release.splice(0)) {
resolve()
}
})
it('does not let accepted frames refill the recovery budget indefinitely', async () => {
const { internals, mux, session } = prepareSession()
const reattach = vi.fn().mockResolvedValue(true)
internals.reattachRejectedPty = reattach
getSshPtyProviderMock.mockReturnValue({ hasPty: () => true } as unknown as SshPtyProvider)
const onTerminalError = vi.fn()
session.setOnTerminalRelayError(onTerminalError)
// Why alternating, with a fresh bad token each round: every rejection retires its own delivery,
// so a flapping PTY only keeps asking for recovery by moving onto new ones, and the accepted
// frame in between is what used to clear the budget outright. Each reattach here succeeds, so
// the consecutive budget is cleared legitimately too — only the per-generation ceiling can stop
// it.
for (let round = 0; round < 40; round++) {
await internals.acceptPtyData(
rejectedPayload({
source: source({
spanId: `token-bad-${round}:0:4`,
deliveryToken: `token-bad-${round}`,
clientGeneration: 9
})
})
)
await internals.acceptPtyData(
rejectedPayload({
sourceRejected: undefined,
source: source({
spanId: `token-good:${round * 4}:${round * 4 + 4}`,
deliveryToken: 'token-good',
sourceStartSu: round * 4,
sourceEndSu: round * 4 + 4
})
})
)
await Promise.resolve()
}
await vi.waitFor(() => expect(mux.dispose).toHaveBeenCalledOnce(), { timeout: 2000 })
expect(onTerminalError).not.toHaveBeenCalled()
expect(acceptOutputDataMock).toHaveBeenCalledTimes(40)
expect(reattach.mock.calls.length).toBeLessThanOrEqual(12)
})
})

View File

@ -1,5 +1,6 @@
import { vi, type Mock } from 'vitest'
import type { BrowserWindow } from 'electron'
import { PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR } from '../../shared/pty-consumer-session'
import type { SshConnection } from './ssh-connection'
import type { Store } from '../persistence'
import type { SshPortForwardManager } from './ssh-port-forward'
@ -53,3 +54,9 @@ export function mockDeploySuccess(): void {
platform: 'linux-x64'
})
}
export function createMismatchedOwnerRecoveryError(): unknown {
return Object.assign(new Error('Owner recovery lease is stale'), {
code: PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR
})
}

View File

@ -79,6 +79,7 @@ import {
import type { Store } from '../persistence'
import type { OrcaRuntimeService } from '../runtime/orca-runtime'
import { DEFAULT_PTY_SOURCE_WINDOW_SU } from '../../shared/pty-source-credit-contract'
import { PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR } from '../../shared/pty-consumer-session'
import {
isSshOwnerAdmissionBlocked,
retrySshOwnerRecoveryWhileBlocked
@ -121,6 +122,8 @@ import {
getSshPtyConsumerRecovery,
rememberSshPtyConsumerRecovery
} from './ssh-pty-consumer-recovery'
import { classifySshPtyFrameRejection, SshPtyFrameRejectionLog } from './ssh-pty-frame-rejection'
import { SshPtyTargetedReattachQueue } from './ssh-pty-targeted-reattach-queue'
export type RelaySessionState = 'idle' | 'deploying' | 'ready' | 'reconnecting' | 'disposed'
@ -131,6 +134,12 @@ const SSH_PTY_REATTACH_MAX_CONCURRENCY = 8
const SSH_PTY_REATTACH_ATTEMPT_TIMEOUT_MS = 10_000
const SSH_PTY_REATTACH_RETRY_MIN_DELAY_MS = 50
const SSH_PTY_REATTACH_RETRY_JITTER_MS = 200
const SSH_REJECTED_PTY_RECOVERY_MAX_ATTEMPTS = 2
// Why a second ceiling: the consecutive budget resets whenever a reattach succeeds, so a PTY that
// alternates recovered and rejected frames would otherwise reattach forever — each one costs a
// store read, an attach round trip and a store write.
const SSH_REJECTED_PTY_RECOVERY_MAX_GENERATION_ATTEMPTS = 12
const SSH_REJECTED_PTY_RECOVERY_RETRY_DELAY_MS = 150
const SSH_SOURCE_RECOVERY_CANCELLATION_FAILED = 'ssh_source_recovery_cancellation_failed'
// Why: superseded attempts stop quietly; a dead mux still owned by this attempt must enter recovery.
@ -179,6 +188,7 @@ type RemoteCliBridgeEnv = {
}
type ExpectedPtyIdentity = { paneKey?: string; tabId?: string }
type TargetedDeliveryRecovery = 'confirm-existing' | 'fresh-activation'
function expectedIdentityForLease(lease: {
tabId?: string
@ -324,6 +334,20 @@ export class SshRelaySession {
}>
>()
private readonly retiredSourceDeliveries = new SshPtyRetiredSourceDeliveries()
private readonly rejectedPtyRecoveryAttempts = new Map<
string,
{
providerGeneration: number
attempts: number
generationAttempts: number
reported: boolean
}
>()
private readonly rejectedPtyRecoveryRetries = new Set<ReturnType<typeof setTimeout>>()
private readonly rejectedPtyReattaches = new SshPtyTargetedReattachQueue(
SSH_PTY_REATTACH_MAX_CONCURRENCY
)
private readonly ptyFrameRejectionLog = new SshPtyFrameRejectionLog()
private readonly ptyConsumerClientInstanceId: string
private ptyConsumerSessionState: SshPtyConsumerSessionState | null = null
private activeCompatibilityAttachmentIds = new Set<string>()
@ -1072,13 +1096,28 @@ export class SshRelaySession {
ownsAttempt: () => boolean
): Promise<SshPtyConsumerSessionState> {
const previousOwner = this.recoverablePtyConsumerOwner(serverBuildId)
const options = {
const options: OpenSshPtyConsumerSessionOptions = {
clientInstanceId: this.ptyConsumerClientInstanceId,
expectedServerBuildId: serverBuildId,
allowSameBuildLegacyFallback: true,
outputFlowControl: { requestedWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU }
}
const admission = await this.admitPtyConsumerOwner(mux, previousOwner, options, ownsAttempt)
let admission: SshPtyConsumerAdmission
try {
admission = await this.admitPtyConsumerOwner(mux, previousOwner, options, ownsAttempt)
} catch (error) {
if (
!previousOwner ||
(error as { code?: unknown }).code !== PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR
) {
throw error
}
this.voidPtyConsumerCheckpoints(previousOwner, ownsAttempt)
if (!ownsAttempt()) {
throw new Error('Session disposed during owner recovery')
}
admission = await openSshPtyConsumerSession(mux, options)
}
if (previousOwner && !admission.resumed) {
this.voidPtyConsumerCheckpoints(previousOwner, ownsAttempt)
}
@ -1535,6 +1574,13 @@ export class SshRelaySession {
unregisterSshGitProvider(this.targetId)
this.sourceIdentityByRelayPtyId.clear()
this.retiredSourceDeliveries.clear()
this.rejectedPtyRecoveryAttempts.clear()
for (const timer of this.rejectedPtyRecoveryRetries) {
clearTimeout(timer)
}
this.rejectedPtyRecoveryRetries.clear()
this.rejectedPtyReattaches.clear()
this.ptyFrameRejectionLog.clear()
for (const pending of this.pendingPtyReattaches.values()) {
for (const resolve of pending.recoveryWaiters) {
resolve()
@ -1637,6 +1683,24 @@ export class SshRelaySession {
}
void this.acceptPtyData(payload).catch(() => {})
})
ptyProvider.onRejectedData?.((payload) => {
if (
this.mux !== mux ||
this.activePtyProviderGeneration !== providerGeneration ||
payload.providerGeneration !== providerGeneration
) {
return
}
const pending = this.pendingPtyReattaches.get(payload.id)
if (pending) {
pending.restoreRequired = payload.sourceMalformed
? 'recoverySourceMalformed'
: 'recoverySourceUnadmitted'
this.wakeRecovery(pending)
return
}
void this.acceptPtyData(payload).catch(() => {})
})
ptyProvider.onReplay((payload) => {
if (this.mux !== mux || this.activePtyProviderGeneration !== providerGeneration) {
return
@ -1677,22 +1741,19 @@ export class SshRelaySession {
) {
return Promise.resolve()
}
if (
consumerOwner?.outputFlowControl &&
(!offeredSource ||
payload.sourceMalformed ||
offeredSource.clientGeneration !== consumerOwner.clientGeneration ||
offeredSource.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 rejection = classifySshPtyFrameRejection(payload, consumerOwner)
if (rejection) {
if (offeredSource) {
this.retiredSourceDeliveries.retire(payload.providerGeneration, offeredSource)
}
this.ptyFrameRejectionLog.record(payload, consumerOwner, rejection)
if (rejection.action === 'retire-and-reattach-delivery') {
this.recoverRejectedPtyDelivery(payload, offeredSource)
}
return Promise.resolve()
}
const source = consumerOwner?.outputFlowControl ? offeredSource : undefined
if (source) {
if (source && consumerOwner) {
const current = this.sourceIdentityByRelayPtyId.get(source.relayPtyId)
if (
source.sourceEndSu <= source.sourceStartSu ||
@ -1703,12 +1764,14 @@ export class SshRelaySession {
current.ptyIncarnation !== payload.ptyIncarnation ||
(current.nextSourceSu !== undefined && current.nextSourceSu !== source.sourceStartSu)))
) {
closeSshPtyOutputGeneration(
payload.providerGeneration,
'ssh_source_frame_stale_or_non_contiguous'
)
this.mux?.dispose('connection_lost')
return Promise.reject(new Error('ssh_source_frame_stale_or_non_contiguous'))
const rejection = {
reason: 'source-range-invalid',
action: 'retire-and-reattach-delivery'
} as const
this.retiredSourceDeliveries.retire(payload.providerGeneration, source)
this.ptyFrameRejectionLog.record(payload, consumerOwner, rejection)
this.recoverRejectedPtyDelivery(payload, source)
return Promise.resolve()
}
this.sourceIdentityByRelayPtyId.set(source.relayPtyId, {
deliveryToken: source.deliveryToken,
@ -1731,6 +1794,163 @@ export class SshRelaySession {
})
}
private recoverRejectedPtyDelivery(
payload: SshPtyDataPayload,
source: SshPtyDataPayload['source']
): void {
const mux = this.mux
const providerGeneration = this.activePtyProviderGeneration
let relayPtyId: string
try {
relayPtyId = source?.relayPtyId ?? toRelaySshPtyId(this.targetId, payload.id)
} catch {
return
}
const appPtyId = toAppSshPtyId(this.targetId, relayPtyId)
if (
payload.id !== appPtyId ||
!mux ||
mux.isDisposed() ||
providerGeneration !== payload.providerGeneration ||
this.pendingPtyReattaches.has(appPtyId) ||
this.rejectedPtyReattaches.has(appPtyId)
) {
return
}
if (payload.rejectedSourceRecovery === 'reconnect-channel') {
console.warn(
`[ssh-relay-session] PTY ${relayPtyId} delivery identity could not be retired safely for ${this.targetId}; dropping the relay channel to reconnect`
)
mux.dispose('connection_lost')
return
}
const previous = this.rejectedPtyRecoveryAttempts.get(appPtyId)
const attempt =
previous?.providerGeneration === providerGeneration
? previous
: { providerGeneration, attempts: 0, generationAttempts: 0, reported: false }
if (
attempt.attempts >= SSH_REJECTED_PTY_RECOVERY_MAX_ATTEMPTS ||
attempt.generationAttempts >= SSH_REJECTED_PTY_RECOVERY_MAX_GENERATION_ATTEMPTS
) {
if (!attempt.reported) {
attempt.reported = true
console.warn(
`[ssh-relay-session] PTY ${relayPtyId} delivery recovery exhausted for ${this.targetId}; dropping the relay channel to reconnect`
)
// Why a channel drop and not a terminal relay error: a terminal error clears the reconnect
// backoff, rotates provider authority (aborting every in-flight fs and git request on the
// target) and parks the target in a manual-recovery state — over one PTY's delivery. Losing
// the channel is the recoverable escalation, and it is what this path did before targeted
// recovery existed.
mux.dispose('connection_lost')
}
return
}
attempt.attempts++
attempt.generationAttempts++
this.rejectedPtyRecoveryAttempts.set(appPtyId, attempt)
void this.rejectedPtyReattaches
.run(appPtyId, () =>
this.reattachRejectedPty(
relayPtyId,
mux,
providerGeneration,
payload.rejectedSourceRecovery === 'fresh-activation'
? 'fresh-activation'
: 'confirm-existing'
)
)
.then(
(recovered) => {
if (recovered) {
// Why only a completed reattach clears this: an accepted frame proves nothing about the
// delivery that was rejected, and resetting on one lets a flapping PTY reattach forever.
attempt.attempts = 0
return
}
this.retryRejectedPtyDelivery(payload, source, appPtyId)
},
(error: unknown) => {
console.warn(`[ssh-relay-session] PTY ${relayPtyId} targeted delivery recovery failed`, {
providerGeneration,
error: error instanceof Error ? error.message : String(error)
})
this.retryRejectedPtyDelivery(payload, source, appPtyId)
}
)
}
// Why liveness is checked before retrying: reattachKnownPty resolves without claiming the lease
// when the PTY exited mid-attach, which is indistinguishable from a failed reattach at the call
// site. Retrying that race twice would drop the relay channel over an ordinary PTY exit.
private retryRejectedPtyDelivery(
payload: SshPtyDataPayload,
source: SshPtyDataPayload['source'],
appPtyId: string
): void {
const ptyProvider = getSshPtyProvider(this.targetId) as SshPtyProvider | undefined
if (!ptyProvider || typeof ptyProvider.hasPty !== 'function' || !ptyProvider.hasPty(appPtyId)) {
this.rejectedPtyRecoveryAttempts.delete(appPtyId)
return
}
const timer = setTimeout(() => {
this.rejectedPtyRecoveryRetries.delete(timer)
this.recoverRejectedPtyDelivery(payload, source)
}, SSH_REJECTED_PTY_RECOVERY_RETRY_DELAY_MS)
timer.unref?.()
this.rejectedPtyRecoveryRetries.add(timer)
}
private async reattachRejectedPty(
relayPtyId: string,
mux: SshChannelMultiplexer,
providerGeneration: number,
targetedDeliveryRecovery: TargetedDeliveryRecovery
): Promise<boolean> {
const shouldContinue = () =>
this.mux === mux &&
!mux.isDisposed() &&
this.activePtyProviderGeneration === providerGeneration
const ptyProvider = getSshPtyProvider(this.targetId) as SshPtyProvider | undefined
// Why re-checked here: this can have waited for a queue slot, and a superseded generation must
// not pay for a lease read or an attach round trip.
if (!ptyProvider || !shouldContinue()) {
return false
}
const activeLease = this.store
.getSshRemotePtyLeases(this.targetId)
.find(
(lease) =>
lease.ptyId === relayPtyId && lease.state !== 'terminated' && lease.state !== 'expired'
)
const activeLeaseByPtyId = activeLease
? new Map<string, SshPtyLease>([[relayPtyId, activeLease]])
: new Map<string, SshPtyLease>()
const expectedIdentity = activeLease ? expectedIdentityForLease(activeLease) : undefined
const attachedLeaseIds = new Set<string>()
await this.reattachKnownPty({
ptyProvider,
ptyId: relayPtyId,
activeLeaseByPtyId,
expectedIdentityByPtyId: expectedIdentity
? new Map([[relayPtyId, expectedIdentity]])
: new Map(),
attachedLeaseIds,
mux,
providerGeneration,
shouldContinue,
targetedDeliveryRecovery
})
if (attachedLeaseIds.size > 0 && shouldContinue()) {
await this.store.markSshRemotePtyLeasesAttachedAsync(
this.targetId,
Array.from(attachedLeaseIds)
)
}
return attachedLeaseIds.has(relayPtyId)
}
private quarantineReattachData(pending: PendingPtyReattach, payload: SshPtyDataPayload): void {
this.observePrivateRecoveryFrame(pending, payload)
if (pending.restoreRequired) {
@ -1943,6 +2163,7 @@ export class SshRelaySession {
this.retiredSourceDeliveries.activate(relayPtyId)
clearProviderPtyState(payload.id)
deletePtyOwnership(payload.id)
this.rejectedPtyRecoveryAttempts.delete(payload.id)
getSshPtyConsumerRecovery(this.targetId)?.checkpointsByAppPtyId.delete(payload.id)
getSshPtyConsumerRecovery(this.targetId)?.checkpointsByAppPtyId.delete(
toRelaySshPtyId(this.targetId, payload.id)
@ -2051,6 +2272,7 @@ export class SshRelaySession {
mux: SshChannelMultiplexer
providerGeneration: number
shouldContinue: () => boolean
targetedDeliveryRecovery?: TargetedDeliveryRecovery
}): Promise<void> {
const {
ptyProvider,
@ -2060,7 +2282,8 @@ export class SshRelaySession {
attachedLeaseIds,
mux,
providerGeneration,
shouldContinue
shouldContinue,
targetedDeliveryRecovery
} = args
const appPtyId = toAppSshPtyId(this.targetId, ptyId)
const pendingReattach: PendingPtyReattach = {
@ -2079,7 +2302,10 @@ export class SshRelaySession {
let sourceActivationLease: SshPtyAttachResult['sourceActivationLease']
let recoveryActivationLease: SshPtyRecoveryActivationLease | undefined
try {
const recoveryRequest = await this.sourceRecoveryRequest(appPtyId)
const recoveryRequest =
targetedDeliveryRecovery === 'fresh-activation'
? undefined
: await this.sourceRecoveryRequest(appPtyId)
const attachResult = await this.attachPtyWithRetry(
ptyProvider,
ptyId,
@ -2105,7 +2331,27 @@ export class SshRelaySession {
await this.acceptPtyExit(exitDuringAttach)
return
}
if (recoveryRequest) {
const existingDeliveryConfirmed =
targetedDeliveryRecovery === 'confirm-existing' &&
recoveryRequest?.status === 'checkpoint' &&
!attachResult.sourceRecovery &&
Boolean(
attachResult.sourceActivation &&
this.sameSourceDelivery(attachResult.sourceActivation, recoveryRequest)
)
if (targetedDeliveryRecovery) {
const owner = this.activePtyConsumerOwner()
const activation = attachResult.sourceActivation
if (
!owner?.outputFlowControl ||
!activation ||
activation.clientGeneration !== owner.clientGeneration ||
activation.ownerGeneration !== owner.ownerGeneration
) {
return
}
}
if (recoveryRequest && !existingDeliveryConfirmed) {
const recovered = await this.finishSourceRecovery(
ptyId,
appPtyId,
@ -2177,6 +2423,18 @@ export class SshRelaySession {
pendingReattach.activated = true
recoveryActivationLease?.commit()
recoveryActivationLease = undefined
if (targetedDeliveryRecovery) {
if (targetedDeliveryRecovery === 'fresh-activation') {
this.retiredSourceDeliveries.activate(ptyId)
this.sourceIdentityByRelayPtyId.delete(ptyId)
getSshPtyConsumerRecovery(this.targetId)?.checkpointsByAppPtyId.delete(appPtyId)
getSshPtyConsumerRecovery(this.targetId)?.checkpointsByAppPtyId.delete(ptyId)
}
while (pendingReattach.queuedData.length > 0) {
await this.acceptPtyData(pendingReattach.queuedData.shift()!)
}
pendingReattach.livePassthrough = true
}
const exitAfterActivation = pendingReattach.exits.find(
(exit) =>
!exit.incarnationId ||
@ -2187,7 +2445,7 @@ export class SshRelaySession {
await this.acceptPtyExit(exitAfterActivation)
return
}
if (!recoveryRequest) {
if (!recoveryRequest && !targetedDeliveryRecovery) {
this.forwardReattachReplay(appPtyId, attachResult.replay ?? '')
}
sourceActivationLease?.commit()

View File

@ -5,6 +5,7 @@ import {
type DispatcherWriterEntry,
type DispatcherWriterLane
} from './dispatcher-writer-admission'
import { DispatcherWriterDrainArm } from './dispatcher-writer-drain-arm'
import { DispatcherWriterLaneScheduler } from './dispatcher-writer-lane-scheduler'
import {
DispatcherWriterSink,
@ -33,11 +34,11 @@ export class DispatcherClientWriter {
private readonly settleOnDrain = new Set<DispatcherWriterEntry>()
private readonly capacityListeners = new Set<() => void>()
private readonly idleWaiters = new Set<() => void>()
private readonly drain = new DispatcherWriterDrainArm()
private saturated = false
private drainArmed = false
private removeDrainListener: (() => void) | null = null
private livenessBypassOutstanding = false
private pumping = false
private retiredCapacity = false
private closed = false
private closeNotified = false
@ -88,7 +89,8 @@ export class DispatcherClientWriter {
encode: () => Buffer,
estimatedBytes: number,
onSettled: (result: SinkWriteSettlement) => void = () => {},
overflowIsNonFatal = false
overflowIsNonFatal = false,
isStillAdmitted?: () => boolean
): boolean {
if (this.closed) {
onSettled({ ok: false, error: new Error('Relay writer is closed') })
@ -97,6 +99,7 @@ export class DispatcherClientWriter {
const entry: DispatcherWriterEntry = {
lane,
encode,
isStillAdmitted,
estimatedBytes,
onSettled: onceDispatcherWriterSettlement(onSettled),
settled: false,
@ -126,8 +129,7 @@ export class DispatcherClientWriter {
}
this.closed = true
this.saturated = false
this.drainArmed = false
this.clearDrainListener()
this.drain.disarm()
for (const entry of this.admission.takeQueued()) {
this.releaseEntry(entry, { ok: false, error })
}
@ -167,6 +169,12 @@ export class DispatcherClientWriter {
} finally {
this.pumping = false
}
// Why deferred: retiring a deep queue drops every entry in one pass, and notifying per drop
// fans out to each listener O(queue) times for capacity that only changed once.
if (this.retiredCapacity && !this.closed) {
this.retiredCapacity = false
this.notifyCapacity()
}
}
private selectNext(): DispatcherWriterEntry | undefined {
@ -197,6 +205,11 @@ export class DispatcherClientWriter {
}
private writeEntry(entry: DispatcherWriterEntry): void {
if (entry.isStillAdmitted && !entry.isStillAdmitted()) {
this.releaseEntry(entry, { ok: false, error: new Error('PTY publication retired') })
this.retiredCapacity = true
return
}
this.laneScheduler.recordWrite(entry.lane)
this.inFlight.add(entry)
let callbackResult: SinkWriteSettlement | undefined
@ -247,30 +260,21 @@ export class DispatcherClientWriter {
}
private armDrain(): void {
if (this.drainArmed) {
return
}
this.drainArmed = true
try {
const registration = this.sink.registerDrain(() => this.handleDrain())
if (!registration.registered) {
this.drainArmed = false
} else if (this.drainArmed) {
this.removeDrainListener = registration.remove
} else {
registration.remove()
}
this.drain.arm(
(onDrain) => this.sink.registerDrain(onDrain),
() => this.handleDrain()
)
} catch (error) {
this.close(error instanceof Error ? error : new Error(String(error)))
}
}
private handleDrain(): void {
if (this.closed || !this.drainArmed) {
if (this.closed || !this.drain.isArmed) {
return
}
this.drainArmed = false
this.clearDrainListener()
this.drain.disarm()
this.saturated = false
this.livenessBypassOutstanding = false
for (const entry of Array.from(this.settleOnDrain)) {
@ -299,12 +303,6 @@ export class DispatcherClientWriter {
}
}
private clearDrainListener(): void {
const remove = this.removeDrainListener
this.removeDrainListener = null
remove?.()
}
private isIdle = (): boolean => this.inFlight.size === 0 && this.admission.queuedEntries === 0
private notifyIdle(): void {

View File

@ -12,6 +12,7 @@ export type DispatcherWriterSettlement = { ok: true } | { ok: false; error: Erro
export type DispatcherWriterEntry = {
lane: DispatcherWriterLane
encode: () => Buffer
isStillAdmitted?: () => boolean
estimatedBytes: number
onSettled: (result: DispatcherWriterSettlement) => void
settled: boolean

View File

@ -0,0 +1,39 @@
type DrainRegistration = { registered: boolean; remove: () => void }
/**
* Holds the writer's single outstanding drain registration on its sink.
*
* Why it is re-checked after registering: a sink may invoke the drain callback synchronously from
* inside registerDrain, which disarms this arm before registration returns. Storing the remove
* handle then would cancel a later, still-wanted arm instead of the one that already fired.
*/
export class DispatcherWriterDrainArm {
private armed = false
private remove: (() => void) | null = null
get isArmed(): boolean {
return this.armed
}
arm(register: (onDrain: () => void) => DrainRegistration, onDrain: () => void): void {
if (this.armed) {
return
}
this.armed = true
const registration = register(onDrain)
if (!registration.registered) {
this.armed = false
} else if (this.armed) {
this.remove = registration.remove
} else {
registration.remove()
}
}
disarm(): void {
this.armed = false
const remove = this.remove
this.remove = null
remove?.()
}
}

View File

@ -55,6 +55,11 @@ export type RelayClientSourceOptions = {
resumeReads?: () => void
}
export type PtyDataPublicationAdmission = (
clientId: number,
params: Readonly<Record<string, unknown>>
) => boolean
export type MethodHandler = (
params: Record<string, unknown>,
context: RequestContext
@ -108,6 +113,7 @@ export class RelayDispatcher {
private disposeListeners = new Set<() => void>()
private legacyCapacityListeners = new Set<() => void>()
private clientCapacityListeners = new Map<number, Set<() => void>>()
private ptyDataPublicationAdmission: PtyDataPublicationAdmission | null = null
private publicationTransactionDepth = 0
private deferredLegacyCapacity = false
private deferredForcedLegacyCapacity = false
@ -205,6 +211,20 @@ export class RelayDispatcher {
return () => this.disposeListeners.delete(listener)
}
// Why single-slot rather than a listener set: admission is a veto, so two registrations would have
// to agree on precedence. One owner (the PTY consumer session) holds it for the dispatcher's life.
registerPtyDataPublicationAdmission(admission: PtyDataPublicationAdmission): () => void {
if (this.ptyDataPublicationAdmission) {
throw new Error('PTY data publication admission is already registered')
}
this.ptyDataPublicationAdmission = admission
return () => {
if (this.ptyDataPublicationAdmission === admission) {
this.ptyDataPublicationAdmission = null
}
}
}
onLegacyPtyCapacity(listener: () => void): () => void {
this.legacyCapacityListeners.add(listener)
return () => this.legacyCapacityListeners.delete(listener)
@ -281,7 +301,9 @@ export class RelayDispatcher {
data: string,
limit = data.length
): number {
const clients = this.activeClients()
const clients = this.activeClients().filter((client) =>
this.admitsPtyDataPublication(client.id, params)
)
const max = Math.min(data.length, limit)
if (clients.length === 0) {
return max
@ -331,7 +353,7 @@ export class RelayDispatcher {
params
}
return this.tryPublishToClients(
this.activeClients(),
this.activeClients().filter((client) => this.admitsPtyDataPublication(client.id, params)),
msg,
options.interactive ? 'interactive' : 'ordinary'
)
@ -346,7 +368,9 @@ export class RelayDispatcher {
return false
}
return this.tryPublishToClients(
this.activeClients().filter((client) => matchesClient(client.id)),
this.activeClients().filter(
(client) => matchesClient(client.id) && this.admitsPtyDataPublication(client.id, params)
),
{ jsonrpc: '2.0', method: 'pty.data', params },
options.interactive ? 'interactive' : 'ordinary'
)
@ -361,7 +385,9 @@ export class RelayDispatcher {
return false
}
return this.projectToClients(
this.activeClients().filter((client) => matchesClient(client.id)),
this.activeClients().filter(
(client) => matchesClient(client.id) && this.admitsPtyDataPublication(client.id, params)
),
{ jsonrpc: '2.0', method: 'pty.data', params },
options.interactive ? 'interactive' : 'ordinary'
)
@ -381,6 +407,10 @@ export class RelayDispatcher {
onSettled({ ok: false, error: new Error('Relay client is not connected') })
return false
}
if (!this.admitsPtyDataPublication(clientId, params)) {
onSettled({ ok: false, error: new Error('PTY publication is not admitted') })
return false
}
return this.publishToClient(
client,
{ jsonrpc: '2.0', method: 'pty.data', params },
@ -547,6 +577,9 @@ export class RelayDispatcher {
if (client.closed) {
continue
}
if (method === 'pty.data' && !this.admitsPtyDataPublication(client.id, params ?? {})) {
continue
}
if (method === 'pty.replay') {
// Why: replay is never re-sent, so it takes the control lane where overflow is fatal — the
// writer closes the client and reconnect reloads history rather than stranding a short buffer.
@ -584,6 +617,9 @@ export class RelayDispatcher {
method,
...(params !== undefined ? { params } : {})
}
if (method === 'pty.data' && !this.admitsPtyDataPublication(client.id, params ?? {})) {
return false
}
const frameBytes = this.estimateFrameBytes(msg)
if (this.publishToClient(client, msg, 'ordinary', undefined, frameBytes)) {
return true
@ -1055,12 +1091,17 @@ export class RelayDispatcher {
const seq = client.nextOutgoingSeq++
return encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq)
}
const isStillAdmitted =
'method' in msg && msg.method === 'pty.data'
? () => this.admitsPtyDataPublication(client.id, msg.params ?? {})
: undefined
return client.writer.enqueue(
lane,
encode,
frameBytes,
onSettled,
lane === 'control' && controlOverflow === 'reject'
lane === 'control' && controlOverflow === 'reject',
isStillAdmitted
)
}
@ -1091,6 +1132,13 @@ export class RelayDispatcher {
return Array.from(this.clients.values()).filter((client) => !client.closed)
}
private admitsPtyDataPublication(
clientId: number,
params: Readonly<Record<string, unknown>>
): boolean {
return this.ptyDataPublicationAdmission?.(clientId, params) ?? true
}
private activeClientKeys(): string[] {
return this.activeClients().map((client) => this.clientKey(client))
}

View File

@ -221,12 +221,18 @@ describe('PtyHandler negotiated source publication', () => {
expect(exitFrames()).toHaveLength(1)
})
function attachSubscriber(): Buffer[] {
async function attachSubscriber(
holdDataSettlement?: (settle: (result: SinkWriteSettlement) => void) => boolean
): Promise<Buffer[]> {
const subscriberWrites: Buffer[] = []
dispatcher.attachClient(
const clientId = dispatcher.attachClient(
(data, settle) => {
subscriberWrites.push(Buffer.from(data))
if (notification(data)?.method === 'pty.exit') {
const frame = notification(data)
if (frame?.method === 'pty.data' && holdDataSettlement?.(settle)) {
return true
}
if (frame?.method === 'pty.exit') {
// Why: real sockets never settle inside write(); see the primary sink above.
queueMicrotask(() => settle({ ok: true }))
return true
@ -234,8 +240,18 @@ describe('PtyHandler negotiated source publication', () => {
settle({ ok: true })
return true
},
{ supportsWriteCallback: true }
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
clientId,
requestFrame(20, 'pty.openClient', {
protocolVersion: 1,
clientInstanceId: 'legacy-subscriber',
requestedRole: 'subscriber'
})
)
await vi.advanceTimersByTimeAsync(0)
return subscriberWrites
}
@ -248,7 +264,7 @@ describe('PtyHandler negotiated source publication', () => {
it('never re-delivers the exit to subscribers when a cancel retires the record', async () => {
await spawn({})
const spawnResult = writes.map((buffer) => responseResult(buffer, 2)).find(Boolean)!
const subscriberWrites = attachSubscriber()
const subscriberWrites = await attachSubscriber()
dataCallback!('prompt')
await vi.advanceTimersByTimeAsync(8)
@ -293,6 +309,13 @@ describe('PtyHandler negotiated source publication', () => {
await spawn({})
const spawnResult = writes.map((buffer) => responseResult(buffer, 2)).find(Boolean)!
await cancelSourceDelivery(spawnResult)
const subscriberWrites = await attachSubscriber((settle) => {
if (!holdDataSettlements) {
return false
}
heldDataSettlements.push(settle)
return true
})
holdDataSettlements = true
dataCallback!('first')
@ -311,7 +334,7 @@ describe('PtyHandler negotiated source publication', () => {
heldDataSettlements[0]({ ok: true })
await vi.advanceTimersByTimeAsync(8)
const frames = writes
const frames = subscriberWrites
.map(notification)
.filter(
(frame): frame is Notification =>
@ -351,7 +374,7 @@ describe('PtyHandler negotiated source publication', () => {
await spawn({})
const spawnResult = writes.map((buffer) => responseResult(buffer, 2)).find(Boolean)!
const id = String(spawnResult.id)
const subscriberWrites = attachSubscriber()
const subscriberWrites = await attachSubscriber()
const stderr = vi.spyOn(process.stderr, 'write').mockReturnValue(true)
const publishOwnerExit = vi
.spyOn(dispatcher, 'tryNotifyPtyExitToClient')
@ -381,7 +404,7 @@ describe('PtyHandler negotiated source publication', () => {
await spawn({})
const spawnResult = writes.map((buffer) => responseResult(buffer, 2)).find(Boolean)!
const id = String(spawnResult.id)
const subscriberWrites = attachSubscriber()
const subscriberWrites = await attachSubscriber()
const stderr = vi.spyOn(process.stderr, 'write').mockReturnValue(true)
const settleOwnerExit = vi.spyOn(adapter, 'settleExitPublication').mockImplementation(() => {
throw new Error('exit settlement failed')
@ -408,7 +431,7 @@ describe('PtyHandler negotiated source publication', () => {
it('lets a retired record re-target its own exit instead of broadcasting a duplicate', async () => {
await spawn({})
const spawnResult = writes.map((buffer) => responseResult(buffer, 2)).find(Boolean)!
const subscriberWrites = attachSubscriber()
const subscriberWrites = await attachSubscriber()
const publishExitAfterRetire = vi.fn(() => true)
handler.setSourcePublication(stubPublication({ accepts: () => false, publishExitAfterRetire }))
@ -680,20 +703,50 @@ describe('PtyHandler negotiated source publication', () => {
await spawn({})
const detached: number[] = []
const healthyWrites: Buffer[] = []
let saturateSubscriber = false
dispatcher.onClientDetached((clientId) => detached.push(clientId))
const saturatedId = dispatcher.attachClient(() => false, {
supportsWriteCallback: true,
writableLength: () => 16 * 1024,
writableHighWaterMark: () => 4 * 1024 * 1024
})
const saturatedId = dispatcher.attachClient(
(_data, settle) => {
if (saturateSubscriber) {
return false
}
settle({ ok: true })
return true
},
{
supportsWriteCallback: true,
writableLength: () => 16 * 1024,
writableHighWaterMark: () => 4 * 1024 * 1024
},
endpointIdentity
)
const healthyId = dispatcher.attachClient(
(data, settle) => {
healthyWrites.push(Buffer.from(data))
settle({ ok: true })
return true
},
{ supportsWriteCallback: true }
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
saturatedId,
requestFrame(20, 'pty.openClient', {
protocolVersion: 1,
clientInstanceId: 'saturated-subscriber',
requestedRole: 'subscriber'
})
)
dispatcher.feedClient(
healthyId,
requestFrame(21, 'pty.openClient', {
protocolVersion: 1,
clientInstanceId: 'healthy-subscriber',
requestedRole: 'subscriber'
})
)
await vi.advanceTimersByTimeAsync(0)
saturateSubscriber = true
const payload = 's'.repeat(16 * 1024)
let admitted = 0
while (

View File

@ -121,7 +121,7 @@ describe('relay PTY consumer owner displacement', () => {
expect(adapter.deliveryMode(reconnectClientId)).toBe('source-owner')
// Why: revocation is what bounds the takeover — the stale owner cannot open or drive deliveries again
// even though its socket never closed.
expect(adapter.deliveryMode(1)).toBe('subscriber')
expect(adapter.deliveryMode(1)).toBe('unadmitted')
expect(adapter.openDelivery(1, 'pty-2', 'incarnation-1')).toBeNull()
expect(detached).toContain(1)
expect(closeStaleTransport).toHaveBeenCalledOnce()
@ -244,7 +244,7 @@ describe('relay PTY consumer owner displacement', () => {
grantSettlement!({ ok: false, error: new Error('reconnect socket closed mid-response') })
expect(adapter.deliveryMode(1)).toBe('source-owner')
expect(adapter.deliveryMode(reconnectClientId)).toBe('subscriber')
expect(adapter.deliveryMode(reconnectClientId)).toBe('unadmitted')
// Why: the rolled-back attempt must leave the incumbent's lease reclaimable by the next reconnect.
const retryWrites: Buffer[] = []
@ -275,7 +275,7 @@ describe('relay PTY consumer owner displacement', () => {
role: 'session-owner',
ownerGeneration: 3
})
expect(adapter.deliveryMode(1)).toBe('subscriber')
expect(adapter.deliveryMode(1)).toBe('unadmitted')
expect(detached).toContain(1)
})
@ -349,7 +349,7 @@ describe('relay PTY consumer owner displacement', () => {
code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
})
expect(adapter.deliveryMode(firstReconnectClientId)).toBe('source-owner')
expect(adapter.deliveryMode(retryClientId)).toBe('subscriber')
expect(adapter.deliveryMode(retryClientId)).toBe('unadmitted')
dispatcher.detachClient(firstReconnectClientId)
dispatcher.feedClient(

View File

@ -0,0 +1,219 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
RelayDispatcher,
type RelayClientSessionIdentity,
type RelayClientSinkOptions,
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'
}
type RpcMessage = {
id?: number
method?: string
params?: Record<string, unknown>
result?: Record<string, unknown>
}
function message(buffer: Buffer): RpcMessage | null {
if (buffer[0] !== MessageType.Regular) {
return null
}
const length = buffer.readUInt32BE(9)
return JSON.parse(buffer.subarray(13, 13 + length).toString('utf8'))
}
function openFrame(
id: number,
requestedRole: 'session-owner' | 'subscriber',
outputFlowControl = false
): Buffer {
return encodeJsonRpcFrame(
{
jsonrpc: '2.0',
id,
method: 'pty.openClient',
params: {
protocolVersion: 1,
clientInstanceId: `client-${id}`,
requestedRole,
...(outputFlowControl
? {
capabilities: {
outputFlowControl: { versions: [1], requestedWindowSu: 16 }
}
}
: {})
}
},
id,
0
)
}
async function flushRequests(): Promise<void> {
await new Promise((resolve) => setImmediate(resolve))
}
class SaturatedSink {
readonly writes: Buffer[] = []
private readonly drainWaiters = new Set<() => void>()
private writableBytes = 0
saturateNext = false
readonly options: RelayClientSinkOptions = {
supportsWriteCallback: true,
writableLength: () => this.writableBytes,
writableHighWaterMark: () => 8 * 1024 * 1024,
waitWriteDrain: (callback) => {
this.drainWaiters.add(callback)
return () => this.drainWaiters.delete(callback)
}
}
write = (data: Buffer, onSettled: (result: SinkWriteSettlement) => void): boolean => {
this.writes.push(Buffer.from(data))
this.writableBytes += data.length
onSettled({ ok: true })
if (!this.saturateNext) {
return true
}
this.saturateNext = false
return false
}
drain(): void {
this.writableBytes = 0
for (const callback of Array.from(this.drainWaiters)) {
callback()
}
}
}
describe('relay PTY publication admission', () => {
let dispatcher: RelayDispatcher | null = null
afterEach(() => {
dispatcher?.dispose()
dispatcher = null
})
it('retires queued legacy data when the source-credit grant response overtakes it', async () => {
const sink = new SaturatedSink()
dispatcher = new RelayDispatcher(sink.write, sink.options, endpointIdentity)
sink.saturateNext = true
dispatcher.notify('test.blocker')
const settled = vi.fn<(result: SinkWriteSettlement) => void>()
const legacyData = 'x'.repeat(1024 * 1024 + 128)
expect(dispatcher.tryNotifyPtyDataToClient(1, { id: 'pty-1', data: legacyData }, settled)).toBe(
true
)
expect(dispatcher.legacyRetentionBelowLowWater).toBe(false)
const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a')
expect(adapter.deliveryMode(1)).toBe('unadmitted')
dispatcher.feed(openFrame(1, 'session-owner', true))
await flushRequests()
sink.drain()
expect(sink.writes.map(message).some((entry) => entry?.method === 'pty.data')).toBe(false)
expect(sink.writes.map(message).find((entry) => entry?.id === 1)?.result).toMatchObject({
role: 'session-owner'
})
expect(settled).toHaveBeenCalledOnce()
expect(settled.mock.calls[0][0]).toMatchObject({ ok: false })
expect(dispatcher.legacyRetentionBelowLowWater).toBe(true)
})
it('does not publish legacy data to a source-credit owner before PTY attachment', async () => {
const writes: Buffer[] = []
dispatcher = new RelayDispatcher(
(data, onSettled) => {
writes.push(Buffer.from(data))
onSettled({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a')
dispatcher.feed(openFrame(1, 'session-owner', true))
await flushRequests()
expect(adapter.deliveryMode(1)).toBe('source-owner')
expect(dispatcher.tryNotifyPtyData({ id: 'pty-1', data: 'legacy' })).toBe(true)
expect(writes.map(message).filter((entry) => entry?.method === 'pty.data')).toEqual([])
})
it('publishes current source output to the admitted source-credit owner', async () => {
const writes: Buffer[] = []
dispatcher = new RelayDispatcher(
(data, onSettled) => {
writes.push(Buffer.from(data))
onSettled({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a')
const publication = new RelayPtySourcePublication(dispatcher, adapter, () => {})
dispatcher.feed(openFrame(1, 'session-owner', true))
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 })
expect(publication.publish('pty-1', { data: 'current' }, false)).toBe(true)
const data = writes.map(message).find((entry) => entry?.method === 'pty.data')?.params
expect(data).toMatchObject({
id: 'pty-1',
data: 'current',
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-1'
})
expect(data?.deliveryToken).toEqual(expect.any(String))
})
it('publishes legacy output to an admitted subscriber', async () => {
const writes: Buffer[] = []
dispatcher = new RelayDispatcher(
(data, onSettled) => {
writes.push(Buffer.from(data))
onSettled({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a')
dispatcher.feed(openFrame(1, 'subscriber'))
await flushRequests()
expect(adapter.deliveryMode(1)).toBe('subscriber')
expect(dispatcher.tryNotifyPtyData({ id: 'pty-1', data: 'legacy' })).toBe(true)
expect(writes.map(message).find((entry) => entry?.method === 'pty.data')?.params).toMatchObject(
{ id: 'pty-1', data: 'legacy' }
)
})
})

View File

@ -58,7 +58,10 @@ export function projectPtySourceOutputToLegacy(
interactive: boolean
): boolean {
return dispatcher.projectPtyDataToMatchingClients(
(clientId) => session.deliveryMode(clientId) !== 'source-owner',
(clientId) => {
const mode = session.deliveryMode(clientId)
return mode === 'legacy-owner' || mode === 'subscriber'
},
{
id,
data: output.data,

View File

@ -223,14 +223,24 @@ describe('RelayPtySourcePublication', () => {
it('keeps mixed legacy and V1 clients on distinct frame authority', async () => {
const harness = await createHarness(8)
const legacyWrites: Buffer[] = []
dispatcher!.attachClient(
const legacyClientId = dispatcher!.attachClient(
(data, onSettled) => {
legacyWrites.push(Buffer.from(data))
onSettled({ ok: true })
return true
},
{ supportsWriteCallback: true }
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher!.feedClient(
legacyClientId,
requestFrame(2, 'pty.openClient', {
protocolVersion: 1,
clientInstanceId: 'legacy-client',
requestedRole: 'subscriber'
})
)
await flushRequests()
harness.publication.publish('pty-1', { data: 'data' }, false)
@ -289,10 +299,15 @@ describe('RelayPtySourcePublication', () => {
const saturatedWrites: Buffer[] = []
const healthyWrites: Buffer[] = []
const heldSettlements: ((result: SinkWriteSettlement) => void)[] = []
let saturateSubscriber = false
dispatcher!.onClientDetached((clientId) => detached.push(clientId))
const saturatedId = dispatcher!.attachClient(
(data, onSettled) => {
saturatedWrites.push(Buffer.from(data))
if (!saturateSubscriber) {
onSettled({ ok: true })
return true
}
heldSettlements.push(onSettled)
return false
},
@ -300,7 +315,8 @@ describe('RelayPtySourcePublication', () => {
supportsWriteCallback: true,
writableLength: () => 128 * 1024,
writableHighWaterMark: () => 4 * 1024 * 1024
}
},
endpointIdentity
)
const healthyId = dispatcher!.attachClient(
(data, onSettled) => {
@ -308,8 +324,27 @@ describe('RelayPtySourcePublication', () => {
onSettled({ ok: true })
return true
},
{ supportsWriteCallback: true }
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher!.feedClient(
saturatedId,
requestFrame(2, 'pty.openClient', {
protocolVersion: 1,
clientInstanceId: 'saturated-subscriber',
requestedRole: 'subscriber'
})
)
dispatcher!.feedClient(
healthyId,
requestFrame(3, 'pty.openClient', {
protocolVersion: 1,
clientInstanceId: 'healthy-subscriber',
requestedRole: 'subscriber'
})
)
await flushRequests()
saturateSubscriber = true
const saturatedPayload = 's'.repeat(128 * 1024)
let admitted = 0
while (
@ -328,7 +363,9 @@ describe('RelayPtySourcePublication', () => {
expect(detached).toEqual([saturatedId])
expect(detached).not.toContain(healthyId)
expect(saturatedWrites).toHaveLength(1)
expect(
saturatedWrites.map(notification).filter((frame) => frame?.method === 'pty.data')
).toHaveLength(1)
expect(heldSettlements).toHaveLength(1)
expect(
healthyWrites.map(notification).filter((frame) => frame?.method === 'pty.data')

View File

@ -71,7 +71,7 @@ export class RelayPtySourcePublication {
}
const mode = this.session.deliveryMode(context.clientId)
let current = this.deliveries.get(id)
if (mode === 'subscriber') {
if (mode === 'unadmitted' || mode === 'subscriber') {
this.sender.releaseRotationFence(current)
return false
}

View File

@ -2,8 +2,7 @@ import {
PTY_CONSUMER_SESSION_PROTOCOL_VERSION,
PtyConsumerSession,
type PtyConsumerSessionAdmission,
type PtyConsumerSessionGrant,
type PtyConsumerSessionHello
type PtyConsumerSessionGrant
} from '../shared/pty-consumer-session'
import { DEFAULT_PTY_SOURCE_WINDOW_SU } from '../shared/pty-source-credit-contract'
import type {
@ -12,63 +11,18 @@ import type {
PtySourceSpan,
PtySourceTransform
} from '../shared/pty-source-credit-contract'
import type { RelayClientSessionIdentity, RelayDispatcher, RequestContext } from './dispatcher'
import type { RelayDispatcher, RequestContext } from './dispatcher'
import type { PtySourceSendReservation } from './pty-source-credit-ledger'
import { SshPtySourceCreditAdapter } from './ssh-pty-source-credit-adapter'
import {
admitsSshPtyDataPublication,
sshPtyDeliveryMode,
type SshPtyDeliveryMode
} from './ssh-pty-data-publication-admission'
import { parseOpenClientParams, requireIdentity } from './ssh-pty-open-client-request'
export const SSH_PTY_OPEN_CLIENT_METHOD = 'pty.openClient'
type OpenClientParams = PtyConsumerSessionHello & {
protocolVersion: number
}
function parseOpenClientParams(params: Record<string, unknown>): OpenClientParams {
const resume =
typeof params.resume === 'object' && params.resume !== null
? (params.resume as Record<string, unknown>)
: undefined
const capabilities =
typeof params.capabilities === 'object' && params.capabilities !== null
? (params.capabilities as Record<string, unknown>)
: undefined
const outputFlowControl =
typeof capabilities?.outputFlowControl === 'object' && capabilities.outputFlowControl !== null
? (capabilities.outputFlowControl as Record<string, unknown>)
: undefined
return {
protocolVersion: Number(params.protocolVersion),
clientInstanceId: String(params.clientInstanceId ?? ''),
requestedRole: String(params.requestedRole ?? '') as PtyConsumerSessionHello['requestedRole'],
...(resume
? {
resume: {
ownerGeneration: Number(resume.ownerGeneration),
ownerLease: String(resume.ownerLease ?? '')
}
}
: {}),
...(outputFlowControl
? {
capabilities: {
outputFlowControl: {
versions: Array.isArray(outputFlowControl.versions)
? outputFlowControl.versions.map(Number)
: [],
requestedWindowSu: Number(outputFlowControl.requestedWindowSu)
}
}
}
: {})
}
}
function requireIdentity(context: RequestContext): RelayClientSessionIdentity {
if (!context.sessionIdentity) {
throw new Error('SSH PTY consumer transport identity is unavailable')
}
return context.sessionIdentity
}
export class SshPtyConsumerSessionAdapter {
private readonly session: PtyConsumerSession
private readonly sourceCredit: SshPtySourceCreditAdapter
@ -92,15 +46,26 @@ export class SshPtyConsumerSessionAdapter {
serverBuildId,
outputFlowControl: { versions: [1], maxWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU }
})
// Why the admission is consulted again at drain time (see isStillAdmitted): a frame can sit
// queued behind a saturated socket long enough for the grant or the delivery to be retired, and
// publishing it then hands the client output from an owner it no longer is.
dispatcher.registerPtyDataPublicationAdmission((clientId, params) =>
admitsSshPtyDataPublication(
this.session.activeGrant(String(clientId)),
params,
this.sourceCredit
)
)
dispatcher.onRequest(SSH_PTY_OPEN_CLIENT_METHOD, (params, context) =>
this.openClient(params, context)
)
dispatcher.onClientDetached((clientId, cause) => {
const grant = this.session.activeGrant(String(clientId))
const connectionKey = String(clientId)
const grant = this.session.activeGrant(connectionKey)
if (grant) {
this.clearPausedForGrant(grant)
}
this.session.close(String(clientId), cause)
this.session.close(connectionKey, cause)
if (grant) {
this.sourceCredit.retainOrCloseOnDetach(grant)
}
@ -239,12 +204,8 @@ export class SshPtyConsumerSessionAdapter {
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'
deliveryMode(clientId: number): SshPtyDeliveryMode {
return sshPtyDeliveryMode(this.session.activeGrant(String(clientId)))
}
private async openClient(

View File

@ -0,0 +1,38 @@
import type { PtyConsumerSessionGrant } from '../shared/pty-consumer-session'
import type { SshPtySourceCreditAdapter } from './ssh-pty-source-credit-adapter'
export type SshPtyDeliveryMode = 'unadmitted' | 'source-owner' | 'legacy-owner' | 'subscriber'
export function sshPtyDeliveryMode(
grant: Readonly<PtyConsumerSessionGrant> | null
): SshPtyDeliveryMode {
if (!grant) {
return 'unadmitted'
}
if (grant.role !== 'session-owner') {
return 'subscriber'
}
return grant.capabilities?.outputFlowControl ? 'source-owner' : 'legacy-owner'
}
export function admitsSshPtyDataPublication(
grant: Readonly<PtyConsumerSessionGrant> | null,
params: Readonly<Record<string, unknown>>,
sourceCredit: Pick<SshPtySourceCreditAdapter, 'ownsDelivery'>
): boolean {
if (!grant) {
return false
}
if (grant.role !== 'session-owner' || !grant.capabilities?.outputFlowControl) {
return true
}
const id = typeof params.id === 'string' ? params.id : ''
const token = typeof params.deliveryToken === 'string' ? params.deliveryToken : ''
const identity = token ? sourceCredit.ownsDelivery(token, grant, id) : null
return (
identity !== null &&
params.clientGeneration === identity.clientGeneration &&
params.ownerGeneration === identity.ownerGeneration &&
params.ptyIncarnation === identity.ptyIncarnation
)
}

View File

@ -0,0 +1,56 @@
import type { PtyConsumerSessionHello } from '../shared/pty-consumer-session'
import type { RelayClientSessionIdentity, RequestContext } from './dispatcher'
export type OpenClientParams = PtyConsumerSessionHello & {
protocolVersion: number
}
// Why every field is coerced rather than trusted: these arrive off the wire from a peer that may be
// a different build, so a missing or wrong-typed member has to degrade into a value the session can
// reject by contract instead of throwing somewhere further in.
export function parseOpenClientParams(params: Record<string, unknown>): OpenClientParams {
const resume =
typeof params.resume === 'object' && params.resume !== null
? (params.resume as Record<string, unknown>)
: undefined
const capabilities =
typeof params.capabilities === 'object' && params.capabilities !== null
? (params.capabilities as Record<string, unknown>)
: undefined
const outputFlowControl =
typeof capabilities?.outputFlowControl === 'object' && capabilities.outputFlowControl !== null
? (capabilities.outputFlowControl as Record<string, unknown>)
: undefined
return {
protocolVersion: Number(params.protocolVersion),
clientInstanceId: String(params.clientInstanceId ?? ''),
requestedRole: String(params.requestedRole ?? '') as PtyConsumerSessionHello['requestedRole'],
...(resume
? {
resume: {
ownerGeneration: Number(resume.ownerGeneration),
ownerLease: String(resume.ownerLease ?? '')
}
}
: {}),
...(outputFlowControl
? {
capabilities: {
outputFlowControl: {
versions: Array.isArray(outputFlowControl.versions)
? outputFlowControl.versions.map(Number)
: [],
requestedWindowSu: Number(outputFlowControl.requestedWindowSu)
}
}
}
: {})
}
}
export function requireIdentity(context: RequestContext): RelayClientSessionIdentity {
if (!context.sessionIdentity) {
throw new Error('SSH PTY consumer transport identity is unavailable')
}
return context.sessionIdentity
}

View File

@ -0,0 +1,84 @@
import {
PTY_CONSUMER_OWNER_HELD_ATTACHED_ERROR,
PTY_CONSUMER_OWNER_HELD_DISCONNECTED_ERROR,
PTY_CONSUMER_OWNER_HELD_GRACE_FLOOR_MS,
PTY_CONSUMER_OWNER_HELD_SELF_ERROR,
type PtyConsumerCloseCause,
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
} from './pty-consumer-session-contract'
type HeldOwner = {
state: 'pending' | 'active' | 'disconnected'
disconnectedAt?: number
disconnectCause?: PtyConsumerCloseCause
}
export type RefuseHeldPtyConsumerOwnerOptions = {
ownerGraceMs: number
now: number
sameClient: boolean
// Why a callback instead of mutating the record: the owner record is reachable from `replaces`
// chains and from displaced-owner snapshots already handed to callers. Clamping in place would
// rewrite those retroactively, so the session that owns the record applies it copy-on-write.
clampGraceTo: (disconnectedAt: number) => void
}
function refuse(message: string, code: number): never {
throw Object.assign(new Error(message), { code })
}
/**
* Why an owner-capable request is refused rather than demoted: a subscriber grant is unusable to a
* client that needs to drive the PTY, and it arrives shaped like success. A coded refusal lets the
* caller retry the transient case and stop on the blocked one.
*/
export function refuseHeldPtyConsumerOwner(
owner: Readonly<HeldOwner>,
options: RefuseHeldPtyConsumerOwnerOptions
): never {
if (owner.state === 'pending') {
refuse('Owner grant publication is still pending', PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR)
}
if (owner.state === 'active') {
// Why identity without the lease: a client that lost its proof — a fresh process, a dropped
// recovery record — still knows who it is. Against an incumbent carrying its own instance id
// the honest answer is "your other connection is still registered", which resolves itself once
// the relay notices that socket. Blocking here strands the single-app case forever.
if (options.sameClient) {
refuse(
"PTY session owner is held by this client's own earlier connection",
PTY_CONSUMER_OWNER_HELD_SELF_ERROR
)
}
refuse(
'PTY session owner is held by an attached connection',
PTY_CONSUMER_OWNER_HELD_ATTACHED_ERROR
)
}
clampDisconnectedOwnerGrace(owner, options)
refuse(
'PTY session owner is held by a disconnected connection within its grace period',
PTY_CONSUMER_OWNER_HELD_DISCONNECTED_ERROR
)
}
// Why only a peer-closed disconnect may shorten this: the floor is a bet that the incumbent is gone,
// and the relay tears a client's socket down for its own reasons too — a full lane queue is the
// signature of an owner that is alive but not draining fast enough. No owner completes a reconnect
// ladder in 250 ms, so clamping on a teardown we initiated hands a live owner's admission away and
// it can never get it back. Expiring a record never stops the remote PTY, but it does cost the user
// every route back to it.
function clampDisconnectedOwnerGrace(
owner: Readonly<HeldOwner>,
options: RefuseHeldPtyConsumerOwnerOptions
): void {
if (owner.disconnectCause !== 'peer-closed') {
return
}
const floorStart =
options.now - Math.max(options.ownerGraceMs - PTY_CONSUMER_OWNER_HELD_GRACE_FLOOR_MS, 0)
if ((owner.disconnectedAt ?? 0) <= floorStart) {
return
}
options.clampGraceTo(floorStart)
}

View File

@ -1,11 +1,6 @@
import { randomUUID } from 'node:crypto'
import {
PTY_CONSUMER_OWNER_GRACE_MS,
PTY_CONSUMER_OWNER_HELD_ATTACHED_ERROR,
PTY_CONSUMER_OWNER_HELD_DISCONNECTED_ERROR,
PTY_CONSUMER_OWNER_HELD_GRACE_FLOOR_MS,
PTY_CONSUMER_OWNER_HELD_SELF_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_SESSION_PROTOCOL_VERSION,
type PtyConsumerAuthentication,
type PtyConsumerCloseCause,
@ -25,6 +20,7 @@ import {
isPtyConsumerOwnerSameClient,
matchesPtyConsumerOwnerClaim
} from './pty-consumer-owner-recovery'
import { refuseHeldPtyConsumerOwner } from './pty-consumer-owner-admission'
export * from './pty-consumer-session-contract'
@ -49,10 +45,6 @@ type OwnerRecord = {
replaces?: OwnerRecord
}
function throwOwnerError(message: string, code: number): never {
throw Object.assign(new Error(message), { code })
}
export class PtyConsumerSession {
private readonly clients = new Map<string, ClientRecord>()
private readonly now: () => number
@ -224,7 +216,14 @@ export class PtyConsumerSession {
return this.newOwner(hello, authentication, null)
}
if (!matchesPtyConsumerOwnerClaim(hello, authentication, current)) {
this.refuseHeldOwner(hello, authentication, current)
refuseHeldPtyConsumerOwner(current, {
ownerGraceMs: this.ownerGraceMs,
now: this.now(),
sameClient: isPtyConsumerOwnerSameClient(hello, authentication, current),
clampGraceTo: (disconnectedAt) => {
this.owner = { ...current, disconnectedAt }
}
})
}
assertPtyConsumerOwnerRecovery(hello, authentication, current)
// Why an active owner is displaced rather than refused: the resume proof matched this owner's
@ -234,61 +233,6 @@ export class PtyConsumerSession {
return this.newOwner(hello, authentication, current)
}
// Why an owner-capable request is refused rather than demoted: a subscriber grant is unusable to a
// client that needs to drive the PTY, and it arrives shaped like success. A coded refusal lets the
// caller retry the transient case and stop on the blocked one.
private refuseHeldOwner(
hello: PtyConsumerSessionHello,
authentication: PtyConsumerAuthentication,
current: OwnerRecord
): never {
if (current.state === 'pending') {
throwOwnerError(
'Owner grant publication is still pending',
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
)
}
if (current.state === 'active') {
// Why identity without the lease: a client that lost its proof — a fresh process, a dropped
// recovery record — still knows who it is. Against an incumbent carrying its own instance id
// the honest answer is "your other connection is still registered", which resolves itself once
// the relay notices that socket. Blocking here strands the single-app case forever.
if (isPtyConsumerOwnerSameClient(hello, authentication, current)) {
throwOwnerError(
"PTY session owner is held by this client's own earlier connection",
PTY_CONSUMER_OWNER_HELD_SELF_ERROR
)
}
throwOwnerError(
'PTY session owner is held by an attached connection',
PTY_CONSUMER_OWNER_HELD_ATTACHED_ERROR
)
}
this.clampDisconnectedOwnerGrace(current)
throwOwnerError(
'PTY session owner is held by a disconnected connection within its grace period',
PTY_CONSUMER_OWNER_HELD_DISCONNECTED_ERROR
)
}
// Why only a peer-closed disconnect may shorten this: the floor is a bet that the incumbent is gone,
// and the relay tears a client's socket down for its own reasons too — a full lane queue is the
// signature of an owner that is alive but not draining fast enough. No owner completes a reconnect
// ladder in 250 ms, so clamping on a teardown we initiated hands a live owner's admission away and
// it can never get it back. Expiring a record never stops the remote PTY, but it does cost the user
// every route back to it.
private clampDisconnectedOwnerGrace(current: OwnerRecord): void {
if (current.disconnectCause !== 'peer-closed') {
return
}
const floorStart =
this.now() - Math.max(this.ownerGraceMs - PTY_CONSUMER_OWNER_HELD_GRACE_FLOOR_MS, 0)
if ((current.disconnectedAt ?? 0) <= floorStart) {
return
}
this.owner = { ...current, disconnectedAt: floorStart }
}
private displacedOwnerFor(
owner: OwnerRecord | null
): Readonly<PtyConsumerDisplacedOwner> | undefined {