fix: support async subscribe in useSWRSubscription

src/subscription/index.ts

@@ -72,13 +72,30 @@ export const subscription = (<Data = any, Error = any>(useSWRNext: SWRHook) =>
       subscriptions.set(subscriptionKey, refCount + 1)
 
       if (!refCount) {
-        const dispose = subscribe(args, { next })
-        if (typeof dispose !== 'function') {
+        const result = subscribe(args, { next })
+
+        if (result && typeof (result as any).then === 'function') {
+          // Race condition guard: if cleanup runs before the async subscribe
+          // resolves, the flag tells the resolver to dispose immediately.
+          let shouldDisposeOnResolve = false
+          ;(result as Promise<(() => void) | void>).then(dispose => {
+            if (shouldDisposeOnResolve) {
+              if (typeof dispose === 'function') dispose()
+            } else if (typeof dispose === 'function') {
+              disposers.set(subscriptionKey!, dispose)
+            }
+          })
+
+          disposers.set(subscriptionKey, () => {
+            shouldDisposeOnResolve = true
+          })
+        } else if (typeof result === 'function') {
+          disposers.set(subscriptionKey, result)
+        } else if (typeof result !== 'undefined') {
           throw new Error(
             'The `subscribe` function must return a function to unsubscribe.'
           )
         }
-        disposers.set(subscriptionKey, dispose)
       }
 
       return () => {

src/subscription/types.ts

@@ -4,16 +4,22 @@ export type SWRSubscriptionOptions<Data = any, Error = any> = {
   next: (err?: Error | null, data?: Data | MutatorCallback<Data>) => void
 }
 
+type SWRSubscribeReturn = (() => void) | void
+type SWRSubscribeFn<Arg, Data, Error> = (
+  key: Arg,
+  { next }: SWRSubscriptionOptions<Data, Error>
+) => SWRSubscribeReturn | Promise<SWRSubscribeReturn>
+
 export type SWRSubscription<
   SWRSubKey extends Key = Key,
   Data = any,
   Error = any
 > = SWRSubKey extends () => infer Arg | null | undefined | false
-  ? (key: Arg, { next }: SWRSubscriptionOptions<Data, Error>) => void
+  ? SWRSubscribeFn<Arg, Data, Error>
   : SWRSubKey extends null | undefined | false
   ? never
   : SWRSubKey extends infer Arg
-  ? (key: Arg, { next }: SWRSubscriptionOptions<Data, Error>) => void
+  ? SWRSubscribeFn<Arg, Data, Error>
   : never
 
 export type SWRSubscriptionResponse<Data = any, Error = any> = {

test/type/subscription.ts

@@ -95,4 +95,20 @@ export function useTestSubscription() {
   const { data: data2, error: error2 } = useSWRSubscription('key', sub)
   expectType<string | undefined>(data2)
   expectType<Error | undefined>(error2)
+
+  // Async subscribe should be accepted.
+  useSWRSubscription(
+    'key',
+    async (_key, { next: _ }: SWRSubscriptionOptions<string, Error>) => {
+      return () => {}
+    }
+  )
+
+  const asyncSub: SWRSubscription<string, string, Error> = async (
+    _,
+    { next: __ }
+  ) => {
+    return () => {}
+  }
+  useSWRSubscription('key', asyncSub)
 }

test/use-swr-subscription.test.tsx

@@ -293,6 +293,79 @@ describe('useSWRSubscription', () => {
     await screen.findByText(`data: 3`)
   })
 
+  it('should support async subscribe', async () => {
+    const swrKey = createKey()
+    let emitter: ((data: string) => void) | null = null
+    let disposed = false
+
+    async function subscribe(_key, { next }) {
+      await sleep(50)
+      emitter = (data: string) => next(undefined, data)
+      return () => {
+        disposed = true
+        emitter = null
+      }
+    }
+
+    function Page() {
+      const { data } = useSWRSubscription(swrKey, subscribe, {
+        fallbackData: 'fallback'
+      })
+      return <div>{'data:' + data}</div>
+    }
+
+    renderWithConfig(<Page />)
+    screen.getByText('data:fallback')
+
+    // Wait for async subscribe to resolve.
+    await act(() => sleep(100))
+    act(() => emitter?.('hello'))
+    await act(() => sleep(10))
+    screen.getByText('data:hello')
+
+    expect(disposed).toBe(false)
+  })
+
+  it('should clean up async subscribe on unmount', async () => {
+    const swrKey = createKey()
+    let disposed = false
+
+    async function subscribe(_key, { next }) {
+      await sleep(100)
+      next(undefined, 'connected')
+      return () => {
+        disposed = true
+      }
+    }
+
+    function Page() {
+      const [show, setShow] = useState(true)
+      return (
+        <>
+          {show ? <Child /> : null}
+          <button onClick={() => setShow(false)}>unmount</button>
+        </>
+      )
+    }
+    function Child() {
+      const { data } = useSWRSubscription(swrKey, subscribe, {
+        fallbackData: 'fallback'
+      })
+      return <div>{'data:' + data}</div>
+    }
+
+    renderWithConfig(<Page />)
+    screen.getByText('data:fallback')
+
+    // Unmount before async subscribe resolves.
+    await act(() => sleep(10))
+    fireEvent.click(screen.getByText('unmount'))
+
+    // After the Promise resolves, the dispose should still be called.
+    await act(() => sleep(200))
+    expect(disposed).toBe(true)
+  })
+
   it('should require a dispose function', async () => {
     jest.spyOn(console, 'error').mockImplementation(() => {})
 
@@ -303,6 +376,7 @@ describe('useSWRSubscription', () => {
     }
 
     function Page() {
+      // @ts-expect-error -- intentionally passing an invalid subscribe function
       useSWRSubscription(swrKey, subscribe)
       return null
     }