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

package/src/convertAdkMessages.test.ts

@@ -1,4 +1,5 @@
 import { describe, it, expect } from "vitest";
+import { convertExternalMessages } from "@assistant-ui/core/react";
 import { convertAdkMessage } from "./convertAdkMessages";
 import type { AdkMessage } from "./types";
 
@@ -27,6 +28,47 @@ describe("convertAdkMessage - human messages", () => {
     }
   });
 
+  it("skips media parts without their data, url, or mime type", () => {
+    const msg = {
+      id: "m1",
+      type: "human",
+      content: [
+        { type: "image", mimeType: "image/png" },
+        { type: "image", data: "iVBORw0KGgo=" },
+        { type: "image_url" },
+        { type: "file", mimeType: "application/pdf" },
+        { type: "file", data: "QUJD" },
+        { type: "file_url", mimeType: "application/pdf" },
+        { type: "text", text: "Hello" },
+      ],
+    } as AdkMessage;
+    expect(convertAdkMessage(msg, {})).toMatchObject({
+      role: "user",
+      content: [{ type: "text", text: "Hello" }],
+    });
+  });
+
+  it("skips a reasoning part on a human message", () => {
+    const messages = convertExternalMessages<AdkMessage>(
+      [
+        {
+          id: "m1",
+          type: "human",
+          content: [
+            { type: "reasoning", text: "thinking" },
+            { type: "text", text: "Hello" },
+          ],
+        } satisfies AdkMessage,
+      ],
+      (message) => convertAdkMessage(message, {}),
+      false,
+      {},
+    );
+    expect(messages).toMatchObject([
+      { role: "user", content: [{ type: "text", text: "Hello" }] },
+    ]);
+  });
+
   it("converts a human message with text content parts", () => {
     const msg: AdkMessage = {
       id: "m1",
@@ -40,6 +82,19 @@ describe("convertAdkMessage - human messages", () => {
     });
   });
 
+  it("preserves activity messages as visible text", () => {
+    const msg: AdkMessage = {
+      id: "m1",
+      type: "human",
+      content: [{ type: "activity", message: "Working on it" }],
+    };
+
+    expect(convertAdkMessage(msg, {})).toMatchObject({
+      role: "user",
+      content: [{ type: "text", text: "Working on it" }],
+    });
+  });
+
   it("restores a file_url part as a file part stamped with sourceType url", () => {
     const msg: AdkMessage = {
       id: "m1",
@@ -341,6 +396,76 @@ describe("convertAdkMessage - ai messages", () => {
     });
   });
 
+  it("skips null entries in tool_calls", () => {
+    const msg = {
+      id: "m1",
+      type: "ai",
+      content: [],
+      tool_calls: [null, { id: "tc-1", name: "search", args: {} }],
+    } as unknown as AdkMessage;
+    expect(convertAdkMessage(msg, {})).toMatchObject({
+      content: [{ type: "tool-call", toolCallId: "tc-1", toolName: "search" }],
+    });
+  });
+
+  it("skips a tool call without a name, so its named result does not break the thread", () => {
+    const messages = convertExternalMessages<AdkMessage>(
+      [
+        {
+          id: "m1",
+          type: "ai",
+          content: "Checking",
+          tool_calls: [{ id: "tc-1", args: {} }],
+        } as unknown as AdkMessage,
+        {
+          id: "m2",
+          type: "tool",
+          tool_call_id: "tc-1",
+          name: "search",
+          content: "{}",
+        },
+      ],
+      (message) => convertAdkMessage(message, {}),
+      false,
+      {},
+    );
+    expect(messages).toMatchObject([
+      { role: "assistant", content: [{ type: "text", text: "Checking" }] },
+    ]);
+  });
+
+  it.each([
+    ["an empty", ""],
+    ["a non-string", 42],
+  ])(
+    "skips a tool call with %s name, so its named result does not break the thread",
+    (_, name) => {
+      const messages = convertExternalMessages<AdkMessage>(
+        [
+          {
+            id: "m1",
+            type: "ai",
+            content: "Checking",
+            tool_calls: [{ id: "tc-1", name, args: {} }],
+          } as AdkMessage,
+          {
+            id: "m2",
+            type: "tool",
+            tool_call_id: "tc-1",
+            name: "search",
+            content: "{}",
+          },
+        ],
+        (message) => convertAdkMessage(message, {}),
+        false,
+        {},
+      );
+      expect(messages).toMatchObject([
+        { role: "assistant", content: [{ type: "text", text: "Checking" }] },
+      ]);
+    },
+  );
+
   it("includes status when present", () => {
     const msg: AdkMessage = {
       id: "m1",
@@ -414,3 +539,20 @@ describe("convertAdkMessage - tool messages", () => {
     expect(result).toMatchObject({ isError: true });
   });
 });
+
+describe("convertAdkMessage - unknown messages", () => {
+  it("skips a message without a type", () => {
+    const messages = convertExternalMessages<AdkMessage>(
+      [
+        { id: "m1", type: "human", content: "Hello" },
+        { id: "m2", content: "no type" } as unknown as AdkMessage,
+      ],
+      (message) => convertAdkMessage(message, {}),
+      false,
+      {},
+    );
+    expect(messages).toMatchObject([
+      { role: "user", content: [{ type: "text", text: "Hello" }] },
+    ]);
+  });
+});

package/src/convertAdkMessages.ts

@@ -42,18 +42,30 @@ const contentToParts = (
             text: typeof part.text === "string" ? part.text : "",
           };
         case "reasoning":
+          if (role === "user") return null;
           return {
             type: "reasoning",
             text: typeof part.text === "string" ? part.text : "",
           };
         case "image":
+          if (
+            typeof part.mimeType !== "string" ||
+            typeof part.data !== "string"
+          )
+            return null;
           return {
             type: "image",
             image: `data:${part.mimeType};base64,${part.data}`,
           };
         case "image_url":
+          if (typeof part.url !== "string") return null;
           return { type: "image", image: part.url };
         case "file":
+          if (
+            typeof part.mimeType !== "string" ||
+            typeof part.data !== "string"
+          )
+            return null;
           return {
             type: "file",
             data: part.data,
@@ -61,6 +73,7 @@ const contentToParts = (
             ...(part.filename != null && { filename: part.filename }),
           };
         case "file_url":
+          if (typeof part.url !== "string") return null;
           if (role === "user") {
             return {
               type: "file",
@@ -89,6 +102,11 @@ const contentToParts = (
             name: "code_execution_result",
             data: { output: part.output, outcome: part.outcome },
           };
+        case "activity":
+          return {
+            type: "text",
+            text: typeof part.message === "string" ? part.message : "",
+          };
         default:
           return null;
       }
@@ -113,7 +131,8 @@ export const createAdkMessageConverter =
 
       case "ai": {
         const toolCallParts: ToolCallMessagePart[] =
-          message.tool_calls?.map((tc) => {
+          message.tool_calls?.flatMap((tc) => {
+            if (typeof tc?.name !== "string" || tc.name.length === 0) return [];
             const approval = approvals.get(tc.id);
             return {
               type: "tool-call",
@@ -164,6 +183,10 @@ export const createAdkMessageConverter =
           isError: message.status === "error",
         };
       }
+
+      default:
+        message satisfies never;
+        return [];
     }
   };
 

package/src/convertToAdkMessages.test.ts

@@ -81,6 +81,20 @@ describe("getPendingToolCalls", () => {
 });
 
 describe("getPendingCancellations", () => {
+  it("skips a null tool_calls entry and cancels the rest", () => {
+    const messages = [
+      {
+        id: "ai-1",
+        type: "ai",
+        content: [],
+        tool_calls: [null, { id: "tc-1", name: "tool_a", args: {} }],
+      },
+    ] as unknown as AdkMessage[];
+    expect(getPendingCancellations(messages, [])).toMatchObject([
+      { type: "tool", name: "tool_a", tool_call_id: "tc-1", status: "error" },
+    ]);
+  });
+
   it("emits a {cancelled:true} tool message for every pending tool call", () => {
     const messages: AdkMessage[] = [
       aiWithToolCalls("ai-1", [{ id: "tc-1", name: "tool_a" }]),
@@ -156,6 +170,12 @@ describe("getPendingCancellations", () => {
 });
 
 describe("getMessageContent", () => {
+  it("sends activity messages as text parts", () => {
+    expect(
+      contentToParts([{ type: "activity", message: "Working on it" }]),
+    ).toEqual([{ text: "Working on it" }]);
+  });
+
   it("serializes data URL images as inline data", () => {
     const content = getMessageContent(
       makeAppendMessage([

package/src/convertToAdkMessages.ts

@@ -1,3 +1,4 @@
+import { isRecord } from "@assistant-ui/core/internal";
 import {
   generateId,
   getExternalStoreMessages,
@@ -94,7 +95,9 @@ export const getPendingToolCalls = (messages: AdkMessage[]) => {
     messages,
     (message) => {
       if (message.type === "ai") {
-        return { toolCalls: message.tool_calls ?? [] };
+        return {
+          toolCalls: (message.tool_calls ?? []).filter(isRecord),
+        };
       }
       if (message.type === "tool") {
         return { toolCallId: message.tool_call_id };

package/src/tests/settleOutsideAct.ts

@@ -0,0 +1,15 @@
+import { act, waitFor } from "@testing-library/react";
+
+// React 18's act holds renders until its callback settles, so a task that waits on a render starts inside act and settles outside it, where React can commit what the task waits on.
+export async function settleOutsideAct<T>(task: () => Promise<T>): Promise<T> {
+  let pending!: Promise<T>;
+  await act(async () => {
+    pending = task();
+    pending.catch(() => {});
+  });
+  await waitFor(() => pending.then(settled, settled));
+  await act(async () => {});
+  return pending;
+}
+
+const settled = () => {};

package/src/useAdkMessages.fast-refresh.test.tsx

@@ -0,0 +1,114 @@
+// @vitest-environment jsdom
+
+import { act, StrictMode, useEffect } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+let renderer: RendererInternals | undefined;
+const fiberRoots = new Set<unknown>();
+vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+  supportsFiber: true,
+  inject: (internals: RendererInternals) => {
+    renderer = internals;
+    return 1;
+  },
+  onScheduleFiberRoot: () => {},
+  onCommitFiberRoot: (_id: number, root: unknown) => fiberRoots.add(root),
+  onCommitFiberUnmount: () => {},
+});
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+const { useAdkMessages } = await import("./useAdkMessages");
+
+afterEach(() => {
+  cleanup();
+  renderer?.setRefreshHandler(() => undefined);
+  fiberRoots.clear();
+});
+afterAll(() => vi.unstubAllGlobals());
+
+const refresh = async (Before: unknown, After: unknown) => {
+  const family: Family = { current: After };
+  renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const root of fiberRoots) {
+      renderer!.scheduleRefresh(root, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+};
+
+it("keeps a direct useAdkMessages stream through Fast Refresh and cancels on unmount", async () => {
+  let signal: AbortSignal | undefined;
+  let send: (() => void) | undefined;
+  let rendered: string | undefined;
+  const stream = vi.fn((_messages, config: { abortSignal: AbortSignal }) => {
+    signal = config.abortSignal;
+    return new Promise<never>(() => {});
+  });
+  const host = (name: string) => () => {
+    rendered = name;
+    const { sendMessage } = useAdkMessages({ stream: stream as never });
+    send = () => {
+      void sendMessage([{ id: "user", type: "human", content: "hello" }], {});
+    };
+    return null;
+  };
+  const Before = host("before");
+  const After = host("after");
+  const view = render(<Before />);
+
+  act(() => send!());
+  await waitFor(() => expect(stream).toHaveBeenCalledTimes(1));
+  expect(signal!.aborted).toBe(false);
+
+  await refresh(Before, After);
+  expect(rendered).toBe("after");
+  expect(signal!.aborted).toBe(false);
+  expect(stream).toHaveBeenCalledTimes(1);
+
+  view.unmount();
+  await act(async () => {});
+  expect(signal!.aborted).toBe(true);
+});
+
+it("keeps a stream started during StrictMode's effect replay", async () => {
+  let signal: AbortSignal | undefined;
+  let launched = false;
+  const stream = vi.fn((_messages, config: { abortSignal: AbortSignal }) => {
+    signal = config.abortSignal;
+    return new Promise<never>(() => {});
+  });
+  const Host = () => {
+    const { sendMessage } = useAdkMessages({ stream: stream as never });
+    useEffect(() => {
+      if (launched) return;
+      launched = true;
+      void sendMessage([{ id: "user", type: "human", content: "hello" }], {});
+    }, [sendMessage]);
+    return null;
+  };
+  const view = render(
+    <StrictMode>
+      <Host />
+    </StrictMode>,
+  );
+  await waitFor(() => expect(stream).toHaveBeenCalledTimes(1));
+  expect(signal!.aborted).toBe(false);
+  view.unmount();
+  await act(async () => {});
+  expect(signal!.aborted).toBe(true);
+});

package/src/useAdkMessages.test.ts

@@ -77,6 +77,40 @@ describe("optimistic tool outcomes", () => {
 });
 
 describe("ADK runtime callbacks", () => {
+  it("reports the same agent transfer again in a later run", async () => {
+    const onAgentTransfer = vi.fn();
+    const stream: AdkStreamCallback = async function* () {
+      yield {
+        id: "transfer",
+        actions: { transferToAgent: "researcher" },
+      };
+      yield {
+        id: "transfer-duplicate",
+        actions: { transferToAgent: "researcher" },
+      };
+    };
+    const { result } = renderHook(() =>
+      useAdkMessages({ stream, eventHandlers: { onAgentTransfer } }),
+    );
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "first" }],
+        {},
+      );
+      expect(onAgentTransfer).toHaveBeenCalledTimes(1);
+
+      await result.current.sendMessage(
+        [{ id: "user-2", type: "human", content: "second" }],
+        {},
+      );
+    });
+
+    expect(onAgentTransfer).toHaveBeenCalledTimes(2);
+    expect(onAgentTransfer).toHaveBeenNthCalledWith(1, "researcher");
+    expect(onAgentTransfer).toHaveBeenNthCalledWith(2, "researcher");
+  });
+
   it.each(["onAgentTransfer", "onCustomEvent", "onError"] as const)(
     "continues streaming when %s throws",
     async (callbackName) => {
@@ -305,6 +339,7 @@ describe("ADK stream lifecycle", () => {
     await started;
 
     unmount();
+    await act(async () => {});
 
     expect(runSignal?.aborted).toBe(true);
     await expect(sendPromise).resolves.toBeUndefined();
@@ -765,6 +800,19 @@ describe("optimistic multi-message sends", () => {
 });
 
 describe("messageToEvent (contentToParts)", () => {
+  it("skips a null tool_calls entry", () => {
+    const event = messageToEvent({
+      id: "ai-1",
+      type: "ai",
+      content: [],
+      tool_calls: [null, { id: "tc-1", name: "search", args: { q: "x" } }],
+    } as unknown as AdkMessage);
+
+    expect(event.content?.parts).toEqual([
+      { functionCall: { name: "search", id: "tc-1", args: { q: "x" } } },
+    ]);
+  });
+
   it.each([
     ["scalar", "false", { result: false }],
     ["array", "[1,2]", { results: [1, 2] }],

package/src/useAdkMessages.ts

@@ -1,12 +1,13 @@
+import { isRecord } from "@assistant-ui/core/internal";
 import {
   useState,
   useCallback,
-  useEffect,
   useInsertionEffect,
   useRef,
   useMemo,
 } from "react";
 import { generateId } from "@assistant-ui/core";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 import { useAui } from "@assistant-ui/store";
 import {
   abortableIterable,
@@ -71,7 +72,6 @@ export const useAdkMessages = ({
   const [messageMetadata, setMessageMetadata] = useState<
     Map<string, AdkMessageMetadata>
   >(new Map());
-  const lastTransferToAgentRef = useRef<string | undefined>(undefined);
   // setMessagesImmediate and setLongRunningToolIds are the only writers of their state and publish these refs with it, so neither ref trails a commit.
   const messagesRef = useRef(messages);
   const longRunningToolIdsRef = useRef(longRunningToolIds);
@@ -168,10 +168,14 @@ export const useAdkMessages = ({
       for (const event of messagesToEvents(newMessagesWithId)) {
         accumulator.processEvent(event);
       }
-      setMessagesImmediate(accumulator.getMessages());
-      setLongRunningToolIds(accumulator.getLongRunningToolIds());
+      const initialMessages = accumulator.getMessages();
+      const initialMessageIds = new Set(initialMessages.map((m) => m.id));
+      const initialLongRunningToolIds = accumulator.getLongRunningToolIds();
+      setMessagesImmediate(initialMessages);
+      setLongRunningToolIds(initialLongRunningToolIds);
       setToolConfirmations(accumulator.getToolConfirmations());
       setAuthRequests(accumulator.getAuthRequests());
+      let lastTransferToAgent: string | undefined;
 
       // Google ADK replaces active runs, while React LangGraph queues sends.
       abortControllerRef.current?.abort();
@@ -226,8 +230,8 @@ export const useAdkMessages = ({
           }
 
           const transfer = accumulator.getLastTransferToAgent();
-          if (transfer && transfer !== lastTransferToAgentRef.current) {
-            lastTransferToAgentRef.current = transfer;
+          if (transfer && transfer !== lastTransferToAgent) {
+            lastTransferToAgent = transfer;
             invokeAdkRuntimeCallback(
               "onAgentTransfer",
               onAgentTransfer,
@@ -265,6 +269,33 @@ export const useAdkMessages = ({
         }
       } finally {
         if (abortControllerRef.current === abortController) {
+          if (abortController.signal.aborted) {
+            setLongRunningToolIds(
+              accumulator
+                .getLongRunningToolIds()
+                .filter((id) => initialLongRunningToolIds.includes(id)),
+            );
+            const updatedMessages = messagesRef.current;
+            const lastAssistantMessage = updatedMessages.findLast(
+              (m) => m.type === "ai",
+            );
+            if (
+              lastAssistantMessage &&
+              !initialMessageIds.has(lastAssistantMessage.id) &&
+              !lastAssistantMessage.status
+            ) {
+              setMessagesImmediate(
+                updatedMessages.map((m) =>
+                  m === lastAssistantMessage
+                    ? {
+                        ...lastAssistantMessage,
+                        status: { type: "incomplete", reason: "cancelled" },
+                      }
+                    : m,
+                ),
+              );
+            }
+          }
           abortControllerRef.current = null;
         }
       }
@@ -286,7 +317,7 @@ export const useAdkMessages = ({
     }
   }, []);
 
-  useEffect(() => cancel, [cancel]);
+  useReplaySafeEffect(() => cancel, []);
 
   return {
     messages,
@@ -395,9 +426,9 @@ export const messageToEvent = (msg: AdkMessage): AdkEvent => {
     role: "model",
     parts: [
       ...contentToParts(msg.content),
-      ...(msg.tool_calls?.map((tc) => ({
+      ...(msg.tool_calls ?? []).filter(isRecord).map((tc) => ({
         functionCall: { name: tc.name, id: tc.id, args: { ...tc.args } },
-      })) ?? []),
+      })),
     ],
   };
   return result;