@assistant-ui/core 0.3.14 → 0.3.16

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 () => ({
@@ -542,6 +638,202 @@ describe("RemoteThreadList", () => {
     handle.destroy();
   });
 
+  it("resets selection when reload follows a different adapter instance", async () => {
+    const methodsA = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          { status: "regular" as const, remoteId: "thread-a", title: "A" },
+        ],
+      })),
+    });
+    const methodsB = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          { status: "regular" as const, remoteId: "thread-b", title: "B" },
+        ],
+      })),
+    });
+    let adapter: RemoteThreadListAdapter = { ...methodsA };
+    const listeners = new Set<() => void>();
+    const source: AssistantConfigSource = {
+      getConfig: () =>
+        AuiConfig({
+          threads: RemoteThreadList({
+            adapter,
+            thread: (id) => StubThread({ threadId: id }) as never,
+          }),
+        }),
+      subscribe: (listener) => {
+        listeners.add(listener);
+        return () => {
+          listeners.delete(listener);
+        };
+      },
+    };
+    const handle = createAssistantClient(source);
+    handle.subscribe(() => {});
+    await handle.getClient().threads.getLoadThreadsPromise();
+    flushTapSync(() => handle.getClient().threads.switchToThread("thread-a"));
+    await vi.waitFor(() => {
+      expect(handle.getClient().threads.item("main").getState().remoteId).toBe(
+        "thread-a",
+      );
+    });
+
+    adapter = { ...methodsB };
+    for (const listener of listeners) listener();
+    await vi.waitFor(async () => {
+      flushTapSync(() => {});
+      await handle.getClient().threads.reload();
+      expect(handle.getClient().threads.getState().threadIds).toEqual([
+        "thread-b",
+      ]);
+    });
+    expect(
+      handle.getClient().threads.item("main").getState().remoteId,
+    ).not.toBe("thread-a");
+    expect(handle.getClient().threads.item("main").getState().status).toBe(
+      "new",
+    );
+    expect(methodsB.list).toHaveBeenCalled();
+    handle.destroy();
+  });
+
+  it("does not reset again when retrying a failed replacement load", async () => {
+    const methodsA = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          { status: "regular" as const, remoteId: "thread-a", title: "A" },
+        ],
+      })),
+    });
+    const methodsB = makeAdapter({
+      list: vi
+        .fn()
+        .mockRejectedValueOnce(new Error("network"))
+        .mockResolvedValueOnce({
+          threads: [
+            { status: "regular" as const, remoteId: "thread-b", title: "B" },
+          ],
+        }),
+    });
+    let adapter: RemoteThreadListAdapter = { ...methodsA };
+    const listeners = new Set<() => void>();
+    const source: AssistantConfigSource = {
+      getConfig: () =>
+        AuiConfig({
+          threads: RemoteThreadList({
+            adapter,
+            thread: (id) => StubThread({ threadId: id }) as never,
+          }),
+        }),
+      subscribe: (listener) => {
+        listeners.add(listener);
+        return () => {
+          listeners.delete(listener);
+        };
+      },
+    };
+    const handle = createAssistantClient(source);
+    handle.subscribe(() => {});
+    await handle.getClient().threads.getLoadThreadsPromise();
+    flushTapSync(() => handle.getClient().threads.switchToThread("thread-a"));
+    await vi.waitFor(() => {
+      expect(handle.getClient().threads.item("main").getState().remoteId).toBe(
+        "thread-a",
+      );
+    });
+
+    adapter = { ...methodsB };
+    for (const listener of listeners) listener();
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    try {
+      await vi.waitFor(async () => {
+        flushTapSync(() => {});
+        await handle.getClient().threads.reload();
+        expect(handle.getClient().threads.item("main").getState().status).toBe(
+          "new",
+        );
+      });
+      const mainAfterFailure = handle
+        .getClient()
+        .threads.getState().mainThreadId;
+      await handle.getClient().threads.reload();
+      await vi.waitFor(() => {
+        expect(handle.getClient().threads.getState().threadIds).toEqual([
+          "thread-b",
+        ]);
+      });
+      expect(handle.getClient().threads.getState().mainThreadId).toBe(
+        mainAfterFailure,
+      );
+    } finally {
+      consoleError.mockRestore();
+    }
+    handle.destroy();
+  });
+
+  it("does not send a superseded rename through the replacement adapter", async () => {
+    const initializeRequest = deferred<{
+      remoteId: string;
+      externalId: string | undefined;
+    }>();
+    const methodsA = makeAdapter({
+      list: vi.fn(async () => ({ threads: [] })),
+      initialize: vi.fn(() => initializeRequest.promise),
+    });
+    const methodsB = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          { status: "regular" as const, remoteId: "thread-b", title: "B" },
+        ],
+      })),
+    });
+    let adapter: RemoteThreadListAdapter = { ...methodsA };
+    const listeners = new Set<() => void>();
+    const source: AssistantConfigSource = {
+      getConfig: () =>
+        AuiConfig({
+          threads: RemoteThreadList({
+            adapter,
+            thread: (id) => StubThread({ threadId: id }) as never,
+          }),
+        }),
+      subscribe: (listener) => {
+        listeners.add(listener);
+        return () => {
+          listeners.delete(listener);
+        };
+      },
+    };
+    const handle = createAssistantClient(source);
+    handle.subscribe(() => {});
+    await handle.getClient().threads.getLoadThreadsPromise();
+    const initializeTask = handle.getClient().threads.item("main").initialize();
+    const renameTask = handle.getClient().threads.item("main").rename("leaked");
+
+    adapter = { ...methodsB };
+    for (const listener of listeners) listener();
+    await vi.waitFor(async () => {
+      flushTapSync(() => {});
+      await handle.getClient().threads.reload();
+      expect(handle.getClient().threads.getState().threadIds).toEqual([
+        "thread-b",
+      ]);
+    });
+    initializeRequest.resolve({
+      remoteId: "leaked",
+      externalId: undefined,
+    });
+    await expect(initializeTask).rejects.toThrow("adapter changed");
+    await expect(renameTask).rejects.toThrow("adapter changed");
+    expect(methodsA.rename).not.toHaveBeenCalled();
+    expect(methodsB.rename).not.toHaveBeenCalled();
+    handle.destroy();
+  });
+
   it("exposes useAdapters adapters to the thread factory", async () => {
     const history = dummyHistory();
     const capture: { adapters: RuntimeAdapters | null } = { adapters: null };