@assistant-ui/react 0.15.28 → 0.15.30

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

@@ -0,0 +1,145 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import { AssistantFrameHost, ModelContextRegistry } from "@assistant-ui/core";
+import { expect, it, vi } from "vitest";
+import { useAssistantFrameHost } from "./useAssistantFrameHost";
+
+const attachFrame = () => {
+  const iframe = document.createElement("iframe");
+  document.body.append(iframe);
+  vi.spyOn(iframe.contentWindow!, "postMessage").mockImplementation(() => {});
+  return iframe;
+};
+
+const publishTool = (iframe: HTMLIFrameElement, name: string) => {
+  window.dispatchEvent(
+    new MessageEvent("message", {
+      source: iframe.contentWindow,
+      origin: window.location.origin,
+      data: {
+        channel: "assistant-ui-frame",
+        message: {
+          type: "model-context-update",
+          context: {
+            tools: {
+              [name]: { parameters: { type: "object", properties: {} } },
+            },
+          },
+        },
+      },
+    }),
+  );
+};
+
+it("withdraws a detached frame's tools and calls, then serves an attached frame", async () => {
+  const registry = new ModelContextRegistry();
+  const firstFrame = attachFrame();
+  const iframeRef = { current: firstFrame };
+  const register = (host: Parameters<typeof registry.addProvider>[0]) =>
+    registry.addProvider(host).remove;
+  const { unmount } = renderHook(() =>
+    useAssistantFrameHost({ iframeRef, register }),
+  );
+
+  try {
+    publishTool(firstFrame, "search");
+    const search = registry.getModelContext().tools?.search;
+    expect(search).toBeDefined();
+    const pending = search!.execute!({}, {} as any);
+    const outcome = pending.then(
+      () => undefined,
+      (error: unknown) => error,
+    );
+
+    firstFrame.remove();
+    await vi.waitFor(() => {
+      expect(registry.getModelContext().tools).toBeUndefined();
+    });
+    expect(await outcome).toMatchObject({
+      message: "AssistantFrameHost has been disposed",
+    });
+
+    const secondFrame = attachFrame();
+    iframeRef.current = secondFrame;
+    await vi.waitFor(() => {
+      publishTool(secondFrame, "replacement");
+      expect(registry.getModelContext().tools?.replacement).toBeDefined();
+    });
+    secondFrame.remove();
+  } finally {
+    unmount();
+    firstFrame.remove();
+  }
+});
+
+it("serves a replacement frame before reporting a disposal failure", () => {
+  const disposalError = new Error("tool cancellation failed");
+  const dispose = vi
+    .spyOn(AssistantFrameHost.prototype, "dispose")
+    .mockImplementationOnce(() => {
+      throw disposalError;
+    });
+  const mutationCallbacks: MutationCallback[] = [];
+  vi.stubGlobal(
+    "MutationObserver",
+    class {
+      constructor(callback: MutationCallback) {
+        mutationCallbacks.push(callback);
+      }
+      observe() {}
+      disconnect() {}
+    },
+  );
+
+  const registry = new ModelContextRegistry();
+  const firstFrame = attachFrame();
+  const secondFrame = attachFrame();
+  const iframeRef = { current: firstFrame };
+  const register = (host: Parameters<typeof registry.addProvider>[0]) =>
+    registry.addProvider(host).remove;
+  const { unmount } = renderHook(() =>
+    useAssistantFrameHost({ iframeRef, register }),
+  );
+
+  try {
+    publishTool(firstFrame, "search");
+    expect(registry.getModelContext().tools?.search).toBeDefined();
+
+    iframeRef.current = secondFrame;
+    expect(() => mutationCallbacks[0]!([], {} as MutationObserver)).toThrow(
+      disposalError,
+    );
+
+    expect(registry.getModelContext().tools?.search).toBeUndefined();
+    publishTool(secondFrame, "replacement");
+    expect(registry.getModelContext().tools?.replacement).toBeDefined();
+  } finally {
+    unmount();
+    firstFrame.remove();
+    secondFrame.remove();
+    dispose.mockRestore();
+    vi.unstubAllGlobals();
+  }
+});
+
+it("disposes a host whose registration throws", () => {
+  const dispose = vi.spyOn(AssistantFrameHost.prototype, "dispose");
+  const iframe = attachFrame();
+  const register = vi.fn((): (() => void) => {
+    throw new Error("register failed");
+  });
+
+  try {
+    expect(() =>
+      renderHook(() =>
+        useAssistantFrameHost({ iframeRef: { current: iframe }, register }),
+      ),
+    ).toThrow("register failed");
+    expect(register).toHaveBeenCalledOnce();
+    expect(dispose).toHaveBeenCalledOnce();
+  } finally {
+    dispose.mockRestore();
+    iframe.remove();
+  }
+});

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/model-context/makeAssistantVisible.test.tsx

