@assistant-ui/core 0.2.21 → 0.2.23

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);
 
@@ -341,22 +356,36 @@ export class RemoteThreadListThreadListRuntimeCore
         [remoteMetadata.remoteId]: mappingId,
       };
 
-      // Filter both arrays first so a concurrent `list()` can't leave the id
-      // duplicated or under the wrong status.
+      // A concurrent `list()` may already have placed this thread; keep that
+      // position and only merge metadata. A genuinely absent thread stays
+      // appended: it may live on an unloaded page, and a prepend would pin it
+      // above newer threads permanently since `classifyThreads` skips ids it
+      // has already seen. Filtering both arrays first still prevents
+      // duplication or a wrong-status entry from `list()`.
+      const remoteId = remoteMetadata.remoteId;
+      const wasInTarget =
+        remoteMetadata.status === "regular"
+          ? state.threadIds.includes(remoteId)
+          : state.archivedThreadIds.includes(remoteId);
+
       const threadIdsWithoutRemote = state.threadIds.filter(
-        (id) => id !== remoteMetadata.remoteId,
+        (id) => id !== remoteId,
       );
       const archivedThreadIdsWithoutRemote = state.archivedThreadIds.filter(
-        (id) => id !== remoteMetadata.remoteId,
+        (id) => id !== remoteId,
       );
 
       const newThreadIds =
         remoteMetadata.status === "regular"
-          ? [...threadIdsWithoutRemote, remoteMetadata.remoteId]
+          ? wasInTarget
+            ? state.threadIds
+            : [...threadIdsWithoutRemote, remoteId]
           : threadIdsWithoutRemote;
       const newArchivedThreadIds =
         remoteMetadata.status === "archived"
-          ? [...archivedThreadIdsWithoutRemote, remoteMetadata.remoteId]
+          ? wasInTarget
+            ? state.archivedThreadIds
+            : [...archivedThreadIdsWithoutRemote, remoteId]
           : archivedThreadIdsWithoutRemote;
 
       this._state.update({
@@ -380,8 +409,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 +421,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 +467,7 @@ export class RemoteThreadListThreadListRuntimeCore
       });
     }
 
-    return this.switchToThread(id);
+    return this._switchToThread(id, undefined, generation);
   }
 
   public initialize = async (threadId: string) => {
@@ -436,10 +476,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 +523,7 @@ export class RemoteThreadListThreadListRuntimeCore
         };
       },
     });
+    return { remoteId, externalId };
   };
 
   public generateTitle = async (threadId: string) => {
@@ -604,8 +646,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;
+        }
+      }
     }
   }
 

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it } from "vitest";
+import { extractAuiV0 } from "./AssistantCloudThreadHistoryAdapter";
+
+const auiV0Message = (status: { type: string; reason?: string }) => ({
+  role: "assistant",
+  status,
+  content: [
+    {
+      type: "tool-call",
+      toolCallId: "call-1",
+      toolName: "send_email",
+      args: {},
+      argsText: "{}",
+    },
+  ],
+  metadata: { steps: [{ usage: { inputTokens: 1, outputTokens: 2 } }] },
+});
+
+describe("extractAuiV0", () => {
+  it("returns null for a requires-action message so paused writes never report", () => {
+    expect(
+      extractAuiV0(
+        auiV0Message({ type: "requires-action", reason: "tool-calls" }),
+      ),
+    ).toBeNull();
+  });
+
+  it("reports a terminal message once with its usage", () => {
+    const result = extractAuiV0(auiV0Message({ type: "complete" }));
+    expect(result?.status).toBe("completed");
+    expect(result?.inputTokens).toBe(1);
+    expect(result?.outputTokens).toBe(2);
+  });
+});

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.test.tsx

