refactor(server): extract terminal stream state manager from daemon client

Move terminal stream buffering, handler fanout, and ack offset state out of DaemonClient into a focused TerminalStreamManager.\n\nReasoning:\n- daemon-client.ts mixed RPC transport concerns with terminal stream state machine details.\n- terminal buffering/ack logic is cohesive and testable in isolation, so splitting it lowers logic density and makes future stream changes safer.\n\nWhat changed:\n- Added daemon-client-terminal-stream-manager.ts to own stream handlers, buffered chunks, and ack monotonicity.\n- Updated DaemonClient to delegate terminal stream state transitions (subscribe, receive chunk, clear stream/all, note ack).\n- Preserved public API by re-exporting TerminalStreamChunk from daemon-client.ts.\n- Hardened daemon-client tests to assert behavior (no delivery/ack after detach or stream_exit) rather than private internal maps.\n- Added dedicated manager unit tests for buffering, eviction, handler-failure behavior, and stale ack suppression.\n\nVerification:\n- npm run test --workspace=@getpaseo/server -- src/client/daemon-client.test.ts src/client/daemon-client-terminal-stream-manager.test.ts\n- npm run typecheck\n\nNext-agent notes:\n- Accomplishment: terminal stream state is now isolated; daemon-client.ts is slimmer and easier to extend.\n- Challenge: daemon-client.ts is still large (~3k lines) with other cohesive seams remaining (dictation flow and checkout RPC groups).
This commit is contained in:
Mohamed Boudra
2026-02-14 09:51:49 +07:00
parent c87c206900
commit d9e14186ce
4 changed files with 366 additions and 139 deletions
@@ -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 });
});
});
@@ -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<number, Set<TerminalChunkHandler>> = new Map();
private readonly bufferedChunks: Map<number, BufferedTerminalStreamQueue> = new Map();
private readonly ackOffsets: Map<number, number> = 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
}
}
}
@@ -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<number, Set<(chunk: unknown) => void>>;
bufferedTerminalStreamChunks: Map<number, { chunks: unknown[]; bytes: number }>;
terminalStreamAckOffsets: Map<number, number>;
};
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<number, { chunks: unknown[]; bytes: number }>;
terminalStreamAckOffsets: Map<number, number>;
};
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 () => {
+24 -116
View File
@@ -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<number, BufferedTerminalStreamQueue> = new Map();
private terminalStreamAckOffsets: Map<number, number> = 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) {