diff --git a/packages/server/src/client/daemon-client-terminal-stream-manager.test.ts b/packages/server/src/client/daemon-client-terminal-stream-manager.test.ts new file mode 100644 index 000000000..9170a4c0a --- /dev/null +++ b/packages/server/src/client/daemon-client-terminal-stream-manager.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, test, vi } from "vitest"; +import { + TerminalStreamManager, + type TerminalStreamChunk, +} from "./daemon-client-terminal-stream-manager.js"; + +function createChunk(input: { + streamId: number; + offset: number; + data: string; +}): TerminalStreamChunk { + const bytes = new TextEncoder().encode(input.data); + return { + streamId: input.streamId, + offset: input.offset, + endOffset: input.offset + bytes.byteLength, + replay: false, + data: bytes, + }; +} + +describe("TerminalStreamManager", () => { + test("buffers chunks and flushes with ack when handler subscribes", () => { + const sendAck = vi.fn(); + const manager = new TerminalStreamManager({ sendAck }); + const seen: string[] = []; + + manager.receiveChunk({ + chunk: createChunk({ streamId: 7, offset: 4, data: "hello" }), + }); + expect(sendAck).not.toHaveBeenCalled(); + + manager.subscribe({ + streamId: 7, + handler: (chunk) => { + seen.push(new TextDecoder().decode(chunk.data)); + }, + }); + + expect(seen).toEqual(["hello"]); + expect(sendAck).toHaveBeenCalledTimes(1); + expect(sendAck).toHaveBeenCalledWith({ streamId: 7, offset: 9 }); + }); + + test("does not ack when every handler throws", () => { + const sendAck = vi.fn(); + const manager = new TerminalStreamManager({ sendAck }); + + manager.subscribe({ + streamId: 11, + handler: () => { + throw new Error("boom"); + }, + }); + + manager.receiveChunk({ + chunk: createChunk({ streamId: 11, offset: 0, data: "x" }), + }); + + expect(sendAck).not.toHaveBeenCalled(); + }); + + test("evicts oldest buffered chunks when max buffered chunk count is exceeded", () => { + const sendAck = vi.fn(); + const manager = new TerminalStreamManager({ + sendAck, + maxBufferedChunks: 2, + maxBufferedBytes: 1024, + }); + + manager.receiveChunk({ + chunk: createChunk({ streamId: 5, offset: 0, data: "A" }), + }); + manager.receiveChunk({ + chunk: createChunk({ streamId: 5, offset: 1, data: "B" }), + }); + manager.receiveChunk({ + chunk: createChunk({ streamId: 5, offset: 2, data: "C" }), + }); + + const seen: string[] = []; + manager.subscribe({ + streamId: 5, + handler: (chunk) => { + seen.push(new TextDecoder().decode(chunk.data)); + }, + }); + + expect(seen).toEqual(["B", "C"]); + expect(sendAck).toHaveBeenNthCalledWith(1, { streamId: 5, offset: 2 }); + expect(sendAck).toHaveBeenNthCalledWith(2, { streamId: 5, offset: 3 }); + }); + + test("tracks explicit ack offsets and skips stale auto-acks", () => { + const sendAck = vi.fn(); + const manager = new TerminalStreamManager({ sendAck }); + + manager.subscribe({ + streamId: 13, + handler: () => { + // no-op + }, + }); + manager.noteAck({ streamId: 13, offset: 10 }); + + manager.receiveChunk({ + chunk: createChunk({ streamId: 13, offset: 0, data: "abc" }), + }); + expect(sendAck).not.toHaveBeenCalled(); + + manager.receiveChunk({ + chunk: createChunk({ streamId: 13, offset: 11, data: "z" }), + }); + expect(sendAck).toHaveBeenCalledTimes(1); + expect(sendAck).toHaveBeenCalledWith({ streamId: 13, offset: 12 }); + }); +}); diff --git a/packages/server/src/client/daemon-client-terminal-stream-manager.ts b/packages/server/src/client/daemon-client-terminal-stream-manager.ts new file mode 100644 index 000000000..df683df81 --- /dev/null +++ b/packages/server/src/client/daemon-client-terminal-stream-manager.ts @@ -0,0 +1,180 @@ +export type TerminalStreamChunk = { + streamId: number; + offset: number; + endOffset: number; + replay: boolean; + data: Uint8Array; +}; + +type TerminalChunkHandler = (chunk: TerminalStreamChunk) => void; + +type BufferedTerminalStreamQueue = { + chunks: TerminalStreamChunk[]; + bytes: number; +}; + +type TerminalStreamAck = { + streamId: number; + offset: number; +}; + +export type TerminalStreamManagerConfig = { + sendAck: (ack: TerminalStreamAck) => void; + maxBufferedChunks?: number; + maxBufferedBytes?: number; +}; + +const DEFAULT_MAX_BUFFERED_CHUNKS = 2048; +const DEFAULT_MAX_BUFFERED_BYTES = 2 * 1024 * 1024; + +export class TerminalStreamManager { + private readonly handlers: Map> = new Map(); + private readonly bufferedChunks: Map = new Map(); + private readonly ackOffsets: Map = new Map(); + private readonly maxBufferedChunks: number; + private readonly maxBufferedBytes: number; + + constructor(private readonly config: TerminalStreamManagerConfig) { + this.maxBufferedChunks = + config.maxBufferedChunks ?? DEFAULT_MAX_BUFFERED_CHUNKS; + this.maxBufferedBytes = + config.maxBufferedBytes ?? DEFAULT_MAX_BUFFERED_BYTES; + } + + clearAll(): void { + this.handlers.clear(); + this.bufferedChunks.clear(); + this.ackOffsets.clear(); + } + + clearStream(input: { streamId: number }): void { + this.handlers.delete(input.streamId); + this.bufferedChunks.delete(input.streamId); + this.ackOffsets.delete(input.streamId); + } + + subscribe(input: { + streamId: number; + handler: TerminalChunkHandler; + }): () => void { + const { streamId, handler } = input; + if (!this.handlers.has(streamId)) { + this.handlers.set(streamId, new Set()); + } + + const streamHandlers = this.handlers.get(streamId)!; + streamHandlers.add(handler); + this.flushBufferedChunks({ streamId, handler }); + + return () => { + streamHandlers.delete(handler); + if (streamHandlers.size === 0) { + this.handlers.delete(streamId); + } + }; + } + + receiveChunk(input: { chunk: TerminalStreamChunk }): void { + const { chunk } = input; + const streamHandlers = this.handlers.get(chunk.streamId); + if (!streamHandlers || streamHandlers.size === 0) { + this.bufferChunk({ chunk }); + return; + } + + let delivered = false; + for (const handler of streamHandlers) { + try { + handler(chunk); + delivered = true; + } catch { + // no-op + } + } + if (delivered) { + this.maybeAckChunk({ + streamId: chunk.streamId, + endOffset: chunk.endOffset, + }); + } + } + + noteAck(input: TerminalStreamAck): void { + const normalizedOffset = Math.max(0, Math.floor(input.offset)); + const previousAck = this.ackOffsets.get(input.streamId) ?? -1; + if (normalizedOffset > previousAck) { + this.ackOffsets.set(input.streamId, normalizedOffset); + } + } + + private flushBufferedChunks(input: { + streamId: number; + handler: TerminalChunkHandler; + }): void { + const buffered = this.bufferedChunks.get(input.streamId); + if (!buffered || buffered.chunks.length === 0) { + return; + } + + for (const chunk of buffered.chunks) { + try { + input.handler(chunk); + this.maybeAckChunk({ + streamId: input.streamId, + endOffset: chunk.endOffset, + }); + } catch { + // no-op + } + } + this.bufferedChunks.delete(input.streamId); + } + + private bufferChunk(input: { chunk: TerminalStreamChunk }): void { + const queue = this.bufferedChunks.get(input.chunk.streamId) ?? { + chunks: [], + bytes: 0, + }; + + queue.chunks.push(input.chunk); + queue.bytes += input.chunk.data.byteLength; + + while ( + queue.chunks.length > this.maxBufferedChunks || + queue.bytes > this.maxBufferedBytes + ) { + const removed = queue.chunks.shift(); + if (!removed) { + break; + } + queue.bytes -= removed.data.byteLength; + if (queue.bytes < 0) { + queue.bytes = 0; + } + } + + this.bufferedChunks.set(input.chunk.streamId, queue); + } + + private maybeAckChunk(input: { streamId: number; endOffset: number }): void { + if (!Number.isFinite(input.endOffset) || input.endOffset < 0) { + return; + } + + const normalizedEndOffset = Math.floor(input.endOffset); + const previousAck = this.ackOffsets.get(input.streamId) ?? -1; + if (normalizedEndOffset <= previousAck) { + return; + } + + this.ackOffsets.set(input.streamId, normalizedEndOffset); + try { + this.config.sendAck({ + streamId: input.streamId, + offset: normalizedEndOffset, + }); + } catch { + // no-op + } + } +} diff --git a/packages/server/src/client/daemon-client.test.ts b/packages/server/src/client/daemon-client.test.ts index b3b5b13e7..948a2a4e2 100644 --- a/packages/server/src/client/daemon-client.test.ts +++ b/packages/server/src/client/daemon-client.test.ts @@ -707,7 +707,7 @@ describe("DaemonClient", () => { expect(ackFrame?.offset).toBe(8); }); - test("cleans terminal stream local state on terminal_stream_exit", async () => { + test("stops delivering and acking after terminal_stream_exit", async () => { const logger = createMockLogger(); const mock = createMockTransport(); @@ -723,16 +723,22 @@ describe("DaemonClient", () => { mock.triggerOpen(); await connectPromise; - const unsubscribe = client.onTerminalStreamData(23, () => { - // no-op + const seen: string[] = []; + const unsubscribe = client.onTerminalStreamData(23, (chunk) => { + seen.push(new TextDecoder().decode(chunk.data)); }); - const internal = client as unknown as { - terminalStreamHandlers: Map void>>; - bufferedTerminalStreamChunks: Map; - terminalStreamAckOffsets: Map; - }; - internal.bufferedTerminalStreamChunks.set(23, { chunks: [], bytes: 0 }); - internal.terminalStreamAckOffsets.set(23, 100); + mock.triggerMessage( + encodeBinaryMuxFrame({ + channel: BinaryMuxChannel.Terminal, + messageType: TerminalBinaryMessageType.OutputUtf8, + streamId: 23, + offset: 0, + payload: new TextEncoder().encode("before-exit"), + }) + ); + expect(seen).toEqual(["before-exit"]); + expect(mock.sent).toHaveLength(1); + mock.sent.length = 0; mock.triggerMessage( wrapSessionMessage({ @@ -744,9 +750,18 @@ describe("DaemonClient", () => { }) ); - expect(internal.terminalStreamHandlers.has(23)).toBe(false); - expect(internal.bufferedTerminalStreamChunks.has(23)).toBe(false); - expect(internal.terminalStreamAckOffsets.has(23)).toBe(false); + mock.triggerMessage( + encodeBinaryMuxFrame({ + channel: BinaryMuxChannel.Terminal, + messageType: TerminalBinaryMessageType.OutputUtf8, + streamId: 23, + offset: 12, + payload: new TextEncoder().encode("after-exit"), + }) + ); + + expect(seen).toEqual(["before-exit"]); + expect(mock.sent).toHaveLength(0); unsubscribe(); }); @@ -766,15 +781,10 @@ describe("DaemonClient", () => { mock.triggerOpen(); await connectPromise; - client.onTerminalStreamData(31, () => { - // no-op + const seen: string[] = []; + client.onTerminalStreamData(31, (chunk) => { + seen.push(new TextDecoder().decode(chunk.data)); }); - const internal = client as unknown as { - bufferedTerminalStreamChunks: Map; - terminalStreamAckOffsets: Map; - }; - internal.bufferedTerminalStreamChunks.set(31, { chunks: [], bytes: 0 }); - internal.terminalStreamAckOffsets.set(31, 55); const detachPromise = client.detachTerminalStream(31, "detach-31"); const detachRequest = JSON.parse(String(mock.sent[0])) as { @@ -797,8 +807,20 @@ describe("DaemonClient", () => { const payload = await detachPromise; expect(payload.success).toBe(false); - expect(internal.bufferedTerminalStreamChunks.has(31)).toBe(false); - expect(internal.terminalStreamAckOffsets.has(31)).toBe(false); + + mock.sent.length = 0; + mock.triggerMessage( + encodeBinaryMuxFrame({ + channel: BinaryMuxChannel.Terminal, + messageType: TerminalBinaryMessageType.OutputUtf8, + streamId: 31, + offset: 0, + payload: new TextEncoder().encode("after-detach"), + }) + ); + + expect(seen).toEqual([]); + expect(mock.sent).toHaveLength(0); }); test("parses canonical agent_stream tool_call payloads without crashing", async () => { diff --git a/packages/server/src/client/daemon-client.ts b/packages/server/src/client/daemon-client.ts index e5d72d43c..de0a57fad 100644 --- a/packages/server/src/client/daemon-client.ts +++ b/packages/server/src/client/daemon-client.ts @@ -68,6 +68,10 @@ import { encodeTerminalKeyInput, type TerminalKeyInput, } from "../shared/terminal-key-input.js"; +import { + TerminalStreamManager, + type TerminalStreamChunk, +} from "./daemon-client-terminal-stream-manager.js"; import { createRelayE2eeTransportFactory, createWebSocketTransportFactory, @@ -102,6 +106,7 @@ export type { WebSocketFactory, WebSocketLike, } from "./daemon-client-transport.js"; +export type { TerminalStreamChunk } from "./daemon-client-terminal-stream-manager.js"; export type ConnectionState = | { status: "idle" } @@ -233,22 +238,6 @@ export type FetchAgentTimelineOptions = { requestId?: string; }; -export type TerminalStreamChunk = { - streamId: number; - offset: number; - endOffset: number; - replay: boolean; - data: Uint8Array; -}; - -type BufferedTerminalStreamQueue = { - chunks: TerminalStreamChunk[]; - bytes: number; -}; - -const TERMINAL_STREAM_MAX_BUFFERED_CHUNKS = 2048; -const TERMINAL_STREAM_MAX_BUFFERED_BYTES = 2 * 1024 * 1024; - type AgentRefreshedStatusPayload = z.infer< typeof AgentRefreshedStatusPayloadSchema >; @@ -361,15 +350,21 @@ export class DaemonClient { private logger: Logger; private pendingSendQueue: PendingSend[] = []; private relayClientId: string | null = null; - private terminalStreamHandlers: Map< - number, - Set<(chunk: TerminalStreamChunk) => void> - > = new Map(); - private bufferedTerminalStreamChunks: Map = new Map(); - private terminalStreamAckOffsets: Map = new Map(); + private terminalStreams: TerminalStreamManager; constructor(private config: DaemonClientConfig) { this.logger = config.logger ?? consoleLogger; + this.terminalStreams = new TerminalStreamManager({ + sendAck: (ack) => { + this.sendBinaryFrame({ + channel: BinaryMuxChannel.Terminal, + messageType: TerminalBinaryMessageType.Ack, + streamId: ack.streamId, + offset: ack.offset, + payload: new Uint8Array(0), + }); + }, + }); // Relay requires a clientId so the daemon can create an independent // socket + E2EE channel per connected client. Generate one per DaemonClient // instance (stable across reconnects in this tab/app session). @@ -584,9 +579,7 @@ export class DaemonClient { } this.disposeTransport(1000, "Client closed"); this.clearWaiters(new Error("Daemon client closed")); - this.terminalStreamHandlers.clear(); - this.bufferedTerminalStreamChunks.clear(); - this.terminalStreamAckOffsets.clear(); + this.terminalStreams.clearAll(); this.updateConnectionState({ status: "disconnected", reason: "client_closed", @@ -2529,9 +2522,7 @@ export class DaemonClient { timeout: 10000, options: { skipQueue: true }, }); - this.terminalStreamHandlers.delete(streamId); - this.bufferedTerminalStreamChunks.delete(streamId); - this.terminalStreamAckOffsets.delete(streamId); + this.terminalStreams.clearStream({ streamId }); return payload; } @@ -2539,31 +2530,7 @@ export class DaemonClient { streamId: number, handler: (chunk: TerminalStreamChunk) => void ): () => void { - if (!this.terminalStreamHandlers.has(streamId)) { - this.terminalStreamHandlers.set(streamId, new Set()); - } - const handlers = this.terminalStreamHandlers.get(streamId)!; - handlers.add(handler); - - const buffered = this.bufferedTerminalStreamChunks.get(streamId); - if (buffered && buffered.chunks.length > 0) { - for (const chunk of buffered.chunks) { - try { - handler(chunk); - this.maybeAckTerminalStreamChunk(streamId, chunk.endOffset); - } catch { - // no-op - } - } - this.bufferedTerminalStreamChunks.delete(streamId); - } - - return () => { - handlers.delete(handler); - if (handlers.size === 0) { - this.terminalStreamHandlers.delete(streamId); - } - }; + return this.terminalStreams.subscribe({ streamId, handler }); } async waitForTerminalStreamData( @@ -2609,10 +2576,7 @@ export class DaemonClient { sendTerminalStreamAck(streamId: number, offset: number): void { const normalizedOffset = Math.max(0, Math.floor(offset)); - const previousAck = this.terminalStreamAckOffsets.get(streamId) ?? -1; - if (normalizedOffset > previousAck) { - this.terminalStreamAckOffsets.set(streamId, normalizedOffset); - } + this.terminalStreams.noteAck({ streamId, offset: normalizedOffset }); this.sendBinaryFrame({ channel: BinaryMuxChannel.Terminal, messageType: TerminalBinaryMessageType.Ack, @@ -2720,7 +2684,6 @@ export class DaemonClient { frame.channel === BinaryMuxChannel.Terminal && frame.messageType === TerminalBinaryMessageType.OutputUtf8 ) { - const handlers = this.terminalStreamHandlers.get(frame.streamId); const chunk: TerminalStreamChunk = { streamId: frame.streamId, offset: frame.offset, @@ -2728,62 +2691,11 @@ export class DaemonClient { replay: Boolean((frame.flags ?? 0) & TerminalBinaryFlags.Replay), data: frame.payload ?? new Uint8Array(0), }; - if (!handlers || handlers.size === 0) { - const queue = this.bufferedTerminalStreamChunks.get(frame.streamId) ?? { - chunks: [], - bytes: 0, - }; - queue.chunks.push(chunk); - queue.bytes += chunk.data.byteLength; - while ( - queue.chunks.length > TERMINAL_STREAM_MAX_BUFFERED_CHUNKS || - queue.bytes > TERMINAL_STREAM_MAX_BUFFERED_BYTES - ) { - const removed = queue.chunks.shift(); - if (!removed) { - break; - } - queue.bytes -= removed.data.byteLength; - if (queue.bytes < 0) { - queue.bytes = 0; - } - } - this.bufferedTerminalStreamChunks.set(frame.streamId, queue); - return; - } - let delivered = false; - for (const handler of handlers) { - try { - handler(chunk); - delivered = true; - } catch { - // no-op - } - } - if (delivered) { - this.maybeAckTerminalStreamChunk(frame.streamId, chunk.endOffset); - } + this.terminalStreams.receiveChunk({ chunk }); return; } } - private maybeAckTerminalStreamChunk(streamId: number, endOffset: number): void { - if (!Number.isFinite(endOffset) || endOffset < 0) { - return; - } - const normalizedEndOffset = Math.floor(endOffset); - const previousAck = this.terminalStreamAckOffsets.get(streamId) ?? -1; - if (normalizedEndOffset <= previousAck) { - return; - } - this.terminalStreamAckOffsets.set(streamId, normalizedEndOffset); - try { - this.sendTerminalStreamAck(streamId, normalizedEndOffset); - } catch { - // no-op - } - } - private updateConnectionState(next: ConnectionState): void { this.connectionState = next; for (const listener of this.connectionListeners) { @@ -2824,9 +2736,7 @@ export class DaemonClient { // and responses from the previous connection will never arrive. this.clearWaiters(new Error(reason ?? "Connection lost")); this.rejectPendingSendQueue(new Error(reason ?? "Connection lost")); - this.terminalStreamHandlers.clear(); - this.bufferedTerminalStreamChunks.clear(); - this.terminalStreamAckOffsets.clear(); + this.terminalStreams.clearAll(); this.updateConnectionState({ status: "disconnected", @@ -2843,9 +2753,7 @@ export class DaemonClient { private handleSessionMessage(msg: SessionOutboundMessage): void { if (msg.type === "terminal_stream_exit") { - this.terminalStreamHandlers.delete(msg.payload.streamId); - this.bufferedTerminalStreamChunks.delete(msg.payload.streamId); - this.terminalStreamAckOffsets.delete(msg.payload.streamId); + this.terminalStreams.clearStream({ streamId: msg.payload.streamId }); } if (this.rawMessageListeners.size > 0) {