fix(streaming): stop close() hanging when no Termination frame arrives

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## [4.36.7]
+
+- Fix `StreamingTranscriber.close()` and `RealtimeTranscriber.close()` hanging forever when the socket closes without a `Termination` message — `onclose` now releases the pending wait, and the wait is bounded by a new optional `terminationTimeout` parameter on `close()` (5000ms default; `0` waits indefinitely). The socket is closed either way
+
 ## [4.36.6]
 
 - Add `effort` to the speech understanding feature requests (`SpeakerIdentificationRequest`, `TranslationRequest`, `CustomFormattingRequest`) — `"low"` (default) or `"medium"`, set per task, typed as the new `SpeechUnderstandingEffort`. The field was already accepted by the API but missing from the SDK types, so setting it failed to type check

package.json

@@ -1,6 +1,6 @@
 {
   "name": "assemblyai",
-  "version": "4.36.6",
+  "version": "4.36.7",
   "description": "The AssemblyAI JavaScript SDK provides an easy-to-use interface for interacting with the AssemblyAI API, which supports async and real-time transcription, as well as the latest LeMUR models.",
   "engines": {
     "node": ">=18"

src/services/realtime/service.ts

@@ -25,6 +25,12 @@ const defaultRealtimeUrl = "wss://api.assemblyai.com/v2/realtime/ws";
 const forceEndOfUtteranceMessage = `{"force_end_utterance":true}`;
 const terminateSessionMessage = `{"terminate_session":true}`;
 
+/**
+ * How long `close()` waits for the server's `SessionTerminated` message before
+ * giving up. The socket is closed either way, so this only bounds the wait.
+ */
+const DEFAULT_TERMINATION_TIMEOUT_MS = 5000;
+
 type BufferLike =
   | string
   | Buffer
@@ -222,6 +228,9 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
             reason = RealtimeErrorMessages[code as RealtimeErrorTypeCodes];
           }
         }
+        // The socket is gone, so no `SessionTerminated` message is coming.
+        // Release a `close()` that is waiting for one.
+        this.resolveSessionTermination();
         this.listeners.close?.(code, reason);
       };
 
@@ -265,7 +274,7 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
             break;
           }
           case "SessionTerminated": {
-            this.sessionTerminatedResolve?.();
+            this.resolveSessionTermination();
             break;
           }
         }
@@ -316,20 +325,65 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
     this.socket.send(data);
   }
 
+  /**
+   * Releases a `close()` that is waiting for the session to terminate.
+   *
+   * Called from the `SessionTerminated` message and from `onclose`. Without the
+   * `onclose` call, a socket that dies without that message leaves `close()`
+   * awaiting forever.
+   */
+  private resolveSessionTermination(): void {
+    const resolve = this.sessionTerminatedResolve;
+    this.sessionTerminatedResolve = undefined;
+    resolve?.();
+  }
+
+  /** Awaits the session-termination message, bounded by `timeoutMs`. */
+  private async waitForTermination(
+    sessionTerminated: Promise<void>,
+    timeoutMs: number,
+  ): Promise<void> {
+    if (timeoutMs <= 0) {
+      await sessionTerminated;
+      return;
+    }
+
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    try {
+      await Promise.race([
+        sessionTerminated,
+        new Promise<void>((resolve) => {
+          timer = setTimeout(resolve, timeoutMs);
+        }),
+      ]);
+    } finally {
+      if (timer) clearTimeout(timer);
+      this.sessionTerminatedResolve = undefined;
+    }
+  }
+
   /**
    * Close the connection to the server.
    * @param waitForSessionTermination - If true, the method will wait for the session to be terminated before closing the connection.
    * While waiting for the session to be terminated, you will receive the final transcript and session information.
+   * @param terminationTimeout - How long to wait for that message, in
+   * milliseconds. `0` waits indefinitely. The socket closes either way.
    */
