From c5d5441c9c1c0cc6572d88212da74495acab400a Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 10 Aug 2026 20:03:28 +0200 Subject: [PATCH] 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 --- .../dictation/dictation-stream-sender.test.ts | 192 +++++++++++++++++- .../src/dictation/dictation-stream-sender.ts | 135 ++++++++++-- packages/app/src/hooks/use-dictation.ts | 1 + 3 files changed, 302 insertions(+), 26 deletions(-) diff --git a/packages/app/src/dictation/dictation-stream-sender.test.ts b/packages/app/src/dictation/dictation-stream-sender.test.ts index 0684dd2b1..1af06997a 100644 --- a/packages/app/src/dictation/dictation-stream-sender.test.ts +++ b/packages/app/src/dictation/dictation-stream-sender.test.ts @@ -1,6 +1,5 @@ -import { describe, expect, it } from "vitest"; -import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; - +import { describe, expect, it, vi } from "vitest"; +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(); + + get rawMessageListenerCount(): number { + return this.rawMessageListeners.size; + } async startDictationStream(dictationId: string, format: string): Promise { 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 (promise: Promise, 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", }); diff --git a/packages/app/src/dictation/dictation-stream-sender.ts b/packages/app/src/dictation/dictation-stream-sender.ts index e4562e0b9..238845884 100644 --- a/packages/app/src/dictation/dictation-stream-sender.ts +++ b/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 => new Promise((resolve) => setTimeout(resolve, 0)); +type DictationStreamAckMessage = Extract; + +interface DrainWaitResult { + kind: "signaled" | "timed-out"; +} + +export interface DictationStreamClient { + readonly isConnected: boolean; + startDictationStream(dictationId: string, format: string): Promise; + 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 | null = null; private drainWaiters: Array<() => void> = []; + private clientCleanup: (() => void) | null = null; private startGeneration = 0; private startPromise: Promise | 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.dictationId = null; - this.sendSeq = 0; + this.resetStreamForReplay(); 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 { - while (this.hasPendingSegments()) { + private async waitForFlushDrain(finalSeq: number): Promise { + 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((resolve) => { - this.drainWaiters.push(resolve); - }); + const result = await this.waitForDrainSignal(); + if (result.kind === "timed-out") { + this.invalidateStream(); + throw new Error("Timed out waiting for dictation stream delivery"); + } await waitForNextFlushTurn(); } } + private waitForDrainSignal(): Promise { + let timeoutHandle: ReturnType | 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 = []; diff --git a/packages/app/src/hooks/use-dictation.ts b/packages/app/src/hooks/use-dictation.ts index 09442b4c7..dde72964b 100644 --- a/packages/app/src/hooks/use-dictation.ts +++ b/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]);