diff --git a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-connection.ts b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-connection.ts index b6a000fc50e..87e962b215d 100644 --- a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-connection.ts +++ b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-connection.ts @@ -11,6 +11,7 @@ import type { TerminalMultiplexSubscribeFrame } from './stream-schemas' import type { TerminalMultiplexStream } from './terminal-stream-types' +import type { MultiplexPtyWaitAdmission } from './terminal-multiplex-pty-wait-admission' import type { TerminalStreamEndVerdict } from '../../../../../shared/terminal-stream-end-verdict' export type MultiplexSubscribeRequest = z.infer @@ -28,6 +29,7 @@ export type TerminalMultiplexConnectionBase = { cursor: number streams: Map pendingPtyWaitControllers: Map> + ptyWaitAdmission: MultiplexPtyWaitAdmission ackTotalInFlightBytes: number ackTotalWindowBytes: number ackFlushCursorStreamId: number | null diff --git a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-method.ts b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-method.ts index 8f72390c4a4..bd54b343a91 100644 --- a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-method.ts +++ b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-method.ts @@ -1,6 +1,10 @@ import { defineStreamingMethod } from '../../core' import { TerminalStreamOpcode } from '../../../../../shared/terminal-stream-protocol' -import { TERMINAL_MULTIPLEX_ACK_TOTAL_INITIAL_WINDOW_BYTES } from '../../../../../shared/terminal-multiplex-flow-control' +import { + TERMINAL_MULTIPLEX_ACK_TOTAL_INITIAL_WINDOW_BYTES, + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION, + TERMINAL_MULTIPLEX_MAX_QUEUED_PTY_WAITS_PER_CONNECTION +} from '../../../../../shared/terminal-multiplex-flow-control' import { TerminalMultiplex } from '../../../../../shared/rpc-contract/terminal-stream-params' import type { TerminalMultiplexConnection } from './terminal-multiplex-connection' import type { TerminalMultiplexStream } from './terminal-stream-types' @@ -9,6 +13,7 @@ import { installMultiplexFlowControl } from './terminal-multiplex-flow-control' import { installMultiplexCleanup } from './terminal-multiplex-cleanup' import { installMultiplexSlotFrames } from './terminal-multiplex-slot-frames' import { installMultiplexSubscribeFrame } from './terminal-multiplex-subscribe-frame' +import { createMultiplexPtyWaitAdmission } from './terminal-multiplex-pty-wait-admission' export const TERMINAL_MULTIPLEX_METHODS = [ defineStreamingMethod({ @@ -40,6 +45,10 @@ export const TERMINAL_MULTIPLEX_METHODS = [ cursor: 0, streams: new Map(), pendingPtyWaitControllers: new Map>(), + ptyWaitAdmission: createMultiplexPtyWaitAdmission( + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION, + TERMINAL_MULTIPLEX_MAX_QUEUED_PTY_WAITS_PER_CONNECTION + ), ackTotalInFlightBytes: 0, ackTotalWindowBytes: TERMINAL_MULTIPLEX_ACK_TOTAL_INITIAL_WINDOW_BYTES, ackFlushCursorStreamId: null, diff --git a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-pty-wait-admission.ts b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-pty-wait-admission.ts new file mode 100644 index 00000000000..52ba0798f46 --- /dev/null +++ b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-pty-wait-admission.ts @@ -0,0 +1,57 @@ +export type MultiplexPtyWaitAdmission = { + /** False when both the running slots and the queue are full. */ + hasRoom: () => boolean + /** Resolves true once a running slot is held, or false if `signal` aborts first. */ + acquire: (signal: AbortSignal) => Promise + release: () => void +} + +// Why FIFO: a restore subscribes panes in layout order, and queued waits must not start their PTY timeout until admitted. +export function createMultiplexPtyWaitAdmission( + maxRunning: number, + maxQueued: number +): MultiplexPtyWaitAdmission { + let running = 0 + const queue: { admit: () => void }[] = [] + + const admitNext = (): void => { + while (running < maxRunning && queue.length > 0) { + running += 1 + queue.shift()?.admit() + } + } + + return { + hasRoom: () => running < maxRunning || queue.length < maxQueued, + acquire: (signal) => { + if (signal.aborted) { + return Promise.resolve(false) + } + if (running < maxRunning && queue.length === 0) { + running += 1 + return Promise.resolve(true) + } + return new Promise((resolve) => { + const entry = { + admit: (): void => { + signal.removeEventListener('abort', onAbort) + resolve(true) + } + } + const onAbort = (): void => { + const index = queue.indexOf(entry) + if (index !== -1) { + queue.splice(index, 1) + } + resolve(false) + } + signal.addEventListener('abort', onAbort, { once: true }) + queue.push(entry) + }) + }, + release: () => { + running = Math.max(0, running - 1) + admitNext() + } + } +} diff --git a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-subscribe-resolution.ts b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-subscribe-resolution.ts index b920cf05c50..ac9d15ed640 100644 --- a/src/main/runtime/rpc/methods/terminal/terminal-multiplex-subscribe-resolution.ts +++ b/src/main/runtime/rpc/methods/terminal/terminal-multiplex-subscribe-resolution.ts @@ -1,7 +1,6 @@ import { TerminalStreamOpcode } from '../../../../../shared/terminal-stream-protocol' import { TERMINAL_MULTIPLEX_MAX_ACTIVE_STREAMS_PER_CONNECTION, - TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION, TERMINAL_MULTIPLEX_STREAM_LIMIT_ERROR } from '../../../../../shared/terminal-multiplex-flow-control' import type { @@ -36,7 +35,14 @@ export function resolveMultiplexSubscribePty( state: TerminalMultiplexConnection, request: MultiplexSubscribeRequest ): string | null | Promise { - const { runtime, pendingPtyWaitControllers, registerBinaryStreamHandler, signal, emit } = state + const { + runtime, + pendingPtyWaitControllers, + ptyWaitAdmission, + registerBinaryStreamHandler, + signal, + emit + } = state state.detachStream(request.streamId, null) state.cancelPendingPtyWaits(request.streamId) @@ -52,14 +58,12 @@ export function resolveMultiplexSubscribePty( if (leaf?.ptyId || !request.client) { return finalizeResolvedMultiplexPty(state, request, leaf) } - if (pendingPtyWaitControllers.size >= TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION) { + if (!ptyWaitAdmission.hasRoom()) { state.sendStreamError(request.streamId, TERMINAL_MULTIPLEX_STREAM_LIMIT_ERROR) emit({ type: 'end', streamId: request.streamId }) return null } - // Why: a never-mounted tab has no graph leaf to await; mounting the exact tab attaches its PTY without activating the worktree. - runtime.requestRendererTerminalTabMount(request.terminal) const waitController = new AbortController() const pendingControllers = pendingPtyWaitControllers.get(request.streamId) ?? new Set() pendingControllers.add(waitController) @@ -75,7 +79,14 @@ export function resolveMultiplexSubscribePty( } }) return (async () => { + let admitted = false try { + admitted = await ptyWaitAdmission.acquire(waitController.signal) + if (!admitted || waitController.signal.aborted) { + return null + } + // Why: a never-mounted tab has no graph leaf to await; mounting the exact tab attaches its PTY without activating the worktree. + runtime.requestRendererTerminalTabMount(request.terminal) const ptyId = await runtime.waitForLeafPtyId(request.terminal, 10_000, waitController.signal) leaf = { ptyId } } catch { @@ -84,6 +95,9 @@ export function resolveMultiplexSubscribePty( } // Fall through to the explicit no_connected_pty error below. } finally { + if (admitted) { + ptyWaitAdmission.release() + } const currentControllers = pendingPtyWaitControllers.get(request.streamId) currentControllers?.delete(waitController) if (currentControllers?.size === 0) { diff --git a/src/main/runtime/rpc/terminal-multiplex-pty-wait-capacity.test.ts b/src/main/runtime/rpc/terminal-multiplex-pty-wait-capacity.test.ts index bbd4c7eca5e..244c1255761 100644 --- a/src/main/runtime/rpc/terminal-multiplex-pty-wait-capacity.test.ts +++ b/src/main/runtime/rpc/terminal-multiplex-pty-wait-capacity.test.ts @@ -355,7 +355,7 @@ describe('terminal multiplex RPC', () => { expect(viewSubscriberCount).toBe(0) }) - it('reserves PTY wait capacity independently from active streams', async () => { + it('queues PTY waits past the running cap independently from active streams', async () => { const activeStreamCount = 44 const waitSignals: AbortSignal[] = [] const resolveWaits: ((ptyId: string) => void)[] = [] @@ -419,32 +419,28 @@ describe('terminal multiplex RPC', () => { await vi.waitFor(() => expect(waitSignals).toHaveLength(TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION) ) - const rejectedStreamId = - activeStreamCount + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + 1 - await vi.waitFor(() => { - const results = harness.messages.map((message) => JSON.parse(message).result) - expect(results).toContainEqual({ - type: 'error', - streamId: rejectedStreamId, - message: 'terminal_stream_limit_exceeded' - }) - expect(results).toContainEqual({ type: 'end', streamId: rejectedStreamId }) - }) + // Why: the wait past the running cap queues instead of failing the restore (STA-2886). + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(waitSignals).toHaveLength(TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION) + resolveWaits[0]?.('pty-pending-1') + await vi.waitFor(() => + expect(waitSignals).toHaveLength(TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + 1) + ) for (const [index, resolve] of resolveWaits.entries()) { resolve(`pty-pending-${index + 1}`) } await vi.waitFor(() => { const results = harness.messages.map((message) => JSON.parse(message).result) expect(results.filter((result) => result?.type === 'subscribed')).toHaveLength( - activeStreamCount + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + activeStreamCount + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + 1 ) expect( results.filter( (result) => result?.type === 'error' && result.message === 'terminal_stream_limit_exceeded' ) - ).toHaveLength(1) + ).toHaveLength(0) }) harness.registry.cleanupSubscription('terminal-multiplex:conn-desktop-first-paint') diff --git a/src/main/runtime/rpc/terminal-multiplex-pty-wait-queue.test.ts b/src/main/runtime/rpc/terminal-multiplex-pty-wait-queue.test.ts new file mode 100644 index 00000000000..a50c2dc32cd --- /dev/null +++ b/src/main/runtime/rpc/terminal-multiplex-pty-wait-queue.test.ts @@ -0,0 +1,134 @@ +import './unused-default-rpc-methods.test-fixture' +import { describe, expect, it, vi } from 'vitest' +import { + TerminalStreamOpcode, + decodeTerminalStreamFrame, + encodeTerminalStreamFrame, + encodeTerminalStreamJson +} from '../../../shared/terminal-stream-protocol' +import { + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION, + TERMINAL_MULTIPLEX_MAX_QUEUED_PTY_WAITS_PER_CONNECTION +} from '../../../shared/terminal-multiplex-flow-control' +import { createMultiplexPtyWaitAdmission } from './methods/terminal/terminal-multiplex-pty-wait-admission' +import { startDesktopMultiplexSubscribe } from './terminal-multiplex-test-harness' + +function startPendingPtyHarness() { + const waits = new Map void }>() + const requestRendererTerminalTabMount = vi.fn() + const harness = startDesktopMultiplexSubscribe({ + resolveLiveLeafForHandle: vi.fn().mockReturnValue({ ptyId: null }), + requestRendererTerminalTabMount, + waitForLeafPtyId: vi.fn( + (handle: string, _timeoutMs?: number, signal?: AbortSignal) => + new Promise((resolve, reject) => { + if (signal) { + waits.set(handle, { signal, resolve }) + } + signal?.addEventListener('abort', () => reject(new Error('request_aborted')), { + once: true + }) + }) + ) + }) + const send = (opcode: TerminalStreamOpcode, streamId: number, terminal?: string): void => { + const subscribe = opcode === TerminalStreamOpcode.Subscribe + harness.handlers.get(subscribe ? 0 : streamId)?.( + decodeTerminalStreamFrame( + encodeTerminalStreamFrame({ + opcode, + streamId: subscribe ? 0 : streamId, + seq: streamId, + payload: subscribe + ? encodeTerminalStreamJson({ + streamId, + terminal, + client: { id: 'desktop-1', type: 'desktop' }, + capabilities: { ackOutput: 1 } + }) + : new Uint8Array() + }) + )! + ) + } + const results = (): { type?: string; streamId?: number; message?: string }[] => + harness.messages.map((message) => JSON.parse(message).result) + return { harness, waits, requestRendererTerminalTabMount, send, results } +} + +describe('terminal multiplex PTY wait queue', () => { + it('rejects only past the queue bound and admits queued waits in order', async () => { + const { harness, waits, requestRendererTerminalTabMount, send, results } = + startPendingPtyHarness() + await vi.waitFor(() => expect(results().some((r) => r?.type === 'ready')).toBe(true)) + const total = + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + + TERMINAL_MULTIPLEX_MAX_QUEUED_PTY_WAITS_PER_CONNECTION + for (let streamId = 1; streamId <= total + 1; streamId += 1) { + send(TerminalStreamOpcode.Subscribe, streamId, `pane-${streamId}`) + } + await vi.waitFor(() => + expect(results()).toContainEqual({ + type: 'error', + streamId: total + 1, + message: 'terminal_stream_limit_exceeded' + }) + ) + expect( + results().filter((r) => r?.type === 'error' && r.message === 'terminal_stream_limit_exceeded') + ).toHaveLength(1) + await vi.waitFor(() => + expect(waits.size).toBe(TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION) + ) + // Why: a queued pane's mount is part of its paced admission, not fired at subscribe. + expect(requestRendererTerminalTabMount).toHaveBeenCalledTimes( + TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + ) + + waits.get('pane-1')?.resolve('pty-1') + const next = `pane-${TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + 1}` + await vi.waitFor(() => expect(waits.has(next)).toBe(true)) + await vi.waitFor(() => + expect(results()).toContainEqual(expect.objectContaining({ type: 'subscribed', streamId: 1 })) + ) + + harness.registry.cleanupSubscription('terminal-multiplex:conn-desktop-first-paint') + await vi.waitFor(() => expect(waits.get(next)?.signal.aborted).toBe(true)) + await harness.dispatchPromise + }) + + it('drops a queued wait on unsubscribe without spending a running slot', async () => { + const { harness, waits, send, results } = startPendingPtyHarness() + await vi.waitFor(() => expect(results().some((r) => r?.type === 'ready')).toBe(true)) + const cap = TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION + for (let streamId = 1; streamId <= cap + 2; streamId += 1) { + send(TerminalStreamOpcode.Subscribe, streamId, `pane-${streamId}`) + } + await vi.waitFor(() => expect(waits.size).toBe(cap)) + send(TerminalStreamOpcode.Unsubscribe, cap + 1) + waits.get('pane-1')?.resolve('pty-1') + + await vi.waitFor(() => expect(waits.has(`pane-${cap + 2}`)).toBe(true)) + expect(waits.has(`pane-${cap + 1}`)).toBe(false) + + harness.registry.cleanupSubscription('terminal-multiplex:conn-desktop-first-paint') + await harness.dispatchPromise + }) +}) + +describe('createMultiplexPtyWaitAdmission', () => { + it('admits FIFO, skips aborted waiters, and reports room against both bounds', async () => { + const admission = createMultiplexPtyWaitAdmission(1, 2) + const first = await admission.acquire(new AbortController().signal) + expect(first).toBe(true) + const aborted = new AbortController() + const second = admission.acquire(aborted.signal) + const third = admission.acquire(new AbortController().signal) + expect(admission.hasRoom()).toBe(false) + aborted.abort() + await expect(second).resolves.toBe(false) + expect(admission.hasRoom()).toBe(true) + admission.release() + await expect(third).resolves.toBe(true) + }) +}) diff --git a/src/shared/terminal-multiplex-flow-control.ts b/src/shared/terminal-multiplex-flow-control.ts index 33187e6b195..cbb520071b4 100644 --- a/src/shared/terminal-multiplex-flow-control.ts +++ b/src/shared/terminal-multiplex-flow-control.ts @@ -12,4 +12,7 @@ export const TERMINAL_MULTIPLEX_RECOVERY_SCROLLBACK_ROWS = 1000 // 128 covers large paired clients while fixed ACK windows and per-stream queues bound pressure. export const TERMINAL_MULTIPLEX_MAX_ACTIVE_STREAMS_PER_CONNECTION = 128 export const TERMINAL_MULTIPLEX_MAX_PENDING_PTY_WAITS_PER_CONNECTION = 32 +// Why: waits past the running cap queue instead of failing, so restoring more panes than 32 never errors; sized to the active-stream cap. +export const TERMINAL_MULTIPLEX_MAX_QUEUED_PTY_WAITS_PER_CONNECTION = + TERMINAL_MULTIPLEX_MAX_ACTIVE_STREAMS_PER_CONNECTION export const TERMINAL_MULTIPLEX_STREAM_LIMIT_ERROR = 'terminal_stream_limit_exceeded'