@@ -136,6 +136,21 @@ const runEditTool = async (
   await task;
 };
 
+const expectDefaultSettleDelay = async (start: () => Promise<unknown>) => {
+  await act(async () => {
+    const settled = vi.fn();
+    const task = start();
+    void task.then(settled);
+
+    await vi.advanceTimersByTimeAsync(1999);
+    expect(settled).not.toHaveBeenCalled();
+
+    await vi.advanceTimersByTimeAsync(1);
+    expect(settled).toHaveBeenCalledOnce();
+    await task;
+  });
+};
+
 describe("makeAssistantVisible", () => {
   beforeEach(() => {
     registerMock.mockClear();
@@ -191,21 +206,14 @@ describe("makeAssistantVisible", () => {
     });
     const input = render(<EditableInput />).getByRole("textbox");
     const editId = (input as HTMLInputElement).dataset.editId!;
-    const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout");
 
     expect(
       (input as HTMLInputElement).dataset.actionSettleDelay,
     ).toBeUndefined();
 
-    const task = getRegisteredTools().edit!.execute({
-      editId,
-      value: "updated",
-    });
-
-    expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 2000);
-    await vi.runAllTimersAsync();
-    await task;
-    setTimeoutSpy.mockRestore();
+    await expectDefaultSettleDelay(() =>
+      getRegisteredTools().edit!.execute({ editId, value: "updated" }),
+    );
   });
 
   it("uses the default when the rendered settle delay is invalid", async () => {
@@ -218,14 +226,9 @@ describe("makeAssistantVisible", () => {
     );
     const clickId = (button as HTMLButtonElement).dataset.clickId!;
     (button as HTMLButtonElement).dataset.actionSettleDelay = "invalid";
-    const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout");
-
-    const task = getRegisteredTools().click!.execute({ clickId });
-
-    expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 2000);
-    await vi.runAllTimersAsync();
-    await task;
-    setTimeoutSpy.mockRestore();
+    await expectDefaultSettleDelay(() =>
+      getRegisteredTools().click!.execute({ clickId }),
+    );
   });
 
   it.each(["", "   "])(
@@ -239,29 +242,14 @@ describe("makeAssistantVisible", () => {
       const input = render(<VisibleInput />).getByRole("textbox");
       const clickId = (input as HTMLInputElement).dataset.clickId!;
       const editId = (input as HTMLInputElement).dataset.editId!;
-      const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout");
       input.dataset.actionSettleDelay = renderedDelay;
 
-      const clickTask = getRegisteredTools().click!.execute({ clickId });
-      expect(setTimeoutSpy).toHaveBeenLastCalledWith(
-        expect.any(Function),
-        2000,
+      await expectDefaultSettleDelay(() =>
+        getRegisteredTools().click!.execute({ clickId }),
       );
-      await vi.runAllTimersAsync();
-      await clickTask;
-
-      setTimeoutSpy.mockClear();
-      const editTask = getRegisteredTools().edit!.execute({
-        editId,
-        value: "updated",
-      });
-      expect(setTimeoutSpy).toHaveBeenLastCalledWith(
-        expect.any(Function),
-        2000,
+      await expectDefaultSettleDelay(() =>
+        getRegisteredTools().edit!.execute({ editId, value: "updated" }),
       );
-      await vi.runAllTimersAsync();
-      await editTask;
-      setTimeoutSpy.mockRestore();
     },
   );
 

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"]);
+});