@assistant-ui/react-google-adk 0.0.34 → 0.0.36

package/src/useAdkRuntime.ts

@@ -4,6 +4,7 @@ import {
   useMemo,
   useRef,
   useState,
+  useSyncExternalStore,
 } from "react";
 import {
   pickExternalStoreSharedOptions,
@@ -21,6 +22,9 @@ import {
 import {
   createAbortableThreadLoad,
   createCloudThreadListAdapterCreateFallback,
+  isRecord,
+  RunLeases,
+  type RunLease,
 } from "@assistant-ui/core/internal";
 import {
   useCloudThreadListAdapter,
@@ -41,7 +45,7 @@ import type {
   OnAdkCustomEventCallback,
   OnAdkAgentTransferCallback,
 } from "./types";
-import { useAdkMessages } from "./useAdkMessages";
+import { useAdkMessagesInternal } from "./useAdkMessages";
 import {
   convertAdkMessage,
   createAdkMessageConverter,
@@ -71,6 +75,7 @@ export type UseAdkRuntimeOptions = ExternalStoreSharedOptions & {
    */
   onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
   autoCancelPendingToolCalls?: boolean | undefined;
+  /** @deprecated Experimental since 2025-01-03. Not scheduled for removal; the API may change in any release. */
   unstable_allowCancellation?: boolean | undefined;
   getCheckpointId?: (
     threadId: string,
@@ -105,6 +110,11 @@ export type UseAdkRuntimeOptions = ExternalStoreSharedOptions & {
       }
     | undefined;
   cloud?: AssistantCloud | undefined;
+  /**
+   * Stable identity for the account or workspace owning Cloud runtime state.
+   * Provide it from the first render and change it when that scope changes.
+   */
+  scopeId?: string | undefined;
   /**
    * A `RemoteThreadListAdapter` to use instead of the cloud adapter.
    * Use with `createAdkSessionAdapter` for ADK session-backed persistence.
@@ -123,7 +133,64 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     eventHandlers,
   } = options;
   const aui = useAui();
+  const runConfigByToolCallIdRef = useRef(new Map<string, unknown>());
+
+  const rememberMessageOwnership = useCallback(
+    (newMessages: AdkMessage[], runConfig: unknown) => {
+      const toolOwnership = runConfigByToolCallIdRef.current;
+      for (const message of newMessages) {
+        if (message.type !== "ai") continue;
+        for (const toolCall of message.tool_calls ?? []) {
+          if (!isRecord(toolCall)) continue;
+          if (!toolOwnership.has(toolCall.id)) {
+            toolOwnership.set(toolCall.id, runConfig);
+          }
+        }
+      }
+    },
+    [],
+  );
+
+  const seedMessageOwnership = useCallback((history: AdkMessage[]) => {
+    const currentOwnership = runConfigByToolCallIdRef.current;
+    const nextOwnership = new Map<string, unknown>();
+    for (const message of history) {
+      if (message.type !== "ai") continue;
+      for (const toolCall of message.tool_calls ?? []) {
+        if (!isRecord(toolCall)) continue;
+        // Loaded ids must remain present even without a local owner because
+        // streamed event windows use has() to avoid attributing them later.
+        nextOwnership.set(
+          toolCall.id,
+          currentOwnership.has(toolCall.id)
+            ? currentOwnership.get(toolCall.id)
+            : undefined,
+        );
+      }
+    }
+    runConfigByToolCallIdRef.current = nextOwnership;
+  }, []);
+
+  const pruneMessageOwnership = useCallback((history: AdkMessage[]) => {
+    const toolCallIds = new Set<string>();
+    for (const message of history) {
+      if (message.type !== "ai") continue;
+      for (const toolCall of message.tool_calls ?? []) {
+        if (!isRecord(toolCall)) continue;
+        toolCallIds.add(toolCall.id);
+      }
+    }
+    for (const id of runConfigByToolCallIdRef.current.keys()) {
+      if (!toolCallIds.has(id)) runConfigByToolCallIdRef.current.delete(id);
+    }
+  }, []);
+
+  const getToolRunConfig = useCallback((toolCallId: string) => {
+    return runConfigByToolCallIdRef.current.get(toolCallId);
+  }, []);
+
   const {
+    controller,
     messages,
     stateDelta,
     agentInfo,
@@ -136,11 +203,12 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     sendMessage,
     cancel,
     setMessages,
-    replaceMessages,
-    applySnapshot,
-  } = useAdkMessages({
+    replaceMessages: replaceAdkMessages,
+    applySnapshot: applyAdkSnapshot,
+  } = useAdkMessagesInternal({
     stream,
     ...(eventHandlers && { eventHandlers }),
+    onMessages: rememberMessageOwnership,
   });
 
   const loadRef = useRef(load);
@@ -148,10 +216,40 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     loadRef.current = load;
   }, [load]);
   const [loadController] = useState(createAbortableThreadLoad);
+  const initialLoadRef = useRef<{
+    promise: Promise<void>;
+    active: boolean;
+    snapshot: AdkThreadSnapshot | undefined;
+  } | null>(null);
+  const waitForInitialLoad = () => {
+    const load = initialLoadRef.current;
+    if (!load) return undefined;
+    return load.promise.then(() => ({
+      active:
+        load.active &&
+        (!threadListItem ||
+          aui.threads.getState().mainThreadId === threadListItem.getState().id),
+      snapshot: load.snapshot,
+    }));
+  };
   const messagesRef = useRef(messages);
   useInsertionEffect(() => {
     messagesRef.current = messages;
   }, [messages]);
+  const applySnapshot = useCallback(
+    (snapshot: AdkThreadSnapshot) => {
+      seedMessageOwnership(snapshot.messages);
+      applyAdkSnapshot(snapshot);
+    },
+    [applyAdkSnapshot, seedMessageOwnership],
+  );
+  const replaceMessages = useCallback(
+    (nextMessages: AdkMessage[]) => {
+      pruneMessageOwnership(nextMessages);
+      replaceAdkMessages(nextMessages);
+    },
+    [pruneMessageOwnership, replaceAdkMessages],
+  );
   const [isLoadingThread, setIsLoadingThread] = useState(
     () =>
       load !== undefined && aui.threadListItem.getState().externalId != null,
@@ -169,21 +267,47 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   useInsertionEffect(() => {
     isRunningRef.current = effectiveIsRunning;
   }, [effectiveIsRunning]);
-  const runGenerationRef = useRef(0);
+  const [runLeases] = useState(() => new RunLeases());
+  const reloadLookupRef = useRef<{
+    lease: RunLease;
+    beforeReload: AdkThreadSnapshot;
+  } | null>(null);
 
-  const handleSendMessage = async (
-    msgs: AdkMessage[],
-    config: AdkSendMessageConfig,
+  const runExclusive = async (
+    run: (isCurrent: () => boolean) => Promise<void>,
   ) => {
-    const generation = ++runGenerationRef.current;
+    const lease = runLeases.begin();
     try {
       setIsRunning(true);
-      await sendMessage(msgs, config);
+      await run(lease.isCurrent);
     } finally {
-      if (runGenerationRef.current === generation) setIsRunning(false);
+      if (lease.isCurrent()) setIsRunning(false);
     }
   };
 
+  const handleSendMessage = (
+    msgs: AdkMessage[],
+    config: AdkSendMessageConfig,
+  ) => {
+    const isToolContinuation =
+      msgs.length > 0 && msgs.every((msg) => msg.type === "tool");
+    const continuationConfig =
+      isToolContinuation && config.runConfig === undefined
+        ? {
+            ...config,
+            runConfig: getToolRunConfig(msgs[0]!.tool_call_id),
+          }
+        : config;
+
+    return runExclusive(() => sendMessage(msgs, continuationConfig));
+  };
+
+  const stopRun = () => {
+    runLeases.invalidate();
+    setIsRunning(false);
+    cancel();
+  };
+
   const { approvals: toolApprovals, key: toolApprovalsKey } =
     projectAdkToolApprovals(messages);
   // The messageConverter memo below reads this during render, where the ref
@@ -233,43 +357,20 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     adkMessagesRef.current = messages;
   }, [messages]);
 
-  const stagedMessagesRef = useRef(
-    new Map<
-      string,
-      {
-        message: AdkMessage & { id: string };
-        runConfig: AppendMessage["runConfig"];
-      }
-    >(),
+  const stagedMessageCount = useSyncExternalStore(
+    controller.subscribe,
+    controller.getStagedMessageCount,
+    controller.getStagedMessageCount,
   );
-  const [stagedMessageCount, setStagedMessageCount] = useState(0);
   const hasStagedMessages = stagedMessageCount > 0;
 
-  const getStagedRun = (parentId: string | null) => {
-    if (!parentId || !stagedMessagesRef.current.has(parentId)) return null;
-
-    const staged: AdkMessage[] = [];
-    for (const message of adkMessagesRef.current) {
-      if (message.id && stagedMessagesRef.current.has(message.id)) {
-        staged.push(stagedMessagesRef.current.get(message.id)!.message);
-      }
-      if (message.id === parentId) break;
-    }
-
-    return {
-      messages: staged,
-      runConfig: stagedMessagesRef.current.get(parentId)!.runConfig,
-    };
-  };
-
   const stageUserMessage = (msg: AppendMessage) => {
     const stagedMessage = toAdkUserMessage(msg);
-    stagedMessagesRef.current.set(stagedMessage.id, {
-      message: stagedMessage,
-      runConfig: msg.runConfig,
+    controller.dispatch({
+      type: "staged.stage",
+      entry: { message: stagedMessage, runConfig: msg.runConfig },
     });
-    setStagedMessageCount(stagedMessagesRef.current.size);
-    const nextMessages = [...adkMessagesRef.current, stagedMessage];
+    const nextMessages = [...controller.getState().messages, stagedMessage];
     adkMessagesRef.current = nextMessages;
     setMessages(nextMessages);
   };
@@ -312,7 +413,14 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
               messagesRef.current !== messagesAtLoadStart)
           )
             return;
+          reloadLookupRef.current = null;
           applySnapshot(snapshot);
+          messagesRef.current = snapshot.messages;
+          adkMessagesRef.current = snapshot.messages;
+          longRunningToolIdsRef.current = snapshot.longRunningToolIds ?? [];
+          if (purpose === "initial" && initialLoadRef.current) {
+            initialLoadRef.current.snapshot = snapshot;
+          }
         },
         onSettled: () => {
           setIsLoadingThread(false);
@@ -326,11 +434,26 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   );
 
   useReplaySafeEffect(() => {
-    runLoad();
+    let release!: () => void;
+    const barrier = {
+      promise: new Promise<void>((resolve) => {
+        release = resolve;
+      }),
+      active: true,
+      snapshot: undefined as AdkThreadSnapshot | undefined,
+    };
+    initialLoadRef.current = barrier;
+    const settle = () => {
+      if (initialLoadRef.current === barrier) initialLoadRef.current = null;
+      release();
+    };
+    void runLoad().then(settle, settle);
     return () => {
+      barrier.active = false;
       // Whatever is current, not this effect's own controller: a refetch swaps
       // the ref, and one in flight at unmount must be aborted too.
       loadController.abort();
+      settle();
       setIsLoadingThread(false);
     };
   }, [threadListItem]);
@@ -353,9 +476,17 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       authRequests,
       escalated,
       messageMetadata,
-      send: handleSendMessage,
+      send: (messages, config) => {
+        const initialLoad = waitForInitialLoad();
+        if (!initialLoad) return handleSendMessage(messages, config);
+        return initialLoad.then(({ active }) =>
+          active ? handleSendMessage(messages, config) : undefined,
+        );
+      },
     }),
     onNew: async (msg) => {
+      const initialLoad = await waitForInitialLoad();
+      if (initialLoad && !initialLoad.active) return;
       if (!(msg.startRun ?? msg.role === "user")) {
         stageUserMessage(msg);
         return;
@@ -363,7 +494,11 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
 
       const cancellations =
         autoCancelPendingToolCalls !== false
-          ? getPendingCancellations(messages, longRunningToolIds)
+          ? getPendingCancellations(
+              initialLoad?.snapshot?.messages ?? messagesRef.current,
+              initialLoad?.snapshot?.longRunningToolIds ??
+                longRunningToolIdsRef.current,
+            )
           : [];
 
       return handleSendMessage(
@@ -380,6 +515,9 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     },
     onEdit: getCheckpointId
       ? async (msg) => {
+          const initialLoad = waitForInitialLoad();
+          if (initialLoad && !(await initialLoad).active) return;
+          stopRun();
           const truncated = truncateAdkMessages(
             threadMessagesRef.current,
             msg.parentId,
@@ -387,44 +525,41 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           replaceMessages(truncated);
           if (!(msg.startRun ?? msg.role === "user")) {
             const stagedMessage = toAdkUserMessage(msg);
-            stagedMessagesRef.current.set(stagedMessage.id, {
-              message: stagedMessage,
-              runConfig: msg.runConfig,
+            controller.dispatch({
+              type: "staged.stage",
+              entry: { message: stagedMessage, runConfig: msg.runConfig },
             });
-            setStagedMessageCount(stagedMessagesRef.current.size);
             const nextMessages = [...truncated, stagedMessage];
             adkMessagesRef.current = nextMessages;
             setMessages(nextMessages);
             return;
           }
+          const editedMessage = toAdkUserMessage(msg);
+          setMessages([...truncated, editedMessage]);
           const externalId = aui.threadListItem.getState().externalId;
-          const checkpointId = externalId
-            ? await getCheckpointId(externalId, truncated)
-            : null;
-          return handleSendMessage(
-            [
-              {
-                id: generateId(),
-                type: "human",
-                content: getMessageContent(msg),
-              },
-            ],
-            {
+          return runExclusive(async (isCurrent) => {
+            const checkpointId = externalId
+              ? await getCheckpointId(externalId, truncated)
+              : null;
+            if (!isCurrent()) return;
+            await sendMessage([editedMessage], {
               runConfig: msg.runConfig,
               ...(checkpointId && { checkpointId }),
-            },
-          );
+            });
+          });
         }
       : undefined,
     ...(getCheckpointId || hasStagedMessages
       ? {
           onReload: async (parentId, config) => {
-            const stagedRun = getStagedRun(parentId);
+            const initialLoad = waitForInitialLoad();
+            if (initialLoad && !(await initialLoad).active) return;
+            const stagedRun = controller.getStagedRun(parentId);
             if (stagedRun) {
-              for (const message of stagedRun.messages) {
-                stagedMessagesRef.current.delete(message.id);
-              }
-              setStagedMessageCount(stagedMessagesRef.current.size);
+              controller.dispatch({
+                type: "staged.unstage",
+                ids: stagedRun.messages.map((message) => message.id!),
+              });
               return handleSendMessage(stagedRun.messages, {
                 runConfig: config.runConfig ?? stagedRun.runConfig,
               });
@@ -433,18 +568,48 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
             if (!getCheckpointId)
               throw new Error("Runtime does not support reloading messages.");
 
+            stopRun();
+            const beforeReload: AdkThreadSnapshot = {
+              messages: adkMessagesRef.current,
+              longRunningToolIds,
+              toolConfirmations,
+              authRequests,
+              escalated,
+              messageMetadata,
+              stateDelta,
+              artifactDelta,
+              agentInfo,
+            };
             const truncated = truncateAdkMessages(
               threadMessagesRef.current,
               parentId,
             );
             replaceMessages(truncated);
             const externalId = aui.threadListItem.getState().externalId;
-            const checkpointId = externalId
-              ? await getCheckpointId(externalId, truncated)
-              : null;
-            return handleSendMessage([], {
-              runConfig: config.runConfig,
-              ...(checkpointId && { checkpointId }),
+            return runExclusive(async (isCurrent) => {
+              const lookup = {
+                lease: runLeases.current(),
+                beforeReload,
+              };
+              reloadLookupRef.current = lookup;
+              let checkpointId: string | null;
+              try {
+                checkpointId = externalId
+                  ? await getCheckpointId(externalId, truncated)
+                  : null;
+              } catch (error) {
+                if (isCurrent() && reloadLookupRef.current === lookup)
+                  applySnapshot(beforeReload);
+                throw error;
+              } finally {
+                if (reloadLookupRef.current === lookup)
+                  reloadLookupRef.current = null;
+              }
+              if (!isCurrent()) return;
+              await sendMessage([], {
+                runConfig: config.runConfig,
+                ...(checkpointId && { checkpointId }),
+              });
             });
           },
         }
@@ -456,6 +621,8 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       isError,
       artifact,
     }) => {
+      const initialLoad = waitForInitialLoad();
+      if (initialLoad && !(await initialLoad).active) return;
       await handleSendMessage(
         [
           {
@@ -472,6 +639,8 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       );
     },
     onRespondToToolApproval: async (options) => {
+      const initialLoad = waitForInitialLoad();
+      if (initialLoad && !(await initialLoad).active) return;
       await handleSendMessage(
         [
           toAdkToolConfirmationReply(
@@ -484,7 +653,14 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     },
     onCancel: unstable_allowCancellation
       ? async () => {
-          cancel();
+          const lookup = reloadLookupRef.current;
+          const beforeReload = lookup?.lease.isCurrent()
+            ? lookup.beforeReload
+            : undefined;
+          stopRun();
+          // A reload stopped before it sent leaves the ADK session holding the
+          // turn it removed, so the thread shows that turn again.
+          if (beforeReload) applySnapshot(beforeReload);
         }
       : undefined,
     ...(load !== undefined && {
@@ -497,6 +673,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
 
 export const useAdkRuntime = ({
   cloud,
+  scopeId,
   sessionAdapter,
   create,
   delete: deleteFn,
@@ -507,6 +684,7 @@ export const useAdkRuntime = ({
   const cloudAdapter = useCloudThreadListAdapter({
     sdk: ADK_SDK,
     cloud,
+    scopeId,
     create: createCloudThreadListAdapterCreateFallback(
       create,
       aui.threadListItem,