@assistant-ui/core 0.3.17 → 0.3.19

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -7,6 +7,7 @@ import {
   WritableSubscribable,
 } from "../../subscribable/subscribable";
 import { useSubscribable } from "../../store/runtime-clients/useSubscribable";
+import { nullProtoRecord } from "../../utils/record";
 import { OptimisticState } from "../../runtimes/remote-thread-list/optimistic-state";
 import { EMPTY_THREAD_CORE } from "../../runtimes/remote-thread-list/empty-thread-core";
 import type {
@@ -20,6 +21,8 @@ import {
   createThreadMappingId,
   getThreadData,
   normalizeCursor,
+  reconcileInitializedThread,
+  promoteNewThreadReducer,
   updateStatusReducer,
   preserveMidLoadTransitions,
   seedNewThread,
@@ -32,7 +35,17 @@ import type {
 } from "../../runtimes/remote-thread-list/types";
 import { ThreadListAdapterChangedError } from "../../runtimes/remote-thread-list/adapter-changed";
 import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
-import { isTitleSourceMessage } from "../../runtimes/remote-thread-list/title";
+import {
+  applyTitleStream,
+  isTitleSourceMessage,
+} from "../../runtimes/remote-thread-list/title";
+import {
+  clearThreadTitleState,
+  finishThreadTitleRename,
+  runThreadTitleGeneration,
+  startThreadTitleRename,
+  type ThreadTitleState,
+} from "../../runtimes/remote-thread-list/title-generation";
 import {
   type ComponentType,
   type FC,
@@ -42,10 +55,10 @@ import {
   useId,
 } from "react";
 import { useAui } from "@assistant-ui/store";
-import { AssistantMessageStream } from "assistant-stream";
 import type { ModelContextProvider } from "../../model-context/types";
 import { RuntimeAdapterProvider } from "./RuntimeAdapterProvider";
 import { useStableRuntimeAdapters } from "./useRuntimeAdapters";
+import { invokeUserCallback } from "../../utils/invoke-user-callback";
 
 const threadNotFoundError = (threadIdOrRemoteId: string, action: string) =>
   new Error(`Thread "${threadIdOrRemoteId}" not found while ${action}.`);
@@ -77,6 +90,7 @@ export class RemoteThreadListThreadListRuntimeCore
   private _staleThreadIdsOnReplace: ReadonlySet<string> | undefined;
   private _switchGeneration = 0;
   private _switchTask: Promise<void> | undefined;
+  private readonly _titleStates = new Map<string, ThreadTitleState>();
 
   private _mainThreadId!: string;
   private readonly _state = new OptimisticState<RemoteThreadState>(
@@ -137,6 +151,7 @@ export class RemoteThreadListThreadListRuntimeCore
             return {
               ...state,
               isLoading: true,
+              loadError: undefined,
             };
           },
           then: (state, l) => {
@@ -146,7 +161,7 @@ export class RemoteThreadListThreadListRuntimeCore
               this._replaceListOnNextLoad = false;
               replacedList = true;
               return this._replaceWithThreads(
-                state,
+                { ...state, loadError: undefined },
                 l.threads,
                 normalizeCursor(l.nextCursor),
               );
@@ -161,6 +176,7 @@ export class RemoteThreadListThreadListRuntimeCore
             const merged = {
               ...state,
               isLoading: false,
+              loadError: undefined,
               cursor: normalizeCursor(l.nextCursor),
               threadIds: fresh.threadIds,
               archivedThreadIds: fresh.archivedThreadIds,
@@ -178,13 +194,18 @@ export class RemoteThreadListThreadListRuntimeCore
             this._state.update({
               ...this._state.baseValue,
               isLoading: false,
+              loadError: error,
             });
             return;
           }
           this._replaceListOnNextLoad = false;
           replacedList = true;
           this._state.update(
-            this._replaceWithThreads(this._state.baseValue, [], undefined),
+            this._replaceWithThreads(
+              { ...this._state.baseValue, loadError: error },
+              [],
+              undefined,
+            ),
           );
         })
         .then(() => {
@@ -320,7 +341,9 @@ export class RemoteThreadListThreadListRuntimeCore
       this._state.update({
         ...this._state.baseValue,
         cursor: undefined,
+        loadError: undefined,
       });
+      this._titleStates.clear();
     }
 
     if (controlledThreadIdChanged) {
@@ -347,8 +370,13 @@ export class RemoteThreadListThreadListRuntimeCore
   ): RemoteThreadState {
     const carried: RemoteThreadData[] = [];
     if (state.newThreadId) {
-      const mappingId = state.threadIdMap[state.newThreadId];
-      const draft = mappingId ? state.threadData[mappingId] : undefined;
+      const mappingId = Object.hasOwn(state.threadIdMap, state.newThreadId)
+        ? state.threadIdMap[state.newThreadId]
+        : undefined;
+      const draft =
+        mappingId && Object.hasOwn(state.threadData, mappingId)
+          ? state.threadData[mappingId]
+          : undefined;
       if (draft?.status === "new") carried.push(draft);
     }
 
@@ -365,12 +393,12 @@ export class RemoteThreadListThreadListRuntimeCore
     const seed: ClassifyAccumulator = {
       threadIds: [],
       archivedThreadIds: [],
-      threadIdMap: {},
-      threadData: {},
+      threadIdMap: nullProtoRecord(),
+      threadData: nullProtoRecord(),
     };
     for (const item of carried) {
       const mappingId = createThreadMappingId(item.id);
-      if (seed.threadData[mappingId]) continue;
+      if (Object.hasOwn(seed.threadData, mappingId)) continue;
       seed.threadIdMap[item.id] = mappingId;
       if (item.remoteId !== undefined) {
         seed.threadIdMap[item.remoteId] = mappingId;
@@ -383,7 +411,10 @@ export class RemoteThreadListThreadListRuntimeCore
 
     for (const item of carried) {
       if (item.remoteId === undefined) continue;
-      const current = threadData[createThreadMappingId(item.id)];
+      const currentMappingId = createThreadMappingId(item.id);
+      const current = Object.hasOwn(threadData, currentMappingId)
+        ? threadData[currentMappingId]
+        : undefined;
       if (current === undefined) continue;
       if (current.status === "regular" && !threadIds.includes(current.id)) {
         threadIds.push(current.id);
@@ -405,7 +436,7 @@ export class RemoteThreadListThreadListRuntimeCore
       threadData,
       newThreadId:
         state.newThreadId !== undefined &&
-        threadIdMap[state.newThreadId] === undefined
+        !Object.hasOwn(threadIdMap, state.newThreadId)
           ? undefined
           : state.newThreadId,
     };
@@ -501,6 +532,10 @@ export class RemoteThreadListThreadListRuntimeCore
     return this._state.value.isLoading;
   }
 
+  public get loadError() {
+    return this._state.value.loadError;
+  }
+
   public get isLoadingMore() {
     return this._state.value.isLoadingMore;
   }
@@ -539,7 +574,12 @@ export class RemoteThreadListThreadListRuntimeCore
     if (this._lastNotifiedThreadId === threadId) return;
     this._lastNotifiedThreadId = threadId;
     if (emit) {
-      this._options.onThreadIdChange?.(threadId);
+      invokeUserCallback(
+        "assistant-ui",
+        "onThreadIdChange",
+        this._options.onThreadIdChange,
+        threadId,
+      );
     }
   }
 
@@ -778,71 +818,40 @@ export class RemoteThreadListThreadListRuntimeCore
       return { remoteId, externalId };
     }
 
+    this._requireAdapterGeneration(adapterGeneration);
+    const initializeTask = adapter.initialize(threadId);
+    let removedMappingId: string | undefined;
     const { remoteId, externalId } = await this._state.optimisticUpdate({
-      execute: () => {
-        this._requireAdapterGeneration(adapterGeneration);
-        return adapter.initialize(threadId);
-      },
-      optimistic: (state) => {
-        return updateStatusReducer(state, threadId, "regular");
-      },
-      loading: (state, task) => {
-        const mappingId = createThreadMappingId(threadId);
-        return {
-          ...state,
-          threadData: {
-            ...state.threadData,
-            [mappingId]: {
-              ...state.threadData[mappingId],
-              initializeTask: task,
-            },
-          },
-        };
-      },
+      execute: () => initializeTask,
+      optimistic: (state) =>
+        promoteNewThreadReducer(state, threadId, initializeTask),
       then: (state, { remoteId, externalId }) => {
         if (adapterGeneration !== this._adapterGeneration) return state;
-        const data = getThreadData(state, threadId);
-        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 }),
+        const reconciliation = reconcileInitializedThread(
+          state,
+          threadId,
           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,
-        };
+          threadId,
+        );
+        removedMappingId = reconciliation.removedMappingId;
+        if (removedMappingId === this._mainThreadId) {
+          this._mainThreadId = reconciliation.survivorMappingId;
+        }
+        return reconciliation.state;
       },
     });
     this._requireAdapterGeneration(adapterGeneration);
+    if (removedMappingId !== undefined) {
+      this._hookManager.stopThreadRuntime(removedMappingId);
+    }
     return { remoteId, externalId };
   };
 
