Improve event stream resilience

This commit is contained in:
Bohdan Triapitsyn
2026-04-24 12:30:11 +03:00
parent f0c3d5e326
commit 9d9af7e263
8 changed files with 240 additions and 2 deletions
@@ -1,6 +1,9 @@
export const MESSAGE_STREAM_GLOBAL_WS_PATH = '/api/global/event/ws';
export const MESSAGE_STREAM_DIRECTORY_WS_PATH = '/api/event/ws';
export const MESSAGE_STREAM_WS_HEARTBEAT_INTERVAL_MS = 15 * 1000;
// Per-client pending outbound WS buffer, not a payload or stream-size limit.
// Healthy clients stay near 0; this only trips when a client is far behind.
export const MESSAGE_STREAM_WS_MAX_BUFFERED_BYTES = 4 * 1024 * 1024;
export function parseSseEventEnvelope(block) {
if (!block || typeof block !== 'string') {
@@ -64,8 +67,23 @@ export function sendMessageStreamWsFrame(socket, payload) {
return false;
}
if (typeof socket.bufferedAmount === 'number' && socket.bufferedAmount > MESSAGE_STREAM_WS_MAX_BUFFERED_BYTES) {
try {
socket.close(1013, 'Message stream client is too slow');
} catch {
}
return false;
}
try {
socket.send(JSON.stringify(payload));
if (typeof socket.bufferedAmount === 'number' && socket.bufferedAmount > MESSAGE_STREAM_WS_MAX_BUFFERED_BYTES) {
try {
socket.close(1013, 'Message stream client is too slow');
} catch {
}
return false;
}
return true;
} catch {
return false;
@@ -3,6 +3,7 @@ import { describe, expect, it } from 'vitest';
import {
MESSAGE_STREAM_DIRECTORY_WS_PATH,
MESSAGE_STREAM_GLOBAL_WS_PATH,
MESSAGE_STREAM_WS_MAX_BUFFERED_BYTES,
parseSseEventEnvelope,
sendMessageStreamWsEvent,
sendMessageStreamWsFrame,
@@ -63,6 +64,30 @@ describe('event stream protocol helpers', () => {
expect(rawPayload).toBe('{"type":"ready"}');
});
it('closes slow websocket clients instead of adding more buffered data', () => {
let closeCall = null;
let sendCalls = 0;
const socket = {
readyState: 1,
bufferedAmount: MESSAGE_STREAM_WS_MAX_BUFFERED_BYTES + 1,
send() {
sendCalls += 1;
},
close(code, reason) {
closeCall = { code, reason };
},
};
const sent = sendMessageStreamWsFrame(socket, { type: 'ready' });
expect(sent).toBe(false);
expect(sendCalls).toBe(0);
expect(closeCall).toEqual({
code: 1013,
reason: 'Message stream client is too slow',
});
});
it('serializes event frames with routing metadata', () => {
let rawPayload = null;
const socket = {
+41 -1
View File
@@ -6,6 +6,43 @@ import {
shouldForwardProxyResponseHeader,
} from '../../proxy-headers.js';
export const waitForSseDrain = (res, signal) => new Promise((resolve) => {
if (signal?.aborted || res.writableEnded || res.destroyed) {
resolve();
return;
}
const cleanup = () => {
res.off?.('drain', onDone);
res.off?.('close', onDone);
res.off?.('error', onDone);
signal?.removeEventListener?.('abort', onDone);
};
const onDone = () => {
cleanup();
resolve();
};
res.once?.('drain', onDone);
res.once?.('close', onDone);
res.once?.('error', onDone);
signal?.addEventListener?.('abort', onDone, { once: true });
});
export const writeSseChunkWithBackpressure = async (res, value, signal) => {
if (!value || value.length === 0 || signal?.aborted || res.writableEnded || res.destroyed) {
return false;
}
const flushed = res.write(value);
if (flushed !== false) {
return true;
}
await waitForSseDrain(res, signal);
return !signal?.aborted && !res.writableEnded && !res.destroyed;
};
export const registerOpenCodeProxy = (app, deps) => {
const {
fs,
@@ -131,7 +168,10 @@ export const registerOpenCodeProxy = (app, deps) => {
break;
}
if (value && value.length > 0) {
res.write(value);
const canContinue = await writeSseChunkWithBackpressure(res, value, abortController.signal);
if (!canContinue) {
break;
}
}
}
@@ -209,7 +209,22 @@ export const registerScheduledTaskRoutes = (app, dependencies) => {
} catch {
}
const heartbeat = setInterval(() => {
try {
writeSseEvent(res, {
type: 'openchamber:heartbeat',
properties: {
timestamp: Date.now(),
},
});
} catch {
clearInterval(heartbeat);
clients.delete(res);
}
}, 25_000);
req.on('close', () => {
clearInterval(heartbeat);
clients.delete(res);
});
});
+23 -1
View File
@@ -1,8 +1,9 @@
import { afterEach, describe, expect, it } from 'vitest';
import { EventEmitter } from 'node:events';
import express from 'express';
import path from 'path';
import { registerOpenCodeProxy } from './lib/opencode/proxy.js';
import { registerOpenCodeProxy, writeSseChunkWithBackpressure } from './lib/opencode/proxy.js';
const listen = (app, host = '127.0.0.1') => new Promise((resolve, reject) => {
const server = app.listen(0, host, () => resolve(server));
@@ -81,6 +82,27 @@ describe('OpenCode proxy SSE forwarding', () => {
expect(seenAuthorization).toBe('Bearer test-token');
});
it('waits for drain when writing to a slow SSE response', async () => {
const writes = [];
const res = new EventEmitter();
res.writableEnded = false;
res.destroyed = false;
res.write = (value) => {
writes.push(value);
return false;
};
const controller = new AbortController();
const write = writeSseChunkWithBackpressure(res, Buffer.from('data: {"ok":true}\n\n'), controller.signal);
await new Promise((resolve) => setTimeout(resolve, 0));
expect(writes).toHaveLength(1);
res.emit('drain');
await expect(write).resolves.toBe(true);
});
it('routes generic API requests through external OpenCode base URL', async () => {
const upstream = express();
upstream.get('/config/providers', (_req, res) => {