fix(gateway): tolerate reconnect bursts without dropping sessions (#100)

Relax the proxy frame limiter to allow normal startup traffic while preserving abuse protection, and slow reconnect retries after policy-violation disconnects so remote gateways can recover cleanly.

Made-with: Cursor

Co-authored-by: iamlukethedev <lucas.guilherme@smartwayslfl.com>
This commit is contained in:
Luke The Dev
2026-04-03 23:19:59 -05:00
committed by GitHub
parent a18c8c630c
commit b573646f2d
4 changed files with 277 additions and 41 deletions
+46 -24
View File
@@ -2,17 +2,24 @@ import { describe, expect, it } from "vitest";
import { resolveGatewayAutoRetryDelayMs } from "@/lib/gateway/GatewayClient";
const baseParams = {
status: "disconnected" as const,
didAutoConnect: true,
hasConnectedOnce: true,
wasManualDisconnect: false,
gatewayUrl: "wss://remote.example",
errorMessage: null as string | null,
connectErrorCode: null as string | null,
lastDisconnectCode: null as number | null,
attempt: 0,
};
describe("resolveGatewayAutoRetryDelayMs", () => {
it("does not retry when upstream gateway url is missing on Studio host", () => {
const delay = resolveGatewayAutoRetryDelayMs({
status: "disconnected",
didAutoConnect: true,
hasConnectedOnce: true,
wasManualDisconnect: false,
gatewayUrl: "wss://remote.example",
...baseParams,
errorMessage: "Gateway error (studio.gateway_url_missing): Upstream gateway URL is missing.",
connectErrorCode: "studio.gateway_url_missing",
attempt: 0,
});
expect(delay).toBeNull();
@@ -20,15 +27,10 @@ describe("resolveGatewayAutoRetryDelayMs", () => {
it("does not retry when the upstream websocket upgrade fails", () => {
const delay = resolveGatewayAutoRetryDelayMs({
status: "disconnected",
didAutoConnect: true,
hasConnectedOnce: true,
wasManualDisconnect: false,
gatewayUrl: "wss://remote.example",
...baseParams,
errorMessage:
"Gateway error (studio.upstream_error): Failed to connect to upstream gateway WebSocket.",
connectErrorCode: "studio.upstream_error",
attempt: 0,
});
expect(delay).toBeNull();
@@ -36,15 +38,10 @@ describe("resolveGatewayAutoRetryDelayMs", () => {
it("does not retry when the upstream websocket handshake times out", () => {
const delay = resolveGatewayAutoRetryDelayMs({
status: "disconnected",
didAutoConnect: true,
hasConnectedOnce: true,
wasManualDisconnect: false,
gatewayUrl: "wss://remote.example",
...baseParams,
errorMessage:
"Gateway error (studio.upstream_timeout): Timed out connecting Studio to the upstream gateway WebSocket.",
connectErrorCode: "studio.upstream_timeout",
attempt: 0,
});
expect(delay).toBeNull();
@@ -52,18 +49,43 @@ describe("resolveGatewayAutoRetryDelayMs", () => {
it("does not retry when the upstream gateway explicitly rejects pairing", () => {
const delay = resolveGatewayAutoRetryDelayMs({
status: "disconnected",
didAutoConnect: true,
hasConnectedOnce: true,
wasManualDisconnect: false,
gatewayUrl: "wss://remote.example",
...baseParams,
errorMessage:
"Gateway error (studio.upstream_rejected): Upstream gateway rejected connect (1008): pairing required.",
connectErrorCode: "studio.upstream_rejected",
attempt: 0,
});
expect(delay).toBeNull();
});
it("uses a longer base delay when disconnected by rate limiting (code 1008)", () => {
const delay = resolveGatewayAutoRetryDelayMs({
...baseParams,
lastDisconnectCode: 1008,
attempt: 0,
});
expect(delay).toBe(15_000);
});
it("applies exponential backoff on top of rate-limit base delay", () => {
const delay = resolveGatewayAutoRetryDelayMs({
...baseParams,
lastDisconnectCode: 1008,
attempt: 1,
});
expect(delay).toBe(22_500);
});
it("uses standard base delay for normal disconnects", () => {
const delay = resolveGatewayAutoRetryDelayMs({
...baseParams,
lastDisconnectCode: 1012,
attempt: 0,
});
expect(delay).toBe(2_000);
});
});
+161
View File
@@ -712,4 +712,165 @@ describe("createGatewayProxy", () => {
}
});
it("allows short bursts of post-connect traffic without closing the socket", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "host-token-456" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-burst-ok",
method: "connect",
params: { auth: {} },
})
);
await waitForEvent(browser, "message");
for (let index = 0; index < 80; index += 1) {
browser.send(
JSON.stringify({
type: "req",
id: `burst-ok-${index}`,
method: "noop",
params: { index },
})
);
}
await new Promise((resolve) => setTimeout(resolve, 25));
expect(browser.readyState).toBe(WebSocket.OPEN);
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("still rate limits abusive bursts that exceed the token bucket", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "host-token-456" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-burst-limit",
method: "connect",
params: { auth: {} },
})
);
await waitForEvent(browser, "message");
const closePromise = waitForEvent<[number, Buffer]>(browser, "close");
for (let index = 0; index < 200; index += 1) {
browser.send(
JSON.stringify({
type: "req",
id: `burst-limit-${index}`,
method: "noop",
params: { index },
})
);
}
const [closeCode, closeReason] = await closePromise;
expect(closeCode).toBe(1008);
expect(closeReason.toString()).toBe("rate limit exceeded");
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
});