Preserve dictation recordings across WebSocket interruptions (#3159)

* fix(app): preserve dictation audio across transport failures

Keep captured PCM in the existing local segment buffer when a WebSocket send fails. Finish only after server acknowledgements and fail an idle drain so the retry path can resend the complete recording.

* fix(app): dispose dictation stream sender cleanly

packages/app/src/dictation/dictation-stream-sender.test.ts

@@ -1,6 +1,5 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
-import type { DaemonClient } from "@getpaseo/client/internal/daemon-client";
+import type { SessionOutboundMessage } from "@getpaseo/protocol/messages";
-
 import { DictationStreamSender } from "@/dictation/dictation-stream-sender";
 
 interface FakeFinish {
@@ -18,19 +17,66 @@ interface FakeChunk {
   format: string;
 }
 
+interface FakeAck {
+  type: "dictation_stream_ack";
+  payload: {
+    dictationId: string;
+    ackSeq: number;
+  };
+}
+
+type FakeRawMessageListener = (message: SessionOutboundMessage) => void;
+
 class FakeDaemonClient {
   isConnected = true;
+  autoAck = true;
+  throwOnNextChunk = false;
+  disconnectOnChunkError = false;
   starts: FakeStart[] = [];
   chunks: FakeChunk[] = [];
   finishes: FakeFinish[] = [];
   cancels: string[] = [];
+  private readonly rawMessageListeners = new Set<FakeRawMessageListener>();
+
+  get rawMessageListenerCount(): number {
+    return this.rawMessageListeners.size;
+  }
 
   async startDictationStream(dictationId: string, format: string): Promise<void> {
     this.starts.push({ dictationId, format });
   }
 
   sendDictationStreamChunk(dictationId: string, seq: number, audio: string, format: string): void {
+    if (this.throwOnNextChunk) {
+      this.throwOnNextChunk = false;
+      if (this.disconnectOnChunkError) {
+        this.isConnected = false;
+      }
+      throw new Error("WebSocket is closing");
+    }
     this.chunks.push({ dictationId, seq, audio, format });
+    if (this.autoAck) {
+      this.emitAck(dictationId, seq);
+    }
+  }
+
+  subscribeRawMessages(handler: FakeRawMessageListener): () => void {
+    this.rawMessageListeners.add(handler);
+    return () => this.rawMessageListeners.delete(handler);
+  }
+
+  setConnected(isConnected: boolean): void {
+    this.isConnected = isConnected;
+  }
+
+  emitAck(dictationId: string, ackSeq: number): void {
+    const message: FakeAck = {
+      type: "dictation_stream_ack",
+      payload: { dictationId, ackSeq },
+    };
+    for (const listener of this.rawMessageListeners) {
+      listener(message);
+    }
   }
 
   async finishDictationStream(
@@ -51,12 +97,22 @@ const tick = async () => {
   await Promise.resolve();
 };
 
+const promiseOutcome = async <T>(promise: Promise<T>, timeoutMs: number) => {
+  return Promise.race([
+    promise.then(
+      () => "resolved" as const,
+      () => "rejected" as const,
+    ),
+    new Promise<"timed-out">((resolve) => setTimeout(() => resolve("timed-out"), timeoutMs)),
+  ]);
+};
+
 describe("DictationStreamSender", () => {
   it("enqueues segments and sends them after stream start", async () => {
     const client = new FakeDaemonClient();
     const ids = ["d1"];
     const sender = new DictationStreamSender({
-      client: client as unknown as DaemonClient,
+      client,
       format: "audio/pcm;rate=16000;bits=16",
       createDictationId: () => ids.shift() ?? "dX",
     });
@@ -77,7 +133,7 @@ describe("DictationStreamSender", () => {
     const client = new FakeDaemonClient();
     const ids = ["d1", "d2"];
     const sender = new DictationStreamSender({
-      client: client as unknown as DaemonClient,
+      client,
       format: "audio/pcm;rate=16000;bits=16",
       createDictationId: () => ids.shift() ?? "dX",
     });
@@ -100,7 +156,7 @@ describe("DictationStreamSender", () => {
     const client = new FakeDaemonClient();
     const ids = ["d1"];
     const sender = new DictationStreamSender({
-      client: client as unknown as DaemonClient,
+      client,
       format: "audio/pcm;rate=16000;bits=16",
       createDictationId: () => ids.shift() ?? "dX",
     });
@@ -121,7 +177,7 @@ describe("DictationStreamSender", () => {
     client.isConnected = false;
     const ids = ["d1"];
     const sender = new DictationStreamSender({
-      client: client as unknown as DaemonClient,
+      client,
       format: "audio/pcm;rate=16000;bits=16",
       createDictationId: () => ids.shift() ?? "dX",
     });
@@ -138,11 +194,131 @@ describe("DictationStreamSender", () => {
     expect(client.chunks.map((c) => c.seq)).toEqual([0, 1]);
   });
 
+  it("keeps capture safe and buffers later segments when a chunk send throws", async () => {
+    const client = new FakeDaemonClient();
+    const sender = new DictationStreamSender({
+      client,
+      format: "audio/pcm;rate=16000;bits=16",
+      createDictationId: () => "d1",
+    });
+
+    sender.enqueueSegment("before-outage");
+    await tick();
+
+    client.throwOnNextChunk = true;
+    client.disconnectOnChunkError = true;
+
+    expect(() => sender.enqueueSegment("during-outage")).not.toThrow();
+    expect(() => sender.enqueueSegment("after-outage")).not.toThrow();
+    expect(sender.getSegmentCount()).toBe(3);
+  });
+
+  it("removes its raw-message listener when disposed", () => {
+    const client = new FakeDaemonClient();
+    const sender = new DictationStreamSender({
+      client,
+      format: "audio/pcm;rate=16000;bits=16",
+    });
+
+    expect(client.rawMessageListenerCount).toBe(1);
+    sender.dispose();
+    expect(client.rawMessageListenerCount).toBe(0);
+  });
+
+  it("replays segments captured before, during, and after an outage from seq 0", async () => {
+    const client = new FakeDaemonClient();
+    const ids = ["d1", "d2"];
+    const sender = new DictationStreamSender({
+      client,
+      format: "audio/pcm;rate=16000;bits=16",
+      createDictationId: () => ids.shift() ?? "dX",
+    });
+
+    sender.enqueueSegment("before-outage");
+    await tick();
+
+    client.throwOnNextChunk = true;
+    client.disconnectOnChunkError = true;
+    expect(() => sender.enqueueSegment("during-outage")).not.toThrow();
+    sender.enqueueSegment("after-outage");
+
+    client.setConnected(true);
+    await sender.restartStream("reconnect");
+
+    const replay = client.chunks
+      .filter((chunk) => chunk.dictationId === "d2")
+      .map((chunk) => [chunk.seq, chunk.audio]);
+    expect(replay).toEqual([
+      [0, "before-outage"],
+      [1, "during-outage"],
+      [2, "after-outage"],
+    ]);
+  });
+
+  it("fails finish after a transport transition without losing buffered audio", async () => {
+    vi.useFakeTimers();
+    try {
+      const client = new FakeDaemonClient();
+      client.setConnected(false);
+      const ids = ["d1", "d2"];
+      const sender = new DictationStreamSender({
+        client,
+        format: "audio/pcm;rate=16000;bits=16",
+        createDictationId: () => ids.shift() ?? "dX",
+      });
+
+      for (let seq = 0; seq < 129; seq += 1) {
+        sender.enqueueSegment(`segment-${seq}`);
+      }
+
+      client.setConnected(true);
+      const finish = sender.finish(sender.getFinalSeq());
+      const outcome = promiseOutcome(finish, 1);
+
+      await tick();
+      client.setConnected(false);
+      await vi.advanceTimersByTimeAsync(1);
+
+      await expect(outcome).resolves.toBe("rejected");
+      expect(sender.getSegmentCount()).toBe(129);
+      expect(sender.getFinalSeq()).toBe(128);
+
+      client.setConnected(true);
+      sender.resetStreamForReplay();
+      await sender.restartStream("retry");
+      await expect(sender.finish(128)).resolves.toEqual({ dictationId: "d2", text: "ok" });
+      expect(client.chunks.filter((chunk) => chunk.dictationId === "d2")).toHaveLength(129);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("waits for server delivery acknowledgement before finishing", async () => {
+    const client = new FakeDaemonClient();
+    client.autoAck = false;
+    const sender = new DictationStreamSender({
+      client,
+      format: "audio/pcm;rate=16000;bits=16",
+      createDictationId: () => "d1",
+    });
+
+    sender.enqueueSegment("seg0");
+    sender.enqueueSegment("seg1");
+    await tick();
+
+    const finish = sender.finish(1);
+    await tick();
+    expect(client.finishes).toEqual([]);
+
+    client.emitAck("d1", 1);
+    await expect(finish).resolves.toEqual({ dictationId: "d1", text: "ok" });
+  });
+
   it("does not replay long buffered native dictation in one synchronous burst", async () => {
     const client = new FakeDaemonClient();
     client.isConnected = false;
     const sender = new DictationStreamSender({
-      client: client as unknown as DaemonClient,
+      client,
       format: "audio/pcm;rate=16000;bits=16",
       createDictationId: () => "d1",
     });

packages/app/src/dictation/dictation-stream-sender.ts

@@ -1,13 +1,32 @@
 import { generateMessageId } from "@/types/stream";
-import type { DaemonClient } from "@getpaseo/client/internal/daemon-client";
+import type { SessionOutboundMessage } from "@getpaseo/protocol/messages";
 import { i18n } from "@/i18n/i18next";
 
 const MAX_CHUNKS_PER_FLUSH_TURN = 128;
+const DICTATION_DRAIN_IDLE_TIMEOUT_MS = 30_000;
 
 const waitForNextFlushTurn = (): Promise<void> => new Promise((resolve) => setTimeout(resolve, 0));
 
+type DictationStreamAckMessage = Extract<SessionOutboundMessage, { type: "dictation_stream_ack" }>;
+
+interface DrainWaitResult {
+  kind: "signaled" | "timed-out";
+}
+
+export interface DictationStreamClient {
+  readonly isConnected: boolean;
+  startDictationStream(dictationId: string, format: string): Promise<void>;
+  sendDictationStreamChunk(dictationId: string, seq: number, audio: string, format: string): void;
+  finishDictationStream(
+    dictationId: string,
+    finalSeq: number,
+  ): Promise<{ dictationId: string; text: string }>;
+  cancelDictationStream(dictationId: string): void;
+  subscribeRawMessages(handler: (message: SessionOutboundMessage) => void): () => void;
+}
+
 export interface DictationStreamSenderParams {
-  client: DaemonClient | null;
+  client: DictationStreamClient | null;
   format: string;
   createDictationId?: () => string;
 }
@@ -30,28 +49,50 @@ interface DictationFinishResult {
  * so enqueues can't "miss" a flush due to in-flight await/coalescing bugs.
  */
 export class DictationStreamSender {
-  private client: DaemonClient | null;
+  private client: DictationStreamClient | null = null;
   private readonly format: string;
   private readonly createDictationId: () => string;
 
   private dictationId: string | null = null;
   private sendSeq = 0;
+  private acknowledgedSeq = -1;
   private segments: string[] = [];
   private streamReady = false;
   private flushTimer: ReturnType<typeof setTimeout> | null = null;
   private drainWaiters: Array<() => void> = [];
+  private clientCleanup: (() => void) | null = null;
 
   private startGeneration = 0;
   private startPromise: Promise<void> | null = null;
 
   constructor(params: DictationStreamSenderParams) {
-    this.client = params.client;
     this.format = params.format;
     this.createDictationId = params.createDictationId ?? generateMessageId;
+    this.setClient(params.client);
   }
 
-  setClient(client: DaemonClient | null): void {
+  setClient(client: DictationStreamClient | null): void {
+    if (this.client === client) {
+      return;
+    }
+
+    if (this.clientCleanup) {
+      this.clientCleanup();
+      this.clientCleanup = null;
+    }
+    this.resetStreamForReplay();
     this.client = client;
+
+    if (!client) {
+      return;
+    }
+
+    this.clientCleanup = client.subscribeRawMessages((message) => {
+      if (message.type !== "dictation_stream_ack") {
+        return;
+      }
+      this.handleAck(message);
+    });
   }
 
   getDictationId(): string | null {
@@ -62,6 +103,10 @@ export class DictationStreamSender {
     return this.segments.length;
   }
 
+  dispose(): void {
+    this.setClient(null);
+  }
+
   getFinalSeq(): number {
     return this.segments.length - 1;
   }
@@ -71,22 +116,19 @@ export class DictationStreamSender {
   }
 
   clearAll(): void {
-    this.clearScheduledFlush();
+    this.resetStreamForReplay();
-    this.dictationId = null;
-    this.sendSeq = 0;
     this.segments = [];
-    this.streamReady = false;
-    this.startPromise = null;
-    this.startGeneration += 1;
   }
 
   resetStreamForReplay(): void {
     this.clearScheduledFlush();
     this.dictationId = null;
     this.sendSeq = 0;
+    this.acknowledgedSeq = -1;
     this.streamReady = false;
     this.startPromise = null;
     this.startGeneration += 1;
+    this.resolveDrainWaiters();
   }
 
   enqueueSegment(base64Pcm: string): void {
@@ -113,6 +155,9 @@ export class DictationStreamSender {
     const client = this.client;
     const dictationId = this.dictationId;
     if (!client?.isConnected || !dictationId || !this.streamReady) {
+      if (!client?.isConnected && this.hasPendingSegments()) {
+        this.invalidateStream();
+      }
       return 0;
     }
 
@@ -120,10 +165,16 @@ export class DictationStreamSender {
     while (this.sendSeq < this.segments.length && sent < MAX_CHUNKS_PER_FLUSH_TURN) {
       const seq = this.sendSeq;
       const audio = this.segments[seq];
-      client.sendDictationStreamChunk(dictationId, seq, audio, this.format);
+      try {
+        client.sendDictationStreamChunk(dictationId, seq, audio, this.format);
+      } catch {
+        this.invalidateStream();
+        return sent;
+      }
       this.sendSeq = seq + 1;
       sent += 1;
     }
+
     if (this.hasPendingSegments()) {
       this.scheduleFlush();
     } else {
@@ -144,6 +195,7 @@ export class DictationStreamSender {
     const dictationId = this.createDictationId();
     this.dictationId = dictationId;
     this.sendSeq = 0;
+    this.acknowledgedSeq = -1;
     this.streamReady = false;
 
     const start = (async () => {
@@ -161,7 +213,10 @@ export class DictationStreamSender {
         // If starting failed, keep the segments for retry but clear the stream so finish can error cleanly.
         if (this.startGeneration === generation && this.dictationId === dictationId) {
           this.dictationId = null;
+          this.sendSeq = 0;
+          this.acknowledgedSeq = -1;
           this.streamReady = false;
+          this.resolveDrainWaiters();
         }
         throw error;
       })
@@ -198,7 +253,7 @@ export class DictationStreamSender {
     }
 
     this.flush();
-    await this.waitForFlushDrain();
+    await this.waitForFlushDrain(finalSeq);
     return client.finishDictationStream(dictationId, finalSeq);
   }
 
@@ -215,6 +270,25 @@ export class DictationStreamSender {
     return this.sendSeq < this.segments.length;
   }
 
+  private hasPendingDelivery(finalSeq: number): boolean {
+    return this.hasPendingSegments() || this.acknowledgedSeq < finalSeq;
+  }
+
+  private handleAck(message: DictationStreamAckMessage): void {
+    if (message.payload.dictationId !== this.dictationId) {
+      return;
+    }
+    if (message.payload.ackSeq <= this.acknowledgedSeq) {
+      return;
+    }
+    this.acknowledgedSeq = message.payload.ackSeq;
+    this.resolveDrainWaiters();
+  }
+
+  private invalidateStream(): void {
+    this.resetStreamForReplay();
+  }
+
   private scheduleFlush(): void {
     if (this.flushTimer) {
       return;
@@ -233,19 +307,44 @@ export class DictationStreamSender {
     this.flushTimer = null;
   }
 
-  private async waitForFlushDrain(): Promise<void> {
+  private async waitForFlushDrain(finalSeq: number): Promise<void> {
-    while (this.hasPendingSegments()) {
+    while (this.hasPendingDelivery(finalSeq)) {
       const client = this.client;
       if (!client?.isConnected || !this.dictationId || !this.streamReady) {
         throw new Error("Failed to flush dictation stream");
       }
-      await new Promise<void>((resolve) => {
+      const result = await this.waitForDrainSignal();
-        this.drainWaiters.push(resolve);
+      if (result.kind === "timed-out") {
-      });
+        this.invalidateStream();
+        throw new Error("Timed out waiting for dictation stream delivery");
+      }
       await waitForNextFlushTurn();
     }
   }
 
+  private waitForDrainSignal(): Promise<DrainWaitResult> {
+    let timeoutHandle: ReturnType<typeof setTimeout> | null = null;
+    let waiter: (() => void) | null = null;
+    const signal = new Promise<"signaled">((resolve) => {
+      waiter = () => resolve("signaled");
+      this.drainWaiters.push(waiter);
+    });
+    const timeout = new Promise<"timed-out">((resolve) => {
+      timeoutHandle = setTimeout(() => resolve("timed-out"), DICTATION_DRAIN_IDLE_TIMEOUT_MS);
+    });
+
+    return Promise.race([signal, timeout])
+      .finally(() => {
+        if (timeoutHandle) {
+          clearTimeout(timeoutHandle);
+        }
+        if (waiter) {
+          this.drainWaiters = this.drainWaiters.filter((candidate) => candidate !== waiter);
+        }
+      })
+      .then((kind) => ({ kind }));
+  }
+
   private resolveDrainWaiters(): void {
     const waiters = this.drainWaiters;
     this.drainWaiters = [];

packages/app/src/hooks/use-dictation.ts

@@ -449,6 +449,7 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
       attemptGuard.cancel();
       stopDurationTracking();
       void audioStop.current().catch(() => undefined);
+      senderRef.current?.dispose();
     };
   }, [stopDurationTracking]);