Keep liveness pings responsive during slow requests (#1723)

This commit is contained in:
Mohamed Boudra
2026-06-25 19:16:54 +08:00
committed by GitHub
parent 9960772fe9
commit 3b5e34fe12
2 changed files with 201 additions and 42 deletions
@@ -105,7 +105,6 @@ import type { SpeechReadinessSnapshot } from "./speech/speech-runtime.js";
interface WebSocketServerInternals {
attachSocket(ws: unknown, req: unknown): Promise<void>;
socketMessageQueues: Map<unknown, Promise<void>>;
}
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<void> {
await asInternals<WebSocketServerInternals>(server).socketMessageQueues.get(socket);
function holdNextSessionMessage(session: (typeof sessionMock.instances)[number]): {
finish: () => void;
} {
let finish = () => {};
session.handleMessage.mockImplementationOnce(
() =>
new Promise<void>((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();
+19 -28
View File
@@ -350,7 +350,6 @@ export class VoiceAssistantWebSocketServer {
private readonly pendingConnections: Map<WebSocketLike, PendingConnection> = new Map();
private readonly sessions: Map<WebSocketLike, SessionConnection> = new Map();
private readonly externalSessionsByKey: Map<string, SessionConnection> = new Map();
private readonly socketMessageQueues: Map<WebSocketLike, Promise<void>> = 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> {
}): 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> {
): 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 });