-  public generateTitle = async (threadId: string) => {
+  public generateTitle = async (
+    threadId: string,
+    options?: { automatic?: boolean },
+  ) => {
     this._requireAdapterSettled();
     const adapter = this._options.adapter;
     const adapterGeneration = this._adapterGeneration;
@@ -861,30 +870,40 @@ export class RemoteThreadListThreadListRuntimeCore
     // would make the payload race-dependent; the title reads settled
     // messages only, matching the trigger's readiness gate.
     const messages = runtimeCore.messages.filter(isTitleSourceMessage);
-    const stream = await adapter.generateTitle(remoteId, messages);
-    const messageStream = AssistantMessageStream.fromAssistantStream(stream);
-    for await (const result of messageStream) {
-      if (adapterGeneration !== this._adapterGeneration) return;
-      const newTitle = result.parts.filter((c) => c.type === "text")[0]?.text;
-      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,
+    await runThreadTitleGeneration({
+      states: this._titleStates,
+      threadId: data.id,
+      automatic: options?.automatic === true,
+      generate: async (onTitle) => {
+        const stream = await adapter.generateTitle(remoteId, messages);
+        this._requireAdapterGeneration(adapterGeneration);
+        await applyTitleStream(stream, onTitle);
+      },
+      rename: async (title) => {
+        this._requireAdapterGeneration(adapterGeneration);
+        await adapter.rename(remoteId, title);
+      },
+      applyTitle: async (title) => {
+        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,
+                },
               },
-            },
-          };
-        },
-      });
-    }
+            };
+          },
+        });
+      },
+    });
   };
 
   public async rename(
@@ -899,28 +918,36 @@ export class RemoteThreadListThreadListRuntimeCore
     if (data.status === "new")
       throw threadStatusError(threadIdOrRemoteId, data.status, "be renamed");
 
-    return this._state.optimisticUpdate({
-      execute: async () => {
-        const { remoteId } = await data.initializeTask;
-        this._requireAdapterGeneration(adapterGeneration);
-        return adapter.rename(remoteId, newTitle);
-      },
-      optimistic: (state) => {
-        const data = getThreadData(state, threadIdOrRemoteId);
-        if (!data) return state;
+    const claim = startThreadTitleRename(this._titleStates, data.id, newTitle);
+    try {
+      const result = await this._state.optimisticUpdate({
+        execute: async () => {
+          const { remoteId } = await data.initializeTask;
+          this._requireAdapterGeneration(adapterGeneration);
+          return adapter.rename(remoteId, newTitle);
+        },
+        optimistic: (state) => {
+          const currentData = getThreadData(state, threadIdOrRemoteId);
+          if (!currentData) return state;
 
-        return {
-          ...state,
-          threadData: {
-            ...state.threadData,
-            [data.id]: {
-              ...data,
-              title: newTitle,
+          return {
+            ...state,
+            threadData: {
+              ...state.threadData,
+              [currentData.id]: {
+                ...currentData,
+                title: newTitle,
+              },
             },
-          },
-        };
-      },
-    });
+          };
+        },
+      });
+      finishThreadTitleRename(this._titleStates, data.id, claim, true);
+      return result;
+    } catch (error) {
+      finishThreadTitleRename(this._titleStates, data.id, claim, false);
+      throw error;
+    }
   }
 
   public async updateCustom(
@@ -1063,9 +1090,7 @@ export class RemoteThreadListThreadListRuntimeCore
 
     await this._ensureThreadIsNotMain(data.id);
     this._requireAdapterGeneration(adapterGeneration);
-    this._hookManager.stopThreadRuntime(data.id);
-
-    return this._state.optimisticUpdate({
+    const result = await this._state.optimisticUpdate({
       execute: async () => {
         const { remoteId } = await data.initializeTask;
         this._requireAdapterGeneration(adapterGeneration);
@@ -1075,6 +1100,12 @@ export class RemoteThreadListThreadListRuntimeCore
         return updateStatusReducer(state, data.id, "deleted");
       },
     });
+    // The optimistic layer survives an adapter swap, so a resolved deletion has
+    // dropped the slot from `threadData`, where `_replaceWithThreads` would
+    // otherwise have found it to stop.
+    this._hookManager.stopThreadRuntime(data.id);
+    clearThreadTitleState(this._titleStates, data.id);
+    return result;
   }
 
   public __internal_dispose() {

package/src/react/runtimes/RemoteThreadResource.test.ts

@@ -53,6 +53,9 @@ describe("subscribeToTitleGeneration", () => {
 
     await vi.waitFor(() => {
       expect(itemRuntime.generateTitle).toHaveBeenCalledOnce();
+      expect(itemRuntime.generateTitle).toHaveBeenCalledWith({
+        automatic: true,
+      });
       expect(consoleError).toHaveBeenCalledWith(
         "[assistant-ui] Thread title generation failed",
         error,
@@ -97,6 +100,9 @@ describe("subscribeToTitleGeneration", () => {
 
     await vi.waitFor(() => {
       expect(itemRuntime.generateTitle).toHaveBeenCalledOnce();
+      expect(itemRuntime.generateTitle).toHaveBeenCalledWith({
+        automatic: true,
+      });
     });
   });
 
@@ -136,6 +142,9 @@ describe("subscribeToTitleGeneration", () => {
 
     await vi.waitFor(() => {
       expect(itemRuntime.generateTitle).toHaveBeenCalledOnce();
+      expect(itemRuntime.generateTitle).toHaveBeenCalledWith({
+        automatic: true,
+      });
     });
   });
 });

package/src/react/runtimes/RemoteThreadResource.ts

@@ -3,6 +3,7 @@ import { resource } from "@assistant-ui/tap";
 import {
   useAssistantContextProvider,
   useConfiguredAui,
+  useDestroySignalProvider,
 } from "@assistant-ui/store/client";
 import { ThreadListItemClient } from "../../store/internal";
 import type { AssistantRuntime } from "../../runtime/api/assistant-runtime";
@@ -38,6 +39,7 @@ export type RemoteThreadResourceProps = {
     runtime: ThreadRuntimeCore,
     generation: number,
   ) => void;
+  destroySignal: AbortSignal;
 };
 
 export const subscribeToTitleGeneration = (
@@ -45,10 +47,12 @@ export const subscribeToTitleGeneration = (
   itemRuntime: ThreadListItemRuntime,
 ) => {
   const generate = () =>
-    void itemRuntime.generateTitle().catch((error: unknown) => {
-      if (isSilentRuntimeAction(error)) return;
-      console.error("[assistant-ui] Thread title generation failed", error);
-    });
+    void itemRuntime
+      .generateTitle({ automatic: true })
+      .catch((error: unknown) => {
+        if (isSilentRuntimeAction(error)) return;
+        console.error("[assistant-ui] Thread title generation failed", error);
+      });
 
   const hasTitleSource = () =>
     threadRuntime.getState().messages.some(isTitleSourceMessage);
@@ -178,6 +182,7 @@ const useRemoteThreadResource = ({
   parentClient,
   adapters,
   publish,
+  destroySignal,
 }: RemoteThreadResourceProps) => {
   const itemRuntime = useMemo(
     () =>
@@ -214,14 +219,16 @@ const useRemoteThreadResource = ({
     threadListItem: ThreadListItemClient({ runtime: itemRuntime }),
   });
 
-  return useAssistantContextProvider(client, function useThreadClient() {
-    return useRuntimeAdaptersProvider(adapters, function useBoundRuntime() {
-      return useRemoteThreadBinder({
-        threadId,
-        generation,
-        runtimeHook,
-        publish,
-        itemRuntime,
+  return useDestroySignalProvider(destroySignal, function useOwnedThread() {
+    return useAssistantContextProvider(client, function useThreadClient() {
+      return useRuntimeAdaptersProvider(adapters, function useBoundRuntime() {
+        return useRemoteThreadBinder({
+          threadId,
+          generation,
+          runtimeHook,
+          publish,
+          itemRuntime,
+        });
       });
     });
   });

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

@@ -28,8 +28,8 @@ describe("extractAuiV0", () => {
   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);
+    expect(result?.usage?.inputTokens).toBe(1);
+    expect(result?.usage?.outputTokens).toBe(2);
   });
 
   it("sums step usage reported under the AI SDK v7 token details", () => {
@@ -57,22 +57,28 @@ describe("extractAuiV0", () => {
       },
     });
 
-    expect(result?.inputTokens).toBe(15);
-    expect(result?.outputTokens).toBe(6);
-    expect(result?.cachedInputTokens).toBe(9);
-    expect(result?.reasoningTokens).toBe(2);
+    expect(result?.usage).toEqual({
+      inputTokens: 15,
+      outputTokens: 6,
+      cachedInputTokens: 9,
+      reasoningTokens: 2,
+    });
     expect(result?.steps).toEqual([
       {
-        input_tokens: 10,
-        output_tokens: 4,
-        reasoning_tokens: 1,
-        cached_input_tokens: 6,
+        usage: {
+          inputTokens: 10,
+          outputTokens: 4,
+          inputTokenDetails: { cacheReadTokens: 6 },
+          outputTokenDetails: { reasoningTokens: 1 },
+        },
       },
       {
-        input_tokens: 5,
-        output_tokens: 2,
-        reasoning_tokens: 1,
-        cached_input_tokens: 3,
+        usage: {
+          inputTokens: 5,
+          outputTokens: 2,
+          inputTokenDetails: { cacheReadTokens: 3 },
+          outputTokenDetails: { reasoningTokens: 1 },
+        },
       },
     ]);
   });