@assistant-ui/react 0.15.28 → 0.15.29

package/src/model-context/frame/useAssistantFrameHost.reload.test.tsx

@@ -0,0 +1,106 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import { ModelContextRegistry } from "@assistant-ui/core";
+import { expect, it, vi } from "vitest";
+import { useAssistantFrameHost } from "./useAssistantFrameHost";
+
+it("drops tools and pending calls when a frame loads a page without a provider", async () => {
+  vi.useFakeTimers();
+  const iframe = document.createElement("iframe");
+  document.body.append(iframe);
+  vi.spyOn(iframe.contentWindow!, "postMessage").mockImplementation(() => {});
+  const registry = new ModelContextRegistry();
+  const register = (host: Parameters<typeof registry.addProvider>[0]) =>
+    registry.addProvider(host).remove;
+  const { unmount } = renderHook(() =>
+    useAssistantFrameHost({ iframeRef: { current: iframe }, register }),
+  );
+
+  try {
+    iframe.dispatchEvent(new Event("load"));
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        source: iframe.contentWindow,
+        origin: window.location.origin,
+        data: {
+          channel: "assistant-ui-frame",
+          message: {
+            type: "model-context-update",
+            context: { tools: { search: { parameters: { type: "object" } } } },
+          },
+        },
+      }),
+    );
+    const search = registry.getModelContext().tools?.search;
+    expect(search).toBeDefined();
+    const outcome = Promise.resolve(search!.execute!({}, {} as any)).catch(
+      (error: unknown) => error,
+    );
+
+    iframe.dispatchEvent(new Event("load"));
+    expect(registry.getModelContext().tools?.search).toBeDefined();
+    vi.advanceTimersByTime(2000);
+    expect(registry.getModelContext().tools?.search).toBeUndefined();
+    expect(await outcome).toBeInstanceOf(Error);
+  } finally {
+    unmount();
+    iframe.remove();
+    vi.restoreAllMocks();
+    vi.useRealTimers();
+  }
+});
+
+it("clears a reloaded frame when the hook mounted after its initial load", () => {
+  vi.useFakeTimers();
+  const iframe = document.createElement("iframe");
+  document.body.append(iframe);
+  const postMessage = vi
+    .spyOn(iframe.contentWindow!, "postMessage")
+    .mockImplementation(() => {});
+  const registry = new ModelContextRegistry();
+  const register = (host: Parameters<typeof registry.addProvider>[0]) =>
+    registry.addProvider(host).remove;
+  const { unmount } = renderHook(() =>
+    useAssistantFrameHost({ iframeRef: { current: iframe }, register }),
+  );
+
+  try {
+    window.dispatchEvent(
+      new MessageEvent("message", {
+        source: iframe.contentWindow,
+        origin: window.location.origin,
+        data: {
+          channel: "assistant-ui-frame",
+          message: {
+            type: "model-context-update",
+            documentId: "first",
+            context: { tools: { search: { parameters: { type: "object" } } } },
+          },
+        },
+      }),
+    );
+    expect(registry.getModelContext().tools?.search).toBeDefined();
+    postMessage.mockClear();
+
+    iframe.dispatchEvent(new Event("load"));
+    expect(registry.getModelContext().tools?.search).toBeDefined();
+    vi.advanceTimersByTime(2000);
+    expect(registry.getModelContext().tools?.search).toBeUndefined();
+    expect(postMessage).toHaveBeenCalledWith(
+      {
+        channel: "assistant-ui-frame",
+        message: {
+          type: "model-context-request",
+          requestId: expect.any(String),
+        },
+      },
+      window.location.origin,
+    );
+  } finally {
+    unmount();
+    iframe.remove();
+    vi.restoreAllMocks();
+    vi.useRealTimers();
+  }
+});

package/src/model-context/frame/useAssistantFrameHost.test.tsx

