From cf2c2d571225e87c0971831fc0add4126ea6456d Mon Sep 17 00:00:00 2001 From: MAC Date: Fri, 26 Jun 2026 00:10:56 +0100 Subject: [PATCH] fix(api): globalize WS connection/subscription caps via SharedStore [BUG-005] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MAX_WS_CONNECTIONS and MAX_GLOBAL_SUBSCRIPTIONS were enforced against process-local state (clients.size, a plain globalSubscriptionCount variable) — under N horizontally-scaled replicas, the true fleet-wide cap became N×limit instead of limit, since each replica only sees its own slice. /ws/stats and /health's WS-utilization check had the same problem: totalConnections only reflected the answering replica's own clients.size. PR #189 already solved this exact class of bug for per-IP connection counts, rate-limit buckets, and auth-failure bans via a SharedStore abstraction (Upstash Redis when configured, in-memory fallback otherwise). This extends that same mechanism to the two caps it didn't cover: - Added addConnectionCount(key, delta) to SharedStore (both InMemoryStore and UpstashStore) — a signed-delta batch update, so the subscribe/unsubscribe loops (up to 50 channels per message) can flush one update after the loop instead of up to 50 sequential round-trips. incrementConnectionCount/decrementConnectionCount are unchanged and remain the right choice for the connection cap, which only changes by ±1 per connect/disconnect. - The subscription cap is checked once per message (read the global count at message start, track an in-loop local delta, compare against count+delta per channel) rather than once per channel — avoiding the round-trip cost while keeping the same non-atomic check-then-act semantics the per-IP caps in this file already use elsewhere. - getWebSocketMetrics() is now async and reports the true global totalConnections/totalSubscriptions from the shared store (plus the previously-unreported maxGlobalSubscriptions limit). connectionsPerSlab, messagesPerSec, and bytesPerSec stay local — they're inherently per-replica observability data, not safety caps, and aggregating them is a different, larger problem than this fix's scope. Also fixed clients.size >= MAX_WS_CONNECTIONS (line ~609), found while scoping this: identical per-replica-only bug, same file, same mechanism — fixing it alongside rather than leaving an identical bug next to the one being fixed. Added regression tests that seed the shared store directly (simulating "other replicas already at the cap") while this replica's own local state stays empty, proving the cap is now enforced fleet-wide rather than per-replica. Verified all three fail against the pre-fix code and pass against the fix. Co-authored-by: Claude Sonnet 4.6 --- src/middleware/shared-store.ts | 41 +++++ src/routes/health.ts | 4 +- src/routes/ws.ts | 128 +++++++++++---- tests/middleware/shared-store.test.ts | 45 ++++++ tests/routes/ws-global-caps.test.ts | 216 ++++++++++++++++++++++++++ 5 files changed, 403 insertions(+), 31 deletions(-) create mode 100644 tests/routes/ws-global-caps.test.ts diff --git a/src/middleware/shared-store.ts b/src/middleware/shared-store.ts index d8763aa..bb8ccec 100644 --- a/src/middleware/shared-store.ts +++ b/src/middleware/shared-store.ts @@ -85,6 +85,15 @@ export interface SharedStore { */ decrementConnectionCount(key: string): Promise; + /** + * Atomically add `delta` (positive or negative) to the counter for `key`. + * For batch updates where calling increment/decrement N times would cost + * N round-trips under Upstash (e.g. N subscriptions acquired/released in + * one WebSocket message). Floors at 0 and deletes the key, same as + * decrementConnectionCount. + */ + addConnectionCount(key: string, delta: number): Promise; + /** * Record an auth failure for an IP. * Returns the updated record so callers can decide whether to ban. @@ -189,6 +198,15 @@ export class InMemoryStore implements SharedStore { } } + async addConnectionCount(key: string, delta: number): Promise { + const next = (this.connCounts.get(key) ?? 0) + delta; + if (next <= 0) { + this.connCounts.delete(key); + } else { + this.connCounts.set(key, next); + } + } + async recordAuthFailure( ip: string, windowMs: number, @@ -416,6 +434,29 @@ export class UpstashStore implements SharedStore { } } + async addConnectionCount(key: string, delta: number): Promise { + try { + // Same 24h safety-net TTL as increment/decrementConnectionCount. + const script = ` + local val = redis.call("INCRBY", KEYS[1], ARGV[1]) + if val <= 0 then + redis.call("DEL", KEYS[1]) + else + redis.call("EXPIRE", KEYS[1], 86400) + end + return val + `; + await this.eval(script, [this.connKey(key)], [delta]); + } catch (err) { + logger.warn("UpstashStore.addConnectionCount failed, using fallback", { + key, + delta, + error: err instanceof Error ? err.message : String(err), + }); + await this.fallback.addConnectionCount(key, delta); + } + } + /** * Auth-failure record stored as a Redis HASH: * pcl:af: → { count, windowStart, bannedUntil } diff --git a/src/routes/health.ts b/src/routes/health.ts index 92050f3..be282d3 100644 --- a/src/routes/health.ts +++ b/src/routes/health.ts @@ -56,7 +56,7 @@ export function healthRoutes(): Hono { // Check WebSocket subsystem — saturated WS means new clients can't connect try { - const wsMetrics = getWebSocketMetrics(); + const wsMetrics = await getWebSocketMetrics(); const utilization = wsMetrics.totalConnections / wsMetrics.limits.maxGlobalConnections; checks.ws = utilization < 0.95; // degraded if >95% of connection slots used } catch { @@ -83,7 +83,7 @@ export function healthRoutes(): Hono { app.get("/ws/stats", requireApiKey(), async (c) => { try { - const metrics = getWebSocketMetrics(); + const metrics = await getWebSocketMetrics(); return c.json(metrics); } catch (err) { logger.error("Failed to get WebSocket metrics", { error: truncateErrorMessage(err instanceof Error ? err.message : err, 120) }); diff --git a/src/routes/ws.ts b/src/routes/ws.ts index b2ee32e..37c070d 100644 --- a/src/routes/ws.ts +++ b/src/routes/ws.ts @@ -110,8 +110,13 @@ interface WsClient { msgWindowStart: number; // Start of current rate-limit window } -// Track global subscription count across all clients -let globalSubscriptionCount = 0; +// Global connection/subscription caps are enforced via the SharedStore (same +// backend as the per-IP counters below) so the limits hold across replicas +// under horizontal scaling. Without this, MAX_WS_CONNECTIONS/ +// MAX_GLOBAL_SUBSCRIPTIONS are each enforced independently per replica — the +// true fleet-wide cap becomes N×limit instead of limit, for N replicas. +const GLOBAL_CONN_KEY = "ws:global-connections"; +const GLOBAL_SUB_KEY = "ws:global-subscriptions"; // Auth failure rate limiting per IP (issue #839: connection flood from repeat auth failures) // Tracks recent auth failures to temporarily ban repeat offenders. @@ -450,19 +455,34 @@ function flushPriceUpdate(slabAddress: string): void { } /** - * Get WebSocket metrics for /ws/stats endpoint + * Get WebSocket metrics for /ws/stats endpoint. + * + * totalConnections and totalSubscriptions are read from the SharedStore and + * reflect the true fleet-wide total across all replicas (see GLOBAL_CONN_KEY/ + * GLOBAL_SUB_KEY). connectionsPerSlab, messagesPerSec, and bytesPerSec remain + * local to this replica — they're inherently per-process observability data + * (this replica's own observed traffic/distribution), not safety caps, so + * aggregating them across replicas is a different and much larger problem + * than fixing the cap-enforcement bug this fixes; left out of scope. */ -export function getWebSocketMetrics(): any { +export async function getWebSocketMetrics(): Promise { const now = Date.now(); const elapsedSec = (now - metrics.lastResetTime) / 1000 || 1; + const store = getSharedStore(); + const [totalConnections, totalSubscriptions] = await Promise.all([ + store.getConnectionCount(GLOBAL_CONN_KEY), + store.getConnectionCount(GLOBAL_SUB_KEY), + ]); return { - totalConnections: metrics.totalConnections, + totalConnections, + totalSubscriptions, connectionsPerSlab: Object.fromEntries(metrics.connectionsPerSlab), messagesPerSec: parseFloat((metrics.messagesReceived / elapsedSec).toFixed(2)), bytesPerSec: parseInt((metrics.bytesSent / elapsedSec).toFixed(0), 10), limits: { maxGlobalConnections: MAX_WS_CONNECTIONS, + maxGlobalSubscriptions: MAX_GLOBAL_SUBSCRIPTIONS, maxConnectionsPerSlab: MAX_CONNECTIONS_PER_SLAB, maxConnectionsPerIp: MAX_CONNECTIONS_PER_IP, maxUnauthConnectionsPerIp: MAX_UNAUTHENTICATED_CONNECTIONS_PER_IP, @@ -605,9 +625,10 @@ export function setupWebSocket(server: Server): WebSocketServer { return; } - // H2: Reject if at max connections - if (clients.size >= MAX_WS_CONNECTIONS) { - logger.warn("Max global WS connections reached", { ip: clientIp }); + // H2: Reject if at max connections (global across all replicas, via SharedStore) + const globalConnCount = await getSharedStore().getConnectionCount(GLOBAL_CONN_KEY); + if (globalConnCount >= MAX_WS_CONNECTIONS) { + logger.warn("Max global WS connections reached", { ip: clientIp, count: globalConnCount }); ws.close(1008, "Connection limit reached"); return; } @@ -683,8 +704,9 @@ export function setupWebSocket(server: Server): WebSocketServer { }; clients.add(client); metrics.totalConnections = clients.size; - - logger.info("WebSocket connection established", { + await getSharedStore().incrementConnectionCount(GLOBAL_CONN_KEY); + + logger.info("WebSocket connection established", { ip: clientIp, authenticated, totalClients: clients.size @@ -870,7 +892,19 @@ export function setupWebSocket(server: Server): WebSocketServer { const subscribed: string[] = []; const errors: string[] = []; - + + // Check the global subscription cap once per message (not once + // per channel) and track local accepted-count as a delta, then + // flush a single batched update after the loop. This avoids up to + // MAX_CHANNELS_PER_MESSAGE (50) shared-store round-trips per + // message. The existing per-IP caps in this file are already + // non-atomic check-then-act, so the small intra-message overshoot + // window this introduces (bounded by MAX_CHANNELS_PER_MESSAGE) is + // not a new class of weaker guarantee. + const subStore = getSharedStore(); + const globalSubsAtStart = await subStore.getConnectionCount(GLOBAL_SUB_KEY); + let newSubsThisMessage = 0; + for (const channel of msg.channels) { if (typeof channel !== "string") continue; const safeChannel = channel.slice(0, 100); @@ -917,31 +951,37 @@ export function setupWebSocket(server: Server): WebSocketServer { continue; } - // Cap global subscriptions to prevent DoS - if (globalSubscriptionCount >= MAX_GLOBAL_SUBSCRIPTIONS) { + // Cap global subscriptions to prevent DoS (checked against the + // count as of this message's start plus any accepted so far in + // this same loop — see comment above the loop). + if (globalSubsAtStart + newSubsThisMessage >= MAX_GLOBAL_SUBSCRIPTIONS) { errors.push("Server subscription limit reached"); break; } - + // Cap subscriptions per client if (client.subscriptions.size >= MAX_SUBSCRIPTIONS_PER_CLIENT) { errors.push("Subscription limit per connection reached"); break; } - + // Check per-slab connection limit const slabClients = connectionsPerSlab.get(sanitized); if (slabClients && slabClients.size >= MAX_CONNECTIONS_PER_SLAB) { errors.push("Connection limit for this market reached"); continue; } - + client.subscriptions.add(fullChannel); - globalSubscriptionCount++; + newSubsThisMessage++; addClientToSlab(client, sanitized); subscribed.push(fullChannel); } - + + if (newSubsThisMessage > 0) { + await subStore.addConnectionCount(GLOBAL_SUB_KEY, newSubsThisMessage); + } + if (subscribed.length > 0) { safeSend(ws, { type: "subscribed", channels: subscribed }); @@ -1043,17 +1083,27 @@ export function setupWebSocket(server: Server): WebSocketServer { message: "Please use channels array. Subscribing to all channels for this slab.", }); + // Same batched cap-check pattern as the modern subscribe path + // above (at most 3 channels here, but kept consistent). + const legacySubStore = getSharedStore(); + const legacyGlobalSubsAtStart = await legacySubStore.getConnectionCount(GLOBAL_SUB_KEY); + let legacyNewSubs = 0; + const subscribed: string[] = []; for (const channel of channels) { if (client.subscriptions.has(channel)) continue; - if (globalSubscriptionCount >= MAX_GLOBAL_SUBSCRIPTIONS) break; + if (legacyGlobalSubsAtStart + legacyNewSubs >= MAX_GLOBAL_SUBSCRIPTIONS) break; if (client.subscriptions.size >= MAX_SUBSCRIPTIONS_PER_CLIENT) break; client.subscriptions.add(channel); - globalSubscriptionCount++; + legacyNewSubs++; subscribed.push(channel); } + if (legacyNewSubs > 0) { + await legacySubStore.addConnectionCount(GLOBAL_SUB_KEY, legacyNewSubs); + } + if (subscribed.length > 0) { addClientToSlab(client, sanitized); } @@ -1070,12 +1120,13 @@ export function setupWebSocket(server: Server): WebSocketServer { } const unsubscribed: string[] = []; - + let removedSubs = 0; + for (const channel of msg.channels) { if (client.subscriptions.delete(channel)) { - globalSubscriptionCount--; + removedSubs++; unsubscribed.push(channel); - + // Extract slab and remove from slab tracking if no more subs for this slab const slab = extractSlabFromChannel(channel); if (slab) { @@ -1088,7 +1139,11 @@ export function setupWebSocket(server: Server): WebSocketServer { } } } - + + if (removedSubs > 0) { + await getSharedStore().addConnectionCount(GLOBAL_SUB_KEY, -removedSubs); + } + if (unsubscribed.length > 0) { safeSend(ws, { type: "unsubscribed", channels: unsubscribed }); } @@ -1099,14 +1154,19 @@ export function setupWebSocket(server: Server): WebSocketServer { if (sanitized) { const channels = [`price:${sanitized}`, `trades:${sanitized}`, `funding:${sanitized}`]; const unsubscribed: string[] = []; - + let legacyRemovedSubs = 0; + for (const channel of channels) { if (client.subscriptions.delete(channel)) { - globalSubscriptionCount--; + legacyRemovedSubs++; unsubscribed.push(channel); } } - + + if (legacyRemovedSubs > 0) { + await getSharedStore().addConnectionCount(GLOBAL_SUB_KEY, -legacyRemovedSubs); + } + removeClientFromSlab(client, sanitized); safeSend(ws, { type: "unsubscribed", slabAddress: sanitized, channels: unsubscribed }); } @@ -1153,8 +1213,18 @@ export function setupWebSocket(server: Server): WebSocketServer { } // H2: O(1) removal with Set - // Decrement global subscription count for all client subscriptions - globalSubscriptionCount = Math.max(0, globalSubscriptionCount - client.subscriptions.size); + // Decrement global subscription + connection counts via the shared + // store (fire-and-forget — this handler is synchronous, matching the + // per-IP decrement pattern above). + const subsToRelease = client.subscriptions.size; + if (subsToRelease > 0) { + closeStore.addConnectionCount(GLOBAL_SUB_KEY, -subsToRelease).catch((err) => + logger.warn("addConnectionCount error on close (subscriptions)", { ip: client.ip, error: String(err) }) + ); + } + closeStore.decrementConnectionCount(GLOBAL_CONN_KEY).catch((err) => + logger.warn("decrementConnectionCount error on close (global connections)", { ip: client.ip, error: String(err) }) + ); clients.delete(client); metrics.totalConnections = clients.size; diff --git a/tests/middleware/shared-store.test.ts b/tests/middleware/shared-store.test.ts index 9deec45..ef1ce29 100644 --- a/tests/middleware/shared-store.test.ts +++ b/tests/middleware/shared-store.test.ts @@ -109,6 +109,41 @@ describe("InMemoryStore — connection counts", () => { }); }); +// --------------------------------------------------------------------------- +// InMemoryStore — addConnectionCount (batched delta, BUG-005) +// --------------------------------------------------------------------------- + +describe("InMemoryStore — addConnectionCount", () => { + it("adds a positive delta in one call", async () => { + const store = new InMemoryStore(); + await store.addConnectionCount("ws:global-subscriptions", 5); + expect(await store.getConnectionCount("ws:global-subscriptions")).toBe(5); + }); + + it("adds a negative delta to release a batch", async () => { + const store = new InMemoryStore(); + await store.addConnectionCount("ws:global-subscriptions", 5); + await store.addConnectionCount("ws:global-subscriptions", -3); + expect(await store.getConnectionCount("ws:global-subscriptions")).toBe(2); + }); + + it("floors at 0 and removes the key rather than going negative", async () => { + const store = new InMemoryStore(); + await store.addConnectionCount("ws:global-subscriptions", 2); + await store.addConnectionCount("ws:global-subscriptions", -10); + expect(await store.getConnectionCount("ws:global-subscriptions")).toBe(0); + }); + + it("is interoperable with incrementConnectionCount/decrementConnectionCount on the same key", async () => { + const store = new InMemoryStore(); + await store.incrementConnectionCount("ws:global-connections"); + await store.addConnectionCount("ws:global-connections", 4); + expect(await store.getConnectionCount("ws:global-connections")).toBe(5); + await store.decrementConnectionCount("ws:global-connections"); + expect(await store.getConnectionCount("ws:global-connections")).toBe(4); + }); +}); + // --------------------------------------------------------------------------- // InMemoryStore — auth failure ban tests // --------------------------------------------------------------------------- @@ -233,6 +268,16 @@ describe("UpstashStore — graceful fallback", () => { globalThis.fetch = originalFetch; }); + it("falls back gracefully for addConnectionCount", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network unreachable")); + + const store = new UpstashStore("https://test.upstash.io", "token"); + await expect(store.addConnectionCount("ws:global-subscriptions", 5)).resolves.toBeUndefined(); + + globalThis.fetch = originalFetch; + }); + it("falls back gracefully for auth ban checks", async () => { const originalFetch = globalThis.fetch; globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network unreachable")); diff --git a/tests/routes/ws-global-caps.test.ts b/tests/routes/ws-global-caps.test.ts new file mode 100644 index 0000000..c4e0b21 --- /dev/null +++ b/tests/routes/ws-global-caps.test.ts @@ -0,0 +1,216 @@ +/** + * Regression for BUG-005: MAX_WS_CONNECTIONS and MAX_GLOBAL_SUBSCRIPTIONS + * must be enforced via the SharedStore so they hold fleet-wide across + * replicas, not just against this process's own (possibly empty) local + * `clients` Set / subscription counter. + * + * These tests simulate "other replicas already at the cap" by bumping the + * shared store directly, then prove THIS replica — whose own local state is + * empty — still rejects new connections/subscriptions, which would not have + * been true before this fix. + */ +import { describe, it, expect, afterEach, vi } from "vitest"; +import http from "node:http"; +import WebSocket from "ws"; + +vi.mock("@percolator/shared", () => ({ + createLogger: vi.fn(() => ({ + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + })), + eventBus: { on: vi.fn() }, + getSupabase: vi.fn(() => ({ + from: vi.fn(() => ({ + select: vi.fn(() => ({ + eq: vi.fn(() => ({ + single: vi.fn(() => Promise.resolve({ data: null, error: null })), + })), + })), + })), + })), + sanitizeSlabAddress: vi.fn((s: string) => s), + sendInfoAlert: vi.fn(), +})); + +function waitForOpen(ws: WebSocket): Promise { + return new Promise((resolve, reject) => { + if (ws.readyState === WebSocket.OPEN) return resolve(); + ws.once("open", resolve); + ws.once("error", reject); + }); +} + +function waitForClose(ws: WebSocket): Promise { + return new Promise((resolve) => { + if (ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) { + return resolve(ws.readyState); + } + ws.once("close", (code) => resolve(code)); + }); +} + +function waitForMessage(ws: WebSocket, predicate: (msg: any) => boolean, timeoutMs = 1500): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error("timeout waiting for message")), timeoutMs); + const onMsg = (raw: unknown) => { + const parsed = JSON.parse(String(raw)); + if (predicate(parsed)) { + clearTimeout(timer); + ws.off("message", onMsg); + resolve(parsed); + } + }; + ws.on("message", onMsg); + }); +} + +interface TestServer { + server: http.Server; + port: number; + sharedStoreModule: typeof import("../../src/middleware/shared-store.js"); +} + +async function startServer(env: Record): Promise { + Object.assign(process.env, env); + vi.resetModules(); + const { setupWebSocket } = await import("../../src/routes/ws.js"); + // Import shared-store AFTER ws.js has resolved it, so this test gets the + // same module instance (and same _store singleton) ws.ts uses internally — + // required to seed/inspect the global counters it reads from. + const sharedStoreModule = await import("../../src/middleware/shared-store.js"); + const server = http.createServer((_req, res) => { + res.writeHead(200); + res.end("ok"); + }); + setupWebSocket(server as unknown as import("node:http").Server); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as { port: number }; + return { server, port, sharedStoreModule }; +} + +async function stopServer(ts: TestServer): Promise { + await new Promise((resolve, reject) => + ts.server.close((err) => (err ? reject(err) : resolve())), + ); +} + +describe("WS global connection cap is shared across replicas (BUG-005)", () => { + let ts: TestServer; + + afterEach(async () => { + if (ts) await stopServer(ts); + delete process.env.MAX_WS_CONNECTIONS; + delete process.env.WS_AUTH_REQUIRED; + }); + + it("rejects a new connection once the global count is at the cap, even though this replica's own local client set is empty", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + MAX_WS_CONNECTIONS: "2", + }); + + // Simulate 2 connections already established on OTHER replicas by + // bumping the shared store directly — this replica's own `clients` Set + // is still empty, so the pre-fix (local-only) check would have allowed + // a new connection here. + const store = ts.sharedStoreModule.getSharedStore(); + await store.incrementConnectionCount("ws:global-connections"); + await store.incrementConnectionCount("ws:global-connections"); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + const code = await waitForClose(ws); + expect(code).toBe(1008); + }); + + it("allows a connection when under the global cap", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + MAX_WS_CONNECTIONS: "2", + }); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + await waitForOpen(ws); + expect(ws.readyState).toBe(WebSocket.OPEN); + ws.close(); + await waitForClose(ws); + }); +}); + +describe("WS global subscription cap is shared across replicas (BUG-005)", () => { + let ts: TestServer; + + afterEach(async () => { + if (ts) await stopServer(ts); + delete process.env.WS_AUTH_REQUIRED; + }); + + it("rejects a new subscription once the global subscription count is at the cap (MAX_GLOBAL_SUBSCRIPTIONS = 1000, not env-overridable)", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + }); + + // Simulate 1000 subscriptions already held across OTHER replicas. + const store = ts.sharedStoreModule.getSharedStore(); + await store.addConnectionCount("ws:global-subscriptions", 1000); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + await waitForOpen(ws); + + ws.send(JSON.stringify({ type: "subscribe", channels: ["price:SLAB-GLOBAL-CAP"] })); + const errorMsg = await waitForMessage(ws, (m) => m.type === "error"); + expect(errorMsg.message).toMatch(/server subscription limit reached/i); + + ws.close(); + await waitForClose(ws); + }); + + it("allows a subscription when under the global cap", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + }); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + await waitForOpen(ws); + + ws.send(JSON.stringify({ type: "subscribe", channels: ["price:SLAB-UNDER-CAP"] })); + const subscribedMsg = await waitForMessage(ws, (m) => m.type === "subscribed"); + expect(subscribedMsg.channels).toContain("price:SLAB-UNDER-CAP"); + + ws.close(); + await waitForClose(ws); + }); +}); + +describe("getWebSocketMetrics reports the global (shared-store) count, not just this replica's local count (BUG-005)", () => { + let ts: TestServer; + + afterEach(async () => { + if (ts) await stopServer(ts); + delete process.env.WS_AUTH_REQUIRED; + }); + + it("totalConnections reflects connections established on other replicas, not seen locally", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + }); + + const { getWebSocketMetrics } = await import("../../src/routes/ws.js"); + const store = ts.sharedStoreModule.getSharedStore(); + + // Simulate 7 connections established on other replicas — this replica's + // own `clients` Set is still empty. + for (let i = 0; i < 7; i++) { + await store.incrementConnectionCount("ws:global-connections"); + } + + const metrics = await getWebSocketMetrics(); + expect(metrics.totalConnections).toBe(7); + }); +});