All files / src/hooks/notifications useNotificationConnection.ts

81.63% Statements 40/49
81.81% Branches 9/11
40% Functions 2/5
89.47% Lines 34/38

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 87 88 89 90                              103x 103x 103x 103x 103x 103x 103x 103x 103x 103x 103x       881x 881x 881x       881x 881x     881x     287438x       287438x 287438x   287438x   5662x   9187x         7847x       1340x 538400x     287438x 1989x         287438x           2150x   289588x 291738x       287438x   292563x    
// 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 { useFlowsheetId } 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",
  RESULT_SUMMARY_CALCULATING = "result-summary/calculating",
  RESULT_SUMMARY_AVAILABLE = "result-summary/available",
  RESULT_SUMMARY_FAILED = "result-summary/failed",
  AUTO_REVISION_SKIPPED_ACTIVE_TASKS = "flowsheet-version/auto-save-skipped-active-tasks",
}
 
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 = useFlowsheetId();
 
  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 };
}