-  async close(waitForSessionTermination = true) {
+  async close(
+    waitForSessionTermination = true,
+    terminationTimeout = DEFAULT_TERMINATION_TIMEOUT_MS,
+  ) {
     if (this.socket) {
       if (this.socket.readyState === this.socket.OPEN) {
         if (waitForSessionTermination) {
           const sessionTerminatedPromise = new Promise<void>((resolve) => {
             this.sessionTerminatedResolve = resolve;
           });
           this.socket.send(terminateSessionMessage);
-          await sessionTerminatedPromise;
+          await this.waitForTermination(
+            sessionTerminatedPromise,
+            terminationTimeout,
+          );
         } else {
           this.socket.send(terminateSessionMessage);
         }

src/services/streaming/service.ts

@@ -65,6 +65,14 @@ const defaultStreamingUrl = "wss://streaming.assemblyai.com/v3/ws";
 const terminateSessionMessage = `{"type":"Terminate"}`;
 
 const DEFAULT_CONNECT_TIMEOUT_MS = 1000;
+
+/**
+ * How long `close()` waits for the server's `Termination` frame before giving
+ * up. The socket is closed either way, so this only bounds the wait. A server
+ * that drops the connection without replying resolves the wait immediately
+ * through `onclose`; this timeout covers a server that goes silent instead.
+ */
+const DEFAULT_TERMINATION_TIMEOUT_MS = 5000;
 const DEFAULT_MAX_CONNECTION_RETRIES = 2;
 const DEFAULT_CONNECTION_RETRY_DELAY_MS = 500;
 
@@ -233,6 +241,19 @@ export class StreamingTranscriber {
     }
   }
 
+  /**
+   * Releases a `close()` that is waiting for the session to terminate.
+   *
+   * Called from the `Termination` frame and from `onclose`. Without the
+   * `onclose` call, a socket that dies without a `Termination` frame leaves
+   * `close()` awaiting forever.
+   */
+  private resolveSessionTermination(): void {
+    const resolve = this.sessionTerminatedResolve;
+    this.sessionTerminatedResolve = undefined;
+    resolve?.();
+  }
+
   private connectionUrl(): URL {
     const url = new URL(this.params.websocketBaseUrl ?? "");
 
@@ -625,6 +646,9 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
           clearInterval(this.flushTimer);
           this.flushTimer = undefined;
         }
+        // The socket is gone, so no `Termination` frame is coming. Release a
+        // `close()` that is waiting for one.
+        this.resolveSessionTermination();
         this.listeners.close?.(code, reason);
       };
 
@@ -713,7 +737,7 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
             break;
           }
           case "Termination": {
-            this.sessionTerminatedResolve?.();
+            this.resolveSessionTermination();
             break;
           }
         }
@@ -1025,14 +1049,49 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
     this.send(JSON.stringify(message));
   }
 
+  /** Awaits the session-termination frame, bounded by `timeoutMs`. */
+  private async waitForTermination(
+    sessionTerminated: Promise<void>,
+    timeoutMs: number,
+  ): Promise<void> {
+    if (timeoutMs <= 0) {
+      await sessionTerminated;
+      return;
+    }
+
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    try {
+      await Promise.race([
+        sessionTerminated,
+        new Promise<void>((resolve) => {
+          timer = setTimeout(resolve, timeoutMs);
+        }),
+      ]);
+    } finally {
+      if (timer) clearTimeout(timer);
+      this.sessionTerminatedResolve = undefined;
+    }
+  }
+
   private send(data: BufferLike) {
     if (!this.socket || this.socket.readyState !== this.socket.OPEN) {
       throw new Error("Socket is not open for communication");
     }
     this.socket.send(data);
   }
 
