diff --git a/packages/app/src/hooks/use-session-rpc.ts b/packages/app/src/hooks/use-session-rpc.ts index e90c2963f..2c6650181 100644 --- a/packages/app/src/hooks/use-session-rpc.ts +++ b/packages/app/src/hooks/use-session-rpc.ts @@ -14,15 +14,23 @@ type RpcState = | { status: "success"; requestId: string; data: T } | { status: "error"; requestId: string | null; error: Error }; -type ResponseWithRequestId = ResponseOf & { +type ResponseWithEnvelope = ResponseOf extends { payload: { requestId?: string; error?: string }; -}; +} + ? ResponseOf + : never; -type SelectResponse = (message: ResponseWithRequestId) => TData; +type EnsureEnvelope = ResponseWithEnvelope extends never ? never : TType; -export function useSessionRpc["payload"]>(options: { +type ResponsePayload = ResponseWithEnvelope extends { payload: infer P } + ? P + : never; + +type SelectResponse = (message: ResponseWithEnvelope) => TData; + +export function useSessionRpc>(options: { ws: UseWebSocketReturn; - type: TType; + type: EnsureEnvelope; select?: SelectResponse; }) { const { ws, type, select } = options; @@ -31,9 +39,16 @@ export function useSessionRpc void) | null>(null); const rejectRef = useRef<((reason?: any) => void) | null>(null); + const clearActiveRequest = useCallback(() => { + activeRequestIdRef.current = null; + resolveRef.current = null; + rejectRef.current = null; + }, []); + useEffect(() => { const unsubscribe = ws.on(type, (message) => { - const payload = (message as ResponseWithRequestId).payload; + const typedMessage = message as ResponseWithEnvelope; + const payload = typedMessage.payload; if (!payload || payload.requestId !== activeRequestIdRef.current) { return; } @@ -42,47 +57,71 @@ export function useSessionRpc) : ((payload as unknown) as TData); + 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); - activeRequestIdRef.current = null; + clearActiveRequest(); }); return () => { unsubscribe(); }; - }, [select, type, ws]); + }, [clearActiveRequest, select, type, ws]); + + useEffect(() => { + if (ws.isConnected || !activeRequestIdRef.current) { + return; + } + const error = new Error("WebSocket disconnected"); + setState({ status: "error", requestId: activeRequestIdRef.current, error }); + rejectRef.current?.(error); + clearActiveRequest(); + }, [clearActiveRequest, ws.isConnected]); const send = useCallback( (params: Omit, "type" | "requestId">) => { - const requestId = generateMessageId(); - activeRequestIdRef.current = requestId; - setState({ status: "loading", requestId }); - - const request = { - type, - ...params, - requestId, - } as RequestOf; - - ws.send({ type: "session", message: request }); - return new Promise((resolve, reject) => { + if (!ws.isConnected) { + const error = new Error("WebSocket is disconnected"); + setState({ status: "error", requestId: null, error }); + reject(error); + return; + } + + const requestId = generateMessageId(); + activeRequestIdRef.current = requestId; resolveRef.current = resolve; rejectRef.current = reject; + setState({ status: "loading", requestId }); + + const request = { + type, + ...params, + requestId, + } as RequestOf; + + try { + ws.send({ type: "session", message: request }); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + setState({ status: "error", requestId: requestId, error: err }); + reject(err); + clearActiveRequest(); + } }); }, - [type, ws] + [clearActiveRequest, type, ws] ); const reset = useCallback(() => { - activeRequestIdRef.current = null; + clearActiveRequest(); setState({ status: "idle", requestId: null }); - }, []); + }, [clearActiveRequest]); return useMemo(() => ({ state, send, reset }), [reset, send, state]); }