@assistant-ui/core 0.3.15 → 0.3.17

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -1,18 +1,29 @@
-import type { ThreadListRuntimeCore } from "../../runtime/interfaces/thread-list-runtime-core";
-import { generateId } from "../../utils/id";
-import { BaseSubscribable } from "../../subscribable/subscribable";
+import type {
+  ThreadListRuntimeCore,
+  ThreadListRuntimeEvent,
+} from "../../runtime/interfaces/thread-list-runtime-core";
+import {
+  BaseSubscribable,
+  WritableSubscribable,
+} from "../../subscribable/subscribable";
+import { useSubscribable } from "../../store/runtime-clients/useSubscribable";
 import { OptimisticState } from "../../runtimes/remote-thread-list/optimistic-state";
 import { EMPTY_THREAD_CORE } from "../../runtimes/remote-thread-list/empty-thread-core";
 import type {
+  ClassifyAccumulator,
   RemoteThreadData,
   RemoteThreadState,
 } from "../../runtimes/remote-thread-list/remote-thread-state";
 import {
   classifyThreads,
+  createEmptyRemoteThreadState,
   createThreadMappingId,
   getThreadData,
   normalizeCursor,
   updateStatusReducer,
+  preserveMidLoadTransitions,
+  seedNewThread,
+  statusSnapshot,
 } from "../../runtimes/remote-thread-list/remote-thread-state";
 import type {
   RemoteThreadListAdapter,
@@ -21,7 +32,7 @@ import type {
 } from "../../runtimes/remote-thread-list/types";
 import { ThreadListAdapterChangedError } from "../../runtimes/remote-thread-list/adapter-changed";
 import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
-import { isTitleSourceMessage } from "./RemoteThreadResource";
+import { isTitleSourceMessage } from "../../runtimes/remote-thread-list/title";
 import {
   type ComponentType,
   type FC,
@@ -31,7 +42,6 @@ import {
   useId,
 } from "react";
 import { useAui } from "@assistant-ui/store";
-import { create } from "zustand";
 import { AssistantMessageStream } from "assistant-stream";
 import type { ModelContextProvider } from "../../model-context/types";
 import { RuntimeAdapterProvider } from "./RuntimeAdapterProvider";
@@ -49,47 +59,7 @@ const threadStatusError = (
     `Thread "${threadIdOrRemoteId}" has status "${status}", so it cannot ${action}.`,
   );
 
-const EMPTY_REMOTE_STATE: RemoteThreadState = {
-  isLoading: true,
-  isLoadingMore: false,
-  cursor: undefined,
-  newThreadId: undefined,
-  threadIds: [],
-  archivedThreadIds: [],
-  threadIdMap: {},
-  threadData: {},
-};
-
-const addNewThread = (state: RemoteThreadState) => {
-  let id: string;
-  do {
-    id = `__LOCALID_${generateId()}`;
-  } while (state.threadIdMap[id]);
-
-  const mappingId = createThreadMappingId(id);
-  return {
-    id,
-    state: {
-      ...state,
-      newThreadId: id,
-      threadIdMap: {
-        ...state.threadIdMap,
-        [id]: mappingId,
-      },
-      threadData: {
-        ...state.threadData,
-        [mappingId]: {
-          status: "new",
-          id,
-          remoteId: undefined,
-          externalId: undefined,
-          title: undefined,
-          custom: undefined,
-        } satisfies RemoteThreadData,
-      },
-    },
-  };
-};
+const EMPTY_REMOTE_STATE = createEmptyRemoteThreadState();
 
 export class RemoteThreadListThreadListRuntimeCore
   extends BaseSubscribable
@@ -158,10 +128,12 @@ export class RemoteThreadListThreadListRuntimeCore
     if (!this._loadThreadsPromise) {
       const generation = this._loadGeneration;
       let replacedList = false;
+      const statusAtRequest = statusSnapshot(this._state.baseValue);
       this._loadThreadsPromise = this._state
         .optimisticUpdate({
           execute: () => this._options.adapter.list(),
           loading: (state) => {
+            if (generation !== this._loadGeneration) return state;
             return {
               ...state,
               isLoading: true,
@@ -183,18 +155,19 @@ export class RemoteThreadListThreadListRuntimeCore
             const fresh = classifyThreads(l.threads, {
               threadIds: [],
               archivedThreadIds: [],
-              threadIdMap: {},
-              threadData: {},
+              threadIdMap: { ...state.threadIdMap },
+              threadData: { ...state.threadData },
             });
-            return {
+            const merged = {
               ...state,
               isLoading: false,
               cursor: normalizeCursor(l.nextCursor),
               threadIds: fresh.threadIds,
               archivedThreadIds: fresh.archivedThreadIds,
-              threadIdMap: { ...state.threadIdMap, ...fresh.threadIdMap },
-              threadData: { ...state.threadData, ...fresh.threadData },
+              threadIdMap: fresh.threadIdMap,
+              threadData: fresh.threadData,
             };
+            return preserveMidLoadTransitions(merged, state, statusAtRequest);
           },
         })
         .catch((error: unknown) => {
@@ -241,7 +214,10 @@ export class RemoteThreadListThreadListRuntimeCore
     const dedup = this._state
       .optimisticUpdate({
         execute: () => adapter.list({ after: cursor }),
-        loading: (state) => ({ ...state, isLoadingMore: true }),
+        loading: (state) => {
+          if (generation !== this._loadGeneration) return state;
+          return { ...state, isLoadingMore: true };
+        },
         then: (state, l) => {
           if (generation !== this._loadGeneration) return state;
           if (adapter !== this._options.adapter) return state;
@@ -302,15 +278,15 @@ export class RemoteThreadListThreadListRuntimeCore
       // synchronous notify would re-enter store consumers mid-render.
       queueMicrotask(() => this._notifySubscribers());
     });
-    this.useProvider = create(() => ({
-      Provider: this.resolveProvider(options.adapter),
-    }));
+    this.providerStore = new WritableSubscribable(
+      this.resolveProvider(options.adapter),
+    );
     this.__internal_setOptions(options);
     this.switchToNewThread();
   }
 
   private _initialThreadLoaded = false;
-  private useProvider;
+  private providerStore;
 
   public __internal_setOptions(options: RemoteThreadListOptions) {
     if (this._options === options) return;
@@ -324,10 +300,7 @@ export class RemoteThreadListThreadListRuntimeCore
 
     this._options = options;
 
-    const Provider = this.resolveProvider(options.adapter);
-    if (Provider !== this.useProvider.getState().Provider) {
-      this.useProvider.setState({ Provider }, true);
-    }
+    this.providerStore.setState(this.resolveProvider(options.adapter));
 
     this._hookManager.setRuntimeHook(options.runtimeHook);
 
@@ -372,24 +345,11 @@ export class RemoteThreadListThreadListRuntimeCore
     threads: readonly RemoteThreadMetadata[],
     cursor: string | undefined,
   ): RemoteThreadState {
-    const fresh = classifyThreads(threads, {
-      threadIds: [],
-      archivedThreadIds: [],
-      threadIdMap: {},
-      threadData: {},
-    });
-    const threadIdMap = { ...fresh.threadIdMap };
-    const threadData = { ...fresh.threadData };
-    const threadIds = [...fresh.threadIds];
-    const archivedThreadIds = [...fresh.archivedThreadIds];
-
+    const carried: RemoteThreadData[] = [];
     if (state.newThreadId) {
       const mappingId = state.threadIdMap[state.newThreadId];
       const draft = mappingId ? state.threadData[mappingId] : undefined;
-      if (draft?.status === "new") {
-        threadIdMap[state.newThreadId] = mappingId!;
-        threadData[mappingId!] = draft;
-      }
+      if (draft?.status === "new") carried.push(draft);
     }
 
     const stale = this._staleThreadIdsOnReplace;
@@ -397,26 +357,44 @@ export class RemoteThreadListThreadListRuntimeCore
       for (const item of Object.values(state.threadData)) {
         if (stale.has(item.id)) continue;
         if (item.status !== "new" && item.remoteId === undefined) continue;
-        const mappingId = createThreadMappingId(item.id);
-        if (threadData[mappingId]) continue;
-        threadIdMap[item.id] = mappingId;
-        if (item.remoteId !== undefined) {
-          threadIdMap[item.remoteId] = mappingId;
-        }
-        threadData[mappingId] = item;
-        if (item.remoteId === undefined) continue;
-        if (item.status === "regular" && !threadIds.includes(item.remoteId)) {
-          threadIds.push(item.remoteId);
-        } else if (
-          item.status === "archived" &&
-          !archivedThreadIds.includes(item.remoteId)
-        ) {
-          archivedThreadIds.push(item.remoteId);
-        }
+        carried.push(item);
       }
     }
     this._staleThreadIdsOnReplace = undefined;
 
+    const seed: ClassifyAccumulator = {
+      threadIds: [],
+      archivedThreadIds: [],
+      threadIdMap: {},
+      threadData: {},
+    };
+    for (const item of carried) {
+      const mappingId = createThreadMappingId(item.id);
+      if (seed.threadData[mappingId]) continue;
+      seed.threadIdMap[item.id] = mappingId;
+      if (item.remoteId !== undefined) {
+        seed.threadIdMap[item.remoteId] = mappingId;
+      }
+      seed.threadData[mappingId] = item;
+    }
+
+    const { threadIds, archivedThreadIds, threadIdMap, threadData } =
+      classifyThreads(threads, seed);
+
+    for (const item of carried) {
+      if (item.remoteId === undefined) continue;
+      const current = threadData[createThreadMappingId(item.id)];
+      if (current === undefined) continue;
+      if (current.status === "regular" && !threadIds.includes(current.id)) {
+        threadIds.push(current.id);
+      } else if (
+        current.status === "archived" &&
+        !archivedThreadIds.includes(current.id)
+      ) {
+        archivedThreadIds.push(current.id);
+      }
+    }
+
     let nextState: RemoteThreadState = {
       ...state,
       isLoading: false,
@@ -437,7 +415,7 @@ export class RemoteThreadListThreadListRuntimeCore
       if (preservedDraft !== undefined) {
         this._mainThreadId = preservedDraft;
       } else {
-        const seeded = addNewThread(nextState);
+        const seeded = seedNewThread(nextState);
         this._mainThreadId = seeded.id;
         nextState = seeded.state;
       }
@@ -590,6 +568,12 @@ export class RemoteThreadListThreadListRuntimeCore
     return this._hookManager.__internal_isThreadRunning(data.id);
   }
 
+  public unstable_subscribeThreadEvents(
+    callback: (event: ThreadListRuntimeEvent) => void,
+  ) {
+    return this._hookManager.__internal_subscribeThreadEvents(callback);
+  }
+
   public getItemById(threadIdOrRemoteId: string) {
     return getThreadData(this._state.value, threadIdOrRemoteId);
   }
@@ -664,9 +648,8 @@ export class RemoteThreadListThreadListRuntimeCore
       // 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()`.
+      // above newer threads permanently. Filtering both arrays first still
+      // prevents duplication or a wrong-status entry from `list()`.
       const remoteId = remoteMetadata.remoteId;
       const wasInTarget =
         remoteMetadata.status === "regular"
@@ -719,11 +702,22 @@ export class RemoteThreadListThreadListRuntimeCore
 
     if (generation !== this._switchGeneration) return;
 
-    if (data.status === "archived" && options?.unarchive !== false) {
-      await this.unarchive(data.id);
+    let current = this.getItemById(data.id);
+    if (current?.id !== data.id) return;
+
+    if (current.status === "archived" && options?.unarchive !== false) {
+      await current.initializeTask;
       if (generation !== this._switchGeneration) return;
+      current = this.getItemById(data.id);
+      if (current?.id !== data.id) return;
+      if (current.status === "archived") {
+        await this.unarchive(current.id);
+        if (generation !== this._switchGeneration) return;
+        current = this.getItemById(data.id);
+        if (current?.id !== data.id) return;
+      }
     }
-    this._mainThreadId = data.id;
+    this._mainThreadId = current.id;
 
     this._notifySubscribers();
     this._notifyThreadIdChange(emitThreadIdChange);
@@ -762,7 +756,7 @@ export class RemoteThreadListThreadListRuntimeCore
     const state = this._state.baseValue;
     let id: string | undefined = this._state.value.newThreadId;
     if (id === undefined) {
-      const next = addNewThread(state);
+      const next = seedNewThread(state);
       id = next.id;
       this._state.update(next.state);
     }
@@ -811,21 +805,36 @@ export class RemoteThreadListThreadListRuntimeCore
         if (!data) return state;
 
         const mappingId = createThreadMappingId(threadId);
+        // A list() response that landed while this initialize was in flight
+        // could not know the remote id yet, so it may have minted its own slot
+        // for it; that slot collapses into this one.
+        const listedMappingId = state.threadIdMap[remoteId];
+        const orphan =
+          listedMappingId !== undefined && listedMappingId !== mappingId
+            ? state.threadData[listedMappingId]
+            : undefined;
+
+        const threadData = { ...state.threadData };
+        if (orphan !== undefined) delete threadData[listedMappingId!];
+        threadData[mappingId] = {
+          ...data,
+          initializeTask: Promise.resolve({ remoteId, externalId }),
+          remoteId,
+          externalId,
+        } as RemoteThreadData;
+
+        const rewire = (ids: readonly string[]) =>
+          orphan === undefined ? ids : ids.filter((id) => id !== orphan.id);
+
         return {
           ...state,
+          threadIds: rewire(state.threadIds),
+          archivedThreadIds: rewire(state.archivedThreadIds),
           threadIdMap: {
             ...state.threadIdMap,
             [remoteId]: mappingId,
           },
-          threadData: {
-            ...state.threadData,
-            [mappingId]: {
-              ...data,
-              initializeTask: Promise.resolve({ remoteId, externalId }),
-              remoteId,
-              externalId,
-            },
-          },
+          threadData,
         };
       },
     });
@@ -857,17 +866,22 @@ export class RemoteThreadListThreadListRuntimeCore
     for await (const result of messageStream) {
       if (adapterGeneration !== this._adapterGeneration) return;
       const newTitle = result.parts.filter((c) => c.type === "text")[0]?.text;
-      const state = this._state.baseValue;
-      const currentData = getThreadData(state, data.id);
-      if (!currentData) continue;
-      this._state.update({
-        ...state,
-        threadData: {
-          ...state.threadData,
-          [currentData.id]: {
-            ...currentData,
-            title: newTitle,
-          },
+      await this._state.optimisticUpdate({
+        execute: async () => {},
+        optimistic: (state) => {
+          if (adapterGeneration !== this._adapterGeneration) return state;
+          const currentData = getThreadData(state, data.id);
+          if (!currentData) return state;
+          return {
+            ...state,
+            threadData: {
+              ...state.threadData,
+              [currentData.id]: {
+                ...currentData,
+                title: newTitle,
+              },
+            },
+          };
         },
       });
     }
@@ -1079,19 +1093,21 @@ export class RemoteThreadListThreadListRuntimeCore
     this._hookManager.stopThreadRuntime(data.id);
   }
 
-  private useBoundIds = create<string[]>(() => []);
+  private boundIdsStore = new WritableSubscribable<readonly string[]>([]);
 
   public __internal_RenderComponent: FC = () => {
     const id = useId();
     useEffect(() => {
-      this.useBoundIds.setState((s) => [...s, id], true);
+      this.boundIdsStore.setState([...this.boundIdsStore.getState(), id]);
       return () => {
-        this.useBoundIds.setState((s) => s.filter((i) => i !== id), true);
+        this.boundIdsStore.setState(
+          this.boundIdsStore.getState().filter((i) => i !== id),
+        );
       };
     }, [id]);
 
-    const boundIds = this.useBoundIds();
-    const { Provider } = this.useProvider();
+    const boundIds = useSubscribable(this.boundIdsStore);
+    const Provider = useSubscribable(this.providerStore);
     const aui = useAui();
     const enabled = boundIds.length === 0 || boundIds[0] === id;
 

package/src/react/runtimes/RemoteThreadResource.ts

@@ -22,6 +22,7 @@ import {
   type RuntimeAdapters,
 } from "./RuntimeAdapterProvider";
 import { isSilentRuntimeAction } from "../../utils/silent-runtime-action";
+import { isTitleSourceMessage } from "../../runtimes/remote-thread-list/title";
 
 export type RemoteThreadListHook = () => AssistantRuntime;
 
@@ -39,10 +40,6 @@ export type RemoteThreadResourceProps = {
   ) => void;
 };
 
-export const isTitleSourceMessage = (message: {
-  status?: { type: string } | undefined;
-}) => message.status?.type !== "running";
-
 export const subscribeToTitleGeneration = (
   threadRuntime: AssistantRuntime["thread"],
   itemRuntime: ThreadListItemRuntime,

package/src/react/runtimes/assistant-transport/runManager.ts

@@ -1,4 +1,5 @@
 import { useCallback, useEffect, useRef, useState } from "react";
+import { invokeUserCallback } from "../../../utils/invoke-user-callback";
 import { useLatestRef } from "./useLatestRef";
 
 export type RunManager = Readonly<{
@@ -11,25 +12,11 @@ const disposeReason = Symbol("assistant-transport-dispose");
 
 type LifecycleCallbackName = "onCancel" | "onError" | "onFinish";
 
-const reportCallbackError = (name: LifecycleCallbackName, error: unknown) => {
-  console.error(
-    `[assistant-ui] Assistant transport ${name} callback threw an error`,
-    error,
-  );
-};
-
-const invokeCallback = async (
+const invokeCallback = (
   name: LifecycleCallbackName,
   callback: (() => unknown) | undefined,
-): Promise<void> => {
-  try {
-    await callback?.();
-  } catch (error) {
-    try {
-      reportCallbackError(name, error);
-    } catch {}
-  }
-};
+) =>
+  invokeUserCallback("assistant-ui", `Assistant transport ${name}`, callback);
 
 export function useRunManager(config: {
   onRun: (signal: AbortSignal) => Promise<void>;

package/src/react/runtimes/assistant-transport/transport-scheduling.test.ts

@@ -1,8 +1,8 @@
 // @vitest-environment jsdom
 
-import { act, renderHook, waitFor } from "@testing-library/react";
+import { act, render, renderHook, waitFor } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
-import { createElement, StrictMode, useRef } from "react";
+import { createElement, StrictMode, useLayoutEffect, useRef } from "react";
 import { useCommandQueue } from "./commandQueue";
 import { useRunManager } from "./runManager";
 import type { AssistantTransportCommand } from "./types";
@@ -85,6 +85,55 @@ const useTransportSchedulingHarness = (
 };
 
 describe("assistant transport scheduling contracts", () => {
+  it("uses current callbacks when scheduled by a descendant layout effect", async () => {
+    const onRunA = vi.fn(async () => {});
+    const onRunB = vi.fn(async () => {});
+    // act flushes passive effects before real microtasks, so preserve the
+    // browser's microtask-before-passive-effect ordering at the layout boundary.
+    const queueMicrotaskSpy = vi
+      .spyOn(globalThis, "queueMicrotask")
+      .mockImplementation((callback) => callback());
+
+    const Scheduler = ({
+      schedule,
+      enabled,
+    }: {
+      schedule: () => void;
+      enabled: boolean;
+    }) => {
+      useLayoutEffect(() => {
+        if (enabled) schedule();
+      }, [enabled, schedule]);
+      return null;
+    };
+    const Probe = ({
+      enabled,
+      onRun,
+    }: {
+      enabled: boolean;
+      onRun: (signal: AbortSignal) => Promise<void>;
+    }) => {
+      const runManager = useRunManager({ onRun });
+      return createElement(Scheduler, {
+        enabled,
+        schedule: runManager.schedule,
+      });
+    };
+
+    try {
+      const view = render(
+        createElement(Probe, { enabled: false, onRun: onRunA }),
+      );
+      view.rerender(createElement(Probe, { enabled: true, onRun: onRunB }));
+      await act(async () => {});
+
+      expect(onRunB).toHaveBeenCalledTimes(1);
+      expect(onRunA).not.toHaveBeenCalled();
+    } finally {
+      queueMicrotaskSpy.mockRestore();
+    }
+  });
+
   it("runs in single-flight mode and schedules exactly one follow-up run", async () => {
     const gate = createDeferred();
     const { result } = renderHook(() =>

package/src/react/runtimes/assistant-transport/useLatestRef.ts

@@ -1,8 +1,8 @@
-import { useEffect, useRef } from "react";
+import { useInsertionEffect, useRef } from "react";
 
 export function useLatestRef<T>(value: T) {
   const ref = useRef(value);
-  useEffect(() => {
+  useInsertionEffect(() => {
     ref.current = value;
   }, [value]);
   return ref as { current: T };

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

@@ -31,4 +31,49 @@ describe("extractAuiV0", () => {
     expect(result?.inputTokens).toBe(1);
     expect(result?.outputTokens).toBe(2);
   });
+
+  it("sums step usage reported under the AI SDK v7 token details", () => {
+    const result = extractAuiV0({
+      ...auiV0Message({ type: "complete" }),
+      metadata: {
+        steps: [
+          {
+            usage: {
+              inputTokens: 10,
+              outputTokens: 4,
+              inputTokenDetails: { cacheReadTokens: 6 },
+              outputTokenDetails: { reasoningTokens: 1 },
+            },
+          },
+          {
+            usage: {
+              inputTokens: 5,
+              outputTokens: 2,
+              inputTokenDetails: { cacheReadTokens: 3 },
+              outputTokenDetails: { reasoningTokens: 1 },
+            },
+          },
+        ],
+      },
+    });
+
+    expect(result?.inputTokens).toBe(15);
+    expect(result?.outputTokens).toBe(6);
+    expect(result?.cachedInputTokens).toBe(9);
+    expect(result?.reasoningTokens).toBe(2);
+    expect(result?.steps).toEqual([
+      {
+        input_tokens: 10,
+        output_tokens: 4,
+        reasoning_tokens: 1,
+        cached_input_tokens: 6,
+      },
+      {
+        input_tokens: 5,
+        output_tokens: 2,
+        reasoning_tokens: 1,
+        cached_input_tokens: 3,
+      },
+    ]);
+  });
 });