feat(dictation): migrate transcription to session rpc hook

2025-12-02 04:30:38 +00:00
parent 5b920a2b13
commit 71b238c0cd
3 changed files with 137 additions and 76 deletions
@@ -702,7 +702,8 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid

  const initializeAgentRpc = useSessionRpc({
    ws,
-    type: "initialize_agent_request",
+    requestType: "initialize_agent_request",
+    responseType: "initialize_agent_request",
  });

  useEffect(() => {
@@ -2,9 +2,9 @@ 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 { UseWebSocketReturn } from "@/hooks/use-websocket";
 import { generateMessageId } from "@/types/stream";
-import type { SessionOutboundMessage } from "@server/server/messages";

 export type UseDictationOptions = {
  agentId: string;
@@ -65,6 +65,12 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
  const [pendingRequestId, setPendingRequestId] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);

+  const { waitForResponse: waitForTranscriptionResponse, reset: resetTranscriptionRpc } = useSessionRpc({
+    ws,
+    requestType: "send_agent_audio",
+    responseType: "transcription_result",
+  });
+
  const handleAudioLevel = useCallback((level: number) => {
    setVolume(level);
  }, []);
@@ -263,9 +269,27 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
      const requestId = generateMessageId();
      pendingRequestIdRef.current = requestId;
      setPendingRequestId(requestId);
+      const transcription = await waitForTranscriptionResponse({
+        requestId,
+        dispatch: () => sendAgentAudio(agentId, audioData, requestId, { mode }),
+      });

-      await sendAgentAudio(agentId, audioData, requestId, { mode });
-      console.log("[useDictation] sendAgentAudio resolved", { requestId });
+      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,
+      });
    } catch (err) {
      pendingRequestIdRef.current = null;
      setPendingRequestId(null);
@@ -273,6 +297,7 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
      isRecordingRef.current = false;
      setIsRecording(false);
      setVolume(0);
+      resetTranscriptionRpc();
      reportError(err, "Failed to complete dictation");
    }
  }, [
@@ -280,10 +305,13 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
    canConfirm,
    isProcessing,
    mode,
+    onTranscriptRef,
    reportError,
+    resetTranscriptionRpc,
    sendAgentAudio,
    stopDurationTracking,
    stopRecorder,
+    waitForTranscriptionResponse,
  ]);

  const reset = useCallback(() => {
@@ -296,7 +324,8 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
    setDuration(0);
    setVolume(0);
    setError(null);
-  }, [stopDurationTracking]);
+    resetTranscriptionRpc();
+  }, [resetTranscriptionRpc, stopDurationTracking]);

  const cancelRef = useRef<(() => void) | null>(null);
  useEffect(() => {
@@ -323,40 +352,8 @@ export function useDictation(options: UseDictationOptions): UseDictationResult {
    pendingRequestIdRef.current = null;
    setPendingRequestId(null);
    setIsProcessing(false);
-  }, [agentId]);
-
-  useEffect(() => {
-    const unsubscribe = ws.on("transcription_result", (message: SessionOutboundMessage) => {
-      if (message.type !== "transcription_result") {
-        return;
-      }
-      const expectedId = pendingRequestIdRef.current;
-      if (!expectedId || message.payload.requestId !== expectedId) {
-        return;
-      }
-
-      pendingRequestIdRef.current = null;
-      setPendingRequestId(null);
-      setIsProcessing(false);
-
-      const transcriptText = message.payload.text?.trim();
-      if (!transcriptText) {
-        return;
-      }
-
-      console.log("[useDictation] transcription_result received", {
-        requestId: expectedId,
-        textLength: transcriptText.length,
-      });
-      onTranscriptRef.current?.(transcriptText, {
-        requestId: message.payload.requestId ?? expectedId,
-      });
-    });
-
-    return () => {
-      unsubscribe();
-    };
-  }, [ws]);
+    resetTranscriptionRpc();
+  }, [agentId, resetTranscriptionRpc]);

  useEffect(() => {
    return () => {
@@ -3,10 +3,11 @@ import type { SessionInboundMessage, SessionOutboundMessage } from "@server/serv
 import type { UseWebSocketReturn } from "./use-websocket";
 import { generateMessageId } from "@/types/stream";

-type SharedType = SessionInboundMessage["type"] & SessionOutboundMessage["type"];
+type RequestType = SessionInboundMessage["type"];
+type ResponseType = SessionOutboundMessage["type"];

-type RequestOf<TType extends SharedType> = Extract<SessionInboundMessage, { type: TType }>;
-type ResponseOf<TType extends SharedType> = Extract<SessionOutboundMessage, { type: TType }>;
+type RequestOf<TType extends RequestType> = Extract<SessionInboundMessage, { type: TType }>;
+type ResponseOf<TType extends ResponseType> = Extract<SessionOutboundMessage, { type: TType }>;

 type RpcState<T> =
  | { status: "idle"; requestId: null }
@@ -14,30 +15,54 @@ type RpcState<T> =
  | { status: "success"; requestId: string; data: T }
  | { status: "error"; requestId: string | null; error: Error };

-type ResponseWithEnvelope<TType extends SharedType> = ResponseOf<TType> extends {
-  payload: { requestId?: string; error?: string };
-}
-  ? ResponseOf<TType>
-  : never;
+type ResponseWithEnvelope<TType extends ResponseType> = Extract<
+  ResponseOf<TType>,
+  { payload: { requestId?: string } }
+>;

-type EnsureEnvelope<TType extends SharedType> = ResponseWithEnvelope<TType> extends never ? never : TType;
+type EnsureEnvelope<TType extends ResponseType> = ResponseWithEnvelope<TType> extends never ? never : TType;

-type ResponsePayload<TType extends SharedType> = ResponseWithEnvelope<TType> extends { payload: infer P }
+type ResponsePayload<TType extends ResponseType> = ResponseWithEnvelope<TType> extends { payload: infer P }
  ? P
  : never;

-type SelectResponse<TType extends SharedType, TData> = (message: ResponseWithEnvelope<TType>) => TData;
+type SelectResponse<TType extends ResponseType, TData> = (message: ResponseWithEnvelope<TType>) => TData;

-export function useSessionRpc<TType extends SharedType, TData = ResponsePayload<TType>>(options: {
+type DispatchRequest<TType extends RequestType> = (request: RequestOf<TType>) => void | Promise<void>;
+
+type WaitForResponseOptions = {
+  requestId: string;
+  dispatch?: (requestId: string) => void | Promise<void>;
+};
+
+type UseSessionRpcReturn<TRequest extends RequestType, TData> = {
+  state: RpcState<TData>;
+  send: (params: Omit<RequestOf<TRequest>, "type" | "requestId">) => Promise<TData>;
+  waitForResponse: (options: WaitForResponseOptions) => Promise<TData>;
+  reset: () => void;
+};
+
+export function useSessionRpc<
+  TRequest extends RequestType,
+  TResponse extends ResponseType,
+  TData = ResponsePayload<TResponse>
+>(options: {
  ws: UseWebSocketReturn;
-  type: EnsureEnvelope<TType>;
-  select?: SelectResponse<TType, TData>;
-}) {
-  const { ws, type, select } = options;
+  requestType: TRequest;
+  responseType: EnsureEnvelope<TResponse>;
+  select?: SelectResponse<TResponse, TData>;
+  dispatch?: DispatchRequest<TRequest>;
+}): UseSessionRpcReturn<TRequest, TData> {
+  const { ws, requestType, responseType, select, dispatch } = options;
  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 dispatchRef = useRef<DispatchRequest<TRequest> | undefined>(dispatch);
+
+  useEffect(() => {
+    dispatchRef.current = dispatch;
+  }, [dispatch]);

  const clearActiveRequest = useCallback(() => {
    activeRequestIdRef.current = null;
@@ -46,22 +71,27 @@ export function useSessionRpc<TType extends SharedType, TData = ResponsePayload<
  }, []);

  useEffect(() => {
-    const unsubscribe = ws.on(type, (message) => {
-      const typedMessage = message as ResponseWithEnvelope<TType>;
+    const unsubscribe = ws.on(responseType, (message) => {
+      const typedMessage = message as ResponseWithEnvelope<TResponse>;
      const payload = typedMessage.payload;
-      if (!payload || payload.requestId !== activeRequestIdRef.current) {
+      const expectedId = activeRequestIdRef.current;
+      if (!payload || !expectedId || payload.requestId !== expectedId) {
        return;
      }

-      if (payload.error) {
-        const error = new Error(payload.error);
+      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();
        return;
      }

-      const baseData = typedMessage.payload as ResponsePayload<TType>;
+      const baseData = typedMessage.payload as ResponsePayload<TResponse>;
      const data = select ? select(typedMessage) : (baseData as unknown as TData);
      setState({ status: "success", requestId: payload.requestId ?? null, data });
      resolveRef.current?.(data);
@@ -71,7 +101,7 @@ export function useSessionRpc<TType extends SharedType, TData = ResponsePayload<
    return () => {
      unsubscribe();
    };
-  }, [clearActiveRequest, select, type, ws]);
+  }, [clearActiveRequest, responseType, select, ws]);

  useEffect(() => {
    if (ws.isConnected || !activeRequestIdRef.current) {
@@ -83,8 +113,8 @@ export function useSessionRpc<TType extends SharedType, TData = ResponsePayload<
    clearActiveRequest();
  }, [clearActiveRequest, ws.isConnected]);

-  const send = useCallback(
-    (params: Omit<RequestOf<TType>, "type" | "requestId">) => {
+  const waitForResponse = useCallback(
+    ({ requestId, dispatch: dispatchOverride }: WaitForResponseOptions) => {
      return new Promise<TData>((resolve, reject) => {
        if (!ws.isConnected) {
          const error = new Error("WebSocket is disconnected");
@@ -93,29 +123,54 @@ export function useSessionRpc<TType extends SharedType, TData = ResponsePayload<
          return;
        }

-        const requestId = generateMessageId();
        activeRequestIdRef.current = requestId;
        resolveRef.current = resolve;
        rejectRef.current = reject;
        setState({ status: "loading", requestId });

-        const request = {
-          type,
-          ...params,
-          requestId,
-        } as RequestOf<TType>;
+        if (!dispatchOverride) {
+          return;
+        }

-        try {
-          ws.send({ type: "session", message: request });
-        } catch (error) {
+        const handleDispatchError = (error: unknown) => {
          const err = error instanceof Error ? error : new Error(String(error));
-          setState({ status: "error", requestId: requestId, error: err });
+          setState({ status: "error", requestId, error: err });
          reject(err);
          clearActiveRequest();
+        };
+
+        try {
+          const maybePromise = dispatchOverride(requestId);
+          if (maybePromise && typeof (maybePromise as Promise<unknown>).then === "function") {
+            (maybePromise as Promise<unknown>).catch(handleDispatchError);
+          }
+        } catch (error) {
+          handleDispatchError(error);
        }
      });
    },
-    [clearActiveRequest, type, ws]
+    [clearActiveRequest, ws.isConnected]
+  );
+
+  const send = useCallback(
+    (params: Omit<RequestOf<TRequest>, "type" | "requestId">) => {
+      const dispatchRequest = dispatchRef.current;
+      return waitForResponse({
+        requestId: generateMessageId(),
+        dispatch: (generatedId) => {
+          const request = {
+            type: requestType,
+            ...params,
+            requestId: generatedId,
+          } as RequestOf<TRequest>;
+          if (dispatchRequest) {
+            return dispatchRequest(request);
+          }
+          ws.send({ type: "session", message: request });
+        },
+      });
+    },
+    [requestType, waitForResponse, ws]
  );

  const reset = useCallback(() => {
@@ -123,5 +178,13 @@ export function useSessionRpc<TType extends SharedType, TData = ResponsePayload<
    setState({ status: "idle", requestId: null });
  }, [clearActiveRequest]);

-  return useMemo(() => ({ state, send, reset }), [reset, send, state]);
+  return useMemo(
+    () => ({
+      state,
+      send,
+      waitForResponse,
+      reset,
+    }),
+    [reset, send, state, waitForResponse]
+  );
 }