@assistant-ui/core 0.3.15 → 0.3.16

package/src/model-context/registry.ts

@@ -2,12 +2,12 @@ import type { Tool } from "assistant-stream";
 import {
   type ModelContext,
   type ModelContextProvider,
-  mergeModelContexts,
   type AssistantToolProps,
   type AssistantInstructionsConfig,
 } from "./types";
-import type { Unsubscribe } from "../types/unsubscribe";
 import { notifySubscribers as notifyStateSubscribers } from "../subscribable/subscribable";
+import { CompositeContextProvider } from "../utils/composite-context-provider";
+import type { Unsubscribe } from "../types/unsubscribe";
 import type {
   ModelContextRegistryToolHandle,
   ModelContextRegistryInstructionHandle,
@@ -17,9 +17,14 @@ import type {
 export class ModelContextRegistry implements ModelContextProvider {
   private _tools = new Map<symbol, AssistantToolProps<any, any>>();
   private _instructions = new Map<symbol, string>();
-  private _providers = new Map<symbol, ModelContextProvider>();
+  private _contextProviders = new CompositeContextProvider();
   private _subscribers = new Set<() => void>();
-  private _providerUnsubscribes = new Map<symbol, Unsubscribe | undefined>();
+
+  constructor() {
+    this._contextProviders.subscribe(() => {
+      this.notifySubscribers();
+    });
+  }
 
   getModelContext(): ModelContext {
     const instructions = Array.from(this._instructions.values()).filter(
@@ -35,9 +40,7 @@ export class ModelContextRegistry implements ModelContextProvider {
       tools[toolName] = tool;
     }
 
-    const providerContexts = mergeModelContexts(
-      new Set(this._providers.values()),
-    );
+    const providerContexts = this._contextProviders.getModelContext();
 
     const context: ModelContext = {
       system,
@@ -139,36 +142,11 @@ export class ModelContextRegistry implements ModelContextProvider {
   addProvider(
     provider: ModelContextProvider,
   ): ModelContextRegistryProviderHandle {
-    const id = Symbol();
-
-    this._providers.set(id, provider);
-
-    let unsubscribe: Unsubscribe | undefined;
-    try {
-      unsubscribe = provider.subscribe?.(() => {
-        this.notifySubscribers();
-      });
-    } catch (error) {
-      this._providers.delete(id);
-      try {
-        this.notifySubscribers();
-      } catch (notifyError) {
-        console.error(notifyError);
-      }
-      throw error;
-    }
-    this._providerUnsubscribes.set(id, unsubscribe);
-
-    this.notifySubscribers();
+    const unregister =
+      this._contextProviders.registerModelContextProvider(provider);
 
     return {
-      remove: () => {
-        this._providers.delete(id);
-        const unsubscribe = this._providerUnsubscribes.get(id);
-        unsubscribe?.();
-        this._providerUnsubscribes.delete(id);
-        this.notifySubscribers();
-      },
+      remove: unregister,
     };
   }
 }

package/src/react/client/InMemoryThreadList.test.tsx

@@ -146,3 +146,52 @@ describe("InMemoryThreadList delete", () => {
     expect(state.mainThreadId).not.toBe("main");
   });
 });
+
+describe("InMemoryThreadList item index selectors", () => {
+  it("resolves index selectors within the archived and regular subsets", async () => {
+    const { getAui } = setup();
+    await act(async () => {});
+
+    await act(async () => {
+      getAui().threads.switchToNewThread();
+    });
+    await act(async () => {});
+    const b = getAui().threads.getState().mainThreadId;
+
+    await act(async () => {
+      getAui().threads.switchToNewThread();
+    });
+    await act(async () => {});
+
+    await act(async () => {
+      getAui().threads.item({ id: b }).archive();
+    });
+    await act(async () => {});
+
+    const state = getAui().threads.getState();
+    expect(state.archivedThreadIds).toEqual([b]);
+    expect(state.threadIds).toHaveLength(2);
+
+    for (const [index, id] of state.archivedThreadIds.entries()) {
+      expect(
+        getAui().threads.item({ index, archived: true }).getState().id,
+      ).toBe(id);
+    }
+    for (const [index, id] of state.threadIds.entries()) {
+      expect(getAui().threads.item({ index }).getState().id).toBe(id);
+      expect(
+        getAui().threads.item({ index, archived: false }).getState().id,
+      ).toBe(id);
+    }
+
+    expect(() =>
+      getAui().threads.item({ index: state.threadIds.length }),
+    ).toThrow("out of bounds");
+    expect(() =>
+      getAui().threads.item({
+        index: state.archivedThreadIds.length,
+        archived: true,
+      }),
+    ).toThrow("out of bounds");
+  });
+});

package/src/react/client/InMemoryThreadList.ts

@@ -236,7 +236,13 @@ const useInMemoryThreadList = (
         const index = threads.findIndex((t) => t.id === selector.id);
         return threadListItems.get({ index });
       }
-      return threadListItems.get(selector);
+      // The lookup is keyed over the combined thread array, while index
+      // selectors address the regular/archived subset the primitives render.
+      const ids = selector.archived ? state.archivedThreadIds : state.threadIds;
+      const id = ids[selector.index];
+      if (id === undefined) return threadListItems.get({ index: -1 });
+      const index = threads.findIndex((t) => t.id === id);
+      return threadListItems.get({ index });
     },
     thread: () => mainThreadClient.methods,
   };

package/src/react/client/Interactables.ts

@@ -24,8 +24,7 @@ import {
   interactableToolName,
 } from "../../model-context/interactable-composer-metadata";
 import { notifySubscribers as notifyStateSubscribers } from "../../subscribable/subscribable";
-
-const PERSISTENCE_DEBOUNCE_MS = 500;
+import { useInteractablePersistenceQueue } from "../interactables-shared/useInteractablePersistenceQueue";
 
 type RestorePersistedStateOptions = {
   stash: Map<string, unknown>;
@@ -109,26 +108,6 @@ const useInteractablesResource = ({
   const adapterRef = useRef<
     Unstable_InteractablePersistenceAdapter | undefined
   >(undefined);
-  const debounceTimerRef = useRef<ReturnType<typeof setTimeout> | undefined>(
-    undefined,
-  );
-  const syncSeqRef = useRef(0);
-  const latestSyncSeqByIdRef = useRef(new Map<string, number>());
-  const inFlightPersistenceRef = useRef(0);
-  const flushResolversRef = useRef<Array<() => void>>([]);
-  const dirtyIdsRef = useRef(new Set<string>());
-
-  type PersistenceBatch = {
-    adapter: Unstable_InteractablePersistenceAdapter;
-    payload: Unstable_InteractablePersistedState;
-    dirtyIds: Set<string>;
-    seq: number;
-  };
-
-  const outgoingQueueRef = useRef<PersistenceBatch[]>([]);
-  const runPersistenceRef = useRef<(batch?: PersistenceBatch) => void>(
-    () => {},
-  );
 
   const setStateAndRef = useCallback(
     (
@@ -152,149 +131,28 @@ const useInteractablesResource = ({
     return result;
   }, []);
 
-  const takeDirtyBatch = useCallback(
+  const updatePersistenceStatus = useCallback(
     (
-      adapter: Unstable_InteractablePersistenceAdapter,
-    ): PersistenceBatch | undefined => {
-      if (dirtyIdsRef.current.size === 0) return;
-      const dirtyIds = new Set(dirtyIdsRef.current);
-      dirtyIdsRef.current.clear();
-      const seq = ++syncSeqRef.current;
-      for (const id of dirtyIds) latestSyncSeqByIdRef.current.set(id, seq);
-      return { adapter, payload: exportState(), dirtyIds, seq };
-    },
-    [exportState],
-  );
-
-  const enqueuePersistence = useCallback(
-    (adapter: Unstable_InteractablePersistenceAdapter) => {
-      const batch = takeDirtyBatch(adapter);
-      if (!batch) return;
-      if (inFlightPersistenceRef.current === 0) {
-        runPersistenceRef.current(batch);
-      } else {
-        outgoingQueueRef.current.push(batch);
-      }
-    },
-    [takeDirtyBatch],
-  );
-
-  const runPersistence = useCallback(
-    async (batch?: PersistenceBatch) => {
-      const resolved =
-        batch ??
-        (adapterRef.current ? takeDirtyBatch(adapterRef.current) : undefined);
-      if (!resolved) {
-        if (inFlightPersistenceRef.current === 0) {
-          for (const resolve of flushResolversRef.current) resolve();
-          flushResolversRef.current = [];
-        }
-        return;
-      }
-
-      const { adapter, payload, dirtyIds, seq } = resolved;
-      inFlightPersistenceRef.current += 1;
-
-      setStateAndRef((prev) => ({
-        ...prev,
-        persistence: {
-          ...prev.persistence,
-          ...Object.fromEntries(
-            [...dirtyIds].map((id) => [
-              id,
-              { isPending: true, error: undefined },
-            ]),
-          ),
-        },
-      }));
-
-      try {
-        await adapter.save(payload);
-        setStateAndRef((prev) => {
-          let changed = false;
-          const persistence = { ...prev.persistence };
-          for (const id of dirtyIds) {
-            if (
-              latestSyncSeqByIdRef.current.get(id) !== seq ||
-              dirtyIdsRef.current.has(id)
-            )
-              continue;
-            latestSyncSeqByIdRef.current.delete(id);
-            delete persistence[id];
-            changed = true;
-          }
-          return changed ? { ...prev, persistence } : prev;
-        });
-      } catch (e) {
-        setStateAndRef((prev) => {
-          let changed = false;
-          const persistence = { ...prev.persistence };
-          for (const id of dirtyIds) {
-            if (
-              latestSyncSeqByIdRef.current.get(id) !== seq ||
-              dirtyIdsRef.current.has(id)
-            )
-              continue;
-            latestSyncSeqByIdRef.current.delete(id);
-            persistence[id] = { isPending: false, error: e };
-            changed = true;
-          }
-          return changed ? { ...prev, persistence } : prev;
-        });
-      } finally {
-        inFlightPersistenceRef.current -= 1;
-        const next =
-          outgoingQueueRef.current.shift() ??
-          (adapterRef.current && dirtyIdsRef.current.size > 0
-            ? takeDirtyBatch(adapterRef.current)
-            : undefined);
-        if (next) {
-          if (debounceTimerRef.current !== undefined) {
-            clearTimeout(debounceTimerRef.current);
-            debounceTimerRef.current = undefined;
-          }
-          runPersistenceRef.current(next);
-        } else if (inFlightPersistenceRef.current === 0) {
-          for (const resolve of flushResolversRef.current) resolve();
-          flushResolversRef.current = [];
-        }
-      }
+      updater: (
+        prev: Unstable_InteractablesState["persistence"],
+      ) => Unstable_InteractablesState["persistence"],
+    ) => {
+      setStateAndRef((prev) => {
+        const persistence = updater(prev.persistence);
+        return persistence === prev.persistence
+          ? prev
+          : { ...prev, persistence };
+      });
     },
-    [setStateAndRef, takeDirtyBatch],
+    [setStateAndRef],
   );
-  runPersistenceRef.current = (nextBatch) => {
-    void runPersistence(nextBatch);
-  };
 
-  const flushIfPending = useCallback(() => {
-    if (debounceTimerRef.current !== undefined) {
-      clearTimeout(debounceTimerRef.current);
-      debounceTimerRef.current = undefined;
-    }
-    if (adapterRef.current) enqueuePersistence(adapterRef.current);
-  }, [enqueuePersistence]);
-
-  const schedulePersistence = useCallback(
-    (id: string) => {
-      if (!adapterRef.current) return;
-      dirtyIdsRef.current.add(id);
-      if (debounceTimerRef.current !== undefined) {
-        clearTimeout(debounceTimerRef.current);
-      }
-      debounceTimerRef.current = setTimeout(() => {
-        debounceTimerRef.current = undefined;
-        if (inFlightPersistenceRef.current === 0 && adapterRef.current) {
-          enqueuePersistence(adapterRef.current);
-        } else {
-          debounceTimerRef.current = setTimeout(() => {
-            debounceTimerRef.current = undefined;
-            if (adapterRef.current) enqueuePersistence(adapterRef.current);
-          }, PERSISTENCE_DEBOUNCE_MS);
-        }
-      }, PERSISTENCE_DEBOUNCE_MS);
-    },
-    [enqueuePersistence],
-  );
+  const { flushIfPending, schedulePersistence, flush } =
+    useInteractablePersistenceQueue({
+      adapterRef,
+      snapshot: exportState,
+      updatePersistenceStatus,
+    });
 
   const restorePersistedState = useCallback(
     (
@@ -394,23 +252,6 @@ const useInteractablesResource = ({
     };
   }, [persistence, setPersistenceAdapter]);
 
-  const flush = useCallback(async () => {
-    if (debounceTimerRef.current !== undefined) {
-      clearTimeout(debounceTimerRef.current);
-      debounceTimerRef.current = undefined;
-    }
-    const hasWork =
-      inFlightPersistenceRef.current > 0 ||
-      dirtyIdsRef.current.size > 0 ||
-      outgoingQueueRef.current.length > 0;
-    if (!hasWork) return;
-    const p = new Promise<void>((resolve) => {
-      flushResolversRef.current.push(resolve);
-    });
-    if (adapterRef.current) enqueuePersistence(adapterRef.current);
-    return p;
-  }, [enqueuePersistence]);
-
   const setDefState = useCallback(
     (id: string, updater: (prev: unknown) => unknown) => {
       touchedIdsRef.current.add(id);

package/src/react/client/RemoteThreadList.test.ts

@@ -234,6 +234,37 @@ describe("RemoteThreadList", () => {
     handle.destroy();
   });
 
+  it("keeps a thread initialized during the list() flight after switching away", async () => {
+    const listDeferred =
+      deferred<Awaited<ReturnType<RemoteThreadListAdapter["list"]>>>();
+    const adapter = makeAdapter({
+      list: vi.fn(() => listDeferred.promise),
+    });
+    const { handle } = mountList(adapter);
+    const aui = handle.getClient();
+
+    const loadPromise = aui.threads.getLoadThreadsPromise();
+    const initializedId = aui.threads.getState().mainThreadId;
+    await aui.threads.item("main").initialize();
+    await vi.waitFor(() => {
+      expect(aui.threads.getState().threadIds).toContain(initializedId);
+    });
+
+    listDeferred.resolve({
+      threads: [{ status: "regular" as const, remoteId: "t1", title: "One" }],
+    });
+    await loadPromise;
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    await vi.waitFor(() => {
+      expect(aui.threads.getState().mainThreadId).not.toBe(initializedId);
+    });
+
+    expect(aui.threads.getState().threadIds[0]).toBe(initializedId);
+    expect(aui.threads.getState().threadIds).toContain("t1");
+    handle.destroy();
+  });
+
   it("switches to a listed thread and back to a new thread", async () => {
     const adapter = makeAdapter({
       list: vi.fn(async () => ({
@@ -331,6 +362,71 @@ describe("RemoteThreadList", () => {
     handle.destroy();
   });
 
+  it("preserves generated titles across an overlapping reload", async () => {
+    const reload = deferred<{
+      threads: {
+        status: "regular";
+        remoteId: string;
+        title: string;
+      }[];
+    }>();
+    const adapter = makeAdapter({
+      list: vi
+        .fn()
+        .mockResolvedValueOnce({
+          threads: [
+            { status: "regular" as const, remoteId: "t1", title: "One" },
+          ],
+        })
+        .mockImplementationOnce(() => reload.promise),
+      generateTitle: vi.fn(
+        async () =>
+          new ReadableStream({
+            start(controller) {
+              controller.enqueue({
+                type: "part-start",
+                path: [],
+                part: { type: "text" },
+              });
+              controller.enqueue({
+                type: "text-delta",
+                path: [0],
+                textDelta: "Generated",
+              });
+              controller.enqueue({ type: "part-finish", path: [0] });
+              controller.close();
+            },
+          }) as never,
+      ),
+    });
+    const { handle } = mountList(adapter);
+    const aui = handle.getClient();
+    await aui.threads.getLoadThreadsPromise();
+    flushTapSync(() => aui.threads.switchToThread("t1"));
+    await vi.waitFor(() => {
+      expect(aui.threads.getState().mainThreadId).toBe("t1");
+    });
+
+    const reloadPromise = aui.threads.reload();
+    await vi.waitFor(() => {
+      expect(adapter.list).toHaveBeenCalledTimes(2);
+    });
+    await aui.threads.item({ id: "t1" }).rename("Renamed");
+    await aui.threads.item({ id: "t1" }).generateTitle();
+    await vi.waitFor(() => {
+      expect(aui.threads.item({ id: "t1" }).getState().title).toBe("Generated");
+    });
+
+    reload.resolve({
+      threads: [{ status: "regular", remoteId: "t1", title: "One" }],
+    });
+    await reloadPromise;
+    await vi.waitFor(() => {
+      expect(aui.threads.item({ id: "t1" }).getState().title).toBe("Generated");
+    });
+    handle.destroy();
+  });
+
   it("opens a controlled threadId without echoing it back", async () => {
     const adapter = makeAdapter({
       list: vi.fn(async () => ({

package/src/react/client/RemoteThreadList.ts

@@ -13,16 +13,19 @@ import {
   useClientResource,
 } from "@assistant-ui/store/client";
 import { isDevelopment, useThreadSelectionEvents } from "../../store/internal";
-import { generateId } from "../../utils/id";
 import { OptimisticState } from "../../runtimes/remote-thread-list/optimistic-state";
 import {
   classifyThreads,
+  createEmptyRemoteThreadState,
   createThreadMappingId,
   getThreadData,
   normalizeCursor,
   updateStatusReducer,
   type RemoteThreadData,
   type RemoteThreadState,
+  preserveMidLoadTransitions,
+  seedNewThread,
+  statusSnapshot,
 } from "../../runtimes/remote-thread-list/remote-thread-state";
 import type {
   RemoteThreadInitializeResponse,
@@ -40,16 +43,7 @@ import { AdaptedRemoteThread } from "./AdaptedRemoteThread";
 
 const RESOLVED_PROMISE = Promise.resolve();
 
-const EMPTY_LIST: RemoteThreadState = {
-  isLoading: true,
-  isLoadingMore: false,
-  cursor: undefined,
-  newThreadId: undefined,
-  threadIds: [],
-  archivedThreadIds: [],
-  threadIdMap: {},
-  threadData: {},
-};
+const EMPTY_LIST = createEmptyRemoteThreadState();
 
 export type RemoteThreadListProps = {
   /**
@@ -91,46 +85,14 @@ const toInitializeResult = (
 
 const applyTitleStream = async (
   stream: Parameters<typeof AssistantMessageStream.fromAssistantStream>[0],
-  onTitle: (title: string | undefined) => void,
+  onTitle: (title: string | undefined) => Promise<void>,
 ) => {
   const messageStream = AssistantMessageStream.fromAssistantStream(stream);
   for await (const result of messageStream) {
-    onTitle(result.parts.filter((part) => part.type === "text")[0]?.text);
+    await onTitle(result.parts.filter((part) => part.type === "text")[0]?.text);
   }
 };
 
-const seedNewThread = (
-  state: RemoteThreadState,
-): { id: string; 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,
-        },
-      },
-    },
-  };
-};
-
 const useThreadListItemClient = (props: {
   data: RemoteThreadData;
   isRunning: boolean;
@@ -424,6 +386,7 @@ const useRemoteThreadList = (
     if (session.loadPromise) return session.loadPromise;
     const generation = session.loadGeneration;
     const adapter = session.adapter;
+    const statusAtRequest = statusSnapshot(store.baseValue);
     session.loadPromise = store
       .optimisticUpdate({
         execute: () => adapter.list(),
@@ -437,7 +400,7 @@ const useRemoteThreadList = (
             threadIdMap: {},
             threadData: {},
           });
-          return {
+          const merged = {
             ...state,
             isLoading: false,
             cursor: normalizeCursor(page.nextCursor),
@@ -452,6 +415,7 @@ const useRemoteThreadList = (
               ...fresh.threadData,
             },
           };
+          return preserveMidLoadTransitions(merged, state, statusAtRequest);
         },
       })
       .catch((error: unknown) => {
@@ -985,19 +949,23 @@ const useRemoteThreadList = (
       if (!messages) return;
       const stream = await currentAdapter.generateTitle(remoteId, messages);
       requireAdapterGeneration(adapterGeneration);
-      await applyTitleStream(stream, (newTitle) => {
-        if (adapterGeneration !== session.adapterGeneration) return;
-        const state = store.baseValue;
-        const current = getThreadData(state, data.id);
-        if (!current) return;
-        store.update({
-          ...state,
-          threadData: {
-            ...state.threadData,
-            [current.id]: {
-              ...current,
-              title: newTitle,
-            },
+      await applyTitleStream(stream, async (newTitle) => {
+        await store.optimisticUpdate({
+          execute: async () => {},
+          optimistic: (state) => {
+            if (adapterGeneration !== session.adapterGeneration) return state;
+            const current = getThreadData(state, data.id);
+            if (!current) return state;
+            return {
+              ...state,
+              threadData: {
+                ...state.threadData,
+                [current.id]: {
+                  ...current,
+                  title: newTitle,
+                },
+              },
+            };
           },
         });
       });