@assistant-ui/react 0.15.16 → 0.15.18

package/src/unstable/useModelContextSnapshot.test.tsx

@@ -0,0 +1,132 @@
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import type { AssistantClient } from "@assistant-ui/store";
+import { describe, expect, it, vi } from "vitest";
+import {
+  shallowEqualRecords,
+  useModelContextSnapshot,
+  type ModelContextSnapshotSource,
+} from "./useModelContextSnapshot";
+
+const EMPTY: Readonly<Record<string, string>> = Object.freeze({});
+
+const createSource = (
+  overrides: Partial<
+    ModelContextSnapshotSource<Readonly<Record<string, string>>>
+  > = {},
+) => {
+  const listeners = new Set<() => void>();
+  const state = { tools: EMPTY as Record<string, string> };
+  const unsubscribe = vi.fn();
+  const source: ModelContextSnapshotSource<Readonly<Record<string, string>>> = {
+    empty: EMPTY,
+    read: vi.fn(() => ({ ...state.tools })),
+    subscribe: vi.fn((_aui, onChange) => {
+      listeners.add(onChange);
+      return () => {
+        listeners.delete(onChange);
+        unsubscribe();
+      };
+    }),
+    isEqual: shallowEqualRecords,
+    ...overrides,
+  };
+  return {
+    source,
+    state,
+    unsubscribe,
+    notify: () => listeners.forEach((listener) => listener()),
+  };
+};
+
+const aui = {} as AssistantClient;
+
+describe("useModelContextSnapshot", () => {
+  it("seeds from the source on the first render when enabled", () => {
+    const { source } = createSource();
+    (source.read as ReturnType<typeof vi.fn>).mockReturnValue({ a: "1" });
+
+    const { result } = renderHook(() =>
+      useModelContextSnapshot(aui, true, source),
+    );
+
+    expect(result.current).toEqual({ a: "1" });
+  });
+
+  it("neither reads nor subscribes while disabled", () => {
+    const { source } = createSource();
+
+    const { result } = renderHook(() =>
+      useModelContextSnapshot(aui, false, source),
+    );
+
+    expect(result.current).toBe(source.empty);
+    expect(source.read).not.toHaveBeenCalled();
+    expect(source.subscribe).not.toHaveBeenCalled();
+  });
+
+  it("keeps the previous reference when isEqual reports no change", () => {
+    const { source, state, notify } = createSource();
+    state.tools = { a: "1" };
+
+    const { result } = renderHook(() =>
+      useModelContextSnapshot(aui, true, source),
+    );
+    const first = result.current;
+
+    act(() => {
+      state.tools = { a: "1" };
+      notify();
+    });
+    expect(result.current).toBe(first);
+
+    act(() => {
+      state.tools = { a: "2" };
+      notify();
+    });
+    expect(result.current).not.toBe(first);
+    expect(result.current).toEqual({ a: "2" });
+  });
+
+  it("takes every notification when the source omits isEqual", () => {
+    const { source, state, notify } = createSource({ isEqual: undefined });
+    state.tools = { a: "1" };
+
+    const { result } = renderHook(() =>
+      useModelContextSnapshot(aui, true, source),
+    );
+    const first = result.current;
+
+    act(() => notify());
+
+    expect(result.current).not.toBe(first);
+    expect(result.current).toEqual({ a: "1" });
+  });
+
+  it("unsubscribes on unmount", () => {
+    const { source, unsubscribe } = createSource();
+    const { unmount } = renderHook(() =>
+      useModelContextSnapshot(aui, true, source),
+    );
+
+    unmount();
+
+    expect(unsubscribe).toHaveBeenCalled();
+  });
+});
+
+describe("shallowEqualRecords", () => {
+  it("compares own keys by identity", () => {
+    expect(shallowEqualRecords({ a: 1 }, { a: 1 })).toBe(true);
+    expect(shallowEqualRecords({ a: 1 }, { a: 2 })).toBe(false);
+    expect(shallowEqualRecords({ a: 1 }, { a: 1, b: 2 })).toBe(false);
+    expect(shallowEqualRecords({ a: undefined }, { b: undefined })).toBe(false);
+  });
+
+  it("does not match a key inherited from Object.prototype", () => {
+    expect(
+      shallowEqualRecords({ toString: Object.prototype.toString }, { a: 1 }),
+    ).toBe(false);
+  });
+});

package/src/unstable/useModelContextSnapshot.ts

