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

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,4 +1,4 @@
-import { useState, useCallback, useRef, useMemo } from "react";
+import { useState, useCallback, useEffect, useRef, useMemo } from "react";
 import { v4 as uuidv4 } from "uuid";
 import { useAui } from "@assistant-ui/store";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
@@ -136,9 +136,16 @@ export const useAdkMessages = ({
         m.id ? m : { ...m, id: uuidv4() },
       ) 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 +245,8 @@ export const useAdkMessages = ({
     }
   }, []);
 
+  useEffect(() => cancel, [cancel]);
+
   return {
     messages,
     stateDelta,
@@ -256,6 +265,55 @@ 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 ?? uuidv4() };
+  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") {

package/src/useAdkRuntime.ts

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
   getExternalStoreMessages,
   pickExternalStoreSharedOptions,
@@ -33,7 +33,14 @@ import type {
   OnAdkAgentTransferCallback,
 } from "./types";
 import { useAdkMessages } from "./useAdkMessages";
-import { convertAdkMessage } from "./convertAdkMessages";
+import {
+  convertAdkMessage,
+  createAdkMessageConverter,
+} from "./convertAdkMessages";
+import {
+  projectAdkToolApprovals,
+  toAdkToolConfirmationReply,
+} from "./adkToolApproval";
 import { adkExtras } from "./adkExtras";
 import { v4 as uuidv4 } from "uuid";
 
@@ -285,8 +292,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,
   });
@@ -556,6 +576,12 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
         {},
       );
     },
+    onRespondToToolApproval: async (options) => {
+      await handleSendMessage(
+        [toAdkToolConfirmationReply(options, toolApprovalsRef.current)],
+        {},
+      );
+    },
     onCancel: unstable_allowCancellation
       ? async () => {
           cancel();