@assistant-ui/react 0.15.17 → 0.15.19

package/src/unstable/useMentionAdapter.test.tsx

@@ -0,0 +1,383 @@
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import {
+  createTapRoot,
+  flushTapSync,
+  resource,
+  useResource,
+} from "@assistant-ui/tap";
+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,
+  type Unstable_UseMentionAdapterOptions,
+} from "./useMentionAdapter";
+
+const runtime = vi.hoisted(() => {
+  const state = {
+    tools: {} as Record<string, { description?: string }>,
+  };
+  const listeners = new Map<string, Set<() => void>>();
+  const events: string[] = [];
+  return {
+    state,
+    listeners,
+    events,
+    client: {
+      thread: {
+        getModelContext: () => ({ tools: state.tools }),
+      },
+      on: (
+        selector: string | { scope: string; event: string },
+        listener: () => void,
+      ) => {
+        const event = typeof selector === "string" ? selector : selector.event;
+        events.push(event);
+        const eventListeners = listeners.get(event) ?? new Set();
+        eventListeners.add(listener);
+        listeners.set(event, eventListeners);
+        return () => eventListeners.delete(listener);
+      },
+    },
+    emit: (event: string) =>
+      listeners.get(event)?.forEach((listener) => listener()),
+  };
+});
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => runtime.client,
+}));
+
+const TestNavigationResource = resource(function useTestNavigationResource(
+  initialAdapter: Unstable_TriggerAdapter,
+) {
+  const [adapter, setAdapter] = useState(initialAdapter);
+  const navigation = useResource(
+    TriggerNavigationResource({ adapter, query: "", open: true }),
+  );
+  return { navigation, setAdapter };
+});
+
+describe("unstable_useMentionAdapter", () => {
+  beforeEach(() => {
+    runtime.state.tools = {};
+    runtime.listeners.clear();
+    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" },
+    };
+
+    const { result } = renderHook(() =>
+      unstable_useMentionAdapter({
+        categories: [{ id: "people", label: "People", items: [] }],
+        includeModelContextTools: true,
+      }),
+    );
+    const initialAdapter = result.current.adapter;
+    const navigationRoot = createTapRoot(function NavigationRoot() {
+      return useResource(TestNavigationResource(initialAdapter));
+    });
+
+    expect(navigationRoot.getValue().navigation.categories).toEqual([
+      { id: "people", label: "People" },
+      { id: "tools", label: "Tools" },
+    ]);
+    expect(runtime.events).toContain("thread.modelContextUpdate");
+
+    act(() =>
+      flushTapSync(() =>
+        navigationRoot.getValue().navigation.selectCategory("tools"),
+      ),
+    );
+
+    expect(navigationRoot.getValue().navigation.items).toEqual([
+      {
+        id: "searchDocs",
+        type: "tool",
+        label: "searchDocs",
+        description: "Search documentation",
+      },
+    ]);
+
+    act(() => {
+      runtime.state.tools = {
+        createIssue: { description: "Create an issue" },
+      };
+      runtime.emit("thread.modelContextUpdate");
+    });
+
+    expect(result.current.adapter).not.toBe(initialAdapter);
+    act(() =>
+      flushTapSync(() =>
+        navigationRoot.getValue().setAdapter(result.current.adapter),
+      ),
+    );
+    expect(navigationRoot.getValue().navigation.items).toEqual([
+      {
+        id: "createIssue",
+        type: "tool",
+        label: "createIssue",
+        description: "Create an issue",
+      },
+    ]);
+    act(() => {
+      runtime.state.tools = {};
+      runtime.emit("thread.modelContextUpdate");
+    });
+    act(() =>
+      flushTapSync(() =>
+        navigationRoot.getValue().setAdapter(result.current.adapter),
+      ),
+    );
+
+    expect(navigationRoot.getValue().navigation.categories).toEqual([
+      { id: "people", label: "People" },
+    ]);
+    expect(navigationRoot.getValue().navigation.items).toEqual([]);
+    navigationRoot.unmount();
+  });
+
+  it("keeps flat model-context tools current while navigation is open", () => {
+    runtime.state.tools = {
+      searchDocs: { description: "Search documentation" },
+    };
+
+    const { result } = renderHook(() => unstable_useMentionAdapter());
+    const initialAdapter = result.current.adapter;
+    const navigationRoot = createTapRoot(function NavigationRoot() {
+      return useResource(TestNavigationResource(initialAdapter));
+    });
+
+    expect(runtime.events).toContain("thread.modelContextUpdate");
+
+    expect(navigationRoot.getValue().navigation.items).toEqual([
+      {
+        id: "searchDocs",
+        type: "tool",
+        label: "searchDocs",
+        description: "Search documentation",
+      },
+    ]);
+
+    act(() => {
+      runtime.state.tools = {
+        createIssue: { description: "Create an issue" },
+      };
+      runtime.emit("thread.modelContextUpdate");
+    });
+
+    expect(result.current.adapter).not.toBe(initialAdapter);
+    act(() =>
+      flushTapSync(() =>
+        navigationRoot.getValue().setAdapter(result.current.adapter),
+      ),
+    );
+    expect(navigationRoot.getValue().navigation.items).toEqual([
+      {
+        id: "createIssue",
+        type: "tool",
+        label: "createIssue",
+        description: "Create an issue",
+      },
+    ]);
+    navigationRoot.unmount();
+  });
+
+  it("refreshes tool mentions when the selected thread changes", () => {
+    runtime.state.tools = {
+      searchDocs: { description: "Search documentation" },
+    };
+
+    const { result } = renderHook(() => unstable_useMentionAdapter());
+    const initialAdapter = result.current.adapter;
+
+    expect(runtime.events).toContain("threads.selectionChanged");
+
+    act(() => {
+      runtime.state.tools = {
+        createIssue: { description: "Create an issue" },
+      };
+      runtime.emit("threads.selectionChanged");
+    });
+
+    expect(result.current.adapter).not.toBe(initialAdapter);
+    expect(result.current.adapter.search?.("")).toEqual([
+      {
+        id: "createIssue",
+        type: "tool",
+        label: "createIssue",
+        description: "Create an issue",
+      },
+    ]);
+  });
+
+  it("keeps adapter identity when tool mentions are unchanged", () => {
+    runtime.state.tools = {
+      searchDocs: { description: "Search documentation" },
+    };
+
+    const { result } = renderHook(() => unstable_useMentionAdapter());
+    const initialAdapter = result.current.adapter;
+
+    act(() => {
+      runtime.state.tools = {
+        searchDocs: { description: "Search documentation" },
+      };
+      runtime.emit("thread.modelContextUpdate");
+    });
+
+    expect(result.current.adapter).toBe(initialAdapter);
+  });
+
+  it("observes a description mutated in place on a stable tool object", () => {
+    const tool = { description: "Search documentation" };
+    runtime.state.tools = { searchDocs: tool };
+
+    const { result } = renderHook(() => unstable_useMentionAdapter());
+    const initialAdapter = result.current.adapter;
+
+    act(() => {
+      tool.description = "Search the docs site";
+      runtime.emit("thread.modelContextUpdate");
+    });
+
+    expect(result.current.adapter).not.toBe(initialAdapter);
+    expect(result.current.adapter.search?.("")).toEqual([
+      {
+        id: "searchDocs",
+        type: "tool",
+        label: "searchDocs",
+        description: "Search the docs site",
+      },
+    ]);
+  });
+
+  it("uses an explicit tool category without custom categories", () => {
+    runtime.state.tools = {
+      searchDocs: { description: "Search documentation" },
+    };
+
+    const { result } = renderHook(() =>
+      unstable_useMentionAdapter({
+        includeModelContextTools: {
+          category: { id: "actions", label: "Actions" },
+        },
+      }),
+    );
+
+    expect(result.current.adapter.categories()).toEqual([
+      { id: "actions", label: "Actions" },
+    ]);
+    expect(result.current.adapter.categoryItems("actions")).toEqual([
+      {
+        id: "searchDocs",
+        type: "tool",
+        label: "searchDocs",
+        description: "Search documentation",
+      },
+    ]);
+  });
+
+  it("keeps an async explicit item list flat when a tool category is configured", () => {
+    runtime.state.tools = {
+      searchDocs: { description: "Search documentation" },
+      createIssue: { description: "Create an issue" },
+    };
+
+    const { result, rerender } = renderHook(
+      ({
+        items,
+      }: {
+        items: readonly {
+          id: string;
+          type: string;
+          label: string;
+        }[];
+      }) =>
+        unstable_useMentionAdapter({
+          items,
+          includeModelContextTools: {
+            category: { id: "actions", label: "Actions" },
+          },
+        }),
+      { initialProps: { items: [] } },
+    );
+
+    expect(result.current.adapter.categories()).toEqual([]);
+
+    rerender({
+      items: [
+        {
+          id: "searchDocs",
+          type: "person",
+          label: "Documentation owner",
+        },
+      ],
+    });
+
+    expect(result.current.adapter.categories()).toEqual([]);
+    expect(result.current.adapter.search?.("")).toEqual([
+      {
+        id: "searchDocs",
+        type: "person",
+        label: "Documentation owner",
+      },
+      {
+        id: "createIssue",
+        type: "tool",
+        label: "createIssue",
+        description: "Create an issue",
+      },
+    ]);
+  });
+});

