From d15e55f227cef54910eb4b2de80e82671fdd898a Mon Sep 17 00:00:00 2001 From: Edgar Durand Date: Thu, 29 Oct 2026 14:26:00 -0300 Subject: [PATCH] perf(cli): coalesce streaming-chunk posts per chunkId (latest wins, one in flight) 2026-10-29 p99 alert: one CLI posted 10-30 cumulative streaming-chunk snapshots per second for the same chunkId. A new StreamingChunkCoalescer paces both publishers (ACP AcpPublisher and the PTY StreamingEmitterService): the first snapshot goes out at once, then at most one request in flight and >=100 ms between posts per chunkId, latest content wins. The isFinal frame supersedes anything pending, skips the spacing wait and is always sent after the in-flight request, so it is the last post of its chunkId; late non-final snapshots after it are dropped. chunkIds never wait on each other. Big snapshots (>=8 KB) keep their 300 ms window. Shutdown flushes pending content. Measured on a simulated 20 updates/s stream (10 s): 201 -> 102 posts at 60 ms request latency, 201 -> 35 at 300 ms. Bead: codeagent-p6nu Co-Authored-By: Claude Opus 5.5 --- .../acp.publisher.streaming-chunk-cap.test.ts | 32 ++- .../services/chunk-coalescer.test.ts | 201 ++++++++++++++++++ apps/cli/src/agents/acp/publisher.ts | 61 +++--- apps/cli/src/agents/acp/runner.ts | 6 + .../src/services/streaming-emitter.service.ts | 16 +- .../src/services/streaming/chunk-coalescer.ts | 136 ++++++++++++ 6 files changed, 409 insertions(+), 43 deletions(-) create mode 100644 apps/cli/__tests__/services/chunk-coalescer.test.ts create mode 100644 apps/cli/src/services/streaming/chunk-coalescer.ts diff --git a/apps/cli/__tests__/agents/acp.publisher.streaming-chunk-cap.test.ts b/apps/cli/__tests__/agents/acp.publisher.streaming-chunk-cap.test.ts index 95ed1239..855ecb2c 100644 --- a/apps/cli/__tests__/agents/acp.publisher.streaming-chunk-cap.test.ts +++ b/apps/cli/__tests__/agents/acp.publisher.streaming-chunk-cap.test.ts @@ -5,6 +5,7 @@ import { STREAMING_CHUNK_SNAPSHOT_MAX_CHARS, boundStreamingChunkContent, } from '../../src/agents/acp/publisher'; +import { STREAMING_CHUNK_MIN_INTERVAL_MS } from '../../src/services/streaming/chunk-coalescer'; import * as transport from '../../src/services/streaming/transport'; /** @@ -84,16 +85,35 @@ describe('AcpPublisher — streaming-chunk snapshot cap', () => { expect(payloads()[1].content.length).toBeLessThan(64 * 1024); }); - it('small snapshots go out on every delta (no behavior change below 8 KB)', async () => { + it('small snapshots: first at once, then latest-wins every 100 ms (2026-10-29 flood)', async () => { for (let i = 0; i < 5; i++) { - await publisher.publishStreamingChunk({ + void publisher.publishStreamingChunk({ chunkId: 'c2', kind: 'text', content: 'a'.repeat(100 * (i + 1)), isFinal: false, }); } - expect(postSpy).toHaveBeenCalledTimes(5); + await vi.advanceTimersByTimeAsync(0); + expect(postSpy).toHaveBeenCalledTimes(1); + await vi.advanceTimersByTimeAsync(STREAMING_CHUNK_MIN_INTERVAL_MS); + expect(postSpy).toHaveBeenCalledTimes(2); + expect(payloads()[1].content).toBe('a'.repeat(500)); + await publisher.publishStreamingChunk({ + chunkId: 'c2', + kind: 'text', + content: 'a'.repeat(600), + isFinal: true, + }); + expect(postSpy).toHaveBeenCalledTimes(3); + expect(payloads()[2]).toMatchObject({ isFinal: true, content: 'a'.repeat(600) }); + }); + + it('flushStreamingChunks sends the coalesced snapshot still waiting', async () => { + void publisher.publishStreamingChunk({ chunkId: 'c4', kind: 'text', content: 'x', isFinal: false }); + void publisher.publishStreamingChunk({ chunkId: 'c4', kind: 'text', content: 'xy', isFinal: false }); + await publisher.flushStreamingChunks(); + expect(payloads().map((p) => p.content)).toEqual(['x', 'xy']); }); it('big snapshots are coalesced to one trailing POST per window, latest content wins; final flushes immediately', async () => { @@ -105,13 +125,13 @@ describe('AcpPublisher — streaming-chunk snapshot cap', () => { isFinal: false, }); expect(postSpy).toHaveBeenCalledTimes(1); // first one goes straight out - await publisher.publishStreamingChunk({ + void publisher.publishStreamingChunk({ chunkId: 'c3', kind: 'thinking', content: base + '1', isFinal: false, }); - await publisher.publishStreamingChunk({ + void publisher.publishStreamingChunk({ chunkId: 'c3', kind: 'thinking', content: base + '12', @@ -121,7 +141,7 @@ describe('AcpPublisher — streaming-chunk snapshot cap', () => { await vi.advanceTimersByTimeAsync(STREAMING_CHUNK_COALESCE_WINDOW_MS + 5); expect(postSpy).toHaveBeenCalledTimes(2); expect(payloads()[1].content.endsWith('12')).toBe(true); // latest snapshot, not the first pending one - await publisher.publishStreamingChunk({ + void publisher.publishStreamingChunk({ chunkId: 'c3', kind: 'thinking', content: base + '123', diff --git a/apps/cli/__tests__/services/chunk-coalescer.test.ts b/apps/cli/__tests__/services/chunk-coalescer.test.ts new file mode 100644 index 00000000..429df67c --- /dev/null +++ b/apps/cli/__tests__/services/chunk-coalescer.test.ts @@ -0,0 +1,201 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { StreamingChunkEvent } from '@codeam/shared'; +import { + STREAMING_CHUNK_MIN_INTERVAL_MS, + StreamingChunkCoalescer, +} from '../../src/services/streaming/chunk-coalescer'; + +/** + * 2026-10-29 p99 alert: one CLI posted 10–30 cumulative streaming-chunk + * snapshots per second for the same chunkId. The coalescer keeps latest-wins, + * one request in flight and ≥100 ms spacing per chunkId, and always delivers + * the final frame last. + */ + +interface Sent { + event: StreamingChunkEvent; + resolve: () => void; +} + +function harness(opts: { autoResolve: boolean; latencyMs?: number }) { + const sent: Sent[] = []; + const send = vi.fn((event: StreamingChunkEvent) => { + return new Promise((resolve) => { + sent.push({ event, resolve }); + if (opts.autoResolve) { + if (opts.latencyMs) setTimeout(resolve, opts.latencyMs); + else resolve(); + } + }); + }); + const coalescer = new StreamingChunkCoalescer({ send }); + return { sent, send, coalescer }; +} + +const chunk = (chunkId: string, content: string, isFinal = false): StreamingChunkEvent => ({ + chunkId, + kind: 'text', + content, + isFinal, +}); + +describe('StreamingChunkCoalescer', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + afterEach(() => { + vi.useRealTimers(); + }); + + it('sends the first snapshot of a chunk immediately', async () => { + const { sent, coalescer } = harness({ autoResolve: true }); + void coalescer.submit(chunk('a', 'h')); + expect(sent).toHaveLength(1); + expect(sent[0].event.content).toBe('h'); + }); + + it('a burst of 30 updates becomes a few posts, the last one carrying the last content', async () => { + const { sent, coalescer } = harness({ autoResolve: true }); + let content = ''; + for (let i = 0; i < 30; i++) { + content += `w${i} `; + void coalescer.submit(chunk('a', content)); + } + await vi.advanceTimersByTimeAsync(STREAMING_CHUNK_MIN_INTERVAL_MS * 2); + expect(sent).toHaveLength(2); // the first one at once + one trailing + expect(sent[1].event.content).toBe(content); + }); + + it('done is always delivered, after the in-flight post, and is the last post', async () => { + const { sent, coalescer } = harness({ autoResolve: false }); + void coalescer.submit(chunk('a', 'one')); + void coalescer.submit(chunk('a', 'one two')); + let finalSettled = false; + const done = coalescer.submit(chunk('a', 'one two three', true)).then(() => { + finalSettled = true; + }); + void coalescer.submit(chunk('a', 'stale after final')); + await vi.advanceTimersByTimeAsync(1000); + // Still only the first request: the final waits for it. + expect(sent).toHaveLength(1); + expect(finalSettled).toBe(false); + sent[0].resolve(); + await vi.advanceTimersByTimeAsync(0); + expect(sent).toHaveLength(2); + expect(sent[1].event).toEqual(chunk('a', 'one two three', true)); + sent[1].resolve(); + await done; + expect(finalSettled).toBe(true); + await vi.advanceTimersByTimeAsync(1000); + expect(sent).toHaveLength(2); + }); + + it('the final skips the spacing wait (no added latency at the end of a message)', async () => { + const { sent, coalescer } = harness({ autoResolve: true }); + void coalescer.submit(chunk('a', 'x')); + void coalescer.submit(chunk('a', 'xy')); + await vi.advanceTimersByTimeAsync(0); + void coalescer.submit(chunk('a', 'xyz', true)); + await vi.advanceTimersByTimeAsync(0); + expect(sent.map((s) => s.event.content)).toEqual(['x', 'xyz']); + expect(sent[1].event.isFinal).toBe(true); + }); + + it('in-flight + newer snapshots → only the latest goes out after the in-flight one', async () => { + const { sent, coalescer } = harness({ autoResolve: false }); + void coalescer.submit(chunk('a', '1')); + void coalescer.submit(chunk('a', '12')); + void coalescer.submit(chunk('a', '123')); + await vi.advanceTimersByTimeAsync(500); + expect(sent).toHaveLength(1); + sent[0].resolve(); + await vi.advanceTimersByTimeAsync(0); + expect(sent.map((s) => s.event.content)).toEqual(['1', '123']); + }); + + it('never posts older content after newer (monotonic per chunkId)', async () => { + const { sent, coalescer } = harness({ autoResolve: true, latencyMs: 37 }); + let content = ''; + for (let i = 0; i < 100; i++) { + content += 'x'; + void coalescer.submit(chunk('a', content)); + await vi.advanceTimersByTimeAsync(7); + } + const done = coalescer.submit(chunk('a', content, true)); + await vi.advanceTimersByTimeAsync(500); + await done; + const lengths = sent.map((s) => s.event.content.length); + for (let i = 1; i < lengths.length; i++) { + expect(lengths[i]).toBeGreaterThanOrEqual(lengths[i - 1]); + } + expect(sent[sent.length - 1].event).toEqual(chunk('a', content, true)); + }); + + it('two interleaved chunkIds do not block each other', async () => { + const { sent, coalescer } = harness({ autoResolve: false }); + void coalescer.submit(chunk('a', 'a1')); + void coalescer.submit(chunk('a', 'a2')); + // `a` has a request stuck in flight; a new message's chunk still goes at once. + void coalescer.submit(chunk('b', 'b1')); + expect(sent.map((s) => s.event.content)).toEqual(['a1', 'b1']); + sent[1].resolve(); + void coalescer.submit(chunk('b', 'b2', true)); + await vi.advanceTimersByTimeAsync(0); + expect(sent.map((s) => s.event.content)).toEqual(['a1', 'b1', 'b2']); + sent[0].resolve(); + await vi.advanceTimersByTimeAsync(STREAMING_CHUNK_MIN_INTERVAL_MS); + expect(sent.map((s) => s.event.content)).toEqual(['a1', 'b1', 'b2', 'a2']); + }); + + it('flush() sends pending latest content at once and waits for it', async () => { + const { sent, coalescer } = harness({ autoResolve: true, latencyMs: 20 }); + void coalescer.submit(chunk('a', 'a1')); + void coalescer.submit(chunk('a', 'a1 a2')); + void coalescer.submit(chunk('b', 'b1')); + void coalescer.submit(chunk('b', 'b1 b2')); + let flushed = false; + const flush = coalescer.flush().then(() => { + flushed = true; + }); + await vi.advanceTimersByTimeAsync(45); // two 20 ms round-trips, well under the 100 ms spacing + await flush; + expect(flushed).toBe(true); + expect(sent.map((s) => s.event.content)).toEqual(['a1', 'b1', 'a1 a2', 'b1 b2']); + }); + + it('a send failure does not wedge the chunk', async () => { + const send = vi + .fn<(event: StreamingChunkEvent) => Promise>() + .mockRejectedValueOnce(new Error('boom')) + .mockResolvedValue(undefined); + const coalescer = new StreamingChunkCoalescer({ send }); + void coalescer.submit(chunk('a', '1')); + await coalescer.submit(chunk('a', '12', true)); + expect(send.mock.calls.map((c) => c[0].content)).toEqual(['1', '12']); + }); + + it.each([60, 300])( + 'measured: a 20 updates/s stream for 10 s with %i ms request latency', + async (latencyMs) => { + const { sent, coalescer } = harness({ autoResolve: true, latencyMs }); + const updates = 200; + let content = ''; + for (let i = 0; i < updates; i++) { + content += `token${i} `; + void coalescer.submit(chunk('a', content)); + await vi.advanceTimersByTimeAsync(50); + } + const done = coalescer.submit(chunk('a', content, true)); + await vi.advanceTimersByTimeAsync(1000); + await done; + const posts = sent.length; + // Before: one POST per update + the final (201). Now ≤ one per 100 ms, + // fewer when the server is slow (one in flight at a time). + expect(posts).toBeLessThanOrEqual(updates / 2 + 2); + expect(sent[posts - 1].event).toEqual(chunk('a', content, true)); + console.warn( + `[coalescer] 20/s, ${latencyMs} ms latency: ${updates + 1} submits → ${posts} posts`, + ); + }, + ); +}); diff --git a/apps/cli/src/agents/acp/publisher.ts b/apps/cli/src/agents/acp/publisher.ts index 062672a7..bf5a5236 100644 --- a/apps/cli/src/agents/acp/publisher.ts +++ b/apps/cli/src/agents/acp/publisher.ts @@ -29,6 +29,10 @@ import { _transport } from '../../services/streaming/transport'; import { resolveApiBaseUrl } from '@codeam/shared'; import type { AwaitingAnswerEvent, SessionBlocker, StreamingChunkEvent } from '@codeam/shared'; import { log } from '../../services/logger'; +import { + STREAMING_CHUNK_MIN_INTERVAL_MS, + StreamingChunkCoalescer, +} from '../../services/streaming/chunk-coalescer'; export interface AcpPublisherOptions { sessionId: string; @@ -65,9 +69,9 @@ export interface AcpPublisherOptions { export const STREAMING_CHUNK_SNAPSHOT_MAX_CHARS = 56 * 1024; /** - * Above this size a chunk's snapshots are coalesced to at most one POST per - * window: a 50 KB tool result re-sent on each of 20 deltas/s is 1 MB/s that - * mobile immediately replaces with the next snapshot anyway. + * Above this size a chunk's snapshots are spaced by the wider window (smaller + * ones by {@link STREAMING_CHUNK_MIN_INTERVAL_MS}): a 50 KB tool result re-sent + * on each of 20 deltas/s is 1 MB/s that mobile immediately replaces anyway. */ export const STREAMING_CHUNK_COALESCE_MIN_CHARS = 8 * 1024; export const STREAMING_CHUNK_COALESCE_WINDOW_MS = 300; @@ -131,11 +135,15 @@ export class AcpPublisher { /** chunkIds whose truncated snapshot already went out — nothing new to show * until the terminal frame. */ private readonly truncatedChunks = new Set(); - private readonly lastChunkPostAt = new Map(); - private readonly pendingChunkPosts = new Map< - string, - { timer: ReturnType; event: StreamingChunkEvent } - >(); + /** Latest-wins, one-in-flight-per-chunkId pacing of the streaming-chunk + * feed (2026-10-29: one CLI posted 10–30 cumulative snapshots/s). */ + private readonly chunkCoalescer = new StreamingChunkCoalescer({ + send: (event) => this.postStreamingChunk(event), + intervalFor: (event) => + event.content.length >= STREAMING_CHUNK_COALESCE_MIN_CHARS + ? STREAMING_CHUNK_COALESCE_WINDOW_MS + : STREAMING_CHUNK_MIN_INTERVAL_MS, + }); /** Latched on an unrecoverable 401/403 — every surface stops posting * (2026-06-28 incident: a dead-token publisher spammed 401 ×34 while * the agent's replies silently never reached the phone). */ @@ -320,15 +328,13 @@ export class AcpPublisher { * ({@link publishOutput}) is independent and carries `type:'text'` * for the main chat bubble; this feed is purely additive for the * richer SessionDetail surface. + * + * POSTs are paced per chunkId by {@link StreamingChunkCoalescer}: latest + * snapshot wins, one request in flight, and the `isFinal` frame always goes + * last. The promise settles once this content (or a newer one) is sent. */ async publishStreamingChunk(event: StreamingChunkEvent): Promise { const { chunkId, isFinal } = event; - // A newer snapshot supersedes any coalesced one still waiting. - const pending = this.pendingChunkPosts.get(chunkId); - if (pending) { - clearTimeout(pending.timer); - this.pendingChunkPosts.delete(chunkId); - } const bounded = boundStreamingChunkContent(event.content); if (bounded.truncated) { if (!isFinal && this.truncatedChunks.has(chunkId)) return; // already told mobile; nothing new to show @@ -340,29 +346,18 @@ export class AcpPublisher { } this.truncatedChunks.add(chunkId); } - if (isFinal) { - this.truncatedChunks.delete(chunkId); - this.lastChunkPostAt.delete(chunkId); - } else if (!bounded.truncated && bounded.content.length >= STREAMING_CHUNK_COALESCE_MIN_CHARS) { - const last = this.lastChunkPostAt.get(chunkId) ?? 0; - const wait = STREAMING_CHUNK_COALESCE_WINDOW_MS - (Date.now() - last); - if (wait > 0) { - // Keep the latest snapshot; one trailing post per window. - const timer = setTimeout(() => { - const latest = this.pendingChunkPosts.get(chunkId); - this.pendingChunkPosts.delete(chunkId); - if (latest) void this.postStreamingChunk(latest.event); - }, wait); - this.pendingChunkPosts.set(chunkId, { timer, event }); - return; - } - } - await this.postStreamingChunk({ ...event, content: bounded.content }); + if (isFinal) this.truncatedChunks.delete(chunkId); + await this.chunkCoalescer.submit({ ...event, content: bounded.content }); + } + + /** Send any coalesced streaming-chunk snapshot still waiting, and wait for + * the in-flight ones — so shutdown never loses the latest content. */ + flushStreamingChunks(): Promise { + return this.chunkCoalescer.flush(); } private async postStreamingChunk(event: StreamingChunkEvent): Promise { const url = `${this.apiBase}/api/sessions/${encodeURIComponent(this.opts.sessionId)}/streaming-chunk`; - if (!event.isFinal) this.lastChunkPostAt.set(event.chunkId, Date.now()); try { const { statusCode, body } = await this.postWithReauth( url, diff --git a/apps/cli/src/agents/acp/runner.ts b/apps/cli/src/agents/acp/runner.ts index 14351ec1..e5b0312e 100644 --- a/apps/cli/src/agents/acp/runner.ts +++ b/apps/cli/src/agents/acp/runner.ts @@ -2360,6 +2360,12 @@ export async function runAcpSession(opts: AcpRunnerOptions): Promise { const shutdown = async (signal: NodeJS.Signals) => { showInfo(`Shutting down ACP session (${signal})…`); clearTimeout(prewarmTimer); + // Coalesced streaming-chunk snapshots still waiting go out before exit; + // bounded so a hung request can't stall the shutdown. + await Promise.race([ + publisher.flushStreamingChunks(), + new Promise((resolve) => setTimeout(resolve, 2_000).unref()), + ]); // AWAITED goodbye — a fire-and-forget `online:false` never made it out // before `process.exit` below, so mobile kept the session ONLINE. await stopRelayWithGoodbye(relay); diff --git a/apps/cli/src/services/streaming-emitter.service.ts b/apps/cli/src/services/streaming-emitter.service.ts index cd64dfed..a4b2d066 100644 --- a/apps/cli/src/services/streaming-emitter.service.ts +++ b/apps/cli/src/services/streaming-emitter.service.ts @@ -10,6 +10,7 @@ import type { import { resolveApiBaseUrl, renderToLines } from '@codeam/shared'; import { log } from './logger'; import { _transport } from './streaming/transport'; +import { StreamingChunkCoalescer } from './streaming/chunk-coalescer'; import type { RuntimeStrategy } from '../agents/strategy'; /** @@ -41,7 +42,9 @@ import type { RuntimeStrategy } from '../agents/strategy'; * * Throttle: we coalesce within a 50 ms window — mirrors the backend's * SSE-coalescing logic and keeps a hot Claude tool-loop from - * generating one HTTP roundtrip per render-tick. + * generating one HTTP roundtrip per render-tick. The POSTs themselves go + * through `StreamingChunkCoalescer` (latest wins, one in flight and ≥100 ms + * apart per chunkId, the final frame always last). * * Best-effort transport: every POST is fire-and-forget with a small * retry budget; we never block the agent on emission. Failures land @@ -141,6 +144,10 @@ export class StreamingEmitterService { */ private lastSelectorSignature: string | null = null; private selectorFirstSeenAt = 0; + /** Latest-wins, one-in-flight-per-chunkId pacing of the chunk POSTs. */ + private readonly chunkCoalescer = new StreamingChunkCoalescer({ + send: (event) => this.postChunk(event), + }); constructor(private readonly opts: StreamingEmitterOptions) { this.apiBase = opts.apiBaseUrl ?? API_BASE; @@ -183,13 +190,14 @@ export class StreamingEmitterService { const finalContent = this.activeChunk.currentContent; const chunk = this.activeChunk; this.activeChunk = null; - await this.postChunk({ + await this.chunkCoalescer.submit({ chunkId: chunk.chunkId, kind: chunk.kind, content: finalContent, isFinal: true, }); } + await this.chunkCoalescer.flush(); log.trace('streamingEmitter', `stopped session=${this.opts.sessionId.slice(0, 8)}`); } @@ -265,7 +273,7 @@ export class StreamingEmitterService { if (this.activeChunk) { const closing = this.activeChunk; this.activeChunk = null; - void this.postChunk({ + void this.chunkCoalescer.submit({ chunkId: closing.chunkId, kind: closing.kind, content: closing.currentContent, @@ -290,7 +298,7 @@ export class StreamingEmitterService { if (chunk.currentContent === chunk.emittedContent) return; chunk.emittedContent = chunk.currentContent; chunk.lastEmitAt = now; - void this.postChunk({ + void this.chunkCoalescer.submit({ chunkId: chunk.chunkId, kind: chunk.kind, content: chunk.currentContent, diff --git a/apps/cli/src/services/streaming/chunk-coalescer.ts b/apps/cli/src/services/streaming/chunk-coalescer.ts new file mode 100644 index 00000000..30daa71b --- /dev/null +++ b/apps/cli/src/services/streaming/chunk-coalescer.ts @@ -0,0 +1,136 @@ +import type { StreamingChunkEvent } from '@codeam/shared'; + +/** + * Minimum spacing between two POSTs of the SAME chunkId. Every snapshot is + * CUMULATIVE, so a newer one fully replaces the older on mobile — sending one + * per ACP delta (10–30/s for a single chunk on 2026-10-29) only loads the API. + */ +export const STREAMING_CHUNK_MIN_INTERVAL_MS = 100; + +/** How many finalised chunkIds are remembered to drop late non-final snapshots. */ +const FINALIZED_MEMORY = 512; + +interface Pending { + event: StreamingChunkEvent; + waiters: Array<() => void>; +} + +interface ChunkState { + pending: Pending | null; + inFlight: Promise | null; + timer: ReturnType | null; + lastSentAt: number; +} + +export interface StreamingChunkCoalescerOptions { + /** Performs one POST. Must not throw; retries/logging live inside it. */ + send: (event: StreamingChunkEvent) => Promise; + /** Spacing for a given snapshot (lets big snapshots use a wider window). */ + intervalFor?: (event: StreamingChunkEvent) => number; +} + +/** + * Per-chunkId latest-wins coalescing for the streaming-chunk feed: + * - the first snapshot of a chunkId goes out immediately; + * - at most ONE request in flight per chunkId, posts spaced by the interval; + * snapshots arriving meanwhile collapse into the latest one; + * - an `isFinal` snapshot supersedes anything pending, skips the spacing wait + * and is sent right after the in-flight request, so it is always the last + * POST of its chunkId; later non-final snapshots of that id are dropped; + * - chunkIds never wait on each other. + * The promise from {@link submit} settles once the submitted content (or a + * newer snapshot that superseded it) has been sent. + */ +export class StreamingChunkCoalescer { + private readonly chunks = new Map(); + private readonly finalized = new Set(); + + constructor(private readonly opts: StreamingChunkCoalescerOptions) {} + + submit(event: StreamingChunkEvent): Promise { + const { chunkId } = event; + if (this.finalized.has(chunkId)) return Promise.resolve(); + const state = this.chunks.get(chunkId) ?? this.open(chunkId); + const settled = new Promise((resolve) => { + const waiters = state.pending ? state.pending.waiters : []; + waiters.push(resolve); + state.pending = { event, waiters }; + }); + if (event.isFinal) { + this.markFinalized(chunkId); + this.drain(chunkId, state); + } else { + this.schedule(chunkId, state); + } + return settled; + } + + /** Send every pending snapshot now (ignoring the spacing) and wait for all + * requests to settle. Used on shutdown so the latest content is not lost. */ + async flush(): Promise { + const work = Array.from(this.chunks.entries()).map(async ([chunkId, state]) => { + while (state.inFlight || state.pending) { + if (!state.inFlight) this.drain(chunkId, state); + await state.inFlight; + } + }); + await Promise.all(work); + } + + private open(chunkId: string): ChunkState { + const state: ChunkState = { pending: null, inFlight: null, timer: null, lastSentAt: 0 }; + this.chunks.set(chunkId, state); + return state; + } + + private schedule(chunkId: string, state: ChunkState): void { + if (state.inFlight || state.timer || !state.pending) return; + const wait = this.intervalFor(state.pending.event) - (Date.now() - state.lastSentAt); + if (wait <= 0) { + this.drain(chunkId, state); + return; + } + state.timer = setTimeout(() => { + state.timer = null; + this.drain(chunkId, state); + }, wait); + } + + /** Send the pending snapshot now, unless a request is still in flight + * (its completion re-enters here or {@link schedule}). */ + private drain(chunkId: string, state: ChunkState): void { + if (state.inFlight || !state.pending) return; + if (state.timer) { + clearTimeout(state.timer); + state.timer = null; + } + const { event, waiters } = state.pending; + state.pending = null; + state.lastSentAt = Date.now(); + state.inFlight = this.opts + .send(event) + .catch(() => undefined) + .then(() => { + state.inFlight = null; + for (const resolve of waiters) resolve(); + if (!state.pending) { + if (event.isFinal) this.chunks.delete(chunkId); + return; + } + if (state.pending.event.isFinal) this.drain(chunkId, state); + else this.schedule(chunkId, state); + }); + } + + private intervalFor(event: StreamingChunkEvent): number { + return this.opts.intervalFor?.(event) ?? STREAMING_CHUNK_MIN_INTERVAL_MS; + } + + private markFinalized(chunkId: string): void { + this.finalized.add(chunkId); + if (this.finalized.size > FINALIZED_MEMORY) { + const oldest = this.finalized.values().next().value; + if (oldest !== undefined) this.finalized.delete(oldest); + } + } +}