@assistant-ui/react 0.15.18 → 0.15.19

package/src/unstable/useLiveCompletionAdapter.test.tsx

@@ -1,5 +1,5 @@
 /** @vitest-environment jsdom */
-import { startTransition, Suspense } from "react";
+import { startTransition, Suspense, useLayoutEffect } from "react";
 import { act, render, renderHook } from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { Unstable_TriggerItem } from "@assistant-ui/core";
@@ -62,6 +62,159 @@ describe("unstable_useLiveCompletionAdapter", () => {
     expect(fetcher).not.toHaveBeenCalled();
   });
 
+  it("does not start a queued fetch after unmount", async () => {
+    const fetcher = vi.fn(async () => []);
+    const { result, unmount } = renderHook(() =>
+      unstable_useLiveCompletionAdapter({ fetcher, debounceMs: 0 }),
+    );
+
+    result.current.adapter.search!("alice");
+    unmount();
+    await Promise.resolve();
+
+    expect(vi.getTimerCount()).toBe(0);
+    await vi.runAllTimersAsync();
+
+    expect(fetcher).not.toHaveBeenCalled();
+  });
+
+  it("keeps an in-flight fetch when the debounce changes", async () => {
+    let resolve!: (items: readonly Unstable_TriggerItem[]) => void;
+    const fetcher = vi.fn(
+      () =>
+        new Promise<readonly Unstable_TriggerItem[]>((r) => {
+          resolve = r;
+        }),
+    );
+    const { result, rerender } = renderHook(
+      ({ debounceMs }) =>
+        unstable_useLiveCompletionAdapter({ fetcher, debounceMs }),
+      { initialProps: { debounceMs: 0 } },
+    );
+
+    await act(async () => {
+      result.current.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(fetcher).toHaveBeenCalledOnce();
+
+    rerender({ debounceMs: 50 });
+    await act(async () => resolve([item("alice")]));
+
+    expect(result.current.adapter.search!("alice")).toEqual([item("alice")]);
+    expect(result.current.isLoading).toBe(false);
+  });
+
+  it("replays a search queued while the adapter is hidden", async () => {
+    const fetcher = vi.fn(async () => []);
+    const suspended = new Promise<never>(() => {});
+    let committed!: ReturnType<typeof unstable_useLiveCompletionAdapter>;
+    const Harness = ({ blocked }: { blocked: boolean }) => {
+      const current = unstable_useLiveCompletionAdapter({
+        fetcher,
+        debounceMs: 0,
+      });
+      useLayoutEffect(() => {
+        committed = current;
+      }, [current]);
+      if (blocked) throw suspended;
+      return null;
+    };
+    const view = (blocked: boolean) => (
+      <Suspense fallback={null}>
+        <Harness blocked={blocked} />
+      </Suspense>
+    );
+    const rendered = render(view(false));
+
+    await act(async () => rendered.rerender(view(true)));
+    committed.adapter.search!("alice");
+    await Promise.resolve();
+
+    expect(vi.getTimerCount()).toBe(0);
+    await act(async () => rendered.rerender(view(false)));
+    await act(async () => vi.advanceTimersByTimeAsync(0));
+
+    expect(fetcher).toHaveBeenCalledOnce();
+    expect(fetcher).toHaveBeenCalledWith("alice");
+  });
+
+  it("drops a hidden search superseded by a cached query", async () => {
+    const fetcher = vi.fn(async (query: string) => [item(query)]);
+    const suspended = new Promise<never>(() => {});
+    let committed!: ReturnType<typeof unstable_useLiveCompletionAdapter>;
+    const Harness = ({ blocked }: { blocked: boolean }) => {
+      const current = unstable_useLiveCompletionAdapter({
+        fetcher,
+        debounceMs: 0,
+      });
+      useLayoutEffect(() => {
+        committed = current;
+      }, [current]);
+      if (blocked) throw suspended;
+      return null;
+    };
+    const view = (blocked: boolean) => (
+      <Suspense fallback={null}>
+        <Harness blocked={blocked} />
+      </Suspense>
+    );
+    const rendered = render(view(false));
+
+    await act(async () => {
+      committed.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(committed.adapter.search!("alice")).toEqual([item("alice")]);
+
+    await act(async () => rendered.rerender(view(true)));
+    committed.adapter.search!("bob");
+    await Promise.resolve();
+    committed.adapter.search!("alice");
+    await Promise.resolve();
+
+    await act(async () => rendered.rerender(view(false)));
+    await act(async () => vi.advanceTimersByTimeAsync(0));
+
+    expect(fetcher).toHaveBeenCalledOnce();
+    expect(fetcher).toHaveBeenCalledWith("alice");
+  });
+
+  it("settles loading when a suspended adapter is shown again", async () => {
+    const fetcher = vi.fn(() => new Promise<never>(() => {}));
+    const suspended = new Promise<never>(() => {});
+    let committed!: ReturnType<typeof unstable_useLiveCompletionAdapter>;
+    const Harness = ({ blocked }: { blocked: boolean }) => {
+      const current = unstable_useLiveCompletionAdapter({
+        fetcher,
+        debounceMs: 0,
+      });
+      useLayoutEffect(() => {
+        committed = current;
+      }, [current]);
+      if (blocked) throw suspended;
+      return <output data-testid="status">{String(current.isLoading)}</output>;
+    };
+    const view = (blocked: boolean) => (
+      <Suspense fallback={null}>
+        <Harness blocked={blocked} />
+      </Suspense>
+    );
+    const rendered = render(view(false));
+
+    await act(async () => {
+      committed.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(fetcher).toHaveBeenCalledOnce();
+    expect(rendered.getByTestId("status").textContent).toBe("true");
+
+    await act(async () => rendered.rerender(view(true)));
+    await act(async () => rendered.rerender(view(false)));
+
+    expect(rendered.getByTestId("status").textContent).toBe("false");
+  });
+
   it("does not fetch when disabled and clears cached items", async () => {
     const fetcher = vi.fn(async () => [item("a")]);
     const { result, rerender } = renderHook(
@@ -435,6 +588,54 @@ describe("unstable_useLiveCompletionAdapter", () => {
     expect(result.current.adapter.search!("alice")).toEqual([item("alice")]);
   });
 
+  it("re-arms a failed query when its retry is hidden", async () => {
+    const fetcher = vi
+      .fn<(query: string) => Promise<readonly Unstable_TriggerItem[]>>()
+      .mockRejectedValueOnce(new Error("temporarily unavailable"))
+      .mockImplementationOnce(() => new Promise(() => {}))
+      .mockResolvedValueOnce([item("alice")]);
+    const suspended = new Promise<never>(() => {});
+    let committed!: ReturnType<typeof unstable_useLiveCompletionAdapter>;
+    const Harness = ({ blocked }: { blocked: boolean }) => {
+      const current = unstable_useLiveCompletionAdapter({
+        fetcher,
+        debounceMs: 0,
+      });
+      useLayoutEffect(() => {
+        committed = current;
+      }, [current]);
+      if (blocked) throw suspended;
+      return null;
+    };
+    const view = (blocked: boolean) => (
+      <Suspense fallback={null}>
+        <Harness blocked={blocked} />
+      </Suspense>
+    );
+    const rendered = render(view(false));
+
+    await act(async () => {
+      committed.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(fetcher).toHaveBeenCalledOnce();
+
+    await act(async () => {
+      committed.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+    expect(fetcher).toHaveBeenCalledTimes(2);
+
+    await act(async () => rendered.rerender(view(true)));
+    await act(async () => rendered.rerender(view(false)));
+    await act(async () => {
+      committed.adapter.search!("alice");
+      await vi.advanceTimersByTimeAsync(0);
+    });
+
+    expect(fetcher).toHaveBeenCalledTimes(3);
+  });
+
   it("drops an in-flight fetch when the query returns to a cached value", async () => {
     const resolvers: Record<
       string,

package/src/unstable/useLiveCompletionAdapter.ts

@@ -84,6 +84,8 @@ export function unstable_useLiveCompletionAdapter(
   const pendingQueryRef = useRef<string | null>(null);
   const retryableQueryRef = useRef<string | null>(null);
   const pendingRetryQueryRef = useRef<string | null>(null);
+  const inactiveRef = useRef(true);
+  const deferredQueryRef = useRef<string | null>(null);
 
   const cancelTimer = useCallback(() => {
     if (timerRef.current !== null) {
@@ -102,6 +104,10 @@ export function unstable_useLiveCompletionAdapter(
   const scheduleFetch = useCallback(
     (query: string) => {
       if (!enabled) return;
+      if (inactiveRef.current) {
+        deferredQueryRef.current = query;
+        return;
+      }
       if (pendingQueryRef.current === query) return;
       rearmPendingRetry();
       if (retryableQueryRef.current === query) {
@@ -136,6 +142,11 @@ export function unstable_useLiveCompletionAdapter(
     [enabled, debounceMs, cancelTimer, rearmPendingRetry],
   );
 
+  const scheduleFetchRef = useRef(scheduleFetch);
+  useLayoutEffect(() => {
+    scheduleFetchRef.current = scheduleFetch;
+  }, [scheduleFetch]);
+
   const invalidatePending = useCallback(() => {
     rearmPendingRetry();
     cancelTimer();
@@ -162,7 +173,18 @@ export function unstable_useLiveCompletionAdapter(
     );
   }, [enabled, invalidatePending]);
 
-  useEffect(() => cancelTimer, [cancelTimer]);
+  // Render-time searches can outlive an abandoned render, so they only arm
+  // request work after this hook commits.
+  useLayoutEffect(() => {
+    inactiveRef.current = false;
+    const deferredQuery = deferredQueryRef.current;
+    deferredQueryRef.current = null;
+    if (deferredQuery !== null) scheduleFetchRef.current(deferredQuery);
+    return () => {
+      inactiveRef.current = true;
+      invalidatePending();
+    };
+  }, [invalidatePending]);
 
   // Arm retries only after the failed state commits. Arming during rejection
   // would let the failure render immediately schedule another request.
@@ -179,13 +201,21 @@ export function unstable_useLiveCompletionAdapter(
         // queueMicrotask so they are not dispatched while another component renders.
         if (query !== state.query || retryableQueryRef.current === query) {
           queueMicrotask(() => scheduleFetch(query));
-        } else if (
-          pendingQueryRef.current !== null &&
-          pendingQueryRef.current !== query
-        ) {
-          // the query returned to a cached value while a fetch for a different
-          // query is in flight; drop it so its result cannot overwrite the cache
-          queueMicrotask(invalidatePending);
+        } else {
+          queueMicrotask(() => {
+            if (
+              deferredQueryRef.current !== null &&
+              deferredQueryRef.current !== query
+            ) {
+              deferredQueryRef.current = null;
+            }
+            if (
+              pendingQueryRef.current !== null &&
+              pendingQueryRef.current !== query
+            ) {
+              invalidatePending();
+            }
+          });
         }
         return state.items;
       },

package/src/unstable/useMentionAdapter.test.tsx

@@ -11,7 +11,10 @@ import type { Unstable_TriggerAdapter } from "@assistant-ui/core";
 import { useState } from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { TriggerNavigationResource } from "../primitives/composer/trigger/triggerNavigationResource";
-import { unstable_useMentionAdapter } from "./useMentionAdapter";
+import {
+  unstable_useMentionAdapter,
+  type Unstable_UseMentionAdapterOptions,
+} from "./useMentionAdapter";
 
 const runtime = vi.hoisted(() => {
   const state = {
@@ -66,6 +69,49 @@ describe("unstable_useMentionAdapter", () => {
     runtime.events.length = 0;
   });
 
+  it("ignores flat items when categories is explicitly empty", () => {
+    const { result } = renderHook(() =>
+      unstable_useMentionAdapter({
+        items: [{ id: "alice", type: "person", label: "Alice" }],
+        categories: [],
+      }),
+    );
+
+    expect(result.current.adapter.search?.("")).toEqual([]);
+  });
+
+  it("does not expose flat items when the last category is removed", () => {
+    const items = [{ id: "alice", type: "person", label: "Alice" }];
+    const { result, rerender } = renderHook(
+      (options: Unstable_UseMentionAdapterOptions) =>
+        unstable_useMentionAdapter(options),
+      {
+        initialProps: {
+          items,
+          categories: [
+            {
+              id: "people",
+              label: "People",
+              items: [{ id: "bob", type: "person", label: "Bob" }],
+            },
+          ],
+        },
+      },
+    );
+
+    expect(result.current.adapter.search?.("")).toEqual([
+      { id: "bob", type: "person", label: "Bob" },
+    ]);
+
+    rerender({ items, categories: [] });
+    expect(result.current.adapter.search?.("")).toEqual([]);
+
+    rerender({ items });
+    expect(result.current.adapter.search?.("")).toEqual([
+      { id: "alice", type: "person", label: "Alice" },
+    ]);
+  });
+
   it("keeps categorized model-context tools current", () => {
     runtime.state.tools = {
       searchDocs: { description: "Search documentation" },

package/src/unstable/useMentionAdapter.ts

@@ -95,7 +95,7 @@ const toolMentionSource: ModelContextSnapshotSource<
   read: (aui) => {
     const tools = aui.thread.getModelContext().tools;
     if (!tools) return EMPTY_TOOL_MENTIONS;
-    const mentions: Record<string, string | undefined> = {};
+    const mentions = Object.create(null) as Record<string, string | undefined>;
     for (const [name, tool] of Object.entries(tools)) {
       mentions[name] = tool.description;
     }
@@ -152,7 +152,7 @@ export function unstable_useMentionAdapter(
   const formatter = options?.formatter;
   const onInserted = options?.onInserted;
   const isCategorized =
-    (categories !== undefined && categories.length > 0) ||
+    categories !== undefined ||
     (toolsConfig?.category !== undefined && items === undefined);
   const toolMentions = useModelContextSnapshot(
     aui,

package/src/unstable/useMessageStallDetection.ts

@@ -2,6 +2,7 @@
 
 import { useEffect, useRef, useState } from "react";
 import { useAuiState } from "@assistant-ui/store";
+import { useShallowSelector } from "@assistant-ui/store/internal";
 
 export type Unstable_MessageStallDetectionOptions = {
   /**
@@ -23,11 +24,10 @@ export type Unstable_MessageStallDetection = {
  * @deprecated Under active development and might change without notice.
  *
  * Detects mid-run output stalls on the current message: while the message is
- * running, watches a fingerprint of its content (part count plus text,
- * argument, and result sizes) and reports a stall once the fingerprint stops
- * changing for `thresholdMs`. Useful for re-surfacing a "still working"
- * indicator during tool think-time or provider stalls, after the first
- * tokens have already streamed.
+ * running, watches its text, reasoning, and tool-argument values plus tool-result
+ * availability and reports a stall once they stop changing for `thresholdMs`.
+ * Useful for re-surfacing a "still working" indicator during tool think-time or
+ * provider stalls, after the first tokens have already streamed.
  *
  * Must be used inside a message scope.
  */
@@ -36,20 +36,27 @@ export function unstable_useMessageStallDetection(
 ): Unstable_MessageStallDetection {
   const thresholdMs = options?.thresholdMs ?? 2000;
 
-  const fingerprint = useAuiState((s) => {
-    if (s.message.status?.type !== "running") return undefined;
-    let size = 0;
-    for (const part of s.message.content) {
-      if (part.type === "text" || part.type === "reasoning") {
-        size += part.text.length;
-      } else if (part.type === "tool-call") {
-        size += part.argsText.length + (part.result !== undefined ? 1 : 0);
+  const activity = useAuiState(
+    useShallowSelector((s) => {
+      const running = s.message.status?.type === "running";
+      if (!running) return [false];
+
+      const values: unknown[] = [true, s.message.content.length];
+
+      for (const part of s.message.content) {
+        if (part.type === "text" || part.type === "reasoning") {
+          values.push(part.type, part.text);
+        } else if (part.type === "tool-call") {
+          values.push(part.type, part.argsText, part.result !== undefined);
+        } else {
+          values.push(part.type);
+        }
       }
-    }
-    return `${s.message.content.length}:${size}`;
-  });
+      return values;
+    }),
+  );
 
-  const running = fingerprint !== undefined;
+  const running = activity[0] === true;
   const lastActivityRef = useRef(Date.now());
   const [stalled, setStalled] = useState(false);
   const [, setTick] = useState(0);
@@ -58,7 +65,7 @@ export function unstable_useMessageStallDetection(
     if (!running) return undefined;
     lastActivityRef.current = Date.now();
     return undefined;
-  }, [running, fingerprint]);
+  }, [running, activity]);
 
   useEffect(() => {
     if (!running) {
@@ -75,7 +82,7 @@ export function unstable_useMessageStallDetection(
     setStalled(false);
     const id = setTimeout(() => setStalled(true), thresholdMs - sinceActivity);
     return () => clearTimeout(id);
-  }, [running, fingerprint, thresholdMs]);
+  }, [running, activity, thresholdMs]);
 
   useEffect(() => {
     if (!stalled) return undefined;

package/src/unstable/webmcp/convertTools.test.ts

@@ -267,43 +267,203 @@ describe("toWebMcpTool cancellation", () => {
     });
   });
 
-  it("merges the caller signal with the lifecycle signal", async () => {
-    const lifecycle = new AbortController();
-    const caller = new AbortController();
-    const descriptor = descriptorFor(
-      {
-        execute: async (_args: unknown, context: any) =>
-          new Promise((_resolve, reject) => {
-            context.abortSignal.addEventListener("abort", () =>
-              reject(new Error("aborted")),
-            );
+  it.for(["caller", "lifecycle"] as const)(
+    "settles while async validation is pending when the %s signal aborts",
+    async (abortedSignal) => {
+      const lifecycle = new AbortController();
+      const caller = new AbortController();
+      let finishValidation!: (result: { issues?: readonly unknown[] }) => void;
+      const schema = z.object({ city: z.string() });
+      (schema as any)["~standard"] = {
+        ...schema["~standard"],
+        validate: () =>
+          new Promise<{ issues?: readonly unknown[] }>((resolve) => {
+            finishValidation = resolve;
           }),
-      },
-      lifecycle.signal,
+      };
+      const execute = vi.fn(async () => "never");
+      const pending = descriptorFor(
+        { execute, parameters: schema },
+        lifecycle.signal,
+      ).execute({ city: "Paris" }, { signal: caller.signal });
+
+      (abortedSignal === "caller" ? caller : lifecycle).abort();
+
+      await expect(pending).resolves.toEqual({
+        isError: true,
+        content: [text("Tool execution was cancelled.")],
+      });
+      expect(execute).not.toHaveBeenCalled();
+      finishValidation({});
+    },
+  );
+
+  it("consumes a validator rejection after cancellation", async () => {
+    const caller = new AbortController();
+    let failValidation!: (error: unknown) => void;
+    const schema = z.object({ city: z.string() });
+    (schema as any)["~standard"] = {
+      ...schema["~standard"],
+      validate: () =>
+        new Promise((_resolve, reject) => {
+          failValidation = reject;
+        }),
+    };
+    const execute = vi.fn(async () => "never");
+    const pending = descriptorFor({ execute, parameters: schema }).execute(
+      { city: "Paris" },
+      { signal: caller.signal },
     );
 
-    const pending = descriptor.execute({}, { signal: caller.signal });
-    lifecycle.abort();
+    caller.abort();
+
     await expect(pending).resolves.toEqual({
       isError: true,
-      content: [text("aborted")],
+      content: [text("Tool execution was cancelled.")],
     });
+    failValidation(new Error("late validation failure"));
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(execute).not.toHaveBeenCalled();
   });
 
-  it("returns an error result when the caller signal cannot be merged", async () => {
+  it("prefers cancellation when validation aborts before rejecting", async () => {
+    const caller = new AbortController();
+    const schema = z.object({ city: z.string() });
+    (schema as any)["~standard"] = {
+      ...schema["~standard"],
+      validate: () => {
+        caller.abort();
+        return Promise.reject(new Error("validation failed"));
+      },
+    };
     const execute = vi.fn(async () => "never");
+
+    const result = await descriptorFor({ execute, parameters: schema }).execute(
+      { city: "Paris" },
+      { signal: caller.signal },
+    );
+
+    expect(result).toEqual({
+      isError: true,
+      content: [text("Tool execution was cancelled.")],
+    });
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("does not execute when cancellation follows validation", async () => {
+    const caller = new AbortController();
+    const schema = z.object({ city: z.string() });
+    (schema as any)["~standard"] = {
+      ...schema["~standard"],
+      validate: () => ({
+        then: (resolve: (value: { issues?: readonly unknown[] }) => void) => {
+          resolve({});
+          caller.abort();
+        },
+      }),
+    };
+    const execute = vi.fn(async () => "never");
+
+    const result = await descriptorFor({ execute, parameters: schema }).execute(
+      { city: "Paris" },
+      { signal: caller.signal },
+    );
+
+    expect(result).toEqual({
+      isError: true,
+      content: [text("Tool execution was cancelled.")],
+    });
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it.for(["caller", "lifecycle"] as const)(
+    "merges signals without AbortSignal.any when the %s signal aborts",
+    async (abortedSignal) => {
+      const lifecycle = new AbortController();
+      const caller = new AbortController();
+      const abortSignalConstructor = AbortSignal as typeof AbortSignal & {
+        any?: (signals: Iterable<AbortSignal>) => AbortSignal;
+      };
+      const originalAbortSignalAny = abortSignalConstructor.any;
+      Object.defineProperty(abortSignalConstructor, "any", {
+        configurable: true,
+        value: () => {
+          throw new Error("AbortSignal.any is not available");
+        },
+      });
+      try {
+        const descriptor = descriptorFor(
+          {
+            execute: async (_args: unknown, context: any) =>
+              new Promise((_resolve, reject) => {
+                context.abortSignal.addEventListener("abort", () =>
+                  reject(new Error("aborted")),
+                );
+              }),
+          },
+          lifecycle.signal,
+        );
+
+        const pending = descriptor.execute({}, { signal: caller.signal });
+        (abortedSignal === "caller" ? caller : lifecycle).abort();
+        await expect(pending).resolves.toEqual({
+          isError: true,
+          content: [text("aborted")],
+        });
+      } finally {
+        Object.defineProperty(abortSignalConstructor, "any", {
+          configurable: true,
+          value: originalAbortSignalAny,
+        });
+      }
+    },
+  );
+
+  it("removes merged signal listeners after execution", async () => {
+    const lifecycle = new AbortController();
+    const caller = new AbortController();
+    const callerRemove = vi.spyOn(caller.signal, "removeEventListener");
+    const lifecycleRemove = vi.spyOn(lifecycle.signal, "removeEventListener");
+
+    const result = await descriptorFor(
+      { execute: async () => "ok" },
+      lifecycle.signal,
+    ).execute({}, { signal: caller.signal });
+
+    expect(result).toEqual({ content: [text("ok")] });
+    expect(callerRemove).toHaveBeenCalledTimes(1);
+    expect(lifecycleRemove).toHaveBeenCalledTimes(1);
+  });
+
+  it("merges a caller signal that is not a native AbortSignal", async () => {
+    const listeners: (() => void)[] = [];
     const foreignSignal = {
       aborted: false,
-      addEventListener: () => {},
+      reason: new Error("host cancelled"),
+      addEventListener: (_type: string, listener: () => void) => {
+        listeners.push(listener);
+      },
       removeEventListener: () => {},
     } as unknown as AbortSignal;
 
-    const result = await descriptorFor(
-      { execute },
+    const pending = descriptorFor(
+      {
+        execute: async (_args: unknown, context: any) =>
+          new Promise((_resolve, reject) => {
+            context.abortSignal.addEventListener("abort", () =>
+              reject(new Error("aborted")),
+            );
+          }),
+      },
       new AbortController().signal,
     ).execute({}, { signal: foreignSignal });
-    expect(result.isError).toBe(true);
-    expect(execute).not.toHaveBeenCalled();
+
+    for (const listener of listeners) listener();
+
+    await expect(pending).resolves.toEqual({
+      isError: true,
+      content: [text("aborted")],
+    });
   });
 });