From ed2a97fda8a00da5c5c88cfa9c38bdab717aa23d Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sun, 10 May 2026 15:28:48 +0800 Subject: [PATCH] Extract websocket runtime metrics (#875) --- .../websocket-server.runtime-metrics.test.ts | 369 ------------------ .../server/src/server/websocket-server.ts | 205 ++-------- .../server/websocket/runtime-metrics.test.ts | 202 ++++++++++ .../src/server/websocket/runtime-metrics.ts | 215 ++++++++++ 4 files changed, 441 insertions(+), 550 deletions(-) delete mode 100644 packages/server/src/server/websocket-server.runtime-metrics.test.ts create mode 100644 packages/server/src/server/websocket/runtime-metrics.test.ts create mode 100644 packages/server/src/server/websocket/runtime-metrics.ts diff --git a/packages/server/src/server/websocket-server.runtime-metrics.test.ts b/packages/server/src/server/websocket-server.runtime-metrics.test.ts deleted file mode 100644 index 7d3f433c0..000000000 --- a/packages/server/src/server/websocket-server.runtime-metrics.test.ts +++ /dev/null @@ -1,369 +0,0 @@ -import { describe, expect, it, vi } from "vitest"; -import { z } from "zod"; -import type { Server as HTTPServer } from "http"; -import type pino from "pino"; -import type { AgentManager } from "./agent/agent-manager.js"; -import type { AgentStorage } from "./agent/agent-storage.js"; -import type { DownloadTokenStore } from "./file-download/token-store.js"; -import type { DaemonConfigStore } from "./daemon-config-store.js"; -import type { FileBackedChatService } from "./chat/chat-service.js"; -import type { LoopService } from "./loop-service.js"; -import type { ScheduleService } from "./schedule/service.js"; -import type { CheckoutDiffManager } from "./checkout-diff-manager.js"; -import type { SessionOutboundMessage, WSOutboundMessage } from "./messages.js"; - -const wsModuleMock = vi.hoisted(() => { - class MockWebSocketServer { - readonly handlers = new Map void>(); - - on(event: string, handler: (...args: unknown[]) => void) { - this.handlers.set(event, handler); - return this; - } - - close() { - // no-op - } - } - - return { MockWebSocketServer }; -}); - -vi.mock("ws", () => ({ - WebSocketServer: wsModuleMock.MockWebSocketServer, -})); - -import { VoiceAssistantWebSocketServer } from "./websocket-server.js"; -import { wrapSessionMessage } from "./messages.js"; - -interface WebSocketServerInternals { - flushRuntimeMetrics(options?: { final?: boolean }): void; - sendToClient(ws: unknown, message: WSOutboundMessage): void; - sendBinaryToClient(ws: unknown, frame: Uint8Array): void; - sessions: Map; -} - -const RuntimeMetricsLogSchema = z.object({ - outboundMessageTypesTop: z.array(z.tuple([z.string(), z.number()])), - outboundSessionMessageTypesTop: z.array(z.tuple([z.string(), z.number()])), - outboundAgentStreamTypesTop: z.array(z.tuple([z.string(), z.number()])), - outboundAgentStreamAgentsTop: z.array(z.tuple([z.string(), z.number()])), - outboundBinaryFrameTypesTop: z.array(z.tuple([z.string(), z.number()])), - bufferedAmount: z.object({ - p95: z.number(), - max: z.number(), - }), -}); - -type RuntimeMetricsLog = z.infer; - -interface TestSocket { - readyState: number; - bufferedAmount: number; - sent: Array; - afterSendBufferedAmounts: number[]; - send: (data: string | Uint8Array | ArrayBuffer) => void; - close: () => void; - on: () => void; - once: () => void; -} - -function createLogger() { - const logger = { - child: vi.fn(() => logger), - trace: vi.fn(), - debug: vi.fn(), - info: vi.fn(), - warn: vi.fn(), - error: vi.fn(), - }; - return logger; -} - -function createServer(logger: ReturnType) { - const daemonConfigStore = { - onChange: vi.fn(() => () => {}), - }; - - return new VoiceAssistantWebSocketServer( - {} as unknown as HTTPServer, - logger as unknown as pino.Logger, - "srv-test", - { - setAgentAttentionCallback: vi.fn(), - getAgent: vi.fn(() => null), - getMetricsSnapshot: vi.fn(() => ({ - totalAgents: 0, - idleAgents: 0, - runningAgents: 0, - pendingPermissionAgents: 0, - erroredAgents: 0, - })), - } as unknown as AgentManager, - {} as unknown as AgentStorage, - {} as unknown as DownloadTokenStore, - "/tmp/paseo-test", - daemonConfigStore as unknown as DaemonConfigStore, - null, - { allowedOrigins: new Set() }, - undefined, - undefined, - undefined, - undefined, - undefined, - undefined, - false, - "1.2.3-test", - undefined, - undefined, - undefined, - {} as unknown as FileBackedChatService, - {} as unknown as LoopService, - {} as unknown as ScheduleService, - { - subscribe: vi.fn(), - scheduleRefreshForCwd: vi.fn(), - getMetrics: vi.fn(() => ({ - checkoutDiffTargetCount: 0, - checkoutDiffSubscriptionCount: 0, - checkoutDiffWatcherCount: 0, - checkoutDiffFallbackRefreshTargetCount: 0, - })), - dispose: vi.fn(), - } as unknown as CheckoutDiffManager, - ); -} - -function createSocket(afterSendBufferedAmounts: number[]): TestSocket { - const socket: TestSocket = { - readyState: 1, - bufferedAmount: 0, - sent: [], - afterSendBufferedAmounts, - send: vi.fn((data) => { - socket.sent.push(data); - socket.bufferedAmount = afterSendBufferedAmounts.shift() ?? socket.bufferedAmount; - }), - close: vi.fn(), - on: vi.fn(), - once: vi.fn(), - }; - return socket; -} - -function flushRuntimeMetrics(server: VoiceAssistantWebSocketServer): void { - (server as unknown as WebSocketServerInternals).flushRuntimeMetrics({ final: true }); -} - -function getRuntimeMetricsLog(logger: ReturnType): RuntimeMetricsLog { - const metricsCall = logger.info.mock.calls.find((call) => call[1] === "ws_runtime_metrics"); - expect(metricsCall).toBeDefined(); - return RuntimeMetricsLogSchema.parse(metricsCall![0]); -} - -function sendToClient( - server: VoiceAssistantWebSocketServer, - socket: TestSocket, - message: WSOutboundMessage, -) { - (server as unknown as WebSocketServerInternals).sendToClient(socket, message); -} - -function sendBinaryToClient( - server: VoiceAssistantWebSocketServer, - socket: TestSocket, - frame: Uint8Array, -) { - (server as unknown as WebSocketServerInternals).sendBinaryToClient(socket, frame); -} - -function attachSessionSocket(server: VoiceAssistantWebSocketServer, socket: TestSocket): void { - (server as unknown as WebSocketServerInternals).sessions.set(socket, { - session: {}, - clientId: "client-1", - appVersion: null, - connectionLogger: createLogger(), - sockets: new Set([socket]), - externalDisconnectCleanupTimeout: null, - }); -} - -function agentStreamMessage(params: { - agentId: string; - event: Extract["payload"]["event"]; -}): WSOutboundMessage { - return wrapSessionMessage({ - type: "agent_stream", - payload: { - agentId: params.agentId, - event: params.event, - timestamp: "2026-04-17T00:00:00.000Z", - }, - }); -} - -describe("VoiceAssistantWebSocketServer runtime metrics", () => { - it("records outbound message type counts in the ws runtime metrics window", () => { - const logger = createLogger(); - const server = createServer(logger); - const broadcastSocket = createSocket([0]); - const directSocket = createSocket([0, 0]); - attachSessionSocket(server, broadcastSocket); - - server.broadcast( - agentStreamMessage({ - agentId: "agent-1", - event: { - type: "turn_completed", - provider: "codex", - }, - }), - ); - sendToClient( - server, - directSocket, - wrapSessionMessage({ - type: "status", - payload: { - status: "ok", - message: "ready", - }, - }), - ); - sendToClient(server, directSocket, { type: "pong" }); - - flushRuntimeMetrics(server); - - const metrics = getRuntimeMetricsLog(logger); - expect(metrics.outboundMessageTypesTop).toEqual([ - ["session_message", 2], - ["pong", 1], - ]); - expect(metrics.outboundSessionMessageTypesTop).toEqual([ - ["agent_stream", 1], - ["status", 1], - ]); - }); - - it("records agent_stream subtypes and top agents", () => { - const logger = createLogger(); - const server = createServer(logger); - const socket = createSocket([0, 0, 0, 0]); - - sendToClient( - server, - socket, - agentStreamMessage({ - agentId: "agent-1", - event: { - type: "timeline", - provider: "codex", - item: { type: "assistant_message", text: "hello" }, - }, - }), - ); - sendToClient( - server, - socket, - agentStreamMessage({ - agentId: "agent-1", - event: { - type: "timeline", - provider: "codex", - item: { type: "reasoning", text: "thinking" }, - }, - }), - ); - sendToClient( - server, - socket, - agentStreamMessage({ - agentId: "agent-1", - event: { - type: "turn_completed", - provider: "codex", - }, - }), - ); - sendToClient( - server, - socket, - agentStreamMessage({ - agentId: "agent-2", - event: { - type: "timeline", - provider: "codex", - item: { type: "assistant_message", text: "there" }, - }, - }), - ); - - flushRuntimeMetrics(server); - - const metrics = getRuntimeMetricsLog(logger); - expect(metrics.outboundAgentStreamTypesTop).toEqual([ - ["timeline:assistant_message", 2], - ["timeline:reasoning", 1], - ["turn_completed", 1], - ]); - expect(metrics.outboundAgentStreamAgentsTop).toEqual([ - ["agent-1", 3], - ["agent-2", 1], - ]); - }); - - it("records bufferedAmount p95 and max from samples taken after send", () => { - const logger = createLogger(); - const server = createServer(logger); - const broadcastSocket = createSocket([0]); - const directSocket = createSocket([10, 50]); - const binarySocket = createSocket([100]); - attachSessionSocket(server, broadcastSocket); - - server.broadcast( - agentStreamMessage({ - agentId: "agent-1", - event: { - type: "turn_completed", - provider: "codex", - }, - }), - ); - sendToClient( - server, - directSocket, - wrapSessionMessage({ - type: "status", - payload: { status: "ok" }, - }), - ); - sendToClient(server, directSocket, { type: "pong" }); - sendBinaryToClient(server, binarySocket, new Uint8Array([1, 2, 3])); - - flushRuntimeMetrics(server); - - const metrics = getRuntimeMetricsLog(logger); - expect(metrics.bufferedAmount).toEqual({ - p95: 100, - max: 100, - }); - }); - - it("counts binary frames without decoding", () => { - const logger = createLogger(); - const server = createServer(logger); - const socket = createSocket([12, 24]); - const frame: Uint8Array = new Proxy(Buffer.from([0xff, 0xfe, 0xfd, 0x00, 0xc0]), { - get(_target, property) { - throw new Error(`Binary frame payload was unexpectedly accessed via ${String(property)}`); - }, - }); - - expect(() => sendBinaryToClient(server, socket, frame)).not.toThrow(); - - flushRuntimeMetrics(server); - - const metrics = getRuntimeMetricsLog(logger); - expect(metrics.outboundBinaryFrameTypesTop).toEqual([["binary", 1]]); - }); -}); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 17fe30806..4f63c0fed 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -16,7 +16,6 @@ import type { DaemonConfigStore, MutableDaemonConfig } from "./daemon-config-sto import { applyMutableProviderConfigToOverrides } from "./daemon-config-store.js"; import { type ServerInfoStatusPayload, - type SessionOutboundMessage, type WorkspaceSetupSnapshot, type WSHelloMessage, type WSInboundMessage, @@ -58,6 +57,10 @@ import { isBearerTokenValid, type DaemonAuthConfig, } from "./auth.js"; +import { + WebSocketRuntimeMetricsWindow, + type WebSocketRuntimeCounters, +} from "./websocket/runtime-metrics.js"; const WS_CLOSE_DAEMON_AUTH_FAILED = 4401; @@ -270,24 +273,6 @@ interface SessionConnection { externalDisconnectCleanupTimeout: ReturnType | null; } -interface WebSocketRuntimeCounters { - connectedAwaitingHello: number; - helloResumed: number; - helloNew: number; - pendingDisconnected: number; - sessionDisconnectedWaitingReconnect: number; - sessionSocketDisconnectedAttached: number; - sessionCleanup: number; - validationFailed: number; - binaryBeforeHelloRejected: number; - pendingMessageRejectedBeforeHello: number; - missingConnectionForMessage: number; - unexpectedHelloOnActiveConnection: number; - relayExternalSocketAttached: number; - originRejected: number; - hostRejected: number; -} - const SLOW_REQUEST_THRESHOLD_MS = 500; const EXTERNAL_SESSION_DISCONNECT_GRACE_MS = 90_000; const HELLO_TIMEOUT_MS = 15_000; @@ -382,33 +367,7 @@ export class VoiceAssistantWebSocketServer { | ((workspaceId: string, oldBranch: string | null, newBranch: string | null) => void) | null; private serverCapabilities: ServerCapabilities | undefined; - private runtimeWindowStartedAt = Date.now(); - private readonly runtimeCounters: WebSocketRuntimeCounters = { - connectedAwaitingHello: 0, - helloResumed: 0, - helloNew: 0, - pendingDisconnected: 0, - sessionDisconnectedWaitingReconnect: 0, - sessionSocketDisconnectedAttached: 0, - sessionCleanup: 0, - validationFailed: 0, - binaryBeforeHelloRejected: 0, - pendingMessageRejectedBeforeHello: 0, - missingConnectionForMessage: 0, - unexpectedHelloOnActiveConnection: 0, - relayExternalSocketAttached: 0, - originRejected: 0, - hostRejected: 0, - }; - private readonly inboundMessageCounts = new Map(); - private readonly inboundSessionRequestCounts = new Map(); - private readonly outboundMessageCounts = new Map(); - private readonly outboundSessionMessageCounts = new Map(); - private readonly outboundAgentStreamCounts = new Map(); - private readonly outboundAgentStreamByAgentCounts = new Map(); - private readonly outboundBinaryFrameCounts = new Map(); - private readonly bufferedAmountSamples: number[] = []; - private readonly requestLatencies = new Map(); + private readonly runtimeMetrics = new WebSocketRuntimeMetricsWindow(); private runtimeMetricsInterval: ReturnType | null = null; private unsubscribeSpeechReadiness: (() => void) | null = null; private unsubscribeDaemonConfigChange: (() => void) | null = null; @@ -668,7 +627,7 @@ export class VoiceAssistantWebSocketServer { // WebSocket.OPEN = 1 if (ws.readyState === 1) { ws.send(payload); - this.recordOutboundMessage(message, ws); + this.runtimeMetrics.recordOutboundMessage(message, ws.bufferedAmount); } } } @@ -785,7 +744,7 @@ export class VoiceAssistantWebSocketServer { } try { ws.send(JSON.stringify(message)); - this.recordOutboundMessage(message, ws); + this.runtimeMetrics.recordOutboundMessage(message, ws.bufferedAmount); } catch (err) { this.logger.warn({ err }, "ws_send_failed"); } @@ -797,7 +756,7 @@ export class VoiceAssistantWebSocketServer { } try { ws.send(frame); - this.recordOutboundBinaryFrame(ws); + this.runtimeMetrics.recordOutboundBinaryFrame(ws.bufferedAmount); } catch (err) { this.logger.warn({ err }, "ws_send_binary_failed"); } @@ -1569,116 +1528,19 @@ export class VoiceAssistantWebSocketServer { } private incrementRuntimeCounter(counter: keyof WebSocketRuntimeCounters): void { - this.runtimeCounters[counter] += 1; - } - - private incrementCount(map: Map, key: string): void { - map.set(key, (map.get(key) ?? 0) + 1); + this.runtimeMetrics.incrementCounter(counter); } private recordInboundMessageType(type: string): void { - this.incrementCount(this.inboundMessageCounts, type); + this.runtimeMetrics.recordInboundMessage(type); } private recordInboundSessionRequestType(type: string): void { - this.incrementCount(this.inboundSessionRequestCounts, type); - } - - private recordOutboundMessage(message: WSOutboundMessage, ws: WebSocketLike): void { - if (message.type !== "session") { - this.incrementCount(this.outboundMessageCounts, message.type); - this.recordBufferedAmount(ws); - return; - } - - this.incrementCount(this.outboundMessageCounts, "session_message"); - this.incrementCount(this.outboundSessionMessageCounts, message.message.type); - - if (message.message.type === "agent_stream") { - this.recordOutboundAgentStreamMessage(message.message.payload); - } - - this.recordBufferedAmount(ws); - } - - private recordOutboundAgentStreamMessage( - payload: Extract["payload"], - ): void { - const { agentId, event } = payload; - const eventType = event.type === "timeline" ? `timeline:${event.item.type}` : event.type; - this.incrementCount(this.outboundAgentStreamCounts, eventType); - this.incrementCount(this.outboundAgentStreamByAgentCounts, agentId); - } - - private recordOutboundBinaryFrame(ws: WebSocketLike): void { - this.incrementCount(this.outboundBinaryFrameCounts, "binary"); - this.recordBufferedAmount(ws); - } - - private recordBufferedAmount(ws: WebSocketLike): void { - if (typeof ws.bufferedAmount !== "number") { - return; - } - this.bufferedAmountSamples.push(ws.bufferedAmount); + this.runtimeMetrics.recordInboundSessionRequest(type); } private recordRequestLatency(type: string, durationMs: number): void { - let latencies = this.requestLatencies.get(type); - if (!latencies) { - latencies = []; - this.requestLatencies.set(type, latencies); - } - latencies.push(durationMs); - } - - private getTopCounts(map: Map, limit: number): Array<[string, number]> { - return [...map.entries()].sort((a, b) => b[1] - a[1]).slice(0, limit); - } - - private computeLatencyStats(): Array<{ - type: string; - count: number; - minMs: number; - maxMs: number; - p50Ms: number; - totalMs: number; - }> { - const stats: Array<{ - type: string; - count: number; - minMs: number; - maxMs: number; - p50Ms: number; - totalMs: number; - }> = []; - for (const [type, latencies] of this.requestLatencies) { - if (latencies.length === 0) continue; - latencies.sort((a, b) => a - b); - const count = latencies.length; - const minMs = Math.round(latencies[0]); - const maxMs = Math.round(latencies[count - 1]); - const p50Ms = Math.round(latencies[Math.floor(count / 2)]); - const totalMs = Math.round(latencies.reduce((sum, v) => sum + v, 0)); - stats.push({ type, count, minMs, maxMs, p50Ms, totalMs }); - } - stats.sort((a, b) => b.totalMs - a.totalMs); - return stats.slice(0, 15); - } - - private computeBufferedAmountStats(): { - p95: number; - max: number; - } { - if (this.bufferedAmountSamples.length === 0) { - return { p95: 0, max: 0 }; - } - - const samples = [...this.bufferedAmountSamples].sort((a, b) => a - b); - const p95Index = Math.ceil(samples.length * 0.95) - 1; - return { - p95: samples[p95Index] ?? 0, - max: samples[samples.length - 1] ?? 0, - }; + this.runtimeMetrics.recordRequestLatency(type, durationMs); } private collectSessionRuntimeMetrics(): WebSocketRuntimeMetrics { @@ -1707,8 +1569,7 @@ export class VoiceAssistantWebSocketServer { } private flushRuntimeMetrics(options?: { final?: boolean }): void { - const now = Date.now(); - const windowMs = Math.max(0, now - this.runtimeWindowStartedAt); + const runtimeMetrics = this.runtimeMetrics.snapshotAndReset(); const activeConnections = new Set(this.sessions.values()).size; const activeSockets = this.sessions.size; const pendingConnections = this.pendingConnections.size; @@ -1717,13 +1578,11 @@ export class VoiceAssistantWebSocketServer { connection.sockets.size === 0 && connection.externalDisconnectCleanupTimeout !== null, ).length; const sessionMetrics = this.collectSessionRuntimeMetrics(); - const latencyStats = this.computeLatencyStats(); - const bufferedAmountStats = this.computeBufferedAmountStats(); const agentSnapshot = this.agentManager.getMetricsSnapshot(); this.logger.info( { - windowMs, + windowMs: runtimeMetrics.windowMs, final: Boolean(options?.final), sessions: { activeConnections, @@ -1734,37 +1593,21 @@ export class VoiceAssistantWebSocketServer { activeSockets, pendingConnections, }, - counters: { ...this.runtimeCounters }, - inboundMessageTypesTop: this.getTopCounts(this.inboundMessageCounts, 12), - inboundSessionRequestTypesTop: this.getTopCounts(this.inboundSessionRequestCounts, 20), - outboundMessageTypesTop: this.getTopCounts(this.outboundMessageCounts, 12), - outboundSessionMessageTypesTop: this.getTopCounts(this.outboundSessionMessageCounts, 20), - outboundAgentStreamTypesTop: this.getTopCounts(this.outboundAgentStreamCounts, 20), - outboundAgentStreamAgentsTop: this.getTopCounts(this.outboundAgentStreamByAgentCounts, 20), - outboundBinaryFrameTypesTop: this.getTopCounts(this.outboundBinaryFrameCounts, 12), - bufferedAmount: bufferedAmountStats, + counters: runtimeMetrics.counters, + inboundMessageTypesTop: runtimeMetrics.inboundMessageTypesTop, + inboundSessionRequestTypesTop: runtimeMetrics.inboundSessionRequestTypesTop, + outboundMessageTypesTop: runtimeMetrics.outboundMessageTypesTop, + outboundSessionMessageTypesTop: runtimeMetrics.outboundSessionMessageTypesTop, + outboundAgentStreamTypesTop: runtimeMetrics.outboundAgentStreamTypesTop, + outboundAgentStreamAgentsTop: runtimeMetrics.outboundAgentStreamAgentsTop, + outboundBinaryFrameTypesTop: runtimeMetrics.outboundBinaryFrameTypesTop, + bufferedAmount: runtimeMetrics.bufferedAmount, runtime: sessionMetrics, - latency: latencyStats, + latency: runtimeMetrics.latency, agents: agentSnapshot, }, "ws_runtime_metrics", ); - - for (const counter of Object.keys(this.runtimeCounters) as Array< - keyof WebSocketRuntimeCounters - >) { - this.runtimeCounters[counter] = 0; - } - this.inboundMessageCounts.clear(); - this.inboundSessionRequestCounts.clear(); - this.outboundMessageCounts.clear(); - this.outboundSessionMessageCounts.clear(); - this.outboundAgentStreamCounts.clear(); - this.outboundAgentStreamByAgentCounts.clear(); - this.outboundBinaryFrameCounts.clear(); - this.bufferedAmountSamples.length = 0; - this.requestLatencies.clear(); - this.runtimeWindowStartedAt = now; } private getClientActivityState(session: Session): ClientPresenceState { diff --git a/packages/server/src/server/websocket/runtime-metrics.test.ts b/packages/server/src/server/websocket/runtime-metrics.test.ts new file mode 100644 index 000000000..11cc64037 --- /dev/null +++ b/packages/server/src/server/websocket/runtime-metrics.test.ts @@ -0,0 +1,202 @@ +import { describe, expect, it } from "vitest"; +import type { SessionOutboundMessage, WSOutboundMessage } from "../messages.js"; +import { wrapSessionMessage } from "../messages.js"; +import { WebSocketRuntimeMetricsWindow } from "./runtime-metrics.js"; + +function createMetricsWindow(): { + metrics: WebSocketRuntimeMetricsWindow; + advanceClock(ms: number): void; +} { + let now = 1_000; + return { + metrics: new WebSocketRuntimeMetricsWindow(() => now), + advanceClock(ms: number) { + now += ms; + }, + }; +} + +function agentStreamMessage(params: { + agentId: string; + event: Extract["payload"]["event"]; +}): WSOutboundMessage { + return wrapSessionMessage({ + type: "agent_stream", + payload: { + agentId: params.agentId, + event: params.event, + timestamp: "2026-04-17T00:00:00.000Z", + }, + }); +} + +describe("WebSocketRuntimeMetricsWindow", () => { + it("records outbound message type counts in the runtime metrics window", () => { + const { metrics } = createMetricsWindow(); + + metrics.recordOutboundMessage( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "turn_completed", + provider: "codex", + }, + }), + 0, + ); + metrics.recordOutboundMessage( + wrapSessionMessage({ + type: "status", + payload: { + status: "ok", + message: "ready", + }, + }), + 0, + ); + metrics.recordOutboundMessage({ type: "pong" }, 0); + + const snapshot = metrics.snapshotAndReset(); + + expect(snapshot.outboundMessageTypesTop).toEqual([ + ["session_message", 2], + ["pong", 1], + ]); + expect(snapshot.outboundSessionMessageTypesTop).toEqual([ + ["agent_stream", 1], + ["status", 1], + ]); + }); + + it("records agent_stream subtypes and top agents", () => { + const { metrics } = createMetricsWindow(); + + metrics.recordOutboundMessage( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "hello" }, + }, + }), + 0, + ); + metrics.recordOutboundMessage( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "timeline", + provider: "codex", + item: { type: "reasoning", text: "thinking" }, + }, + }), + 0, + ); + metrics.recordOutboundMessage( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "turn_completed", + provider: "codex", + }, + }), + 0, + ); + metrics.recordOutboundMessage( + agentStreamMessage({ + agentId: "agent-2", + event: { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "there" }, + }, + }), + 0, + ); + + const snapshot = metrics.snapshotAndReset(); + + expect(snapshot.outboundAgentStreamTypesTop).toEqual([ + ["timeline:assistant_message", 2], + ["timeline:reasoning", 1], + ["turn_completed", 1], + ]); + expect(snapshot.outboundAgentStreamAgentsTop).toEqual([ + ["agent-1", 3], + ["agent-2", 1], + ]); + }); + + it("records bufferedAmount p95 and max from samples taken after send", () => { + const { metrics } = createMetricsWindow(); + + metrics.recordOutboundMessage( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "turn_completed", + provider: "codex", + }, + }), + 0, + ); + metrics.recordOutboundMessage( + wrapSessionMessage({ + type: "status", + payload: { status: "ok" }, + }), + 10, + ); + metrics.recordOutboundMessage({ type: "pong" }, 50); + metrics.recordOutboundBinaryFrame(100); + + const snapshot = metrics.snapshotAndReset(); + + expect(snapshot.bufferedAmount).toEqual({ + p95: 100, + max: 100, + }); + }); + + it("counts binary frames without decoding", () => { + const { metrics } = createMetricsWindow(); + + metrics.recordOutboundBinaryFrame(24); + + const snapshot = metrics.snapshotAndReset(); + expect(snapshot.outboundBinaryFrameTypesTop).toEqual([["binary", 1]]); + }); + + it("resets the runtime window after producing a snapshot", () => { + const { metrics, advanceClock } = createMetricsWindow(); + metrics.incrementCounter("helloNew"); + metrics.recordInboundMessage("session"); + metrics.recordInboundSessionRequest("send"); + metrics.recordRequestLatency("send", 12.4); + advanceClock(250); + + const firstSnapshot = metrics.snapshotAndReset(); + const secondSnapshot = metrics.snapshotAndReset(); + + expect(firstSnapshot.windowMs).toBe(250); + expect(firstSnapshot.counters.helloNew).toBe(1); + expect(firstSnapshot.inboundMessageTypesTop).toEqual([["session", 1]]); + expect(firstSnapshot.inboundSessionRequestTypesTop).toEqual([["send", 1]]); + expect(firstSnapshot.latency).toEqual([ + { + type: "send", + count: 1, + minMs: 12, + maxMs: 12, + p50Ms: 12, + totalMs: 12, + }, + ]); + expect(secondSnapshot.windowMs).toBe(0); + expect(secondSnapshot.counters.helloNew).toBe(0); + expect(secondSnapshot.inboundMessageTypesTop).toEqual([]); + expect(secondSnapshot.inboundSessionRequestTypesTop).toEqual([]); + expect(secondSnapshot.latency).toEqual([]); + }); +}); diff --git a/packages/server/src/server/websocket/runtime-metrics.ts b/packages/server/src/server/websocket/runtime-metrics.ts new file mode 100644 index 000000000..76fbc39bd --- /dev/null +++ b/packages/server/src/server/websocket/runtime-metrics.ts @@ -0,0 +1,215 @@ +import type { SessionOutboundMessage, WSOutboundMessage } from "../messages.js"; + +export interface WebSocketRuntimeCounters { + connectedAwaitingHello: number; + helloResumed: number; + helloNew: number; + pendingDisconnected: number; + sessionDisconnectedWaitingReconnect: number; + sessionSocketDisconnectedAttached: number; + sessionCleanup: number; + validationFailed: number; + binaryBeforeHelloRejected: number; + pendingMessageRejectedBeforeHello: number; + missingConnectionForMessage: number; + unexpectedHelloOnActiveConnection: number; + relayExternalSocketAttached: number; + originRejected: number; + hostRejected: number; +} + +export interface WebSocketRuntimeMetricsSnapshot { + windowMs: number; + counters: WebSocketRuntimeCounters; + inboundMessageTypesTop: Array<[string, number]>; + inboundSessionRequestTypesTop: Array<[string, number]>; + outboundMessageTypesTop: Array<[string, number]>; + outboundSessionMessageTypesTop: Array<[string, number]>; + outboundAgentStreamTypesTop: Array<[string, number]>; + outboundAgentStreamAgentsTop: Array<[string, number]>; + outboundBinaryFrameTypesTop: Array<[string, number]>; + bufferedAmount: { + p95: number; + max: number; + }; + latency: Array<{ + type: string; + count: number; + minMs: number; + maxMs: number; + p50Ms: number; + totalMs: number; + }>; +} + +type Clock = () => number; + +export class WebSocketRuntimeMetricsWindow { + private windowStartedAt: number; + private readonly counters: WebSocketRuntimeCounters = createRuntimeCounters(); + private readonly inboundMessageCounts = new Map(); + private readonly inboundSessionRequestCounts = new Map(); + private readonly outboundMessageCounts = new Map(); + private readonly outboundSessionMessageCounts = new Map(); + private readonly outboundAgentStreamCounts = new Map(); + private readonly outboundAgentStreamByAgentCounts = new Map(); + private readonly outboundBinaryFrameCounts = new Map(); + private readonly bufferedAmountSamples: number[] = []; + private readonly requestLatencies = new Map(); + + constructor(private readonly clock: Clock = Date.now) { + this.windowStartedAt = this.clock(); + } + + incrementCounter(counter: keyof WebSocketRuntimeCounters): void { + this.counters[counter] += 1; + } + + recordInboundMessage(type: string): void { + incrementCount(this.inboundMessageCounts, type); + } + + recordInboundSessionRequest(type: string): void { + incrementCount(this.inboundSessionRequestCounts, type); + } + + recordOutboundMessage(message: WSOutboundMessage, bufferedAmount?: number): void { + if (message.type !== "session") { + incrementCount(this.outboundMessageCounts, message.type); + this.recordBufferedAmount(bufferedAmount); + return; + } + + incrementCount(this.outboundMessageCounts, "session_message"); + incrementCount(this.outboundSessionMessageCounts, message.message.type); + + if (message.message.type === "agent_stream") { + this.recordOutboundAgentStreamMessage(message.message.payload); + } + + this.recordBufferedAmount(bufferedAmount); + } + + recordOutboundBinaryFrame(bufferedAmount?: number): void { + incrementCount(this.outboundBinaryFrameCounts, "binary"); + this.recordBufferedAmount(bufferedAmount); + } + + recordRequestLatency(type: string, durationMs: number): void { + let latencies = this.requestLatencies.get(type); + if (!latencies) { + latencies = []; + this.requestLatencies.set(type, latencies); + } + latencies.push(durationMs); + } + + snapshotAndReset(): WebSocketRuntimeMetricsSnapshot { + const now = this.clock(); + const snapshot: WebSocketRuntimeMetricsSnapshot = { + windowMs: Math.max(0, now - this.windowStartedAt), + counters: { ...this.counters }, + inboundMessageTypesTop: getTopCounts(this.inboundMessageCounts, 12), + inboundSessionRequestTypesTop: getTopCounts(this.inboundSessionRequestCounts, 20), + outboundMessageTypesTop: getTopCounts(this.outboundMessageCounts, 12), + outboundSessionMessageTypesTop: getTopCounts(this.outboundSessionMessageCounts, 20), + outboundAgentStreamTypesTop: getTopCounts(this.outboundAgentStreamCounts, 20), + outboundAgentStreamAgentsTop: getTopCounts(this.outboundAgentStreamByAgentCounts, 20), + outboundBinaryFrameTypesTop: getTopCounts(this.outboundBinaryFrameCounts, 12), + bufferedAmount: this.computeBufferedAmountStats(), + latency: this.computeLatencyStats(), + }; + + this.reset(now); + return snapshot; + } + + private recordOutboundAgentStreamMessage( + payload: Extract["payload"], + ): void { + const { agentId, event } = payload; + const eventType = event.type === "timeline" ? `timeline:${event.item.type}` : event.type; + incrementCount(this.outboundAgentStreamCounts, eventType); + incrementCount(this.outboundAgentStreamByAgentCounts, agentId); + } + + private recordBufferedAmount(bufferedAmount: number | undefined): void { + if (typeof bufferedAmount !== "number") { + return; + } + this.bufferedAmountSamples.push(bufferedAmount); + } + + private computeLatencyStats(): WebSocketRuntimeMetricsSnapshot["latency"] { + const stats: WebSocketRuntimeMetricsSnapshot["latency"] = []; + for (const [type, latencies] of this.requestLatencies) { + if (latencies.length === 0) continue; + const sortedLatencies = [...latencies].sort((a, b) => a - b); + const count = sortedLatencies.length; + const minMs = Math.round(sortedLatencies[0]); + const maxMs = Math.round(sortedLatencies[count - 1]); + const p50Ms = Math.round(sortedLatencies[Math.floor(count / 2)]); + const totalMs = Math.round(sortedLatencies.reduce((sum, value) => sum + value, 0)); + stats.push({ type, count, minMs, maxMs, p50Ms, totalMs }); + } + stats.sort((a, b) => b.totalMs - a.totalMs); + return stats.slice(0, 15); + } + + private computeBufferedAmountStats(): WebSocketRuntimeMetricsSnapshot["bufferedAmount"] { + if (this.bufferedAmountSamples.length === 0) { + return { p95: 0, max: 0 }; + } + + const samples = [...this.bufferedAmountSamples].sort((a, b) => a - b); + const p95Index = Math.ceil(samples.length * 0.95) - 1; + return { + p95: samples[p95Index] ?? 0, + max: samples[samples.length - 1] ?? 0, + }; + } + + private reset(windowStartedAt: number): void { + for (const counter of Object.keys(this.counters) as Array) { + this.counters[counter] = 0; + } + this.inboundMessageCounts.clear(); + this.inboundSessionRequestCounts.clear(); + this.outboundMessageCounts.clear(); + this.outboundSessionMessageCounts.clear(); + this.outboundAgentStreamCounts.clear(); + this.outboundAgentStreamByAgentCounts.clear(); + this.outboundBinaryFrameCounts.clear(); + this.bufferedAmountSamples.length = 0; + this.requestLatencies.clear(); + this.windowStartedAt = windowStartedAt; + } +} + +function createRuntimeCounters(): WebSocketRuntimeCounters { + return { + connectedAwaitingHello: 0, + helloResumed: 0, + helloNew: 0, + pendingDisconnected: 0, + sessionDisconnectedWaitingReconnect: 0, + sessionSocketDisconnectedAttached: 0, + sessionCleanup: 0, + validationFailed: 0, + binaryBeforeHelloRejected: 0, + pendingMessageRejectedBeforeHello: 0, + missingConnectionForMessage: 0, + unexpectedHelloOnActiveConnection: 0, + relayExternalSocketAttached: 0, + originRejected: 0, + hostRejected: 0, + }; +} + +function incrementCount(map: Map, key: string): void { + map.set(key, (map.get(key) ?? 0) + 1); +} + +function getTopCounts(map: Map, limit: number): Array<[string, number]> { + return [...map.entries()].sort((a, b) => b[1] - a[1]).slice(0, limit); +}