@assistant-ui/react-ag-ui 0.0.65 → 0.0.67

package/src/useAgUiRuntime.ts

@@ -1,7 +1,6 @@
 "use client";
 
 import {
-  useCallback,
   useEffect,
   useLayoutEffect,
   useMemo,
@@ -61,10 +60,9 @@ export function useAgUiRuntime(
   options: UseAgUiRuntimeOptions,
 ): AgUiAssistantRuntime {
   const logger = useMemo(() => makeLogger(options.logger), [options.logger]);
-  const [_version, setVersion] = useState(0);
-  const notifyUpdate = useCallback(() => setVersion((v) => v + 1), []);
   const coreRef = useRef<AgUiThreadRuntimeCore | null>(null);
   const threadSwitchGenerationRef = useRef(0);
+  const switchingGenerationRef = useRef<number | null>(null);
   const runtimeAdapters = useRuntimeAdapters();
 
   const historyAdapter = options.adapters?.history ?? runtimeAdapters?.history;
@@ -80,11 +78,16 @@ export function useAgUiRuntime(
       ...(options.onError && { onError: options.onError }),
       ...(options.onCancel && { onCancel: options.onCancel }),
       ...(historyAdapter && { history: historyAdapter }),
-      notifyUpdate,
+      isThreadSwitching: () => switchingGenerationRef.current !== null,
     });
   }
 
   const core = coreRef.current;
