diff --git a/src/main/runtime/rpc/methods/terminal.ts b/src/main/runtime/rpc/methods/terminal.ts index 687f5f11e0e..f30e1b2f591 100644 --- a/src/main/runtime/rpc/methods/terminal.ts +++ b/src/main/runtime/rpc/methods/terminal.ts @@ -27,6 +27,9 @@ const TERMINAL_STREAM_CHUNK_BYTES = 48 * 1024 const TERMINAL_OUTPUT_FLUSH_MS = 5 // Why: output batches become binary stream payloads; byte size is the transport cost. const TERMINAL_OUTPUT_BATCH_MAX_BYTES = 64 * 1024 +// Why: remote clients can apply output pressure without pausing runtime PTY ingestion. +const TERMINAL_MULTIPLEX_ACK_STREAM_HIGH_WATER_BYTES = 512 * 1024 +const TERMINAL_MULTIPLEX_ACK_TOTAL_HIGH_WATER_BYTES = 2 * 1024 * 1024 // Why: pending output is held for later binary frames, so cap the encoded // payload bytes rather than UTF-16 code units. const TERMINAL_MULTIPLEX_PENDING_MAX_BYTES = 256 * 1024 @@ -68,7 +71,10 @@ type TerminalMultiplexStream = { ptyId: string client: TerminalViewportClient | undefined isMobile: boolean + ackOutput: boolean + ackInFlightBytes: number buffering: boolean + ackPendingOutput: TerminalOutputFrameChunk[] pendingOutput: TerminalOutputChunk[] pendingOutputBytes: number pendingOutputOverflowed: boolean @@ -538,7 +544,16 @@ const TerminalMultiplexSubscribeFrame = TerminalHandle.extend({ type: z.enum(['mobile', 'desktop']).default('desktop') }) .optional(), - viewport: TerminalViewport.optional() + viewport: TerminalViewport.optional(), + capabilities: z + .object({ + ackOutput: z.literal(1).optional() + }) + .optional() +}) + +const TerminalMultiplexAckFrame = z.object({ + bytes: z.number().int().nonnegative() }) const TerminalMultiplexSnapshotRequestFrame = z.object({ @@ -880,6 +895,7 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [ let closed = false let cursor = 0 const streams = new Map() + let ackTotalInFlightBytes = 0 let resolveMultiplex = (): void => {} const multiplexClosed = new Promise((resolve) => { resolveMultiplex = resolve @@ -906,6 +922,63 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [ sendFrame(streamId, TerminalStreamOpcode.Error, encodeTerminalStreamText(message)) emit({ type: 'error', streamId, message }) } + const canSendAckGatedOutput = (stream: TerminalMultiplexStream, bytes: number): boolean => { + if (!stream.ackOutput) { + return true + } + return ( + stream.ackInFlightBytes + bytes <= TERMINAL_MULTIPLEX_ACK_STREAM_HIGH_WATER_BYTES && + ackTotalInFlightBytes + bytes <= TERMINAL_MULTIPLEX_ACK_TOTAL_HIGH_WATER_BYTES + ) + } + const sendAckGatedOutput = ( + stream: TerminalMultiplexStream, + chunk: TerminalOutputFrameChunk + ): void => { + sendFrame(stream.streamId, TerminalStreamOpcode.Output, chunk.bytes, chunk.seq) + if (stream.ackOutput) { + stream.ackInFlightBytes += chunk.bytes.byteLength + ackTotalInFlightBytes += chunk.bytes.byteLength + } + } + const queueOrSendOutput = ( + stream: TerminalMultiplexStream, + chunk: TerminalOutputFrameChunk + ): void => { + if (closed || streams.get(stream.streamId) !== stream) { + return + } + if ( + stream.ackPendingOutput.length > 0 || + !canSendAckGatedOutput(stream, chunk.bytes.byteLength) + ) { + stream.ackPendingOutput.push(chunk) + return + } + sendAckGatedOutput(stream, chunk) + } + const flushAckPendingOutput = (stream: TerminalMultiplexStream): void => { + let flushed = 0 + while ( + flushed < stream.ackPendingOutput.length && + canSendAckGatedOutput(stream, stream.ackPendingOutput[flushed]!.bytes.byteLength) + ) { + sendAckGatedOutput(stream, stream.ackPendingOutput[flushed]!) + flushed += 1 + } + if (flushed > 0) { + stream.ackPendingOutput.splice(0, flushed) + } + } + const acknowledgeOutput = (stream: TerminalMultiplexStream, bytes: number): void => { + if (!stream.ackOutput || bytes <= 0) { + return + } + const acknowledged = Math.min(stream.ackInFlightBytes, bytes) + stream.ackInFlightBytes -= acknowledged + ackTotalInFlightBytes = Math.max(0, ackTotalInFlightBytes - acknowledged) + flushAckPendingOutput(stream) + } const detachStream = (streamId: number, emitEnd: boolean): void => { const stream = streams.get(streamId) if (!stream) { @@ -913,6 +986,9 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [ } stream.outputBatcher.flush() stream.outputBatcher.dispose() + ackTotalInFlightBytes = Math.max(0, ackTotalInFlightBytes - stream.ackInFlightBytes) + stream.ackInFlightBytes = 0 + stream.ackPendingOutput = [] stream.unsubscribeData() stream.unsubscribeResize() stream.unsubscribeFit() @@ -948,6 +1024,15 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [ detachStream(stream.streamId, false) return } + if (frame.opcode === TerminalStreamOpcode.Ack) { + const parsed = TerminalMultiplexAckFrame.safeParse( + decodeTerminalStreamJson(frame.payload) ?? {} + ) + if (parsed.success) { + acknowledgeOutput(stream, parsed.data.bytes) + } + return + } if (frame.opcode === TerminalStreamOpcode.Input) { const text = decodeTerminalStreamText(frame.payload) if (!text) { @@ -1110,13 +1195,16 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [ ptyId, client: request.client, isMobile, + ackOutput: request.capabilities?.ackOutput === 1, + ackInFlightBytes: 0, buffering: true, + ackPendingOutput: [], pendingOutput: [], pendingOutputBytes: 0, pendingOutputOverflowed: false, outputBatcher: createTerminalOutputBatcher((data, meta) => { for (const chunk of splitTerminalOutputFrameChunks(data, meta)) { - sendFrame(request.streamId, TerminalStreamOpcode.Output, chunk.bytes, chunk.seq) + queueOrSendOutput(stream, chunk) } }), unsubscribeData: () => {}, diff --git a/src/main/runtime/rpc/terminal-multiplex.test.ts b/src/main/runtime/rpc/terminal-multiplex.test.ts index 2376c36c43d..169ab9a3ed7 100644 --- a/src/main/runtime/rpc/terminal-multiplex.test.ts +++ b/src/main/runtime/rpc/terminal-multiplex.test.ts @@ -347,6 +347,148 @@ describe('terminal multiplex RPC', () => { } }) + it('holds ACK-capable multiplex output over budget until the client acknowledges bytes', async () => { + const messages: string[] = [] + const binaryFrames: Uint8Array[] = [] + const handlers = new Map< + number, + (frame: NonNullable>) => void + >() + const cleanups = new Map void>() + const dataListenerRef: { + current?: (data: string, meta?: { seq?: number; rawLength?: number }) => void + } = {} + const runtime = stubRuntime({ + resolveLeafForHandle: vi.fn().mockReturnValue({ ptyId: 'pty-1' }), + readTerminal: vi.fn().mockResolvedValue({ tail: [], truncated: false }), + serializeTerminalBuffer: vi.fn().mockResolvedValue({ + data: 'snapshot', + cols: 120, + rows: 40 + }), + getTerminalSize: vi.fn().mockReturnValue({ cols: 120, rows: 40 }), + getMobileDisplayMode: vi.fn().mockReturnValue('auto'), + getLayout: vi.fn().mockReturnValue({ seq: 1 }), + subscribeToTerminalData: vi.fn( + ( + _: string, + listener: (data: string, meta?: { seq?: number; rawLength?: number }) => void + ) => { + dataListenerRef.current = listener + return vi.fn() + } + ), + subscribeToTerminalResize: vi.fn().mockReturnValue(vi.fn()), + subscribeToFitOverrideChanges: vi.fn().mockReturnValue(vi.fn()), + subscribeToDriverChanges: vi.fn().mockReturnValue(vi.fn()), + getTerminalFitOverride: vi.fn().mockReturnValue(null), + getDriver: vi.fn().mockReturnValue({ kind: 'idle' }), + registerSubscriptionCleanup: vi.fn((id: string, cleanup: () => void) => { + cleanups.set(id, cleanup) + }), + cleanupSubscription: vi.fn((id: string) => { + const cleanup = cleanups.get(id) + cleanups.delete(id) + cleanup?.() + }), + waitForTerminal: vi.fn(() => new Promise(() => {})), + sendTerminal: vi.fn().mockResolvedValue({ accepted: true }), + updateDesktopViewport: vi.fn().mockResolvedValue(true) + }) + const dispatcher = new RpcDispatcher({ runtime, methods: TERMINAL_METHODS }) + + const dispatchPromise = dispatcher.dispatchStreaming( + makeRequest('terminal.multiplex', {}), + (msg) => messages.push(msg), + { + connectionId: 'conn-ack-gated', + sendBinary: (bytes) => binaryFrames.push(bytes), + registerBinaryStreamHandler: (streamId, handler) => { + handlers.set(streamId, handler) + return () => handlers.delete(streamId) + } + } + ) + + await vi.waitFor(() => + expect(messages.some((msg) => JSON.parse(msg).result?.type === 'ready')).toBe(true) + ) + handlers.get(0)?.( + decodeTerminalStreamFrame( + encodeTerminalStreamFrame({ + opcode: TerminalStreamOpcode.Subscribe, + streamId: 0, + seq: 1, + payload: encodeTerminalStreamJson({ + streamId: 16, + terminal: 'terminal-1', + client: { id: 'desktop-1', type: 'desktop' }, + viewport: { cols: 120, rows: 40 }, + capabilities: { ackOutput: 1 } + }) + }) + )! + ) + await vi.waitFor(() => + expect(messages.some((msg) => JSON.parse(msg).result?.type === 'subscribed')).toBe(true) + ) + binaryFrames.splice(0) + + const output = 'x'.repeat(700 * 1024) + dataListenerRef.current?.(output, { seq: output.length, rawLength: output.length }) + + const initialOutputFrames = binaryFrames + .map((frame) => decodeTerminalStreamFrame(frame)) + .filter((frame) => frame?.opcode === TerminalStreamOpcode.Output) + const initialBytes = initialOutputFrames.reduce( + (total, frame) => total + (frame?.payload.byteLength ?? 0), + 0 + ) + expect(initialBytes).toBeLessThanOrEqual(512 * 1024) + expect(initialOutputFrames.length).toBeGreaterThan(0) + const initialOutput = initialOutputFrames + .map((frame) => (frame ? decodeTerminalStreamText(frame.payload) : '')) + .join('') + expect(initialOutput.length).toBeLessThan(output.length) + + handlers.get(16)?.( + decodeTerminalStreamFrame( + encodeTerminalStreamFrame({ + opcode: TerminalStreamOpcode.Input, + streamId: 16, + seq: 2, + payload: encodeTerminalStreamText('still interactive\r') + }) + )! + ) + await vi.waitFor(() => + expect(runtime.sendTerminal).toHaveBeenCalledWith('terminal-1', { + text: 'still interactive\r', + enter: false, + interrupt: false + }) + ) + + handlers.get(16)?.( + decodeTerminalStreamFrame( + encodeTerminalStreamFrame({ + opcode: TerminalStreamOpcode.Ack, + streamId: 16, + seq: 3, + payload: encodeTerminalStreamJson({ bytes: initialBytes }) + }) + )! + ) + + const flushedOutputFrames = binaryFrames + .map((frame) => decodeTerminalStreamFrame(frame)) + .filter((frame) => frame?.opcode === TerminalStreamOpcode.Output) + expect(flushedOutputFrames.length).toBeGreaterThan(initialOutputFrames.length) + + runtime.cleanupSubscription('terminal-multiplex:conn-ack-gated') + await dispatchPromise + }) + it('marks multiplex fallback snapshots truncated when the uncursored read is limited', async () => { const messages: string[] = [] const binaryFrames: Uint8Array[] = [] diff --git a/src/renderer/src/runtime/remote-runtime-terminal-multiplexer.ts b/src/renderer/src/runtime/remote-runtime-terminal-multiplexer.ts index bd2f13fc6d2..0d70a8906a2 100644 --- a/src/renderer/src/runtime/remote-runtime-terminal-multiplexer.ts +++ b/src/renderer/src/runtime/remote-runtime-terminal-multiplexer.ts @@ -70,6 +70,7 @@ type RemoteRuntimeMultiplexedTerminalState = { streamId: number terminal: string callbacks: RemoteRuntimeMultiplexedTerminalCallbacks + acknowledgeOutput: boolean snapshotChunks: Uint8Array[] snapshotBytes: number snapshotOverflowed: boolean @@ -138,6 +139,7 @@ class RemoteRuntimeTerminalMultiplexer { streamId, terminal: args.terminal, callbacks: args.callbacks, + acknowledgeOutput: args.client.type === 'desktop', snapshotChunks: [], snapshotBytes: 0, snapshotOverflowed: false, @@ -181,7 +183,8 @@ class RemoteRuntimeTerminalMultiplexer { streamId, terminal: args.terminal, client: args.client, - viewport: args.viewport + viewport: args.viewport, + capabilities: args.client.type === 'desktop' ? { ackOutput: 1 } : undefined }) ) if (!sent) { @@ -329,10 +332,20 @@ class RemoteRuntimeTerminalMultiplexer { } if (frame.opcode === TerminalStreamOpcode.Output) { const data = decodeTerminalStreamText(frame.payload) - stream.callbacks.onData(data, { - seq: typeof frame.seq === 'number' && frame.seq > 0 ? frame.seq : undefined, - rawLength: data.length - }) + try { + stream.callbacks.onData(data, { + seq: typeof frame.seq === 'number' && frame.seq > 0 ? frame.seq : undefined, + rawLength: data.length + }) + } finally { + if (stream.acknowledgeOutput) { + this.sendFrame( + stream.streamId, + TerminalStreamOpcode.Ack, + encodeTerminalStreamJson({ bytes: frame.payload.byteLength }) + ) + } + } return } if (frame.opcode === TerminalStreamOpcode.SnapshotStart) { diff --git a/src/renderer/src/runtime/runtime-terminal-stream.test.ts b/src/renderer/src/runtime/runtime-terminal-stream.test.ts index 0d0a59212f0..820e59ea4a0 100644 --- a/src/renderer/src/runtime/runtime-terminal-stream.test.ts +++ b/src/renderer/src/runtime/runtime-terminal-stream.test.ts @@ -98,8 +98,12 @@ describe('remote runtime terminal data subscriptions', () => { const subscribeFrame = decodeTerminalStreamFrame(sendBinary.mock.calls[0][0]) expect(subscribeFrame?.opcode).toBe(TerminalStreamOpcode.Subscribe) const subscribePayload = - subscribeFrame && decodeTerminalStreamJson<{ streamId: number }>(subscribeFrame.payload) + subscribeFrame && + decodeTerminalStreamJson<{ streamId: number; capabilities?: { ackOutput?: 1 } }>( + subscribeFrame.payload + ) expect(subscribePayload?.streamId).toEqual(expect.any(Number)) + expect(subscribePayload?.capabilities).toEqual({ ackOutput: 1 }) callbacks?.onBinary?.( encodeTerminalStreamFrame({ @@ -111,6 +115,12 @@ describe('remote runtime terminal data subscriptions', () => { ) expect(watcher).toHaveBeenCalledWith('live') + const ackFrame = sendBinary.mock.calls + .slice(1) + .map((call) => decodeTerminalStreamFrame(call[0])) + .find((frame) => frame?.opcode === TerminalStreamOpcode.Ack) + expect(ackFrame?.streamId).toBe(subscribePayload!.streamId) + expect(ackFrame && decodeTerminalStreamJson(ackFrame.payload)).toEqual({ bytes: 4 }) expect(_getRemoteRuntimeTerminalMultiplexerCountForTest()).toBe(1) dispose() expect(unsubscribe).toHaveBeenCalled() diff --git a/src/shared/terminal-stream-protocol.test.ts b/src/shared/terminal-stream-protocol.test.ts index 626973b1051..3df867ef4b2 100644 --- a/src/shared/terminal-stream-protocol.test.ts +++ b/src/shared/terminal-stream-protocol.test.ts @@ -109,6 +109,21 @@ describe('terminal-stream-protocol', () => { expect(unsubscribe?.streamId).toBe(12) }) + it('round-trips output acknowledgement frames', () => { + const ack = decodeTerminalStreamFrame( + encodeTerminalStreamFrame({ + opcode: TerminalStreamOpcode.Ack, + streamId: 12, + seq: 4, + payload: encodeTerminalStreamJson({ bytes: 4096 }) + }) + ) + + expect(ack?.opcode).toBe(TerminalStreamOpcode.Ack) + expect(ack?.streamId).toBe(12) + expect(ack && decodeTerminalStreamJson(ack.payload)).toEqual({ bytes: 4096 }) + }) + it('rejects unknown frame versions and opcodes', () => { const encoded = encodeTerminalStreamFrame({ opcode: TerminalStreamOpcode.Output, diff --git a/src/shared/terminal-stream-protocol.ts b/src/shared/terminal-stream-protocol.ts index 9a993bee99a..57bb9fefe8d 100644 --- a/src/shared/terminal-stream-protocol.ts +++ b/src/shared/terminal-stream-protocol.ts @@ -13,7 +13,8 @@ export enum TerminalStreamOpcode { Resize = 8, Subscribe = 9, Unsubscribe = 10, - SnapshotRequest = 11 + SnapshotRequest = 11, + Ack = 12 } export type TerminalStreamFrame = { @@ -92,6 +93,7 @@ function isTerminalStreamOpcode(value: number): value is TerminalStreamOpcode { value === TerminalStreamOpcode.Resize || value === TerminalStreamOpcode.Subscribe || value === TerminalStreamOpcode.Unsubscribe || - value === TerminalStreamOpcode.SnapshotRequest + value === TerminalStreamOpcode.SnapshotRequest || + value === TerminalStreamOpcode.Ack ) }