444 use websocket to trigger reloads (#458)

* [#444] added refresh on sync token calendar update

* [#444] added refresh on register and tests

* [#444] extracted logic for useEffects and updates

* [#444] refactored code structure to use absolute paths

src/websocket/WebSocketGate.tsx

@@ -1,103 +1,79 @@
-import { useEffect, useRef, useState } from "react";
-import { useAppSelector } from "../app/hooks";
-import { useSelectedCalendars } from "../utils/storage/useSelectedCalendars";
-import {
-  createWebSocketConnection,
-  WebSocketWithCleanup,
-} from "./createWebSocketConnection";
-import { registerToCalendars } from "./ws/registerToCalendars";
-import { unregisterToCalendars } from "./ws/unregisterToCalendars";
+import { useAppDispatch, useAppSelector } from "@/app/hooks";
+import { useSelectedCalendars } from "@/utils/storage/useSelectedCalendars";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import type { WebSocketWithCleanup } from "./connection";
+import { closeWebSocketConnection } from "./connection/lifecycle/closeWebSocketConnection";
+import { establishWebSocketConnection } from "./connection/lifecycle/establishWebSocketConnection";
+import { updateCalendars } from "./messaging";
+import { syncCalendarRegistrations } from "./operations";
 
 export function WebSocketGate() {
   const socketRef = useRef<WebSocketWithCleanup | null>(null);
   const previousCalendarListRef = useRef<string[]>([]);
-
+  const dispatch = useAppDispatch();
   const isAuthenticated = useAppSelector((state) =>
     Boolean(state.user.userData && state.user.tokens)
   );
+
   const [isSocketOpen, setIsSocketOpen] = useState(false);
 
   const calendarList = useSelectedCalendars();
+  const onMessage = useCallback(
+    (message: unknown) => {
+      updateCalendars(message, dispatch);
+    },
+    [dispatch]
+  );
+
+  const onClose = useCallback((event: CloseEvent) => {
+    // Socket already cleaned up by internal handler before this callback fires
+    socketRef.current = null;
+    setIsSocketOpen(false);
+    // TODO: Add reconnection logic here
+  }, []);
+
+  const onError = useCallback((error: Event) => {
+    console.error("WebSocket error:", error);
+  }, []);
+
+  const callBacks = useMemo(
+    () => ({
+      onMessage,
+      onClose,
+      onError,
+    }),
+    [onMessage, onClose, onError]
+  );
 
   // Manage WebSocket connection
   useEffect(() => {
+    const abortController = new AbortController();
     if (!isAuthenticated) {
-      if (socketRef.current) {
-        socketRef.current.cleanup();
-        socketRef.current.close();
-        socketRef.current = null;
-        setIsSocketOpen(false);
-      }
+      closeWebSocketConnection(socketRef, setIsSocketOpen);
       return;
     }
 
-    const connect = async () => {
-      try {
-        const socket = await createWebSocketConnection();
-        socketRef.current = socket;
-        socket.addEventListener("close", () => {
-          setIsSocketOpen(false);
-          socketRef.current = null;
-        });
-        // Check if socket closed during setup
-        if (socket.readyState === WebSocket.OPEN) {
-          setIsSocketOpen(true);
-        }
-      } catch (error) {
-        console.error("Failed to create WebSocket connection:", error);
-      }
-    };
-
-    connect();
+    establishWebSocketConnection(
+      callBacks,
+      socketRef,
+      setIsSocketOpen,
+      abortController.signal
+    );
 
     return () => {
-      if (socketRef.current) {
-        socketRef.current.cleanup();
-        socketRef.current.close();
-        socketRef.current = null;
-        setIsSocketOpen(false);
-      }
+      abortController.abort();
+      closeWebSocketConnection(socketRef, setIsSocketOpen);
     };
-  }, [isAuthenticated]);
+  }, [isAuthenticated, callBacks]);
 
   // Register using a diff with previous calendars
   useEffect(() => {
-    if (
-      !isSocketOpen ||
-      !socketRef.current ||
-      socketRef.current.readyState !== WebSocket.OPEN
-    )
-      return;
-
-    const currentPaths = calendarList.map((cal) => `/calendars/${cal}`);
-    const previousPaths = previousCalendarListRef.current.map(
-      (cal) => `/calendars/${cal}`
+    syncCalendarRegistrations(
+      isSocketOpen,
+      socketRef,
+      calendarList,
+      previousCalendarListRef
     );
-
-    // calendars to register
-    const toRegister = currentPaths.filter(
-      (path) => !previousPaths.includes(path)
-    );
-
-    // calendars to unregister
-    const toUnregister = previousPaths.filter(
-      (path) => !currentPaths.includes(path)
-    );
-
-    try {
-      if (toRegister.length > 0) {
-        registerToCalendars(socketRef.current, toRegister);
-      }
-
-      if (toUnregister.length > 0) {
-        unregisterToCalendars(socketRef.current, toUnregister);
-      }
-
-      // Only update the ref if operations succeeded
-      previousCalendarListRef.current = calendarList;
-    } catch (error) {
-      console.error("Failed to update calendar registrations:", error);
-    }
   }, [isSocketOpen, calendarList]);
 
   return null;

src/websocket/connection/createConnection.ts

@@ -1,16 +1,15 @@
-import { fetchWebSocketTicket } from "./api/fetchWebSocketTicket";
-import { WS_INBOUND_EVENTS } from "./protocols";
+import { fetchWebSocketTicket } from "../api/fetchWebSocketTicket";
+import { WS_INBOUND_EVENTS } from "../protocols";
+import { WebSocketCallbacks, WebSocketWithCleanup } from "./types";
 
-export interface WebSocketWithCleanup extends WebSocket {
-  cleanup: () => void;
-}
-
-export async function createWebSocketConnection(): Promise<WebSocketWithCleanup> {
+export async function createWebSocketConnection(
+  callbacks: WebSocketCallbacks
+): Promise<WebSocketWithCleanup> {
   const wsBaseUrl =
     (window as any).WEBSOCKET_URL ??
     (window as any).CALENDAR_BASE_URL?.replace(
       /^http(s)?:/,
-      (_: boolean, s: boolean) => (s ? "wss:" : "ws:")
+      (_: string, s: string | undefined) => (s ? "wss:" : "ws:")
     ) ??
     "";
 
@@ -36,6 +35,7 @@ export async function createWebSocketConnection(): Promise<WebSocketWithCleanup>
       socket.close();
       reject(new Error("WebSocket connection timed out"));
     }, CONNECTION_TIMEOUT_MS);
+
     const openHandler = () => {
       console.log("WebSocket connection opened");
       clearTimeout(timeoutId);
@@ -66,9 +66,7 @@ export async function createWebSocketConnection(): Promise<WebSocketWithCleanup>
   const messageHandler = (event: MessageEvent) => {
     try {
       const message = JSON.parse(event.data);
-      console.log("WebSocket message received:", message);
-
-      // TODO: Handle different message types
+      callbacks.onMessage(message);
     } catch (error) {
       console.error("Failed to parse WebSocket message:", error);
     }
@@ -76,13 +74,13 @@ export async function createWebSocketConnection(): Promise<WebSocketWithCleanup>
 
   const closeHandler = (event: CloseEvent) => {
     console.log("WebSocket closed:", event.code, event.reason);
-    // Clean up all event listeners when socket closes
     cleanup();
-    // TODO: Add reconnection logic
+    callbacks.onClose?.(event);
   };
 
   const errorHandler = (error: Event) => {
     console.error("WebSocket error:", error);
+    callbacks.onError?.(error);
   };
 
   // Cleanup function to remove all event listeners

src/websocket/connection/index.ts

@@ -0,0 +1,2 @@
+export { createWebSocketConnection } from "./createConnection";
+export type { WebSocketWithCleanup, WebSocketCallbacks } from "./types";

src/websocket/connection/lifecycle/closeWebSocketConnection.ts

@@ -0,0 +1,13 @@
+import { WebSocketWithCleanup } from "../types";
+
+export function closeWebSocketConnection(
+  socketRef: React.MutableRefObject<WebSocketWithCleanup | null>,
+  setIsSocketOpen: (value: boolean) => void
+) {
+  if (socketRef.current) {
+    socketRef.current.cleanup();
+    socketRef.current.close();
+    socketRef.current = null;
+    setIsSocketOpen(false);
+  }
+}

src/websocket/connection/lifecycle/establishWebSocketConnection.ts

@@ -0,0 +1,28 @@
+import { createWebSocketConnection } from "../createConnection";
+import { WebSocketCallbacks, WebSocketWithCleanup } from "../types";
+
+export async function establishWebSocketConnection(
+  callbacks: WebSocketCallbacks,
+  socketRef: React.MutableRefObject<WebSocketWithCleanup | null>,
+  setIsSocketOpen: (value: boolean) => void,
+  signal?: AbortSignal
+) {
+  try {
+    const socket = await createWebSocketConnection(callbacks);
+
+    if (signal?.aborted) {
+      socket.cleanup();
+      socket.close();
+      return;
+    }
+
+    socketRef.current = socket;
+
+    if (socket.readyState === WebSocket.OPEN) {
+      setIsSocketOpen(true);
+    }
+  } catch (error) {
+    console.error("Failed to create WebSocket connection:", error);
+    setIsSocketOpen(false);
+  }
+}

src/websocket/connection/types.ts

@@ -0,0 +1,9 @@
+export interface WebSocketWithCleanup extends WebSocket {
+  cleanup: () => void;
+}
+
+export interface WebSocketCallbacks {
+  onMessage: (data: unknown) => void;
+  onClose?: (event: CloseEvent) => void;
+  onError?: (error: Event) => void;
+}

src/websocket/index.ts

@@ -0,0 +1,2 @@
+export { WebSocketGate } from "./WebSocketGate";
+export type { WebSocketWithCleanup, WebSocketCallbacks } from "./connection";

src/websocket/messaging/index.ts

@@ -0,0 +1,3 @@
+export { parseMessage } from "./parseMessage";
+export { updateCalendars } from "./updateCalendars";
+export { parseCalendarPath } from "./parseCalendarPath";

src/websocket/messaging/parseCalendarPath.ts

@@ -0,0 +1,11 @@
+const CALENDAR_PATH_REGEX = /^\/calendars\/[a-zA-Z0-9_-]+\/[a-zA-Z0-9_-]+$/;
+
+export function parseCalendarPath(key: string) {
+  if (!CALENDAR_PATH_REGEX.test(key)) {
+    return null;
+  }
+
+  const [, , userId, calendarId] = key.split("/");
+
+  return `${userId}/${calendarId}`;
+}

src/websocket/messaging/parseMessage.ts

@@ -0,0 +1,30 @@
+import { WS_INBOUND_EVENTS } from "../protocols";
+
+export function parseMessage(message: unknown) {
+  console.log("WebSocket message received:", message);
+  const calendarsToRefresh = new Set<string>();
+  const calendarsToHide = new Set<string>();
+  if (typeof message !== "object" || message === null) {
+    return { calendarsToRefresh, calendarsToHide };
+  }
+
+  for (const [key, value] of Object.entries(message)) {
+    switch (key) {
+      case WS_INBOUND_EVENTS.CLIENT_REGISTERED:
+        if (Array.isArray(value)) {
+          value.forEach((cal: string) => calendarsToRefresh.add(cal));
+        }
+        break;
+      case WS_INBOUND_EVENTS.CLIENT_UNREGISTERED:
+        if (Array.isArray(value)) {
+          value.forEach((cal: string) => calendarsToHide.add(cal));
+        }
+        break;
+      default: {
+        calendarsToRefresh.add(key);
+      }
+    }
+  }
+
+  return { calendarsToRefresh, calendarsToHide };
+}

src/websocket/messaging/updateCalendars.ts

@@ -0,0 +1,45 @@
+import type { AppDispatch } from "@/app/store";
+import { store } from "@/app/store";
+import { refreshCalendarWithSyncToken } from "@/features/Calendars/services/refreshCalendar";
+import { findCalendarById, getDisplayedCalendarRange } from "@/utils";
+import { setSelectedCalendars } from "@/utils/storage/setSelectedCalendars";
+import { parseCalendarPath } from "./parseCalendarPath";
+import { parseMessage } from "./parseMessage";
+
+export function updateCalendars(message: unknown, dispatch: AppDispatch) {
+  const currentRange = getDisplayedCalendarRange();
+  const state = store.getState();
+  const { calendarsToRefresh, calendarsToHide } = parseMessage(message);
+  calendarsToRefresh.forEach((calendarPath) => {
+    const calendarId = parseCalendarPath(calendarPath);
+    if (!calendarId) {
+      console.warn("Invalid calendar path received:", calendarPath);
+      return;
+    }
+    const calendar = findCalendarById(state, calendarId);
+    if (calendar) {
+      dispatch(
+        refreshCalendarWithSyncToken({
+          calendar: calendar.calendar,
+          calType: calendar.type,
+          calendarRange: currentRange,
+        })
+      );
+    } else {
+      console.warn("Calendar not found for id:", calendarId);
+    }
+  });
+  const currentSelectedCalendars = JSON.parse(
+    localStorage.getItem("selectedCalendars") ?? "[]"
+  ) as string[];
+
+  const calendarIdsToHide = [...calendarsToHide]
+    .map(parseCalendarPath)
+    .filter((id): id is string => Boolean(id));
+
+  const updatedSelectedCalendars = currentSelectedCalendars.filter(
+    (id) => !calendarIdsToHide.includes(id)
+  );
+
+  setSelectedCalendars(updatedSelectedCalendars);
+}

src/websocket/operations/index.ts

@@ -0,0 +1,3 @@
+export { registerToCalendars } from "./registerToCalendars";
+export { unregisterToCalendars } from "./unregisterToCalendars";
+export { syncCalendarRegistrations } from "./syncCalendarRegistrations";

src/websocket/operations/syncCalendarRegistrations.ts

@@ -0,0 +1,48 @@
+import { WebSocketWithCleanup } from "../connection";
+import { registerToCalendars } from "./registerToCalendars";
+import { unregisterToCalendars } from "./unregisterToCalendars";
+
+export function syncCalendarRegistrations(
+  isSocketOpen: boolean,
+  socketRef: React.MutableRefObject<WebSocketWithCleanup | null>,
+  calendarList: string[],
+  previousCalendarListRef: React.MutableRefObject<string[]>
+) {
+  if (
+    !isSocketOpen ||
+    !socketRef.current ||
+    socketRef.current.readyState !== WebSocket.OPEN
+  ) {
+    return;
+  }
+
+  const currentPaths = calendarList.map((cal) => `/calendars/${cal}`);
+  const previousPaths = previousCalendarListRef.current.map(
+    (cal) => `/calendars/${cal}`
+  );
+
+  const toRegister = currentPaths.filter(
+    (path) => !previousPaths.includes(path)
+  );
+  const toUnregister = previousPaths.filter(
+    (path) => !currentPaths.includes(path)
+  );
+
+  try {
+    if (toRegister.length > 0) {
+      registerToCalendars(socketRef.current, toRegister);
+    }
+  } catch (error) {
+    console.error("Failed to register calendar:", error);
+    return;
+  }
+  try {
+    if (toUnregister.length > 0) {
+      unregisterToCalendars(socketRef.current, toUnregister);
+    }
+  } catch (error) {
+    console.error("Failed to unregister calendar:", error);
+    return;
+  }
+  previousCalendarListRef.current = calendarList;
+}

src/websocket/protocols.ts

@@ -1,13 +0,0 @@
-// WebSocket event listeners (browser events)
-export const WS_INBOUND_EVENTS = {
-  CONNECTION_OPENED: "open",
-  MESSAGE: "message",
-  ERROR: "error",
-  CONNECTION_CLOSED: "close",
-} as const;
-
-// WebSocket message types sent to server
-export const WS_OUTBOUND_EVENTS = {
-  REGISTER_CLIENT: "register",
-  UNREGISTER_CLIENT: "unregister",
-} as const;

src/websocket/protocols/index.ts

@@ -0,0 +1,8 @@
+export const WS_INBOUND_EVENTS = {
+  CONNECTION_OPENED: "open",
+  MESSAGE: "message",
+  ERROR: "error",
+  CONNECTION_CLOSED: "close",
+  CLIENT_REGISTERED: "registered",
+  CLIENT_UNREGISTERED: "unregistered",
+} as const;