refactor: More robust handling of jsonrpc calls (#915)

Co-authored-by: Marc Brooks <IDisposable@gmail.com>

Author: adamshiervani

ui/src/hooks/stores.ts

@@ -116,7 +116,7 @@ export interface RTCState {
   peerConnection: RTCPeerConnection | null;
   setPeerConnection: (pc: RTCState["peerConnection"]) => void;
 
-  setRpcDataChannel: (channel: RTCDataChannel) => void;
+  setRpcDataChannel: (channel: RTCDataChannel | null) => void;
   rpcDataChannel: RTCDataChannel | null;
 
   hidRpcDisabled: boolean;
@@ -178,79 +178,80 @@ export const useRTCStore = create<RTCState>(set => ({
   setPeerConnection: (pc: RTCState["peerConnection"]) => set({ peerConnection: pc }),
 
   rpcDataChannel: null,
-  setRpcDataChannel: (channel: RTCDataChannel) => set({ rpcDataChannel: channel }),
+  setRpcDataChannel: channel => set({ rpcDataChannel: channel }),
 
   hidRpcDisabled: false,
-  setHidRpcDisabled: (disabled: boolean) => set({ hidRpcDisabled: disabled }),
+  setHidRpcDisabled: disabled => set({ hidRpcDisabled: disabled }),
 
   rpcHidProtocolVersion: null,
-  setRpcHidProtocolVersion: (version: number | null) => set({ rpcHidProtocolVersion: version }),
+  setRpcHidProtocolVersion: version => set({ rpcHidProtocolVersion: version }),
 
   rpcHidChannel: null,
-  setRpcHidChannel: (channel: RTCDataChannel) => set({ rpcHidChannel: channel }),
+  setRpcHidChannel: channel => set({ rpcHidChannel: channel }),
 
   rpcHidUnreliableChannel: null,
-  setRpcHidUnreliableChannel: (channel: RTCDataChannel) => set({ rpcHidUnreliableChannel: channel }),
+  setRpcHidUnreliableChannel: channel => set({ rpcHidUnreliableChannel: channel }),
 
   rpcHidUnreliableNonOrderedChannel: null,
-  setRpcHidUnreliableNonOrderedChannel: (channel: RTCDataChannel) => set({ rpcHidUnreliableNonOrderedChannel: channel }),
+  setRpcHidUnreliableNonOrderedChannel: channel =>
+    set({ rpcHidUnreliableNonOrderedChannel: channel }),
 
   transceiver: null,
-  setTransceiver: (transceiver: RTCRtpTransceiver) => set({ transceiver }),
+  setTransceiver: transceiver => set({ transceiver }),
 
   peerConnectionState: null,
-  setPeerConnectionState: (state: RTCPeerConnectionState) => set({ peerConnectionState: state }),
+  setPeerConnectionState: state => set({ peerConnectionState: state }),
 
   mediaStream: null,
-  setMediaStream: (stream: MediaStream) => set({ mediaStream: stream }),
+  setMediaStream: stream => set({ mediaStream: stream }),
 
   videoStreamStats: null,
-  appendVideoStreamStats: (stats: RTCInboundRtpStreamStats) => set({ videoStreamStats: stats }),
+  appendVideoStreamStats: stats => set({ videoStreamStats: stats }),
   videoStreamStatsHistory: new Map(),
 
   isTurnServerInUse: false,
-  setTurnServerInUse: (inUse: boolean) => set({ isTurnServerInUse: inUse }),
+  setTurnServerInUse: inUse => set({ isTurnServerInUse: inUse }),
 
   inboundRtpStats: new Map(),
-  appendInboundRtpStats: (stats: RTCInboundRtpStreamStats) => {
+  appendInboundRtpStats: stats => {
     set(prevState => ({
       inboundRtpStats: appendStatToMap(stats, prevState.inboundRtpStats),
     }));
   },
   clearInboundRtpStats: () => set({ inboundRtpStats: new Map() }),
 
   candidatePairStats: new Map(),
-  appendCandidatePairStats: (stats: RTCIceCandidatePairStats) => {
+  appendCandidatePairStats: stats => {
     set(prevState => ({
       candidatePairStats: appendStatToMap(stats, prevState.candidatePairStats),
     }));
   },
   clearCandidatePairStats: () => set({ candidatePairStats: new Map() }),
 
   localCandidateStats: new Map(),
-  appendLocalCandidateStats: (stats: RTCIceCandidateStats) => {
+  appendLocalCandidateStats: stats => {
     set(prevState => ({
       localCandidateStats: appendStatToMap(stats, prevState.localCandidateStats),
     }));
   },
 
   remoteCandidateStats: new Map(),
-  appendRemoteCandidateStats: (stats: RTCIceCandidateStats) => {
+  appendRemoteCandidateStats: stats => {
     set(prevState => ({
       remoteCandidateStats: appendStatToMap(stats, prevState.remoteCandidateStats),
     }));
   },
 
   diskDataChannelStats: new Map(),
-  appendDiskDataChannelStats: (stats: RTCDataChannelStats) => {
+  appendDiskDataChannelStats: stats => {
     set(prevState => ({
       diskDataChannelStats: appendStatToMap(stats, prevState.diskDataChannelStats),
     }));
   },
 
   // Add these new properties to the store implementation
   terminalChannel: null,
-  setTerminalChannel: (channel: RTCDataChannel) => set({ terminalChannel: channel }),
+  setTerminalChannel: channel => set({ terminalChannel: channel }),
 }));
 
 export interface MouseMove {
@@ -270,12 +271,20 @@ export interface MouseState {
 export const useMouseStore = create<MouseState>(set => ({
   mouseX: 0,
   mouseY: 0,
-  setMouseMove: (move?: MouseMove) => set({ mouseMove: move }),
-  setMousePosition: (x: number, y: number) => set({ mouseX: x, mouseY: y }),
+  setMouseMove: move => set({ mouseMove: move }),
+  setMousePosition: (x, y) => set({ mouseX: x, mouseY: y }),
 }));
 
-export type HdmiStates = "ready" | "no_signal" | "no_lock" | "out_of_range" | "connecting";
-export type HdmiErrorStates = Extract<VideoState["hdmiState"], "no_signal" | "no_lock" | "out_of_range">
+export type HdmiStates =
+  | "ready"
+  | "no_signal"
+  | "no_lock"
+  | "out_of_range"
+  | "connecting";
+export type HdmiErrorStates = Extract<
+  VideoState["hdmiState"],
+  "no_signal" | "no_lock" | "out_of_range"
+>;
 
 export interface HdmiState {
   ready: boolean;
@@ -290,10 +299,7 @@ export interface VideoState {
   setClientSize: (width: number, height: number) => void;
   setSize: (width: number, height: number) => void;
   hdmiState: HdmiStates;
-  setHdmiState: (state: {
-    ready: boolean;
-    error?: HdmiErrorStates;
-  }) => void;
+  setHdmiState: (state: { ready: boolean; error?: HdmiErrorStates }) => void;
 }
 
 export const useVideoStore = create<VideoState>(set => ({
@@ -304,7 +310,8 @@ export const useVideoStore = create<VideoState>(set => ({
   clientHeight: 0,
 
   // The video element's client size
-  setClientSize: (clientWidth: number, clientHeight: number) => set({ clientWidth, clientHeight }),
+  setClientSize: (clientWidth: number, clientHeight: number) =>
+    set({ clientWidth, clientHeight }),
 
   // Resolution
   setSize: (width: number, height: number) => set({ width, height }),
@@ -451,13 +458,15 @@ export interface MountMediaState {
 
 export const useMountMediaStore = create<MountMediaState>(set => ({
   remoteVirtualMediaState: null,
-  setRemoteVirtualMediaState: (state: MountMediaState["remoteVirtualMediaState"]) => set({ remoteVirtualMediaState: state }),
+  setRemoteVirtualMediaState: (state: MountMediaState["remoteVirtualMediaState"]) =>
+    set({ remoteVirtualMediaState: state }),
 
   modalView: "mode",
   setModalView: (view: MountMediaState["modalView"]) => set({ modalView: view }),
 
   isMountMediaDialogOpen: false,
-  setIsMountMediaDialogOpen: (isOpen: MountMediaState["isMountMediaDialogOpen"]) => set({ isMountMediaDialogOpen: isOpen }),
+  setIsMountMediaDialogOpen: (isOpen: MountMediaState["isMountMediaDialogOpen"]) =>
+    set({ isMountMediaDialogOpen: isOpen }),
 
   uploadedFiles: [],
   addUploadedFile: (file: { name: string; size: string; uploadedAt: string }) =>
@@ -474,7 +483,7 @@ export interface KeyboardLedState {
   compose: boolean;
   kana: boolean;
   shift: boolean; // Optional, as not all keyboards have a shift LED
-};
+}
 
 export const hidKeyBufferSize = 6;
 export const hidErrorRollOver = 0x01;
@@ -509,14 +518,23 @@ export interface HidState {
 }
 
 export const useHidStore = create<HidState>(set => ({
-  keyboardLedState: { num_lock: false, caps_lock: false, scroll_lock: false, compose: false, kana: false, shift: false } as KeyboardLedState,
-  setKeyboardLedState: (ledState: KeyboardLedState): void => set({ keyboardLedState: ledState }),
+  keyboardLedState: {
+    num_lock: false,
+    caps_lock: false,
+    scroll_lock: false,
+    compose: false,
+    kana: false,
+    shift: false,
+  } as KeyboardLedState,
+  setKeyboardLedState: (ledState: KeyboardLedState): void =>
+    set({ keyboardLedState: ledState }),
 
   keysDownState: { modifier: 0, keys: [0, 0, 0, 0, 0, 0] } as KeysDownState,
   setKeysDownState: (state: KeysDownState): void => set({ keysDownState: state }),
 
   isVirtualKeyboardEnabled: false,
-  setVirtualKeyboardEnabled: (enabled: boolean): void => set({ isVirtualKeyboardEnabled: enabled }),
+  setVirtualKeyboardEnabled: (enabled: boolean): void =>
+    set({ isVirtualKeyboardEnabled: enabled }),
 
   isPasteInProgress: false,
   setPasteModeEnabled: (enabled: boolean): void => set({ isPasteInProgress: enabled }),
@@ -568,7 +586,7 @@ export interface OtaState {
 
   systemUpdateProgress: number;
   systemUpdatedAt: string | null;
-};
+}
 
 export interface UpdateState {
   isUpdatePending: boolean;
@@ -580,7 +598,7 @@ export interface UpdateState {
   otaState: OtaState;
   setOtaState: (state: OtaState) => void;
 
-  modalView: UpdateModalViews
+  modalView: UpdateModalViews;
   setModalView: (view: UpdateModalViews) => void;
 
   updateErrorMessage: string | null;
@@ -620,12 +638,11 @@ export const useUpdateStore = create<UpdateState>(set => ({
   setModalView: (view: UpdateModalViews) => set({ modalView: view }),
 
   updateErrorMessage: null,
-  setUpdateErrorMessage: (errorMessage: string) => set({ updateErrorMessage: errorMessage }),
+  setUpdateErrorMessage: (errorMessage: string) =>
+    set({ updateErrorMessage: errorMessage }),
 }));
 
-export type UsbConfigModalViews =
-  | "updateUsbConfig"
-  | "updateUsbConfigSuccess";
+export type UsbConfigModalViews = "updateUsbConfig" | "updateUsbConfigSuccess";
 
 export interface UsbConfigModalState {
   modalView: UsbConfigModalViews;
@@ -833,12 +850,12 @@ export interface MacrosState {
   loadMacros: () => Promise<void>;
   saveMacros: (macros: KeySequence[]) => Promise<void>;
   sendFn:
-  | ((
-    method: string,
-    params: unknown,
-    callback?: ((resp: JsonRpcResponse) => void) | undefined,
-  ) => void)
-  | null;
+    | ((
+        method: string,
+        params: unknown,
+        callback?: ((resp: JsonRpcResponse) => void) | undefined,
+      ) => void)
+    | null;
   setSendFn: (
     sendFn: (
       method: string,
@@ -978,5 +995,5 @@ export const useMacrosStore = create<MacrosState>((set, get) => ({
     } finally {
       set({ loading: false });
     }
-  }
+  },
 }));

ui/src/routes/devices.$id.tsx

@@ -557,8 +557,9 @@ export default function KvmIdRoute() {
       clearCandidatePairStats();
       setSidebarView(null);
       setPeerConnection(null);
+      setRpcDataChannel(null);
     };
-  }, [clearCandidatePairStats, clearInboundRtpStats, setPeerConnection, setSidebarView]);
+  }, [clearCandidatePairStats, clearInboundRtpStats, setPeerConnection, setSidebarView, setRpcDataChannel]);
 
   // TURN server usage detection
   useEffect(() => {

ui/src/utils/jsonrpc.ts

@@ -24,17 +24,47 @@ export interface JsonRpcCallResponse<T = unknown> {
 let rpcCallCounter = 0;
 
 // Helper: wait for RTC data channel to be ready
+// This waits indefinitely for the channel to be ready, only aborting via the signal
+// Throws if the channel instance changed while waiting (stale connection detected)
 async function waitForRtcReady(signal: AbortSignal): Promise<RTCDataChannel> {
   const pollInterval = 100;
+  let lastSeenChannel: RTCDataChannel | null = null;
 
   while (!signal.aborted) {
     const state = useRTCStore.getState();
-    if (state.rpcDataChannel?.readyState === "open") {
-      return state.rpcDataChannel;
+    const currentChannel = state.rpcDataChannel;
+
+    // Channel instance changed (new connection replaced old one)
+    if (lastSeenChannel && currentChannel && lastSeenChannel !== currentChannel) {
+      console.debug("[waitForRtcReady] Channel instance changed, aborting wait");
+      throw new Error("RTC connection changed while waiting for readiness");
+    }
+
+    // Channel was removed from store (connection closed)
+    if (lastSeenChannel && !currentChannel) {
+      console.debug("[waitForRtcReady] Channel was removed from store, aborting wait");
+      throw new Error("RTC connection was closed while waiting for readiness");
+    }
+
+    // No channel yet, keep waiting
+    if (!currentChannel) {
+      await sleep(pollInterval);
+      continue;
+    }
+
+    // Track this channel instance
+    lastSeenChannel = currentChannel;
+
+    // Channel is ready!
+    if (currentChannel.readyState === "open") {
+      return currentChannel;
     }
+
     await sleep(pollInterval);
   }
 
+  // Signal was aborted for some reason
+  console.debug("[waitForRtcReady] Aborted via signal");
   throw new Error("RTC readiness check aborted");
 }
 
@@ -97,25 +127,26 @@ export async function callJsonRpc<T = unknown>(
   const timeout = options.attemptTimeoutMs || 5000;
 
   for (let attempt = 0; attempt < maxAttempts; attempt++) {
-    const abortController = new AbortController();
-    const timeoutId = setTimeout(() => abortController.abort(), timeout);
-
     // Exponential backoff for retries that starts at 500ms up to a maximum of 10 seconds
     const backoffMs = Math.min(500 * Math.pow(2, attempt), 10000);
+    let timeoutId: ReturnType<typeof setTimeout> | null = null;
 
     try {
-      // Wait for RTC readiness
-      const rpcDataChannel = await waitForRtcReady(abortController.signal);
+      // Wait for RTC readiness without timeout - this allows time for WebRTC to connect
+      const readyAbortController = new AbortController();
+      const rpcDataChannel = await waitForRtcReady(readyAbortController.signal);
+
+      // Now apply timeout only to the actual RPC request/response
+      const rpcAbortController = new AbortController();
+      timeoutId = setTimeout(() => rpcAbortController.abort(), timeout);
 
       // Send RPC request and wait for response
       const response = await sendRpcRequest<T>(
         rpcDataChannel,
         options,
-        abortController.signal,
+        rpcAbortController.signal,
       );
 
-      clearTimeout(timeoutId);
-
       // Retry on error if attempts remain
       if (response.error && attempt < maxAttempts - 1) {
         await sleep(backoffMs);
@@ -124,8 +155,6 @@ export async function callJsonRpc<T = unknown>(
 
       return response;
     } catch (error) {
-      clearTimeout(timeoutId);
-
       // Retry on timeout/error if attempts remain
       if (attempt < maxAttempts - 1) {
         await sleep(backoffMs);
@@ -135,6 +164,10 @@ export async function callJsonRpc<T = unknown>(
       throw error instanceof Error
         ? error
         : new Error(`JSON-RPC call failed after ${timeout}ms`);
+    } finally {
+      if (timeoutId !== null) {
+        clearTimeout(timeoutId);
+      }
     }
   }