@assistant-ui/react-google-adk 0.0.33 → 0.0.35

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",
@@ -118,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",
@@ -169,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" }]),

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/useAdkMessages.test.ts

@@ -17,6 +17,7 @@ import {
   messageToEvent,
   messagesToEvents,
   useAdkMessages,
+  useAdkMessagesInternal,
 } from "./useAdkMessages";
 import { projectAdkToolApprovals } from "./adkToolApproval";
 import { createAdkStream } from "./AdkClient";
@@ -212,6 +213,47 @@ describe("ADK runtime callbacks", () => {
 });
 
 describe("ADK stream lifecycle", () => {
+  it("reports streamed tool calls with the run config that produced them", async () => {
+    const runConfig = { custom: { model: "model-a" } };
+    const onMessages = vi.fn();
+    const stream: AdkStreamCallback = async function* () {
+      yield {
+        id: "event-1",
+        author: "agent",
+        content: {
+          role: "model",
+          parts: [
+            { functionCall: { id: "tool-1", name: "lookup", args: {} } },
+            { functionCall: { id: "tool-2", name: "search", args: {} } },
+          ],
+        },
+      };
+    };
+    const { result } = renderHook(() =>
+      useAdkMessagesInternal({ stream, onMessages }),
+    );
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "look it up" }],
+        { runConfig },
+      );
+    });
+
+    expect(onMessages).toHaveBeenLastCalledWith(
+      expect.arrayContaining([
+        expect.objectContaining({
+          type: "ai",
+          tool_calls: [
+            expect.objectContaining({ id: "tool-1" }),
+            expect.objectContaining({ id: "tool-2" }),
+          ],
+        }),
+      ]),
+      runConfig,
+    );
+  });
+
   it("settles a superseded send while its stream is still opening", async () => {
     const signals: AbortSignal[] = [];
     const parked = new Promise<AsyncGenerator<AdkEvent>>(() => {});
@@ -800,6 +842,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,3 +1,4 @@
+import { isRecord } from "@assistant-ui/core/internal";
 import {
   useState,
   useCallback,
@@ -39,6 +40,10 @@ export type UseAdkMessagesOptions = {
   };
 };
 
+type UseAdkMessagesInternalOptions = UseAdkMessagesOptions & {
+  onMessages?: (messages: AdkMessage[], runConfig: unknown) => void;
+};
+
 type AdkRuntimeCallbackName = "onError" | "onCustomEvent" | "onAgentTransfer";
 
 const invokeAdkRuntimeCallback = <TArgs extends readonly unknown[]>(
@@ -49,10 +54,11 @@ const invokeAdkRuntimeCallback = <TArgs extends readonly unknown[]>(
   void invokeUserCallback("react-google-adk", name, callback, ...args);
 };
 
-export const useAdkMessages = ({
+const useAdkMessagesInternal = ({
   stream,
   eventHandlers,
-}: UseAdkMessagesOptions) => {
+  onMessages,
+}: UseAdkMessagesInternalOptions) => {
   const [messages, _setMessages] = useState<AdkMessage[]>([]);
   const [stateDelta, setStateDelta] = useState<Record<string, unknown>>({});
   const [agentInfo, setAgentInfo] = useState<{
@@ -167,8 +173,11 @@ 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;
@@ -202,6 +211,17 @@ export const useAdkMessages = ({
             break;
           }
           const updatedMessages = accumulator.processEvent(event);
+          // Each event part can append at most one message, and a function call
+          // stays on the current assistant message until a later part finalizes
+          // it, so every message touched by this event is within this tail.
+          const affectedMessageCount = Math.max(
+            event.content?.parts?.length ?? 0,
+            1,
+          );
+          const affectedMessages = updatedMessages.slice(-affectedMessageCount);
+          if (affectedMessages.length > 0) {
+            onMessages?.(affectedMessages, config.runConfig);
+          }
           setMessagesImmediate(updatedMessages);
           setStateDelta({
             ...stateDeltaRef.current,
@@ -265,6 +285,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;
         }
       }
@@ -277,6 +324,7 @@ export const useAdkMessages = ({
       onError,
       onCustomEvent,
       onAgentTransfer,
+      onMessages,
     ],
   );
 
@@ -306,6 +354,17 @@ export const useAdkMessages = ({
   };
 };
 
+export const useAdkMessages = ({
+  stream,
+  eventHandlers,
+}: UseAdkMessagesOptions) =>
+  useAdkMessagesInternal({
+    stream,
+    ...(eventHandlers !== undefined && { eventHandlers }),
+  });
+
+export { useAdkMessagesInternal };
+
 /**
  * Transport sends every human and tool message of one `send` call as a single
  * ADK `Content`, and ADK parses that event's function responses before running
@@ -395,9 +454,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;