@assistant-ui/core 0.2.21 → 0.2.23

package/src/model-context/frame/provider.test.ts

@@ -0,0 +1,83 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { AssistantFrameProvider } from "./provider";
+import { FRAME_MESSAGE_CHANNEL } from "./types";
+
+describe("AssistantFrameProvider", () => {
+  let messageHandler: ((event: MessageEvent) => void) | undefined;
+  let parentWindow: Window;
+
+  beforeEach(() => {
+    parentWindow = {
+      postMessage: vi.fn(),
+    } as unknown as Window;
+
+    Object.defineProperty(window, "parent", {
+      value: parentWindow,
+      configurable: true,
+    });
+
+    vi.spyOn(window, "addEventListener").mockImplementation(
+      (event, listener) => {
+        if (event === "message" && typeof listener === "function") {
+          messageHandler = listener as (event: MessageEvent) => void;
+        }
+      },
+    );
+    vi.spyOn(window, "removeEventListener").mockImplementation(() => {});
+  });
+
+  afterEach(() => {
+    AssistantFrameProvider.dispose();
+    vi.restoreAllMocks();
+  });
+
+  it("only accepts tool calls from the parent window", async () => {
+    const execute = vi.fn(async () => "result");
+    AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({
+          tools: {
+            sensitiveTool: { execute },
+          },
+        }),
+      },
+      "https://parent.example",
+    );
+
+    const toolCall = {
+      channel: FRAME_MESSAGE_CHANNEL,
+      message: {
+        type: "tool-call",
+        id: "tool-call-1",
+        toolName: "sensitiveTool",
+        args: {},
+      },
+    };
+    const otherWindow = {
+      postMessage: vi.fn(),
+    } as unknown as Window;
+
+    messageHandler?.(
+      new MessageEvent("message", {
+        data: toolCall,
+        origin: "https://parent.example",
+        source: otherWindow,
+      }),
+    );
+
+    expect(execute).not.toHaveBeenCalled();
+
+    messageHandler?.(
+      new MessageEvent("message", {
+        data: toolCall,
+        origin: "https://parent.example",
+        source: parentWindow,
+      }),
+    );
+
+    await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
+  });
+});

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

@@ -59,6 +59,7 @@ export class AssistantFrameProvider {
   private handleMessage(event: MessageEvent) {
     if (this._targetOrigin !== "*" && event.origin !== this._targetOrigin)
       return;
+    if (event.source !== window.parent) return;
     if (event.data?.channel !== FRAME_MESSAGE_CHANNEL) return;
 
     const message = event.data.message as FrameMessage;

package/src/model-context/types.ts

@@ -55,6 +55,12 @@ export type AssistantContextConfig = {
   disabled?: boolean | undefined;
 };
 
+const stripOverwrite = (tool: Tool<any, any>): Tool<any, any> => {
+  if (!tool.overwrite) return tool;
+  const { overwrite: _, ...rest } = tool;
+  return rest as Tool<any, any>;
+};
+
 export const mergeModelContexts = (
   configSet: Set<ModelContextProvider>,
 ): ModelContext => {
@@ -79,25 +85,29 @@ export const mergeModelContexts = (
         if (existing && existing !== tool) {
           const existingPriority = toolPriorities[name]!;
           if (existingPriority === priority) {
-            throw new Error(
-              `You tried to define a tool with the name ${name}, but it already exists.`,
-            );
+            if (!tool.overwrite) {
+              throw new Error(
+                `You tried to define a tool with the name ${name}, but it already exists.`,
+              );
+            }
+            acc.tools![name] = stripOverwrite(tool);
+            continue;
           }
 
           const higherPriorityTool =
             existingPriority > priority ? existing : tool;
           const lowerPriorityTool =
             existingPriority > priority ? tool : existing;
-          acc.tools![name] = {
+          acc.tools![name] = stripOverwrite({
             ...lowerPriorityTool,
             ...higherPriorityTool,
-          } as Tool<any, any>;
+          } as Tool<any, any>);
           toolPriorities[name] = Math.max(existingPriority, priority);
           continue;
         }
 
         if (!acc.tools) acc.tools = {};
-        acc.tools[name] = tool;
+        acc.tools[name] = stripOverwrite(tool);
         toolPriorities[name] ??= priority;
       }
     }

package/src/react/adapters/LocalStorageThreadListAdapter.test.ts

@@ -207,6 +207,88 @@ describe("createLocalStorageAdapter", () => {
     ]);
   });
 
