@assistant-ui/react-google-adk 0.0.25 → 0.0.27

package/src/useAdkRuntime.ts

@@ -1,6 +1,5 @@
-import { useCallback, useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
-  getExternalStoreMessages,
   pickExternalStoreSharedOptions,
   type AttachmentAdapter,
   type DictationAdapter,
@@ -9,11 +8,13 @@ import {
   type RealtimeVoiceAdapter,
   type SpeechSynthesisAdapter,
   type AppendMessage,
-  type ThreadMessage,
   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,
@@ -33,136 +34,21 @@ import type {
   OnAdkAgentTransferCallback,
 } from "./types";
 import { useAdkMessages } from "./useAdkMessages";
-import { convertAdkMessage } from "./convertAdkMessages";
+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;
@@ -250,11 +136,7 @@ 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);
+  const loadController = useMemo(createAbortableThreadLoad, []);
   const messagesRef = useRef(messages);
   messagesRef.current = messages;
   const [isLoadingThread, setIsLoadingThread] = useState(
@@ -285,8 +167,21 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     }
   };
 
+  const { approvals: toolApprovals, key: toolApprovalsKey } =
+    projectAdkToolApprovals(messages);
+  const toolApprovalsRef = useRef(toolApprovals);
+  toolApprovalsRef.current = toolApprovals;
+
+  const messageConverter = useMemo(
+    () =>
+      toolApprovalsKey === ""
+        ? convertAdkMessage
+        : createAdkMessageConverter(toolApprovalsRef.current),
+    [toolApprovalsKey],
+  );
+
   const threadMessages = useExternalMessageConverter({
-    callback: convertAdkMessage,
+    callback: messageConverter,
     messages,
     isRunning: effectiveIsRunning,
   });
