new

src/models/useSSE.ts

@@ -0,0 +1,232 @@
+import {
+  EventStreamContentType,
+  fetchEventSource,
+} from '@microsoft/fetch-event-source';
+import { useCallback, useRef, useState } from 'react';
+import { authFaiedToLoginPage } from '@/access';
+export type SSEStatus =
+  | 'idle'
+  | 'connecting'
+  | 'open'
+  | 'error'
+  | 'closed'
+  | 'replaced';
+interface SSEOptions {
+  method?: 'GET' | 'POST';
+  headers?: Record<string, string>;
+  body?: any;
+  maxRetries?: number;
+  retryInterval?: number;
+  openWhenHidden?: boolean;
+  onMessage?: (data: string) => void;
+  onError?: (error: any) => void;
+  onStatusChange?: (status: SSEStatus) => void;
+}
+
+export default () => {
+  const [status, setStatus] = useState<SSEStatus>('idle');
+  const [retryCount, setRetryCount] = useState(0);
+
+  // Use ref to synchronize internal logic to avoid logical judgment errors caused by status update delays.
+  const internalCountRef = useRef(0);
+  const abortControllerRef = useRef<AbortController | null>(null);
+
+  // Save retry parameters for onclose use
+  const retryParamsRef = useRef<{
+    url: string;
+    options: SSEOptions;
+    maxRetries: number;
+    retryInterval: number;
+  } | null>(null);
+
+  const updateStatus = useCallback((s: SSEStatus) => {
+    setStatus(s);
+  }, []);
+
+  /**
+   * Execute retry logic
+   * @returns true means you can continue to retry, false means the maximum number of retries has been reached
+   */
+  const doRetry = useCallback(
+    (err?: any): boolean => {
+      const { maxRetries = 10, options } = retryParamsRef.current || {};
+
+      if (!retryParamsRef.current) return false;
+
+      internalCountRef.current += 1;
+      setRetryCount(internalCountRef.current);
+
+      // Maximum number of retries reached
+      if (internalCountRef.current >= maxRetries) {
+        updateStatus('error');
+        options?.onError?.(
+          err || new Error('SSE connection failed after max retries'),
+        );
+        return false;
+      }
+      // Continue retrying
+      updateStatus('connecting');
+      return true;
+    },
+    [updateStatus],
+  );
+
+  const connectSSE = useCallback(
+    (url: string, options: SSEOptions) => {
+      const {
+        method = 'GET',
+        headers,
+        body,
+        maxRetries = 3,
+        retryInterval = 3000,
+        openWhenHidden = true,
+        onMessage,
+      } = options;
+
+      // Save retry parameters
+      retryParamsRef.current = { url, options, maxRetries, retryInterval };
+
+      // Clean up old connections
+      if (abortControllerRef.current) {
+        abortControllerRef.current.abort();
+      }
+
+      abortControllerRef.current = new AbortController();
+      // Reset count when first connecting (if not a retry scenario)
+      if (status === 'idle' || status === 'error' || status === 'closed') {
+        internalCountRef.current = 0;
+        setRetryCount(0);
+      }
+      updateStatus('connecting');
+
+      fetchEventSource(url, {
+        method,
+        headers: { ...headers },
+        body: body ? JSON.stringify(body) : undefined,
+        signal: abortControllerRef.current.signal,
+        openWhenHidden,
+        async onopen(response) {
+          if (
+            response.ok &&
+            response.headers
+              .get('content-type')
+              ?.includes(EventStreamContentType)
+          ) {
+            internalCountRef.current = 0;
+            setRetryCount(0);
+            updateStatus('open');
+            return;
+          }
+
+          // 401 Handling: Clear token and redirect to login page
+          if (response.status === 401) {
+            await authFaiedToLoginPage();
+            return;
+          }
+
+          throw new Error(`Connection failed with status ${response.status}`);
+        },
+
+        onmessage(msg) {
+          const event = msg.event.toLowerCase();
+          // Filter notification events
+          if (msg.data && event === 'data') {
+            onMessage?.(msg.data);
+          }
+
+          /**
+           * Other clients log in and connect to SSE. The current logic only makes one connection per user.
+           * The previous connection will wait for the default timeout in the background before disconnecting and will not trigger a disconnect event after testing.
+           * So it is agreed to actively disconnect the connection with this event and notify the user, allowing them to decide whether to reconnect.
+           */
+          if (event === 'replaced') {
+            disconnectSSE('replaced');
+          }
+
+          if (event === 'error') {
+            console.log('backend notice see closed');
+            const shouldRetry = doRetry();
+            if (shouldRetry) {
+              console.log(
+                `Retrying connection in ${retryInterval}ms... (attempt ${internalCountRef.current})`,
+              );
+              // Use setTimeout to delay retries
+              setTimeout(() => {
+                // Check if actively disconnected
+                if (abortControllerRef.current) {
+                  connectSSE(url, options);
+                }
+              }, retryInterval);
+            } else {
+              updateStatus('error');
+            }
+          }
+        },
+
+        onclose() {
+          console.log('SSE connection closed');
+          // Try to reconnect
+          const shouldRetry = doRetry();
+
+          if (shouldRetry) {
+            console.log(
+              `Retrying connection in ${retryInterval}ms... (attempt ${internalCountRef.current})`,
+            );
+            // Use setTimeout to delay retries
+            setTimeout(() => {
+              // Check if actively disconnected
+              if (abortControllerRef.current) {
+                connectSSE(url, options);
+              }
+            }, retryInterval);
+          } else {
+            updateStatus('error');
+          }
+        },
+
+        onerror(err) {
+          console.error('SSE connection error:');
+
+          // User actively disconnects, do not count
+          if (err.name === 'AbortError') {
+            updateStatus('closed');
+            throw err;
+          }
+
+          const shouldRetry = doRetry(err);
+
+          if (!shouldRetry) {
+            throw err; // Maximum number of retries reached, throw exception to stop automatic reconnection
+          }
+
+          updateStatus('connecting'); // Mark as reconnecting
+          return retryInterval;
+        },
+      });
+    },
+    [updateStatus, doRetry],
+  );
+
+  const disconnectSSE = useCallback(
+    (status: SSEStatus = 'error') => {
+      console.log('disconnect ...');
+
+      // Clear the retry parameters to prevent continued retries in onclose
+      retryParamsRef.current = null;
+
+      if (abortControllerRef.current) {
+        abortControllerRef.current.abort();
+        abortControllerRef.current = null;
+        updateStatus(status);
+      }
+    },
+    [updateStatus],
+  );
+
+  return {
+    status,
+    retryCount,
+    connectSSE,
+    disconnectSSE,
+  };
+};