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

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

@@ -3,9 +3,33 @@ import { useCallback, useEffect, useRef, useState } from "react";
 import type { SessionContextValue } from "@/contexts/session-context";
 import { useAudioRecorder } from "@/hooks/use-audio-recorder";
 import { useSessionRpc } from "@/hooks/use-session-rpc";
+import type { RpcFailureReason, RpcRetryAttemptEvent } from "@/hooks/use-session-rpc";
 import type { UseWebSocketReturn } from "@/hooks/use-websocket";
 import { generateMessageId } from "@/types/stream";
 
+export type DictationStatus = "idle" | "recording" | "uploading" | "retrying" | "failed";
+
+export type DictationRetryInfo = {
+  attempt: number;
+  maxAttempts: number;
+  reason: RpcFailureReason;
+  errorMessage: string;
+  nextRetryMs: number;
+};
+
+export type FailedDictationRecording = {
+  requestId: string;
+  durationSeconds: number;
+  sizeBytes: number;
+  format: string;
+  recordedAt: number;
+  errorMessage: string;
+};
+
+export type DictationOutcome =
+  | { type: "success"; requestId: string; timestamp: number }
+  | { type: "failure"; requestId: string; errorMessage: string; timestamp: number };
+
 export type UseDictationOptions = {
   agentId: string;
   sendAgentAudio: SessionContextValue["sendAgentAudio"];
@@ -13,6 +37,8 @@ export type UseDictationOptions = {
   mode?: "transcribe_only" | "auto_run";
   onTranscript: (text: string, meta: { requestId: string }) => void;
   onError?: (error: Error) => void;
+  onRetryAttempt?: (info: DictationRetryInfo) => void;
+  onPermanentFailure?: (error: Error, context: { requestId: string }) => void;
   canStart?: () => boolean;
   canConfirm?: () => boolean;
   autoStopWhenHidden?: { isVisible: boolean };
@@ -26,13 +52,27 @@ export type UseDictationResult = {
   duration: number;
   pendingRequestId: string | null;
   error: string | null;
+  status: DictationStatus;
+  retryAttempt: number;
+  maxRetryAttempts: number;
+  retryInfo: DictationRetryInfo | null;
+  failedRecording: FailedDictationRecording | null;
+  lastOutcome: DictationOutcome | null;
   startDictation: () => Promise<void>;
   cancelDictation: () => Promise<void>;
   confirmDictation: () => Promise<void>;
+  retryFailedDictation: () => Promise<void>;
+  discardFailedDictation: () => void;
   reset: () => void;
 };
 
 const DURATION_TICK_MS = 1000;
+const MAX_AUTO_RETRY_ATTEMPTS = 5;
+const RETRY_BASE_DELAY_MS = 2000;
+const RETRY_MAX_DELAY_MS = 12000;
+const RETRY_BACKOFF_FACTOR = 1.8;
+const RETRY_JITTER_MS = 400;
+const TRANSCRIPTION_TIMEOUT_MS = 120000;
 
 const toError = (error: unknown): Error => {
   if (error instanceof Error) {
@@ -44,6 +84,35 @@ const toError = (error: unknown): Error => {
   return new Error("An unexpected error occurred while handling dictation.");
 };
 
+type CapturedAudioPayload = {
+  blob: Blob;
+  format: string;
+  sizeBytes: number;
+  durationSeconds: number;
+  recordedAt: number;
+};
+
+const deriveFormatFromMime = (mimeType?: string): string => {
+  if (!mimeType || mimeType.length === 0) {
+    return "webm";
+  }
+  const slashIndex = mimeType.indexOf("/");
+  let formatPart = slashIndex >= 0 ? mimeType.slice(slashIndex + 1) : mimeType;
+  const semicolonIndex = formatPart.indexOf(";");
+  if (semicolonIndex >= 0) {
+    formatPart = formatPart.slice(0, semicolonIndex);
+  }
+  return formatPart.trim().length > 0 ? formatPart.trim() : "webm";
+};
+
+const buildCapturedAudioPayload = (blob: Blob, durationSeconds: number): CapturedAudioPayload => ({
+  blob,
+  format: deriveFormatFromMime(blob.type),
+  sizeBytes: typeof blob.size === "number" ? blob.size : 0,
+  durationSeconds,
+  recordedAt: Date.now(),
+});
+
 export function useDictation(options: UseDictationOptions): UseDictationResult {
   const {
     agentId,
@@ -52,6 +121,8 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     mode = "transcribe_only",
     onTranscript,
     onError,
+    onRetryAttempt,
+    onPermanentFailure,
     canStart,
     canConfirm,
     autoStopWhenHidden,
@@ -64,6 +135,12 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
   const [duration, setDuration] = useState(0);
   const [pendingRequestId, setPendingRequestId] = useState<string | null>(null);
   const [error, setError] = useState<string | null>(null);
+  const [status, setStatus] = useState<DictationStatus>("idle");
+  const [retryAttempt, setRetryAttempt] = useState(0);
+  const [retryInfo, setRetryInfo] = useState<DictationRetryInfo | null>(null);
+  const [failedRecording, setFailedRecording] = useState<FailedDictationRecording | null>(null);
+  const [lastOutcome, setLastOutcome] = useState<DictationOutcome | null>(null);
+  const maxRetryAttempts = MAX_AUTO_RETRY_ATTEMPTS;
 
   const { waitForResponse: waitForTranscriptionResponse, reset: resetTranscriptionRpc } = useSessionRpc({
     ws,
@@ -71,10 +148,16 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     responseType: "transcription_result",
   });
 
+  const pendingAudioRef = useRef<CapturedAudioPayload | null>(null);
   const handleAudioLevel = useCallback((level: number) => {
     setVolume(level);
   }, []);
 
+  const durationRef = useRef(0);
+  useEffect(() => {
+    durationRef.current = duration;
+  }, [duration]);
+
   const recorder = useAudioRecorder({ onAudioLevel: handleAudioLevel });
   const recorderRef = useRef(recorder);
   useEffect(() => {
@@ -91,6 +174,16 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     onErrorRef.current = onError;
   }, [onError]);
 
+  const onRetryAttemptRef = useRef(onRetryAttempt);
+  useEffect(() => {
+    onRetryAttemptRef.current = onRetryAttempt;
+  }, [onRetryAttempt]);
+
+  const onPermanentFailureRef = useRef(onPermanentFailure);
+  useEffect(() => {
+    onPermanentFailureRef.current = onPermanentFailure;
+  }, [onPermanentFailure]);
+
   const isRecordingRef = useRef(isRecording);
   useEffect(() => {
     isRecordingRef.current = isRecording;
@@ -160,6 +253,137 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     return stopPromise;
   }, []);
 
+  const transmitDictation = useCallback(
+    async (requestId: string) => {
+      const capturedAudio = pendingAudioRef.current;
+      if (!capturedAudio) {
+        throw new Error("No recorded audio available for transcription");
+      }
+
+      setRetryAttempt(1);
+      setRetryInfo(null);
+
+      console.info("[useDictation] sending transcription request", {
+        requestId,
+        attempt: 1,
+        agentId,
+        size: capturedAudio.sizeBytes,
+        durationSeconds: capturedAudio.durationSeconds,
+      });
+
+      const transcription = await waitForTranscriptionResponse({
+        requestId,
+        dispatch: async (_id, attempt) => {
+          setStatus("uploading");
+          setRetryAttempt(attempt);
+          await sendAgentAudio(agentId, capturedAudio.blob, requestId, { mode });
+        },
+        retry: {
+          maxAttempts: MAX_AUTO_RETRY_ATTEMPTS,
+          baseDelayMs: RETRY_BASE_DELAY_MS,
+          maxDelayMs: RETRY_MAX_DELAY_MS,
+          backoffFactor: RETRY_BACKOFF_FACTOR,
+          jitterMs: RETRY_JITTER_MS,
+          shouldRetry: ({ attempt, maxAttempts }) => attempt < maxAttempts,
+          onRetryAttempt: (event: RpcRetryAttemptEvent) => {
+            setStatus("retrying");
+            setRetryAttempt(event.attempt);
+            const info: DictationRetryInfo = {
+              attempt: event.attempt,
+              maxAttempts: event.maxAttempts,
+              reason: event.reason,
+              errorMessage: event.error.message,
+              nextRetryMs: event.nextDelayMs,
+            };
+            setRetryInfo(info);
+            onRetryAttemptRef.current?.(info);
+            console.warn("[useDictation] retry scheduled", {
+              requestId,
+              attempt: event.attempt,
+              maxAttempts: event.maxAttempts,
+              reason: event.reason,
+              error: event.error.message,
+              nextRetryMs: event.nextDelayMs,
+            });
+          },
+        },
+        timeoutMs: TRANSCRIPTION_TIMEOUT_MS,
+      });
+
+      return transcription;
+    },
+    [agentId, mode, onRetryAttemptRef, sendAgentAudio, waitForTranscriptionResponse]
+  );
+
+  const handleTranscriptionSuccess = useCallback(
+    (transcription: Awaited<ReturnType<typeof waitForTranscriptionResponse>>, requestId: string) => {
+      pendingRequestIdRef.current = null;
+      setPendingRequestId(null);
+      setIsProcessing(false);
+      setStatus("idle");
+      setRetryAttempt(0);
+      setRetryInfo(null);
+      setFailedRecording(null);
+      pendingAudioRef.current = null;
+      setLastOutcome({ type: "success", requestId, timestamp: Date.now() });
+
+      const transcriptText = transcription.text?.trim();
+      if (!transcriptText) {
+        return;
+      }
+
+      console.log("[useDictation] transcription_result received", {
+        requestId,
+        textLength: transcriptText.length,
+      });
+      onTranscriptRef.current?.(transcriptText, {
+        requestId: transcription.requestId ?? requestId,
+      });
+    },
+    [onTranscriptRef]
+  );
+
+  const handleDictationFailure = useCallback(
+    (failure: unknown, requestId: string | null) => {
+      const normalized = toError(failure);
+      pendingRequestIdRef.current = null;
+      setPendingRequestId(null);
+      setIsProcessing(false);
+      isRecordingRef.current = false;
+      setIsRecording(false);
+      setVolume(0);
+      setRetryInfo(null);
+
+      const capturedAudio = pendingAudioRef.current;
+      if (capturedAudio) {
+        setStatus("failed");
+        setFailedRecording({
+          requestId: requestId ?? generateMessageId(),
+          durationSeconds: capturedAudio.durationSeconds,
+          sizeBytes: capturedAudio.sizeBytes,
+          format: capturedAudio.format,
+          recordedAt: capturedAudio.recordedAt,
+          errorMessage: normalized.message,
+        });
+        if (requestId) {
+          onPermanentFailureRef.current?.(normalized, { requestId });
+        }
+      } else {
+        setStatus("idle");
+      }
+
+      setRetryAttempt(0);
+      setLastOutcome({
+        type: "failure",
+        requestId: requestId ?? generateMessageId(),
+        errorMessage: normalized.message,
+        timestamp: Date.now(),
+      });
+      reportError(normalized, "Failed to complete dictation");
+    },
+    [onPermanentFailureRef, reportError]
+  );
+
   const startDictation = useCallback(async () => {
     console.log("[useDictation] startDictation requested", {
       isRecording: isRecordingRef.current,
@@ -182,6 +406,12 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     setVolume(0);
     setDuration(0);
     setIsProcessing(false);
+    setStatus("recording");
+    setRetryAttempt(0);
+    setRetryInfo(null);
+    setFailedRecording(null);
+    pendingAudioRef.current = null;
+    setLastOutcome(null);
     pendingRequestIdRef.current = null;
     setPendingRequestId(null);
 
@@ -234,6 +464,12 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
       setIsRecording(false);
       setIsProcessing(false);
       setVolume(0);
+      setStatus("idle");
+      setRetryAttempt(0);
+      setRetryInfo(null);
+      pendingAudioRef.current = null;
+      setFailedRecording(null);
+      setLastOutcome(null);
     }
   }, [reportError, stopDurationTracking, stopRecorder]);
 
@@ -259,61 +495,85 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     stopDurationTracking();
     setDuration(0);
     setIsProcessing(true);
+    setRetryInfo(null);
+    setRetryAttempt(0);
+    setLastOutcome(null);
 
+    let requestId: string | null = null;
     try {
       const audioData = await stopRecorder();
+      const recordedDurationSeconds = durationRef.current;
+      pendingAudioRef.current = buildCapturedAudioPayload(audioData, recordedDurationSeconds);
+      setStatus("uploading");
       isRecordingRef.current = false;
       setIsRecording(false);
       setVolume(0);
 
-      const requestId = generateMessageId();
+      requestId = generateMessageId();
       pendingRequestIdRef.current = requestId;
       setPendingRequestId(requestId);
-      const transcription = await waitForTranscriptionResponse({
-        requestId,
-        dispatch: () => sendAgentAudio(agentId, audioData, requestId, { mode }),
-      });
 
-      pendingRequestIdRef.current = null;
-      setPendingRequestId(null);
-      setIsProcessing(false);
-
-      const transcriptText = transcription.text?.trim();
-      if (!transcriptText) {
-        return;
-      }
-
-      console.log("[useDictation] transcription_result received", {
-        requestId,
-        textLength: transcriptText.length,
-      });
-      onTranscriptRef.current?.(transcriptText, {
-        requestId: transcription.requestId ?? requestId,
-      });
+      const transcription = await transmitDictation(requestId);
+      handleTranscriptionSuccess(transcription, requestId);
     } catch (err) {
-      pendingRequestIdRef.current = null;
-      setPendingRequestId(null);
-      setIsProcessing(false);
-      isRecordingRef.current = false;
-      setIsRecording(false);
-      setVolume(0);
       resetTranscriptionRpc();
-      reportError(err, "Failed to complete dictation");
+      handleDictationFailure(err, requestId);
     }
   }, [
     agentId,
     canConfirm,
     isProcessing,
     mode,
-    onTranscriptRef,
-    reportError,
+    handleDictationFailure,
+    handleTranscriptionSuccess,
     resetTranscriptionRpc,
-    sendAgentAudio,
     stopDurationTracking,
     stopRecorder,
-    waitForTranscriptionResponse,
+    transmitDictation,
   ]);
 
+  const retryFailedDictation = useCallback(async () => {
+    if (!pendingAudioRef.current) {
+      return;
+    }
+    setError(null);
+    setRetryInfo(null);
+    setRetryAttempt(0);
+    setStatus("uploading");
+    setIsProcessing(true);
+    setLastOutcome(null);
+
+    const requestId = generateMessageId();
+    pendingRequestIdRef.current = requestId;
+    setPendingRequestId(requestId);
+
+    try {
+      const transcription = await transmitDictation(requestId);
+      handleTranscriptionSuccess(transcription, requestId);
+    } catch (err) {
+      resetTranscriptionRpc();
+      handleDictationFailure(err, requestId);
+    }
+  }, [
+    handleDictationFailure,
+    handleTranscriptionSuccess,
+    resetTranscriptionRpc,
+    transmitDictation,
+  ]);
+
+  const discardFailedDictation = useCallback(() => {
+    pendingAudioRef.current = null;
+    pendingRequestIdRef.current = null;
+    setPendingRequestId(null);
+    setIsProcessing(false);
+    setFailedRecording(null);
+    setStatus("idle");
+    setRetryAttempt(0);
+    setRetryInfo(null);
+    setError(null);
+    setLastOutcome(null);
+  }, []);
+
   const reset = useCallback(() => {
     pendingRequestIdRef.current = null;
     setPendingRequestId(null);
@@ -324,6 +584,12 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     setDuration(0);
     setVolume(0);
     setError(null);
+    setStatus("idle");
+    setRetryAttempt(0);
+    setRetryInfo(null);
+    setFailedRecording(null);
+    pendingAudioRef.current = null;
+    setLastOutcome(null);
     resetTranscriptionRpc();
   }, [resetTranscriptionRpc, stopDurationTracking]);
 
@@ -352,6 +618,11 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     pendingRequestIdRef.current = null;
     setPendingRequestId(null);
     setIsProcessing(false);
+    pendingAudioRef.current = null;
+    setStatus("idle");
+    setRetryAttempt(0);
+    setRetryInfo(null);
+    setFailedRecording(null);
     resetTranscriptionRpc();
   }, [agentId, resetTranscriptionRpc]);
 
@@ -377,9 +648,17 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
     duration,
     pendingRequestId,
     error,
+    status,
+    retryAttempt,
+    maxRetryAttempts,
+    retryInfo,
+    failedRecording,
+    lastOutcome,
     startDictation,
     cancelDictation,
     confirmDictation,
+    retryFailedDictation,
+    discardFailedDictation,
     reset,
   };
 }

packages/app/src/hooks/use-session-rpc.ts

@@ -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]

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