@@ -0,0 +1,67 @@
+"use client";
+
+import { useEffect, useState } from "react";
+import type { AssistantClient } from "@assistant-ui/store";
+import type { Unsubscribe } from "@assistant-ui/core";
+
+/**
+ * Where a snapshot comes from and what counts as a change to it. Declared once
+ * per consumer at module scope: the subscription effect keys on this object, so
+ * an inline source without an `isEqual` re-reads itself into a render loop.
+ */
+export type ModelContextSnapshotSource<T> = {
+  /** Seeds the snapshot when it starts disabled. */
+  readonly empty: T;
+  /** Projects the fields the consumer observes out of the live model context. */
+  read(aui: AssistantClient): T;
+  subscribe(aui: AssistantClient, onChange: () => void): Unsubscribe;
+  /**
+   * An equal pair keeps the previous reference. Omitted by a consumer that
+   * re-derives on every notification, where a caller republishing an unchanged
+   * tool set is how it asks to be re-run.
+   */
+  isEqual?(previous: T, next: T): boolean;
+};
+
+/**
+ * `getModelContext()` rebuilds its result on every call, so a projection of the
+ * fields a consumer observes is held in state and refreshed on notification
+ * rather than read during render. Keeping the previous reference for an equal
+ * projection is what stops a notification carrying no observable change from
+ * invalidating everything derived from it.
+ */
+export const useModelContextSnapshot = <T>(
+  aui: AssistantClient,
+  enabled: boolean,
+  source: ModelContextSnapshotSource<T>,
+): T => {
+  const [snapshot, setSnapshot] = useState(() =>
+    enabled ? source.read(aui) : source.empty,
+  );
+
+  useEffect(() => {
+    if (!enabled) return undefined;
+    const read = () => {
+      const next = source.read(aui);
+      setSnapshot((previous) =>
+        source.isEqual?.(previous, next) ? previous : next,
+      );
+    };
+    read();
+    return source.subscribe(aui, read);
+  }, [aui, enabled, source]);
+
+  return snapshot;
+};
+
+/** Shallow own-key comparison, for a projection whose values are records. */
+export const shallowEqualRecords = (
+  previous: Readonly<Record<string, unknown>>,
+  next: Readonly<Record<string, unknown>>,
+): boolean => {
+  const keys = Object.keys(previous);
+  if (keys.length !== Object.keys(next).length) return false;
+  return keys.every(
+    (key) => Object.hasOwn(next, key) && previous[key] === next[key],
+  );
+};

package/src/unstable/useSlashCommandAdapter.test.tsx

@@ -0,0 +1,141 @@
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { startTransition, Suspense, type PropsWithChildren } from "react";
+import { describe, expect, it, vi } from "vitest";
+import {
+  unstable_useSlashCommandAdapter,
+  type Unstable_SlashCommand,
+} from "./useSlashCommandAdapter";
+
+const command = (
+  label: string,
+  execute: () => void,
+): Unstable_SlashCommand => ({
+  id: "run",
+  label,
+  execute,
+});
+
+describe("unstable_useSlashCommandAdapter", () => {
+  it("keeps commands scoped to committed renders", () => {
+    const executeA = vi.fn();
+    const executeB = vi.fn();
+    const commandsA = [command("Workspace A", executeA)];
+    const commandsB = [command("Workspace B", executeB)];
+    const renderedLabel = vi.fn();
+    const interruptedRender = vi.fn();
+    const pending = new Promise<never>(() => {});
+    let blocked = false;
+
+    const Blocker = () => {
+      if (blocked) {
+        interruptedRender();
+        throw pending;
+      }
+      return null;
+    };
+    const Wrapper = ({ children }: PropsWithChildren) => (
+      <Suspense fallback={null}>
+        {children}
+        <Blocker />
+      </Suspense>
+    );
+
+    const { result, rerender } = renderHook(
+      ({ commands }) => {
+        renderedLabel(commands[0]?.label);
+        return unstable_useSlashCommandAdapter({ commands });
+      },
+      {
+        initialProps: { commands: commandsA },
+        wrapper: Wrapper,
+      },
+    );
+
+    act(() => {
+      blocked = true;
+      startTransition(() => rerender({ commands: commandsB }));
+    });
+
+    expect(interruptedRender).toHaveBeenCalled();
+    expect(renderedLabel).toHaveBeenCalledWith("Workspace B");
+    expect(result.current.adapter.search?.("")).toEqual([
+      { id: "run", type: "command", label: "Workspace A" },
+    ]);
+
+    result.current.action.onExecute({
+      id: "run",
+      type: "command",
+      label: "Workspace A",
+    });
+    expect(executeA).toHaveBeenCalledOnce();
+    expect(executeB).not.toHaveBeenCalled();
+  });
+
+  it("updates search while keeping actions stable after commands commit", () => {
+    const executeA = vi.fn();
+    const executeB = vi.fn();
+    const commandsA = [command("Workspace A", executeA)];
+    const commandsB = [command("Workspace B", executeB)];
+    const { result, rerender } = renderHook(
+      ({ commands }) => unstable_useSlashCommandAdapter({ commands }),
+      { initialProps: { commands: commandsA } },
+    );
+    const initial = result.current;
+
+    rerender({ commands: commandsB });
+
+    expect(result.current.adapter).not.toBe(initial.adapter);
+    expect(result.current.action).toBe(initial.action);
+    expect(result.current.adapter.search?.("")).toEqual([
+      { id: "run", type: "command", label: "Workspace B" },
+    ]);
+
+    result.current.action.onExecute({
+      id: "run",
+      type: "command",
+      label: "Workspace B",
+    });
+    expect(executeA).not.toHaveBeenCalled();
+    expect(executeB).toHaveBeenCalledOnce();
+  });
+
+  it("matches the displayed label for a command with no explicit label", () => {
+    const { result } = renderHook(() =>
+      unstable_useSlashCommandAdapter({
+        commands: [{ id: "summarize", execute: vi.fn() }],
+      }),
+    );
+
+    expect(result.current.adapter.search?.("/")).toEqual([
+      { id: "summarize", type: "command", label: "/summarize" },
+    ]);
+    expect(result.current.adapter.search?.("sum")).toEqual([
+      { id: "summarize", type: "command", label: "/summarize" },
+    ]);
+    expect(result.current.adapter.search?.("nope")).toEqual([]);
+  });
+
+  it("keeps the adapter stable for equivalent inline commands", () => {
+    const executeA = vi.fn();
+    const executeB = vi.fn();
+    const { result, rerender } = renderHook(
+      ({ execute }) =>
+        unstable_useSlashCommandAdapter({
+          commands: [{ id: "run", label: "Run", execute }],
+        }),
+      { initialProps: { execute: executeA } },
+    );
+    const initialAdapter = result.current.adapter;
+
+    rerender({ execute: executeB });
+
+    expect(result.current.adapter).toBe(initialAdapter);
+    act(() =>
+      result.current.action.onExecute(result.current.adapter.search!("")[0]!),
+    );
+    expect(executeA).not.toHaveBeenCalled();
+    expect(executeB).toHaveBeenCalledOnce();
+  });
+});

