@assistant-ui/react-google-adk 0.0.20 → 0.0.21

package/src/useAdkRuntime.test.ts

@@ -4,6 +4,7 @@ import {
   getPendingCancellations,
   getPendingToolCalls,
 } from "./useAdkRuntime";
+import { convertAdkMessage } from "./convertAdkMessages";
 import type { AppendMessage } from "@assistant-ui/core";
 import type { AdkMessage } from "./types";
 
@@ -209,6 +210,89 @@ describe("getMessageContent", () => {
     ]);
   });
 
+  it("emits a file_url part for file parts with sourceType url", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          mimeType: "application/pdf",
+          data: "gs://bucket/report.pdf",
+          filename: "report.pdf",
+          sourceType: "url",
+        },
+      ]),
+    );
+    expect(result).toEqual([
+      {
+        type: "file_url",
+        url: "gs://bucket/report.pdf",
+        mimeType: "application/pdf",
+      },
+    ]);
+  });
+
+  it("keeps file parts inline without sourceType", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          mimeType: "application/pdf",
+          data: "gs://bucket/report.pdf",
+        },
+      ]),
+    );
+    expect(result).toEqual([
+      {
+        type: "file",
+        mimeType: "application/pdf",
+        data: "gs://bucket/report.pdf",
+      },
+    ]);
+  });
+
+  it("round-trips a file_url part through convert and edit-resend", () => {
+    const converted = convertAdkMessage(
+      {
+        id: "m1",
+        type: "human",
+        content: [
+          {
+            type: "file_url",
+            url: "gs://bucket/report.pdf",
+            mimeType: "application/pdf",
+          },
+        ],
+      },
+      {},
+    );
+    const content = (converted as { content: AppendMessage["content"] })
+      .content;
+    const result = getMessageContent(makeAppendMessage(content));
+    expect(result).toEqual([
+      {
+        type: "file_url",
+        url: "gs://bucket/report.pdf",
+        mimeType: "application/pdf",
+      },
+    ]);
+  });
+
+  it("ignores sourceType id on file parts", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          mimeType: "application/pdf",
+          data: "file-abc123",
+          sourceType: "id",
+        },
+      ]),
+    );
+    expect(result).toEqual([
+      { type: "file", mimeType: "application/pdf", data: "file-abc123" },
+    ]);
+  });
+
   it("forwards an audio part as a file block with the format-derived mime type", () => {
     const result = getMessageContent(
       makeAppendMessage([
@@ -245,6 +329,83 @@ describe("getMessageContent", () => {
     ]);
   });
 
+  it("strips a data URL envelope from file data", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          data: "data:application/pdf;base64,QUJD",
+          mimeType: "application/pdf",
+          filename: "a.pdf",
+        },
+      ]),
+    );
+    expect(result).toEqual([
+      {
+        type: "file",
+        mimeType: "application/pdf",
+        data: "QUJD",
+        filename: "a.pdf",
+      },
+    ]);
+  });
+
+  it("emits a file_url part for an unmarked http source", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          data: "https://cdn.example.com/a.pdf",
+          mimeType: "application/pdf",
+        },
+      ]),
+    );
+    expect(result).toEqual([
+      {
+        type: "file_url",
+        url: "https://cdn.example.com/a.pdf",
+        mimeType: "application/pdf",
+      },
+    ]);
+  });
+
+  it("leaves bare base64 file data untouched", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        { type: "file", data: "QUJD", mimeType: "application/pdf" },
+      ]),
+    );
+    expect(result).toEqual([
+      { type: "file", mimeType: "application/pdf", data: "QUJD" },
+    ]);
+  });
+
+  it("round-trips an audio file part through both converters", () => {
+    const outbound = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          data: "data:audio/mp3;base64,QUJD",
+          mimeType: "audio/mp3",
+        },
+      ]),
+    );
+
+    expect(outbound).toEqual([
+      { type: "file", mimeType: "audio/mp3", data: "QUJD" },
+    ]);
+
+    const inbound = convertAdkMessage(
+      { id: "m1", type: "human", content: outbound } as never,
+      {},
+    );
+
+    expect(inbound).toMatchObject({
+      role: "user",
+      content: [{ type: "file", data: "QUJD", mimeType: "audio/mp3" }],
+    });
+  });
+
   it("skips data parts while keeping surrounding text", () => {
     const result = getMessageContent(
       makeAppendMessage([

package/src/useAdkRuntime.ts

@@ -1,4 +1,4 @@
-import { useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
 import {
   getExternalStoreMessages,
   pickExternalStoreSharedOptions,
@@ -13,7 +13,7 @@ import {
   type ToolExecutionStatus,
   generateId,
 } from "@assistant-ui/core";
-import { parseDataUrl } from "@assistant-ui/core/internal";
+import { httpUrlPattern, parseDataUrl } from "@assistant-ui/core/internal";
 import {
   useCloudThreadListAdapter,
   useRemoteThreadListRuntime,
@@ -25,6 +25,7 @@ import type { AssistantCloud } from "assistant-cloud";
 import type { RemoteThreadListAdapter } from "@assistant-ui/core";
 import type {
   AdkMessage,
+  AdkThreadSnapshot,
   AdkSendMessageConfig,
   AdkStreamCallback,
   OnAdkErrorCallback,
@@ -50,10 +51,19 @@ export const getMessageContent = (msg: AppendMessage) => {
       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,
-          data: part.data,
+          // 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": {
@@ -171,7 +181,16 @@ export type UseAdkRuntimeOptions = ExternalStoreSharedOptions & {
     threadId: string,
     parentMessages: AdkMessage[],
   ) => Promise<string | null>;
-  load?: (threadId: string) => Promise<{ messages: AdkMessage[] }>;
+  /**
+   * Loads a thread's stored state. Called when the thread opens, and again for
+   * `threads.reloadMainThread()`, which refetches in place rather than
+   * remounting the runtime; the signal aborts a load the runtime no longer
+   * needs.
+   */
+  load?: (
+    threadId: string,
+    options?: { signal?: AbortSignal | undefined },
+  ) => Promise<AdkThreadSnapshot>;
   create?: () => Promise<{ externalId: string }>;
   delete?: (threadId: string) => Promise<void>;
   adapters?:
@@ -223,11 +242,26 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     cancel,
     setMessages,
     replaceMessages,
+    applySnapshot,
   } = useAdkMessages({
     stream,
     ...(eventHandlers && { eventHandlers }),
   });
 
+  const loadRef = useRef(load);
+  loadRef.current = load;
+  const loadControllerRef = useRef<{
+    controller: AbortController;
+    purpose: "initial" | "reload";
+    promise?: Promise<void> | undefined;
+  } | null>(null);
+  const messagesRef = useRef(messages);
+  messagesRef.current = messages;
+  const [isLoadingThread, setIsLoadingThread] = useState(
+    () =>
+      load !== undefined && aui.threadListItem.getState().externalId != null,
+  );
+
   const [isRunning, setIsRunning] = useState(false);
   const [toolStatuses, setToolStatuses] = useState<
     Record<string, ToolExecutionStatus>
@@ -236,6 +270,8 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     (s) => s?.type === "executing",
   );
   const effectiveIsRunning = isRunning || hasExecutingTools;
+  const isRunningRef = useRef(effectiveIsRunning);
+  isRunningRef.current = effectiveIsRunning;
 
   const handleSendMessage = async (
     msgs: AdkMessage[],
@@ -302,9 +338,93 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     setMessages(nextMessages);
   };
 
+  // The scoped client, not `aui` itself: useAui returns a render-bound
+  // instance, so depending on it would re-run the load on every render.
+  const threadListItem =
+    aui.threadListItem.source !== null ? aui.threadListItem : undefined;
+
+  const runLoad = useCallback(
+    (purpose: "initial" | "reload" = "initial") => {
+      const loadFn = loadRef.current;
+      if (!loadFn || !threadListItem) return Promise.resolve();
+
+      const externalId = threadListItem.getState().externalId;
+      if (externalId == null) return Promise.resolve();
+
+      // 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,
+        purpose,
+      };
+      loadControllerRef.current = record;
+
+      const messagesAtLoadStart = messagesRef.current;
+      if (purpose === "initial") setIsLoadingThread(true);
+
+      const task = loadFn(externalId, { signal: controller.signal })
+        .then((snapshot) => {
+          if (controller.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,
+          // so there is nothing here that could merge the two. A refetch that
+          // raced a run therefore defers to the run, whether the run started
+          // during the load or was already streaming when it began.
+          if (
+            purpose === "reload" &&
+            (isRunningRef.current ||
+              messagesRef.current !== messagesAtLoadStart)
+          )
+            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;
+          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);
+      });
+    },
+    [threadListItem, applySnapshot],
+  );
+
+  useEffect(() => {
+    runLoad();
+    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();
+      setIsLoadingThread(false);
+    };
+  }, [runLoad]);
+
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),
     isRunning: effectiveIsRunning,
+    isLoading: isLoadingThread,
     messages: threadMessages,
     unstable_enableToolInvocations: true,
     setToolStatuses,
@@ -441,32 +561,11 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           cancel();
         }
       : undefined,
+    ...(load !== undefined && {
+      onRefetchThread: () => runLoad("reload"),
+    }),
   });
 
-  {
-    const loadRef = useRef(load);
-    useEffect(() => {
-      loadRef.current = load;
-    });
-
-    useEffect(() => {
-      const loadFn = loadRef.current;
-      if (!loadFn) return;
-
-      const externalId = aui.threadListItem.getState().externalId;
-      if (externalId == null) return;
-
-      loadFn(externalId).then(
-        ({ messages: msgs }) => {
-          replaceMessages(msgs);
-        },
-        (e) => {
-          console.warn("Failed to load ADK session:", e);
-        },
-      );
-    }, [aui, replaceMessages]);
-  }
-
   return runtime;
 };