From e733dffe10eb41eecc2be2e7c901ef01f4998203 Mon Sep 17 00:00:00 2001 From: 3ndetz Date: Wed, 28 Oct 2026 13:00:43 +0300 Subject: [PATCH] Prevent duplicate replay fanout and bound stalled SSE writes (#2223) * Prevent replay fanout and bound stalled SSE writes * Handle streams aborted during replay * Preserve repeated consumption acknowledgements * Bound SSE queues by size while allowing short bursts --- .../handlers/cli/sessionHandlers.test.ts | 22 +++++++++++ .../socket/handlers/cli/sessionHandlers.ts | 4 ++ hub/src/sse/sseManager.test.ts | 20 ++++++++++ hub/src/sse/sseManager.ts | 37 ++++++++++++++++++- hub/src/web/routes/events.ts | 2 + 5 files changed, 83 insertions(+), 2 deletions(-) diff --git a/hub/src/socket/handlers/cli/sessionHandlers.test.ts b/hub/src/socket/handlers/cli/sessionHandlers.test.ts index 62c5d35e..d22e2816 100644 --- a/hub/src/socket/handlers/cli/sessionHandlers.test.ts +++ b/hub/src/socket/handlers/cli/sessionHandlers.test.ts @@ -27,6 +27,28 @@ class FakeSocket { } } +it('does not rebroadcast stored messages and retains repeated consumption acknowledgements', () => { + const store = new Store(':memory:') + const session = store.sessions.getOrCreateSession('dedupe-replay', {}, null, 'default') + const socket = new FakeSocket() + const events: SyncEvent[] = [] + const progress = mock(() => {}) + registerSessionHandlers(socket as unknown as CliSocketWithData, { + store, resolveSessionAccess: () => ({ ok: true, value: session }), emitAccessError() {}, + onWebappEvent: event => events.push(event), onAgentProgress: progress + }) + const message = { sid: session.id, localId: 'stable-1', message: { role: 'user', content: { type: 'text', text: 'hello' } } } + socket.trigger('message', message) + socket.trigger('messages-consumed', { sid: session.id, localIds: ['stable-1'] }) + socket.trigger('message', message) + socket.trigger('messages-consumed', { sid: session.id, localIds: ['stable-1'] }) + expect(events.filter(event => event.type === 'message-received')).toHaveLength(1) + expect(events.filter(event => event.type === 'messages-consumed')).toHaveLength(2) + expect(progress).toHaveBeenCalledTimes(1) + expect(store.messages.getAllMessages(session.id)).toHaveLength(1) + store.close() +}) + function redundantGoalStatusContent(message: string): unknown { return { role: 'agent', diff --git a/hub/src/socket/handlers/cli/sessionHandlers.ts b/hub/src/socket/handlers/cli/sessionHandlers.ts index 548c5675..680272dd 100644 --- a/hub/src/socket/handlers/cli/sessionHandlers.ts +++ b/hub/src/socket/handlers/cli/sessionHandlers.ts @@ -172,7 +172,11 @@ export function registerSessionHandlers(socket: CliSocketWithData, deps: Session return } + const wasStored = Boolean(localId && store.messages.lookupQueuedMessage(sid, localId).status !== 'absent') const msg = store.messages.addMessage(sid, content, localId, undefined, createdAt) + // The durable row already exists; replay must not be mistaken for + // fresh agent progress or trigger another round of UI updates. + if (wasStored) return // A reasoning stream arrives as a series of growing snapshots under one // stable id, so a stream should cost one row rather than one per diff --git a/hub/src/sse/sseManager.test.ts b/hub/src/sse/sseManager.test.ts index affad20a..8f29ad3a 100644 --- a/hub/src/sse/sseManager.test.ts +++ b/hub/src/sse/sseManager.test.ts @@ -1,6 +1,26 @@ import { describe, expect, it } from 'bun:test' import { SSEManager } from './sseManager' import type { SyncEvent } from '../sync/syncEngine' + +it('closes a stalled subscriber once pending writes reach the bound', () => { + const manager = new SSEManager(0, new VisibilityTracker()) + let closed = 0 + manager.subscribe({ id: 'slow', namespace: 'default', all: true, + send: () => new Promise(() => {}), sendHeartbeat: () => {}, close: () => { closed++ } }) + for (let i = 0; i < 4200; i++) manager.broadcast({ type: 'session-updated', sessionId: 's', namespace: 'default' }) + expect(closed).toBe(1) + expect(manager.hasSubscription('slow')).toBe(false) +}) + +it('also bounds stalled subscriber writes by bytes', () => { + const manager = new SSEManager(0, new VisibilityTracker()) + let closed = 0 + manager.subscribe({ id: 'slow-bytes', namespace: 'default', all: true, + send: () => new Promise(() => {}), sendHeartbeat: () => {}, close: () => { closed++ } }) + for (let i = 0; i < 10; i++) manager.broadcast({ type: 'session-updated', sessionId: 'x'.repeat(512 * 1024), namespace: 'default' }) + expect(closed).toBe(1) + expect(manager.hasSubscription('slow-bytes')).toBe(false) +}) import { VisibilityTracker } from '../visibility/visibilityTracker' describe('SSEManager namespace filtering', () => { diff --git a/hub/src/sse/sseManager.ts b/hub/src/sse/sseManager.ts index cabd4f7a..ae6b6060 100644 --- a/hub/src/sse/sseManager.ts +++ b/hub/src/sse/sseManager.ts @@ -14,6 +14,10 @@ export type SSESubscription = { type SSEConnection = SSESubscription & { send: (event: SyncEvent, eventId?: string) => void | Promise sendHeartbeat: () => void | Promise + close?: () => void + inflight: number + inflightBytes: number + pendingBytes: number /** * Non-null while a resumed connection is still writing its replay: live * broadcasts land here instead of on the wire so replayed events keep @@ -73,6 +77,7 @@ export class SSEManager { resumeFrom?: string | null send: (event: SyncEvent, eventId?: string) => void | Promise sendHeartbeat: () => void | Promise + close?: () => void }): SSESubscription & SSEResumeResult { const subscription: SSEConnection = { id: options.id, @@ -82,6 +87,10 @@ export class SSEManager { machineId: options.machineId ?? null, send: options.send, sendHeartbeat: options.sendHeartbeat, + close: options.close, + inflight: 0, + inflightBytes: 0, + pendingBytes: 0, pending: null } @@ -122,6 +131,7 @@ export class SSEManager { } while (connection.pending) { const batch = connection.pending.splice(0) + connection.pendingBytes = 0 if (batch.length === 0) { // No await between this check and the assignment, so no // broadcast can slip into the queue and be dropped. @@ -250,6 +260,7 @@ export class SSEManager { broadcast(event: SyncEvent): void { const seq = this.recordEvent(event) + const eventBytes = JSON.stringify(event).length * 2 for (const connection of this.connections.values()) { if (!this.shouldSend(connection, event)) { continue @@ -259,13 +270,35 @@ export class SSEManager { // the receiving namespace so the cursor stays bound to it. const eventId = this.eventIdFor(seq, connection.namespace) if (connection.pending) { + if (connection.pending.length >= 4096 || connection.pendingBytes + eventBytes > 8 * 1024 * 1024) { + this.unsubscribe(connection.id) + connection.close?.() + continue + } connection.pending.push({ event, eventId }) + connection.pendingBytes += eventBytes continue } - void Promise.resolve(connection.send(event, eventId)).catch(() => { + if (connection.inflight >= 4096 || connection.inflightBytes + eventBytes > 8 * 1024 * 1024) { this.unsubscribe(connection.id) - }) + connection.close?.() + continue + } + try { + const result = connection.send(event, eventId) + if (result) { + connection.inflight++ + connection.inflightBytes += eventBytes + void Promise.resolve(result).catch(() => { + this.unsubscribe(connection.id) + connection.close?.() + }).finally(() => { connection.inflight--; connection.inflightBytes -= eventBytes }) + } + } catch { + this.unsubscribe(connection.id) + connection.close?.() + } } } diff --git a/hub/src/web/routes/events.ts b/hub/src/web/routes/events.ts index 0ccf28a6..5ac49032 100644 --- a/hub/src/web/routes/events.ts +++ b/hub/src/web/routes/events.ts @@ -97,6 +97,7 @@ export function createEventsRoutes( visibility, resumeFrom, send: (event, eventId) => stream.writeSSE({ data: JSON.stringify(event), id: eventId }), + close: () => stream.abort(), sendHeartbeat: async () => { await stream.writeSSE({ data: JSON.stringify({ @@ -132,6 +133,7 @@ export function createEventsRoutes( await new Promise((resolve) => { const done = () => resolve() + if (stream.aborted || c.req.raw.signal.aborted) { done(); return } c.req.raw.signal.addEventListener('abort', done, { once: true }) stream.onAbort(done) })