diff --git a/src/main/ssh/ssh-connection.test.ts b/src/main/ssh/ssh-connection.test.ts index 50fbea0eabe..9433ece028d 100644 --- a/src/main/ssh/ssh-connection.test.ts +++ b/src/main/ssh/ssh-connection.test.ts @@ -167,7 +167,8 @@ import { writeFileViaSystemSsh } from './ssh-system-fallback' import { getRemoteHostPlatform } from './ssh-remote-platform' -import type { SshTarget } from '../../shared/ssh-types' +import { CONNECT_TIMEOUT_MS } from './ssh-connection-utils' +import { MIN_SSH_RELAY_GRACE_PERIOD_SECONDS, type SshTarget } from '../../shared/ssh-types' function createTarget(overrides?: Partial): SshTarget { return { @@ -534,6 +535,61 @@ describe('SshConnection', () => { } }) + it('retries a saturated flap streak while the remote relay is still in grace', async () => { + vi.useFakeTimers() + try { + const conn = new SshConnection(createTarget(), createCallbacks()) + const connected = conn.connect() + await vi.advanceTimersByTimeAsync(1) + await connected + + // Saturate the delay ladder on flaps alone; 45s reconnects each drop while staying under + // STABLE_CONNECTION_MS, so the ladder never resets to the head. + for (let drop = 0; drop < 12; drop++) { + emitSshEvent('close') + await vi.advanceTimersByTimeAsync(45_000) + } + expect(conn.getState().status).toBe('connected') + + const before = clientInstances.length + emitSshEvent('close') + // The retry plus a worst-case handshake must land inside the shortest configurable relay grace, + // or the remote daemon shuts down and takes every PTY on that host with it. + await vi.advanceTimersByTimeAsync( + MIN_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 - CONNECT_TIMEOUT_MS - 1 + ) + expect(clientInstances.length).toBeGreaterThan(before) + } finally { + vi.useRealTimers() + } + }) + + it('logs the delay ladder position separately from the failure streak', async () => { + vi.useFakeTimers() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + const conn = new SshConnection(createTarget(), createCallbacks()) + const connected = conn.connect() + await vi.advanceTimersByTimeAsync(1) + await connected + + for (let drop = 0; drop < 12; drop++) { + emitSshEvent('close') + await vi.advanceTimersByTimeAsync(45_000) + } + + const lastReconnectLog = warn.mock.calls + .map((call) => String(call[0])) + .findLast((line) => line.includes('Reconnecting to')) + // A saturated flap ladder must not read like the connection is one step from giving up. + expect(lastReconnectLog).toContain('delay step 9/9') + expect(lastReconnectLog).toContain('failed handshakes 0/9') + } finally { + warn.mockRestore() + vi.useRealTimers() + } + }) + it('keeps retrying when a flap streak is followed by one handshake failure', async () => { vi.useFakeTimers() try { diff --git a/src/main/ssh/ssh-connection.ts b/src/main/ssh/ssh-connection.ts index 070494050d9..603f583de58 100644 --- a/src/main/ssh/ssh-connection.ts +++ b/src/main/ssh/ssh-connection.ts @@ -1241,7 +1241,7 @@ export class SshConnection { this.state.reconnectAttempt = decision.attemptIndex this.setState('reconnecting') console.warn( - `[ssh] Reconnecting to ${this.target.label} in ${decision.delayMs}ms (step ${decision.attemptIndex + 1}/${RECONNECT_BACKOFF_MS.length})` + `[ssh] Reconnecting to ${this.target.label} in ${decision.delayMs}ms (delay step ${decision.attemptIndex + 1}/${RECONNECT_BACKOFF_MS.length}, failed handshakes ${this.reconnectLadder.failedAttemptStreak}/${RECONNECT_BACKOFF_MS.length})` ) this.reconnectTimer = setTimeout(async () => { this.reconnectTimer = null diff --git a/src/main/ssh/ssh-reconnect-ladder.test.ts b/src/main/ssh/ssh-reconnect-ladder.test.ts index 48e94df5a20..994ea698fb1 100644 --- a/src/main/ssh/ssh-reconnect-ladder.test.ts +++ b/src/main/ssh/ssh-reconnect-ladder.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' -import { RECONNECT_BACKOFF_MS } from './ssh-connection-utils' -import { SshReconnectLadder, STABLE_CONNECTION_MS } from './ssh-reconnect-ladder' +import { MIN_SSH_RELAY_GRACE_PERIOD_SECONDS } from '../../shared/ssh-types' +import { CONNECT_TIMEOUT_MS, RECONNECT_BACKOFF_MS } from './ssh-connection-utils' +import { FLAP_DELAY_CAP_MS, SshReconnectLadder, STABLE_CONNECTION_MS } from './ssh-reconnect-ladder' describe('SshReconnectLadder', () => { it('climbs the table across consecutive post-handshake drops', () => { @@ -11,27 +12,72 @@ describe('SshReconnectLadder', () => { ladder.markConnected(60) expect(ladder.next(110)).toEqual({ kind: 'retry', delayMs: 2000, attemptIndex: 1 }) + // Table steps past FLAP_DELAY_CAP_MS are clamped so the retry still fits inside the relay grace floor. const rest = [5000, 5000, 10000, 10000, 10000, 30000, 30000] let now = 200 rest.forEach((delayMs, index) => { ladder.markConnected(now) now += 50 - expect(ladder.next(now)).toEqual({ kind: 'retry', delayMs, attemptIndex: index + 2 }) + expect(ladder.next(now)).toEqual({ + kind: 'retry', + delayMs: Math.min(delayMs, FLAP_DELAY_CAP_MS), + attemptIndex: index + 2 + }) }) // Saturates at the last step instead of going terminal. for (let i = 0; i < 3; i++) { ladder.markConnected(now) now += 50 - expect(ladder.next(now)).toEqual({ kind: 'retry', delayMs: 30000, attemptIndex: 8 }) + expect(ladder.next(now)).toEqual({ + kind: 'retry', + delayMs: Math.min(30000, FLAP_DELAY_CAP_MS), + attemptIndex: 8 + }) } }) + it('caps every flap delay so a full handshake still beats the shortest relay grace', () => { + const ladder = new SshReconnectLadder() + let now = 0 + + for (let i = 0; i < RECONNECT_BACKOFF_MS.length + 4; i++) { + ladder.markConnected(now) + now += 50 + const decision = ladder.next(now) + expect(decision.kind).toBe('retry') + if (decision.kind === 'retry') { + expect(decision.delayMs + CONNECT_TIMEOUT_MS).toBeLessThan( + MIN_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 + ) + now += decision.delayMs + } + } + }) + + it('keeps the uncapped table for a host whose handshakes fail', () => { + const ladder = new SshReconnectLadder() + const delays: number[] = [] + + for (let i = 0; i < RECONNECT_BACKOFF_MS.length - 1; i++) { + ladder.markAttemptFailed() + const decision = ladder.next(i * 1000) + if (decision.kind === 'retry') { + delays.push(decision.delayMs) + } + expect(ladder.failedAttemptStreak).toBe(i + 1) + } + + expect(delays).toEqual(RECONNECT_BACKOFF_MS.slice(0, -1)) + expect(delays.some((delayMs) => delayMs > FLAP_DELAY_CAP_MS)).toBe(true) + }) + it('escalates a post-handshake drop identically to a handshake failure', () => { const dropLadder = new SshReconnectLadder() const failureLadder = new SshReconnectLadder() const dropDelays: number[] = [] - const failureDelays: number[] = [] + const dropIndexes: number[] = [] + const failureIndexes: number[] = [] let now = 0 for (let i = 0; i < RECONNECT_BACKOFF_MS.length; i++) { @@ -43,16 +89,18 @@ describe('SshReconnectLadder', () => { expect(failureDecision.kind).toBe('retry') if (dropDecision.kind === 'retry') { dropDelays.push(dropDecision.delayMs) + dropIndexes.push(dropDecision.attemptIndex) } if (failureDecision.kind === 'retry') { - failureDelays.push(failureDecision.delayMs) + failureIndexes.push(failureDecision.attemptIndex) } failureLadder.markAttemptFailed() now += 50 } - expect(dropDelays).toEqual(RECONNECT_BACKOFF_MS) - expect(failureDelays).toEqual(dropDelays) + // Same ladder position; the flap path only clamps the wall-clock delay to the grace budget. + expect(dropIndexes).toEqual(failureIndexes) + expect(dropDelays).toEqual(RECONNECT_BACKOFF_MS.map((d) => Math.min(d, FLAP_DELAY_CAP_MS))) }) it('never reaches give-up on a flap streak, even with a later handshake failure', () => { diff --git a/src/main/ssh/ssh-reconnect-ladder.ts b/src/main/ssh/ssh-reconnect-ladder.ts index d635c5225b6..00c937bb0eb 100644 --- a/src/main/ssh/ssh-reconnect-ladder.ts +++ b/src/main/ssh/ssh-reconnect-ladder.ts @@ -1,8 +1,21 @@ -import { RECONNECT_BACKOFF_MS } from './ssh-connection-utils' +import { MIN_SSH_RELAY_GRACE_PERIOD_SECONDS } from '../../shared/ssh-types' +import { CONNECT_TIMEOUT_MS, RECONNECT_BACKOFF_MS } from './ssh-connection-utils' // A connection that held this long is treated as healthy, so the next drop restarts the delay ladder. export const STABLE_CONNECTION_MS = 60_000 +/** + * Why: a flap leaves the remote relay in grace, and the shortest grace a target can configure is + * MIN_SSH_RELAY_GRACE_PERIOD_SECONDS. Retrying later than this cap would let even a full-length + * handshake (CONNECT_TIMEOUT_MS) finish after grace expires, so the relay would shut down and take + * every remote PTY with it. Raising RECONNECT_BACKOFF_MS past this bound therefore changes nothing. + */ +export const FLAP_DELAY_CAP_MS = (() => { + const budgetMs = MIN_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 - CONNECT_TIMEOUT_MS + const fitting = RECONNECT_BACKOFF_MS.filter((delayMs) => delayMs < budgetMs) + return fitting.length > 0 ? Math.max(...fitting) : RECONNECT_BACKOFF_MS[0] +})() + export type SshReconnectDecision = | { kind: 'retry'; delayMs: number; attemptIndex: number } | { kind: 'give-up' } @@ -30,7 +43,18 @@ export class SshReconnectLadder { } const attemptIndex = Math.min(this.delayIndex, RECONNECT_BACKOFF_MS.length - 1) this.delayIndex = attemptIndex + 1 - return { kind: 'retry', delayMs: RECONNECT_BACKOFF_MS[attemptIndex], attemptIndex } + const tableDelayMs = RECONNECT_BACKOFF_MS[attemptIndex] + // Why: only a flap (no failed handshake) means the remote relay is alive and burning grace; a dead + // host keeps the full table because nothing over there is waiting for us. + const delayMs = + this.consecutiveFailedAttempts === 0 + ? Math.min(tableDelayMs, FLAP_DELAY_CAP_MS) + : tableDelayMs + return { kind: 'retry', delayMs, attemptIndex } + } + + get failedAttemptStreak(): number { + return this.consecutiveFailedAttempts } markConnected(nowMs: number): void { diff --git a/src/relay/agent-hook-envelope-publication.test.ts b/src/relay/agent-hook-envelope-publication.test.ts index 83b6594ad26..2be4d2f0301 100644 --- a/src/relay/agent-hook-envelope-publication.test.ts +++ b/src/relay/agent-hook-envelope-publication.test.ts @@ -10,19 +10,34 @@ type BoundedClient = { closes: number options: RelayClientSinkOptions write: (data: Buffer) => boolean + /** Mimics a socket whose kernel buffer is full: writes are accepted but never settle. */ + blocked: boolean + drain: () => void } function makeBoundedClient(highWaterMark: number): BoundedClient { + const drainListeners = new Set<() => void>() const client: BoundedClient = { frames: [], closes: 0, + blocked: false, write: (data: Buffer) => { client.frames.push(Buffer.from(data)) - return true + return !client.blocked + }, + drain: () => { + for (const listener of Array.from(drainListeners)) { + drainListeners.delete(listener) + listener() + } }, options: { writableHighWaterMark: () => highWaterMark, writableLength: () => 0, + waitWriteDrain: (callback: () => void) => { + drainListeners.add(callback) + return () => drainListeners.delete(callback) + }, close: () => { client.closes++ } @@ -31,16 +46,38 @@ function makeBoundedClient(highWaterMark: number): BoundedClient { return client } +/** Fills the producer queue so the next publish is rejected for backpressure, not for frame size. */ +function saturateProducerQueue(dispatcher: RelayDispatcher, client: BoundedClient): void { + client.blocked = true + // Shrinking chunks close the headroom a rejected 8 KB frame leaves behind, so even a small + // hook envelope is refused. + for (const size of [8_000, 512, 16]) { + const data = 'x'.repeat(size) + let filled = false + for (let attempt = 0; attempt < 5_000 && !filled; attempt++) { + filled = !dispatcher.tryNotifyPtyData({ id: 'pty-1', data }) + } + if (!filled) { + throw new Error('producer queue never filled') + } + } +} + function decodePayload(frame: Buffer): Record { const length = frame.readUInt32BE(9) return JSON.parse(frame.subarray(13, 13 + length).toString('utf-8')) } -function decodeEnvelopes(client: BoundedClient): AgentHookRelayEnvelope[] { +/** `shedFields` is relay-emitted wire metadata; the typed envelope does not declare it yet. */ +type PublishedEnvelope = AgentHookRelayEnvelope & { shedFields?: string[] } + +function decodeEnvelopes(client: BoundedClient): PublishedEnvelope[] { + // Keepalive frames carry no payload; only regular frames decode as JSON-RPC. return client.frames + .filter((frame) => frame.readUInt32BE(9) > 0) .map((frame) => decodePayload(frame)) .filter((msg) => msg.method === AGENT_HOOK_NOTIFICATION_METHOD) - .map((msg) => msg.params as unknown as AgentHookRelayEnvelope) + .map((msg) => msg.params as unknown as PublishedEnvelope) } function makeSubagents(count: number): AgentSubagentSnapshot[] { @@ -197,18 +234,38 @@ describe('publishAgentHookEnvelope', () => { } }) - it('measures the frame budget once up front and once per real shed', () => { + it('never probes the budget for a single sink that accepts the frame', () => { const primary = makeBoundedClient(16384) const dispatcher = new RelayDispatcher(primary.write, primary.options) const probe = vi.spyOn(dispatcher, 'producerEnvelopeBudget') try { + // One sink writes nothing when it rejects, so the publish attempt is itself the measurement. publishAgentHookEnvelope(dispatcher, makeEnvelope({ lastAssistantMessage: 128 })) - expect(probe).toHaveBeenCalledTimes(1) + expect(probe).not.toHaveBeenCalled() - probe.mockClear() // Two of the three ladder steps are absent here; skipping them must not cost an extra encode. publishAgentHookEnvelope(dispatcher, makeEnvelope({ interactivePrompt: 20_000 })) + expect(probe).toHaveBeenCalledTimes(1) + expect(decodeEnvelopes(primary)).toHaveLength(2) + } finally { + probe.mockRestore() + dispatcher.dispose() + } + }) + + it('measures a fan-out up front so both sinks get the same payload', () => { + const primary = makeBoundedClient(65536) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + const secondary = makeBoundedClient(16384) + const probe = vi.spyOn(dispatcher, 'producerEnvelopeBudget') + try { + dispatcher.attachClient(secondary.write, secondary.options) + publishAgentHookEnvelope(dispatcher, makeEnvelope({ lastAssistantMessage: 20_000 })) + // Probing first is what keeps the larger sink from being written a frame we then have to shed. expect(probe).toHaveBeenCalledTimes(2) + for (const client of [primary, secondary]) { + expect(decodeEnvelopes(client)).toHaveLength(1) + } } finally { probe.mockRestore() dispatcher.dispose() @@ -290,3 +347,162 @@ describe('publishAgentHookEnvelope', () => { } }) }) + +describe('publishAgentHookEnvelope shed marker', () => { + it('names every shed field so a consumer can tell a shed roster from an absent one', () => { + const primary = makeBoundedClient(16384) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + try { + publishAgentHookEnvelope( + dispatcher, + makeEnvelope({ lastAssistantMessage: 8_000, interactivePrompt: 6_000, subagents: 40 }) + ) + + const published = decodeEnvelopes(primary) + expect(published).toHaveLength(1) + // Without this marker the renderer overwrites a populated roster with undefined, blanking live + // subagent rows and unblocking hibernation for a pane that is still working. + expect(published[0].shedFields).toEqual(['lastAssistantMessage', 'subagents']) + expect(published[0].payload.subagents).toBeUndefined() + expect(published[0].payload.interactivePrompt).toBe('q'.repeat(6_000)) + } finally { + dispatcher.dispose() + } + }) + + it('omits the marker when nothing was shed and never stamps the caller envelope', () => { + const primary = makeBoundedClient(65536) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + try { + const envelope = makeEnvelope({ lastAssistantMessage: 128, subagents: 2 }) + publishAgentHookEnvelope(dispatcher, envelope) + expect(decodeEnvelopes(primary)[0].shedFields).toBeUndefined() + + const shedding = makeEnvelope({ lastAssistantMessage: 20_000, subagents: 2 }) + publishAgentHookEnvelope(dispatcher, shedding) + // The hook server replays this exact object, so the marker must live on the wire copy only. + expect((shedding as PublishedEnvelope).shedFields).toBeUndefined() + expect(shedding.payload.lastAssistantMessage).toBe('a'.repeat(20_000)) + } finally { + dispatcher.dispose() + } + }) +}) + +describe('publishAgentHookEnvelope redelivery', () => { + beforeEach(() => { + vi.useFakeTimers() + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('redelivers a backpressure-rejected envelope once the producer queue drains', () => { + const primary = makeBoundedClient(16384) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + try { + saturateProducerQueue(dispatcher, primary) + const envelope = makeEnvelope({}) + envelope.payload.state = 'done' + publishAgentHookEnvelope(dispatcher, envelope) + expect(decodeEnvelopes(primary)).toHaveLength(0) + + primary.blocked = false + primary.drain() + // Nothing in the pipeline republishes a hook envelope: it is fire-and-forget, and the only + // other delivery path is the reattach replay. + expect(decodeEnvelopes(primary)).toHaveLength(0) + + vi.advanceTimersByTime(250) + const published = decodeEnvelopes(primary) + expect(published).toHaveLength(1) + expect(published[0].payload.state).toBe('done') + expect(published[0].paneKey).toBe('tab-1:4f1b0f4e-0000-4000-8000-000000000001') + expect(primary.closes).toBe(0) + } finally { + dispatcher.dispose() + } + }) + + it('keeps only the newest snapshot per pane instead of queueing every drop', () => { + const primary = makeBoundedClient(16384) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + try { + saturateProducerQueue(dispatcher, primary) + for (const state of ['working', 'waiting', 'done'] as const) { + const envelope = makeEnvelope({}) + envelope.payload.state = state + publishAgentHookEnvelope(dispatcher, envelope) + } + + primary.blocked = false + primary.drain() + vi.advanceTimersByTime(1_000) + + const published = decodeEnvelopes(primary) + expect(published).toHaveLength(1) + expect(published[0].payload.state).toBe('done') + } finally { + dispatcher.dispose() + } + }) + + it('stops retrying a wedged client instead of pinning the payload forever', () => { + const primary = makeBoundedClient(16384) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + try { + const idleTimers = vi.getTimerCount() + saturateProducerQueue(dispatcher, primary) + publishAgentHookEnvelope(dispatcher, makeEnvelope({})) + expect(vi.getTimerCount()).toBeGreaterThan(idleTimers) + + vi.advanceTimersByTime(250 * 41) + expect(vi.getTimerCount()).toBe(idleTimers) + + primary.blocked = false + primary.drain() + vi.advanceTimersByTime(10_000) + expect(decodeEnvelopes(primary)).toHaveLength(0) + } finally { + dispatcher.dispose() + } + }) + + it('drops pending redeliveries when the dispatcher is disposed', () => { + const primary = makeBoundedClient(16384) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + const idleTimers = vi.getTimerCount() + saturateProducerQueue(dispatcher, primary) + publishAgentHookEnvelope(dispatcher, makeEnvelope({})) + + dispatcher.dispose() + primary.blocked = false + vi.advanceTimersByTime(10_000) + + expect(vi.getTimerCount()).toBeLessThanOrEqual(idleTimers) + expect(decodeEnvelopes(primary)).toHaveLength(0) + }) + + it('cancels the retry when a later snapshot for the same pane gets through', () => { + const primary = makeBoundedClient(16384) + const dispatcher = new RelayDispatcher(primary.write, primary.options) + try { + const idleTimers = vi.getTimerCount() + saturateProducerQueue(dispatcher, primary) + publishAgentHookEnvelope(dispatcher, makeEnvelope({})) + + primary.blocked = false + primary.drain() + const accepted = makeEnvelope({}) + accepted.payload.state = 'done' + publishAgentHookEnvelope(dispatcher, accepted) + expect(vi.getTimerCount()).toBe(idleTimers) + + vi.advanceTimersByTime(10_000) + expect(decodeEnvelopes(primary)).toHaveLength(1) + } finally { + dispatcher.dispose() + } + }) +}) diff --git a/src/relay/agent-hook-envelope-publication.ts b/src/relay/agent-hook-envelope-publication.ts index c5ffbb6fea5..8aa4b4558ef 100644 --- a/src/relay/agent-hook-envelope-publication.ts +++ b/src/relay/agent-hook-envelope-publication.ts @@ -7,39 +7,182 @@ import type { RelayDispatcher } from './dispatcher' // Why: shed the biggest, most reconstructible fields first — state/paneKey must survive or the pane // stays stuck on its last spinner, and an over-capacity frame is now dropped rather than sent. // Why: interactivePrompt is last on purpose — a blocking question is load-bearing (without its card -// the pane sits on `waiting` with no way to answer from web or mobile) while the roster is cosmetic. +// the pane sits on `waiting` with no way to answer from web or mobile). The roster is shed earlier +// only because it is cheaper to rebuild, NOT because it is cosmetic: a missing roster blanks live +// subagent rows and unblocks pane hibernation, which is why shed fields are named on the wire. const SHED_ORDER = ['lastAssistantMessage', 'subagents', 'interactivePrompt'] as const -function fitsProducerFrame(dispatcher: RelayDispatcher, envelope: AgentHookRelayEnvelope): boolean { - return ( - dispatcher.producerEnvelopeBudget( - AGENT_HOOK_NOTIFICATION_METHOD, - envelope as unknown as Record - ) >= 0 - ) +/** Wire marker listing the fields dropped to fit the frame, so a consumer can tell "shed" from + * "genuinely absent" and keep its current value instead of overwriting it with a gap. */ +const SHED_FIELDS_KEY = 'shedFields' + +// Why: these envelopes are fire-and-forget state snapshots — no ack, no producer-side retry — and +// the only other delivery path is the reattach replay. A queue-full drop would otherwise leave the +// pane on a stale spinner until the SSH link cycles, so a rejected snapshot is retried briefly. +const REDELIVERY_INTERVAL_MS = 250 +const MAX_REDELIVERY_ATTEMPTS = 40 +// Why: a wedged link must not pin snapshots for panes that are already gone; recency-evict the rest. +const MAX_PENDING_PANES = 64 + +type PendingEnvelope = { + params: Record + clientIds: readonly number[] + attemptsLeft: number +} + +type RedeliveryState = { + pending: Map + timer: ReturnType | null +} + +const redeliveryByDispatcher = new WeakMap() + +function fitsProducerFrame(dispatcher: RelayDispatcher, params: Record): boolean { + return dispatcher.producerEnvelopeBudget(AGENT_HOOK_NOTIFICATION_METHOD, params) >= 0 +} + +function toNotificationParams( + envelope: AgentHookRelayEnvelope, + shedFields: readonly string[] +): Record { + const params = envelope as unknown as Record + return shedFields.length === 0 ? params : { ...params, [SHED_FIELDS_KEY]: [...shedFields] } +} + +/** Returns the ids of the clients that rejected the frame; a rejected client is never written to. */ +function publishToClients( + dispatcher: RelayDispatcher, + params: Record, + clientIds: readonly number[] +): number[] { + const rejected: number[] = [] + for (const clientId of clientIds) { + if (!dispatcher.publishProducerNotification(clientId, AGENT_HOOK_NOTIFICATION_METHOD, params)) { + rejected.push(clientId) + } + } + return rejected } /** Publishes an agent-hook envelope, dropping optional detail fields until the frame fits the - * smallest attached sink. Publishes the trimmed envelope even if it still does not fit. */ + * smallest attached sink. A frame the sink queue merely has no room for is retried in the + * background; one that no shedding can fit is dropped. */ export function publishAgentHookEnvelope( dispatcher: RelayDispatcher, envelope: AgentHookRelayEnvelope ): void { + const clientIds = dispatcher.activeClientIds() + if (clientIds.length === 0) { + return + } + // Why: a fan-out must choose one payload before it writes anything, or the first client keeps a + // frame a smaller later client forces us to shed. A single sink writes nothing when it rejects, + // so there the attempt itself is the measurement — one encode instead of a probe plus a publish. + const measureBeforePublish = clientIds.length > 1 let candidate = envelope - // Why: re-encoding is the expensive part, so only re-measure after a field was actually dropped. - let fits = fitsProducerFrame(dispatcher, candidate) - for (const field of SHED_ORDER) { - if (fits) { - break + const shedFields: string[] = [] + let step = 0 + for (;;) { + const params = toNotificationParams(candidate, shedFields) + if (!measureBeforePublish || fitsProducerFrame(dispatcher, params)) { + const rejected = publishToClients(dispatcher, params, clientIds) + if (rejected.length === 0) { + clearPendingEnvelope(dispatcher, envelope.paneKey) + return + } + // Rejection is ambiguous: an over-capacity frame must shed, a full producer queue must wait. + if (measureBeforePublish || fitsProducerFrame(dispatcher, params)) { + setPendingEnvelope(dispatcher, envelope.paneKey, params, rejected) + return + } } - if (candidate.payload[field] === undefined) { - continue + while (step < SHED_ORDER.length && candidate.payload[SHED_ORDER[step]] === undefined) { + step += 1 } + if (step >= SHED_ORDER.length) { + // Nothing left to shed and it still does not fit: waiting cannot make it sendable, and this + // snapshot supersedes any pending one, which is now stale. + clearPendingEnvelope(dispatcher, envelope.paneKey) + return + } + const field = SHED_ORDER[step] + step += 1 // Why: the hook server caches envelopes and replays them after --connect, so shedding in place // would permanently strip the cached copy too. candidate = { ...candidate, payload: { ...candidate.payload } } delete candidate.payload[field] - fits = fitsProducerFrame(dispatcher, candidate) + shedFields.push(field) + } +} + +function setPendingEnvelope( + dispatcher: RelayDispatcher, + paneKey: string, + params: Record, + clientIds: readonly number[] +): void { + let state = redeliveryByDispatcher.get(dispatcher) + if (!state) { + state = { pending: new Map(), timer: null } + redeliveryByDispatcher.set(dispatcher, state) + // A disposed dispatcher can never accept the retry, and pending params pin the cached payload. + dispatcher.onDisposed(() => stopRedelivery(dispatcher)) + } + // Latest-wins: one snapshot per pane, never a growing queue. Delete-then-set keeps Map insertion + // order = recency so the cap below evicts the longest-idle pane. + state.pending.delete(paneKey) + state.pending.set(paneKey, { params, clientIds, attemptsLeft: MAX_REDELIVERY_ATTEMPTS }) + while (state.pending.size > MAX_PENDING_PANES) { + state.pending.delete(state.pending.keys().next().value as string) + } + if (!state.timer) { + const timer = setInterval(() => runRedeliveryPass(dispatcher), REDELIVERY_INTERVAL_MS) + if (typeof timer.unref === 'function') { + timer.unref() + } + state.timer = timer + } +} + +function clearPendingEnvelope(dispatcher: RelayDispatcher, paneKey: string): void { + const state = redeliveryByDispatcher.get(dispatcher) + if (!state?.pending.delete(paneKey) || state.pending.size > 0) { + return + } + stopRedelivery(dispatcher) +} + +function stopRedelivery(dispatcher: RelayDispatcher): void { + const state = redeliveryByDispatcher.get(dispatcher) + if (!state) { + return + } + state.pending.clear() + if (state.timer) { + clearInterval(state.timer) + state.timer = null + } +} + +function runRedeliveryPass(dispatcher: RelayDispatcher): void { + const state = redeliveryByDispatcher.get(dispatcher) + if (!state) { + return + } + const active = new Set(dispatcher.activeClientIds()) + for (const [paneKey, pending] of Array.from(state.pending)) { + // A client that detached gets caught up by the reattach replay, so stop tracking it. + const targets = pending.clientIds.filter((clientId) => active.has(clientId)) + const rejected = + targets.length === 0 ? [] : publishToClients(dispatcher, pending.params, targets) + pending.attemptsLeft -= 1 + if (rejected.length === 0 || pending.attemptsLeft <= 0) { + state.pending.delete(paneKey) + continue + } + pending.clientIds = rejected + } + if (state.pending.size === 0) { + stopRedelivery(dispatcher) } - dispatcher.notify(AGENT_HOOK_NOTIFICATION_METHOD, candidate as unknown as Record) } diff --git a/src/relay/dispatcher-capacity-degradation.test.ts b/src/relay/dispatcher-capacity-degradation.test.ts index c41dcd4d860..5cc96ff88fe 100644 --- a/src/relay/dispatcher-capacity-degradation.test.ts +++ b/src/relay/dispatcher-capacity-degradation.test.ts @@ -1,7 +1,26 @@ import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest' -import { RelayDispatcher, type RelayClientSinkOptions } from './dispatcher' +import { + RelayDispatcher, + type RelayClientSinkOptions, + type SinkWriteSettlement +} from './dispatcher' +import type * as ProtocolModule from './protocol' import { encodeJsonRpcFrame } from './protocol' +// Counts every frame encode (including the estimate-only ones) so a redundant re-encode is observable. +const encodeCalls = vi.hoisted(() => ({ count: 0 })) + +vi.mock('./protocol', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + encodeJsonRpcFrame: (...args: Parameters) => { + encodeCalls.count++ + return actual.encodeJsonRpcFrame(...args) + } + } +}) + type BoundedClient = { frames: Buffer[] closes: number @@ -90,22 +109,41 @@ describe('RelayDispatcher bounded-capacity degradation', () => { 12288 - envelopeBytes ) expect(bounded.producerEnvelopeBudget('fs.changed', params, 999)).toBe( - Number.MAX_SAFE_INTEGER + Number.MIN_SAFE_INTEGER ) } finally { bounded.dispose() } }) + it('producerEnvelopeBudget reports no capacity for a named client that is gone', () => { + const primary = makeBoundedClient(65536) + const secondary = makeBoundedClient(65536) + const bounded = new RelayDispatcher(primary.write, primary.options) + try { + const secondaryId = bounded.attachClient(secondary.write, secondary.options) + const params = { events: ['x'.repeat(64)] } + expect(bounded.producerEnvelopeBudget('fs.changed', params, secondaryId)).toBeGreaterThan(0) + + // A client that detached between activeClientIds() and publication must never look like it fits. + bounded.detachClient(secondaryId) + expect(bounded.producerEnvelopeBudget('fs.changed', params, secondaryId)).toBeLessThan(0) + expect(bounded.producerEnvelopeBudget('fs.changed', params, 999)).toBeLessThan(0) + // The no-target broadcast case still reports unbounded room. + expect(bounded.producerEnvelopeBudget('fs.changed', params)).toBeGreaterThan(0) + } finally { + bounded.dispose() + } + }) + it('producerDataBudget keeps its pre-delegation value for a pty.data envelope', () => { const primary = makeBoundedClient(65536) const bounded = new RelayDispatcher(primary.write, primary.options) try { // Snapshot of the shipped formula: 49152-byte producer capacity minus the 96-byte empty-data frame. expect(bounded.producerDataBudget('pty.data', { ptyId: 'pty-7', seq: 42 })).toBe(49056) - expect(bounded.producerDataBudget('pty.data', { ptyId: 'pty-7', seq: 42 }, 999)).toBe( - Number.MAX_SAFE_INTEGER - ) + // A vanished client leaves no room for data at all. + expect(bounded.producerDataBudget('pty.data', { ptyId: 'pty-7', seq: 42 }, 999)).toBe(0) } finally { bounded.dispose() } @@ -241,6 +279,50 @@ describe('RelayDispatcher bounded-capacity degradation', () => { } }) + it('publishProducerNotification reports its drops like notify', () => { + const primary = makeBoundedClient(65536) + const secondary = makeBoundedClient(16384) + const stderr = vi.spyOn(process.stderr, 'write').mockReturnValue(true) + const bounded = new RelayDispatcher(primary.write, primary.options) + try { + const secondaryId = bounded.attachClient(secondary.write, secondary.options) + const flood = { events: ['x'.repeat(20_000)] } + + expect(bounded.publishProducerNotification(secondaryId, 'fs.changed', flood)).toBe(false) + expect(stderr).toHaveBeenCalledTimes(1) + expect(stderr.mock.calls[0][0]).toMatch( + /^\[relay\] Dropped fs\.changed \(\d+B > producer frame capacity 12288B\)\n$/ + ) + + // Same one-line-per-generation budget as notify(). + expect(bounded.publishProducerNotification(secondaryId, 'fs.changed', flood)).toBe(false) + expect(stderr).toHaveBeenCalledTimes(1) + } finally { + stderr.mockRestore() + bounded.dispose() + } + }) + + it('does not re-encode a dropped frame when its log line is suppressed', () => { + const primary = makeBoundedClient(16384) + const stderr = vi.spyOn(process.stderr, 'write').mockReturnValue(true) + const bounded = new RelayDispatcher(primary.write, primary.options) + try { + const flood = { events: ['x'.repeat(20_000)] } + bounded.notify('fs.changed', flood) + expect(stderr).toHaveBeenCalledTimes(1) + + // The suppressed drop must size the frame once, not once to publish and again to log. + encodeCalls.count = 0 + bounded.notify('fs.changed', flood) + expect(stderr).toHaveBeenCalledTimes(1) + expect(encodeCalls.count).toBe(1) + } finally { + stderr.mockRestore() + bounded.dispose() + } + }) + it('sendResponse substitutes a JSON-RPC error instead of closing when the legacy lane rejects', async () => { const primary = makeBoundedClient(65536) const bounded = new RelayDispatcher(primary.write, primary.options) @@ -263,6 +345,30 @@ describe('RelayDispatcher bounded-capacity degradation', () => { } }) + it('settles the response fence as failed when the substitute error replaces the result', async () => { + const primary = makeBoundedClient(65536) + const bounded = new RelayDispatcher(primary.write, primary.options) + try { + const settlements: SinkWriteSettlement[] = [] + bounded.onRequest('fs.listFiles', async (_params, context) => { + context.onResponseSettled?.((result) => settlements.push(result)) + return { paths: 'x'.repeat(3 * 1024 * 1024) } + }) + bounded.feed(encodeJsonRpcFrame({ jsonrpc: '2.0', id: 79, method: 'fs.listFiles' }, 1, 0)) + await vi.advanceTimersByTimeAsync(0) + + expect(primary.closes).toBe(0) + expect(settlements).toHaveLength(1) + // The peer got a capacity error, not the result — the fence must not claim delivery. + expect(settlements[0]).toEqual({ + ok: false, + error: new Error('Relay response exceeded the bounded transport capacity') + }) + } finally { + bounded.dispose() + } + }) + it('sendResponse leaves an already-failed response rejected without closing', async () => { const primary = makeBoundedClient(65536) const bounded = new RelayDispatcher(primary.write, primary.options) diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 77a4e6867a7..a51869f9d8d 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -77,6 +77,8 @@ type RelayClient = { // grow it inside one generation — cap it well above the fixed relay method vocabulary. const DROPPED_NOTIFICATION_LOG_KEY_LIMIT = 64 +const RESPONSE_OVER_CAPACITY_MESSAGE = 'Relay response exceeded the bounded transport capacity' + type DroppedProducerNotificationLog = { generation: number loggedKeys: Set @@ -385,10 +387,19 @@ export class RelayDispatcher { params: Record, clientId?: number ): number { - const targets = - clientId === undefined - ? this.activeClients() - : [this.clients.get(clientId)].filter((client): client is RelayClient => !!client) + if (clientId !== undefined) { + const client = this.clients.get(clientId) + // Why: a detached or closed target has no room at all — reporting infinite capacity passes a + // fits-check and the frame is then dropped by the publish seam instead. + if (!client || client.closed) { + return Number.MIN_SAFE_INTEGER + } + return ( + client.writer.producerFrameCapacity - + this.estimateFrameBytes({ jsonrpc: '2.0', method, params }) + ) + } + const targets = this.activeClients() if (targets.length === 0) { return Number.MAX_SAFE_INTEGER } @@ -433,19 +444,20 @@ export class RelayDispatcher { method, ...(params !== undefined ? { params } : {}) } + const frameBytes = this.estimateFrameBytes(msg) this.runPublicationTransaction(() => { for (const client of this.clients.values()) { if (client.closed) { continue } if (method === 'pty.replay') { - this.enqueueFrame(client, msg, 'control') + this.enqueueFrame(client, msg, 'control', undefined, frameBytes) continue } // Why: closing can never make an oversized frame sendable — the producer regenerates it after // reattach and re-kills the link, turning a recoverable drop into an endless reconnect loop. - if (!this.publishToClient(client, msg, 'ordinary')) { - this.logDroppedProducerNotification(client, method, msg) + if (!this.publishToClient(client, msg, 'ordinary', undefined, frameBytes)) { + this.logDroppedProducerNotification(client, method, frameBytes) } } }) @@ -465,25 +477,23 @@ export class RelayDispatcher { if (!client || client.closed) { return false } - return this.publishToClient( - client, - { - jsonrpc: '2.0', - method, - ...(params !== undefined ? { params } : {}) - }, - 'ordinary' - ) + const msg: JsonRpcNotification = { + jsonrpc: '2.0', + method, + ...(params !== undefined ? { params } : {}) + } + const frameBytes = this.estimateFrameBytes(msg) + if (this.publishToClient(client, msg, 'ordinary', undefined, frameBytes)) { + return true + } + // Why: same diagnostics as notify() — a producer that drops here must not do so silently. + this.logDroppedProducerNotification(client, method, frameBytes) + return false } // Why: one line per generation, method and drop reason — a flooding producer retries every batch and would // spam stderr, but a transient queue-full drop must not consume the slot a real over-capacity drop needs. - private logDroppedProducerNotification( - client: RelayClient, - method: string, - msg: JsonRpcNotification - ): void { - const bytes = this.estimateFrameBytes(msg) + private logDroppedProducerNotification(client: RelayClient, method: string, bytes: number): void { const capacity = client.writer.producerFrameCapacity const overCapacity = bytes > capacity const key = `${method}:${overCapacity ? 'over-capacity' : 'queue-full'}` @@ -900,10 +910,17 @@ export class RelayDispatcher { { jsonrpc: '2.0', id, - error: { code: -32010, message: 'Relay response exceeded the bounded transport capacity' } + error: { code: -32010, message: RESPONSE_OVER_CAPACITY_MESSAGE } }, 'control', - onSettled + // Why: writing the substitute is not delivering the result — a settlement fence must never read + // the capacity error's successful write as "the peer received your result". + (settlement) => + onSettled( + settlement.ok + ? { ok: false, error: new Error(RESPONSE_OVER_CAPACITY_MESSAGE) } + : settlement + ) ) } @@ -1027,9 +1044,11 @@ export class RelayDispatcher { client: RelayClient, msg: JsonRpcNotification, lane: 'interactive' | 'ordinary' | 'fixed-bulk' | 'bulk', - onSettled: (result: SinkWriteSettlement) => void = () => {} + onSettled: (result: SinkWriteSettlement) => void = () => {}, + // Why: broadcast callers size the frame once for every client; avoid a redundant encode. + estimatedBytes?: number ): boolean { - const bytes = this.estimateFrameBytes(msg) + const bytes = estimatedBytes ?? this.estimateFrameBytes(msg) const fixedBlocked = lane === 'fixed-bulk' && (client.writer.retainedProducerBytes > 0 || bytes > client.writer.fixedFrameCapacity) diff --git a/src/relay/fs-handler-file-read.ts b/src/relay/fs-handler-file-read.ts index e56cdbdaa0a..452bf10c8f0 100644 --- a/src/relay/fs-handler-file-read.ts +++ b/src/relay/fs-handler-file-read.ts @@ -1,8 +1,14 @@ import { open, readFile, stat } from 'node:fs/promises' +import type { FileHandle } from 'node:fs/promises' import { extname } from 'node:path' import type { RelayDispatcher, RequestContext } from './dispatcher' -import { STREAM_ACK_WINDOW_CHUNKS, STREAM_CHUNK_SIZE, RelayErrorCode } from './protocol' -import type { RelayStreamRegistry, TooManyStreamsError } from './fs-stream-registry' +import { + MAX_CONCURRENT_STREAMS, + STREAM_ACK_WINDOW_CHUNKS, + STREAM_CHUNK_SIZE, + RelayErrorCode +} from './protocol' +import { TooManyStreamsError, type RelayStreamRegistry } from './fs-stream-registry' import { BINARY_PROBE_BYTES, IMAGE_MIME_TYPES, @@ -102,12 +108,17 @@ export async function readRelayFileStreamMetadata( return { totalSize: 0, isBinary: true, empty: true } } - const handle = await open(filePath, 'r') + // Why: reserved before the fd opens so a refusal costs nothing, and released only + // once the terminal frame settles — see reserveTerminalFrameSlot. + const releaseTerminalFrameSlot = reserveTerminalFrameSlot(registry) + let handle: FileHandle | undefined let streamId: number try { + handle = await open(filePath, 'r') streamId = registry.register(handle) } catch (err) { - await handle.close() + await handle?.close() + releaseTerminalFrameSlot() throw err } @@ -118,7 +129,15 @@ export async function readRelayFileStreamMetadata( // sees the response before the first chunk frame. const resolvedPumpOptions = pumpOptions ?? { paceWithAcks: false } setImmediate(() => { - void pumpChunks(streamId, stats.size, dispatcher, registry, context, resolvedPumpOptions) + void pumpChunks( + streamId, + stats.size, + dispatcher, + registry, + context, + resolvedPumpOptions, + releaseTerminalFrameSlot + ) }) return { @@ -132,16 +151,46 @@ export async function readRelayFileStreamMetadata( } } +/** + * Terminal frames (fs.streamEnd/fs.streamError) ride the control lane, which does not + * drop on overflow — it destroys the link at 256 queued frames / 1 MB. A stream's + * registry slot is gone the moment its last chunk is read, so on a socket that is not + * draining, back-to-back reads could stack one undelivered terminal frame each until + * that budget blew. Holding the slot until the frame settles keeps the number of queued + * terminal frames at MAX_CONCURRENT_STREAMS, far below the killing threshold; the + * overflow now costs one refused read (TooManyStreams, which clients already handle) + * instead of the whole connection. + */ +const pendingTerminalFrames = new WeakMap() + +function reserveTerminalFrameSlot(registry: RelayStreamRegistry): () => void { + const pending = pendingTerminalFrames.get(registry) ?? 0 + if (pending >= MAX_CONCURRENT_STREAMS) { + throw new TooManyStreamsError() + } + pendingTerminalFrames.set(registry, pending + 1) + let released = false + return () => { + if (released) { + return + } + released = true + pendingTerminalFrames.set(registry, (pendingTerminalFrames.get(registry) ?? 1) - 1) + } +} + async function pumpChunks( streamId: number, totalSize: number, dispatcher: RelayDispatcher, registry: RelayStreamRegistry, context: RequestContext, - pumpOptions: StreamPumpOptions + pumpOptions: StreamPumpOptions, + releaseTerminalFrameSlot: () => void ): Promise { const entry = registry.get(streamId) if (!entry) { + releaseTerminalFrameSlot() return } const buffer = Buffer.allocUnsafe(STREAM_CHUNK_SIZE) @@ -150,6 +199,7 @@ async function pumpChunks( let endReason: 'end' | 'aborted' | 'stale' | 'error' = 'end' let errorCode: string | null = null let errorMessage: string | null = null + let slotReleaseDeferred = false try { try { @@ -226,14 +276,22 @@ async function pumpChunks( } try { - // Why: the terminal frame must not be droppable by the producer-lane capacity check, and the - // per-chunk await already settled every chunk so a control frame cannot overtake stream data. + // Why: a dropped terminal frame hangs the reader forever, so it takes the control lane, which + // never drops — but that lane KILLS the link when it overflows, hence the reserved slot held + // until this frame settles. The per-chunk await already settled every chunk, so the control + // frame cannot overtake stream data. const publishTerminal = (method: string, params: Record): void => { - if (pumpOptions.clientId !== undefined) { - dispatcher.notifyClient(pumpOptions.clientId, method, params) - } else { + if (pumpOptions.clientId === undefined) { + // Legacy broadcast path (direct calls/tests): no per-frame settlement to hold the slot on. dispatcher.notifyControl(method, params) + return } + slotReleaseDeferred = dispatcher.tryNotifyClient( + pumpOptions.clientId, + method, + params, + releaseTerminalFrameSlot + ) } if (endReason === 'end') { publishTerminal('fs.streamEnd', { streamId }) @@ -256,7 +314,12 @@ async function pumpChunks( ) } } finally { + // Why: the fd goes back first — a terminal frame that can never be delivered must not + // strand it. Cancelled/stale streams publish nothing, so nothing else frees their slot. await registry.release(streamId) + if (!slotReleaseDeferred) { + releaseTerminalFrameSlot() + } } } diff --git a/src/relay/fs-handler-stream.test.ts b/src/relay/fs-handler-stream.test.ts index 4e1a3330c3d..d387879327e 100644 --- a/src/relay/fs-handler-stream.test.ts +++ b/src/relay/fs-handler-stream.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest' import { FsHandler } from './fs-handler' import { RelayContext } from './context' import type { RelayDispatcher } from './dispatcher' -import { STREAM_CHUNK_SIZE } from './protocol' +import { MAX_CONCURRENT_STREAMS, STREAM_CHUNK_SIZE } from './protocol' import * as fs from 'node:fs/promises' import * as path from 'node:path' import { mkdtempSync, writeFileSync } from 'node:fs' @@ -29,6 +29,8 @@ function createMockDispatcher() { const droppedProducerFrames: Notification[] = [] // Mirrors the real lane split: a saturated producer queue drops frames, control frames still land. let producerSaturated = false + let holdControlSettlement = false + const heldControlSettlements: (() => void)[] = [] return { onRequest: vi.fn( ( @@ -61,6 +63,25 @@ function createMockDispatcher() { notifyClient: vi.fn((clientId: number, method: string, params?: Record) => { notifications.push({ method, params, lane: 'control', clientId }) }), + // Mirrors the real control lane: the frame is queued on enqueue, but settles only once the + // sink actually wrote it — holding settlement models a socket that is not draining. + tryNotifyClient: vi.fn( + ( + clientId: number, + method: string, + params?: Record, + onSettled?: (result: { ok: boolean }) => void + ): boolean => { + notifications.push({ method, params, lane: 'control', clientId }) + const settle = () => onSettled?.({ ok: true }) + if (holdControlSettlement) { + heldControlSettlements.push(settle) + } else { + settle() + } + return true + } + ), notifyControl: vi.fn((method: string, params?: Record) => { notifications.push({ method, params, lane: 'control' }) }), @@ -69,6 +90,15 @@ function createMockDispatcher() { saturateProducerLane() { producerSaturated = true }, + holdControlSettlements() { + holdControlSettlement = true + }, + settleHeldControlFrames() { + holdControlSettlement = false + for (const settle of heldControlSettlements.splice(0)) { + settle() + } + }, callRequest( method: string, params: Record = {}, @@ -421,6 +451,39 @@ describe('FsHandler readFileStream', () => { expect(err).toBeNull() }) + it('caps undelivered terminal frames so they cannot overflow the link-killing control lane', async () => { + const filePath = path.join(tmpDir, 'terminal-budget.png') + writeFileSync(filePath, Buffer.alloc(64 * 1024, 0x42)) // one chunk per stream + // A socket that accepts nothing: control frames queue, and the real writer destroys the + // client once 256 of them pile up. + dispatcher.holdControlSettlements() + + const registry = (handler as unknown as { streamRegistry: { size(): number } }).streamRegistry + const context = { clientId: 5, isStale: () => false } + const terminalFrames = () => + dispatcher._notifications.filter( + (n) => n.method === 'fs.streamEnd' || n.method === 'fs.streamError' + ) + + for (let i = 0; i < MAX_CONCURRENT_STREAMS; i++) { + await dispatcher.callRequest('fs.readFileStream', { filePath }, context) + await waitFor(() => terminalFrames().length === i + 1) + } + // The fd must go back even though its terminal frame is still undelivered. + await waitFor(() => registry.size() === 0) + + await expect( + dispatcher.callRequest('fs.readFileStream', { filePath }, context) + ).rejects.toThrow(/Too many concurrent streams/) + await flush(10) + expect(terminalFrames()).toHaveLength(MAX_CONCURRENT_STREAMS) + + // Draining the sink settles the queued frames and hands the budget back. + dispatcher.settleHeldControlFrames() + await dispatcher.callRequest('fs.readFileStream', { filePath }, context) + await waitFor(() => terminalFrames().length === MAX_CONCURRENT_STREAMS + 1) + }) + it('rejects the 17th concurrent stream with TooManyStreams', async () => { const paths: string[] = [] for (let i = 0; i < 17; i++) { diff --git a/src/relay/relay-watcher-event-emitter.test.ts b/src/relay/relay-watcher-event-emitter.test.ts index 3bf42296a52..8b82eefcfeb 100644 --- a/src/relay/relay-watcher-event-emitter.test.ts +++ b/src/relay/relay-watcher-event-emitter.test.ts @@ -1,7 +1,11 @@ +import { readFileSync } from 'node:fs' +import { fileURLToPath } from 'node:url' import { describe, expect, it } from 'vitest' import type { WatcherProcessEvent } from '../main/ipc/parcel-watcher-process-protocol' import { RelayDispatcher } from './dispatcher' +import { DEFAULT_PRODUCER_QUEUE_MAX_BYTES } from './dispatcher-writer-admission' import type { RelayClientSinkOptions, RelayClientWrite } from './dispatcher-writer-sink' +import { LEGACY_CLIENT_RETAINED_BYTES_LOW } from './legacy-relay-publication-ledger' import { HEADER_LENGTH, parseJsonRpcMessage } from './protocol' import { emitRelayWatcherEvents, emitRelayWatcherOverflow } from './relay-watcher-event-emitter' @@ -98,6 +102,21 @@ function saturateProducerQueue(dispatcher: RelayDispatcher): void { expect(admitted).toBeLessThan(5_000) } +// Parks PTY frames in the producer queue of a stalled sink without filling it. +function parkProducerBytes(dispatcher: RelayDispatcher, bytes: number): number { + const chunk = 'x'.repeat(40_000) + let parked = 0 + while (parked < bytes && dispatcher.tryNotifyPtyData({ paneId: 'pane', data: chunk })) { + parked += chunk.length + } + expect(parked).toBeGreaterThanOrEqual(bytes) + return parked +} + +function isOverflowFrame(frame: Buffer): boolean { + return frameEvents(frame).some((event) => event.kind === 'overflow') +} + describe('relay watcher writer admission', () => { it('keeps watcher batches on the producer lane while overflow markers ride the control lane', () => { const sink = createRecordingSink() @@ -124,6 +143,38 @@ describe('relay watcher writer admission', () => { }) }) +describe('relay watcher overflow suppression key', () => { + it('keeps the source free of NUL bytes so git and grep still see text', () => { + const source = readFileSync( + fileURLToPath(new URL('./relay-watcher-event-emitter.ts', import.meta.url)) + ) + expect(source.includes(0)).toBe(false) + }) + + it('scopes the outstanding marker per client as well as per root', () => { + const primary = createRecordingSink(65536) + const secondary = createRecordingSink(65536) + primary.blockNextWrite() + secondary.blockNextWrite() + const dispatcher = new RelayDispatcher(primary.write, primary.options) + + try { + dispatcher.attachClient(secondary.write, secondary.options) + // Two clients, same root: a key that collapses them would silently desync the second tree. + emitRelayWatcherOverflow(dispatcher, '/workspace', false) + primary.drainWaiters.shift()?.() + secondary.drainWaiters.shift()?.() + + for (const sink of [primary, secondary]) { + expect(sink.frames.filter((frame) => frameMethod(frame) === 'fs.changed')).toHaveLength(1) + expect(sink.closes()).toBe(0) + } + } finally { + dispatcher.dispose() + } + }) +}) + describe('relay watcher batch chunking', () => { it.each([ [65536, 49152], @@ -301,6 +352,52 @@ describe('relay watcher batch chunking', () => { } }) + it('stops the chunk walk while the producer queue sits past its retention reserve', () => { + const sink = createRecordingSink(65536) + sink.blockNextWrite() + const dispatcher = new RelayDispatcher(sink.write, sink.options) + + try { + // Past the retention reserve but far from the 2 MB queue limit: every chunk would still be admitted. + const parked = parkProducerBytes(dispatcher, LEGACY_CLIENT_RETAINED_BYTES_LOW) + expect(parked).toBeLessThan(DEFAULT_PRODUCER_QUEUE_MAX_BYTES / 2 + 256 * 1024) + + emitRelayWatcherEvents(dispatcher, '/workspace', false, watcherBatch(5000)) + sink.drainWaiters.shift()?.() + + const watcherFrames = sink.frames.filter((frame) => frameMethod(frame) === 'fs.changed') + expect(watcherFrames.map(isOverflowFrame)).toEqual([true]) + // The flood injected nothing: the interactive lane keeps the whole reserve instead of ~425 KB less. + const injected = watcherFrames + .filter((frame) => !isOverflowFrame(frame)) + .reduce((total, frame) => total + frame.length, 0) + expect(injected).toBe(0) + expect(injected).toBeLessThan( + DEFAULT_PRODUCER_QUEUE_MAX_BYTES - LEGACY_CLIENT_RETAINED_BYTES_LOW + ) + expect(sink.closes()).toBe(0) + } finally { + dispatcher.dispose() + } + }) + + it('still delivers a flood larger than the reserve while the sink keeps draining', () => { + const sink = createRecordingSink(16384) + const dispatcher = new RelayDispatcher(sink.write, sink.options) + const events = watcherBatch(20_000) + + try { + emitRelayWatcherEvents(dispatcher, '/workspace', false, events) + + const totalBytes = sink.frames.reduce((total, frame) => total + frame.length, 0) + expect(totalBytes).toBeGreaterThan(LEGACY_CLIENT_RETAINED_BYTES_LOW) + expect(sink.frames.some(isOverflowFrame)).toBe(false) + expect(sink.frames.flatMap(frameEvents)).toHaveLength(events.length) + } finally { + dispatcher.dispose() + } + }) + it('chunks per client so a small sink never degrades a healthy one', () => { const healthy = createRecordingSink(65536) const small = createRecordingSink(16384) diff --git a/src/relay/relay-watcher-event-emitter.ts b/src/relay/relay-watcher-event-emitter.ts index b23705975cd..3805e8b4872 100644 Binary files a/src/relay/relay-watcher-event-emitter.ts and b/src/relay/relay-watcher-event-emitter.ts differ