@@ -0,0 +1,69 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import type { AssistantCloud } from "assistant-cloud";
+import { describe, expect, it, vi } from "vitest";
+import { useAssistantCloudThreadHistoryAdapter } from "./AssistantCloudThreadHistoryAdapter";
+
+const mocks = vi.hoisted(() => {
+  const threadListItem = {
+    getState: () => ({ remoteId: "thread-1" }),
+  };
+
+  return {
+    assistantClient: {
+      threadListItem: () => threadListItem,
+    },
+  };
+});
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.assistantClient,
+}));
+
+const makeCloud = () =>
+  ({
+    threads: {
+      messages: {
+        list: vi.fn().mockResolvedValue({ messages: [] }),
+      },
+    },
+  }) as unknown as AssistantCloud;
+
+describe("useAssistantCloudThreadHistoryAdapter", () => {
+  it("refreshes formatted persistence when the Cloud client changes", async () => {
+    const firstCloud = makeCloud();
+    const secondCloud = makeCloud();
+    const cloudRef = { current: firstCloud };
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter(cloudRef),
+    );
+    const formatted = result.current.withFormat<
+      { id: string },
+      Record<string, unknown>
+    >({
+      format: "test",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message) => message.id,
+    });
+
+    await formatted.load();
+    await formatted.load();
+
+    expect(firstCloud.threads.messages.list).toHaveBeenCalledTimes(2);
+
+    cloudRef.current = secondCloud;
+    await formatted.load();
+
+    expect(firstCloud.threads.messages.list).toHaveBeenCalledTimes(2);
+    expect(secondCloud.threads.messages.list).toHaveBeenCalledOnce();
+    expect(secondCloud.threads.messages.list).toHaveBeenCalledWith("thread-1", {
+      format: "test",
+    });
+  });
+});

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.ts

@@ -22,17 +22,20 @@ const globalPersistence = new WeakMap<
 >();
 
 class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
-  constructor(
-    private cloudRef: RefObject<AssistantCloud>,
-    private aui: AssistantClient,
-  ) {}
+  private cloudRef: RefObject<AssistantCloud>;
+  private aui: AssistantClient;
+
+  constructor(cloudRef: RefObject<AssistantCloud>, aui: AssistantClient) {
+    this.cloudRef = cloudRef;
+    this.aui = aui;
+  }
 
   private get _persistence(): CloudMessagePersistence {
     const key = this.aui.threadListItem();
     if (!globalPersistence.has(key)) {
       globalPersistence.set(
         key,
-        new CloudMessagePersistence(this.cloudRef.current),
+        new CloudMessagePersistence(() => this.cloudRef.current),
       );
     }
     return globalPersistence.get(key)!;
@@ -102,6 +105,21 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
     }
   }
 
