@assistant-ui/core 0.2.21 → 0.2.22

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;
@@ -248,6 +279,7 @@ class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
     private storage: AsyncStorageLike,
     private aui: ReturnType<typeof useAui>,
     private prefix: string,
+    private mutationQueue: KeyedMutationQueue,
   ) {}
 
   private _messagesKey(remoteId: string) {
@@ -266,31 +298,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 +346,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 +359,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 +393,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 +470,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({

package/src/react/primitives/message/MessageParts.tsx

@@ -728,50 +728,63 @@ const EMPTY_RUNNING_TEXT_PART: Extract<EnrichedPartState, { type: "text" }> =
  * `<MessagePrimitive.GroupedParts>`. Returns whatever `children`
  * returns — callers decide how to handle a `null` return.
  */
-export const MessagePartChildren: FC<{
+type MessagePartChildrenProps = {
   index: number;
   children: (value: { part: EnrichedPartState }) => ReactNode;
-}> = ({ index, children }) => {
+};
+
+const MessagePartChildrenInner: FC<
+  Pick<MessagePartChildrenProps, "children">
+> = ({ children }) => {
   const aui = useAui();
   // Subscribed (not snapshotted like `tools`) so fallbacks registered
   // after the first render trigger a re-render and `hasUI` re-evaluates.
   const dataRenderers = useAuiState((s) => s.dataRenderers);
 
+  return (
+    <RenderChildrenWithAccessor
+      getItemState={(client) => client.part().getState()}
+    >
+      {(getItem) =>
+        children({
+          get part() {
+            const state = getItem();
+            if (state.type === "tool-call") {
+              const toolsState = aui.tools().getState();
+              const hasUI = resolveToolRender(toolsState, state) !== null;
+              const partMethods = aui.part();
+              return {
+                ...state,
+                toolUI: hasUI ? <RegisteredToolUI /> : null,
+                addResult: partMethods.addToolResult,
+                resume: partMethods.resumeToolCall,
+                respondToApproval: partMethods.respondToToolApproval,
+              };
+            }
+            if (state.type === "data") {
+              const hasUI =
+                getDataRenderer(dataRenderers, state.name, undefined) !==
+                undefined;
+              return {
+                ...state,
+                dataRendererUI: hasUI ? <RegisteredDataRendererUI /> : null,
+              };
+            }
+            return state;
+          },
+        })
+      }
+    </RenderChildrenWithAccessor>
+  );
+};
+
+export const MessagePartChildren: FC<MessagePartChildrenProps> = ({
+  index,
+  children,
+}) => {
   return (
     <PartByIndexProvider index={index}>
-      <RenderChildrenWithAccessor
-        getItemState={(aui) => aui.message().part({ index }).getState()}
-      >
-        {(getItem) =>
-          children({
-            get part() {
-              const state = getItem();
-              if (state.type === "tool-call") {
-                const toolsState = aui.tools().getState();
-                const hasUI = resolveToolRender(toolsState, state) !== null;
-                const partMethods = aui.message().part({ index });
-                return {
-                  ...state,
-                  toolUI: hasUI ? <RegisteredToolUI /> : null,
-                  addResult: partMethods.addToolResult,
-                  resume: partMethods.resumeToolCall,
-                  respondToApproval: partMethods.respondToToolApproval,
-                };
-              }
-              if (state.type === "data") {
-                const hasUI =
-                  getDataRenderer(dataRenderers, state.name, undefined) !==
-                  undefined;
-                return {
-                  ...state,
-                  dataRendererUI: hasUI ? <RegisteredDataRendererUI /> : null,
-                };
-              }
-              return state;
-            },
-          })
-        }
-      </RenderChildrenWithAccessor>
+      <MessagePartChildrenInner>{children}</MessagePartChildrenInner>
     </PartByIndexProvider>
   );
 };

package/src/react/providers/ChainOfThoughtPartByIndexProvider.tsx

@@ -1,16 +1,33 @@
-import type { FC, PropsWithChildren } from "react";
+import { useMemo, type FC, type PropsWithChildren } from "react";
 import { useAui, AuiProvider, Derived } from "@assistant-ui/store";
+import type { PartMethods } from "../../store/scopes/part";
 
 export const ChainOfThoughtPartByIndexProvider: FC<
   PropsWithChildren<{
     index: number;
   }>
 > = ({ index, children }) => {
+  const lastPartRef = useMemo(
+    () => ({ index, current: null as PartMethods | null }),
+    [index],
+  );
   const aui = useAui({
     part: Derived({
       source: "chainOfThought",
       query: { type: "index", index },
-      get: (aui) => aui.chainOfThought().part({ index }),
+      get: (aui) => {
+        const chainOfThought = aui.chainOfThought();
+        if (
+          index >= chainOfThought.getState().parts.length &&
+          lastPartRef.current
+        ) {
+          return lastPartRef.current;
+        }
+
+        const part = chainOfThought.part({ index });
+        lastPartRef.current = part;
+        return part;
+      },
     }),
   });
 

package/src/react/providers/PartByIndexProvider.tsx

@@ -1,16 +1,30 @@
-import type { FC, PropsWithChildren } from "react";
+import { useMemo, type FC, type PropsWithChildren } from "react";
 import { useAui, AuiProvider, Derived } from "@assistant-ui/store";
+import type { PartMethods } from "../../store/scopes/part";
 
 export const PartByIndexProvider: FC<
   PropsWithChildren<{
     index: number;
   }>
 > = ({ index, children }) => {
+  const lastPartRef = useMemo(
+    () => ({ index, current: null as PartMethods | null }),
+    [index],
+  );
   const aui = useAui({
     part: Derived({
       source: "message",
       query: { type: "index", index },
-      get: (aui) => aui.message().part({ index }),
+      get: (aui) => {
+        const message = aui.message();
+        if (index >= message.getState().parts.length && lastPartRef.current) {
+          return lastPartRef.current;
+        }
+
+        const part = message.part({ index });
+        lastPartRef.current = part;
+        return part;
+      },
     }),
   });
 

package/src/react/runtimes/RemoteThreadListHookInstanceManager.test.tsx

@@ -0,0 +1,18 @@
+import { describe, expect, it } from "vitest";
+import type { ThreadListRuntimeCore } from "../../runtime/interfaces/thread-list-runtime-core";
+import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
+
+describe("RemoteThreadListHookInstanceManager", () => {
+  it("rejects a pending start when the thread runtime is stopped", async () => {
+    const manager = new RemoteThreadListHookInstanceManager(() => {
+      throw new Error("Runtime hook should not render during this test");
+    }, {} as ThreadListRuntimeCore);
+
+    const startPromise = manager.startThreadRuntime("thread-1");
+    manager.stopThreadRuntime("thread-1");
+
+    await expect(startPromise).rejects.toThrow(
+      "Thread was deleted before runtime was started",
+    );
+  });
+});

package/src/react/runtimes/RemoteThreadListHookInstanceManager.tsx

@@ -85,6 +85,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   public stopThreadRuntime(threadId: string) {
     this.instances.delete(threadId);
     this.useAliveThreadsKeysChanged.setState({}, true);
+    this._notifySubscribers();
   }
 
   public setRuntimeHook(newRuntimeHook: RemoteThreadListHook) {

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -44,6 +44,8 @@ export class RemoteThreadListThreadListRuntimeCore
   private _loadThreadsPromise: Promise<void> | undefined;
   private _loadMorePromise: Promise<void> | undefined;
   private _loadGeneration = 0;
+  private _switchGeneration = 0;
+  private _switchTask: Promise<void> | undefined;
 
   private _mainThreadId!: string;
   private readonly _state = new OptimisticState<RemoteThreadState>({
@@ -307,15 +309,28 @@ export class RemoteThreadListThreadListRuntimeCore
     return getThreadData(this._state.value, threadIdOrRemoteId);
   }
 
-  public async switchToThread(
+  public switchToThread(
     threadIdOrRemoteId: string,
     options?: { unarchive?: boolean },
+  ): Promise<void> {
+    const generation = ++this._switchGeneration;
+    const task = this._switchToThread(threadIdOrRemoteId, options, generation);
+    this._switchTask = task;
+    return task;
+  }
+
+  private async _switchToThread(
+    threadIdOrRemoteId: string,
+    options: { unarchive?: boolean } | undefined,
+    generation: number,
   ): Promise<void> {
     let data = this.getItemById(threadIdOrRemoteId);
 
     if (!data) {
       const remoteMetadata =
         await this._options.adapter.fetch(threadIdOrRemoteId);
+      if (generation !== this._switchGeneration) return;
+
       const state = this._state.value;
       const mappingId = createThreadMappingId(remoteMetadata.remoteId);
 
@@ -380,8 +395,11 @@ export class RemoteThreadListThreadListRuntimeCore
       task.then(() => this._notifySubscribers());
     }
 
+    if (generation !== this._switchGeneration) return;
+
     if (data.status === "archived" && options?.unarchive !== false) {
       await this.unarchive(data.id);
+      if (generation !== this._switchGeneration) return;
     }
     this._mainThreadId = data.id;
 
@@ -389,13 +407,21 @@ export class RemoteThreadListThreadListRuntimeCore
     this._notifyThreadIdChange();
   }
 
-  public async switchToNewThread(): Promise<void> {
+  public switchToNewThread(): Promise<void> {
+    const generation = ++this._switchGeneration;
+    const task = this._switchToNewThread(generation);
+    this._switchTask = task;
+    return task;
+  }
+
+  private async _switchToNewThread(generation: number): Promise<void> {
     // an initialization transaction is in progress, wait for it to settle
     while (
       this._state.baseValue.newThreadId !== undefined &&
       this._state.value.newThreadId === undefined
     ) {
       await this._state.waitForUpdate();
+      if (generation !== this._switchGeneration) return;
     }
 
     const state = this._state.value;
@@ -427,7 +453,7 @@ export class RemoteThreadListThreadListRuntimeCore
       });
     }
 
-    return this.switchToThread(id);
+    return this._switchToThread(id, undefined, generation);
   }
 
   public initialize = async (threadId: string) => {
@@ -436,10 +462,11 @@ export class RemoteThreadListThreadListRuntimeCore
       if (!data) throw threadNotFoundError(threadId, "initializing it");
       if (data.status === "new")
         throw threadStatusError(threadId, data.status, "be initialized here");
-      return data.initializeTask;
+      const { remoteId, externalId } = await data.initializeTask;
+      return { remoteId, externalId };
     }
 
-    return this._state.optimisticUpdate({
+    const { remoteId, externalId } = await this._state.optimisticUpdate({
       execute: () => {
         return this._options.adapter.initialize(threadId);
       },
@@ -482,6 +509,7 @@ export class RemoteThreadListThreadListRuntimeCore
         };
       },
     });
+    return { remoteId, externalId };
   };
 
   public generateTitle = async (threadId: string) => {
@@ -604,8 +632,21 @@ export class RemoteThreadListThreadListRuntimeCore
     if (threadId === this.newThreadId)
       throw new Error("Cannot ensure new thread is not main");
 
-    if (threadId === this._mainThreadId) {
-      await this.switchToNewThread();
+    let lastAwaitedTask: Promise<void> | undefined;
+
+    while (threadId === this._mainThreadId) {
+      let switchTask = this._switchTask;
+      const startedFallback = !switchTask || switchTask === lastAwaitedTask;
+      if (startedFallback) switchTask = this.switchToNewThread();
+      lastAwaitedTask = switchTask;
+
+      try {
+        await switchTask;
+      } catch (error) {
+        if (startedFallback && this._switchTask === switchTask) {
+          throw error;
+        }
+      }
     }
   }