Improve event stream resilience
This commit is contained in:
@@ -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 = {
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
Reference in New Issue
Block a user