@assistant-ui/react-google-adk 0.0.15 → 0.0.16

package/src/useAdkRuntime.ts

@@ -11,6 +11,7 @@ import {
   type AppendMessage,
   type ThreadMessage,
   type ToolExecutionStatus,
+  generateId,
 } from "@assistant-ui/core";
 import {
   useCloudThreadListAdapter,
@@ -31,7 +32,7 @@ import type {
 } from "./types";
 import { useAdkMessages } from "./useAdkMessages";
 import { convertAdkMessage } from "./convertAdkMessages";
-import { symbolAdkRuntimeExtras, type AdkRuntimeExtras } from "./hooks";
+import { adkExtras } from "./adkExtras";
 import { v4 as uuidv4 } from "uuid";
 
 /** @internal — exported for unit tests. */
@@ -138,8 +139,26 @@ const truncateAdkMessages = (
   return truncated;
 };
 
+const toAdkUserMessage = (
+  msg: AppendMessage,
+  id = generateId(),
+): AdkMessage & { type: "human"; id: string } => ({
+  id,
+  type: "human",
+  content: getMessageContent(msg),
+});
+
 export type UseAdkRuntimeOptions = ExternalStoreSharedOptions & {
   stream: AdkStreamCallback;
+  /**
+   * Called whenever the active thread's canonical (remote) ID changes, so the
+   * value can be treated as a managed/controlled variable (e.g. synced to a URL
+   * query param). Only the settled remote ID is emitted: while a freshly created
+   * thread is still optimistic the value is `undefined`, and the real ID is
+   * emitted once the thread is initialized; the transient local ID is never
+   * surfaced.
+   */
+  onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
   autoCancelPendingToolCalls?: boolean | undefined;
   unstable_allowCancellation?: boolean | undefined;
   getCheckpointId?: (
@@ -196,6 +215,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     messageMetadata,
     sendMessage,
     cancel,
+    setMessages,
     replaceMessages,
   } = useAdkMessages({
     stream,
@@ -232,6 +252,50 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   const threadMessagesRef = useRef(threadMessages);
   threadMessagesRef.current = threadMessages;
 
+  const adkMessagesRef = useRef(messages);
+  adkMessagesRef.current = messages;
+
+  const stagedMessagesRef = useRef(
+    new Map<
+      string,
+      {
+        message: AdkMessage & { id: string };
+        runConfig: AppendMessage["runConfig"];
+      }
+    >(),
+  );
+  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,
+    });
+    setStagedMessageCount(stagedMessagesRef.current.size);
+    const nextMessages = [...adkMessagesRef.current, stagedMessage];
+    adkMessagesRef.current = nextMessages;
+    setMessages(nextMessages);
+  };
+
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),
     isRunning: effectiveIsRunning,
@@ -239,8 +303,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     unstable_enableToolInvocations: true,
     setToolStatuses,
     adapters: { attachments, dictation, feedback, speech, voice },
-    extras: {
-      [symbolAdkRuntimeExtras]: true,
+    extras: adkExtras.provide({
       agentInfo,
       stateDelta,
       artifactDelta,
@@ -250,8 +313,13 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       escalated,
       messageMetadata,
       send: handleSendMessage,
-    } satisfies AdkRuntimeExtras,
+    }),
     onNew: async (msg) => {
+      if (!(msg.startRun ?? msg.role === "user")) {
+        stageUserMessage(msg);
+        return;
+      }
+
       const cancellations =
         autoCancelPendingToolCalls !== false
           ? getPendingCancellations(messages, longRunningToolIds)
@@ -276,6 +344,18 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
             msg.parentId,
           );
           replaceMessages(truncated);
+          if (!(msg.startRun ?? msg.role === "user")) {
+            const stagedMessage = toAdkUserMessage(msg);
+            stagedMessagesRef.current.set(stagedMessage.id, {
+              message: stagedMessage,
+              runConfig: msg.runConfig,
+            });
+            setStagedMessageCount(stagedMessagesRef.current.size);
+            const nextMessages = [...truncated, stagedMessage];
+            adkMessagesRef.current = nextMessages;
+            setMessages(nextMessages);
+            return;
+          }
           const externalId = aui.threadListItem().getState().externalId;
           const checkpointId = externalId
             ? await getCheckpointId(externalId, truncated)
@@ -295,23 +375,39 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           );
         }
       : undefined,
-    onReload: getCheckpointId
-      ? async (parentId, config) => {
-          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 }),
-          });
+    ...(getCheckpointId || hasStagedMessages
+      ? {
+          onReload: async (parentId, config) => {
+            const stagedRun = getStagedRun(parentId);
+            if (stagedRun) {
+              for (const message of stagedRun.messages) {
+                stagedMessagesRef.current.delete(message.id);
+              }
+              setStagedMessageCount(stagedMessagesRef.current.size);
+              return handleSendMessage(stagedRun.messages, {
+                runConfig: config.runConfig ?? stagedRun.runConfig,
+              });
+            }
+
+            if (!getCheckpointId)
+              throw new Error("Runtime does not support reloading messages.");
+
+            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 }),
+            });
+          },
         }
-      : undefined,
+      : {}),
     onAddToolResult: async ({
       toolCallId,
       toolName,
@@ -373,6 +469,7 @@ export const useAdkRuntime = ({
   sessionAdapter,
   create,
   delete: deleteFn,
+  onThreadIdChange,
   ...options
 }: UseAdkRuntimeOptions) => {
   const aui = useAui();
@@ -394,5 +491,6 @@ export const useAdkRuntime = ({
     },
     adapter,
     allowNesting: true,
+    onThreadIdChange,
   });
 };