From 3b5e34fe127bfc86d5ee30fea9f93724bf8f49e9 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 25 Jun 2026 19:16:54 +0800 Subject: [PATCH] Keep liveness pings responsive during slow requests (#1723) --- .../websocket-server.relay-reconnect.test.ts | 196 ++++++++++++++++-- .../server/src/server/websocket-server.ts | 47 ++--- 2 files changed, 201 insertions(+), 42 deletions(-) diff --git a/packages/server/src/server/websocket-server.relay-reconnect.test.ts b/packages/server/src/server/websocket-server.relay-reconnect.test.ts index a1ba6c53d..d4a456ada 100644 --- a/packages/server/src/server/websocket-server.relay-reconnect.test.ts +++ b/packages/server/src/server/websocket-server.relay-reconnect.test.ts @@ -105,7 +105,6 @@ import type { SpeechReadinessSnapshot } from "./speech/speech-runtime.js"; interface WebSocketServerInternals { attachSocket(ws: unknown, req: unknown): Promise; - socketMessageQueues: Map>; } const TEST_DAEMON_VERSION = "1.2.3-test"; @@ -398,7 +397,6 @@ async function attachRelayAndHello(params: { }) { await params.server.attachExternalSocket(params.socket, { transport: "relay" }); params.socket.emit("message", JSON.stringify(createHelloMessage(params.clientId))); - await waitForSocketMessages(params.server, params.socket); expect(params.socket.sent.length).toBeGreaterThan(0); const envelope = parseSentEnvelope(params.socket.sent[0]); expect(envelope.type).toBe("session"); @@ -418,7 +416,6 @@ async function attachDirectAndHello(params: { createDirectRequest(), ); params.socket.emit("message", JSON.stringify(createHelloMessage(params.clientId))); - await waitForSocketMessages(params.server, params.socket); expect(params.socket.sent.length).toBeGreaterThan(0); const envelope = parseSentEnvelope(params.socket.sent[0]); expect(envelope.type).toBe("session"); @@ -428,11 +425,19 @@ async function attachDirectAndHello(params: { return serverInfo!; } -async function waitForSocketMessages( - server: VoiceAssistantWebSocketServer, - socket: MockSocket, -): Promise { - await asInternals(server).socketMessageQueues.get(socket); +function holdNextSessionMessage(session: (typeof sessionMock.instances)[number]): { + finish: () => void; +} { + let finish = () => {}; + session.handleMessage.mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + return { + finish: () => finish(), + }; } describe("relay external socket reconnect behavior", () => { @@ -489,8 +494,6 @@ describe("relay external socket reconnect behavior", () => { }), ), ); - await waitForSocketMessages(server, socket); - expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]; expect(session.args.clientCapabilities).toEqual({ @@ -586,8 +589,6 @@ describe("relay external socket reconnect behavior", () => { }, }), ); - await waitForSocketMessages(server, socket); - expect(closeCode).toBe(4002); expect(["Invalid hello", "Session message before hello"]).toContain(closeReason); expect(sessionMock.instances).toHaveLength(0); @@ -595,6 +596,134 @@ describe("relay external socket reconnect behavior", () => { await server.close(); }); + test("responds to top-level ping while provider diagnostic is still running", async () => { + const server = createServer(); + const socket = new MockSocket(); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-ping-during-provider-diagnostic", + }); + + const session = sessionMock.instances[0]; + const providerDiagnostic = holdNextSessionMessage(session); + + const sentBeforeDiagnostic = socket.sent.length; + socket.emit( + "message", + JSON.stringify({ + type: "session", + message: { + type: "provider_diagnostic_request", + provider: "grok", + requestId: "slow-provider-diagnostic", + }, + }), + ); + await vi.waitFor(() => { + expect(session.handleMessage).toHaveBeenCalledTimes(1); + }); + + socket.emit("message", JSON.stringify({ type: "ping" })); + await Promise.resolve(); + + expect(sentEnvelopes(socket).slice(sentBeforeDiagnostic)).toContainEqual({ type: "pong" }); + + providerDiagnostic.finish(); + await Promise.resolve(); + await server.close(); + }); + + test("routes later session requests while provider diagnostic is still running", async () => { + const server = createServer(); + const socket = new MockSocket(); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-session-request-during-provider-diagnostic", + }); + + const session = sessionMock.instances[0]; + const providerDiagnostic = holdNextSessionMessage(session); + + socket.emit( + "message", + JSON.stringify({ + type: "session", + message: { + type: "provider_diagnostic_request", + provider: "grok", + requestId: "slow-provider-diagnostic", + }, + }), + ); + await vi.waitFor(() => { + expect(session.handleMessage).toHaveBeenCalledTimes(1); + }); + + socket.emit( + "message", + JSON.stringify({ + type: "session", + message: { + type: "ping", + requestId: "second-session-request", + clientSentAt: Date.now(), + }, + }), + ); + await vi.waitFor(() => { + expect(session.handleMessage).toHaveBeenCalledTimes(2); + }); + + providerDiagnostic.finish(); + await Promise.resolve(); + await server.close(); + }); + + test("sends rpc_error when an async session request fails", async () => { + const server = createServer(); + const socket = new MockSocket(); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-session-request-failure", + }); + + const session = sessionMock.instances[0]; + session.handleMessage.mockRejectedValueOnce(new Error("handler exploded")); + + const sentBeforeRequest = socket.sent.length; + socket.emit( + "message", + JSON.stringify({ + type: "session", + message: { + type: "provider_diagnostic_request", + provider: "grok", + requestId: "failing-provider-diagnostic", + }, + }), + ); + + await vi.waitFor(() => { + expect(sentEnvelopes(socket).slice(sentBeforeRequest)).toContainEqual({ + type: "session", + message: { + type: "rpc_error", + payload: { + requestId: "failing-provider-diagnostic", + requestType: "provider_diagnostic_request", + error: "Invalid message", + code: "invalid_message", + }, + }, + }); + }); + + await server.close(); + }); + test("reuses direct session when same clientId reconnects within grace window", async () => { const server = createServer(); const clientId = "cid-direct-reconnect"; @@ -793,8 +922,6 @@ describe("relay external socket reconnect behavior", () => { }), ), ); - await waitForSocketMessages(server, socket); - expect(session.handleBinaryFrame).toHaveBeenCalledTimes(1); const { frame } = BinaryFrameSchema.parse(session.handleBinaryFrame.mock.calls[0]?.[0]); expect(frame.opcode).toBe(TerminalStreamOpcode.Input); @@ -804,6 +931,47 @@ describe("relay external socket reconnect behavior", () => { await server.close(); }); + test("sends status error when async binary frame handling fails", async () => { + const server = createServer(); + + const socket = new MockSocket(); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-binary-inbound-failure", + }); + expect(sessionMock.instances).toHaveLength(1); + const session = sessionMock.instances[0]; + session.handleBinaryFrame.mockRejectedValueOnce(new Error("binary exploded")); + + const sentBeforeFrame = socket.sent.length; + socket.emit( + "message", + Buffer.from( + encodeTerminalStreamFrame({ + opcode: TerminalStreamOpcode.Input, + slot: 11, + payload: new TextEncoder().encode("pwd\r"), + }), + ), + ); + + await vi.waitFor(() => { + expect(sentEnvelopes(socket).slice(sentBeforeFrame)).toContainEqual({ + type: "session", + message: { + type: "status", + payload: { + status: "error", + message: "Invalid message: binary exploded", + }, + }, + }); + }); + + await server.close(); + }); + test("sends outbound terminal frames from session over websocket", async () => { const server = createServer(); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index c8b856e6f..4290ad288 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -350,7 +350,6 @@ export class VoiceAssistantWebSocketServer { private readonly pendingConnections: Map = new Map(); private readonly sessions: Map = new Map(); private readonly externalSessionsByKey: Map = new Map(); - private readonly socketMessageQueues: Map> = new Map(); private readonly serverId: string; private readonly daemonVersion: string; private readonly daemonRuntimeConfig: @@ -1209,7 +1208,7 @@ export class VoiceAssistantWebSocketServer { private bindSocketHandlers(ws: WebSocketLike): void { ws.on("message", (...args: unknown[]) => { const data = args[0] as Buffer | ArrayBuffer | Buffer[] | string; - this.enqueueRawMessage(ws, data); + this.handleRawMessage(ws, data); }); ws.on("close", async (...args: unknown[]) => { @@ -1232,25 +1231,6 @@ export class VoiceAssistantWebSocketServer { }); } - private enqueueRawMessage( - ws: WebSocketLike, - data: Buffer | ArrayBuffer | Buffer[] | string, - ): void { - const previous = this.socketMessageQueues.get(ws) ?? Promise.resolve(); - const next = previous.then( - () => this.handleRawMessage(ws, data), - () => this.handleRawMessage(ws, data), - ); - this.socketMessageQueues.set(ws, next); - void next - .catch(() => undefined) - .finally(() => { - if (this.socketMessageQueues.get(ws) === next) { - this.socketMessageQueues.delete(ws); - } - }); - } - public resolveVoiceSpeakHandler(callerAgentId: string): VoiceSpeakHandler | null { return this.voiceSpeakHandlers.get(callerAgentId) ?? null; } @@ -1430,12 +1410,12 @@ export class VoiceAssistantWebSocketServer { ); } - private async maybeHandleBinaryFrame(params: { + private maybeHandleBinaryFrame(params: { ws: WebSocketLike; buffer: Buffer; activeConnection: SessionConnection | undefined; log: pino.Logger; - }): Promise { + }): boolean { const { ws, buffer, activeConnection, log } = params; const asBytes = asUint8Array(buffer); if (!asBytes) { @@ -1456,7 +1436,16 @@ export class VoiceAssistantWebSocketServer { } return true; } - await activeConnection.session.handleBinaryFrame(decodedFrame); + void Promise.resolve(activeConnection.session.handleBinaryFrame(decodedFrame)).catch( + (error: unknown) => { + this.handleRawMessageError({ + ws, + data: buffer, + error, + log: activeConnection.connectionLogger, + }); + }, + ); return true; } @@ -1490,10 +1479,10 @@ export class VoiceAssistantWebSocketServer { } } - private async handleRawMessage( + private handleRawMessage( ws: WebSocketLike, data: Buffer | ArrayBuffer | Buffer[] | string, - ): Promise { + ): void { const activeConnection = this.sessions.get(ws); const pendingConnection = this.pendingConnections.get(ws); const log = @@ -1501,7 +1490,7 @@ export class VoiceAssistantWebSocketServer { try { const buffer = bufferFromWsData(data); - const binaryHandled = await this.maybeHandleBinaryFrame({ + const binaryHandled = this.maybeHandleBinaryFrame({ ws, buffer, activeConnection, @@ -1564,7 +1553,9 @@ export class VoiceAssistantWebSocketServer { } if (message.type === "session") { - await this.dispatchSessionMessage(activeConnection, message); + void this.dispatchSessionMessage(activeConnection, message).catch((error: unknown) => { + this.handleRawMessageError({ ws, data, error, log: activeConnection.connectionLogger }); + }); } } catch (error) { this.handleRawMessageError({ ws, data, error, log });