-  async close(waitForSessionTermination = true) {
+  /**
+   * Close the connection to the server.
+   * @param waitForSessionTermination - Wait for the server's `Termination`
+   * frame before closing the socket. While waiting you still receive the final
+   * transcript and the session information.
+   * @param terminationTimeout - How long to wait for that frame, in
+   * milliseconds. `0` waits indefinitely. The socket closes either way.
+   */
+  async close(
+    waitForSessionTermination = true,
+    terminationTimeout = DEFAULT_TERMINATION_TIMEOUT_MS,
+  ) {
     if (this.flushTimer) {
       clearInterval(this.flushTimer);
       this.flushTimer = undefined;
@@ -1049,7 +1108,10 @@ Learn more at https://github.com/AssemblyAI/assemblyai-node-sdk/blob/main/docs/c
             this.sessionTerminatedResolve = resolve;
           });
           this.socket.send(terminateSessionMessage);
-          await sessionTerminatedPromise;
+          await this.waitForTermination(
+            sessionTerminatedPromise,
+            terminationTimeout,
+          );
         } else {
           this.socket.send(terminateSessionMessage);
         }

tests/unit/streaming-close.test.ts

@@ -0,0 +1,89 @@
+jest.mock("ws", () => require("./mocks/ws"));
+
+import WS from "jest-websocket-mock";
+import fetchMock from "jest-fetch-mock";
+import { AssemblyAI, StreamingTranscriber } from "../../src";
+import { createClient } from "./utils";
+
+fetchMock.enableMocks();
+
+const websocketBaseUrl = "wss://localhost:1234/v3/ws";
+const sessionBeginsMessage = {
+  type: "Begin",
+  id: "123",
+  expires_at: 123456789,
+};
+
+let server: WS;
+let aai: AssemblyAI;
+let rt: StreamingTranscriber;
+
+// These tests deliberately avoid the shared `close()` helper used elsewhere,
+// which always sends a `Termination` frame. The point is what happens when the
+// server never sends one.
+describe("streaming close", () => {
+  beforeEach(async () => {
+    server = new WS(websocketBaseUrl);
+    aai = createClient();
+    rt = aai.streaming.transcriber({
+      websocketBaseUrl: websocketBaseUrl,
+      apiKey: "123",
+      sampleRate: 16_000,
+      speechModel: "universal-streaming-english",
+    });
+    const connectPromise = rt.connect();
+    await server.connected;
+    server.send(JSON.stringify(sessionBeginsMessage));
+    await connectPromise;
+  });
+
+  afterEach(() => {
+    WS.clean();
+  });
+
+  it("resolves when the server closes without a Termination frame", async () => {
+    const closePromise = rt.close();
+    // The server acknowledges the terminate by dropping the connection, which
+    // is what a crashing or impatient server does.
+    server.close({ code: 1000, reason: "", wasClean: true });
+
+    await expect(closePromise).resolves.toBeUndefined();
+  });
+
+  it("resolves when the server goes silent, bounded by the timeout", async () => {
+    const started = Date.now();
+    // The server neither replies nor closes. Only the timeout can end this.
+    await rt.close(true, 150);
+
+    expect(Date.now() - started).toBeGreaterThanOrEqual(140);
+    expect(Date.now() - started).toBeLessThan(2000);
+  });
+
+  it("resolves as soon as the Termination frame arrives", async () => {
+    const started = Date.now();
+    const closePromise = rt.close(true, 10_000);
+    server.send(JSON.stringify({ type: "Termination" }));
+    await closePromise;
+
+    // Well under the timeout, so the frame ended the wait rather than the timer.
+    expect(Date.now() - started).toBeLessThan(1000);
+  });
+
+  it("does not wait when waitForSessionTermination is false", async () => {
+    const started = Date.now();
+    await rt.close(false);
+
+    expect(Date.now() - started).toBeLessThan(1000);
+  });
+
+  it("resolves when the socket is already dead", async () => {
+    server.close({
+      code: 4031,
+      reason: "Session idle for too long",
+      wasClean: false,
+    });
+    await new Promise((resolve) => setTimeout(resolve, 20));
+
+    await expect(rt.close()).resolves.toBeUndefined();
+  });
+});