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:
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
@@ -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),
|
||||
]);
|
||||
}
|
||||
});
|
||||
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user