package/src/unstable/useMentionAdapter.ts

@@ -10,6 +10,12 @@ import type {
 } from "@assistant-ui/core";
 import { unstable_defaultDirectiveFormatter } from "@assistant-ui/core";
 import type { ReadonlyJSONObject } from "assistant-stream/utils";
+import { matchesTriggerItemQuery } from "../primitives/composer/trigger/matchesTriggerItemQuery";
+import {
+  shallowEqualRecords,
+  useModelContextSnapshot,
+  type ModelContextSnapshotSource,
+} from "./useModelContextSnapshot";
 
 /** Icon component shape consumed by `ComposerTriggerPopover`'s `iconMap`. */
 export type Unstable_IconComponent = FC<{ className?: string }>;
@@ -31,7 +37,10 @@ export type Unstable_MentionCategory = {
 };
 
 export type Unstable_ModelContextToolsOptions = {
-  /** Wrap tools in a dedicated category (drill-down mode). */
+  /**
+   * Wrap tools in a dedicated category. Selects drill-down mode on its own
+   * when `categories` is unset and `items` is omitted.
+   */
   readonly category?: { readonly id: string; readonly label: string };
   /** Format tool name for display. */
   readonly formatLabel?: (toolName: string) => string;
@@ -40,7 +49,10 @@ export type Unstable_ModelContextToolsOptions = {
 };
 
 export type Unstable_UseMentionAdapterOptions = {
-  /** Flat mention list. Ignored when `categories` is set. */
+  /**
+   * Flat mention list. Ignored when `categories` is set, and keeps the
+   * adapter flat when a tool `category` is configured.
+   */
   readonly items?: readonly Unstable_Mention[];
   /** Categorized mentions for drill-down navigation. */
   readonly categories?: readonly Unstable_MentionCategory[];
@@ -49,7 +61,7 @@ export type Unstable_UseMentionAdapterOptions = {
    * - `false`: exclude.
    * - `true`: include (default when no `items`/`categories`; as a category
    *   if `categories` is set, flat otherwise).
-   * - object: explicit config.
+   * - object: explicit config; `category` also selects drill-down mode.
    *
    * Omitted → defaults to `true` iff neither `items` nor `categories`.
    */
@@ -71,6 +83,42 @@ export type Unstable_MentionDirective = {
   readonly onInserted?: ((item: Unstable_TriggerItem) => void) | undefined;
 };
 
+const EMPTY_TOOL_MENTIONS: Readonly<Record<string, string | undefined>> =
+  Object.freeze({});
+
+// The description is copied out rather than reached through the tool, so a
+// provider that returns a stable tool object and edits it in place is observed.
+const toolMentionSource: ModelContextSnapshotSource<
+  Readonly<Record<string, string | undefined>>
+> = {
+  empty: EMPTY_TOOL_MENTIONS,
+  read: (aui) => {
+    const tools = aui.thread.getModelContext().tools;
+    if (!tools) return EMPTY_TOOL_MENTIONS;
+    const mentions = Object.create(null) as Record<string, string | undefined>;
+    for (const [name, tool] of Object.entries(tools)) {
+      mentions[name] = tool.description;
+    }
+    return mentions;
+  },
+  subscribe: (aui, onChange) => {
+    const unsubscribeContext = aui.on("thread.modelContextUpdate", onChange);
+    // Rebinding the thread event subject to a new thread does not replay it,
+    // so a switch between threads carrying different providers is its own
+    // refresh trigger. Subscribed globally because a composer can render
+    // without a thread list scope.
+    const unsubscribeSelection = aui.on(
+      { scope: "*", event: "threads.selectionChanged" },
+      onChange,
+    );
+    return () => {
+      unsubscribeContext();
+      unsubscribeSelection();
+    };
+  },
+  isEqual: shallowEqualRecords,
+};
+
 /**
  * @deprecated Under active development and might change without notice.
  *
@@ -103,50 +151,48 @@ export function unstable_useMentionAdapter(
   const wantsTools = includeTools !== false;
   const formatter = options?.formatter;
   const onInserted = options?.onInserted;
+  const isCategorized =
+    categories !== undefined ||
+    (toolsConfig?.category !== undefined && items === undefined);
+  const toolMentions = useModelContextSnapshot(
+    aui,
+    wantsTools,
+    toolMentionSource,
+  );
 
   const adapter = useMemo<Unstable_TriggerAdapter>(() => {
-    const getModelContextTools = (): Unstable_TriggerItem[] => {
-      if (!wantsTools) return [];
-      const ctx = aui.thread.getModelContext();
-      const tools = ctx.tools;
-      if (!tools) return [];
-      const formatLabel = toolsConfig?.formatLabel;
-      const defaultIcon = toolsConfig?.icon;
-      return Object.entries(tools).map(([name, tool]) =>
-        toTriggerItem({
-          id: name,
-          type: "tool",
-          label: formatLabel ? formatLabel(name) : name,
-          description: tool.description ?? undefined,
-          icon: defaultIcon,
-        }),
-      );
-    };
+    const formatLabel = toolsConfig?.formatLabel;
+    const defaultIcon = toolsConfig?.icon;
+    const toolItems = wantsTools
+      ? Object.entries(toolMentions).map(([name, description]) =>
+          toTriggerItem({
+            id: name,
+            type: "tool",
+            label: formatLabel ? formatLabel(name) : name,
+            description,
+            icon: defaultIcon,
+          }),
+        )
+      : [];
 
     // Categorized: drill-down mode
-    if (categories && categories.length > 0) {
-      const groups = categories.map((cat) => ({
+    if (isCategorized) {
+      const groups = (categories ?? []).map((cat) => ({
         id: cat.id,
         label: cat.label,
         items: cat.items.map(toTriggerItem),
       }));
-
-      let toolCategory: {
-        id: string;
-        label: string;
-        items: Unstable_TriggerItem[];
-      } | null = null;
-      if (wantsTools) {
-        const toolItems = getModelContextTools();
-        if (toolItems.length > 0) {
-          toolCategory = {
-            id: toolsConfig?.category?.id ?? "tools",
-            label: toolsConfig?.category?.label ?? "Tools",
-            items: toolItems,
-          };
-        }
-      }
-      const allGroups = toolCategory ? [...groups, toolCategory] : groups;
+      const allGroups =
+        toolItems.length > 0
+          ? [
+              ...groups,
+              {
+                id: toolsConfig?.category?.id ?? "tools",
+                label: toolsConfig?.category?.label ?? "Tools",
+                items: toolItems,
+              },
+            ]
+          : groups;
 
       return {
         categories: () => allGroups.map(({ id, label }) => ({ id, label })),
@@ -155,30 +201,29 @@ export function unstable_useMentionAdapter(
           const lower = query.toLowerCase();
           return allGroups
             .flatMap((g) => g.items)
-            .filter((item) => matchesQuery(item, lower));
+            .filter((item) => matchesTriggerItemQuery(item, lower));
         },
       };
     }
 
     // Flat: items + (optionally) tools, all in one search pool
     const flatItems = (items ?? []).map(toTriggerItem);
-    const getFlatPool = (): Unstable_TriggerItem[] => {
-      if (!wantsTools) return flatItems;
-      const toolItems = getModelContextTools();
-      // Dedupe by id — explicit items win.
-      const seen = new Set(flatItems.map((i) => i.id));
-      return [...flatItems, ...toolItems.filter((t) => !seen.has(t.id))];
-    };
+    // Dedupe by id — explicit items win.
+    const seen = new Set(flatItems.map((i) => i.id));
+    const flatPool = [
+      ...flatItems,
+      ...toolItems.filter((t) => !seen.has(t.id)),
+    ];
 
     return {
       categories: (): readonly Unstable_TriggerCategory[] => [],
       categoryItems: () => [],
       search: (query) => {
         const lower = query.toLowerCase();
-        return getFlatPool().filter((item) => matchesQuery(item, lower));
+        return flatPool.filter((item) => matchesTriggerItemQuery(item, lower));
       },
     };
-  }, [aui, items, categories, wantsTools, toolsConfig]);
+  }, [items, categories, wantsTools, toolsConfig, isCategorized, toolMentions]);
 
   const directive = useMemo<Unstable_MentionDirective>(
     () => ({
@@ -207,11 +252,3 @@ function toTriggerItem(m: Unstable_Mention): Unstable_TriggerItem {
     ...(metadata !== undefined ? { metadata } : {}),
   };
 }
-
-function matchesQuery(item: Unstable_TriggerItem, lower: string): boolean {
-  if (!lower) return true;
-  if (item.id.toLowerCase().includes(lower)) return true;
-  if (item.label.toLowerCase().includes(lower)) return true;
-  if (item.description?.toLowerCase().includes(lower)) return true;
-  return false;
-}

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;