diff --git a/.env.example b/.env.example index b9bd77b768..d895a940cb 100644 --- a/.env.example +++ b/.env.example @@ -992,6 +992,14 @@ HELP_AND_FAQ_URL=https://librechat.ai # pre-v2 build while v2 generations remain active in Redis. # GENERATION_PROTOCOL_VERSION=2 +# Coalesce streamed model/tool-argument deltas into windowed Redis publications (ms). +# Unset or 0 (default) publishes per delta. 25 is recommended: it batches the publish +# EVAL and the durable append across the window (fewer Redis round trips and lower +# Redis CPU at high token rates) at the cost of up to one window of added delivery +# latency. Enable only after EVERY replica runs a build with batch-frame support; +# older subscribers drop coalesced frames. Values are capped at 1000. +# STREAM_DELTA_COALESCE_MS=25 + # Single Redis instance # REDIS_URI=redis://127.0.0.1:6379 diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index fdf504f886..c63417ae33 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -68,6 +68,7 @@ import { InMemoryEventTransport } from './implementations/InMemoryEventTransport import { InMemoryJobStore } from './implementations/InMemoryJobStore'; import { normalizeResumeRunStepIndices } from '~/agents/hitl/resume'; import { emitChunkWithReceipt } from './internal/chunkPublication'; +import { resolveCoalesceWindowMs } from './internal/coalescing'; import { filterPersistableAbortContent } from './abortContent'; import { toClientPendingAction } from '~/agents/hitl/policy'; import { ApprovalLifecycle, pausePersistenceActionId } from './ApprovalLifecycle'; @@ -79,6 +80,12 @@ const APPROVAL_EXPIRED_ERROR = 'Approval expired before a decision was made'; /** Error surfaced to any client still attached when a stale/hung job is reaped. */ const REAPED_JOB_ERROR = 'Generation timed out'; +/** Un-awaited coalesced publications allowed per stream before the emitter + * awaits a receipt. Healthy settlement keeps outstanding counts in the single + * digits; this trips only when Redis stalls, bounding buffered batches and + * queued commands by pacing the producer instead of growing without limit. */ +const MAX_OUTSTANDING_COALESCED_RECEIPTS = 256; + /** Bounded completed-request replay horizon. It exceeds the default 24-hour * approval window; if a custom/live job outlasts it, `resumeClaimedGeneration` * atomically adopts the reacquired claim instead of replacing that job. */ @@ -323,6 +330,19 @@ function omitAlreadyAppliedSteers(items: SteerQueueItem[], content: unknown[]): return items.filter((item) => !appliedIds.has(item.steerId)); } +/** + * Streaming deltas eligible for windowed publish/append coalescing. High-volume, + * order-preserved by per-event sequences, and consumed for the generation fence + * only — unlike control events, nothing awaits their publication for correctness. + */ +function isCoalescableDeltaEvent(eventType: string | undefined): boolean { + return ( + eventType === 'on_message_delta' || + eventType === 'on_reasoning_delta' || + eventType === 'on_run_step_delta' + ); +} + function getReplayStepId(event: t.ServerSentEvent): unknown { if (!('event' in event) || !event.data || typeof event.data !== 'object') { return undefined; @@ -556,6 +576,8 @@ interface RuntimeJobState { >; /** Prevents later events from overtaking the initial `created` metadata write and publish. */ createdEventPublication?: Promise; + /** Coalesced delta publications emitted but not yet settled by a window flush. */ + outstandingCoalescedReceipts?: number; hasSubscriber: boolean; /** Advances whenever every local SSE subscriber for one attachment generation leaves. */ attachmentGeneration: number; @@ -640,6 +662,10 @@ class GenerationJobManagerClass { /** Whether we're using Redis stores */ private _isRedis = false; + /** Whether streaming-delta publish/append coalescing is enabled (Redis only). + * Off (the default) preserves the awaited per-event emitChunk contract exactly. */ + private _deltaCoalescingEnabled = false; + /** Whether to cleanup event transport immediately on job completion */ private _cleanupOnComplete = true; @@ -743,6 +769,17 @@ class GenerationJobManagerClass { this._steering = new SteeringLifecycle(this.jobStore); this.eventTransport = services.eventTransport; this._isRedis = services.isRedis ?? false; + /** Coalescing needs BOTH configured services to actually batch: the flush + * capabilities are how implementations advertise it. A custom transport + * without them would silently lose the awaited per-event ordering contract + * (its receipts are un-awaited on the coalescable path), and a batching + * transport over a per-event store would let the durable log trail the + * sequence counter by a full window, breaking the resume frontier. */ + this._deltaCoalescingEnabled = + this._isRedis && + resolveCoalesceWindowMs() > 0 && + typeof services.eventTransport.flushPendingChunks === 'function' && + typeof services.jobStore.flushPendingAppends === 'function'; this._cleanupOnComplete = services.cleanupOnComplete ?? true; this.shuttingDown = false; this.syncRunningJobMetrics(); @@ -3009,6 +3046,22 @@ class GenerationJobManagerClass { * a terminal client event before this succeeds: abort, pause, and completion * can all race on the same generation epoch. */ + /** + * Drain both delta coalescers ahead of a terminal status CAS. Coalesced + * deltas still buffered for a stream must land under its live status: + * flushed after the CAS they fence against the generation's own completion, + * and the false receipts retire a healthy runtime and error-close its + * subscribers ahead of the terminal frame. Every terminal transition that + * can interrupt a live emitter (claim, abort, shutdown) must call this + * before its CAS; no-op (two Map lookups) when coalescing is off or idle. + */ + private async flushCoalescedStreamBuffers(streamId: string): Promise { + await Promise.all([ + this.jobStore.flushPendingAppends?.(streamId), + this.eventTransport.flushPendingChunks?.(streamId), + ]); + } + async claimTerminalJob( streamId: string, status: TerminalJobClaim['status'], @@ -3073,6 +3126,7 @@ class GenerationJobManagerClass { const createdAt = jobData.createdAt; const runtime = observedRuntime?.createdAt === createdAt ? observedRuntime : undefined; const terminalError = status === 'error' ? (error ?? 'Generation failed') : undefined; + await this.flushCoalescedStreamBuffers(streamId); const completedAt = Date.now(); const drainedSteers = await this.jobStore.transitionStatusAndDrainSteers(streamId, { from: sourceStatus, @@ -3539,6 +3593,10 @@ class GenerationJobManagerClass { } const runtime = observedRuntime?.createdAt === jobData.createdAt ? observedRuntime : undefined; + /** Abort claims terminal state through its own CAS loop below (not + * claimTerminalJob), so it must drain the coalescers itself — and ahead of + * the content snapshot, so a chunk-log reconstruction sees the window tail. */ + await this.flushCoalescedStreamBuffers(streamId); /** Snapshot before claiming terminal state. This is non-destructive: if a * same-epoch approval resume wins the later CAS, its content and steer * queue remain fully owned by that resumed run. */ @@ -5159,6 +5217,19 @@ class GenerationJobManagerClass { this.saveRunStepFromEvent(streamId, eventData as Record, runtime.createdAt); } + /** + * One decision drives both durable-log and publish batching: the append and + * the sequence allocation for an event must stay tightly coupled in time, or + * the resume frontier (chunk-log snapshot → sequence-counter sync) misreads + * a window's tail as already-delivered or as duplicates. + */ + const coalescableDelta = + this._deltaCoalescingEnabled && + !runtime.startupTelemetry && + options?.durable !== true && + options?.deliveredSteer == null && + isCoalescableDeltaEvent(eventType); + // For Redis mode, persist chunk for later reconstruction (fire-and-forget for resumability) if (this._isRedis) { // The SSE event structure is { event: string, data: unknown, ... } @@ -5170,6 +5241,7 @@ class GenerationJobManagerClass { { event: eventType, data: eventData }, runtime.createdAt, options?.deliveredSteer, + coalescableDelta ? { coalesce: true } : undefined, ); if (options?.durable === true) { @@ -5237,6 +5309,60 @@ class GenerationJobManagerClass { return; } + /** + * Streaming deltas dominate publication volume, and their receipt is consumed + * only for the generation fence (`false` retires the runtime) — never awaited + * for content correctness. Marking them coalescable lets the Redis transport + * batch a window of them into one sequenced frame, and NOT awaiting here takes + * the per-delta publish round trip off the provider-stream consumption path. + * The fence continuation mirrors the fire-and-forget appendChunk fence above. + * Fenced emissions (durable, steer receipts, created) and telemetry-observed + * runs stay on the awaited per-event path below. + */ + if (coalescableDelta) { + const publication = emitChunkWithReceipt( + this.eventTransport, + streamId, + event, + runtime.createdAt, + { coalesce: true }, + ); + if (buffered) { + runtime.earlyEventSequencePromises.push( + publication.then( + (published) => (typeof published === 'number' ? published : undefined), + () => undefined, + ), + ); + } + runtime.outstandingCoalescedReceipts = (runtime.outstandingCoalescedReceipts ?? 0) + 1; + void publication.then( + (published) => { + runtime.outstandingCoalescedReceipts = (runtime.outstandingCoalescedReceipts ?? 1) - 1; + if (published === false) { + this.retireRuntimeAfterDurableFence(streamId, runtime); + } + }, + (err) => { + runtime.outstandingCoalescedReceipts = (runtime.outstandingCoalescedReceipts ?? 1) - 1; + logger.error(`[GenerationJobManager] Failed to publish coalesced chunk:`, err); + }, + ); + /** + * Backpressure only under distress. Healthy settlement is one window plus + * a round trip (~30ms), so outstanding receipts sit in the single digits + * even at hundreds of deltas per second and this await never runs. If + * Redis stalls, the un-awaited path would otherwise accumulate batches, + * resolver closures, and queued commands without bound — awaiting one + * receipt paces the producer to Redis exactly like the flag-off path, + * with memory capped near the threshold instead of one delta. + */ + if (runtime.outstandingCoalescedReceipts >= MAX_OUTSTANDING_COALESCED_RECEIPTS) { + await publication.catch(() => undefined); + } + return; + } + if (!buffered && !runtime.startupTelemetry) { try { const published = await emitChunkWithReceipt( @@ -5320,6 +5446,16 @@ class GenerationJobManagerClass { if (this.runtimeState.get(streamId) !== runtime) { return; } + /** A runtime whose stop signal already landed (cross-replica abort, + * replacement handshake) observes this fence as a consequence of its own + * termination — most often a coalesced window draining after the abort + * CAS. Those flows own terminal delivery and cleanup; the forced teardown + * below would error-close local subscribers in a race with the FINAL + * frame they are about to receive. It remains the backstop for the + * lost-signal case, which is exactly a fence on a NOT-yet-aborted owner. */ + if (runtime.abortController.signal.aborted) { + return; + } runtime.startupTelemetry?.end('replaced'); runtime.startupTelemetry = undefined; runtime.abortController.abort(); @@ -6655,6 +6791,9 @@ class GenerationJobManagerClass { runtime.lastSubscriberCleanupGeneration = runtime.attachmentGeneration; await this.persistSubscriberCleanup(streamId, runtime); } + /** Shutdown interrupts live emitters, so their coalesced window must + * drain before the terminal CAS — same rule as claim/abort. */ + await this.flushCoalescedStreamBuffers(streamId); const finalized = await this.jobStore.transitionStatus(streamId, { from: 'running', to: 'error', diff --git a/packages/api/src/stream/__tests__/deltaCoalescing.manual.spec.ts b/packages/api/src/stream/__tests__/deltaCoalescing.manual.spec.ts new file mode 100644 index 0000000000..0b6a6d635c --- /dev/null +++ b/packages/api/src/stream/__tests__/deltaCoalescing.manual.spec.ts @@ -0,0 +1,286 @@ +/* eslint jest/no-standalone-expect: ["error", { "additionalTestBlockFunctions": ["testRedis"] }] */ +import type { Redis, Cluster } from 'ioredis'; +import type { ServerSentEvent } from '~/types'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; +import { createStreamServices } from '~/stream/createStreamServices'; + +/** + * MANUAL benchmark for streaming-delta coalescing. Not part of any CI suite. + * + * Measures, per (window x event-rate) scenario against a real Redis: + * - Redis EVAL invocations and script time (INFO commandstats) + * - Redis engine CPU (INFO cpu) + * - Node process CPU + * - producer await stall (the round trips production callbacks wait on) + * - subscriber delivery latency percentiles + * - published frame mix (single vs batch, mean batch size) + * + * Run: + * USE_REDIS=true REDIS_URI=redis://127.0.0.1:6379 npx jest \ + * src/stream/__tests__/deltaCoalescing.manual.spec.ts --coverage=false --runInBand --forceExit + */ +describe('delta coalescing benchmark (manual)', () => { + let originalEnv: NodeJS.ProcessEnv; + let ioredisClient: Redis | Cluster | null = null; + const testPrefix = 'DeltaCoalescing-Bench'; + const redisConfigured = process.env.USE_REDIS === 'true'; + const testRedis = redisConfigured ? test : test.skip; + + beforeAll(async () => { + originalEnv = { ...process.env }; + process.env.USE_REDIS = process.env.USE_REDIS ?? 'true'; + process.env.REDIS_URI = process.env.REDIS_URI ?? 'redis://127.0.0.1:6379'; + process.env.REDIS_KEY_PREFIX = testPrefix; + jest.resetModules(); + const redisModule = await import('~/cache/redisClients'); + ioredisClient = redisModule.ioredisClient; + }); + + afterAll(async () => { + if (ioredisClient) { + try { + const keys = await ioredisClient.keys(`${testPrefix}*`); + const streamKeys = await ioredisClient.keys(`stream:*`); + await Promise.all([...keys, ...streamKeys].map((key) => ioredisClient!.del(key))); + await ioredisClient.quit(); + } catch { + /* ignore */ + } + } + process.env = originalEnv; + }); + + interface CommandStat { + calls: number; + usec: number; + } + + function parseCommandStats(info: string): Map { + const stats = new Map(); + for (const line of info.split('\n')) { + const match = line.match(/^cmdstat_([a-z|]+):calls=(\d+),usec=(\d+)/); + if (match) { + stats.set(match[1], { calls: Number(match[2]), usec: Number(match[3]) }); + } + } + return stats; + } + + function parseCpu(info: string): number { + let total = 0; + for (const line of info.split('\n')) { + const match = line.match(/^used_cpu_(?:sys|user):([\d.]+)/); + if (match) { + total += Number(match[1]); + } + } + return total; + } + + function percentile(sorted: number[], p: number): number { + if (sorted.length === 0) { + return 0; + } + const idx = Math.min(sorted.length - 1, Math.ceil((p / 100) * sorted.length) - 1); + return sorted[Math.max(0, idx)]; + } + + interface ScenarioResult { + scenario: string; + windowMs: number; + targetRate: number | 'serial'; + events: number; + achievedRate: number; + wallMs: number; + awaitStallMs: number; + evalCalls: number; + evalUsec: number; + publishCalls: number; + xaddCalls: number; + redisCpuSec: number; + nodeCpuMs: number; + deliveryP50: number; + deliveryP95: number; + deliveryMax: number; + singleFrames: number; + batchFrames: number; + meanBatchSize: number; + delivered: number; + } + + const DELTA_TEXT = 'lorem ipsum token piece that approximates a realistic streamed delta '; + + async function runScenario(options: { + windowMs: number; + rate: number | 'serial'; + events: number; + }): Promise { + const { windowMs, rate, events } = options; + process.env.STREAM_DELTA_COALESCE_MS = String(windowMs); + + const manager = new GenerationJobManagerClass(); + manager.configure(createStreamServices({ useRedis: true, redisClient: ioredisClient! })); + manager.initialize(); + + const streamId = `bench-${windowMs}ms-${rate}-${Date.now()}`; + await manager.createJob(streamId, 'bench-user', streamId); + + const latencies: number[] = []; + let delivered = 0; + const subscription = await manager.subscribe(streamId, (event) => { + delivered++; + const emittedAt = (event as { data?: { t?: number } }).data?.t; + if (typeof emittedAt === 'number') { + latencies.push(Date.now() - emittedAt); + } + }); + + const rawSubscriber = (ioredisClient as Redis).duplicate(); + let singleFrames = 0; + let batchFrames = 0; + let batchedEvents = 0; + rawSubscriber.on('message', (_channel: string, message: string) => { + const parsed = JSON.parse(message) as { type: string; events?: unknown[] }; + if (parsed.type === 'chunk') { + singleFrames++; + } else if (parsed.type === 'chunk_batch') { + batchFrames++; + batchedEvents += parsed.events?.length ?? 0; + } + }); + await rawSubscriber.subscribe(`stream:{${streamId}}:events`); + + await (ioredisClient as Redis).config('RESETSTAT'); + const cpuBefore = parseCpu(await ioredisClient!.info('cpu')); + const nodeCpuBefore = process.cpuUsage(); + const started = Date.now(); + let awaitStallMs = 0; + + if (rate === 'serial') { + for (let i = 0; i < events; i++) { + const before = Date.now(); + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { + id: 'step-1', + t: before, + delta: { content: [{ type: 'text', text: `${DELTA_TEXT}${i}` }] }, + }, + }); + awaitStallMs += Date.now() - before; + } + } else { + const intervalMs = 1000 / rate; + for (let i = 0; i < events; i++) { + const deadline = started + i * intervalMs; + const wait = deadline - Date.now(); + if (wait > 0) { + await new Promise((resolve) => setTimeout(resolve, wait)); + } + const before = Date.now(); + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { + id: 'step-1', + t: before, + delta: { content: [{ type: 'text', text: `${DELTA_TEXT}${i}` }] }, + }, + }); + awaitStallMs += Date.now() - before; + } + } + const wallMs = Date.now() - started; + + await new Promise((resolve) => setTimeout(resolve, Math.max(200, windowMs * 3))); + + const stats = parseCommandStats(await ioredisClient!.info('commandstats')); + const cpuAfter = parseCpu(await ioredisClient!.info('cpu')); + const nodeCpuAfter = process.cpuUsage(nodeCpuBefore); + + /** Terminal-first teardown: a pre-terminal unsubscribe fires the + * disconnected-subscriber persistence chain (XRANGE over the whole chunk + * log + content writes), which would bleed into the NEXT scenario's + * RESETSTAT window. After the final event it is skipped entirely. */ + await manager.emitDone(streamId, { final: true, streamId } as unknown as ServerSentEvent); + await manager.completeJob(streamId); + subscription?.unsubscribe(); + await manager.destroy(); + rawSubscriber.disconnect(); + + await new Promise((resolve) => setTimeout(resolve, 500)); + const staleKeys = [ + ...(await ioredisClient!.keys(`${testPrefix}*`)), + ...(await ioredisClient!.keys(`stream:*`)), + ]; + await Promise.all(staleKeys.map((key) => ioredisClient!.del(key))); + await new Promise((resolve) => setTimeout(resolve, 250)); + + latencies.sort((a, b) => a - b); + return { + scenario: `${windowMs}ms @ ${rate}`, + windowMs, + targetRate: rate, + events, + achievedRate: Math.round((events / wallMs) * 1000), + wallMs, + awaitStallMs, + evalCalls: stats.get('eval')?.calls ?? 0, + evalUsec: stats.get('eval')?.usec ?? 0, + publishCalls: stats.get('publish')?.calls ?? 0, + xaddCalls: stats.get('xadd')?.calls ?? 0, + redisCpuSec: Number((cpuAfter - cpuBefore).toFixed(3)), + nodeCpuMs: Math.round((nodeCpuAfter.user + nodeCpuAfter.system) / 1000), + deliveryP50: percentile(latencies, 50), + deliveryP95: percentile(latencies, 95), + deliveryMax: latencies.length > 0 ? latencies[latencies.length - 1] : 0, + singleFrames, + batchFrames, + meanBatchSize: batchFrames > 0 ? Number((batchedEvents / batchFrames).toFixed(1)) : 0, + delivered, + }; + } + + testRedis( + 'matrix: window x rate', + async () => { + const results: ScenarioResult[] = []; + + /** JIT/connection warmup; discarded so the first recorded rows are not inflated. */ + await runScenario({ windowMs: 0, rate: 200, events: 400 }); + await runScenario({ windowMs: 25, rate: 200, events: 400 }); + + for (const rate of [40, 100, 200] as const) { + for (const windowMs of [0, 20, 25, 50]) { + results.push(await runScenario({ windowMs, rate, events: rate * 8 })); + } + } + for (const windowMs of [0, 25]) { + results.push(await runScenario({ windowMs, rate: 'serial', events: 3000 })); + } + + console.table( + results.map((r) => ({ + scenario: r.scenario, + events: r.events, + 'rate ev/s': r.achievedRate, + 'eval calls': r.evalCalls, + 'eval ms': Math.round(r.evalUsec / 1000), + 'redis cpu s': r.redisCpuSec, + 'node cpu ms': r.nodeCpuMs, + 'stall ms': r.awaitStallMs, + 'p50 ms': r.deliveryP50, + 'p95 ms': r.deliveryP95, + frames: `${r.singleFrames}s/${r.batchFrames}b`, + 'batch avg': r.meanBatchSize, + delivered: r.delivered, + })), + ); + console.log(JSON.stringify(results, null, 2)); + + for (const result of results) { + expect(result.delivered).toBeGreaterThanOrEqual(result.events); + } + }, + 15 * 60_000, + ); +}); diff --git a/packages/api/src/stream/__tests__/deltaCoalescing.spec.ts b/packages/api/src/stream/__tests__/deltaCoalescing.spec.ts new file mode 100644 index 0000000000..42ebfde2b7 --- /dev/null +++ b/packages/api/src/stream/__tests__/deltaCoalescing.spec.ts @@ -0,0 +1,167 @@ +import type { ChunkPublicationReceipt } from '~/stream/internal/chunkPublication'; +import type { ServerSentEvent } from '~/types'; +import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; +import { registerChunkPublicationCapability } from '~/stream/internal/chunkPublication'; +import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; + +jest.spyOn(console, 'log').mockImplementation(); + +const DELTA_EVENT: ServerSentEvent = { + event: 'on_message_delta', + data: { id: 'step-1', delta: { content: [{ type: 'text', text: 'token' }] } }, +} as unknown as ServerSentEvent; + +/** Structural coalescing advertisement: the manager gates on these methods. */ +function advertiseCoalescing( + jobStore: InMemoryJobStore, + eventTransport: InMemoryEventTransport, +): void { + (jobStore as { flushPendingAppends?: (id: string) => Promise }).flushPendingAppends = + async () => undefined; + (eventTransport as { flushPendingChunks?: (id: string) => Promise }).flushPendingChunks = + async () => undefined; +} + +describe('delta coalescing manager gating and backpressure', () => { + const originalWindow = process.env.STREAM_DELTA_COALESCE_MS; + + afterEach(() => { + if (originalWindow === undefined) { + delete process.env.STREAM_DELTA_COALESCE_MS; + } else { + process.env.STREAM_DELTA_COALESCE_MS = originalWindow; + } + }); + + async function createManager(options: { + windowMs: string | undefined; + advertise: boolean; + }): Promise<{ + manager: GenerationJobManagerClass; + publishCalls: Array; + settleAll: (receipt: number | false | undefined) => void; + }> { + if (options.windowMs === undefined) { + delete process.env.STREAM_DELTA_COALESCE_MS; + } else { + process.env.STREAM_DELTA_COALESCE_MS = options.windowMs; + } + + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + if (options.advertise) { + advertiseCoalescing(jobStore, eventTransport); + } + + const publishCalls: Array = []; + const resolvers: Array<(receipt: ChunkPublicationReceipt) => void> = []; + registerChunkPublicationCapability(eventTransport, (...args: unknown[]) => { + publishCalls.push(args); + return new Promise((resolve) => { + resolvers.push(resolve); + }); + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: true }); + manager.initialize(); + return { + manager, + publishCalls, + settleAll: (receipt) => { + for (const resolve of resolvers.splice(0)) { + resolve(receipt); + } + }, + }; + } + + it('does not send coalesce hints when the configured services lack the capability', async () => { + const { manager, publishCalls, settleAll } = await createManager({ + windowMs: '25', + advertise: false, + }); + await manager.createJob('gate-off', 'user-1', 'gate-off'); + + const emission = manager.emitChunk('gate-off', DELTA_EVENT); + await new Promise((resolve) => setImmediate(resolve)); + expect(publishCalls).toHaveLength(1); + /** 3-arg call shape = the awaited per-event path; a 4th options argument + * would mean the manager assumed batching the transport cannot provide. */ + expect(publishCalls[0]).toHaveLength(3); + + /** The per-event path awaits its receipt: the emission must still be + * pending until the publication settles. */ + let emitted = false; + void emission.then(() => { + emitted = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(emitted).toBe(false); + settleAll(0); + await emission; + + await manager.destroy(); + }); + + it('sends coalesce hints and stops awaiting receipts when both services advertise support', async () => { + const { manager, publishCalls, settleAll } = await createManager({ + windowMs: '25', + advertise: true, + }); + await manager.createJob('gate-on', 'user-1', 'gate-on'); + + await manager.emitChunk('gate-on', DELTA_EVENT); + expect(publishCalls).toHaveLength(1); + expect(publishCalls[0][3]).toEqual({ coalesce: true }); + + settleAll(0); + await manager.destroy(); + }); + + it('keeps the awaited path when the window is unset even with capable services', async () => { + const { manager, publishCalls, settleAll } = await createManager({ + windowMs: undefined, + advertise: true, + }); + await manager.createJob('window-off', 'user-1', 'window-off'); + + const emission = manager.emitChunk('window-off', DELTA_EVENT); + await new Promise((resolve) => setImmediate(resolve)); + expect(publishCalls[0]).toHaveLength(3); + settleAll(0); + await emission; + + await manager.destroy(); + }); + + it('applies backpressure once outstanding coalesced receipts hit the cap', async () => { + const { manager, publishCalls, settleAll } = await createManager({ + windowMs: '25', + advertise: true, + }); + await manager.createJob('backpressure', 'user-1', 'backpressure'); + + /** The cap is 256: every emission below it resolves without awaiting the + * (deliberately unsettled) publication receipts. */ + for (let i = 0; i < 255; i++) { + await manager.emitChunk('backpressure', DELTA_EVENT); + } + expect(publishCalls).toHaveLength(255); + + let saturatedResolved = false; + const saturated = manager.emitChunk('backpressure', DELTA_EVENT).then(() => { + saturatedResolved = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(publishCalls).toHaveLength(256); + expect(saturatedResolved).toBe(false); + + settleAll(0); + await saturated; + expect(saturatedResolved).toBe(true); + + await manager.destroy(); + }); +}); diff --git a/packages/api/src/stream/__tests__/deltaCoalescing.stream_integration.spec.ts b/packages/api/src/stream/__tests__/deltaCoalescing.stream_integration.spec.ts new file mode 100644 index 0000000000..ded9014a87 --- /dev/null +++ b/packages/api/src/stream/__tests__/deltaCoalescing.stream_integration.spec.ts @@ -0,0 +1,544 @@ +/* eslint jest/no-standalone-expect: ["error", { "additionalTestBlockFunctions": ["testRedis"] }] */ +import type { Redis, Cluster } from 'ioredis'; +import type { emitChunkWithReceipt as EmitChunkWithReceipt } from '~/stream/internal/chunkPublication'; +import type { ServerSentEvent } from '~/types'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; +import { createStreamServices } from '~/stream/createStreamServices'; + +jest.spyOn(console, 'log').mockImplementation(); + +/** + * Integration tests for streaming-delta coalescing (STREAM_DELTA_COALESCE_MS > 0). + * + * Coalescing batches eligible delta publications (and their durable appends) into + * one windowed frame while preserving per-event sequences, barrier ordering, + * terminal flushing, and the generation fence. + * + * Run with: USE_REDIS=true npx jest deltaCoalescing.stream_integration + */ +describe('Delta coalescing integration', () => { + let originalEnv: NodeJS.ProcessEnv; + let ioredisClient: Redis | Cluster | null = null; + const testPrefix = 'DeltaCoalescing-Integration-Test'; + const redisConfigured = process.env.USE_REDIS === 'true'; + const testRedis = redisConfigured ? test : test.skip; + + beforeAll(async () => { + originalEnv = { ...process.env }; + process.env.USE_REDIS = process.env.USE_REDIS ?? 'true'; + process.env.REDIS_URI = process.env.REDIS_URI ?? 'redis://127.0.0.1:6379'; + process.env.REDIS_KEY_PREFIX = testPrefix; + process.env.STREAM_DELTA_COALESCE_MS = '25'; + + jest.resetModules(); + const redisModule = await import('~/cache/redisClients'); + ioredisClient = redisModule.ioredisClient; + }); + + afterEach(async () => { + if (ioredisClient) { + try { + const keys = await ioredisClient.keys(`${testPrefix}*`); + const streamKeys = await ioredisClient.keys(`stream:*`); + await Promise.all([...keys, ...streamKeys].map((key) => ioredisClient!.del(key))); + } catch { + /* ignore */ + } + } + }); + + afterAll(async () => { + if (ioredisClient) { + try { + await ioredisClient.quit(); + } catch { + try { + ioredisClient.disconnect(); + } catch { + /* ignore */ + } + } + } + process.env = originalEnv; + }); + + function createRedisManager(): GenerationJobManagerClass { + const manager = new GenerationJobManagerClass(); + manager.configure( + createStreamServices({ + useRedis: true, + redisClient: ioredisClient!, + }), + ); + manager.initialize(); + return manager; + } + + /** The dynamically imported transport registers its receipt capability in its + * own module registry; the receipt helper must come from that same registry. */ + async function importFreshTransportModules(): Promise<{ + RedisEventTransport: typeof import('../implementations/RedisEventTransport').RedisEventTransport; + emitChunkWithReceipt: typeof EmitChunkWithReceipt; + }> { + const [{ RedisEventTransport }, { emitChunkWithReceipt }] = await Promise.all([ + import('../implementations/RedisEventTransport'), + import('~/stream/internal/chunkPublication'), + ]); + return { RedisEventTransport, emitChunkWithReceipt }; + } + + async function waitFor(condition: () => boolean, timeoutMs = 2000): Promise { + const deadline = Date.now() + timeoutMs; + while (!condition()) { + if (Date.now() > deadline) { + throw new Error('Timed out waiting for condition'); + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + } + + testRedis( + 'delivers a coalesced window as individually sequenced chunks in order', + async () => { + const { RedisEventTransport, emitChunkWithReceipt } = await importFreshTransportModules(); + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient!, subscriber); + const streamId = `coalesce-order-${Date.now()}`; + + const received: unknown[] = []; + const rawSubscriber = (ioredisClient as Redis).duplicate(); + const rawFrames: Array<{ type: string; count?: number }> = []; + rawSubscriber.on('message', (_channel: string, message: string) => { + const parsed = JSON.parse(message) as { type: string; events?: unknown[] }; + rawFrames.push({ type: parsed.type, count: parsed.events?.length }); + }); + await rawSubscriber.subscribe(`stream:{${streamId}}:events`); + + const subscription = transport.subscribe(streamId, { + onChunk: (event) => received.push(event), + }); + await subscription.ready; + + const receipts = await Promise.all( + Array.from({ length: 5 }, (_, i) => + emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { i } }, + undefined, + { coalesce: true }, + ), + ), + ); + + await waitFor(() => received.length === 5); + expect(received.map((event) => (event as { data: { i: number } }).data.i)).toEqual([ + 0, 1, 2, 3, 4, + ]); + const sequences = receipts.filter((value): value is number => typeof value === 'number'); + expect(sequences).toHaveLength(5); + for (let i = 1; i < sequences.length; i++) { + expect(sequences[i]).toBe(sequences[0] + i); + } + const batchFrames = rawFrames.filter((frame) => frame.type === 'chunk_batch'); + expect(batchFrames).toHaveLength(1); + expect(batchFrames[0].count).toBe(5); + + subscription.unsubscribe(); + transport.destroy(); + rawSubscriber.disconnect(); + }, + 15000, + ); + + testRedis( + 'a non-coalescable publication is a barrier that preserves emission order', + async () => { + const { RedisEventTransport, emitChunkWithReceipt } = await importFreshTransportModules(); + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient!, subscriber); + const streamId = `coalesce-barrier-${Date.now()}`; + + const received: Array<{ event: string; data: { label: string } }> = []; + const subscription = transport.subscribe(streamId, { + onChunk: (event) => received.push(event as (typeof received)[number]), + }); + await subscription.ready; + + const first = emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { label: 'delta-1' } }, + undefined, + { coalesce: true }, + ); + const second = emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { label: 'delta-2' } }, + undefined, + { coalesce: true }, + ); + const barrier = emitChunkWithReceipt(transport, streamId, { + event: 'on_run_step', + data: { label: 'barrier' }, + }); + const trailing = emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { label: 'delta-3' } }, + undefined, + { coalesce: true }, + ); + + const [firstSeq, secondSeq, barrierSeq, trailingSeq] = await Promise.all([ + first, + second, + barrier, + trailing, + ]); + await waitFor(() => received.length === 4); + + expect(received.map((event) => event.data.label)).toEqual([ + 'delta-1', + 'delta-2', + 'barrier', + 'delta-3', + ]); + expect([firstSeq, secondSeq, barrierSeq, trailingSeq]).toEqual([ + firstSeq, + (firstSeq as number) + 1, + (firstSeq as number) + 2, + (firstSeq as number) + 3, + ]); + + subscription.unsubscribe(); + transport.destroy(); + }, + 15000, + ); + + testRedis( + 'emitDone flushes the pending window ahead of the terminal frame', + async () => { + const { RedisEventTransport, emitChunkWithReceipt } = await importFreshTransportModules(); + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient!, subscriber); + const streamId = `coalesce-done-${Date.now()}`; + + const order: string[] = []; + const subscription = transport.subscribe(streamId, { + onChunk: (event) => order.push((event as { data: { label: string } }).data.label), + onDone: () => order.push('done'), + }); + await subscription.ready; + + const receipts = ['tail-1', 'tail-2', 'tail-3'].map((label) => + emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { label } }, + undefined, + { coalesce: true }, + ), + ); + await transport.emitDone(streamId, { final: true }); + await Promise.all(receipts); + + await waitFor(() => order.length === 4); + expect(order).toEqual(['tail-1', 'tail-2', 'tail-3', 'done']); + + subscription.unsubscribe(); + transport.destroy(); + }, + 15000, + ); + + testRedis( + 'a replaced generation fences the whole pending window', + async () => { + const { RedisEventTransport, emitChunkWithReceipt } = await importFreshTransportModules(); + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient!, subscriber); + const store = new RedisJobStore(ioredisClient!); + const streamId = `coalesce-fence-${Date.now()}`; + + const created = await store.createJob(streamId, 'user-1', streamId); + const staleGenerationId = created.createdAt - 1000; + + const receipts = await Promise.all( + Array.from({ length: 3 }, (_, i) => + emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { i } }, + staleGenerationId, + { coalesce: true }, + ), + ), + ); + + expect(receipts).toEqual([false, false, false]); + + transport.destroy(); + await store.destroy(); + }, + 15000, + ); + + testRedis( + 'a throwing subscriber callback loses one event, not the rest of the batch', + async () => { + const { RedisEventTransport, emitChunkWithReceipt } = await importFreshTransportModules(); + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient!, subscriber); + const streamId = `coalesce-throwing-${Date.now()}`; + + const received: number[] = []; + let done = false; + const subscription = transport.subscribe(streamId, { + onChunk: (event) => { + const value = (event as { data: { i: number } }).data.i; + if (value === 1) { + throw new Error('subscriber exploded on event 1'); + } + received.push(value); + }, + onDone: () => { + done = true; + }, + }); + await subscription.ready; + + const receipts = Array.from({ length: 4 }, (_, i) => + emitChunkWithReceipt( + transport, + streamId, + { event: 'on_message_delta', data: { i } }, + undefined, + { coalesce: true }, + ), + ); + await transport.emitDone(streamId, { final: true }); + await Promise.all(receipts); + + /** Event 1's sequence stalls the reorder cursor exactly like a lost + * individual frame; the force-flush recovers the remaining events and + * the terminal after REORDER_TIMEOUT_MS. What must NOT happen is the + * batch tail (2, 3) vanishing without ever entering the buffer. */ + await waitFor(() => done, 3000); + expect(received).toEqual([0, 2, 3]); + + subscription.unsubscribe(); + transport.destroy(); + }, + 15000, + ); + + testRedis( + 'abort with a pending window delivers the tail without error-closing subscribers', + async () => { + const manager = createRedisManager(); + const streamId = `coalesce-abort-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + + const received: string[] = []; + const errors: string[] = []; + const subscription = await manager.subscribe( + streamId, + (event) => { + const data = (event as { data?: { delta?: { content?: Array<{ text?: string }> } } }) + .data; + const text = data?.delta?.content?.[0]?.text; + if (typeof text === 'string') { + received.push(text); + } + }, + undefined, + (error) => errors.push(error), + ); + + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: { + id: 'step-1', + runId: 'run-1', + index: 0, + stepDetails: { type: 'message_creation' }, + }, + }); + for (const text of ['tail-1', 'tail-2']) { + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { id: 'step-1', delta: { content: [{ type: 'text', text }] } }, + }); + } + + /** Abort lands while both deltas are still inside the 25ms window. The + * abort path must drain the coalescers before its terminal CAS; without + * that, the window flush fences (-1) and the false receipts retire the + * runtime, error-closing this subscriber with a reconnect error. + * + * Delivery is observed WHILE the abort runs: the pre-CAS flush publishes + * the tail several round trips before abortJob's terminal cleanup tears + * down local subscription state, but under Redis Cluster the frames + * cross the cluster bus, so waiting until after abortJob returns races + * that teardown (publish receipts acknowledge execution, not delivery). */ + const abortPromise = manager.abortJob(streamId); + await waitFor(() => received.length === 2); + expect(received).toEqual(['tail-1', 'tail-2']); + const abortResult = await abortPromise; + expect(abortResult.success).toBe(true); + + await new Promise((resolve) => setTimeout(resolve, 150)); + expect(errors).toEqual([]); + + subscription?.unsubscribe(); + await manager.destroy(); + }, + 20000, + ); + + testRedis( + 'cross-replica abort does not error-close the owner subscribers over a warm window', + async () => { + const ownerManager = createRedisManager(); + const abortingManager = createRedisManager(); + const streamId = `coalesce-xreplica-abort-${Date.now()}`; + await ownerManager.createJob(streamId, 'user-1', streamId); + + const errors: string[] = []; + let done = false; + const subscription = await ownerManager.subscribe( + streamId, + () => undefined, + () => { + done = true; + }, + (error) => errors.push(error), + ); + + await ownerManager.emitChunk(streamId, { + event: 'on_run_step', + data: { + id: 'step-1', + runId: 'run-1', + index: 0, + stepDetails: { type: 'message_creation' }, + }, + }); + for (const text of ['warm-1', 'warm-2']) { + await ownerManager.emitChunk(streamId, { + event: 'on_message_delta', + data: { id: 'step-1', delta: { content: [{ type: 'text', text }] } }, + }); + } + + /** The abort CAS lands on another replica while the owner's window is + * still warm; the owner cannot flush pre-CAS, so its window flush fences + * against the aborted status. `beforePublish` runs between the CAS and + * the abort FINAL — forcing the owner's flush there pins the race the + * window timer only hits probabilistically. The fenced receipts land on + * a runtime whose abort signal already arrived, and the fence backstop + * must NOT error-close the subscribers awaiting the FINAL frame. */ + const abortResult = await abortingManager.abortJob(streamId, { + beforePublish: async () => { + await Promise.all([ + ( + ownerManager as unknown as { + jobStore: { flushPendingAppends?: (id: string) => Promise }; + } + ).jobStore.flushPendingAppends?.(streamId), + ( + ownerManager as unknown as { + eventTransport: { flushPendingChunks?: (id: string) => Promise }; + } + ).eventTransport.flushPendingChunks?.(streamId), + ]); + await new Promise((resolve) => setTimeout(resolve, 50)); + }, + }); + expect(abortResult.success).toBe(true); + + await waitFor(() => done); + expect(errors).toEqual([]); + + subscription?.unsubscribe(); + await ownerManager.destroy(); + await abortingManager.destroy(); + }, + 20000, + ); + + testRedis( + 'manager streams coalesced deltas end-to-end and completes cleanly', + async () => { + const manager = createRedisManager(); + const streamId = `coalesce-manager-${Date.now()}`; + + await manager.createJob(streamId, 'user-1', streamId); + const received: ServerSentEvent[] = []; + let done = false; + const subscription = await manager.subscribe( + streamId, + (event) => received.push(event), + () => { + done = true; + }, + ); + + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: { + id: 'step-1', + runId: 'run-1', + index: 0, + stepDetails: { type: 'message_creation' }, + }, + }); + for (const text of ['Hello', ' coalesced', ' world']) { + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { id: 'step-1', delta: { content: [{ type: 'text', text }] } }, + }); + } + await manager.emitChunk(streamId, { + event: 'on_run_step_completed', + data: { + id: 'step-1', + runId: 'run-1', + index: 0, + stepDetails: { type: 'message_creation' }, + result: { id: 'step-1' }, + }, + }); + + await waitFor(() => received.length === 5); + const labels = received.map((event) => + 'event' in event ? (event as { event: string }).event : 'created', + ); + expect(labels).toEqual([ + 'on_run_step', + 'on_message_delta', + 'on_message_delta', + 'on_message_delta', + 'on_run_step_completed', + ]); + + /** Same-replica durable read: the append-side window must flush so the + * resume snapshot reflects everything accepted before the read. */ + const resumeState = await manager.getResumeState(streamId); + const aggregated = JSON.stringify(resumeState?.aggregatedContent ?? ''); + expect(aggregated).toContain('Hello coalesced world'); + + await manager.emitDone(streamId, { final: true, streamId } as unknown as ServerSentEvent); + await manager.completeJob(streamId); + await waitFor(() => done); + + subscription?.unsubscribe(); + await manager.destroy(); + }, + 20000, + ); +}); diff --git a/packages/api/src/stream/implementations/RedisEventTransport.ts b/packages/api/src/stream/implementations/RedisEventTransport.ts index fdb881c8d7..1e97b8b52e 100644 --- a/packages/api/src/stream/implementations/RedisEventTransport.ts +++ b/packages/api/src/stream/implementations/RedisEventTransport.ts @@ -2,6 +2,12 @@ import { randomUUID } from 'crypto'; import { logger } from '@librechat/data-schemas'; import type { Redis, Cluster } from 'ioredis'; import type { IEventTransport, PreemptMessage } from '~/stream/interfaces/IJobStore'; +import type { ChunkPublicationOptions } from '~/stream/internal/chunkPublication'; +import { + MAX_COALESCED_BYTES, + MAX_COALESCED_EVENTS, + resolveCoalesceWindowMs, +} from '~/stream/internal/coalescing'; import { registerChunkPublicationCapability } from '~/stream/internal/chunkPublication'; import { instrumentIORedisClient, RedisUseCases } from '~/cache/redisTelemetry'; @@ -33,6 +39,7 @@ const KEYS = { */ const EventTypes = { CHUNK: 'chunk', + CHUNK_BATCH: 'chunk_batch', DONE: 'done', ERROR: 'error', ABORT: 'abort', @@ -46,6 +53,10 @@ interface PubSubMessage { seq?: number; data?: unknown; error?: string; + /** First sequence of a CHUNK_BATCH frame; events[i] owns baseSeq + i. */ + baseSeq?: number; + /** Coalesced chunk payloads of a CHUNK_BATCH frame, in emission order. */ + events?: unknown[]; /** Immutable identity of the generation that emitted the event. */ generationId?: number; /** Opaque nonce linking a replacement abort to its owner acknowledgement. */ @@ -54,6 +65,19 @@ interface PubSubMessage { preempt?: PreemptMessage; } +/** + * Producer-side buffer of coalescable chunk publications for one stream. + * Payloads are pre-serialized at enqueue so a flush only joins strings, and + * each resolver settles its caller's receipt with `baseSeq + index`. + */ +interface PendingChunkBatch { + generationId?: number; + events: string[]; + resolvers: Array<(receipt: number | false | undefined) => void>; + bytes: number; + timer: ReturnType | null; +} + /** * Reorder buffer state for a stream subscription. * Handles out-of-order message delivery in Redis Cluster mode. @@ -109,9 +133,13 @@ interface PreemptRegistration { * expectCreatedAt | "", * allowRetainedEpoch ("0" | "1"), * generationEpochGraceTtl, - * requireActiveJob ("0" | "1") + * requireActiveJob ("0" | "1"), + * sequenceCount * ] - * RETURNS: the 0-indexed seq assigned to this event, or -1 when the generation guard fails + * RETURNS: the 0-indexed first seq assigned to this frame, or -1 when the generation + * guard fails. A single-event frame passes count 1 and splices the seq; a coalesced + * frame passes its event count, reserves that many consecutive sequences in one INCRBY, + * and splices the base — event i in the frame owns base + i. * * During a rolling deployment, a job created by the previous version can expire without * leaving a generation marker. A tagged terminal event may claim that absent marker only @@ -136,7 +164,8 @@ const PUBLISH_SEQ_LUA = 'local currentStatus = redis.call("HGET", KEYS[2], "status") ' + 'if currentStatus ~= "running" and currentStatus ~= "requires_action" then return -1 end ' + 'end ' + - 'local val = redis.call("INCR", KEYS[1]) ' + + 'local count = tonumber(ARGV[9]) ' + + 'local val = redis.call("INCRBY", KEYS[1], count) ' + 'local ttl = tonumber(ARGV[4]) ' + 'local seqTtl = redis.call("TTL", KEYS[1]) ' + 'if seqTtl < math.floor(ttl / 2) then ' + @@ -144,7 +173,7 @@ const PUBLISH_SEQ_LUA = 'if jobTtl > ttl then ttl = jobTtl end ' + 'redis.call("EXPIRE", KEYS[1], ttl) ' + 'end ' + - 'local seq = val - 1 ' + + 'local seq = val - count ' + 'redis.call("PUBLISH", ARGV[1], ARGV[2] .. string.format("%d", seq) .. ARGV[3]) ' + 'return seq'; @@ -237,6 +266,10 @@ export class RedisEventTransport implements IEventTransport { private channelSubscriptions = new Map>(); /** Counter for generating unique subscriber IDs */ private subscriberIdCounter = 0; + /** Coalescable chunk publications awaiting their window flush, per stream */ + private pendingBatches = new Map(); + /** Delta-coalescing window; 0 keeps every publication on the per-event path */ + private readonly coalesceWindowMs: number; private createStreamState(): StreamSubscribers { return { @@ -273,8 +306,9 @@ export class RedisEventTransport implements IEventTransport { constructor(publisher: Redis | Cluster, subscriber: Redis | Cluster) { this.publisher = instrumentIORedisClient(publisher, RedisUseCases.GENERATION_STREAM); this.subscriber = instrumentIORedisClient(subscriber, RedisUseCases.GENERATION_STREAM); - registerChunkPublicationCapability(this, (streamId, event, generationId) => - this.publishChunkWithReceipt(streamId, event, generationId), + this.coalesceWindowMs = resolveCoalesceWindowMs(); + registerChunkPublicationCapability(this, (streamId, event, generationId, publishOptions) => + this.publishChunkWithReceipt(streamId, event, generationId, publishOptions), ); // Set up message handler for all subscriptions @@ -315,6 +349,26 @@ export class RedisEventTransport implements IEventTransport { requireActiveJob = false, ): Promise { const [prefix, suffix] = RedisEventTransport.buildPayloadParts(message); + return this.evalPublishSequenced( + streamId, + prefix, + suffix, + 1, + expectedGenerationId, + allowRetainedEpoch, + requireActiveJob, + ); + } + + private async evalPublishSequenced( + streamId: string, + prefix: string, + suffix: string, + count: number, + expectedGenerationId?: number, + allowRetainedEpoch = false, + requireActiveJob = false, + ): Promise { const seq = await this.publisher.eval( PUBLISH_SEQ_LUA, 3, @@ -329,6 +383,7 @@ export class RedisEventTransport implements IEventTransport { allowRetainedEpoch ? '1' : '0', String(GENERATION_EPOCH_GRACE_TTL_SECONDS), requireActiveJob ? '1' : '0', + String(count), ); return seq as number; } @@ -337,7 +392,18 @@ export class RedisEventTransport implements IEventTransport { streamId: string, event: unknown, generationId?: number, + options?: ChunkPublicationOptions, ): Promise { + if (options?.coalesce === true && this.coalesceWindowMs > 0) { + return this.enqueueCoalescedChunk(streamId, event, generationId); + } + /** A sequenced non-coalescable publication is an ordering barrier: pending + * deltas must be issued first so their reserved sequences stay below this + * frame's. Both EVALs ride the same connection (and, under Cluster, the + * same hash slot), so issue order alone preserves sequence order. */ + if (this.pendingBatches.has(streamId)) { + void this.flushCoalescedChunks(streamId); + } return this.publishWithSequence( streamId, { @@ -359,6 +425,119 @@ export class RedisEventTransport implements IEventTransport { }); } + /** + * Buffer a coalescable chunk for the current window and settle its receipt when the + * batch flushes. The receipt keeps per-event semantics: its own absolute sequence, + * `false` under a generation/status fence, `undefined` on operational failure. + */ + private enqueueCoalescedChunk( + streamId: string, + event: unknown, + generationId?: number, + ): Promise { + let pending = this.pendingBatches.get(streamId); + if (pending && pending.generationId !== generationId) { + void this.flushCoalescedChunks(streamId); + pending = undefined; + } + if (!pending) { + pending = { generationId, events: [], resolvers: [], bytes: 0, timer: null }; + this.pendingBatches.set(streamId, pending); + } + + const batch = pending; + const encoded = JSON.stringify(event); + batch.events.push(encoded); + batch.bytes += encoded.length; + const receipt = new Promise((resolve) => { + batch.resolvers.push(resolve); + }); + + if (batch.events.length >= MAX_COALESCED_EVENTS || batch.bytes >= MAX_COALESCED_BYTES) { + void this.flushCoalescedChunks(streamId); + } else if (batch.timer == null) { + batch.timer = setTimeout(() => { + void this.flushCoalescedChunks(streamId); + }, this.coalesceWindowMs); + } + return receipt; + } + + /** + * Publish the stream's pending coalesced chunks as one CHUNK_BATCH frame. + * + * One INCRBY reserves a consecutive sequence per buffered event, so subscribers + * unpack the frame into individually sequenced chunks and the reorder buffer is + * none the wiser. The events array is spliced from pre-serialized payloads for + * the same reason single frames are: no server-side re-encoding. + */ + private flushCoalescedChunks(streamId: string): Promise { + const pending = this.pendingBatches.get(streamId); + if (!pending) { + return Promise.resolve(); + } + this.pendingBatches.delete(streamId); + if (pending.timer != null) { + clearTimeout(pending.timer); + pending.timer = null; + } + + const { generationId, events, resolvers } = pending; + const prefix = `{"type":${JSON.stringify(EventTypes.CHUNK_BATCH)},"baseSeq":`; + const suffix = + (generationId != null ? `,"generationId":${generationId}` : '') + + `,"events":[${events.join(',')}]}`; + + return this.evalPublishSequenced( + streamId, + prefix, + suffix, + events.length, + generationId, + false, + generationId != null, + ).then( + (baseSeq) => { + if (baseSeq === -1) { + for (const resolve of resolvers) { + resolve(false); + } + return; + } + for (let i = 0; i < resolvers.length; i++) { + resolvers[i](baseSeq + i); + } + }, + (err) => { + logger.error(`[RedisEventTransport] Failed to publish chunk batch:`, err); + for (const resolve of resolvers) { + resolve(undefined); + } + }, + ); + } + + /** Publish a stream's pending coalesced chunks now (pre-transition barrier). */ + async flushPendingChunks(streamId: string): Promise { + await this.flushCoalescedChunks(streamId); + } + + /** Drop a stream's pending coalesced chunks without publishing (teardown path). */ + private discardCoalescedChunks(streamId: string): void { + const pending = this.pendingBatches.get(streamId); + if (!pending) { + return; + } + this.pendingBatches.delete(streamId); + if (pending.timer != null) { + clearTimeout(pending.timer); + pending.timer = null; + } + for (const resolve of pending.resolvers) { + resolve(undefined); + } + } + private ensureChannelSubscription(channel: string): Promise { const existing = this.channelSubscriptions.get(channel); if (existing) { @@ -518,6 +697,30 @@ export class RedisEventTransport implements IEventTransport { } if (parsed.type === EventTypes.CHUNK && parsed.seq != null) { this.handleOrderedChunk(streamId, streamState, parsed); + } else if ( + parsed.type === EventTypes.CHUNK_BATCH && + parsed.baseSeq != null && + Array.isArray(parsed.events) + ) { + /** Unpack at ingress: each coalesced payload owns baseSeq + i, so the + * reorder buffer sees the exact per-event sequences it would have seen + * from individual frames (dup drop, gap buffering, force-flush). Each + * event is isolated: a throwing subscriber callback must degrade like + * a lost individual frame (that sequence stalls until the reorder + * force-flush) instead of discarding the rest of the batch, whose + * sequences are already reserved and would otherwise never arrive. */ + for (let i = 0; i < parsed.events.length; i++) { + try { + this.handleOrderedChunk(streamId, streamState, { + type: EventTypes.CHUNK, + seq: parsed.baseSeq + i, + data: parsed.events[i], + ...(parsed.generationId != null && { generationId: parsed.generationId }), + }); + } catch (err) { + logger.error(`[RedisEventTransport] Failed to deliver coalesced chunk:`, err); + } + } } else if ( (parsed.type === EventTypes.DONE || parsed.type === EventTypes.ERROR) && parsed.seq != null @@ -881,6 +1084,9 @@ export class RedisEventTransport implements IEventTransport { */ async emitDone(streamId: string, event: unknown, generationId?: number): Promise { try { + /** Terminal frames must carry a later sequence than every pending delta, + * or subscribers would close on DONE and drop the coalesced tail. */ + await this.flushCoalescedChunks(streamId); const sequence = await this.publishWithSequence( streamId, { @@ -906,6 +1112,7 @@ export class RedisEventTransport implements IEventTransport { replacedGenerationId: number, creationAttemptId: string, ): Promise { + await this.flushCoalescedChunks(streamId); const [prefix, suffix] = RedisEventTransport.buildPayloadParts({ type: EventTypes.DONE, data: event, @@ -934,6 +1141,7 @@ export class RedisEventTransport implements IEventTransport { */ async emitError(streamId: string, error: string, generationId?: number): Promise { try { + await this.flushCoalescedChunks(streamId); const sequence = await this.publishWithSequence( streamId, { @@ -1233,6 +1441,11 @@ export class RedisEventTransport implements IEventTransport { const channel = CHANNELS.events(streamId); const state = this.streams.get(streamId); + /** Terminal publications flushed ahead of themselves; anything still pending + * here belongs to a torn-down generation and stays recoverable from the + * durable chunk log, matching a dropped per-event publication. */ + this.discardCoalescedChunks(streamId); + if (state) { state.handlers.clear(); state.allSubscribersLeftCallback = undefined; @@ -1261,6 +1474,10 @@ export class RedisEventTransport implements IEventTransport { * Destroy all resources. */ destroy(): void { + for (const streamId of this.pendingBatches.keys()) { + this.discardCoalescedChunks(streamId); + } + // Clear all flush timeouts and buffered messages. // Sequence keys are NOT deleted here — they are shared across replicas. // A shutting-down replica must not nuke the counter for active publishers. diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index ae252af620..bdfc0d3d8d 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -33,6 +33,11 @@ import { PAUSE_PERSISTENCE_TIMEOUT_MS, isPendingActionStale, } from '~/stream/interfaces/IJobStore'; +import { + MAX_COALESCED_BYTES, + MAX_COALESCED_EVENTS, + resolveCoalesceWindowMs, +} from '~/stream/internal/coalescing'; import { instrumentIORedisClient, RedisUseCases } from '~/cache/redisTelemetry'; import { RecoveredSteerPayloadMismatchError } from '~/stream/SteerRecovery'; @@ -738,6 +743,50 @@ const CHUNK_APPEND_LUA = 'redis.call("HSET", KEYS[3], ARGV[4], cjson.encode(receipt)) end ' + 'return 1'; +/** + * Batched CHUNK_APPEND_LUA for plain streaming deltas: identical generation/status/epoch + * guards and extend-only TTL housekeeping, evaluated once per batch, with one XADD per + * event. Steer-delivery settlement is deliberately absent — an append carrying a steer + * receipt is a barrier and stays on the per-event script. + * + * KEYS: [chunks, job, steerReceipts, steerReceiptOrder, claimedSteers, steers, + * parkedSteers, generationEpoch] + * ARGV: [runningTtl, expectCreatedAt | "", nowMs, parkedSteersTtl, + * generationEpochGraceTtl, eventJson...] + */ +const CHUNK_APPEND_BATCH_LUA = + 'local currentCreatedAt = redis.call("HGET", KEYS[2], "createdAt") ' + + 'if not currentCreatedAt then return 0 end ' + + 'if ARGV[2] ~= "" and currentCreatedAt ~= ARGV[2] then return 0 end ' + + 'local currentStatus = redis.call("HGET", KEYS[2], "status") ' + + 'if currentStatus ~= "running" and currentStatus ~= "requires_action" then return 0 end ' + + 'local retainedEpoch = redis.call("GET", KEYS[8]) ' + + 'if retainedEpoch and retainedEpoch ~= currentCreatedAt then return 0 end ' + + 'local run = tonumber(ARGV[1]) ' + + 'local target = run ' + + 'local jobTtl = redis.call("TTL", KEYS[2]) ' + + 'if jobTtl < target then redis.call("EXPIRE", KEYS[2], target) ' + + 'elseif jobTtl > target then target = jobTtl end ' + + 'local recoveryTarget = target ' + + 'if redis.call("HGET", KEYS[2], "recoveredSteerId") then ' + + 'recoveryTarget = target + tonumber(ARGV[4]) ' + + 'local pt = redis.call("TTL", KEYS[7]) ' + + 'if pt >= 0 and pt < recoveryTarget then redis.call("EXPIRE", KEYS[7], recoveryTarget) end end ' + + 'local epochTarget = target + tonumber(ARGV[5]) ' + + 'if retainedEpoch then local epochTtl = redis.call("TTL", KEYS[8]) ' + + 'if epochTtl >= 0 and epochTtl < epochTarget then redis.call("EXPIRE", KEYS[8], epochTarget) end ' + + 'else redis.call("SET", KEYS[8], currentCreatedAt, "EX", epochTarget) end ' + + 'for i = 6, #ARGV do redis.call("XADD", KEYS[1], "*", "event", ARGV[i]) end ' + + 'if currentStatus == "running" then ' + + 'redis.call("HSET", KEYS[2], "lastActiveAt", ARGV[3]) end ' + + 'local cur = redis.call("TTL", KEYS[1]) ' + + 'if cur < target then redis.call("EXPIRE", KEYS[1], target) end ' + + 'for i = 3, 4 do local rt = redis.call("TTL", KEYS[i]) ' + + 'if rt >= 0 and rt < recoveryTarget then redis.call("EXPIRE", KEYS[i], recoveryTarget) end end ' + + 'for i = 5, 6 do local qt = redis.call("TTL", KEYS[i]) ' + + 'if qt >= 0 and qt < target then redis.call("EXPIRE", KEYS[i], target) end end ' + + 'return 1'; + /** * Persist the run-step timeline with the same paused-window TTL as the chunk stream. * `saveRunSteps` SETs (overwrites) the whole array, so unlike the chunk append there's no @@ -1502,10 +1551,27 @@ interface LocalCacheEntry { value: T; } +/** + * Coalescable durable appends buffered for one stream. Events are + * pre-serialized at enqueue; a flush XADDs them in order under one guard pass. + * The whole batch shares one fate, so every resolver settles identically. + */ +interface PendingChunkAppendBatch { + expectedCreatedAt?: number; + events: string[]; + settlers: Array<{ resolve: (appended: boolean) => void; reject: (err: unknown) => void }>; + bytes: number; + timer: ReturnType | null; +} + export class RedisJobStore implements IJobStoreV2 { private redis: Redis | Cluster; private cleanupInterval: NodeJS.Timeout | null = null; private ttl: typeof DEFAULT_TTL; + /** Coalescable chunk appends awaiting their window flush, per stream */ + private pendingAppends = new Map(); + /** Durable-append coalescing window; 0 keeps every append on the per-event path */ + private readonly coalesceWindowMs: number; /** Whether Redis client is in cluster mode (affects pipeline usage) */ private isCluster: boolean; @@ -1536,6 +1602,7 @@ export class RedisJobStore implements IJobStoreV2 { constructor(redis: Redis | Cluster, options?: RedisJobStoreOptions) { this.redis = instrumentIORedisClient(redis, RedisUseCases.GENERATION_STREAM); + this.coalesceWindowMs = resolveCoalesceWindowMs(); this.ttl = { completed: options?.completedTtl ?? DEFAULT_TTL.completed, running: options?.runningTtl ?? DEFAULT_TTL.running, @@ -2781,6 +2848,18 @@ export class RedisJobStore implements IJobStoreV2 { clearInterval(this.cleanupInterval); this.cleanupInterval = null; } + /** Shutdown terminals flushed per stream already; whatever remains did not + * commit, and resolving false lets the owning fence continuations settle. */ + for (const [streamId, pending] of this.pendingAppends) { + this.pendingAppends.delete(streamId); + if (pending.timer != null) { + clearTimeout(pending.timer); + pending.timer = null; + } + for (const settler of pending.settlers) { + settler.resolve(false); + } + } // Clear local caches this.localGraphCache.clear(); this.localContentParts.clear(); @@ -3528,7 +3607,17 @@ export class RedisJobStore implements IJobStoreV2 { event: unknown, expectedCreatedAt?: number, deliveredSteer?: SteerQueueItem, + options?: { coalesce?: boolean }, ): Promise { + if (options?.coalesce === true && deliveredSteer == null && this.coalesceWindowMs > 0) { + return this.enqueueCoalescedAppend(streamId, event, expectedCreatedAt); + } + /** The chunk log is replayed in XADD order, so a per-event append (durable + * control events, steer receipts) is a barrier: pending coalesced deltas + * must be issued first. Same connection, so issue order is land order. */ + if (this.pendingAppends.has(streamId)) { + void this.flushCoalescedAppends(streamId); + } const key = KEYS.chunks(streamId); const jobKey = KEYS.job(streamId); // XADD + derive-and-extend-only EXPIRE in a single atomic eval. Refreshing the TTL on @@ -3563,10 +3652,105 @@ export class RedisJobStore implements IJobStoreV2 { return appended === 1; } + /** + * Buffer a coalescable durable append for the current window. The whole batch + * settles together: `true` on commit, `false` under the generation/status + * fence, and a rejection on operational failure — mirroring the per-event + * appendChunk contract each caller's fence continuation already handles. + */ + private enqueueCoalescedAppend( + streamId: string, + event: unknown, + expectedCreatedAt?: number, + ): Promise { + let pending = this.pendingAppends.get(streamId); + if (pending && pending.expectedCreatedAt !== expectedCreatedAt) { + void this.flushCoalescedAppends(streamId); + pending = undefined; + } + if (!pending) { + pending = { expectedCreatedAt, events: [], settlers: [], bytes: 0, timer: null }; + this.pendingAppends.set(streamId, pending); + } + + const batch = pending; + const encoded = JSON.stringify(event); + batch.events.push(encoded); + batch.bytes += encoded.length; + const settled = new Promise((resolve, reject) => { + batch.settlers.push({ resolve, reject }); + }); + + if (batch.events.length >= MAX_COALESCED_EVENTS || batch.bytes >= MAX_COALESCED_BYTES) { + void this.flushCoalescedAppends(streamId); + } else if (batch.timer == null) { + batch.timer = setTimeout(() => { + void this.flushCoalescedAppends(streamId); + }, this.coalesceWindowMs); + } + return settled; + } + + private flushCoalescedAppends(streamId: string): Promise { + const pending = this.pendingAppends.get(streamId); + if (!pending) { + return Promise.resolve(); + } + this.pendingAppends.delete(streamId); + if (pending.timer != null) { + clearTimeout(pending.timer); + pending.timer = null; + } + + const { expectedCreatedAt, events, settlers } = pending; + return this.redis + .eval( + CHUNK_APPEND_BATCH_LUA, + 8, + KEYS.chunks(streamId), + KEYS.job(streamId), + KEYS.steerReceipts(streamId), + KEYS.steerReceiptOrder(streamId), + KEYS.claimedSteers(streamId), + KEYS.steers(streamId), + KEYS.parkedSteers(streamId), + KEYS.generationEpoch(streamId), + String(this.runningStorageTtlSeconds()), + expectedCreatedAt != null ? String(expectedCreatedAt) : '', + String(Date.now()), + String(this.parkedRecoveryTtlSeconds()), + String(GENERATION_EPOCH_GRACE_TTL_S), + ...events, + ) + .then( + (appended) => { + const committed = appended === 1; + for (const settler of settlers) { + settler.resolve(committed); + } + }, + (err) => { + for (const settler of settlers) { + settler.reject(err); + } + }, + ); + } + + /** Persist a stream's pending coalesced appends now (pre-transition barrier). */ + async flushPendingAppends(streamId: string): Promise { + await this.flushCoalescedAppends(streamId); + } + /** * Get all chunks from Redis Stream. */ private async getChunks(streamId: string, expectedCreatedAt?: number): Promise { + /** A same-replica snapshot read must observe the appends this process has + * already accepted, or a resume during an active window reconstructs + * without the buffered tail. Cross-replica readers keep today's contract: + * the log may trail live emission by up to one window. */ + await this.flushCoalescedAppends(streamId); const rawEntries = expectedCreatedAt == null ? await this.redis.xrange(KEYS.chunks(streamId), '-', '+') diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index ac6452b318..9191d8b6e2 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -661,6 +661,10 @@ export interface IJobStore { expectedCreatedAt?: number, ): Promise; + /** Optional batching capability: persist any coalesced appends buffered for + * this stream now. Stores without append coalescing simply omit it. */ + flushPendingAppends?(streamId: string): Promise; + clearContentState(streamId: string, expectedCreatedAt?: number): void; saveRunSteps?( streamId: string, @@ -870,6 +874,16 @@ export interface IJobStoreV2 extends IJobStore { */ recordActivity?(streamId: string, expectedCreatedAt?: number): void; + /** + * Persist any coalesced chunk appends still buffered for this stream. + * Terminal transitions must flush first so the batch lands under the + * generation's live status instead of fencing against its own completion. + * + * Presence of this method is how a store advertises append-coalescing + * support; see the transport's flushPendingChunks for the pairing rule. + */ + flushPendingAppends?(streamId: string): Promise; + /** Get total job count */ getJobCount(): Promise; @@ -966,6 +980,10 @@ export interface IJobStoreV2 extends IJobStore { /** When present, the same durable write records this drained steer as * delivered. Redis performs both mutations in one same-slot Lua step. */ deliveredSteer?: SteerQueueItem, + /** Hot-path hint: `coalesce` marks a plain streaming delta whose durable + * append may batch with its window peers. Stores without batching (and any + * append carrying a steer receipt) ignore it and stay per-event. */ + options?: { coalesce?: boolean }, ): Promise; /** @@ -1310,6 +1328,18 @@ export interface IEventTransport { */ closeLocalSubscribers?(streamId: string, error: string): void; + /** + * Publish any coalesced chunk publications still buffered for this stream. + * Callers about to transition a generation's status must flush first, or the + * batch would land behind the transition and fence itself. + * + * Presence of this method is how a transport advertises delta-coalescing + * support: the generation manager only sends `coalesce` hints (and only + * stops awaiting per-delta receipts) when both the transport and the job + * store expose their flush capability. + */ + flushPendingChunks?(streamId: string): Promise; + /** Cleanup transport resources for a specific stream */ cleanup(streamId: string): void; diff --git a/packages/api/src/stream/internal/chunkPublication.ts b/packages/api/src/stream/internal/chunkPublication.ts index 8fcced3fe0..2f97256d9e 100644 --- a/packages/api/src/stream/internal/chunkPublication.ts +++ b/packages/api/src/stream/internal/chunkPublication.ts @@ -8,10 +8,22 @@ import type { IEventTransport } from '../interfaces/IJobStore'; */ export type ChunkPublicationReceipt = number | false | void; +/** + * Hot-path hints for transports that can batch sequenced publications. + * `coalesce` marks an event as eligible for windowed batching: the receipt + * promise then resolves when the batch flushes instead of per event. Only + * order-insensitive-to-latency streaming deltas should opt in; fenced + * emissions (durable, steer, created, terminal) must stay per-event. + */ +export interface ChunkPublicationOptions { + coalesce?: boolean; +} + type ChunkPublicationCapability = ( streamId: string, event: unknown, generationId?: number, + options?: ChunkPublicationOptions, ) => Promise; /** @@ -36,10 +48,15 @@ export function emitChunkWithReceipt( streamId: string, event: unknown, generationId?: number, + options?: ChunkPublicationOptions, ): Promise { const capability = chunkPublicationCapabilities.get(transport); if (capability) { - return capability(streamId, event, generationId); + /** Preserve the 3-arg call shape when no hint is given: registered + * capabilities predate the options parameter and spies assert on it. */ + return options === undefined + ? capability(streamId, event, generationId) + : capability(streamId, event, generationId, options); } return Promise.resolve(transport.emitChunk(streamId, event, generationId)); } diff --git a/packages/api/src/stream/internal/coalescing.ts b/packages/api/src/stream/internal/coalescing.ts new file mode 100644 index 0000000000..5440566688 --- /dev/null +++ b/packages/api/src/stream/internal/coalescing.ts @@ -0,0 +1,25 @@ +/** Upper clamp for the delta-coalescing window; larger values only add UI staleness. */ +const MAX_COALESCE_WINDOW_MS = 1000; +/** Coalesced-batch safety caps: a full buffer flushes immediately, ahead of the window. */ +export const MAX_COALESCED_EVENTS = 64; +export const MAX_COALESCED_BYTES = 128 * 1024; + +/** + * Streaming-delta coalescing window shared by the Redis transport (publish + * frames), the Redis job store (durable appends), and the generation manager + * (hinting/backpressure). STREAM_DELTA_COALESCE_MS is deliberately the ONLY + * source: all three read the same value, so one process cannot mix a hinting + * manager with non-batching services or batching services with a silent + * manager. Both sides MUST also buffer on the same window: the subscriber + * resume frontier assumes an event is never visible in the durable chunk log + * meaningfully earlier than its sequence lands on the shared counter, so + * batching one side without the other reopens that race for the full window + * instead of a same-tick skew. 0 (the default) disables coalescing. + */ +export function resolveCoalesceWindowMs(): number { + const raw = Number(process.env.STREAM_DELTA_COALESCE_MS ?? 0); + if (!Number.isFinite(raw) || raw <= 0) { + return 0; + } + return Math.min(Math.floor(raw), MAX_COALESCE_WINDOW_MS); +}