@assistant-ui/core 0.3.13 → 0.3.14

package/src/react/client/Interactables.ts

@@ -4,6 +4,7 @@ import type { ClientOutput } from "@assistant-ui/store";
 import {
   attachTransformScopes,
   useAssistantClientRef,
+  useAssistantScopeEffect,
 } from "@assistant-ui/store/client";
 import type {
   Unstable_InteractablesState,
@@ -22,6 +23,7 @@ import {
   findModelKnownState,
   interactableToolName,
 } from "../../model-context/interactable-composer-metadata";
+import { notifySubscribers as notifyStateSubscribers } from "../../subscribable/subscribable";
 
 const PERSISTENCE_DEBOUNCE_MS = 500;
 
@@ -49,6 +51,12 @@ type InternalInteractableRegistration = Unstable_InteractableRegistration & {
   scope?: "thread" | undefined;
 };
 
+type UpdateToolUIEntry = {
+  count: number;
+  render: NonNullable<Unstable_InteractableRegistration["updateRender"]>;
+  unsubscribe: (() => void) | undefined;
+};
+
 const hasInteractableCreateCall = (
   messages: readonly MessageLike[],
   id: string,
@@ -92,9 +100,7 @@ const useInteractablesResource = ({
   const registrationCountsRef = useRef(new Map<string, number>());
   // One update-tool UI per interactable name, alive while any registrant
   // that supplied an updateRender is mounted.
-  const updateToolUIsRef = useRef(
-    new Map<string, { count: number; unsubscribe: () => void }>(),
-  );
+  const updateToolUIsRef = useRef(new Map<string, UpdateToolUIEntry>());
   // App-scoped state restored via adapter.load(), consumed as components register.
   const loadedStateRef = useRef(new Map<string, unknown>());
   // Ids edited locally this session — a local edit always wins over a slow load.
@@ -107,10 +113,23 @@ const useInteractablesResource = ({
     undefined,
   );
   const syncSeqRef = useRef(0);
-  const hasPendingLocalChangeRef = useRef(false);
+  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(
     (
       updater: (
@@ -133,67 +152,127 @@ const useInteractablesResource = ({
     return result;
   }, []);
 
-  const runPersistence = useCallback(async () => {
-    const adapter = adapterRef.current;
-    if (!adapter) {
-      for (const resolve of flushResolversRef.current) resolve();
-      flushResolversRef.current = [];
-      return;
-    }
+  const takeDirtyBatch = 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 seq = ++syncSeqRef.current;
-    const dirtyIds = new Set(dirtyIdsRef.current);
-    dirtyIdsRef.current.clear();
-    hasPendingLocalChangeRef.current = true;
-
-    // Snapshot before any await so unregistered definitions are still included.
-    const payload = exportState();
-
-    setStateAndRef((prev) => ({
-      ...prev,
-      persistence: {
-        ...prev.persistence,
-        ...Object.fromEntries(
-          [...dirtyIds].map((id) => [
-            id,
-            { isPending: true, error: undefined },
-          ]),
-        ),
-      },
-    }));
+  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;
+      }
 
-    try {
-      await adapter.save(payload);
-      if (syncSeqRef.current === seq) {
-        hasPendingLocalChangeRef.current = false;
+      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) delete persistence[id];
-          return { ...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 = [];
+        }
       }
-    } catch (e) {
-      if (syncSeqRef.current === seq) {
-        hasPendingLocalChangeRef.current = false;
-        setStateAndRef((prev) => ({
-          ...prev,
-          persistence: {
-            ...prev.persistence,
-            ...Object.fromEntries(
-              [...dirtyIds].map((id) => [id, { isPending: false, error: e }]),
-            ),
-          },
-        }));
-      }
-    } finally {
-      if (dirtyIdsRef.current.size > 0 && adapterRef.current) {
-        runPersistence();
-      } else {
-        for (const resolve of flushResolversRef.current) resolve();
-        flushResolversRef.current = [];
-      }
+    },
+    [setStateAndRef, takeDirtyBatch],
+  );
+  runPersistenceRef.current = (nextBatch) => {
+    void runPersistence(nextBatch);
+  };
+
+  const flushIfPending = useCallback(() => {
+    if (debounceTimerRef.current !== undefined) {
+      clearTimeout(debounceTimerRef.current);
+      debounceTimerRef.current = undefined;
     }
-  }, [exportState, setStateAndRef]);
+    if (adapterRef.current) enqueuePersistence(adapterRef.current);
+  }, [enqueuePersistence]);
 
   const schedulePersistence = useCallback(
     (id: string) => {
@@ -204,17 +283,17 @@ const useInteractablesResource = ({
       }
       debounceTimerRef.current = setTimeout(() => {
         debounceTimerRef.current = undefined;
-        if (!hasPendingLocalChangeRef.current) {
-          runPersistence();
+        if (inFlightPersistenceRef.current === 0 && adapterRef.current) {
+          enqueuePersistence(adapterRef.current);
         } else {
           debounceTimerRef.current = setTimeout(() => {
             debounceTimerRef.current = undefined;
-            runPersistence();
+            if (adapterRef.current) enqueuePersistence(adapterRef.current);
           }, PERSISTENCE_DEBOUNCE_MS);
         }
       }, PERSISTENCE_DEBOUNCE_MS);
     },
-    [runPersistence],
+    [enqueuePersistence],
   );
 
   const restorePersistedState = useCallback(
@@ -281,10 +360,11 @@ const useInteractablesResource = ({
 
   const setPersistenceAdapter = useCallback(
     (adapter: Unstable_InteractablePersistenceAdapter | undefined) => {
+      if (adapterRef.current !== adapter) flushIfPending();
       adapterRef.current = adapter;
       if (adapter) void loadFromAdapter(adapter);
     },
-    [loadFromAdapter],
+    [flushIfPending, loadFromAdapter],
   );
 
   const getCurrentThreadId = useCallback((): string | undefined => {
@@ -309,7 +389,7 @@ const useInteractablesResource = ({
     setPersistenceAdapter(persistence);
     return () => {
       if (adapterRef.current === persistence) {
-        adapterRef.current = undefined;
+        setPersistenceAdapter(undefined);
       }
     };
   }, [persistence, setPersistenceAdapter]);
@@ -319,25 +399,17 @@ const useInteractablesResource = ({
       clearTimeout(debounceTimerRef.current);
       debounceTimerRef.current = undefined;
     }
-    if (!adapterRef.current) return;
-    if (!hasPendingLocalChangeRef.current && dirtyIdsRef.current.size === 0)
-      return;
+    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 (!hasPendingLocalChangeRef.current) {
-      runPersistence();
-    }
+    if (adapterRef.current) enqueuePersistence(adapterRef.current);
     return p;
-  }, [runPersistence]);
-
-  const flushIfPending = useCallback(() => {
-    if (adapterRef.current && debounceTimerRef.current !== undefined) {
-      clearTimeout(debounceTimerRef.current);
-      debounceTimerRef.current = undefined;
-      runPersistence();
-    }
-  }, [runPersistence]);
+  }, [enqueuePersistence]);
 
   const setDefState = useCallback(
     (id: string, updater: (prev: unknown) => unknown) => {
@@ -384,12 +456,52 @@ const useInteractablesResource = ({
   );
 
   useEffect(() => {
-    for (const cb of subscribersRef.current) cb();
+    notifyStateSubscribers(subscribersRef.current);
   }, [state]);
 
-  useEffect(() => {
-    return clientRef.current!.modelContext().register(provider);
-  }, [clientRef, provider]);
+  useAssistantScopeEffect(
+    "modelContext",
+    () => clientRef.current!.modelContext().register(provider),
+    [provider],
+  );
+
+  const installUpdateToolUI = useCallback(
+    (name: string, entry: UpdateToolUIEntry) => {
+      const toolsAccessor = clientRef.current?.tools;
+      if (!toolsAccessor || toolsAccessor.source == null) return false;
+      entry.unsubscribe = toolsAccessor().setToolUI(
+        interactableToolName(name),
+        entry.render,
+        { standalone: true },
+      );
+      return true;
+    },
+    [clientRef],
+  );
+
+  // register() installs update-tool UIs against the tools instance bound at
+  // call time; this re-applies the retained entries when that instance is
+  // structurally replaced and installs entries recorded while no tools
+  // scope was available. Each re-apply releases the entry's previous
+  // install first, so it is an orphaned no-op against a replaced instance
+  // and an idempotent replacement against a live one.
+  useAssistantScopeEffect(
+    "tools",
+    () => {
+      for (const [name, entry] of updateToolUIsRef.current) {
+        entry.unsubscribe?.();
+        entry.unsubscribe = undefined;
+        installUpdateToolUI(name, entry);
+      }
+      return () => {
+        for (const entry of updateToolUIsRef.current.values()) {
+          entry.unsubscribe?.();
+          entry.unsubscribe = undefined;
+        }
+      };
+    },
+    [installUpdateToolUI],
+  );
 
   const register = useCallback(
     (def: InternalInteractableRegistration) => {
@@ -423,36 +535,35 @@ const useInteractablesResource = ({
 
       let releaseUpdateToolUI: (() => void) | undefined;
       if (def.updateRender) {
-        const toolsAccessor = clientRef.current?.tools;
-        if (toolsAccessor && toolsAccessor.source != null) {
-          const toolName = interactableToolName(def.name);
-          const existing = updateToolUIsRef.current.get(def.name);
-          if (existing) {
-            existing.count++;
-          } else {
-            updateToolUIsRef.current.set(def.name, {
-              count: 1,
-              unsubscribe: toolsAccessor().setToolUI(
-                toolName,
-                def.updateRender,
-                { standalone: true },
-              ),
-            });
-          }
-          releaseUpdateToolUI = () => {
-            const entry = updateToolUIsRef.current.get(def.name);
-            if (!entry) return;
-            if (--entry.count === 0) {
-              updateToolUIsRef.current.delete(def.name);
-              entry.unsubscribe();
-            }
-          };
-        } else if (process.env.NODE_ENV !== "production") {
+        const existing = updateToolUIsRef.current.get(def.name);
+        const entry = existing ?? {
+          count: 0,
+          render: def.updateRender,
+          unsubscribe: undefined,
+        };
+        entry.count++;
+        if (!existing) updateToolUIsRef.current.set(def.name, entry);
+
+        if (
+          !entry.unsubscribe &&
+          !installUpdateToolUI(def.name, entry) &&
+          process.env.NODE_ENV !== "production"
+        ) {
           console.warn(
             `[Interactables] "${def.name}" supplied an updateRender, but no ` +
-              `tools scope is available to install it into.`,
+              `tools scope is available yet; it will be installed once one appears.`,
           );
         }
+
+        releaseUpdateToolUI = () => {
+          const entry = updateToolUIsRef.current.get(def.name);
+          if (!entry) return;
+          if (--entry.count === 0) {
+            updateToolUIsRef.current.delete(def.name);
+            entry.unsubscribe?.();
+            entry.unsubscribe = undefined;
+          }
+        };
       }
 
       // The same id re-registers once per anchor (its create call + each update_*).
@@ -552,7 +663,13 @@ const useInteractablesResource = ({
         });
       };
     },
-    [flushIfPending, clientRef, getCurrentThreadId, setStateAndRef],
+    [
+      flushIfPending,
+      clientRef,
+      getCurrentThreadId,
+      installUpdateToolUI,
+      setStateAndRef,
+    ],
   );
 
   return {