feat: Support promise in watch util (#825)

* feat: Support promise in watch util
Implements #507

* chore: reverted return type of callback in subscribe, removed mem optimization

* chore: useFakeTimers

* fix: return type + sleep

* chore: code order

---------

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>

src/vanilla/utils/watch.ts

@@ -2,7 +2,9 @@ import { subscribe } from '../../vanilla.ts'
 
 type Cleanup = () => void
 type WatchGet = <T extends object>(proxyObject: T) => T
-type WatchCallback = (get: WatchGet) => Cleanup | void | undefined
+type WatchCallback = (
+  get: WatchGet,
+) => Cleanup | void | Promise<Cleanup | void> | undefined
 type WatchOptions = {
   sync?: boolean
 }
@@ -53,10 +55,12 @@ export function watch(
     }
   }
 
-  const revalidate = () => {
+  const revalidate = async () => {
     if (!alive) {
       return
     }
+
+    // run own cleanups before re-subscribing
     cleanups.forEach((clean) => clean())
     cleanups.clear()
 
@@ -69,35 +73,40 @@ export function watch(
 
     // Ensures that the parent is reset if the callback throws an error.
     try {
-      const cleanupReturn = callback((proxyObject) => {
+      const promiseOrPossibleCleanup = callback((proxyObject) => {
         proxiesToSubscribe.add(proxyObject)
+        // in case the callback is a promise and the watch has ended
+        if (alive && !subscriptions.has(proxyObject)) {
+          // subscribe to new proxy immediately -> this fixes problems when Promises are used due to the callstack
+          const unsubscribe = subscribe(proxyObject, revalidate, options?.sync)
+          subscriptions.set(proxyObject, unsubscribe)
+        }
         return proxyObject
       })
+      const couldBeCleanup =
+        promiseOrPossibleCleanup && promiseOrPossibleCleanup instanceof Promise
+          ? await promiseOrPossibleCleanup
+          : promiseOrPossibleCleanup
 
       // If there's a cleanup, we add this to the cleanups set
-      if (cleanupReturn) {
-        cleanups.add(cleanupReturn)
+      if (couldBeCleanup) {
+        if (alive) {
+          cleanups.add(couldBeCleanup)
+        } else {
+          cleanup()
+        }
       }
     } finally {
       currentCleanups = parent
     }
 
     // Unsubscribe old subscriptions
     subscriptions.forEach((unsubscribe, proxyObject) => {
-      if (proxiesToSubscribe.has(proxyObject)) {
-        // Already subscribed
-        proxiesToSubscribe.delete(proxyObject)
-      } else {
+      if (!proxiesToSubscribe.has(proxyObject)) {
         subscriptions.delete(proxyObject)
         unsubscribe()
       }
     })
-
-    // Subscribe to new proxies
-    proxiesToSubscribe.forEach((proxyObject) => {
-      const unsubscribe = subscribe(proxyObject, revalidate, options?.sync)
-      subscriptions.set(proxyObject, unsubscribe)
-    })
   }
 
   // If there's a parent watch call, we attach this watch's

tests/watch.test.tsx

@@ -1,8 +1,21 @@
-import { describe, expect, it, vi } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { proxy } from 'valtio'
 import { watch } from 'valtio/utils'
 
+const sleep = (ms: number) =>
+  new Promise((resolve) => {
+    setTimeout(resolve, ms)
+  })
+
 describe('watch', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
   it('should re-run for individual proxy updates', async () => {
     const reference = proxy({ value: 'Example' })
 
@@ -96,4 +109,56 @@ describe('watch', () => {
 
     reference.value = 'Update'
   })
+  it('should support promise watchers', async () => {
+    const reference = proxy({ value: 'Example' })
+
+    const callback = vi.fn()
+
+    const waitPromise = sleep(10000)
+    watch(async (get) => {
+      await waitPromise
+      get(reference)
+      callback()
+    })
+
+    vi.runAllTimers()
+    await waitPromise
+
+    expect(callback).toBeCalledTimes(1)
+    // listener will only be attached after one promise callback due to the await stack
+    await Promise.resolve()
+    reference.value = 'Update'
+    // wait for internal promise
+    await Promise.resolve()
+    // wait for next promise resolve call due to promise usage inside of callback
+    await Promise.resolve()
+    expect(callback).toBeCalledTimes(2)
+  })
+
+  it('should not subscribe if the watch is stopped before the promise completes', async () => {
+    const reference = proxy({ value: 'Example' })
+
+    const callback = vi.fn()
+
+    const waitPromise = sleep(10000)
+    const stop = watch(async (get) => {
+      await waitPromise
+      get(reference)
+      callback()
+    })
+    stop()
+
+    vi.runAllTimers()
+    await waitPromise
+
+    expect(callback).toBeCalledTimes(1)
+    // listener will only be attached after one promise callback due to the await stack
+    await Promise.resolve()
+    reference.value = 'Update'
+    // wait for internal promise
+    await Promise.resolve()
+    // wait for next promise resolve call due to promise usage inside of callback
+    await Promise.resolve()
+    expect(callback).toBeCalledTimes(1)
+  })
 })