@@ -6,6 +6,11 @@ import type {
   SessionOutboundMessage,
 } from "@server/server/messages";
 
+export interface ConnectionStatusSnapshot {
+  isConnected: boolean;
+  isConnecting: boolean;
+}
+
 export interface UseWebSocketReturn {
   isConnected: boolean;
   isConnecting: boolean;
@@ -18,6 +23,8 @@ export interface UseWebSocketReturn {
   ) => () => void;
   sendPing: () => void;
   sendUserMessage: (message: string) => void;
+  subscribeConnectionStatus?: (listener: (status: ConnectionStatusSnapshot) => void) => () => void;
+  getConnectionState?: () => ConnectionStatusSnapshot;
 }
 
 const RECONNECT_BASE_DELAY_MS = 1500;
@@ -34,6 +41,25 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
   const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
   const reconnectAttemptRef = useRef(0);
   const shouldReconnectRef = useRef(true);
+  const connectionListenersRef = useRef(new Set<(status: ConnectionStatusSnapshot) => void>());
+  const connectionStateRef = useRef<ConnectionStatusSnapshot>({ isConnected: false, isConnecting: true });
+
+  const notifyConnectionListeners = useCallback((state: ConnectionStatusSnapshot) => {
+    connectionStateRef.current = state;
+    for (const listener of connectionListenersRef.current) {
+      try {
+        listener(state);
+      } catch (error) {
+        console.error("[WS] Connection listener error", error);
+      }
+    }
+  }, []);
+
+  const updateConnectionState = useCallback((state: ConnectionStatusSnapshot) => {
+    setIsConnected(state.isConnected);
+    setIsConnecting(state.isConnecting);
+    notifyConnectionListeners(state);
+  }, [notifyConnectionListeners]);
 
   const connect = useCallback(() => {
     if (wsRef.current && (wsRef.current.readyState === WebSocket.OPEN || wsRef.current.readyState === WebSocket.CONNECTING)) {
@@ -60,8 +86,7 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
         setLastError(reason.trim());
       }
 
-      setIsConnected(false);
-      setIsConnecting(false);
+      updateConnectionState({ isConnected: false, isConnecting: false });
 
       const attempt = reconnectAttemptRef.current;
       const delay = Math.min(RECONNECT_BASE_DELAY_MS * 2 ** attempt, RECONNECT_MAX_DELAY_MS);
@@ -76,22 +101,21 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
         if (!shouldReconnectRef.current) {
           return;
         }
-        setIsConnecting(true);
-        connect();
-      }, delay);
-    };
+      updateConnectionState({ isConnected: false, isConnecting: true });
+      connect();
+    }, delay);
+  };
 
     try {
       // Add conversation ID to URL if provided
       const wsUrl = conversationId ? `${url}?conversationId=${conversationId}` : url;
       const ws = new WebSocket(wsUrl);
       wsRef.current = ws;
-      setIsConnecting(true);
+      updateConnectionState({ isConnected: false, isConnecting: true });
 
       ws.onopen = () => {
         console.log("[WS] Connected to server");
-        setIsConnected(true);
-        setIsConnecting(false);
+        updateConnectionState({ isConnected: true, isConnecting: false });
         setLastError(null);
         reconnectAttemptRef.current = 0;
       };
@@ -102,7 +126,7 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
           typeof event?.reason === "string" && event.reason.trim().length > 0
             ? event.reason.trim()
             : `Socket closed (code ${event?.code ?? "unknown"})`;
-        setIsConnected(false);
+        updateConnectionState({ isConnected: false, isConnecting: false });
         scheduleReconnect(reason);
       };
 
