All files / src/hooks/notifications useNotificationConnection.ts

80% Statements 36/45
81.81% Branches 9/11
40% Functions 2/5
88.23% Lines 30/34

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86                              76x 76x 76x 76x 76x 76x 76x       217x 217x 217x       217x 217x     217x     139359x       139359x 139359x   139359x   2840x   6130x         4885x       1245x 337985x     139359x 1102x         139359x           1102x   140461x 141563x       139359x   141789x    
// Vite 8 no longer unwraps this package's CommonJS default export reliably at runtime,
// so import the concrete hook entrypoint directly while keeping the same hook API.
 
import { useMemo } from "react";
import type { Options } from "react-use-websocket/dist/lib/types";
import { useWebSocket } from "react-use-websocket/dist/lib/use-websocket";
import { useEventCallback } from "usehooks-ts";
import { resolveNotificationEndpointUrl } from "@/api/apiBaseUrl.ts";
import { useProjectId } from "@/hooks/project.ts";
 
export type NotificationMessage = {
  message_type: NotificationMessageType;
  data: never;
};
 
export enum NotificationMessageType {
  TASK_COMPLETED = "task/completed",
  TASK_UPDATED = "task/updated",
  TASK_CANCELLED = "task/cancelled",
  TASK_CANCELLING = "task/cancelling",
  TASK_CHILDREN_CANCELLED = "task/children-cancelled",
  BUILD_STATE_COMPLETED = "build-state/completed",
}
 
function retryHandler(retryCount: number) {
  const baseWaitPeriod = 250; // 0.25 seconds
  const maxWaitPeriod = 10000; // 10 seconds
  const exponentialWaitPeriod = Math.min(
    Math.pow(2, retryCount) * baseWaitPeriod,
    maxWaitPeriod,
  );
  const jitter = Math.random() * 0.1 * exponentialWaitPeriod; // Random jitter of up to 10% of the wait period
  const sign = Math.random() < 0.5 ? -1 : 1; // Randomly choose the sign of the jitter
 
  // Wait = exponentialWaitPeriod +/- 10%
  return exponentialWaitPeriod + sign * jitter;
}
 
export function useNotificationConnection(
  onMessage: (message: never) => void,
  subscribedMessageType?: NotificationMessageType,
) {
  const socketUrl = resolveNotificationEndpointUrl();
  const projectId = useProjectId();
 
  const messageHandler = useEventCallback((message: MessageEvent) => {
    // Ping/pong messages are handled by the library
    Iif (message.data === "pong") return;
 
    const notificationMessage = JSON.parse(message.data) as NotificationMessage;
 
    // Only deliver messages that have been subscribed to
    Iif (
      subscribedMessageType &&
      notificationMessage.message_type !== subscribedMessageType
    )
      return;
 
    onMessage(notificationMessage.data);
  });
 
  const options: Options = useMemo(
    () => ({
      heartbeat: {
        message: "ping",
        returnMessage: "pong",
        timeout: 10000, // 10 seconds
        interval: 5000, // 5 seconds
      },
      onMessage: messageHandler,
      share: true,
      shouldReconnect: () => true, // Always reconnect
      reconnectInterval: retryHandler,
      retryOnError: true,
      queryParams: {
        flowsheetId: projectId,
      },
    }),
    [messageHandler, projectId],
  );
 
  const { readyState } = useWebSocket(socketUrl, options);
 
  return { readyState };
}