diff --git a/config/hooks.ts b/config/hooks.ts index 26ffff428a..6a1f15f541 100644 --- a/config/hooks.ts +++ b/config/hooks.ts @@ -32,6 +32,7 @@ export const menus = [ 'useTextSelection', 'useWebSocket', 'useTheme', + 'useSse', ], }, { diff --git a/packages/hooks/src/index.ts b/packages/hooks/src/index.ts index 55c7232b0d..98e6a0e9e4 100644 --- a/packages/hooks/src/index.ts +++ b/packages/hooks/src/index.ts @@ -73,6 +73,7 @@ import useUpdateEffect from './useUpdateEffect'; import useUpdateLayoutEffect from './useUpdateLayoutEffect'; import useVirtualList from './useVirtualList'; import useWebSocket from './useWebSocket'; +import useSse from './useSse'; import useWhyDidYouUpdate from './useWhyDidYouUpdate'; import useMutationObserver from './useMutationObserver'; import useTheme from './useTheme'; @@ -133,6 +134,7 @@ export { useFavicon, useCountDown, useWebSocket, + useSse, useLockFn, useUnmountedRef, useExternal, diff --git a/packages/hooks/src/useSse/__tests__/index.spec.ts b/packages/hooks/src/useSse/__tests__/index.spec.ts new file mode 100644 index 0000000000..a6c013dd06 --- /dev/null +++ b/packages/hooks/src/useSse/__tests__/index.spec.ts @@ -0,0 +1,190 @@ +// useSse.test.ts +import { act, renderHook } from '@testing-library/react'; +import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; +import useSse, { ReadyState } from '../index'; + +class MockEventSource { + url: string; + withCredentials: boolean; + readyState: number; + onopen: ((this: EventSource, ev: Event) => any) | null = null; + onmessage: ((this: EventSource, ev: MessageEvent) => any) | null = null; + onerror: ((this: EventSource, ev: Event) => any) | null = null; + private listeners: Record void>> = {}; + private openTimeout?: NodeJS.Timeout; + + static CONNECTING = 0; + static OPEN = 1; + static CLOSED = 2; + + constructor(url: string, init?: EventSourceInit) { + this.url = url; + this.withCredentials = Boolean(init?.withCredentials); + this.readyState = MockEventSource.CONNECTING; + + this.openTimeout = setTimeout(() => { + this.readyState = MockEventSource.OPEN; + this.onopen?.(new Event('open')); + }, 10); + } + + addEventListener(type: string, listener: (ev: Event) => void) { + if (!this.listeners[type]) this.listeners[type] = []; + this.listeners[type].push(listener); + } + + dispatchEvent(type: string, event: Event) { + this.listeners[type]?.forEach((l) => l(event)); + } + + emitMessage(data: any) { + if (this.readyState !== MockEventSource.OPEN) return; + this.onmessage?.(new MessageEvent('message', { data })); + } + + emitError() { + this.onerror?.(new Event('error')); + } + + emitRetry(ms: number) { + const ev = new MessageEvent('message', { data: '' }); + (ev as any).retry = ms; + this.onmessage?.(ev); + } + + close() { + this.readyState = MockEventSource.CLOSED; + if (this.openTimeout) clearTimeout(this.openTimeout); + } +} + +describe('useSse Hook', () => { + const OriginalEventSource = (globalThis as any).EventSource; + + beforeEach(() => { + vi.useFakeTimers(); + (globalThis as any).EventSource = MockEventSource; + }); + + afterEach(() => { + vi.runAllTimers(); + vi.useRealTimers(); + (globalThis as any).EventSource = OriginalEventSource; + vi.restoreAllMocks(); + }); + + test('should connect and receive message', () => { + const hook = renderHook(() => useSse('/sse')); + expect(hook.result.current.readyState).toBe(ReadyState.Connecting); + + act(() => vi.advanceTimersByTime(20)); + expect(hook.result.current.readyState).toBe(ReadyState.Open); + + act(() => { + const es = hook.result.current.eventSource as unknown as MockEventSource; + es.emitMessage('hello'); + }); + expect(hook.result.current.latestMessage?.data).toBe('hello'); + + act(() => hook.result.current.disconnect()); + expect(hook.result.current.readyState).toBe(ReadyState.Closed); + }); + + test('manual mode should not auto connect', () => { + const hook = renderHook(() => useSse('/sse', { manual: true })); + expect(hook.result.current.readyState).toBe(ReadyState.Closed); + + act(() => { + hook.result.current.connect(); + vi.advanceTimersByTime(20); + }); + expect(hook.result.current.readyState).toBe(ReadyState.Open); + + act(() => hook.result.current.disconnect()); + }); + + test('should handle custom events', () => { + const onEvent = vi.fn(); + const hook = renderHook(() => useSse('/sse', { onEvent })); + act(() => vi.advanceTimersByTime(20)); + + act(() => { + const es = hook.result.current.eventSource as unknown as MockEventSource; + es.dispatchEvent('custom', new MessageEvent('custom', { data: 'foo' })); + }); + + expect(onEvent).toHaveBeenCalledWith( + 'custom', + expect.objectContaining({ data: 'foo' }), + expect.any(MockEventSource), + ); + + act(() => hook.result.current.disconnect()); + }); + + test('should reconnect on error respecting reconnectLimit', () => { + const hook = renderHook(() => useSse('/sse', { reconnectLimit: 1, reconnectInterval: 5 })); + act(() => vi.advanceTimersByTime(20)); + expect(hook.result.current.readyState).toBe(ReadyState.Open); + + act(() => { + const es = hook.result.current.eventSource as unknown as MockEventSource; + es.emitError(); + vi.advanceTimersByTime(20); + }); + + expect( + [ReadyState.Reconnecting, ReadyState.Open].includes(hook.result.current.readyState), + ).toBe(true); + + act(() => hook.result.current.disconnect()); + }); + + test('should respect server retry when enabled', () => { + const hook = renderHook(() => + useSse('/sse', { reconnectLimit: 1, reconnectInterval: 5, respectServerRetry: true }), + ); + act(() => vi.advanceTimersByTime(20)); + expect(hook.result.current.readyState).toBe(ReadyState.Open); + + act(() => { + const es = hook.result.current.eventSource as unknown as MockEventSource; + es.emitRetry(50); + es.emitError(); + vi.advanceTimersByTime(60); + }); + + expect( + [ReadyState.Reconnecting, ReadyState.Open].includes(hook.result.current.readyState), + ).toBe(true); + + act(() => hook.result.current.disconnect()); + }); + + test('should trigger all callbacks', () => { + const onOpen = vi.fn(); + const onMessage = vi.fn(); + const onError = vi.fn(); + const onReconnect = vi.fn(); + + const hook = renderHook(() => useSse('/sse', { onOpen, onMessage, onError, onReconnect })); + act(() => vi.advanceTimersByTime(20)); + expect(onOpen).toHaveBeenCalled(); + + act(() => { + const es = hook.result.current.eventSource as unknown as MockEventSource; + es.emitMessage('world'); + }); + expect(onMessage).toHaveBeenCalled(); + + act(() => { + const es = hook.result.current.eventSource as unknown as MockEventSource; + es.emitError(); + vi.advanceTimersByTime(20); + }); + expect(onError).toHaveBeenCalled(); + expect(onReconnect).toHaveBeenCalled(); + + act(() => hook.result.current.disconnect()); + }); +}); diff --git a/packages/hooks/src/useSse/demo/demo1.tsx b/packages/hooks/src/useSse/demo/demo1.tsx new file mode 100644 index 0000000000..53919600f5 --- /dev/null +++ b/packages/hooks/src/useSse/demo/demo1.tsx @@ -0,0 +1,39 @@ +import React, { useMemo, useRef } from 'react'; +import { useSse } from 'ahooks'; + +enum ReadyState { + Connecting = 0, + Open = 1, + Closed = 2, +} + +export default () => { + const historyRef = useRef([]); + const { readyState, latestMessage, connect, disconnect } = useSse('/api/sse'); + + historyRef.current = useMemo(() => historyRef.current.concat(latestMessage), [latestMessage]); + + return ( +
+ + +
readyState: {readyState}
+
+

received message:

+ {historyRef.current.map((m, i) => ( +

+ {m?.data} +

+ ))} +
+
+ ); +}; diff --git a/packages/hooks/src/useSse/index.en-US.md b/packages/hooks/src/useSse/index.en-US.md new file mode 100644 index 0000000000..433b9e755b --- /dev/null +++ b/packages/hooks/src/useSse/index.en-US.md @@ -0,0 +1,48 @@ +--- +nav: + path: /hooks +--- + +# useSse + +A hook for Server-Sent Events (SSE), which supports automatic reconnect and message callbacks. + +## Examples + +### Basic Usage + + + +## API + +```typescript +const { readyState, latestMessage, connect, disconnect, eventSource } = useSse( + url: string, + options?: UseSseOptions +) +``` + +### Options + +| Property | Description | Type | Default | +| -------------------- | ------------------------------------------------------- | ---------------------------------------------- | ------------ | +| manual | Whether to connect manually | `boolean` | `false` | +| withCredentials | Whether to send cross-domain requests with credentials | `boolean` | `false` | +| reconnectLimit | Maximum number of reconnection attempts | `number` | `3` | +| reconnectInterval | Reconnection interval (in milliseconds) | `number` | `3000` | +| respectServerRetry | Whether to respect the retry time sent by the server | `boolean` | `false` | +| onOpen | Callback when the connection is successfully established| `(es: EventSource) => void` | - | +| onMessage | Callback when a message is received | `(ev: MessageEvent, es: EventSource) => void` | - | +| onError | Callback when an error occurs | `(ev: Event, es: EventSource) => void` | - | +| onReconnect | Callback when a reconnection occurs | `(attempt: number, es: EventSource) => void` | - | +| onEvent | Callback for custom events | `(event: string, ev: MessageEvent, es: EventSource) => void` | - | + +### Result + +| Property | Description | Type | +| ------------- | ------------------------------------------ | ------------------------------------- | +| readyState | The current connection state | `ReadyState` (0: Connecting, 1: Open, 2: Closed, 3: Reconnecting) | +| latestMessage | The latest message received | `MessageEvent` \| `null` | +| connect | Function to manually connect | `() => void` | +| disconnect | Function to manually disconnect | `() => void` | +| eventSource | The native EventSource instance | `EventSource` \| `null` | \ No newline at end of file diff --git a/packages/hooks/src/useSse/index.ts b/packages/hooks/src/useSse/index.ts new file mode 100644 index 0000000000..ab35bc5db5 --- /dev/null +++ b/packages/hooks/src/useSse/index.ts @@ -0,0 +1,145 @@ +import { useCallback, useEffect, useRef, useState } from 'react'; + +export enum ReadyState { + Connecting = 0, + Open = 1, + Closed = 2, + Reconnecting = 3, +} + +export interface UseSseOptions { + manual?: boolean; // 是否手动连接（默认自动） + withCredentials?: boolean; // 是否携带跨域凭证 + reconnectLimit?: number; // 最大重连次数 + reconnectInterval?: number; // 默认重连间隔（毫秒） + respectServerRetry?: boolean; // 是否遵循服务端下发的 retry 时间 + onOpen?: (es: EventSource) => void; // 连接成功回调 + onMessage?: (ev: MessageEvent, es: EventSource) => void; // 收到消息回调 + onError?: (ev: Event, es: EventSource) => void; // 出错回调 + onReconnect?: (attempt: number, es: EventSource | null) => void; // 发生重连时回调 + onEvent?: (event: string, ev: MessageEvent, es: EventSource) => void; // 自定义事件回调 +} + +/** + * useSse - 一个支持自动重连 & 回调的 SSE Hook + */ +export default function useSse(url: string, options: UseSseOptions = {}) { + const { + manual, + withCredentials, + reconnectLimit = 3, + reconnectInterval = 3000, + respectServerRetry = false, + onOpen, + onMessage, + onError, + onReconnect, + onEvent, + } = options; + + const [readyState, setReadyState] = useState( + manual ? ReadyState.Closed : ReadyState.Connecting, + ); + + const [latestMessage, setLatestMessage] = useState(null); + + const eventSourceRef = useRef(null); + + const reconnectAttempts = useRef(0); + + const reconnectTimer = useRef | null>(null); + + const cleanup = useCallback(() => { + if (eventSourceRef.current) { + eventSourceRef.current.close(); + eventSourceRef.current = null; + } + if (reconnectTimer.current) { + clearTimeout(reconnectTimer.current); + reconnectTimer.current = null; + } + }, []); + + const connect = useCallback(() => { + cleanup(); + setReadyState(ReadyState.Connecting); + + const es = new EventSource(url, { withCredentials }); + eventSourceRef.current = es; + + es.onopen = () => { + reconnectAttempts.current = 0; + setReadyState(ReadyState.Open); + onOpen?.(es); + }; + + es.onmessage = (ev) => { + setLatestMessage(ev); + onMessage?.(ev, es); + }; + + es.onerror = (ev) => { + setReadyState(ReadyState.Closed); + onError?.(ev, es); + + if (reconnectAttempts.current < reconnectLimit) { + reconnectAttempts.current += 1; + + const delay = + respectServerRetry && (ev as any)?.retry ? (ev as any).retry : reconnectInterval; + + setReadyState(ReadyState.Reconnecting); + + onReconnect?.(reconnectAttempts.current, es); + + reconnectTimer.current = setTimeout(() => { + connect(); + }, delay); + } else { + cleanup(); + } + }; + + if (onEvent) { + es.addEventListener('custom', (ev) => { + onEvent('custom', ev as MessageEvent, es); + }); + } + }, [ + url, + withCredentials, + reconnectLimit, + reconnectInterval, + respectServerRetry, + onOpen, + onMessage, + onError, + onReconnect, + onEvent, + cleanup, + ]); + + /** + * 手动断开连接 + */ + const disconnect = useCallback(() => { + cleanup(); + setReadyState(ReadyState.Closed); + }, [cleanup]); + + /** + * 初始化：非 manual 模式下自动连接 + */ + useEffect(() => { + if (!manual) connect(); + return cleanup; + }, [manual, connect, cleanup]); + + return { + readyState, + latestMessage, + eventSource: eventSourceRef.current, + connect, + disconnect, + }; +} diff --git a/packages/hooks/src/useSse/index.zh-CN.md b/packages/hooks/src/useSse/index.zh-CN.md new file mode 100644 index 0000000000..ac187d2230 --- /dev/null +++ b/packages/hooks/src/useSse/index.zh-CN.md @@ -0,0 +1,48 @@ +--- +nav: + path: /hooks +--- + +# useSse + +一个用于 Server-Sent Events (SSE) 的 Hook，支持自动重连和消息回调。 + +## 代码演示 + +### 基础用法 + + + +## API + +```typescript +const { readyState, latestMessage, connect, disconnect, eventSource } = useSse( + url: string, + options?: UseSseOptions +) +``` + +### Options + +| 参数 | 说明 | 类型 | 默认值 | +| -------------------- | -------------------------------------- | ---------------------------------------------- | ------------ | +| manual | 是否手动连接 | `boolean` | `false` | +| withCredentials | 是否携带跨域凭证 | `boolean` | `false` | +| reconnectLimit | 最大重连次数 | `number` | `3` | +| reconnectInterval | 重连间隔（毫秒） | `number` | `3000` | +| respectServerRetry | 是否遵循服务端下发的 retry 时间 | `boolean` | `false` | +| onOpen | 连接成功回调 | `(es: EventSource) => void` | - | +| onMessage | 收到消息回调 | `(ev: MessageEvent, es: EventSource) => void` | - | +| onError | 出错回调 | `(ev: Event, es: EventSource) => void` | - | +| onReconnect | 发生重连时回调 | `(attempt: number, es: EventSource) => void` | - | +| onEvent | 自定义事件回调 | `(event: string, ev: MessageEvent, es: EventSource) => void` | - | + +### Result + +| 参数 | 说明 | 类型 | +| ------------- | ---------------------- | ------------------------------------- | +| readyState | 当前连接状态 | `ReadyState` (0: 连接中, 1: 已连接, 2: 已关闭, 3: 重连中) | +| latestMessage | 收到的最新消息 | `MessageEvent` \| `null` | +| connect | 手动连接函数 | `() => void` | +| disconnect | 手动断开函数 | `() => void` | +| eventSource | 原生 EventSource 实例 | `EventSource` \| `null` |