From c7ea60a01accd4c62bffcff8380466f82e8f46a1 Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Wed, 8 Jul 2026 10:02:35 +0800 Subject: [PATCH] fix(kap-server): buffer and coalesce outbound ws deltas - add a per-connection outbound send buffer in WsConnectionV1; sendFrame enqueues and a flush drains on a 16ms interval or once 64 frames queue - coalesce adjacent volatile assistant/thinking deltas for the same session/agent/turn into one frame, keeping the first offset/seq so client alignment is unchanged - defer flushing while socket.bufferedAmount exceeds a 1MiB high-water mark, merging new deltas into the queued frame instead of growing it - flush remaining frames on close to avoid truncating the tail of a stream - expose flushIntervalMs / maxBatchSize / highWaterMarkBytes via registerWsV1 --- .../src/transport/ws/v1/registerWsV1.ts | 6 + .../src/transport/ws/v1/wsConnectionV1.ts | 181 +++++++++- .../kap-server/test/wsConnectionV1.test.ts | 320 ++++++++++++++++++ 3 files changed, 502 insertions(+), 5 deletions(-) create mode 100644 packages/kap-server/test/wsConnectionV1.test.ts diff --git a/packages/kap-server/src/transport/ws/v1/registerWsV1.ts b/packages/kap-server/src/transport/ws/v1/registerWsV1.ts index 2b3958efd..e7f03cd6a 100644 --- a/packages/kap-server/src/transport/ws/v1/registerWsV1.ts +++ b/packages/kap-server/src/transport/ws/v1/registerWsV1.ts @@ -29,6 +29,9 @@ export interface RegisterWsV1Options { readonly pingIntervalMs?: number; readonly pongTimeoutMs?: number; readonly maxBufferSize?: number; + readonly flushIntervalMs?: number; + readonly maxBatchSize?: number; + readonly highWaterMarkBytes?: number; } export function registerWsV1(core: Scope, opts: RegisterWsV1Options): WebSocketServer { @@ -48,6 +51,9 @@ export function registerWsV1(core: Scope, opts: RegisterWsV1Options): WebSocketS pingIntervalMs: opts.pingIntervalMs, pongTimeoutMs: opts.pongTimeoutMs, maxBufferSize: opts.maxBufferSize, + flushIntervalMs: opts.flushIntervalMs, + maxBatchSize: opts.maxBatchSize, + highWaterMarkBytes: opts.highWaterMarkBytes, }); socket.on('close', () => registry.remove(conn.id)); }); diff --git a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts index 344b84237..11ca0fe23 100644 --- a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts +++ b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts @@ -38,6 +38,15 @@ const DEFAULT_PING_INTERVAL_MS = 30_000; const DEFAULT_PONG_TIMEOUT_MS = 10_000; const DEFAULT_MAX_BUFFER_SIZE = 1000; +// Outbound send buffer — coalesces a burst of frames (notably high-frequency +// volatile text deltas) into fewer `socket.send` calls and applies backpressure +// when the peer is not draining fast enough. See `flush()` / `coalesceFrames`. +const DEFAULT_FLUSH_INTERVAL_MS = 16; +const DEFAULT_MAX_BATCH_SIZE = 64; +const DEFAULT_HIGH_WATER_MARK_BYTES = 1 << 20; // 1 MiB +const DEFAULT_BACKPRESSURE_RETRY_MS = 5; +const DEFAULT_BACKPRESSURE_MAX_DELAY_MS = 100; + interface InboundFrame { type: string; id?: string; @@ -62,6 +71,12 @@ export interface WsConnectionV1Options { readonly pingIntervalMs?: number; readonly pongTimeoutMs?: number; readonly maxBufferSize?: number; + /** Delay before a buffered batch is flushed; coalesces frames within the window. */ + readonly flushIntervalMs?: number; + /** Flush immediately once this many frames are queued, even before the interval. */ + readonly maxBatchSize?: number; + /** `socket.bufferedAmount` above which flushing is deferred (backpressure). */ + readonly highWaterMarkBytes?: number; } export class WsConnectionV1 implements BroadcastTarget { @@ -76,6 +91,9 @@ export class WsConnectionV1 implements BroadcastTarget { private readonly pingIntervalMs: number; private readonly pongTimeoutMs: number; private readonly maxBufferSize: number; + private readonly flushIntervalMs: number; + private readonly maxBatchSize: number; + private readonly highWaterMarkBytes: number; private readonly logger?: JournalLogger; private closed = false; @@ -86,6 +104,13 @@ export class WsConnectionV1 implements BroadcastTarget { private pingTimer?: ReturnType; private pongTimer?: ReturnType; + /** Outbound frames awaiting the next flush. */ + private outbound: unknown[] = []; + private flushTimer?: ReturnType; + private backpressureRetryTimer?: ReturnType; + /** Epoch ms when the current backpressure deferral started; caps the wait. */ + private backpressureSince?: number; + constructor(opts: WsConnectionV1Options) { this.id = `conn_${ulid()}`; this.connectedAt = new Date().toISOString(); @@ -98,6 +123,9 @@ export class WsConnectionV1 implements BroadcastTarget { this.pingIntervalMs = opts.pingIntervalMs ?? DEFAULT_PING_INTERVAL_MS; this.pongTimeoutMs = opts.pongTimeoutMs ?? DEFAULT_PONG_TIMEOUT_MS; this.maxBufferSize = opts.maxBufferSize ?? DEFAULT_MAX_BUFFER_SIZE; + this.flushIntervalMs = opts.flushIntervalMs ?? DEFAULT_FLUSH_INTERVAL_MS; + this.maxBatchSize = opts.maxBatchSize ?? DEFAULT_MAX_BATCH_SIZE; + this.highWaterMarkBytes = opts.highWaterMarkBytes ?? DEFAULT_HIGH_WATER_MARK_BYTES; this.socket.on('message', (data: RawData) => this.onMessage(data)); this.socket.on('close', () => this.onClose()); @@ -325,16 +353,85 @@ export class WsConnectionV1 implements BroadcastTarget { } private sendFrame(msg: unknown): void { - if (this.closed || this.socket.readyState !== this.socket.OPEN) return; - try { - this.socket.send(JSON.stringify(msg)); - } catch { - // best-effort + if (this.closed) return; + this.outbound.push(msg); + if (this.outbound.length >= this.maxBatchSize) { + // Batch is full — flush now rather than wait for the interval. + this.flush(); + return; } + this.scheduleFlush(); + } + + private scheduleFlush(): void { + if (this.flushTimer !== undefined) return; + this.flushTimer = setTimeout(() => { + this.flushTimer = undefined; + this.flush(); + }, this.flushIntervalMs); + this.flushTimer.unref?.(); + } + + /** + * Drain the outbound buffer: coalesce adjacent compatible volatile deltas, + * then write the surviving frames to the socket. When the peer is not + * draining (`bufferedAmount` above the high-water mark) and `force` is not + * set, defer and keep accumulating — later deltas merge into the queued + * ones, so the frame count does not grow while we wait. + */ + private flush(force = false): void { + if (this.flushTimer !== undefined) { + clearTimeout(this.flushTimer); + this.flushTimer = undefined; + } + if (this.outbound.length === 0) return; + if (this.closed || this.socket.readyState !== this.socket.OPEN) { + // Socket is gone — drop queued frames rather than send into a dead pipe. + this.outbound = []; + return; + } + + if (!force && this.socket.bufferedAmount > this.highWaterMarkBytes) { + this.deferForBackpressure(); + return; + } + this.backpressureSince = undefined; + + const frames = coalesceFrames(this.outbound); + this.outbound = []; + for (const frame of frames) { + if (this.closed || this.socket.readyState !== this.socket.OPEN) return; + try { + this.socket.send(JSON.stringify(frame)); + } catch { + // best-effort + } + } + } + + private deferForBackpressure(): void { + const now = Date.now(); + if (this.backpressureSince === undefined) this.backpressureSince = now; + if (now - this.backpressureSince >= DEFAULT_BACKPRESSURE_MAX_DELAY_MS) { + // Peer stayed above the watermark too long — force-flush to avoid + // starving the stream; the socket layer will buffer or drop. + this.flush(true); + return; + } + if (this.backpressureRetryTimer !== undefined) return; + this.backpressureRetryTimer = setTimeout(() => { + this.backpressureRetryTimer = undefined; + this.flush(); + }, DEFAULT_BACKPRESSURE_RETRY_MS); + this.backpressureRetryTimer.unref?.(); } close(code = 1000, reason?: string): void { if (this.closed) return; + // Best-effort: push out any queued frames (e.g. the tail of a delta + // stream) before tearing the socket down, so the client sees a complete + // stream rather than a truncated one. + this.flush(true); try { this.socket.close(code, reason); } catch { @@ -347,6 +444,9 @@ export class WsConnectionV1 implements BroadcastTarget { this.closed = true; if (this.pingTimer !== undefined) clearInterval(this.pingTimer); if (this.pongTimer !== undefined) clearTimeout(this.pongTimer); + if (this.flushTimer !== undefined) clearTimeout(this.flushTimer); + if (this.backpressureRetryTimer !== undefined) clearTimeout(this.backpressureRetryTimer); + this.outbound = []; for (const sid of this.subscriptions) this.broadcaster.unsubscribe(sid, this); // registry removal is handled by registerWsV1 on the socket 'close' event. } @@ -363,3 +463,74 @@ function rawDataToString(data: RawData): string { if (Array.isArray(data)) return Buffer.concat(data).toString('utf8'); return Buffer.from(data as ArrayBuffer).toString('utf8'); } + +// --------------------------------------------------------------------------- +// Outbound coalescing +// --------------------------------------------------------------------------- + +/** A volatile text-delta envelope that can be merged with an adjacent one. */ +interface CoalescableDelta { + type: 'assistant.delta' | 'thinking.delta'; + seq: number; + volatile: true; + offset?: number; + session_id?: string; + timestamp: string; + payload: { + agentId?: string; + turnId?: number; + delta: string; + [key: string]: unknown; + }; +} + +function isCoalescableDelta(frame: unknown): frame is CoalescableDelta { + if (typeof frame !== 'object' || frame === null) return false; + const f = frame as Record; + if (f['volatile'] !== true) return false; + const type = f['type']; + if (type !== 'assistant.delta' && type !== 'thinking.delta') return false; + const payload = f['payload']; + if (typeof payload !== 'object' || payload === null) return false; + return typeof (payload as Record)['delta'] === 'string'; +} + +/** + * Merge adjacent compatible volatile text deltas into a single envelope. + * + * Two adjacent frames merge when both are `volatile` `assistant.delta` / + * `thinking.delta` of the same type, addressed to the same session, agent, + * and turn. The merged frame keeps the first frame's `seq` / `offset` / + * `timestamp` and concatenates `payload.delta` in order — the client's + * offset-based alignment against the in-flight snapshot stays correct + * (the broadcaster's per-session dispatch queue guarantees consecutive deltas + * for a turn carry consecutive offsets). + * + * Durable events, control frames, and non-text deltas are never merged, and + * merging never crosses a non-mergeable frame, so overall ordering is + * preserved. The input frames are not mutated; merged results are fresh + * objects. Exported for unit testing. + */ +export function coalesceFrames(frames: readonly unknown[]): unknown[] { + const out: unknown[] = []; + for (const frame of frames) { + const last = out.at(-1); + if ( + last !== undefined && + isCoalescableDelta(last) && + isCoalescableDelta(frame) && + last.type === frame.type && + last.session_id === frame.session_id && + last.payload.agentId === frame.payload.agentId && + last.payload.turnId === frame.payload.turnId + ) { + out[out.length - 1] = { + ...last, + payload: { ...last.payload, delta: last.payload.delta + frame.payload.delta }, + }; + } else { + out.push(frame); + } + } + return out; +} diff --git a/packages/kap-server/test/wsConnectionV1.test.ts b/packages/kap-server/test/wsConnectionV1.test.ts new file mode 100644 index 000000000..33144e8ae --- /dev/null +++ b/packages/kap-server/test/wsConnectionV1.test.ts @@ -0,0 +1,320 @@ +/** + * `WsConnectionV1` — outbound send buffer: coalescing of high-frequency + * volatile text deltas, batch flush, backpressure deferral, and close flush. + */ + +import type { WebSocket } from 'ws'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { IConnectionRegistry } from '../src/transport/ws/connectionRegistry'; +import type { SessionEventBroadcaster } from '../src/transport/ws/v1/sessionEventBroadcaster'; +import { + type WsConnectionV1Options, + WsConnectionV1, + coalesceFrames, +} from '../src/transport/ws/v1/wsConnectionV1'; + +// --------------------------------------------------------------------------- +// Fakes +// --------------------------------------------------------------------------- + +class FakeSocket { + readonly OPEN = 1; + readonly CLOSED = 3; + readyState = 1; + bufferedAmount = 0; + sent: string[] = []; + closeCalls: Array<{ code?: number; reason?: string }> = []; + private readonly handlers = new Map void>>(); + + on(event: string, cb: (...a: unknown[]) => void): this { + const list = this.handlers.get(event) ?? []; + list.push(cb); + this.handlers.set(event, list); + return this; + } + + send(data: string): void { + this.sent.push(data); + } + + close(code?: number, reason?: string): void { + this.closeCalls.push({ code, reason }); + this.readyState = this.CLOSED; + this.emit('close'); + } + + terminate(): void { + this.readyState = this.CLOSED; + this.emit('close'); + } + + emit(event: string, ...a: unknown[]): void { + for (const cb of this.handlers.get(event) ?? []) cb(...a); + } + + frames(): unknown[] { + return this.sent.map((s) => JSON.parse(s)); + } +} + +function makeBroadcaster(): SessionEventBroadcaster { + return { + subscribe: async () => true, + unsubscribe: () => {}, + getCursor: async () => ({ seq: 0, epoch: '' }), + getBufferedSince: async () => ({ + events: [], + resyncRequired: false, + currentSeq: 0, + epoch: '', + }), + } as unknown as SessionEventBroadcaster; +} + +function makeRegistry(): IConnectionRegistry { + return { + add: () => {}, + remove: () => {}, + get: () => undefined, + values: () => [], + closeAll: () => {}, + size: () => 0, + }; +} + +function makeConn(socket: FakeSocket, opts: Partial = {}): WsConnectionV1 { + return new WsConnectionV1({ + socket: socket as unknown as WebSocket, + broadcaster: makeBroadcaster(), + connectionRegistry: makeRegistry(), + remoteAddress: null, + userAgent: null, + // Keep the heartbeat far from the test window so pings never interfere. + pingIntervalMs: 600_000, + ...opts, + }); +} + +function delta( + sessionId: string, + agentId: string, + turnId: number, + text: string, + offset: number, + type: 'assistant.delta' | 'thinking.delta' = 'assistant.delta', +) { + return { + type, + seq: 1, + volatile: true as const, + offset, + session_id: sessionId, + timestamp: '2026-01-01T00:00:00.000Z', + payload: { type, agentId, sessionId, turnId, delta: text }, + }; +} + +function durable(type: string, sessionId: string, seq: number) { + return { + type, + seq, + session_id: sessionId, + timestamp: '2026-01-01T00:00:00.000Z', + payload: { type, agentId: 'main', sessionId }, + }; +} + +// --------------------------------------------------------------------------- +// coalesceFrames — pure +// --------------------------------------------------------------------------- + +describe('coalesceFrames', () => { + it('merges adjacent compatible assistant deltas', () => { + const out = coalesceFrames([ + delta('s1', 'main', 1, 'Hello', 0), + delta('s1', 'main', 1, ' ', 5), + delta('s1', 'main', 1, 'world', 6), + ]); + expect(out).toHaveLength(1); + const f = out[0] as { offset: number; volatile: boolean; seq: number; payload: { delta: string } }; + expect(f.payload.delta).toBe('Hello world'); + expect(f.offset).toBe(0); + expect(f.volatile).toBe(true); + expect(f.seq).toBe(1); + }); + + it('does not merge across a durable frame', () => { + const out = coalesceFrames([ + delta('s1', 'main', 1, 'a', 0), + durable('turn.ended', 's1', 2), + delta('s1', 'main', 1, 'b', 1), + ]); + expect(out).toHaveLength(3); + expect((out[0] as { payload: { delta: string } }).payload.delta).toBe('a'); + expect((out[1] as { type: string }).type).toBe('turn.ended'); + expect((out[2] as { payload: { delta: string } }).payload.delta).toBe('b'); + }); + + it('does not merge different delta types', () => { + const out = coalesceFrames([ + delta('s1', 'main', 1, 'hi', 0, 'assistant.delta'), + delta('s1', 'main', 1, 'think', 0, 'thinking.delta'), + ]); + expect(out).toHaveLength(2); + }); + + it('does not merge deltas from different sessions / agents / turns', () => { + expect( + coalesceFrames([delta('s1', 'main', 1, 'a', 0), delta('s2', 'main', 1, 'b', 0)]), + ).toHaveLength(2); + expect( + coalesceFrames([delta('s1', 'main', 1, 'a', 0), delta('s1', 'sub', 1, 'b', 0)]), + ).toHaveLength(2); + expect( + coalesceFrames([delta('s1', 'main', 1, 'a', 0), delta('s1', 'main', 2, 'b', 0)]), + ).toHaveLength(2); + }); + + it('leaves non-volatile and non-text frames untouched', () => { + const toolCallDelta = { + type: 'tool.call.delta', + seq: 1, + volatile: true as const, + session_id: 's1', + timestamp: '2026-01-01T00:00:00.000Z', + payload: { type: 'tool.call.delta', agentId: 'main', turnId: 1, args: { x: 1 } }, + }; + expect(coalesceFrames([toolCallDelta, toolCallDelta])).toHaveLength(2); + }); + + it('does not mutate the input frames', () => { + const a = delta('s1', 'main', 1, 'a', 0); + const b = delta('s1', 'main', 1, 'b', 1); + const out = coalesceFrames([a, b]); + expect(out).toHaveLength(1); + expect(a.payload.delta).toBe('a'); + expect(b.payload.delta).toBe('b'); + }); + + it('handles empty and single-element input', () => { + expect(coalesceFrames([])).toEqual([]); + const only = delta('s1', 'main', 1, 'x', 0); + const out = coalesceFrames([only]); + expect(out).toHaveLength(1); + expect(out[0]).toBe(only); + }); +}); + +// --------------------------------------------------------------------------- +// WsConnectionV1 — flush / backpressure / close +// --------------------------------------------------------------------------- + +describe('WsConnectionV1 outbound buffer', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + afterEach(() => { + vi.useRealTimers(); + }); + + it('buffers server_hello and flushes it after the interval', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 16 }); + expect(socket.sent).toHaveLength(0); + await vi.advanceTimersByTimeAsync(16); + expect(socket.frames().map((f) => (f as { type: string }).type)).toContain('server_hello'); + conn.close(); + }); + + it('coalesces adjacent deltas into one socket.send', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 16 }); + await vi.advanceTimersByTimeAsync(16); // flush server_hello + socket.sent = []; + + conn.send(delta('s1', 'main', 1, 'Hello', 0)); + conn.send(delta('s1', 'main', 1, ' ', 5)); + conn.send(delta('s1', 'main', 1, 'world', 6)); + expect(socket.sent).toHaveLength(0); // still buffered + await vi.advanceTimersByTimeAsync(16); + + const frames = socket.frames(); + expect(frames).toHaveLength(1); + const f = frames[0] as { type: string; offset: number; payload: { delta: string } }; + expect(f.type).toBe('assistant.delta'); + expect(f.offset).toBe(0); + expect(f.payload.delta).toBe('Hello world'); + conn.close(); + }); + + it('flushes immediately once the batch reaches maxBatchSize', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 1000, maxBatchSize: 3 }); + // constructor already queued server_hello (1); two deltas bring it to 3. + conn.send(delta('s1', 'main', 1, 'a', 0)); + conn.send(delta('s1', 'main', 1, 'b', 1)); + // No timer advanced — flush must have happened synchronously. + const types = socket.frames().map((f) => (f as { type: string }).type); + expect(types).toEqual(['server_hello', 'assistant.delta']); + conn.close(); + }); + + it('defers flushing while the peer is above the watermark, then coalesces on drain', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { + flushIntervalMs: 16, + highWaterMarkBytes: 100, + }); + await vi.advanceTimersByTimeAsync(16); // flush server_hello + socket.sent = []; + + socket.bufferedAmount = 200; // above the watermark + conn.send(delta('s1', 'main', 1, 'Hello', 0)); + await vi.advanceTimersByTimeAsync(16); // flush attempted → deferred + expect(socket.sent).toHaveLength(0); + + // More deltas arrive while deferred — they merge into the queued frame. + conn.send(delta('s1', 'main', 1, ' world', 5)); + await vi.advanceTimersByTimeAsync(5); // backpressure retry, still high + expect(socket.sent).toHaveLength(0); + + socket.bufferedAmount = 0; // peer drained + await vi.advanceTimersByTimeAsync(5); // retry succeeds + const frames = socket.frames(); + expect(frames).toHaveLength(1); + expect((frames[0] as { payload: { delta: string } }).payload.delta).toBe('Hello world'); + conn.close(); + }); + + it('force-flushes buffered frames on close', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 1000 }); + // server_hello is still buffered (interval not elapsed). + conn.send(delta('s1', 'main', 1, 'tail', 0)); + expect(socket.sent).toHaveLength(0); + + conn.close(); + const types = socket.frames().map((f) => (f as { type: string }).type); + expect(types).toContain('server_hello'); + expect(types).toContain('assistant.delta'); + const tail = socket + .frames() + .find((f) => (f as { type: string }).type === 'assistant.delta') as { + payload: { delta: string }; + }; + expect(tail.payload.delta).toBe('tail'); + }); + + it('drops buffered frames when the socket is already closed at flush time', async () => { + const socket = new FakeSocket(); + const conn = makeConn(socket, { flushIntervalMs: 16 }); + await vi.advanceTimersByTimeAsync(16); // flush server_hello + socket.sent = []; + + socket.readyState = socket.CLOSED; // peer went away + conn.send(delta('s1', 'main', 1, 'lost', 0)); + await vi.advanceTimersByTimeAsync(16); + expect(socket.sent).toHaveLength(0); + }); +});