@@ -164,7 +188,7 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
       const reason = err instanceof Error ? err.message : "Failed to create WebSocket";
       scheduleReconnect(reason);
     }
-  }, [url, conversationId]);
+  }, [updateConnectionState, url, conversationId]);
 
   useEffect(() => {
     shouldReconnectRef.current = true;
@@ -259,6 +283,25 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
     [send]
   );
 
+  const subscribeConnectionStatus = useCallback(
+    (listener: (status: ConnectionStatusSnapshot) => void) => {
+      connectionListenersRef.current.add(listener);
+      listener(connectionStateRef.current);
+      return () => {
+        connectionListenersRef.current.delete(listener);
+      };
+    },
+    []
+  );
+
+  const getConnectionState = useCallback(() => {
+    const readyState = wsRef.current?.readyState ?? WebSocket.CLOSED;
+    return {
+      isConnected: readyState === WebSocket.OPEN,
+      isConnecting: readyState === WebSocket.CONNECTING,
+    } satisfies ConnectionStatusSnapshot;
+  }, []);
+
   return useMemo(
     () => ({
       isConnected,
@@ -269,7 +312,20 @@ export function useWebSocket(url: string, conversationId?: string | null): UseWe
       on,
       sendPing,
       sendUserMessage,
+      subscribeConnectionStatus,
+      getConnectionState,
     }),
-    [isConnected, isConnecting, currentConversationId, lastError, send, on, sendPing, sendUserMessage]
+    [
+      isConnected,
+      isConnecting,
+      currentConversationId,
+      lastError,
+      send,
+      on,
+      sendPing,
+      sendUserMessage,
+      subscribeConnectionStatus,
+      getConnectionState,
+    ]
   );
 }