@@ -26,7 +26,9 @@ describe("useAssistantFrameHost", () => {
     const iframeRef = {
       current: {
         contentWindow: { postMessage: vi.fn() } as unknown as Window,
-      } as HTMLIFrameElement,
+        addEventListener: vi.fn(),
+        removeEventListener: vi.fn(),
+      } as unknown as HTMLIFrameElement,
     };
     const { unmount } = renderHook(() =>
       useAssistantFrameHost({ iframeRef, register }),

package/src/model-context/frame/useAssistantFrameHost.ts

@@ -33,39 +33,87 @@ export const useAssistantFrameHost = ({
   register,
 }: UseAssistantFrameHostOptions): void => {
   useEffect(() => {
-    const iframeWindow = iframeRef.current?.contentWindow;
-    if (!iframeWindow) return;
-
-    const frameHost = new AssistantFrameHost(iframeWindow, targetOrigin);
-
-    const unsubscribe = register(frameHost);
+    let active:
+      | {
+          iframe: HTMLIFrameElement;
+          window: Window;
+          dispose: () => void;
+        }
+      | undefined;
 
-    return () => {
-      let cleanupFailed = false;
-      let cleanupError: unknown;
+    const syncFrame = () => {
+      const iframe = iframeRef.current;
+      const iframeWindow =
+        iframe?.isConnected === false ? null : iframe?.contentWindow;
+      if (active && active.iframe === iframe && active.window === iframeWindow)
+        return;
 
+      const previous = active;
+      active = undefined;
+      let disposalFailed = false;
+      let disposalError: unknown;
       try {
-        frameHost.dispose();
+        previous?.dispose();
       } catch (error) {
-        cleanupFailed = true;
-        cleanupError = error;
+        disposalFailed = true;
+        disposalError = error;
       }
 
-      try {
-        unsubscribe();
-      } catch (error) {
-        if (cleanupFailed) {
-          console.error(
-            "[assistant-ui] AssistantFrameHost unregistration failed.",
-            error,
-          );
-        } else {
-          cleanupFailed = true;
-          cleanupError = error;
+      if (iframe && iframeWindow) {
+        const frameHost = new AssistantFrameHost(iframeWindow, targetOrigin);
+        let unsubscribe: Unsubscribe;
+        try {
+          unsubscribe = register(frameHost);
+        } catch (error) {
+          frameHost.dispose();
+          throw error;
         }
+        const onLoad = () => frameHost.frameLoaded();
+        iframe.addEventListener("load", onLoad);
+        active = {
+          iframe,
+          window: iframeWindow,
+          dispose: () => {
+            iframe.removeEventListener("load", onLoad);
+            let cleanupFailed = false;
+            let cleanupError: unknown;
+
+            try {
+              frameHost.dispose();
+            } catch (error) {
+              cleanupFailed = true;
+              cleanupError = error;
+            }
+
+            try {
+              unsubscribe();
+            } catch (error) {
+              if (cleanupFailed) {
+                console.error(
+                  "[assistant-ui] AssistantFrameHost unregistration failed.",
+                  error,
+                );
+              } else {
+                cleanupFailed = true;
+                cleanupError = error;
+              }
+            }
+
+            if (cleanupFailed) throw cleanupError;
+          },
+        };
       }
 
-      if (cleanupFailed) throw cleanupError;
+      if (disposalFailed) throw disposalError;
+    };
+
+    syncFrame();
+    const observer = new MutationObserver(syncFrame);
+    observer.observe(document, { childList: true, subtree: true });
+
+    return () => {
+      observer.disconnect();
+      active?.dispose();
     };
   }, [iframeRef, targetOrigin, register]);
 };

package/src/primitives/composer/ComposerAddAttachment.test.tsx

@@ -30,6 +30,7 @@ vi.mock("@assistant-ui/store", async (importOriginal) => {
   return {
     ...actual,
     useAui: () => ({
+      optional: {},
       composer: {
         getState: () => ({ attachmentAccept: "*" }),
       },

package/src/primitives/composer/ComposerAddAttachment.thread-scope.test.tsx

@@ -0,0 +1,202 @@
+// @vitest-environment jsdom
+import { act, createElement, useState } from "react";
+import { render, waitFor, cleanup } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { AuiConfig, AuiProvider, Derived, useAui } from "@assistant-ui/store";
+import { resource, useResource } from "@assistant-ui/tap";
+import { InMemoryThreadList } from "@assistant-ui/core/store";
+import {
+  AssistantRuntimeProvider,
+  ExternalThread,
+  SimpleTextAttachmentAdapter,
+  Tools,
+  useLocalRuntime,
+  type AssistantRuntime,
+  type ChatModelAdapter,
+} from "../../index";
+import { ComposerPrimitiveAddAttachment } from "./ComposerAddAttachment";
+
+const attachmentAdapter = new SimpleTextAttachmentAdapter();
+const model: ChatModelAdapter = { async *run() {} };
+
+afterEach(() => {
+  cleanup();
+  document
+    .querySelectorAll('input[type="file"]')
+    .forEach((input) => input.remove());
+  vi.restoreAllMocks();
+});
+
+const StandaloneThread = resource(function useStandaloneThread() {
+  return useResource(ExternalThread({ messages: [], attachmentAdapter }));
+});
+
+const setup = async (mode: "store" | "runtime" | "standalone") => {
+  vi.spyOn(HTMLInputElement.prototype, "click").mockImplementation(() => {});
+  let captured!: ReturnType<typeof useAui>;
+  let runtime: AssistantRuntime | undefined;
+  let changeScope!: () => void;
+  const Capture = () => {
+    captured = useAui();
+    return (
+      <ComposerPrimitiveAddAttachment>Add file</ComposerPrimitiveAddAttachment>
+    );
+  };
+  const Scope = () => {
+    const aui = useAui();
+    const [withTools, setWithTools] = useState(false);
+    changeScope = () => setWithTools(true);
+    const config = AuiConfig(withTools ? { tools: Tools({}) } : {});
+    return (
+      <AuiProvider extends={aui} config={config}>
+        <Capture />
+      </AuiProvider>
+    );
+  };
+  const StoreApp = () => {
+    const aui = useAui({
+      threads: InMemoryThreadList({
+        thread: () => ExternalThread({ messages: [], attachmentAdapter }),
+      }),
+    });
+    return (
+      <AuiProvider value={aui}>
+        <Scope />
+      </AuiProvider>
+    );
+  };
+  const RuntimeApp = () => {
+    runtime = useLocalRuntime(model, {
+      adapters: { attachments: attachmentAdapter },
+    });
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <Scope />
+      </AssistantRuntimeProvider>
+    );
+  };
+  const StandaloneApp = () => {
+    const config = AuiConfig({
+      thread: StandaloneThread(),
+      composer: Derived({
+        source: "thread",
+        query: {},
+        get: (aui) => aui.thread.composer(),
+      }),
+    });
+    return (
+      <AuiProvider config={config}>
+        <Capture />
+      </AuiProvider>
+    );
+  };
+  const App =
+    mode === "store"
+      ? StoreApp
+      : mode === "runtime"
+        ? RuntimeApp
+        : StandaloneApp;
+  const view = render(createElement(App));
+  if (runtime)
+    await act(async () => {
+      await runtime!.threads.mainItem.initialize();
+    });
+  await act(async () => view.getByRole("button").click());
+  const input = document.querySelector<HTMLInputElement>('input[type="file"]')!;
+  const select = async () => {
+    Object.defineProperty(input, "files", {
+      configurable: true,
+      value: [
+        new File(["private draft"], "thread-a.txt", { type: "text/plain" }),
+      ],
+    });
+    const event = new Event("change");
+    Object.defineProperty(event, "target", { value: input });
+    await act(async () => {
+      await input.onchange?.call(input, event);
+    });
+  };
+  return {
+    aui: () => captured,
+    select,
+    changeScope,
+    unmount: view.unmount,
+    rerender: () => view.rerender(createElement(App)),
+  };
+};
+
+describe.each(["store", "runtime"] as const)(
+  "ComposerPrimitive.AddAttachment thread ownership in %s",
+  (mode) => {
+    it("adds a selected file to the thread that opened the picker", async () => {
+      const { aui, select } = await setup(mode);
+      await select();
+      expect(
+        aui()
+          .composer.getState()
+          .attachments.map((item) => item.name),
+      ).toEqual(["thread-a.txt"]);
+    });
+
+    it("does not add an old picker's selection to a different thread", async () => {
+      const { aui, select } = await setup(mode);
+      const originId = aui().threads.getState().mainThreadId;
+      await act(async () => {
+        aui().threads.switchToNewThread();
+      });
+      await waitFor(() =>
+        expect(aui().threads.getState().mainThreadId).not.toBe(originId),
+      );
+      expect(aui().composer.getState().attachments).toHaveLength(0);
+      await select();
+      expect(aui().composer.getState().attachments).toHaveLength(0);
+    });
+
+    it("keeps the picker usable when the same draft changes", async () => {
+      const { aui, select } = await setup(mode);
+      await act(async () => {
+        aui().composer.setText("same conversation");
+      });
+      await select();
+      expect(aui().composer.getState().text).toBe("same conversation");
+      expect(
+        aui()
+          .composer.getState()
+          .attachments.map((item) => item.name),
+      ).toEqual(["thread-a.txt"]);
+    });
+
+    it("keeps the picker usable when an unrelated scope changes", async () => {
+      const { aui, select, changeScope } = await setup(mode);
+      const original = aui();
+      await act(async () => changeScope());
+      expect(Object.is(aui(), original)).toBe(false);
+      expect(Object.is(aui().composer, original.composer)).toBe(true);
+      await select();
+      expect(
+        aui()
+          .composer.getState()
+          .attachments.map((item) => item.name),
+      ).toEqual(["thread-a.txt"]);
+    });
+
+    it("settles a picker result after its provider unmounts", async () => {
+      const { select, unmount } = await setup(mode);
+      unmount();
+      await expect(select()).resolves.toBeUndefined();
+      expect(document.querySelector('input[type="file"]')).toBeNull();
+    });
+  },
+);
+
+it("keeps a standalone composer picker usable after a parent re-render", async () => {
+  const { aui, select, rerender } = await setup("standalone");
+  expect(aui().optional.threadListItem).toBeUndefined();
+  rerender();
+  await select();
+  expect(
+    aui()
+      .composer.getState()
+      .attachments.map((item) => item.name),
+  ).toEqual(["thread-a.txt"]);
+});

package/src/primitives/composer/ComposerAddAttachment.ts

@@ -5,7 +5,7 @@ import {
   type ActionButtonProps,
   createActionButton,
 } from "../../utils/createActionButton";
-import { useCallback } from "react";
+import { useCallback, useInsertionEffect, useRef } from "react";
 import { useAui } from "@assistant-ui/store";
 import { useComposerAddAttachment as useComposerAddAttachmentBehavior } from "@assistant-ui/core/react";
 
@@ -17,8 +17,17 @@ const useComposerAddAttachment = ({
 } = {}) => {
   const { disabled, addAttachment } = useComposerAddAttachmentBehavior();
   const aui = useAui();
+  const scopeGenerationRef = useRef(0);
+
+  useInsertionEffect(
+    () => () => {
+      scopeGenerationRef.current += 1;
+    },
+    [aui.composer, aui.optional.threadListItem],
+  );
 
   const callback = useCallback(() => {
+    const scopeGeneration = scopeGenerationRef.current;
     const input = document.createElement("input");
     input.type = "file";
     input.multiple = multiple;
@@ -37,6 +46,7 @@ const useComposerAddAttachment = ({
 
       const attachmentPromises = Array.from(fileList, async (file) => {
         try {
+          if (scopeGenerationRef.current !== scopeGeneration) return;
           await addAttachment(file);
         } catch {
           // The composer runtime emits composer.attachmentAddError before rejecting.

package/src/primitives/composer/ComposerInput.focus-on-run-start.test.tsx

@@ -0,0 +1,130 @@
+// @vitest-environment jsdom
+
+import { act, fireEvent, render, screen } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { AssistantRuntimeProvider } from "../../context";
+import { useLocalRuntime } from "../../legacy-runtime/runtime-cores/local/useLocalRuntime";
+import { ThreadPrimitiveRoot } from "../thread/ThreadRoot";
+import { ComposerPrimitiveInput } from "./ComposerInput";
+import { ComposerPrimitiveRoot } from "./ComposerRoot";
+
+const adapter = { async *run() {} };
+
+type Runtime = ReturnType<typeof useLocalRuntime>;
+
+const setup = () => {
+  let runtime: Runtime | null = null;
+  const App = () => {
+    const currentRuntime = useLocalRuntime(adapter);
+    runtime = currentRuntime;
+    return (
+      <AssistantRuntimeProvider runtime={currentRuntime}>
+        <input aria-label="Outside thread" />
+        <ThreadPrimitiveRoot>
+          <button type="button">Inside thread</button>
+          <ComposerPrimitiveInput aria-label="Composer" autoFocus />
+          <ThreadPrimitiveRoot>
+            <button type="button">Nested thread</button>
+          </ThreadPrimitiveRoot>
+        </ThreadPrimitiveRoot>
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  render(<App />);
+  const composer = screen.getByRole("textbox", {
+    name: "Composer",
+  }) as HTMLTextAreaElement;
+  const startRun = async () => {
+    await act(async () => {
+      await runtime!.thread.append("Programmatic message");
+    });
+  };
+  return { composer, startRun };
+};
+
+describe("ComposerPrimitiveInput run-start focus", () => {
+  it("focuses the composer when its button is focused without a thread root", async () => {
+    let runtime: Runtime | null = null;
+    const App = () => {
+      const currentRuntime = useLocalRuntime(adapter);
+      runtime = currentRuntime;
+      return (
+        <AssistantRuntimeProvider runtime={currentRuntime}>
+          <ComposerPrimitiveRoot>
+            <ComposerPrimitiveInput aria-label="Floating composer" autoFocus />
+            <button type="button">Send</button>
+          </ComposerPrimitiveRoot>
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    render(<App />);
+    const composer = screen.getByRole("textbox", {
+      name: "Floating composer",
+    });
+    const send = screen.getByRole("button", { name: "Send" });
+    act(() => send.focus());
+
+    await act(async () => {
+      await runtime!.thread.append("Programmatic message");
+    });
+
+    expect(document.activeElement).toBe(composer);
+  });
+
+  it("keeps focus in an input outside the thread", async () => {
+    const { startRun } = setup();
+    const outside = screen.getByRole("textbox", { name: "Outside thread" });
+    act(() => outside.focus());
+
+    await startRun();
+
+    expect(document.activeElement).toBe(outside);
+  });
+
+  it("focuses the composer when nothing is focused", async () => {
+    const { composer, startRun } = setup();
+    act(() => composer.blur());
+    expect(document.activeElement).toBe(document.body);
+
+    await startRun();
+
+    expect(document.activeElement).toBe(composer);
+  });
+
+  it("focuses the composer when another element in its thread is focused", async () => {
+    const { composer, startRun } = setup();
+    const inside = screen.getByRole("button", { name: "Inside thread" });
+    act(() => inside.focus());
+
+    await startRun();
+
+    expect(document.activeElement).toBe(composer);
+  });
+
+  it("keeps focus inside a nested thread root", async () => {
+    const { startRun } = setup();
+    const nested = screen.getByRole("button", { name: "Nested thread" });
+    act(() => nested.focus());
+
+    await startRun();
+
+    expect(document.activeElement).toBe(nested);
+  });
+
+  it("preserves the focused composer's selection", async () => {
+    const { composer, startRun } = setup();
+    fireEvent.change(composer, { target: { value: "hello world" } });
+    act(() => {
+      composer.focus();
+      composer.setSelectionRange(2, 5);
+    });
+
+    await startRun();
+
+    expect(document.activeElement).toBe(composer);
+    expect(composer.selectionStart).toBe(2);
+    expect(composer.selectionEnd).toBe(5);
+  });
+});

package/src/primitives/composer/ComposerInput.tsx

@@ -28,6 +28,7 @@ import { flushTapSync } from "@assistant-ui/tap";
 import { useComposerInputPluginRegistryOptional } from "./ComposerInputPluginContext";
 import { useComposerCompactContextOptional } from "./ComposerCompactContext";
 import { useComposerCancelWithFocus } from "./useComposerCancelWithFocus";
+import { useThreadRootElementRef } from "../thread/ThreadRootElementContext";
 import {
   useComposerInputDisabled,
   useComposerInputValue,
@@ -57,7 +58,10 @@ export namespace ComposerPrimitiveInput {
      */
     cancelOnEscape?: boolean | undefined;
     /**
-     * Whether to automatically focus the input when a new run starts.
+     * Whether to focus the input when a new run starts if nothing is focused
+     * or focus is already inside the same thread or composer. Keeps the current
+     * selection when the input is already focused. Applies when `autoFocus` is
+     * enabled.
      *
      * @deprecated Experimental since 2024-10-12. Not scheduled for removal; the API may change in any release.
      *
@@ -189,6 +193,7 @@ export const ComposerPrimitiveInput = forwardRef<
     forwardedRef,
   ) => {
     const aui = useAui();
+    const threadRootRef = useThreadRootElementRef();
     const pluginRegistry = useComposerInputPluginRegistryOptional();
     const compactContext = useComposerCompactContextOptional();
 
@@ -343,6 +348,27 @@ export const ComposerPrimitiveInput = forwardRef<
       textarea.setSelectionRange(textarea.value.length, textarea.value.length);
     }, [autoFocusEnabled]);
 
+    const focusOnRunStart = useCallback(() => {
+      const textarea = textareaRef.current;
+      if (!textarea) return;
+      const doc = textarea.ownerDocument;
+      const active = doc.activeElement;
+      if (active === textarea) return;
+
+      const root = threadRootRef?.current;
+      const composerRoot = textarea.closest("[data-aui-composer-type]");
+      if (
+        active &&
+        active !== doc.body &&
+        !composerRoot?.contains(active) &&
+        (!root?.contains(active) ||
+          active.closest("[data-aui-thread-root]") !== root)
+      )
+        return;
+
+      focus();
+    }, [focus, threadRootRef]);
+
     useEffect(() => focus(), [focus]);
 
     useOnScrollToBottom(() => {
@@ -361,8 +387,8 @@ export const ComposerPrimitiveInput = forwardRef<
       )
         return undefined;
 
-      return aui.on("thread.runStart", focus);
-    }, [unstable_focusOnRunStart, focus, aui]);
+      return aui.on("thread.runStart", focusOnRunStart);
+    }, [unstable_focusOnRunStart, focusOnRunStart, aui]);
 
     useEffect(() => {
       if (