feat(dictation): add resilient rpc retries and status toasts

2025-12-02 08:12:02 +00:00
parent 71b238c0cd
commit dc26bf713e
8 changed files with 1213 additions and 111 deletions
@@ -3,6 +3,21 @@ import type { SessionInboundMessage, SessionOutboundMessage } from "@server/serv
 import type { UseWebSocketReturn } from "./use-websocket";
 import { generateMessageId } from "@/types/stream";

+const DEFAULT_BASE_DELAY_MS = 1000;
+const DEFAULT_MAX_DELAY_MS = 15000;
+const DEFAULT_BACKOFF_FACTOR = 2;
+const DEFAULT_JITTER_MS = 250;
+
+const toError = (value: unknown): Error => {
+  if (value instanceof Error) {
+    return value;
+  }
+  if (typeof value === "string") {
+    return new Error(value);
+  }
+  return new Error("Unexpected RPC error");
+};
+
 type RequestType = SessionInboundMessage["type"];
 type ResponseType = SessionOutboundMessage["type"];

@@ -30,14 +45,100 @@ type SelectResponse<TType extends ResponseType, TData> = (message: ResponseWithE

 type DispatchRequest<TType extends RequestType> = (request: RequestOf<TType>) => void | Promise<void>;

+type DispatchOverride = (requestId: string, attempt: number) => void | Promise<void>;
+
+export type RpcFailureReason = "dispatch" | "timeout" | "response" | "disconnected";
+
+export interface RpcRetryContext {
+  requestId: string;
+  attempt: number;
+  maxAttempts: number;
+  reason: RpcFailureReason;
+  error: Error;
+}
+
+export interface RpcRetryAttemptEvent extends RpcRetryContext {
+  nextDelayMs: number;
+}
+
+export interface RpcRetryOptions {
+  maxAttempts?: number;
+  baseDelayMs?: number;
+  maxDelayMs?: number;
+  backoffFactor?: number;
+  jitterMs?: number;
+  shouldRetry?: (context: RpcRetryContext) => boolean;
+  onRetryAttempt?: (event: RpcRetryAttemptEvent) => void;
+}
+
+export class RpcRequestError extends Error {
+  public readonly reason: RpcFailureReason;
+  public readonly attempt: number;
+  public readonly maxAttempts: number;
+  public readonly requestId: string | null;
+
+  constructor(message: string, options: { reason: RpcFailureReason; attempt: number; maxAttempts: number; requestId?: string | null; cause?: unknown }) {
+    super(message);
+    this.name = "RpcRequestError";
+    this.reason = options.reason;
+    this.attempt = options.attempt;
+    this.maxAttempts = options.maxAttempts;
+    this.requestId = options.requestId ?? null;
+    if (options.cause !== undefined) {
+      (this as Error & { cause?: unknown }).cause = options.cause;
+    }
+  }
+}
+
+interface ResolvedRetryOptions {
+  maxAttempts: number;
+  baseDelayMs: number;
+  maxDelayMs: number;
+  backoffFactor: number;
+  jitterMs: number;
+  shouldRetry: NonNullable<RpcRetryOptions["shouldRetry"]>;
+  onRetryAttempt?: RpcRetryOptions["onRetryAttempt"];
+}
+
+const resolveRetryOptions = (retry: RpcRetryOptions | undefined, canRetry: boolean): ResolvedRetryOptions => {
+  const maxAttempts = Math.max(1, retry?.maxAttempts ?? 1);
+  const resolved: ResolvedRetryOptions = {
+    maxAttempts: canRetry ? maxAttempts : 1,
+    baseDelayMs: retry?.baseDelayMs ?? DEFAULT_BASE_DELAY_MS,
+    maxDelayMs: retry?.maxDelayMs ?? DEFAULT_MAX_DELAY_MS,
+    backoffFactor: retry?.backoffFactor ?? DEFAULT_BACKOFF_FACTOR,
+    jitterMs: retry?.jitterMs ?? DEFAULT_JITTER_MS,
+    shouldRetry: retry?.shouldRetry ?? (() => true),
+    onRetryAttempt: retry?.onRetryAttempt,
+  };
+  return resolved;
+};
+
+const computeDelayMs = (attempt: number, options: ResolvedRetryOptions): number => {
+  const exponential = options.baseDelayMs * options.backoffFactor ** Math.max(0, attempt - 1);
+  const capped = Math.min(options.maxDelayMs, exponential);
+  if (options.jitterMs <= 0) {
+    return capped;
+  }
+  const jitter = Math.floor(Math.random() * options.jitterMs);
+  return capped + jitter;
+};
+
 type WaitForResponseOptions = {
  requestId: string;
-  dispatch?: (requestId: string) => void | Promise<void>;
+  dispatch?: DispatchOverride;
+  retry?: RpcRetryOptions;
+  timeoutMs?: number | null;
+};
+
+type SendOptions = {
+  retry?: RpcRetryOptions;
+  timeoutMs?: number | null;
 };

 type UseSessionRpcReturn<TRequest extends RequestType, TData> = {
  state: RpcState<TData>;
-  send: (params: Omit<RequestOf<TRequest>, "type" | "requestId">) => Promise<TData>;
+  send: (params: Omit<RequestOf<TRequest>, "type" | "requestId">, options?: SendOptions) => Promise<TData>;
  waitForResponse: (options: WaitForResponseOptions) => Promise<TData>;
  reset: () => void;
 };
@@ -57,18 +158,39 @@ export function useSessionRpc<
  const [state, setState] = useState<RpcState<TData>>({ status: "idle", requestId: null });
  const activeRequestIdRef = useRef<string | null>(null);
  const resolveRef = useRef<((value: TData) => void) | null>(null);
-  const rejectRef = useRef<((reason?: any) => void) | null>(null);
+  const rejectRef = useRef<((error: Error) => void) | null>(null);
  const dispatchRef = useRef<DispatchRequest<TRequest> | undefined>(dispatch);
+  const timeoutHandleRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const retryOptionsRef = useRef<ResolvedRetryOptions | null>(null);
+  const failureHandlerRef = useRef<((reason: RpcFailureReason, error: Error) => void) | null>(null);
+  const currentAttemptRef = useRef(0);

  useEffect(() => {
    dispatchRef.current = dispatch;
  }, [dispatch]);

+  const clearTimeoutHandle = useCallback(() => {
+    if (timeoutHandleRef.current) {
+      clearTimeout(timeoutHandleRef.current);
+      timeoutHandleRef.current = null;
+    }
+  }, []);
+
  const clearActiveRequest = useCallback(() => {
+    clearTimeoutHandle();
    activeRequestIdRef.current = null;
    resolveRef.current = null;
    rejectRef.current = null;
-  }, []);
+    retryOptionsRef.current = null;
+    failureHandlerRef.current = null;
+    currentAttemptRef.current = 0;
+  }, [clearTimeoutHandle]);
+
+  useEffect(() => {
+    return () => {
+      clearActiveRequest();
+    };
+  }, [clearActiveRequest]);

  useEffect(() => {
    const unsubscribe = ws.on(responseType, (message) => {
@@ -79,15 +201,15 @@ export function useSessionRpc<
        return;
      }

+      clearTimeoutHandle();
+
      const payloadError =
        payload && typeof payload === "object" && "error" in payload && typeof (payload as any).error === "string"
          ? ((payload as any).error as string)
          : null;
      if (payloadError) {
        const error = new Error(payloadError);
-        setState({ status: "error", requestId: payload.requestId ?? null, error });
-        rejectRef.current?.(error);
-        clearActiveRequest();
+        failureHandlerRef.current?.("response", error);
        return;
      }

@@ -101,59 +223,129 @@ export function useSessionRpc<
    return () => {
      unsubscribe();
    };
-  }, [clearActiveRequest, responseType, select, ws]);
+  }, [clearActiveRequest, clearTimeoutHandle, responseType, select, ws]);

  useEffect(() => {
-    if (ws.isConnected || !activeRequestIdRef.current) {
-      return;
+    if (ws.subscribeConnectionStatus) {
+      return ws.subscribeConnectionStatus((status) => {
+        if (status.isConnected || !activeRequestIdRef.current || !failureHandlerRef.current) {
+          return;
+        }
+        failureHandlerRef.current("disconnected", new Error("WebSocket disconnected"));
+      });
    }
-    const error = new Error("WebSocket disconnected");
-    setState({ status: "error", requestId: activeRequestIdRef.current, error });
-    rejectRef.current?.(error);
-    clearActiveRequest();
-  }, [clearActiveRequest, ws.isConnected]);
+    if (!ws.isConnected && activeRequestIdRef.current && failureHandlerRef.current) {
+      failureHandlerRef.current("disconnected", new Error("WebSocket disconnected"));
+    }
+  }, [ws.isConnected, ws.subscribeConnectionStatus]);

  const waitForResponse = useCallback(
-    ({ requestId, dispatch: dispatchOverride }: WaitForResponseOptions) => {
+    ({ requestId, dispatch: dispatchOverride, retry, timeoutMs = null }: WaitForResponseOptions) => {
      return new Promise<TData>((resolve, reject) => {
-        if (!ws.isConnected) {
-          const error = new Error("WebSocket is disconnected");
-          setState({ status: "error", requestId: null, error });
-          reject(error);
-          return;
-        }
+        const finalDispatch = dispatchOverride ?? null;
+        const canRetry = typeof finalDispatch === "function";
+        const resolvedRetry = resolveRetryOptions(retry, canRetry);

        activeRequestIdRef.current = requestId;
-        resolveRef.current = resolve;
-        rejectRef.current = reject;
-        setState({ status: "loading", requestId });
-
-        if (!dispatchOverride) {
-          return;
-        }
-
-        const handleDispatchError = (error: unknown) => {
-          const err = error instanceof Error ? error : new Error(String(error));
-          setState({ status: "error", requestId, error: err });
-          reject(err);
+        resolveRef.current = (value) => {
+          resolve(value);
+        };
+        rejectRef.current = (error) => {
+          reject(error);
          clearActiveRequest();
        };
+        retryOptionsRef.current = resolvedRetry;
+        setState({ status: "loading", requestId });

-        try {
-          const maybePromise = dispatchOverride(requestId);
-          if (maybePromise && typeof (maybePromise as Promise<unknown>).then === "function") {
-            (maybePromise as Promise<unknown>).catch(handleDispatchError);
+        const finalizeError = (reason: RpcFailureReason, error: Error, attempt: number) => {
+          const rpcError = new RpcRequestError(error.message, {
+            reason,
+            attempt,
+            maxAttempts: resolvedRetry.maxAttempts,
+            requestId,
+            cause: error,
+          });
+          setState({ status: "error", requestId, error: rpcError });
+          rejectRef.current?.(rpcError);
+        };
+
+        const scheduleAttempt = (attemptNumber: number) => {
+          currentAttemptRef.current = attemptNumber;
+
+          const runDispatch = async () => {
+            if (!ws.isConnected) {
+              throw new Error("WebSocket is disconnected");
+            }
+
+            if (finalDispatch) {
+              await finalDispatch(requestId, attemptNumber);
+            }
+
+            if (timeoutMs !== null) {
+              clearTimeoutHandle();
+              timeoutHandleRef.current = setTimeout(() => {
+                failureHandlerRef.current?.("timeout", new Error("RPC request timed out"));
+              }, timeoutMs);
+            }
+          };
+
+          runDispatch().catch((error) => {
+            failureHandlerRef.current?.("dispatch", toError(error));
+          });
+        };
+
+        const handleFailure = (reason: RpcFailureReason, rawError: Error) => {
+          const attempt = currentAttemptRef.current || 1;
+          const normalized = toError(rawError);
+          const options = retryOptionsRef.current;
+          if (!options) {
+            finalizeError(reason, normalized, attempt);
+            return;
          }
-        } catch (error) {
-          handleDispatchError(error);
-        }
+
+          const withinLimit = attempt < options.maxAttempts;
+          const shouldRetry = withinLimit && options.shouldRetry({
+            requestId,
+            attempt,
+            maxAttempts: options.maxAttempts,
+            reason,
+            error: normalized,
+          });
+
+          if (!shouldRetry) {
+            finalizeError(reason, normalized, attempt);
+            return;
+          }
+
+          const nextAttempt = attempt + 1;
+          const delay = computeDelayMs(nextAttempt, options);
+          options.onRetryAttempt?.({
+            requestId,
+            attempt: nextAttempt,
+            maxAttempts: options.maxAttempts,
+            reason,
+            error: normalized,
+            nextDelayMs: delay,
+          });
+
+          clearTimeoutHandle();
+          setTimeout(() => {
+            scheduleAttempt(nextAttempt);
+          }, delay);
+        };
+
+        failureHandlerRef.current = (reason, error) => {
+          handleFailure(reason, toError(error));
+        };
+
+        scheduleAttempt(1);
      });
    },
-    [clearActiveRequest, ws.isConnected]
+    [clearActiveRequest, clearTimeoutHandle, ws.isConnected]
  );

  const send = useCallback(
-    (params: Omit<RequestOf<TRequest>, "type" | "requestId">) => {
+    (params: Omit<RequestOf<TRequest>, "type" | "requestId">, options?: SendOptions) => {
      const dispatchRequest = dispatchRef.current;
      return waitForResponse({
        requestId: generateMessageId(),
@@ -168,6 +360,8 @@ export function useSessionRpc<
          }
          ws.send({ type: "session", message: request });
        },
+        retry: options?.retry,
+        timeoutMs: options?.timeoutMs,
      });
    },
    [requestType, waitForResponse, ws]