@@ -353,26 +248,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;
-
-      const messagesAtLoadStart = messagesRef.current;
-      if (purpose === "initial") setIsLoadingThread(true);
+        load: async (signal) => {
+          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,
@@ -386,29 +272,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(() => {
@@ -416,14 +289,14 @@ 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,
@@ -455,7 +328,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
         [
           ...cancellations,
           {
-            id: uuidv4(),
+            id: generateId(),
             type: "human",
             content: getMessageContent(msg),
           },
@@ -489,7 +362,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           return handleSendMessage(
             [
               {
-                id: uuidv4(),
+                id: generateId(),
                 type: "human",
                 content: getMessageContent(msg),
               },
@@ -544,7 +417,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       await handleSendMessage(
         [
           {
-            id: uuidv4(),
+            id: generateId(),
             type: "tool",
             name: toolName,
             tool_call_id: toolCallId,
@@ -556,6 +429,12 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
         {},
       );
     },
+    onRespondToToolApproval: async (options) => {
+      await handleSendMessage(
+        [toAdkToolConfirmationReply(options, toolApprovalsRef.current)],
+        {},
+      );
+    },
     onCancel: unstable_allowCancellation
       ? async () => {
           cancel();
@@ -580,11 +459,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,
   });
 

package/src/useAdkRuntimeApproval.test.tsx

@@ -0,0 +1,344 @@
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type {
+  RespondToToolApprovalOptions,
+  ThreadMessage,
+  ToolCallMessagePart,
+} from "@assistant-ui/core";
+import type { AdkMessage } from "./types";
+
+const mocks = vi.hoisted(() => ({
+  adapters: [] as unknown[],
+  sendMessage: vi.fn().mockResolvedValue(undefined),
+  messages: [] as AdkMessage[],
+}));
+
+vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/core/react")>()),
+  useCloudThreadListAdapter: () => ({}),
+  useExternalStoreRuntime: (adapter: unknown) => {
+    mocks.adapters.push(adapter);
+    return {};
+  },
+  useRemoteThreadListRuntime: (options: { runtimeHook: () => unknown }) =>
+    options.runtimeHook(),
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => ({
+    threadListItem: {
+      source: null,
+      getState: () => ({ externalId: undefined }),
+      initialize: vi.fn(),
+    },
+  }),
+}));
+
+vi.mock("./useAdkMessages", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("./useAdkMessages")>()),
+  useAdkMessages: () => ({
+    messages: mocks.messages,
+    stateDelta: {},
+    agentInfo: {},
+    longRunningToolIds: [],
+    artifactDelta: {},
+    // Deliberately empty: the projection must come from the transcript, not
+    // from derived confirmation state that a mid-run accumulator swap drops.
+    toolConfirmations: [],
+    authRequests: [],
+    escalated: false,
+    messageMetadata: new Map(),
+    sendMessage: mocks.sendMessage,
+    cancel: vi.fn(),
+    setMessages: vi.fn(),
+    replaceMessages: vi.fn(),
+    applySnapshot: vi.fn(),
+  }),
+}));
+
+import { AdkEventAccumulator } from "./AdkEventAccumulator";
+import { useAdkRuntime } from "./useAdkRuntime";
+
+type ApprovalAdapter = {
+  messages: readonly ThreadMessage[];
+  onRespondToToolApproval?: (
+    options: RespondToToolApprovalOptions,
+  ) => Promise<void> | void;
+};
+
+const CONFIRMATION_CALL = "adk-confirmation-1";
+
+const latestAdapter = () => mocks.adapters.at(-1) as ApprovalAdapter;
+
+const makeConfirmationRequest = (): AdkMessage => ({
+  id: "ai-1",
+  type: "ai",
+  content: [],
+  tool_calls: [
+    {
+      id: CONFIRMATION_CALL,
+      name: "adk_request_confirmation",
+      args: {
+        originalFunctionCall: { id: "adk-original-1", name: "delete_file" },
+        toolConfirmation: { hint: "Delete /tmp/a?" },
+      },
+    },
+  ],
+});
+
+const approvalPart = () =>
+  latestAdapter()
+    .messages.at(-1)!
+    .content.find(
+      (part) => part.type === "tool-call" && part.approval !== undefined,
+    );
+
+afterEach(() => {
+  mocks.adapters.length = 0;
+  mocks.messages = [];
+  vi.clearAllMocks();
+});
+
+describe("useAdkRuntime tool approvals", () => {
+  it("exposes, answers, and settles the default approval seam across a rerender", async () => {
+    // Retained across the rerender: core caches converted messages by input
+    // object, so only a rebuilt converter can surface the settled decision.
+    const confirmationRequest = makeConfirmationRequest();
+    mocks.messages = [
+      { id: "u-1", type: "human", content: "delete the file" },
+      confirmationRequest,
+    ];
+
+    const { rerender } = renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    expect(latestAdapter().messages.at(-1)!.status).toMatchObject({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+    expect(approvalPart()).toMatchObject({
+      toolCallId: CONFIRMATION_CALL,
+      approval: { id: CONFIRMATION_CALL },
+    });
+
+    await act(async () => {
+      await latestAdapter().onRespondToToolApproval!({
+        approvalId: CONFIRMATION_CALL,
+        approved: false,
+      });
+    });
+
+    expect(mocks.sendMessage.mock.calls.at(-1)![0]).toEqual([
+      expect.objectContaining({
+        type: "tool",
+        tool_call_id: CONFIRMATION_CALL,
+        name: "adk_request_confirmation",
+        content: JSON.stringify({ confirmed: false }),
+      }),
+    ]);
+
+    mocks.messages = [
+      ...mocks.messages,
+      {
+        id: "tool-1",
+        type: "tool",
+        tool_call_id: CONFIRMATION_CALL,
+        name: "adk_request_confirmation",
+        content: JSON.stringify({ confirmed: false }),
+        status: "success",
+      },
+    ];
+    rerender();
+
+    expect(approvalPart()).toMatchObject({
+      approval: { id: CONFIRMATION_CALL, approved: false },
+    });
+    expect(latestAdapter().messages.at(-1)!.status).not.toMatchObject({
+      type: "requires-action",
+    });
+
+    await expect(
+      latestAdapter().onRespondToToolApproval!({
+        approvalId: CONFIRMATION_CALL,
+        approved: true,
+      }),
+    ).rejects.toThrow("No pending ADK tool confirmation");
+  });
+
+  it("keeps a gate answered by an unreadable reply retryable at the runtime seam", async () => {
+    mocks.messages = [
+      { id: "u-1", type: "human", content: "delete the file" },
+      makeConfirmationRequest(),
+      {
+        id: "tool-1",
+        type: "tool",
+        tool_call_id: CONFIRMATION_CALL,
+        name: "adk_request_confirmation",
+        // ADK parses the wrapped text without a `try`, so this reply raises
+        // rather than denying, and the gate stays answerable.
+        content: JSON.stringify({ response: "not-json" }),
+        status: "success",
+      },
+    ];
+
+    renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    expect(latestAdapter().messages.at(-1)!.status).toMatchObject({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+    const part = approvalPart() as { result?: unknown; approval: unknown };
+    expect(part.result).toBeUndefined();
+    expect(part.approval).toEqual({ id: CONFIRMATION_CALL });
+
+    await act(async () => {
+      await latestAdapter().onRespondToToolApproval!({
+        approvalId: CONFIRMATION_CALL,
+        approved: true,
+      });
+    });
+
+    expect(mocks.sendMessage.mock.calls.at(-1)![0]).toEqual([
+      expect.objectContaining({
+        tool_call_id: CONFIRMATION_CALL,
+        content: JSON.stringify({ confirmed: true }),
+      }),
+    ]);
+  });
+
+  it("keeps both gates of an event retryable when one reply is unreadable", () => {
+    const accumulator = new AdkEventAccumulator();
+    const confirmationCall = (id: string, tool: string) => ({
+      functionCall: {
+        id,
+        name: "adk_request_confirmation",
+        args: {
+          originalFunctionCall: { id: `original-${id}`, name: tool },
+          toolConfirmation: { hint: `Run ${tool}?` },
+        },
+      },
+    });
+    accumulator.processEvent({
+      id: "evt-request",
+      author: "agent",
+      longRunningToolIds: ["conf-a", "conf-b"],
+      content: {
+        role: "model",
+        parts: [
+          confirmationCall("conf-a", "delete_file"),
+          confirmationCall("conf-b", "send_email"),
+        ],
+      },
+    });
+    mocks.messages = accumulator.processEvent({
+      id: "evt-reply",
+      author: "user",
+      content: {
+        role: "user",
+        parts: [
+          {
+            functionResponse: {
+              id: "conf-a",
+              name: "adk_request_confirmation",
+              response: { confirmed: true },
+            },
+          },
+          {
+            functionResponse: {
+              id: "conf-b",
+              name: "adk_request_confirmation",
+              response: { response: "not json" },
+            },
+          },
+        ],
+      },
+    });
+
+    renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    const assistant = latestAdapter().messages.find(
+      (message): message is Extract<ThreadMessage, { role: "assistant" }> =>
+        message.role === "assistant",
+    )!;
+    const gates = assistant.content.filter(
+      (part): part is ToolCallMessagePart =>
+        part.type === "tool-call" && part.approval !== undefined,
+    );
+
+    expect(gates.map((gate) => gate.approval)).toEqual([
+      { id: "conf-a" },
+      { id: "conf-b" },
+    ]);
+    expect(gates.map((gate) => gate.result)).toEqual([undefined, undefined]);
+    expect(assistant.status).toMatchObject({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+  });
+
+  it("settles a confirmation reply carried beside user text without an orphan message", () => {
+    const accumulator = new AdkEventAccumulator();
+    accumulator.processEvent({
+      id: "evt-request",
+      author: "agent",
+      longRunningToolIds: [CONFIRMATION_CALL],
+      content: {
+        role: "model",
+        parts: [
+          {
+            functionCall: {
+              id: CONFIRMATION_CALL,
+              name: "adk_request_confirmation",
+              args: {
+                originalFunctionCall: {
+                  id: "adk-original-1",
+                  name: "delete_file",
+                },
+                toolConfirmation: { hint: "Delete /tmp/a?" },
+              },
+            },
+          },
+        ],
+      },
+    });
+    mocks.messages = accumulator.processEvent({
+      id: "evt-reply",
+      author: "user",
+      content: {
+        role: "user",
+        parts: [
+          { text: "go ahead" },
+          {
+            functionResponse: {
+              id: CONFIRMATION_CALL,
+              name: "adk_request_confirmation",
+              response: { confirmed: true },
+            },
+          },
+        ],
+      },
+    });
+
+    renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    const messages = latestAdapter().messages;
+    const assistant = messages.find(
+      (message): message is Extract<ThreadMessage, { role: "assistant" }> =>
+        message.role === "assistant",
+    );
+    const gate = assistant?.content.find(
+      (part) => part.type === "tool-call" && part.approval !== undefined,
+    ) as ToolCallMessagePart | undefined;
+    expect(gate?.approval).toEqual({
+      id: CONFIRMATION_CALL,
+      approved: true,
+    });
+    expect(
+      messages.some(
+        (message) =>
+          message.role === "assistant" && message.content.length === 0,
+      ),
+    ).toBe(false);
+    expect(messages.at(-1)!.role).toBe("user");
+  });
+});