fix useDocHandle

packages/automerge-repo-react-hooks/src/useDocHandle.ts

@@ -1,7 +1,7 @@
-import type { AnyDocumentId, DocHandle } from "@automerge/automerge-repo/slim"
+import { AnyDocumentId, DocHandle } from "@automerge/automerge-repo/slim"
 import { wrapPromise } from "./wrapPromise.js"
 import { useRepo } from "./useRepo.js"
-import { useRef } from "react"
+import { useEffect, useRef, useState } from "react"
 
 // Shared with useDocHandles
 export const wrapperCache = new Map<
@@ -34,6 +34,7 @@ export function useDocHandle<T>(
 ): DocHandle<T> | undefined {
   const repo = useRepo()
   const controllerRef = useRef<AbortController>()
+  const [handle, setHandle] = useState<DocHandle<T> | undefined>()
 
   let wrapper = wrapperCache.get(id)
   if (!wrapper) {
@@ -45,16 +46,22 @@ export function useDocHandle<T>(
     wrapperCache.set(id, wrapper)
   }
 
-  if (suspense === false) {
-    try {
-      return wrapper.read() as DocHandle<T>
-    } catch (e) {
-      if (e instanceof Promise) {
-        return undefined
-      }
-      throw e
+  useEffect(() => {
+    if (suspense === false) {
+      void wrapper.promise
+        .then(handle => {
+          setHandle(handle as DocHandle<T>)
+        })
+        .catch(e => {
+          console.log("handle promise caught", e)
+          setHandle(undefined)
+        })
     }
-  }
+  }, [suspense, wrapper])
 
-  return wrapper.read() as DocHandle<T>
+  if (suspense) {
+    return wrapper.read() as DocHandle<T>
+  } else {
+    return handle || undefined
+  }
 }

packages/automerge-repo-react-hooks/src/useDocHandles.ts

@@ -19,37 +19,41 @@ export function useDocHandles<T>(
   const controllerRef = useRef<AbortController>()
 
   // First, handle suspense outside of effects
-  if (suspense) {
-    controllerRef.current?.abort()
-    controllerRef.current = new AbortController()
+  controllerRef.current?.abort()
+  controllerRef.current = new AbortController()
 
-    // Check if we need any new wrappers
-    const pendingPromises: Promise<unknown>[] = []
+  // Check if we need any new wrappers
+  const pendingPromises: Promise<unknown>[] = []
 
-    for (const id of ids) {
-      if (!wrapperCache.has(id)) {
-        const promise = repo.find<T>(id, {
-          signal: controllerRef.current.signal,
-        })
-        const wrapper = wrapPromise(promise)
-        wrapperCache.set(id, wrapper)
-      }
+  for (const id of ids) {
+    if (!wrapperCache.has(id)) {
+      const promise = repo.find<T>(id, {
+        signal: controllerRef.current.signal,
+      })
+      const wrapper = wrapPromise(promise)
+      wrapperCache.set(id, wrapper)
+    }
 
-      // Try to read each wrapper
-      const wrapper = wrapperCache.get(id)!
-      try {
-        wrapper.read()
-      } catch (e) {
-        if (e instanceof Promise) {
-          pendingPromises.push(e)
-        }
+    // Try to read each wrapper
+    const wrapper = wrapperCache.get(id)!
+    try {
+      wrapper.read()
+      const handle = wrapper.read() as DocHandle<T>
+      setHandleMap(prev => {
+        const next = new Map(prev)
+        next.set(id, handle)
+        return next
+      })
+    } catch (e) {
+      if (e instanceof Promise) {
+        pendingPromises.push(e)
       }
     }
+  }
 
-    // If any promises are pending, suspend with Promise.all
-    if (pendingPromises.length > 0) {
-      throw Promise.all(pendingPromises)
-    }
+  // If any promises are pending, suspend with Promise.all
+  if (suspense && pendingPromises.length > 0) {
+    throw Promise.all(pendingPromises)
   }
 
   useEffect(() => {
@@ -59,22 +63,11 @@ export function useDocHandles<T>(
     }
 
     // Now safely get all available handles
-    ids.forEach(id => {
-      if (!suspense && !wrapperCache.has(id)) {
-        const promise = repo.find<T>(id, {
-          signal: controllerRef.current!.signal,
-        })
-        wrapperCache.set(id, wrapPromise(promise))
-      }
-
+
       const wrapper = wrapperCache.get(id)
       try {
         const handle = wrapper?.read() as DocHandle<T>
-        setHandleMap(prev => {
-          const next = new Map(prev)
-          next.set(id, handle)
-          return next
-        })
+
       } catch (e) {
         if (!(e instanceof Promise)) {
           console.error(`Error loading document ${id}:`, e)

packages/automerge-repo-react-hooks/src/wrapPromise.ts

@@ -1,6 +1,7 @@
 type Status = "pending" | "success" | "error"
 
 type PromiseWrapper<T> = {
+  promise: Promise<T>
   read(): T
 }
 
@@ -21,6 +22,7 @@ export function wrapPromise<T>(promise: Promise<T>): PromiseWrapper<T> {
   )
 
   return {
+    promise,
     read(): T {
       switch (status) {
         case "pending":

packages/automerge-repo-react-hooks/test/useDocHandle.test.tsx

@@ -233,4 +233,109 @@ describe("useHandle", () => {
       expect(onHandle).not.toHaveBeenCalledWith(handleA)
     })
   })
+
+  describe("useHandle with suspense: false", () => {
+    it("returns undefined while loading then resolves to handle", async () => {
+      const { handleA, repo, wrapper } = await setup()
+      const onHandle = vi.fn()
+
+      // Mock find to simulate network delay
+      const originalFind = repo.find.bind(repo)
+      repo.find = vi.fn().mockImplementation(async (...args) => {
+        await new Promise(resolve => setTimeout(resolve, 100))
+        return originalFind(...args)
+      })
+
+      const NonSuspenseComponent = ({
+        url,
+        onHandle,
+      }: {
+        url: AutomergeUrl
+        onHandle: (handle: DocHandle<unknown> | undefined) => void
+      }) => {
+        const handle = useDocHandle(url, { suspense: false })
+        onHandle(handle)
+        return null
+      }
+
+      render(<NonSuspenseComponent url={handleA.url} onHandle={onHandle} />, {
+        wrapper,
+      })
+
+      // Initially should be called with undefined
+      expect(onHandle).toHaveBeenCalledWith(undefined)
+
+      // Wait for handle to load
+      await waitFor(() => {
+        expect(onHandle).toHaveBeenLastCalledWith(handleA)
+      })
+
+      // Restore original find implementation
+      repo.find = originalFind
+    })
+
+    it("handles unavailable documents by returning undefined", async () => {
+      const { repo, wrapper } = await setup()
+      const url = generateAutomergeUrl()
+      const onHandle = vi.fn()
+
+      const NonSuspenseComponent = ({
+        url,
+        onHandle,
+      }: {
+        url: AutomergeUrl
+        onHandle: (handle: DocHandle<unknown> | undefined) => void
+      }) => {
+        const handle = useDocHandle(url, { suspense: false })
+        onHandle(handle)
+        return null
+      }
+
+      render(<NonSuspenseComponent url={url} onHandle={onHandle} />, {
+        wrapper,
+      })
+
+      // Should start with undefined
+      expect(onHandle).toHaveBeenCalledWith(undefined)
+
+      // Should continue to return undefined after attempted load
+      await waitFor(() => {
+        expect(onHandle).toHaveBeenLastCalledWith(undefined)
+      })
+    })
+
+    it("updates the handle when url changes", async () => {
+      const { wrapper, handleA, handleB } = setup()
+      const onHandle = vi.fn()
+
+      const NonSuspenseComponent = ({
+        url,
+        onHandle,
+      }: {
+        url: AutomergeUrl
+        onHandle: (handle: DocHandle<unknown> | undefined) => void
+      }) => {
+        const handle = useDocHandle(url, { suspense: false })
+        onHandle(handle)
+        return null
+      }
+
+      const { rerender } = render(
+        <NonSuspenseComponent url={handleA.url} onHandle={onHandle} />,
+        { wrapper }
+      )
+
+      // Wait for first handle to load
+      await waitFor(() => expect(onHandle).toHaveBeenLastCalledWith(handleA))
+
+      // Change URL
+      rerender(<NonSuspenseComponent url={handleB.url} onHandle={onHandle} />)
+
+      // Should temporarily return to undefined
+      expect(onHandle).toHaveBeenCalledWith(undefined)
+
+      // Then resolve to new handle
+      await waitFor(() => expect(onHandle).toHaveBeenLastCalledWith(handleB))
+    })
+  })
 })