+  async update(item: ExportedMessageRepositoryItem) {
+    if (!this._persistence.isPersisted(item.message.id)) {
+      return this.append(item);
+    }
+    const { message } = item;
+    const remoteId = this.aui.threadListItem().getState().remoteId;
+    if (!remoteId) return;
+    const encoded = auiV0Encode(message);
+    await this._persistence.update(remoteId, message.id, "aui/v0", encoded);
+
+    if (this.cloudRef.current.telemetry.enabled) {
+      this._maybeReportRun(remoteId, "aui/v0", encoded);
+    }
+  }
+
   async delete() {
     throw new Error(
       "Assistant Cloud does not support deleting thread messages yet.",
@@ -341,7 +359,7 @@ const AUI_STATUS_MAP: Record<string, TelemetryData["status"]> = {
   incomplete: "incomplete",
 };
 
-function extractAuiV0<T>(content: T): TelemetryData | null {
+export function extractAuiV0<T>(content: T): TelemetryData | null {
   const msg = content as {
     role?: string;
     status?: { type: string };
@@ -369,6 +387,9 @@ function extractAuiV0<T>(content: T): TelemetryData | null {
   };
 
   if (msg.role !== "assistant") return null;
+  // A paused (requires-action) write is not a finished run; reporting it would
+  // mislabel it "completed" and double-count steps once the terminal write reports.
+  if (msg.status?.type === "requires-action") return null;
 
   const toolCalls = msg.content
     ?.filter((p) => p.type === "tool-call" && p.toolName && p.toolCallId)

package/src/react/runtimes/cloud/CloudFileAttachmentAdapter.test.ts

@@ -57,6 +57,22 @@ describe("CloudFileAttachmentAdapter", () => {
     ]);
   });
 
+  it("uploads when Web Crypto is unavailable", async () => {
+    vi.stubGlobal("crypto", undefined);
+    vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true }));
+    const adapter = new CloudFileAttachmentAdapter(makeCloud());
+
+    const yields = await drain(adapter);
+
+    expect(yields).toHaveLength(2);
+    expect(yields[0]?.id).toBeTruthy();
+    expect(yields[1]?.id).toBe(yields[0]?.id);
+    expect(yields.at(-1)?.status).toEqual({
+      type: "requires-action",
+      reason: "composer-send",
+    });
+  });
+
   it("marks the attachment incomplete when the upload returns an HTTP error", async () => {
     const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
     vi.stubGlobal(

package/src/react/runtimes/cloud/CloudFileAttachmentAdapter.ts

@@ -6,6 +6,7 @@ import type {
 } from "../../../types/attachment";
 import type { ThreadUserMessagePart } from "../../../types/message";
 import type { AttachmentAdapter } from "../../../adapters/attachment";
+import { generateId } from "../../../utils/id";
 
 const guessAttachmentType = (
   contentType: string,
@@ -18,7 +19,11 @@ const guessAttachmentType = (
 export class CloudFileAttachmentAdapter implements AttachmentAdapter {
   public accept = "*";
 
-  constructor(private cloud: AssistantCloud) {}
+  private cloud: AssistantCloud;
+
+  constructor(cloud: AssistantCloud) {
+    this.cloud = cloud;
+  }
 
   private uploadedUrls = new Map<string, string>();
 
@@ -27,7 +32,7 @@ export class CloudFileAttachmentAdapter implements AttachmentAdapter {
   }: {
     file: File;
   }): AsyncGenerator<PendingAttachment, void> {
-    const id = crypto.randomUUID();
+    const id = generateId();
     const type = guessAttachmentType(file.type);
     let attachment: PendingAttachment = {
       id,

package/src/react/runtimes/cloud/auiV0.ts

@@ -6,6 +6,7 @@ import type {
   SourceProviderMetadata,
   ThreadMessage,
   TextMessagePart,
+  ToolApprovalOption,
   Unstable_AudioMessagePart,
 } from "../../../types/message";
 import type { CompleteAttachment } from "../../../types/attachment";
@@ -18,6 +19,16 @@ import type {
 } from "assistant-stream/utils";
 import type { ExportedMessageRepositoryItem } from "../../../runtime/utils/message-repository";
 
+type AuiV0ToolApproval = {
+  readonly id: string;
+  readonly approved?: boolean;
+  readonly reason?: string;
+  readonly isAutomatic?: boolean;
+  readonly options?: readonly ToolApprovalOption[];
+  readonly optionId?: string;
+  readonly resolution?: "cancelled" | "expired";
+};
+
 type AuiV0MessagePart =
   | {
       readonly type: "text";
@@ -51,6 +62,7 @@ type AuiV0MessagePart =
       readonly args: ReadonlyJSONObject;
       readonly result?: ReadonlyJSONValue;
       readonly isError?: true;
+      readonly approval?: AuiV0ToolApproval;
     }
   | {
       readonly type: "tool-call";
@@ -59,6 +71,7 @@ type AuiV0MessagePart =
       readonly argsText: string;
       readonly result?: ReadonlyJSONValue;
       readonly isError?: true;
+      readonly approval?: AuiV0ToolApproval;
     }
   | {
       readonly type: "image";
@@ -216,6 +229,7 @@ export function auiV0Encode(message: ThreadMessage): AuiV0Message {
               ? { result: part.result as ReadonlyJSONValue }
               : undefined),
             ...(part.isError ? { isError: true } : undefined),
+            ...(part.approval ? { approval: part.approval } : undefined),
           };
         }