diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 77747b5ac..c4cf65dda 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -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(() => { diff --git a/packages/app/src/hooks/use-dictation.ts b/packages/app/src/hooks/use-dictation.ts index 353df932e..4ad5b73ca 100644 --- a/packages/app/src/hooks/use-dictation.ts +++ b/packages/app/src/hooks/use-dictation.ts @@ -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(null); const [error, setError] = useState(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 () => { diff --git a/packages/app/src/hooks/use-session-rpc.ts b/packages/app/src/hooks/use-session-rpc.ts index 2c6650181..27b3f5a7c 100644 --- a/packages/app/src/hooks/use-session-rpc.ts +++ b/packages/app/src/hooks/use-session-rpc.ts @@ -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 = Extract; -type ResponseOf = Extract; +type RequestOf = Extract; +type ResponseOf = Extract; type RpcState = | { status: "idle"; requestId: null } @@ -14,30 +15,54 @@ type RpcState = | { status: "success"; requestId: string; data: T } | { status: "error"; requestId: string | null; error: Error }; -type ResponseWithEnvelope = ResponseOf extends { - payload: { requestId?: string; error?: string }; -} - ? ResponseOf - : never; +type ResponseWithEnvelope = Extract< + ResponseOf, + { payload: { requestId?: string } } +>; -type EnsureEnvelope = ResponseWithEnvelope extends never ? never : TType; +type EnsureEnvelope = ResponseWithEnvelope extends never ? never : TType; -type ResponsePayload = ResponseWithEnvelope extends { payload: infer P } +type ResponsePayload = ResponseWithEnvelope extends { payload: infer P } ? P : never; -type SelectResponse = (message: ResponseWithEnvelope) => TData; +type SelectResponse = (message: ResponseWithEnvelope) => TData; -export function useSessionRpc>(options: { +type DispatchRequest = (request: RequestOf) => void | Promise; + +type WaitForResponseOptions = { + requestId: string; + dispatch?: (requestId: string) => void | Promise; +}; + +type UseSessionRpcReturn = { + state: RpcState; + send: (params: Omit, "type" | "requestId">) => Promise; + waitForResponse: (options: WaitForResponseOptions) => Promise; + reset: () => void; +}; + +export function useSessionRpc< + TRequest extends RequestType, + TResponse extends ResponseType, + TData = ResponsePayload +>(options: { ws: UseWebSocketReturn; - type: EnsureEnvelope; - select?: SelectResponse; -}) { - const { ws, type, select } = options; + requestType: TRequest; + responseType: EnsureEnvelope; + select?: SelectResponse; + dispatch?: DispatchRequest; +}): UseSessionRpcReturn { + const { ws, requestType, responseType, select, dispatch } = options; const [state, setState] = useState>({ status: "idle", requestId: null }); const activeRequestIdRef = useRef(null); const resolveRef = useRef<((value: TData) => void) | null>(null); const rejectRef = useRef<((reason?: any) => void) | null>(null); + const dispatchRef = useRef | undefined>(dispatch); + + useEffect(() => { + dispatchRef.current = dispatch; + }, [dispatch]); const clearActiveRequest = useCallback(() => { activeRequestIdRef.current = null; @@ -46,22 +71,27 @@ export function useSessionRpc { - const unsubscribe = ws.on(type, (message) => { - const typedMessage = message as ResponseWithEnvelope; + const unsubscribe = ws.on(responseType, (message) => { + const typedMessage = message as ResponseWithEnvelope; 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; + const baseData = typedMessage.payload as ResponsePayload; 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 { unsubscribe(); }; - }, [clearActiveRequest, select, type, ws]); + }, [clearActiveRequest, responseType, select, ws]); useEffect(() => { if (ws.isConnected || !activeRequestIdRef.current) { @@ -83,8 +113,8 @@ export function useSessionRpc, "type" | "requestId">) => { + const waitForResponse = useCallback( + ({ requestId, dispatch: dispatchOverride }: WaitForResponseOptions) => { return new Promise((resolve, reject) => { if (!ws.isConnected) { const error = new Error("WebSocket is disconnected"); @@ -93,29 +123,54 @@ export function useSessionRpc; + 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).then === "function") { + (maybePromise as Promise).catch(handleDispatchError); + } + } catch (error) { + handleDispatchError(error); } }); }, - [clearActiveRequest, type, ws] + [clearActiveRequest, ws.isConnected] + ); + + const send = useCallback( + (params: Omit, "type" | "requestId">) => { + const dispatchRequest = dispatchRef.current; + return waitForResponse({ + requestId: generateMessageId(), + dispatch: (generatedId) => { + const request = { + type: requestType, + ...params, + requestId: generatedId, + } as RequestOf; + if (dispatchRequest) { + return dispatchRequest(request); + } + ws.send({ type: "session", message: request }); + }, + }); + }, + [requestType, waitForResponse, ws] ); const reset = useCallback(() => { @@ -123,5 +178,13 @@ export function useSessionRpc ({ state, send, reset }), [reset, send, state]); + return useMemo( + () => ({ + state, + send, + waitForResponse, + reset, + }), + [reset, send, state, waitForResponse] + ); }