@assistant-ui/react-google-adk 0.0.26 → 0.0.28

package/src/useAdkRuntime.toolOwnership.test.tsx

@@ -0,0 +1,236 @@
+// @vitest-environment jsdom
+
+import { act, render, waitFor } from "@testing-library/react";
+import { type FC } from "react";
+import { describe, expect, it, vi } from "vitest";
+import {
+  AssistantRuntimeProvider,
+  useAssistantTool,
+} from "@assistant-ui/core/react";
+import type {
+  AssistantRuntime,
+  RemoteThreadListAdapter,
+  ToolCallMessagePart,
+} from "@assistant-ui/core";
+import { useAdkRuntime } from "./useAdkRuntime";
+import type { AdkEvent } from "./types";
+
+const makeThreadListAdapter = (): RemoteThreadListAdapter => ({
+  list: vi.fn(async () => ({
+    threads: [
+      {
+        status: "regular" as const,
+        remoteId: "adk-1",
+        externalId: "adk-1",
+        title: "Existing ADK session",
+      },
+    ],
+  })),
+  initialize: vi.fn(async () => ({ remoteId: "adk-1", externalId: "adk-1" })),
+  rename: vi.fn(async () => {}),
+  archive: vi.fn(async () => {}),
+  unarchive: vi.fn(async () => {}),
+  delete: vi.fn(async () => {}),
+  generateTitle: vi.fn(async () => new ReadableStream() as never),
+  fetch: vi.fn(async () => ({
+    status: "regular" as const,
+    remoteId: "adk-1",
+    externalId: "adk-1",
+    title: "Existing ADK session",
+  })),
+});
+
+const toolCalls = (runtime: AssistantRuntime): ToolCallMessagePart[] =>
+  runtime.thread
+    .getState()
+    .messages.flatMap((message) => message.content as readonly unknown[])
+    .filter(
+      (part): part is ToolCallMessagePart =>
+        (part as ToolCallMessagePart).type === "tool-call",
+    );
+
+const GATED_CALL = {
+  id: "adk-original-1",
+  name: "delete_file",
+  args: { path: "/tmp/a" },
+};
+
+const callEvent = (longRunningToolIds?: string[]): AdkEvent => ({
+  id: "ev-1",
+  invocationId: "inv-1",
+  author: "agent",
+  content: { role: "model", parts: [{ functionCall: GATED_CALL }] },
+  ...(longRunningToolIds && { longRunningToolIds }),
+});
+
+const confirmationEvent = (): AdkEvent => ({
+  id: "ev-2",
+  invocationId: "inv-1",
+  author: "agent",
+  content: {
+    role: "model",
+    parts: [
+      {
+        functionCall: {
+          id: "adk-confirmation-1",
+          name: "adk_request_confirmation",
+          args: {
+            originalFunctionCall: GATED_CALL,
+            toolConfirmation: { hint: "Delete /tmp/a?" },
+          },
+        },
+      },
+    ],
+  },
+  longRunningToolIds: ["adk-confirmation-1"],
+});
+
+const serverResultEvent = (): AdkEvent => ({
+  id: "ev-2",
+  invocationId: "inv-1",
+  author: "agent",
+  content: {
+    role: "user",
+    parts: [
+      {
+        functionResponse: {
+          id: "adk-original-1",
+          name: "delete_file",
+          response: { server: "deleted" },
+        },
+      },
+    ],
+  },
+});
+
+/**
+ * Streams `first`, hands control back so the window between the two events is
+ * observable, then streams `second` once the returned `release` is called.
+ */
+const renderStreamingAdk = async (first: AdkEvent, second: AdkEvent) => {
+  let release!: () => void;
+  const held = new Promise<void>((r) => (release = r));
+  let firstDelivered!: () => void;
+  const firstDeliveredP = new Promise<void>((r) => (firstDelivered = r));
+
+  const stream = vi.fn(async function* (): AsyncGenerator<AdkEvent> {
+    yield first;
+    firstDelivered();
+    await held;
+    yield second;
+  });
+
+  const execute = vi.fn(async () => ({ deleted: true }));
+  const capture: { runtime: AssistantRuntime | null } = { runtime: null };
+
+  const DeleteFileTool = () => {
+    useAssistantTool({
+      toolName: "delete_file",
+      description: "delete a file",
+      parameters: {
+        type: "object" as const,
+        properties: { path: { type: "string" } },
+      },
+      execute,
+    });
+    return null;
+  };
+
+  const Inner: FC = () => {
+    const runtime = useAdkRuntime({
+      stream: stream as never,
+      sessionAdapter: makeThreadListAdapter(),
+    });
+    capture.runtime = runtime;
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <DeleteFileTool />
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  await act(async () => {
+    render(<Inner />);
+  });
+  await waitFor(() => expect(capture.runtime).not.toBeNull());
+  await act(async () => {
+    await capture.runtime!.threads.switchToThread("adk-1");
+  });
+  await act(async () => {
+    await new Promise((r) => setTimeout(r, 0));
+  });
+
+  await act(async () => {
+    void capture.runtime!.thread.append({
+      role: "user",
+      content: [{ type: "text", text: "delete it" }],
+    });
+    await firstDeliveredP;
+  });
+  await act(async () => {
+    await new Promise((r) => setTimeout(r, 0));
+  });
+
+  const settle = async () => {
+    await act(async () => {
+      release();
+      await new Promise((r) => setTimeout(r, 0));
+    });
+    await act(async () => {
+      await new Promise((r) => setTimeout(r, 0));
+    });
+  };
+
+  return { capture, execute, settle };
+};
+
+describe("useAdkRuntime tool ownership", () => {
+  it("does not run a frontend execute on a call ADK is about to gate", async () => {
+    const { capture, execute, settle } = await renderStreamingAdk(
+      callEvent(),
+      confirmationEvent(),
+    );
+
+    const inWindow = toolCalls(capture.runtime!)[0]!;
+    expect(inWindow.argsText).toBe('{"path":"/tmp/a"}');
+    expect(inWindow.approval).toBeUndefined();
+    expect(inWindow.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    const gated = toolCalls(capture.runtime!).find(
+      (part) => part.toolCallId === "adk-original-1",
+    )!;
+    expect(gated.approval).toEqual({ id: "adk-confirmation-1" });
+    expect(gated.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("does not run a frontend execute on a call ADK resolves itself", async () => {
+    const { capture, execute, settle } = await renderStreamingAdk(
+      callEvent(),
+      serverResultEvent(),
+    );
+
+    expect(toolCalls(capture.runtime!)[0]!.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    expect(toolCalls(capture.runtime!)[0]!.result).toBe('{"server":"deleted"}');
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("runs a frontend execute on a call ADK marks long-running", async () => {
+    const { capture, execute } = await renderStreamingAdk(
+      callEvent(["adk-original-1"]),
+      serverResultEvent(),
+    );
+
+    await waitFor(() => expect(execute).toHaveBeenCalledTimes(1));
+    await waitFor(() =>
+      expect(toolCalls(capture.runtime!)[0]!.result).toBeDefined(),
+    );
+  });
+});

package/src/useAdkRuntime.ts

@@ -1,6 +1,12 @@
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
-  getExternalStoreMessages,
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
+import {
   pickExternalStoreSharedOptions,
   type AttachmentAdapter,
   type DictationAdapter,
@@ -9,11 +15,14 @@ import {
   type RealtimeVoiceAdapter,
   type SpeechSynthesisAdapter,
   type AppendMessage,
-  type ThreadMessage,
+  type ToolCallMessagePart,
   type ToolExecutionStatus,
   generateId,
 } from "@assistant-ui/core";
-import { httpUrlPattern, parseDataUrl } from "@assistant-ui/core/internal";
+import {
+  createAbortableThreadLoad,
+  createCloudThreadListAdapterCreateFallback,
+} from "@assistant-ui/core/internal";
 import {
   useCloudThreadListAdapter,
   useRemoteThreadListRuntime,
@@ -37,139 +46,17 @@ import {
   convertAdkMessage,
   createAdkMessageConverter,
 } from "./convertAdkMessages";
+import {
+  getMessageContent,
+  getPendingCancellations,
+  toAdkUserMessage,
+  truncateAdkMessages,
+} from "./convertToAdkMessages";
 import {
   projectAdkToolApprovals,
   toAdkToolConfirmationReply,
 } from "./adkToolApproval";
 import { adkExtras } from "./adkExtras";
-import { v4 as uuidv4 } from "uuid";
-
-/** @internal — exported for unit tests. */
-export const getMessageContent = (msg: AppendMessage) => {
-  const allContent = [
-    ...msg.content,
-    ...(msg.attachments?.flatMap((a) => a.content) ?? []),
-  ];
-  const content = allContent.flatMap((part) => {
-    const type = part.type;
-    switch (type) {
-      case "text":
-        return { type: "text" as const, text: part.text };
-      case "image":
-        return { type: "image_url" as const, url: part.image };
-      case "file":
-        if (part.sourceType === "url" || httpUrlPattern.test(part.data)) {
-          return {
-            type: "file_url" as const,
-            url: part.data,
-            mimeType: part.mimeType,
-          };
-        }
-        return {
-          type: "file" as const,
-          mimeType: part.mimeType,
-          // Lands in Gemini `inlineData.data`, which takes bare base64, so a
-          // data URL envelope is stripped rather than forwarded.
-          data: parseDataUrl(part.data)?.data ?? part.data,
-          ...(part.filename != null && { filename: part.filename }),
-        };
-      case "audio": {
-        const parsed = parseDataUrl(part.audio.data);
-        return {
-          type: "file" as const,
-          mimeType: `audio/${part.audio.format}`,
-          data: parsed?.data ?? part.audio.data,
-        };
-      }
-      case "data":
-        return [];
-
-      case "tool-call":
-        throw new Error("Tool call appends are not supported.");
-
-      default: {
-        const _exhaustiveCheck: "reasoning" | "source" | "generative-ui" = type;
-        throw new Error(
-          `Unsupported append message part type: ${_exhaustiveCheck}`,
-        );
-      }
-    }
-  });
-
-  if (content.length === 1 && content[0]?.type === "text") {
-    return content[0].text ?? "";
-  }
-
-  return content;
-};
-
-/** @internal — exported for unit tests. */
-export const getPendingToolCalls = (messages: AdkMessage[]) => {
-  const pending = new Map<string, { id: string; name: string }>();
-  for (const msg of messages) {
-    if (msg.type === "ai" && msg.tool_calls) {
-      for (const tc of msg.tool_calls) {
-        pending.set(tc.id, tc);
-      }
-    }
-    if (msg.type === "tool") {
-      pending.delete(msg.tool_call_id);
-    }
-  }
-  return [...pending.values()];
-};
-
-/**
- * @internal — exported for unit tests.
- *
- * Returns `{cancelled: true}` tool responses for pending tool calls when the
- * user sends a new turn, EXCEPT for HITL interrupts marked via
- * `long_running_tool_ids` (`adk_request_input`, `adk_request_confirmation`,
- * `adk_request_credential`). Those must be answered through a dedicated tool
- * UI + submit helper, not auto-cancelled.
- */
-export const getPendingCancellations = (
-  messages: AdkMessage[],
-  longRunningToolIds: readonly string[],
-): Array<AdkMessage & { type: "tool" }> => {
-  const longRunningSet = new Set(longRunningToolIds);
-  return getPendingToolCalls(messages)
-    .filter((t) => !longRunningSet.has(t.id))
-    .map(
-      (t) =>
-        ({
-          id: uuidv4(),
-          type: "tool",
-          name: t.name,
-          tool_call_id: t.id,
-          content: JSON.stringify({ cancelled: true }),
-          status: "error",
-        }) satisfies AdkMessage & { type: "tool" },
-    );
-};
-
-const truncateAdkMessages = (
-  threadMessages: readonly ThreadMessage[],
-  parentId: string | null,
-): AdkMessage[] => {
-  if (parentId === null) return [];
-  const parentIndex = threadMessages.findIndex((m) => m.id === parentId);
-  if (parentIndex === -1) return [];
-  const truncated: AdkMessage[] = [];
-  for (let i = 0; i <= parentIndex && i < threadMessages.length; i++) {
-    truncated.push(...getExternalStoreMessages<AdkMessage>(threadMessages[i]!));
-  }
-  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;
@@ -256,14 +143,14 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   });
 
   const loadRef = useRef(load);
-  loadRef.current = load;
-  const loadControllerRef = useRef<{
-    controller: AbortController;
-    purpose: "initial" | "reload";
-    promise?: Promise<void> | undefined;
-  } | null>(null);
+  useInsertionEffect(() => {
+    loadRef.current = load;
+  }, [load]);
+  const loadController = useMemo(createAbortableThreadLoad, []);
   const messagesRef = useRef(messages);
-  messagesRef.current = messages;
+  useInsertionEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
   const [isLoadingThread, setIsLoadingThread] = useState(
     () =>
       load !== undefined && aui.threadListItem.getState().externalId != null,
@@ -278,7 +165,9 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   );
   const effectiveIsRunning = isRunning || hasExecutingTools;
   const isRunningRef = useRef(effectiveIsRunning);
-  isRunningRef.current = effectiveIsRunning;
+  useInsertionEffect(() => {
+    isRunningRef.current = effectiveIsRunning;
+  }, [effectiveIsRunning]);
 
   const handleSendMessage = async (
     msgs: AdkMessage[],
@@ -294,9 +183,26 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
 
   const { approvals: toolApprovals, key: toolApprovalsKey } =
     projectAdkToolApprovals(messages);
+  // The messageConverter memo below reads this during render, where the ref
+  // must carry the same render's approvals; a commit-scoped write would feed
+  // the memo the previous commit's approvals whenever the key changes. No
+  // callback reads it — approval replies project from the committed messages.
   const toolApprovalsRef = useRef(toolApprovals);
   toolApprovalsRef.current = toolApprovals;
 
+  const longRunningToolIdsRef = useRef(longRunningToolIds);
+  useInsertionEffect(() => {
+    longRunningToolIdsRef.current = longRunningToolIds;
+  }, [longRunningToolIds]);
+  // ADK resolves every call it did not mark long-running itself, and yields
+  // that call to the client one or more events before its own response, so
+  // only a long-running call is the client's to execute.
+  const isClientToolCall = useCallback(
+    (toolCall: ToolCallMessagePart) =>
+      longRunningToolIdsRef.current.includes(toolCall.toolCallId),
+    [],
+  );
+
   const messageConverter = useMemo(
     () =>
       toolApprovalsKey === ""
@@ -312,10 +218,17 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   });
 
   const threadMessagesRef = useRef(threadMessages);
-  threadMessagesRef.current = threadMessages;
+  useInsertionEffect(() => {
+    threadMessagesRef.current = threadMessages;
+  }, [threadMessages]);
 
+  // Staging assigns adkMessagesRef.current directly, so the effect must key on
+  // the committed messages alone; a dep-less publication would clobber the
+  // optimistic value on any unrelated commit.
   const adkMessagesRef = useRef(messages);
-  adkMessagesRef.current = messages;
+  useInsertionEffect(() => {
+    adkMessagesRef.current = messages;
+  }, [messages]);
 
   const stagedMessagesRef = useRef(
     new Map<
@@ -373,26 +286,17 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
 
       // The initial load is already fetching what a refetch would ask for, and
       // taking it over strands the thread's history if the refetch then fails.
-      if (
-        purpose === "reload" &&
-        loadControllerRef.current?.purpose === "initial"
-      )
-        return loadControllerRef.current.promise ?? Promise.resolve();
-
-      loadControllerRef.current?.controller.abort();
-      const controller = new AbortController();
-      const record: NonNullable<typeof loadControllerRef.current> = {
-        controller,
+      // Aborting a load the runtime no longer needs is not a failure.
+      // A refetch reports the failure to whoever awaited it; the initial load
+      // has no caller to tell.
+      return loadController.run({
         purpose,
-      };
-      loadControllerRef.current = record;
+        load: async (signal) => {
+          const messagesAtLoadStart = messagesRef.current;
+          if (purpose === "initial") setIsLoadingThread(true);
 
-      const messagesAtLoadStart = messagesRef.current;
-      if (purpose === "initial") setIsLoadingThread(true);
-
-      const task = loadFn(externalId, { signal: controller.signal })
-        .then((snapshot) => {
-          if (controller.signal.aborted) return;
+          const snapshot = await loadFn(externalId, { signal });
+          if (signal.aborted) return;
           // A snapshot the session assembled before a run cannot speak for what
           // that run has since produced, and an ADK id cannot correlate a
           // message sent optimistically with the one the session stored for it,
@@ -406,29 +310,16 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           )
             return;
           applySnapshot(snapshot);
-        })
-        .catch((error: unknown) => {
-          // Aborting a load the runtime no longer needs is not a failure.
-          if (controller.signal.aborted) return;
-          throw error;
-        })
-        .finally(() => {
-          if (loadControllerRef.current?.controller === controller) {
-            loadControllerRef.current = null;
-          }
-          if (controller.signal.aborted) return;
+        },
+        onSettled: () => {
           setIsLoadingThread(false);
-        });
-      record.promise = task;
-
-      // A refetch reports the failure to whoever awaited it; the initial load
-      // has no caller to tell.
-      if (purpose === "reload") return task;
-      return task.catch((e: unknown) => {
-        console.warn("Failed to load ADK session:", e);
+        },
+        onInitialError: (error) => {
+          console.warn("Failed to load ADK session:", error);
+        },
       });
     },
-    [threadListItem, applySnapshot],
+    [threadListItem, loadController, applySnapshot],
   );
 
   useEffect(() => {
@@ -436,17 +327,18 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     return () => {
       // Whatever is current, not this effect's own controller: a refetch swaps
       // the ref, and one in flight at unmount must be aborted too.
-      loadControllerRef.current?.controller.abort();
+      loadController.abort();
       setIsLoadingThread(false);
     };
-  }, [runLoad]);
+  }, [loadController, runLoad]);
 
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),
-    isRunning: effectiveIsRunning,
+    isRunning,
     isLoading: isLoadingThread,
     messages: threadMessages,
     unstable_enableToolInvocations: true,
+    unstable_isClientToolCall: isClientToolCall,
     setToolStatuses,
     adapters: { attachments, dictation, feedback, speech, voice },
     extras: adkExtras.provide({
@@ -475,7 +367,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
         [
           ...cancellations,
           {
-            id: uuidv4(),
+            id: generateId(),
             type: "human",
             content: getMessageContent(msg),
           },
@@ -509,7 +401,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           return handleSendMessage(
             [
               {
-                id: uuidv4(),
+                id: generateId(),
                 type: "human",
                 content: getMessageContent(msg),
               },
@@ -564,7 +456,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       await handleSendMessage(
         [
           {
-            id: uuidv4(),
+            id: generateId(),
             type: "tool",
             name: toolName,
             tool_call_id: toolCallId,
@@ -578,7 +470,12 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     },
     onRespondToToolApproval: async (options) => {
       await handleSendMessage(
-        [toAdkToolConfirmationReply(options, toolApprovalsRef.current)],
+        [
+          toAdkToolConfirmationReply(
+            options,
+            projectAdkToolApprovals(adkMessagesRef.current).approvals,
+          ),
+        ],
         {},
       );
     },
@@ -606,11 +503,10 @@ export const useAdkRuntime = ({
   const aui = useAui();
   const cloudAdapter = useCloudThreadListAdapter({
     cloud,
-    create: async () => {
-      if (create) return create();
-      if (aui.threadListItem.source) return aui.threadListItem.initialize();
-      return { externalId: undefined };
-    },
+    create: createCloudThreadListAdapterCreateFallback(
+      create,
+      aui.threadListItem,
+    ),
     delete: deleteFn,
   });