package/src/unstable/useSlashCommandAdapter.ts

@@ -1,11 +1,12 @@
 "use client";
 
-import { useMemo, useRef } from "react";
+import { useLayoutEffect, useMemo, useRef } from "react";
 import type {
   Unstable_TriggerAdapter,
   Unstable_TriggerItem,
 } from "@assistant-ui/core";
 import type { Unstable_IconComponent } from "./useMentionAdapter";
+import { matchesTriggerItemQuery } from "../primitives/composer/trigger/matchesTriggerItemQuery";
 
 export type Unstable_SlashCommand = {
   readonly id: string;
@@ -61,34 +62,51 @@ export function unstable_useSlashCommandAdapter(
   const { commands, removeOnExecute } = options;
 
   const commandsRef = useRef(commands);
-  commandsRef.current = commands;
+  const committedItemsRef = useRef<readonly Unstable_TriggerItem[]>(undefined);
+  const nextItems = commands.map(toItem);
+  // A referential cache, never state: `items` is always structurally equal to
+  // `nextItems`, so the render-time read cannot change what the adapter shows.
+  const items = areTriggerItemsEqual(committedItemsRef.current, nextItems)
+    ? committedItemsRef.current
+    : nextItems;
 
-  return useMemo(() => {
-    const adapter: Unstable_TriggerAdapter = {
+  // The adapter's callbacks must only observe commands from committed renders.
+  useLayoutEffect(() => {
+    commandsRef.current = commands;
+    committedItemsRef.current = items;
+  }, [commands, items]);
+
+  const adapter = useMemo<Unstable_TriggerAdapter>(
+    () => ({
       categories: () => [],
       categoryItems: () => [],
       search: (query: string) => {
         const lower = query.toLowerCase();
-        return commandsRef.current
-          .filter((c) => matchesQuery(c, lower))
-          .map(toItem);
+        return items.filter((item) => matchesTriggerItemQuery(item, lower));
       },
-    };
+    }),
+    [items],
+  );
 
-    const action: Unstable_SlashCommandAction = {
+  const action = useMemo<Unstable_SlashCommandAction>(
+    () => ({
       onExecute: (item) => {
         commandsRef.current.find((c) => c.id === item.id)?.execute();
       },
       ...(removeOnExecute !== undefined ? { removeOnExecute } : {}),
-    };
+    }),
+    [removeOnExecute],
+  );
 
-    return {
+  return useMemo(
+    () => ({
       adapter,
       action,
       ...(options.iconMap ? { iconMap: options.iconMap } : {}),
       ...(options.fallbackIcon ? { fallbackIcon: options.fallbackIcon } : {}),
-    };
-  }, [removeOnExecute, options.iconMap, options.fallbackIcon]);
+    }),
+    [adapter, action, options.iconMap, options.fallbackIcon],
+  );
 }
 
 function toItem(cmd: Unstable_SlashCommand): Unstable_TriggerItem {
@@ -101,10 +119,20 @@ function toItem(cmd: Unstable_SlashCommand): Unstable_TriggerItem {
   };
 }
 
-function matchesQuery(cmd: Unstable_SlashCommand, lower: string): boolean {
-  if (!lower) return true;
-  if (cmd.id.toLowerCase().includes(lower)) return true;
-  if (cmd.label?.toLowerCase().includes(lower)) return true;
-  if (cmd.description?.toLowerCase().includes(lower)) return true;
-  return false;
+function areTriggerItemsEqual(
+  previous: readonly Unstable_TriggerItem[] | undefined,
+  next: readonly Unstable_TriggerItem[],
+): previous is readonly Unstable_TriggerItem[] {
+  if (!previous || previous.length !== next.length) return false;
+
+  return previous.every((item, index) => {
+    const nextItem = next[index];
+    return (
+      nextItem !== undefined &&
+      item.id === nextItem.id &&
+      item.label === nextItem.label &&
+      item.description === nextItem.description &&
+      item.metadata?.icon === nextItem.metadata?.icon
+    );
+  });
 }

package/src/unstable/webmcp/WebMcpRegistrationResource.ts

@@ -0,0 +1,92 @@
+import { useEffect, useRef, useState } from "react";
+import { resource } from "@assistant-ui/tap";
+import type { Tool } from "assistant-stream";
+import { toWebMcpTool } from "./convertTools";
+import type { WebMcpHost } from "./webmcp-host";
+
+export type WebMcpRegistrationProps = {
+  host: WebMcpHost;
+  name: string;
+  /** Re-registers when it changes; an execute-only edit reads through instead. */
+  signature: string;
+  tool: Tool<any, any>;
+};
+
+// The explainer specifies NotAllowedError for a page whose tools permission is
+// off, which a cross-origin iframe without allow="tools" reaches routinely. It
+// documents no rejection for a name already taken, so that stays a hedge.
+const notPermitted = (error: unknown) =>
+  (error as { name?: unknown } | null | undefined)?.name === "NotAllowedError";
+
+const notPermittedMessage = (name: string) =>
+  `[assistant-ui] WebMCP registration for tool "${name}" was not permitted; the page's tools permission is disabled.`;
+
+const useWebMcpRegistration = ({
+  host,
+  name,
+  signature,
+  tool,
+}: WebMcpRegistrationProps): string | null => {
+  // A refused name is remembered for as long as the tool stays in the model
+  // context, so a permanent collision warns once rather than on every sync.
+  const [refused, setRefused] = useState(false);
+  const toolRef = useRef(tool);
+  // Commit phase, so an abandoned render cannot hand the host a tool the tree
+  // never committed.
+  useEffect(() => {
+    toolRef.current = tool;
+  }, [tool]);
+
+  useEffect(() => {
+    if (refused) return undefined;
+
+    const lifecycle = new AbortController();
+    let live = true;
+    let dispose: (() => void) | undefined;
+
+    const refuse = (message: string, error: unknown) => {
+      if (!live) return;
+      live = false;
+      setRefused(true);
+      console.warn(message, error);
+    };
+
+    try {
+      dispose = host.registerTool(
+        toWebMcpTool(name, () => toolRef.current, lifecycle.signal),
+        (error) =>
+          refuse(
+            notPermitted(error)
+              ? notPermittedMessage(name)
+              : `[assistant-ui] WebMCP registration for tool "${name}" failed (name may already be registered).`,
+            error,
+          ),
+      );
+    } catch (error) {
+      refuse(
+        notPermitted(error)
+          ? notPermittedMessage(name)
+          : `[assistant-ui] Skipping WebMCP registration for tool "${name}": registerTool failed (name may already be registered).`,
+        error,
+      );
+      return undefined;
+    }
+
+    return () => {
+      live = false;
+      lifecycle.abort();
+      try {
+        dispose();
+      } catch (error) {
+        console.warn(
+          `[assistant-ui] Unregistering WebMCP tool "${name}" failed.`,
+          error,
+        );
+      }
+    };
+  }, [host, name, signature, refused]);
+
+  return refused ? null : name;
+};
+
+export const WebMcpRegistrationResource = resource(useWebMcpRegistration);