+  it("preserves concurrent metadata mutations across adapters", async () => {
+    const threadsKey = "@assistant-ui:threads";
+    const values = new Map<string, string>();
+    let metadataReads = 0;
+    let metadataWrites = 0;
+    let markFirstWriteStarted!: () => void;
+    let releaseFirstWrite!: () => void;
+    const firstWriteStarted = new Promise<void>((resolve) => {
+      markFirstWriteStarted = resolve;
+    });
+    const firstWriteCanFinish = new Promise<void>((resolve) => {
+      releaseFirstWrite = resolve;
+    });
+    const storage: AsyncStorageLike = {
+      getItem: async (key) => {
+        if (key === threadsKey) metadataReads += 1;
+        return values.get(key) ?? null;
+      },
+      setItem: async (key, value) => {
+        if (key === threadsKey) {
+          metadataWrites += 1;
+          if (metadataWrites === 1) {
+            markFirstWriteStarted();
+            await firstWriteCanFinish;
+          }
+        }
+        values.set(key, value);
+      },
+      removeItem: async (key) => {
+        values.delete(key);
+      },
+    };
+    const firstAdapter = createLocalStorageAdapter({ storage });
+    const secondAdapter = createLocalStorageAdapter({ storage });
+
+    const firstInitialization = firstAdapter.initialize("thread-1");
+    await firstWriteStarted;
+    const secondInitialization = secondAdapter.initialize("thread-2");
+    const readsWhileFirstWritePending = metadataReads;
+
+    releaseFirstWrite();
+    await Promise.all([firstInitialization, secondInitialization]);
+
+    expect(readsWhileFirstWritePending).toBe(1);
+    expect(JSON.parse(values.get(threadsKey) ?? "")).toEqual([
+      { remoteId: "thread-2", status: "regular" },
+      { remoteId: "thread-1", status: "regular" },
+    ]);
+  });
+
+  it("continues processing mutations after a storage failure", async () => {
+    const threadsKey = "@assistant-ui:threads";
+    const values = new Map<string, string>();
+    let shouldFail = true;
+    const storage: AsyncStorageLike = {
+      getItem: async (key) => values.get(key) ?? null,
+      setItem: async (key, value) => {
+        if (shouldFail) {
+          shouldFail = false;
+          throw new Error("Storage unavailable");
+        }
+        values.set(key, value);
+      },
+      removeItem: async (key) => {
+        values.delete(key);
+      },
+    };
+    const adapter = createLocalStorageAdapter({ storage });
+
+    await expect(adapter.initialize("thread-1")).rejects.toThrow(
+      "Storage unavailable",
+    );
+    await expect(adapter.initialize("thread-2")).resolves.toEqual({
+      remoteId: "thread-2",
+      externalId: undefined,
+    });
+
+    expect(JSON.parse(values.get(threadsKey) ?? "")).toEqual([
+      { remoteId: "thread-2", status: "regular" },
+    ]);
+  });
+
   it("includes the thread id when a stored thread cannot be fetched", async () => {
     const storage = createStorage({
       "@assistant-ui:threads": JSON.stringify([

package/src/react/adapters/LocalStorageThreadListAdapter.tsx

@@ -23,6 +23,37 @@ export type AsyncStorageLike = {
   removeItem(key: string): Promise<void>;
 };
 
+class KeyedMutationQueue {
+  private readonly tails = new Map<string, Promise<void>>();
+
+  run<T>(key: string, mutation: () => Promise<T>): Promise<T> {
+    const previous = this.tails.get(key);
+    const result = previous ? previous.then(mutation) : mutation();
+    const tail = result.then(
+      () => undefined,
+      () => undefined,
+    );
+
+    this.tails.set(key, tail);
+    void tail.then(() => {
+      if (this.tails.get(key) === tail) this.tails.delete(key);
+    });
+
+    return result;
+  }
+}
+
+const mutationQueues = new WeakMap<AsyncStorageLike, KeyedMutationQueue>();
+
+const getMutationQueue = (storage: AsyncStorageLike): KeyedMutationQueue => {
+  let queue = mutationQueues.get(storage);
+  if (!queue) {
+    queue = new KeyedMutationQueue();
+    mutationQueues.set(storage, queue);
+  }
+  return queue;
+};
+
 type LocalStorageAdapterOptions = {
   storage: AsyncStorageLike;
   prefix?: string | undefined;
@@ -244,11 +275,22 @@ export const parseStoredMessageRepository = (
 };
 
 class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
+  private storage: AsyncStorageLike;
+  private aui: ReturnType<typeof useAui>;
+  private prefix: string;
+  private mutationQueue: KeyedMutationQueue;
+
   constructor(
-    private storage: AsyncStorageLike,
-    private aui: ReturnType<typeof useAui>,
-    private prefix: string,
-  ) {}
+    storage: AsyncStorageLike,
+    aui: ReturnType<typeof useAui>,
+    prefix: string,
+    mutationQueue: KeyedMutationQueue,
+  ) {
+    this.storage = storage;
+    this.aui = aui;
+    this.prefix = prefix;
+    this.mutationQueue = mutationQueue;
+  }
 
   private _messagesKey(remoteId: string) {
     return `${this.prefix}messages:${remoteId}`;
@@ -266,31 +308,34 @@ class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
     const { remoteId } = await this.aui.threadListItem().initialize();
 
     const key = this._messagesKey(remoteId);
-    const raw = await this.storage.getItem(key);
-    const repo = parseStoredMessageRepository(raw);
+    await this.mutationQueue.run(key, async () => {
+      const raw = await this.storage.getItem(key);
+      const repo = parseStoredMessageRepository(raw);
+
+      const idx = repo.messages.findIndex(
+        (m) => m.message.id === item.message.id,
+      );
+      if (idx >= 0) {
+        repo.messages[idx] = item;
+      } else {
+        repo.messages.push(item);
+      }
+      repo.headId = item.message.id;
 
-    const idx = repo.messages.findIndex(
-      (m) => m.message.id === item.message.id,
-    );
-    if (idx >= 0) {
-      repo.messages[idx] = item;
-    } else {
-      repo.messages.push(item);
-    }
-    repo.headId = item.message.id;
-
-    await this.storage.setItem(key, JSON.stringify(repo));
+      await this.storage.setItem(key, JSON.stringify(repo));
+    });
   }
 }
 
 const createHistoryProvider = (
   storage: AsyncStorageLike,
   prefix: string,
+  mutationQueue: KeyedMutationQueue,
 ): FC<PropsWithChildren> => {
   const Provider: FC<PropsWithChildren> = ({ children }) => {
     const aui = useAui();
     const history = useMemo(
-      () => new AsyncStorageHistoryAdapter(storage, aui, prefix),
+      () => new AsyncStorageHistoryAdapter(storage, aui, prefix, mutationQueue),
       [aui],
     );
     const adapters = useMemo(() => ({ history }), [history]);
@@ -311,6 +356,7 @@ export const createLocalStorageAdapter = (
 
   const threadsKey = `${prefix}threads`;
   const messagesKey = (threadId: string) => `${prefix}messages:${threadId}`;
+  const mutationQueue = getMutationQueue(storage);
 
   const loadThreadMetadata = async (): Promise<StoredThreadMetadata[]> => {
     const raw = await storage.getItem(threadsKey);
@@ -323,8 +369,22 @@ export const createLocalStorageAdapter = (
     await storage.setItem(threadsKey, JSON.stringify(threads));
   };
 
+  const updateThreadMetadata = async (
+    remoteId: string,
+    update: (thread: StoredThreadMetadata) => void,
+  ): Promise<void> => {
+    await mutationQueue.run(threadsKey, async () => {
+      const threads = await loadThreadMetadata();
+      const thread = threads.find((item) => item.remoteId === remoteId);
+      if (thread) {
+        update(thread);
+        await saveThreadMetadata(threads);
+      }
+    });
+  };
+
   const adapter: RemoteThreadListAdapter = {
-    unstable_Provider: createHistoryProvider(storage, prefix),
+    unstable_Provider: createHistoryProvider(storage, prefix, mutationQueue),
 
     async list(): Promise<RemoteThreadListResponse> {
       const threads = await loadThreadMetadata();
@@ -343,64 +403,57 @@ export const createLocalStorageAdapter = (
       threadId: string,
     ): Promise<RemoteThreadInitializeResponse> {
       const remoteId = threadId;
-      const threads = await loadThreadMetadata();
+      return mutationQueue.run(threadsKey, async () => {
+        const threads = await loadThreadMetadata();
 
-      // Only add if not already present
-      if (!threads.some((t) => t.remoteId === remoteId)) {
-        threads.unshift({
-          remoteId,
-          status: "regular",
-        });
-        await saveThreadMetadata(threads);
-      }
+        // Only add if not already present
+        if (!threads.some((t) => t.remoteId === remoteId)) {
+          threads.unshift({
+            remoteId,
+            status: "regular",
+          });
+          await saveThreadMetadata(threads);
+        }
 
-      return { remoteId, externalId: undefined };
+        return { remoteId, externalId: undefined };
+      });
     },
 
     async rename(remoteId: string, newTitle: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.title = newTitle;
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async updateCustom(
       remoteId: string,
       custom: Record<string, unknown> | undefined,
     ): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.custom = custom;
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async archive(remoteId: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.status = "archived";
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async unarchive(remoteId: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.status = "regular";
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async delete(remoteId: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const filtered = threads.filter((t) => t.remoteId !== remoteId);
-      await saveThreadMetadata(filtered);
-      await storage.removeItem(messagesKey(remoteId));
+      await mutationQueue.run(threadsKey, async () => {
+        const threads = await loadThreadMetadata();
+        const filtered = threads.filter((t) => t.remoteId !== remoteId);
+        await saveThreadMetadata(filtered);
+      });
+      const key = messagesKey(remoteId);
+      await mutationQueue.run(key, () => storage.removeItem(key));
     },
 
     async fetch(threadId: string): Promise<RemoteThreadMetadata> {
@@ -427,12 +480,9 @@ export const createLocalStorageAdapter = (
         const title = await titleGenerator.generateTitle(messages);
 
         // Update the stored title
-        const threads = await loadThreadMetadata();
-        const thread = threads.find((t) => t.remoteId === remoteId);
-        if (thread) {
+        await updateThreadMetadata(remoteId, (thread) => {
           thread.title = title;
-          await saveThreadMetadata(threads);
-        }
+        });
 
         // Return a stream with a single text part
         return createAssistantStream((controller) => {

package/src/react/primitive-hooks/useActionBarCopy.test.ts

@@ -0,0 +1,79 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => {
+  const setIsCopied = vi.fn();
+
+  return {
+    setIsCopied,
+    state: {
+      message: {
+        role: "assistant",
+        status: { type: "complete", reason: "stop" },
+        parts: [{ type: "text", text: "Hello" }],
+        isCopied: false,
+      },
+      composer: {
+        isEditing: false,
+        text: "",
+      },
+    },
+    aui: {
+      message: () => ({
+        getCopyText: () => "Hello",
+        setIsCopied,
+      }),
+    },
+  };
+});
+
+vi.mock("react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("react")>()),
+  useCallback: ((callback: unknown) =>
+    callback) as typeof import("react").useCallback,
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: ((selector: (state: typeof mocks.state) => unknown) =>
+    selector(mocks.state)) as typeof import("@assistant-ui/store").useAuiState,
+}));
+
+import { useActionBarCopy } from "./useActionBarCopy";
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe("useActionBarCopy", () => {
+  it("does not report copy success without a clipboard handler", async () => {
+    const { copy, disabled } = useActionBarCopy();
+
+    copy();
+    await Promise.resolve();
+
+    expect(disabled).toBe(true);
+    expect(mocks.setIsCopied).not.toHaveBeenCalled();
+  });
+
+  it("reports copy success after the clipboard handler resolves", async () => {
+    const copyToClipboard = vi.fn();
+    const { copy } = useActionBarCopy({ copyToClipboard });
+
+    copy();
+    await Promise.resolve();
+
+    expect(copyToClipboard).toHaveBeenCalledWith("Hello");
+    expect(mocks.setIsCopied).toHaveBeenCalledWith(true);
+  });
+
+  it("does not report copy success when the clipboard handler rejects", async () => {
+    const copyToClipboard = vi.fn().mockRejectedValue(new Error("denied"));
+    const { copy } = useActionBarCopy({ copyToClipboard });
+
+    copy();
+    await Promise.resolve();
+
+    expect(mocks.setIsCopied).not.toHaveBeenCalled();
+  });
+});

package/src/react/primitive-hooks/useActionBarCopy.ts

@@ -23,13 +23,14 @@ export const useActionBarCopy = ({
   const composerValue = useAuiState((s) => s.composer.text);
 
   const copy = useCallback(() => {
+    if (!copyToClipboard) return;
+
     const valueToCopy = isEditing ? composerValue : aui.message().getCopyText();
     if (!valueToCopy) return;
 
-    const write = copyToClipboard ?? (() => {});
     // The rejection handler swallows clipboard write failures (permission denied,
     // API unavailable) so they don't surface as unhandled promise rejections.
-    Promise.resolve(write(valueToCopy)).then(
+    Promise.resolve(copyToClipboard(valueToCopy)).then(
       () => {
         aui.message().setIsCopied(true);
         setTimeout(() => aui.message().setIsCopied(false), copiedDuration);
@@ -38,5 +39,5 @@ export const useActionBarCopy = ({
     );
   }, [aui, isEditing, composerValue, copiedDuration, copyToClipboard]);
 
-  return { copy, disabled, isCopied };
+  return { copy, disabled: disabled || !copyToClipboard, isCopied };
 };

package/src/react/primitives/chainOfThought/ChainOfThoughtParts.tsx

@@ -51,9 +51,7 @@ const ChainOfThoughtPrimitivePartsInner: FC<{
       Array.from({ length: partsLength }, (_, index) => (
         <ChainOfThoughtPartByIndexProvider key={index} index={index}>
           <RenderChildrenWithAccessor
-            getItemState={(aui) =>
-              aui.chainOfThought().part({ index }).getState()
-            }
+            getItemState={(aui) => aui.part().getState()}
           >
             {(getItem) =>
               children({