diff --git a/extensions/voice-call/src/media-stream.lifecycle.test.ts b/extensions/voice-call/src/media-stream.lifecycle.test.ts index 62ea111a32cf..a8a6c63553e7 100644 --- a/extensions/voice-call/src/media-stream.lifecycle.test.ts +++ b/extensions/voice-call/src/media-stream.lifecycle.test.ts @@ -70,4 +70,62 @@ describe("MediaStreamHandler lifecycle", () => { await server.close(); } }); + + it("terminates active streams and shares concurrent close completion", async () => { + const closeSession = vi.fn(); + const onConnect = vi.fn(); + const onDisconnect = vi.fn(); + const handler = new MediaStreamHandler({ + transcriptionProvider: { + createSession: () => ({ + connect: async () => {}, + sendAudio: () => {}, + close: closeSession, + isConnected: () => true, + }), + id: "openai", + label: "OpenAI", + isConfigured: () => true, + }, + providerConfig: {}, + shouldAcceptStream: () => true, + onConnect, + onDisconnect, + }); + const server = await startUpgradeWsServer({ + urlPath: "/voice/stream", + onUpgrade: (request, socket, head) => { + handler.handleUpgrade(request, socket, head); + }, + }); + const ws = await connectWs(server.url); + + try { + ws.send( + JSON.stringify({ + event: "start", + streamSid: "MZ-shutdown", + start: { callSid: "CA-shutdown" }, + }), + ); + await vi.waitFor(() => { + expect(onConnect).toHaveBeenCalledWith("CA-shutdown", "MZ-shutdown"); + }); + + const closed = waitForClose(ws); + const firstClose = handler.close(); + const secondClose = handler.close(); + + expect(secondClose).toBe(firstClose); + await firstClose; + expect(await closed).toEqual({ code: 1006, reason: "" }); + expect(closeSession).toHaveBeenCalledTimes(1); + expect(onDisconnect).toHaveBeenCalledWith("CA-shutdown", "MZ-shutdown"); + expect(onDisconnect).toHaveBeenCalledTimes(1); + } finally { + ws.terminate(); + await handler.close(); + await server.close(); + } + }); }); diff --git a/extensions/voice-call/src/media-stream.ts b/extensions/voice-call/src/media-stream.ts index b93d811b88c3..07abbbb5c360 100644 --- a/extensions/voice-call/src/media-stream.ts +++ b/extensions/voice-call/src/media-stream.ts @@ -138,6 +138,8 @@ function parseTwilioMediaMessage(data: RawData): TwilioMediaMessage { */ export class MediaStreamHandler { private wss: WebSocketServer | null = null; + private closePromise: Promise | null = null; + private closing = false; private sessions = new Map(); private config: MediaStreamConfig; /** Pending sockets that have upgraded but not yet sent an accepted `start` frame. */ @@ -172,6 +174,11 @@ export class MediaStreamHandler { * Handle WebSocket upgrade for media stream connections. */ handleUpgrade(request: IncomingMessage, socket: Duplex, head: Buffer): void { + if (this.closing) { + this.rejectUpgrade(socket, 503, "Media stream handler is shutting down"); + return; + } + if (!this.wss) { this.wss = new WebSocketServer({ noServer: true, @@ -221,6 +228,31 @@ export class MediaStreamHandler { } } + close(): Promise { + if (this.closePromise) { + return this.closePromise; + } + + this.closing = true; + const wss = this.wss; + this.wss = null; + this.closePromise = (async () => { + if (!wss) { + return; + } + await new Promise((resolve) => { + wss.close(() => resolve()); + for (const ws of wss.clients) { + ws.terminate(); + } + }); + })().finally(() => { + this.closing = false; + this.closePromise = null; + }); + return this.closePromise; + } + /** * Handle new WebSocket connection from Twilio. */