@assistant-ui/react 0.15.8 → 0.15.9

package/src/unstable/useLiveCompletionAdapter.test.tsx

@@ -129,6 +129,96 @@ describe("unstable_useLiveCompletionAdapter", () => {
     expect(result.current.adapter.search!("ab")).toEqual([item("ab")]);
   });
 
+  it("allows a failed query to be retried", async () => {
+    const fetcher = vi
+      .fn<(query: string) => Promise<readonly Unstable_TriggerItem[]>>()
+      .mockRejectedValueOnce(new Error("temporarily unavailable"))
+      .mockRejectedValueOnce(new Error("still unavailable"))
+      .mockResolvedValueOnce([item("alice")]);
+    const { result } = renderHook(() =>
+      unstable_useLiveCompletionAdapter({ fetcher, debounceMs: 0 }),
+    );
+
+    for (const attempt of [1, 2]) {
+      await act(async () => {
+        result.current.adapter.search!("alice");
+        await vi.advanceTimersByTimeAsync(0);
+      });
+      expect(fetcher).toHaveBeenCalledTimes(attempt);
+      expect(result.current.isLoading).toBe(false);
+    }
+
+    await act(async () => {
+      result.current.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(fetcher).toHaveBeenCalledTimes(3);
+    expect(result.current.adapter.search!("alice")).toEqual([item("alice")]);
+  });
+
+  it("does not automatically retry when search runs during every render", async () => {
+    const fetcher = vi
+      .fn<(query: string) => Promise<readonly Unstable_TriggerItem[]>>()
+      .mockRejectedValueOnce(new Error("temporarily unavailable"))
+      .mockImplementation(
+        () => new Promise<readonly Unstable_TriggerItem[]>(() => {}),
+      );
+    const { result } = renderHook(() => {
+      const completion = unstable_useLiveCompletionAdapter({
+        fetcher,
+        debounceMs: 0,
+      });
+      completion.adapter.search!("alice");
+      return completion;
+    });
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(result.current.isLoading).toBe(false);
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(1_000);
+    });
+    expect(fetcher).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-arms a failed query when its pending retry is interrupted", async () => {
+    const fetcher = vi
+      .fn<(query: string) => Promise<readonly Unstable_TriggerItem[]>>()
+      .mockRejectedValueOnce(new Error("temporarily unavailable"))
+      .mockImplementationOnce(
+        () => new Promise<readonly Unstable_TriggerItem[]>(() => {}),
+      )
+      .mockResolvedValueOnce([item("alice")]);
+    const { result } = renderHook(() =>
+      unstable_useLiveCompletionAdapter({ fetcher, debounceMs: 0 }),
+    );
+
+    await act(async () => {
+      result.current.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+
+    await act(async () => {
+      result.current.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(fetcher).toHaveBeenCalledTimes(2);
+
+    await act(async () => {
+      result.current.adapter.search!("alicex");
+    });
+    await act(async () => {
+      result.current.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+
+    expect(fetcher).toHaveBeenCalledTimes(3);
+    expect(fetcher).toHaveBeenLastCalledWith("alice");
+    expect(result.current.adapter.search!("alice")).toEqual([item("alice")]);
+  });
+
   it("drops an in-flight fetch when the query returns to a cached value", async () => {
     const resolvers: Record<
       string,

package/src/unstable/useLiveCompletionAdapter.ts

@@ -57,7 +57,8 @@ export function unstable_useLiveCompletionAdapter(
   const [state, setState] = useState<{
     query: string;
     items: readonly Unstable_TriggerItem[];
-  }>({ query: NO_QUERY, items: [] });
+    failed: boolean;
+  }>({ query: NO_QUERY, items: [], failed: false });
   const [isLoading, setIsLoading] = useState(false);
 
   const fetcherRef = useRef(fetcher);
@@ -66,6 +67,8 @@ export function unstable_useLiveCompletionAdapter(
   const timerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const tokenRef = useRef(0);
   const pendingQueryRef = useRef<string | null>(null);
+  const retryableQueryRef = useRef<string | null>(null);
+  const pendingRetryQueryRef = useRef<string | null>(null);
 
   const cancelTimer = useCallback(() => {
     if (timerRef.current !== null) {
@@ -74,10 +77,22 @@ export function unstable_useLiveCompletionAdapter(
     }
   }, []);
 
+  const rearmPendingRetry = useCallback(() => {
+    const query = pendingRetryQueryRef.current;
+    if (query === null) return;
+    retryableQueryRef.current = query;
+    pendingRetryQueryRef.current = null;
+  }, []);
+
   const scheduleFetch = useCallback(
     (query: string) => {
       if (!enabled) return;
       if (pendingQueryRef.current === query) return;
+      rearmPendingRetry();
+      if (retryableQueryRef.current === query) {
+        retryableQueryRef.current = null;
+        pendingRetryQueryRef.current = query;
+      }
       pendingQueryRef.current = query;
       cancelTimer();
       const token = ++tokenRef.current;
@@ -87,37 +102,47 @@ export function unstable_useLiveCompletionAdapter(
         fetcherRef.current(query).then(
           (items) => {
             if (token !== tokenRef.current) return;
-            setState({ query, items });
+            pendingRetryQueryRef.current = null;
+            setState({ query, items, failed: false });
             setIsLoading(false);
           },
           () => {
             if (token !== tokenRef.current) return;
-            setState({ query, items: [] });
+            pendingQueryRef.current = null;
+            pendingRetryQueryRef.current = null;
+            setState({ query, items: [], failed: true });
             setIsLoading(false);
           },
         );
       }, debounceMs);
     },
-    [enabled, debounceMs, cancelTimer],
+    [enabled, debounceMs, cancelTimer, rearmPendingRetry],
   );
 
   const invalidatePending = useCallback(() => {
+    rearmPendingRetry();
     cancelTimer();
     pendingQueryRef.current = null;
     tokenRef.current += 1;
     setIsLoading(false);
-  }, [cancelTimer]);
+  }, [cancelTimer, rearmPendingRetry]);
 
   useEffect(() => {
     if (enabled) return;
     invalidatePending();
     setState((s) =>
-      s.query === NO_QUERY ? s : { query: NO_QUERY, items: [] },
+      s.query === NO_QUERY ? s : { query: NO_QUERY, items: [], failed: false },
     );
   }, [enabled, invalidatePending]);
 
   useEffect(() => cancelTimer, [cancelTimer]);
 
+  // Arm retries only after the failed state commits. Arming during rejection
+  // would let the failure render immediately schedule another request.
+  useEffect(() => {
+    retryableQueryRef.current = state.failed ? state.query : null;
+  }, [state]);
+
   const adapter = useMemo<Unstable_TriggerAdapter>(
     () => ({
       categories: () => [],
@@ -125,7 +150,7 @@ export function unstable_useLiveCompletionAdapter(
       search: (query: string) => {
         // search() runs inside the popover's render; defer state updates with
         // queueMicrotask so they are not dispatched while another component renders.
-        if (query !== state.query) {
+        if (query !== state.query || retryableQueryRef.current === query) {
           queueMicrotask(() => scheduleFetch(query));
         } else if (
           pendingQueryRef.current !== null &&