+  const snapshot = useSyncExternalStore(
+    core.subscribe,
+    core.getSnapshot,
+    core.getSnapshot,
+  );
   useEffect(() => {
     core.updateOptions({
       agent: options.agent,
@@ -105,7 +108,7 @@ export function useAgUiRuntime(
   const hasExecutingTools = Object.values(toolStatuses).some(
     (s) => s?.type === "executing",
   );
-  const isRunning = core.isRunning() || hasExecutingTools;
+  const isRunning = snapshot.isRunning || hasExecutingTools;
 
   // The driver sends through the agent core rather than the runtime, because
   // the runtime's append routes every tail append back into this queue.
@@ -122,6 +125,12 @@ export function useAgUiRuntime(
     queueRef.current = createMessageQueue({
       run: (message) => {
         const controller = queueRef.current;
+        // A send dispatched mid-switch is dropped, and the queue released, or
+        // it would wait for an idle edge that never comes.
+        if (switchingGenerationRef.current !== null) {
+          controller?.notifyIdle();
+          return;
+        }
         const edgesAtDispatch = busyEdgesRef.current;
         // The queue drops the item before dispatching and stays busy until an
         // idle edge. An append observed as busy is released by that falling
@@ -205,43 +214,82 @@ export function useAgUiRuntime(
     const { onSwitchToNewThread, onSwitchToThread, ...rest } =
       threadListAdapter;
 
+    const abandonPreviousThread = () => {
+      queueRef.current?.notifyCancelled();
+      if (core.isRunning()) {
+        void core.cancel().catch((error: unknown) => {
+          logger.error?.(
+            "[agui] cancelling the run on thread switch failed",
+            error,
+          );
+        });
+      }
+      queueRef.current?.clear();
+    };
+
+    const releaseSwitch = (generation: number) => {
+      if (switchingGenerationRef.current === generation) {
+        switchingGenerationRef.current = null;
+      }
+    };
+
     return {
       ...rest,
       onSwitchToNewThread: onSwitchToNewThread
         ? async () => {
             const generation = ++threadSwitchGenerationRef.current;
-            // Clear before the thread id flips, or the old messages leak
-            // into the new thread as a sibling branch.
-            core.applyExternalMessages([]);
-            core.resetThreadState();
-            await onSwitchToNewThread();
-            if (generation !== threadSwitchGenerationRef.current) return;
-            core.applyExternalMessages([]);
-            core.resetThreadState();
+            switchingGenerationRef.current = generation;
+            try {
+              abandonPreviousThread();
+              if (generation !== threadSwitchGenerationRef.current) return;
+              // Clear before the thread id flips, or the old messages leak
+              // into the new thread as a sibling branch.
+              core.applyExternalMessages([]);
+              core.resetThreadState();
+              await onSwitchToNewThread();
+              if (generation !== threadSwitchGenerationRef.current) return;
+              abandonPreviousThread();
+              if (generation !== threadSwitchGenerationRef.current) return;
+              core.applyExternalMessages([]);
+              core.resetThreadState();
+            } finally {
+              releaseSwitch(generation);
+            }
           }
         : undefined,
       onSwitchToThread: onSwitchToThread
         ? async (threadId: string) => {
             const generation = ++threadSwitchGenerationRef.current;
-            // Clear before the thread id flips, or the old messages leak
-            // into the new thread as a sibling branch.
-            core.applyExternalMessages([]);
-            core.resetThreadState();
-            const result = await onSwitchToThread(threadId);
-            if (generation !== threadSwitchGenerationRef.current) return;
-            core.applyExternalMessages([]);
-            core.resetThreadState();
-            core.applyExternalMessages(result.messages);
-            if (result.state !== undefined) {
-              core.loadExternalState(result.state);
-            }
-            if (result.unstable_resume) {
-              void core.resumeInFlightRun(result.messages);
+            switchingGenerationRef.current = generation;
+            try {
+              abandonPreviousThread();
+              if (generation !== threadSwitchGenerationRef.current) return;
+              // Clear before the thread id flips, or the old messages leak
+              // into the new thread as a sibling branch.
+              core.applyExternalMessages([]);
+              core.resetThreadState();
+              const result = await onSwitchToThread(threadId);
+              if (generation !== threadSwitchGenerationRef.current) return;
+              abandonPreviousThread();
+              if (generation !== threadSwitchGenerationRef.current) return;
+              core.applyExternalMessages([]);
+              core.resetThreadState();
+              core.applyExternalMessages(result.messages);
+              if (result.state !== undefined) {
+                core.loadExternalState(result.state);
+              }
+              if (generation !== threadSwitchGenerationRef.current) return;
+              releaseSwitch(generation);
+              if (result.unstable_resume) {
+                void core.resumeInFlightRun(result.messages);
+              }
+            } finally {
+              releaseSwitch(generation);
             }
           }
         : undefined,
     };
-  }, [threadListAdapter, core]);
+  }, [threadListAdapter, core, logger]);
 
   const adapters = options.adapters;
   const adapterAdapters = useMemo(
@@ -259,14 +307,12 @@ export function useAgUiRuntime(
   const shared = useExternalStoreSharedOptions(options);
   const store = useMemo(
     () => {
-      void _version; // rerender on version change
-
       return {
         ...shared,
-        isLoading: core.isLoading,
-        messageRepository: core.getMessageRepository(),
-        state: core.getState(),
-        isRunning: core.isRunning(),
+        isLoading: snapshot.isLoading,
+        messageRepository: snapshot.messageRepository,
+        state: snapshot.state,
+        isRunning: snapshot.isRunning,
         extras: agUiExtras.provide({
           interrupts:
             core.getPendingInterrupts()?.interrupts ?? EMPTY_INTERRUPTS,
@@ -274,13 +320,16 @@ export function useAgUiRuntime(
           submitInterruptResponses: (responses) =>
             core.submitInterruptResponses(responses),
           steerAway: (message, responses) => core.steerAway(message, responses),
-          state: core.getState(),
+          state: snapshot.state,
           setState: (next) => core.setState(next),
         }),
         unstable_enableToolInvocations: true,
         unstable_persistsHistory: true,
         setToolStatuses,
-        onNew: (message: AppendMessage) => core.append(message),
+        onNew: (message: AppendMessage) =>
+          switchingGenerationRef.current !== null
+            ? Promise.resolve()
+            : core.append(message),
         onVoiceTranscript: (message) => core.appendVoiceTranscript(message),
         onEdit: (message: AppendMessage) => {
           queueController?.clear();
@@ -313,29 +362,34 @@ export function useAgUiRuntime(
         ...(queueController && { queue: queueController.adapter }),
       } satisfies ExternalStoreAdapter<ThreadMessage>;
     },
-    // _version is intentionally included to trigger re-computation when core state changes via notifyUpdate
     // toolInvocations intentionally excluded: abort/resume use refs internally and work with stale captures
     [
       adapterAdapters,
       core,
-      _version,
       isRunning,
       queueController,
       queueItems,
       steerQueueItems,
       shared,
+      snapshot,
     ],
   );
 
   const baseRuntime = useExternalStoreRuntime(store);
 
   const createRuntime = (): AgUiAssistantRuntime => {
-    const wrapper = Object.create(baseRuntime) as AgUiAssistantRuntime;
-    wrapper.unstable_getPendingInterrupts = () =>
-      core.getPendingInterrupts()?.interrupts ?? [];
-    wrapper.unstable_submitInterruptResponses = (responses) =>
-      core.submitInterruptResponses(responses);
-    return wrapper;
+    return {
+      threads: baseRuntime.threads,
+      get thread() {
+        return baseRuntime.thread;
+      },
+      registerModelContextProvider: (provider) =>
+        baseRuntime.registerModelContextProvider(provider),
+      unstable_getPendingInterrupts: () =>
+        core.getPendingInterrupts()?.interrupts ?? [],
+      unstable_submitInterruptResponses: (responses) =>
+        core.submitInterruptResponses(responses),
+    };
   };
   const [pinnedRuntime, setPinnedRuntime] = useState(() => ({
     baseRuntime,