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

package/src/useAdkMessages.test.ts

@@ -15,13 +15,18 @@ vi.mock("@assistant-ui/store", async (importOriginal) => ({
 
 import {
   messageToEvent,
+  messagesToEvents,
   useAdkMessages,
   type UseAdkMessagesOptions,
 } from "./useAdkMessages";
+import { projectAdkToolApprovals } from "./adkToolApproval";
+import { createAdkStream } from "./AdkClient";
+import { AdkEventAccumulator } from "./AdkEventAccumulator";
 import type { AdkEvent, AdkMessage, AdkStreamCallback } from "./types";
 
 afterEach(() => {
   vi.restoreAllMocks();
+  vi.unstubAllGlobals();
 });
 
 describe("ADK runtime callbacks", () => {
@@ -125,6 +130,314 @@ describe("ADK runtime callbacks", () => {
   });
 });
 
+describe("ADK stream lifecycle", () => {
+  it("aborts the active stream when the hook unmounts", async () => {
+    let runSignal: AbortSignal | undefined;
+    let resolveStarted!: () => void;
+    const started = new Promise<void>((resolve) => {
+      resolveStarted = resolve;
+    });
+    const stream: AdkStreamCallback = async function* (
+      _messages,
+      { abortSignal },
+    ) {
+      runSignal = abortSignal;
+      resolveStarted();
+      await new Promise<void>((resolve) => {
+        abortSignal.addEventListener("abort", () => resolve(), { once: true });
+      });
+    };
+    const { result, unmount } = renderHook(() => useAdkMessages({ stream }));
+
+    let sendPromise!: Promise<void>;
+    act(() => {
+      sendPromise = result.current.sendMessage(
+        [{ id: "user", type: "human", content: "hello" }],
+        {},
+      );
+    });
+    await started;
+
+    unmount();
+
+    expect(runSignal?.aborted).toBe(true);
+    await expect(sendPromise).resolves.toBeUndefined();
+  });
+});
+
+describe("optimistic confirmation replies", () => {
+  const confirmationCall = (id: string): AdkMessage => ({
+    id: `ai-${id}`,
+    type: "ai",
+    content: [],
+    tool_calls: [
+      {
+        id,
+        name: "adk_request_confirmation",
+        args: {
+          originalFunctionCall: { id: `orig-${id}`, name: "delete_file" },
+          toolConfirmation: { hint: "Delete?" },
+        },
+      },
+    ],
+  });
+
+  /**
+   * ADK parses the text inside its client wrapper without a `try`, so a wrapper
+   * holding text that is not JSON raises and abandons the whole event.
+   */
+  const UNREADABLE_REPLY = JSON.stringify({ response: "not json" });
+
+  const confirmationReply = (
+    id: string,
+    toolCallId: string,
+    content: string,
+  ): AdkMessage => ({
+    id,
+    type: "tool",
+    tool_call_id: toolCallId,
+    name: "adk_request_confirmation",
+    content,
+    status: "success",
+  });
+
+  const emptyStream: AdkStreamCallback = async function* () {};
+
+  const renderWithGates = async () => {
+    const { result } = renderHook(() =>
+      useAdkMessages({ stream: emptyStream }),
+    );
+    await act(async () => {
+      result.current.setMessages([
+        confirmationCall("conf-a"),
+        confirmationCall("conf-b"),
+      ]);
+    });
+    return result;
+  };
+
+  it("keeps both gates pending when one send carries an unreadable reply", async () => {
+    const result = await renderWithGates();
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [
+          confirmationReply(
+            "reply-a",
+            "conf-a",
+            JSON.stringify({ confirmed: true }),
+          ),
+          confirmationReply("reply-b", "conf-b", UNREADABLE_REPLY),
+        ],
+        {},
+      );
+    });
+
+    expect([
+      // The confirmation and the call it gates share one approval object.
+      ...new Set(
+        projectAdkToolApprovals(result.current.messages).approvals.values(),
+      ),
+    ]).toEqual([{ id: "conf-a" }, { id: "conf-b" }]);
+  });
+
+  it("keeps both gates pending when an ai message sits between the replies", async () => {
+    const result = await renderWithGates();
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [
+          confirmationReply(
+            "reply-a",
+            "conf-a",
+            JSON.stringify({ confirmed: true }),
+          ),
+          { id: "ai-interleaved", type: "ai", content: "thinking" },
+          confirmationReply("reply-b", "conf-b", UNREADABLE_REPLY),
+        ],
+        {},
+      );
+    });
+
+    expect([
+      // The confirmation and the call it gates share one approval object.
+      ...new Set(
+        projectAdkToolApprovals(result.current.messages).approvals.values(),
+      ),
+    ]).toEqual([{ id: "conf-a" }, { id: "conf-b" }]);
+  });
+
+  it("settles the readable reply when the two replies were sent separately", async () => {
+    const result = await renderWithGates();
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [
+          confirmationReply(
+            "reply-a",
+            "conf-a",
+            JSON.stringify({ confirmed: true }),
+          ),
+        ],
+        {},
+      );
+    });
+    await act(async () => {
+      await result.current.sendMessage(
+        [confirmationReply("reply-b", "conf-b", UNREADABLE_REPLY)],
+        {},
+      );
+    });
+
+    expect([
+      // The confirmation and the call it gates share one approval object.
+      ...new Set(
+        projectAdkToolApprovals(result.current.messages).approvals.values(),
+      ),
+    ]).toEqual([{ id: "conf-a", approved: true }, { id: "conf-b" }]);
+  });
+});
+
+describe("messagesToEvents", () => {
+  const toolReply = (id: string, toolCallId: string): AdkMessage => ({
+    id,
+    type: "tool",
+    tool_call_id: toolCallId,
+    name: "adk_request_confirmation",
+    content: JSON.stringify({ confirmed: true }),
+    status: "success",
+  });
+
+  it("emits nothing for an empty batch", () => {
+    // `onReload` sends no messages. The transport still puts an empty user
+    // content on the wire, but projecting one here would append an empty user
+    // bubble above every regenerated turn.
+    expect(messagesToEvents([])).toEqual([]);
+  });
+
+  it("merges the human/tool run across an interleaved ai message", () => {
+    const events = messagesToEvents([
+      toolReply("reply-a", "conf-a"),
+      { id: "ai-interleaved", type: "ai", content: "thinking" },
+      toolReply("reply-b", "conf-b"),
+    ]);
+
+    expect(events).toHaveLength(2);
+    expect(events[0]?.id).toBe("reply-a");
+    expect(
+      events[0]?.content?.parts?.map((p) => p.functionResponse?.id),
+    ).toEqual(["conf-a", "conf-b"]);
+    expect(events[1]?.id).toBe("ai-interleaved");
+  });
+
+  it("mirrors the transport's empty user content for an ai-only batch", async () => {
+    const aiOnly: AdkMessage[] = [
+      { id: "ai-1", type: "ai", content: "one" },
+      { id: "ai-2", type: "ai", content: "two" },
+    ];
+
+    let sentBody = "";
+    vi.stubGlobal("fetch", async (_url: string, init: RequestInit) => {
+      sentBody = init.body as string;
+      return new Response(
+        new ReadableStream<Uint8Array>({
+          start: (controller) => controller.close(),
+        }),
+        { status: 200, headers: { "Content-Type": "text/event-stream" } },
+      );
+    });
+    const stream = createAdkStream({
+      api: "http://localhost:8000",
+      appName: "app",
+      userId: "user-1",
+    });
+    for await (const _ of await stream(aiOnly, {
+      abortSignal: new AbortController().signal,
+      initialize: async () => ({ remoteId: "r1", externalId: "s1" }),
+    })) {
+      /* drain */
+    }
+    const sent = JSON.parse(sentBody);
+
+    const events = messagesToEvents(aiOnly);
+    const userEvent = events.find((e) => e.author === "user");
+    expect(sent.newMessage.parts).toEqual([{ text: "" }]);
+    expect(userEvent?.content?.parts).toEqual(sent.newMessage.parts);
+    expect(events.map((e) => e.id)).toEqual(["ai-1", "ai-2", userEvent?.id]);
+
+    const accumulator = new AdkEventAccumulator([]);
+    let messages: AdkMessage[] = [];
+    for (const event of events) messages = accumulator.processEvent(event);
+    expect(messages.filter((m) => m.type === "human")).toEqual([
+      { id: userEvent?.id, type: "human", content: "" },
+    ]);
+  });
+
+  it("emits the merged event at the position of the run it replaces", () => {
+    const events = messagesToEvents([
+      { id: "ai-first", type: "ai", content: "thinking" },
+      { id: "human-a", type: "human", content: "a" },
+    ]);
+
+    expect(events.map((e) => e.id)).toEqual(["ai-first", "human-a"]);
+  });
+});
+
+describe("optimistic multi-message sends", () => {
+  const emptyStream: AdkStreamCallback = async function* () {};
+
+  it("does not duplicate later messages of a staged multi-human send", async () => {
+    const staged: AdkMessage[] = [
+      { id: "human-a", type: "human", content: "a" },
+      { id: "human-b", type: "human", content: "b" },
+    ];
+    const { result } = renderHook(() =>
+      useAdkMessages({ stream: emptyStream }),
+    );
+    await act(async () => {
+      result.current.setMessages(staged);
+    });
+
+    await act(async () => {
+      await result.current.sendMessage(staged, {});
+    });
+
+    expect(result.current.messages).toEqual([
+      {
+        id: "human-a",
+        type: "human",
+        content: [
+          { type: "text", text: "a" },
+          { type: "text", text: "b" },
+        ],
+      },
+    ]);
+  });
+
+  it("keeps messages that are not part of the send", async () => {
+    const { result } = renderHook(() =>
+      useAdkMessages({ stream: emptyStream }),
+    );
+    await act(async () => {
+      result.current.setMessages([
+        { id: "earlier", type: "human", content: "earlier" },
+      ]);
+    });
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [{ id: "human-a", type: "human", content: "a" }],
+        {},
+      );
+    });
+
+    expect(result.current.messages.map((m) => m.id)).toEqual([
+      "earlier",
+      "human-a",
+    ]);
+  });
+});
+
 describe("messageToEvent (contentToParts)", () => {
   it("serializes a file content part as inlineData", () => {
     const msg: AdkMessage = {

package/src/useAdkMessages.ts

@@ -1,6 +1,7 @@
-import { useState, useCallback, useRef, useMemo } from "react";
-import { v4 as uuidv4 } from "uuid";
+import { useState, useCallback, useEffect, useRef, useMemo } from "react";
+import { generateId } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
+import { invokeUserCallback } from "@assistant-ui/core/internal";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
 import { contentToParts } from "./contentToParts";
 import type {
@@ -28,24 +29,12 @@ export type UseAdkMessagesOptions = {
 
 type AdkRuntimeCallbackName = "onError" | "onCustomEvent" | "onAgentTransfer";
 
-const reportCallbackError = (name: AdkRuntimeCallbackName, error: unknown) => {
-  console.error(`[react-google-adk] ${name} callback threw an error`, error);
-};
-
-const invokeAdkRuntimeCallback = <TArgs extends unknown[]>(
+const invokeAdkRuntimeCallback = <TArgs extends readonly unknown[]>(
   name: AdkRuntimeCallbackName,
-  callback: ((...args: TArgs) => void | Promise<void>) | undefined,
+  callback: ((...args: TArgs) => unknown) | undefined,
   ...args: TArgs
-) => {
-  if (!callback) return;
-
-  try {
-    void Promise.resolve(callback(...args)).catch((error) => {
-      reportCallbackError(name, error);
-    });
-  } catch (error) {
-    reportCallbackError(name, error);
-  }
+): void => {
+  void invokeUserCallback("react-google-adk", name, callback, ...args);
 };
 
 export const useAdkMessages = ({
@@ -133,12 +122,19 @@ export const useAdkMessages = ({
   const sendMessage = useCallback(
     async (newMessages: AdkMessage[], config: AdkSendMessageConfig) => {
       const newMessagesWithId = newMessages.map((m) =>
-        m.id ? m : { ...m, id: uuidv4() },
+        m.id ? m : { ...m, id: generateId() },
       ) as AdkMessage[];
 
-      const accumulator = new AdkEventAccumulator(messagesRef.current);
-      for (const msg of newMessagesWithId) {
-        accumulator.processEvent(messageToEvent(msg));
+      // A staged message is already in the thread under its own id, and the
+      // merged event below re-emits the whole batch under the first one. Seeding
+      // with the originals would leave every later staged id beside the merged
+      // copy of itself.
+      const resentIds = new Set(newMessagesWithId.map((m) => m.id));
+      const accumulator = new AdkEventAccumulator(
+        messagesRef.current.filter((m) => !resentIds.has(m.id)),
+      );
+      for (const event of messagesToEvents(newMessagesWithId)) {
+        accumulator.processEvent(event);
       }
       setMessagesImmediate(accumulator.getMessages());
 
@@ -238,6 +234,8 @@ export const useAdkMessages = ({
     }
   }, []);
 
+  useEffect(() => cancel, [cancel]);
+
   return {
     messages,
     stateDelta,
@@ -256,11 +254,60 @@ export const useAdkMessages = ({
   };
 };
 
+/**
+ * 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
+ * any tool, so the batch runs whole or not at all. The optimistic projection
+ * has to sit on the same boundary, so a run of those messages becomes one
+ * synthetic event whose parts come from the same per-message conversion.
+ *
+ * The transport drops `ai` messages from that `Content`, so one interleaved
+ * between two replies does not split the batch on the wire and must not split
+ * it here either. It still becomes its own event, placed after the merged one,
+ * so the optimistic projection keeps the assistant turn.
+ *
+ * @internal — exported for unit tests.
+ */
+export const messagesToEvents = (messages: AdkMessage[]): AdkEvent[] => {
+  // A reload sends no messages at all, and the empty user content the transport
+  // puts on the wire for it is not part of the optimistic view: projecting one
+  // would put an empty user bubble above every regenerated turn.
+  if (messages.length === 0) return [];
+
+  const events: AdkEvent[] = [];
+  const run: AdkMessage[] = [];
+  let runIndex = 0;
+
+  for (const msg of messages) {
+    if (msg.type === "ai") {
+      events.push(messageToEvent(msg));
+    } else {
+      if (run.length === 0) runIndex = events.length;
+      run.push(msg);
+    }
+  }
+
+  const parts = run.flatMap((m) => messageToEvent(m).content?.parts ?? []);
+  const human = run.find((m) => m.type === "human");
+
+  // A batch that contributes no part still reaches the wire: the transport
+  // sends an empty user `Content`, which a reload replays as an empty human
+  // message. Emitting it here keeps the optimistic view equal to that replay.
+  if (parts.length === 0) parts.push({ text: "" });
+
+  const event: AdkEvent = { id: (human ?? run[0])?.id ?? generateId() };
+  if (human || run.length === 0) event.author = "user";
+  event.content = { role: "user", parts };
+  events.splice(run.length > 0 ? runIndex : events.length, 0, event);
+
+  return events;
+};
+
 /** @internal — exported for unit tests. */
 export const messageToEvent = (msg: AdkMessage): AdkEvent => {
   if (msg.type === "human") {
     return {
-      id: msg.id ?? uuidv4(),
+      id: msg.id ?? generateId(),
       author: "user",
       content: { role: "user", parts: contentToParts(msg.content) },
     };
@@ -274,7 +321,7 @@ export const messageToEvent = (msg: AdkMessage): AdkEvent => {
       response = msg.content;
     }
     return {
-      id: msg.id ?? uuidv4(),
+      id: msg.id ?? generateId(),
       content: {
         role: "user",
         parts: [
@@ -290,7 +337,7 @@ export const messageToEvent = (msg: AdkMessage): AdkEvent => {
     };
   }
 
-  const result: AdkEvent = { id: msg.id ?? uuidv4() };
+  const result: AdkEvent = { id: msg.id ?? generateId() };
   if (msg.author != null) result.author = msg.author;
   result.content = {
     role: "model",