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
This commit is contained in:
Mohamed Boudra
2026-08-11 02:03:28 +08:00
committed by GitHub
parent ca2b633286
commit c5d5441c9c
3 changed files with 302 additions and 26 deletions
@@ -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<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",
});
@@ -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.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<void> {
while (this.hasPendingSegments()) {
private async waitForFlushDrain(finalSeq: number): Promise<void> {
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) => {
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<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 = [];
+1
View File
@@ -449,6 +449,7 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
attemptGuard.cancel();
stopDurationTracking();
void audioStop.current().catch(() => undefined);
senderRef.current?.dispose();
};
}, [stopDurationTracking]);