@assistant-ui/react-google-adk 0.0.28 → 0.0.30

package/src/useAdkMessages.test.ts

@@ -22,6 +22,7 @@ import {
 import { projectAdkToolApprovals } from "./adkToolApproval";
 import { createAdkStream } from "./AdkClient";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
+import { getPendingCancellations } from "./convertToAdkMessages";
 import type { AdkEvent, AdkMessage, AdkStreamCallback } from "./types";
 
 afterEach(() => {
@@ -29,6 +30,53 @@ afterEach(() => {
   vi.unstubAllGlobals();
 });
 
+describe("optimistic tool outcomes", () => {
+  it.each([false, true])(
+    "preserves failures alongside successful results (batch: %s)",
+    (batch) => {
+      const failed: AdkMessage = {
+        id: "failed",
+        type: "tool",
+        name: "search",
+        tool_call_id: "tc-error",
+        content: "denied",
+        status: "error",
+      };
+      const succeeded: AdkMessage = {
+        id: "succeeded",
+        type: "tool",
+        name: "search",
+        tool_call_id: "tc-ok",
+        content: "found",
+        status: "success",
+      };
+      const events = batch
+        ? messagesToEvents([
+            failed,
+            succeeded,
+            { id: "human", type: "human", content: "continue" },
+          ])
+        : [messageToEvent(failed), messageToEvent(succeeded)];
+      const acc = new AdkEventAccumulator();
+      for (const event of events) acc.processEvent(event);
+      expect(
+        acc.getMessages().filter((message) => message.type === "tool"),
+      ).toMatchObject([
+        {
+          tool_call_id: "tc-error",
+          status: "error",
+          content: JSON.stringify({ error: "denied" }),
+        },
+        {
+          tool_call_id: "tc-ok",
+          status: "success",
+          content: JSON.stringify({ result: "found" }),
+        },
+      ]);
+    },
+  );
+});
+
 describe("ADK runtime callbacks", () => {
   it.each(["onAgentTransfer", "onCustomEvent", "onError"] as const)(
     "continues streaming when %s throws",
@@ -131,6 +179,105 @@ describe("ADK runtime callbacks", () => {
 });
 
 describe("ADK stream lifecycle", () => {
+  it("settles a superseded send while its stream is still opening", async () => {
+    const signals: AbortSignal[] = [];
+    const parked = new Promise<AsyncGenerator<AdkEvent>>(() => {});
+    let calls = 0;
+    const stream = vi.fn(function (_messages, { abortSignal }) {
+      signals.push(abortSignal);
+      if (calls++ === 0) return parked;
+      return (async function* () {
+        yield {
+          id: "event-1",
+          invocationId: "run-1",
+          author: "agent",
+          content: { role: "model", parts: [{ text: "done-1" }] },
+        };
+      })();
+    }) satisfies AdkStreamCallback;
+    const { result } = renderHook(() => useAdkMessages({ stream }));
+
+    let firstSend!: Promise<void>;
+    act(() => {
+      firstSend = result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "first" }],
+        {},
+      );
+    });
+    await vi.waitFor(() => expect(stream).toHaveBeenCalledOnce());
+
+    let secondSend!: Promise<void>;
+    act(() => {
+      secondSend = result.current.sendMessage(
+        [{ id: "user-2", type: "human", content: "second" }],
+        {},
+      );
+    });
+
+    await act(async () => {
+      await Promise.all([firstSend, secondSend]);
+    });
+    expect(signals[0]?.aborted).toBe(true);
+    expect(signals[1]?.aborted).toBe(false);
+    expect(result.current.messages.at(-1)).toMatchObject({
+      type: "ai",
+      content: [{ type: "text", text: "done-1" }],
+    });
+  });
+
+  it("aborts and settles a superseded stream that stops yielding", async () => {
+    const signals: AbortSignal[] = [];
+    const parked = new Promise<void>(() => {});
+    let calls = 0;
+    const stream = vi.fn(function (
+      _messages,
+      { abortSignal },
+    ): AsyncGenerator<AdkEvent> {
+      signals.push(abortSignal);
+      if (calls++ === 0) {
+        return (async function* () {
+          await parked;
+        })();
+      }
+      return (async function* () {
+        yield {
+          id: "event-1",
+          invocationId: "run-1",
+          author: "agent",
+          content: { role: "model", parts: [{ text: "done-1" }] },
+        };
+      })();
+    }) satisfies AdkStreamCallback;
+    const { result } = renderHook(() => useAdkMessages({ stream }));
+
+    let firstSend!: Promise<void>;
+    act(() => {
+      firstSend = result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "first" }],
+        {},
+      );
+    });
+    await vi.waitFor(() => expect(stream).toHaveBeenCalledOnce());
+
+    let secondSend!: Promise<void>;
+    act(() => {
+      secondSend = result.current.sendMessage(
+        [{ id: "user-2", type: "human", content: "second" }],
+        {},
+      );
+    });
+
+    await act(async () => {
+      await Promise.all([firstSend, secondSend]);
+    });
+    expect(signals[0]?.aborted).toBe(true);
+    expect(signals[1]?.aborted).toBe(false);
+    expect(result.current.messages.at(-1)).toMatchObject({
+      type: "ai",
+      content: [{ type: "text", text: "done-1" }],
+    });
+  });
+
   it("aborts the active stream when the hook unmounts", async () => {
     let runSignal: AbortSignal | undefined;
     let resolveStarted!: () => void;
@@ -216,6 +363,86 @@ describe("optimistic confirmation replies", () => {
     return result;
   };
 
+  it("preserves an unanswered gate across a reply run", async () => {
+    let run = 0;
+    const stream: AdkStreamCallback = async function* () {
+      run += 1;
+      if (run === 1) {
+        yield {
+          id: "gates",
+          author: "agent",
+          longRunningToolIds: ["conf-a", "conf-b"],
+          content: {
+            role: "model",
+            parts: [
+              {
+                functionCall: {
+                  id: "conf-a",
+                  name: "adk_request_confirmation",
+                  args: {},
+                },
+              },
+              {
+                functionCall: {
+                  id: "conf-b",
+                  name: "adk_request_confirmation",
+                  args: {},
+                },
+              },
+            ],
+          },
+        } satisfies AdkEvent;
+      } else {
+        yield {
+          id: "rerun",
+          author: "agent",
+          content: {
+            role: "user",
+            parts: [
+              {
+                functionResponse: {
+                  id: "orig-conf-a",
+                  name: "delete_file",
+                  response: { result: "deleted" },
+                },
+              },
+            ],
+          },
+        } satisfies AdkEvent;
+      }
+    };
+    const { result } = renderHook(() => useAdkMessages({ stream }));
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "start" }],
+        {},
+      );
+    });
+    expect(result.current.longRunningToolIds).toEqual(["conf-a", "conf-b"]);
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [
+          confirmationReply(
+            "reply-a",
+            "conf-a",
+            JSON.stringify({ confirmed: true }),
+          ),
+        ],
+        {},
+      );
+    });
+
+    expect(result.current.longRunningToolIds).toEqual(["conf-b"]);
+    expect(
+      getPendingCancellations(
+        result.current.messages,
+        result.current.longRunningToolIds,
+      ),
+    ).toEqual([]);
+  });
+
   it("keeps both gates pending when one send carries an unreadable reply", async () => {
     const result = await renderWithGates();
 
@@ -439,6 +666,26 @@ describe("optimistic multi-message sends", () => {
 });
 
 describe("messageToEvent (contentToParts)", () => {
+  it.each([
+    ["scalar", "false", { result: false }],
+    ["array", "[1,2]", { results: [1, 2] }],
+  ])(
+    "normalizes an optimistic %s tool response",
+    (_label, content, response) => {
+      const event = messageToEvent({
+        id: "tool-1",
+        type: "tool",
+        content,
+        tool_call_id: "call-1",
+        name: "search",
+      });
+
+      expect(event.content?.parts[0]?.functionResponse?.response).toEqual(
+        response,
+      );
+    },
+  );
+
   it("serializes a file content part as inlineData", () => {
     const msg: AdkMessage = {
       id: "m1",

package/src/useAdkMessages.ts

@@ -8,9 +8,14 @@ import {
 } from "react";
 import { generateId } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
-import { invokeUserCallback } from "@assistant-ui/core/internal";
+import {
+  abortableIterable,
+  invokeUserCallback,
+  openAbortableIterable,
+} from "@assistant-ui/core/internal";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
 import { contentToParts } from "./contentToParts";
+import { toAdkFunctionResponse } from "./toAdkFunctionResponse";
 import type {
   AdkEvent,
   AdkMessage,
@@ -54,7 +59,7 @@ export const useAdkMessages = ({
     name?: string | undefined;
     branch?: string | undefined;
   }>({});
-  const [longRunningToolIds, setLongRunningToolIds] = useState<string[]>([]);
+  const [longRunningToolIds, _setLongRunningToolIds] = useState<string[]>([]);
   const [artifactDelta, setArtifactDelta] = useState<Record<string, number>>(
     {},
   );
@@ -67,9 +72,9 @@ export const useAdkMessages = ({
     Map<string, AdkMessageMetadata>
   >(new Map());
   const lastTransferToAgentRef = useRef<string | undefined>(undefined);
-  // setMessagesImmediate is the only writer of the messages state and publishes
-  // this ref with it, so the ref never trails a commit.
+  // 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);
   const stateDeltaRef = useRef(stateDelta);
   useInsertionEffect(() => {
     stateDeltaRef.current = stateDelta;
@@ -87,6 +92,10 @@ export const useAdkMessages = ({
     messagesRef.current = msgs;
     _setMessages(msgs);
   }, []);
+  const setLongRunningToolIds = useCallback((ids: string[]) => {
+    longRunningToolIdsRef.current = ids;
+    _setLongRunningToolIds(ids);
+  }, []);
 
   /**
    * Swap the thread over to a loaded snapshot in one commit. Unlike
@@ -106,7 +115,7 @@ export const useAdkMessages = ({
       setArtifactDelta(snapshot.artifactDelta ?? {});
       setAgentInfo(snapshot.agentInfo ?? {});
     },
-    [setMessagesImmediate],
+    [setLongRunningToolIds, setMessagesImmediate],
   );
 
   // Replace the message list AND reset derived per-turn HITL state.
@@ -122,7 +131,7 @@ export const useAdkMessages = ({
       setEscalated(false);
       setMessageMetadata(new Map());
     },
-    [setMessagesImmediate],
+    [setLongRunningToolIds, setMessagesImmediate],
   );
 
   const abortControllerRef = useRef<AbortController | null>(null);
@@ -144,27 +153,52 @@ export const useAdkMessages = ({
       // with the originals would leave every later staged id beside the merged
       // copy of itself.
       const resentIds = new Set(newMessagesWithId.map((m) => m.id));
+      // The optimistic event for a tool-only batch carries no author, so the accumulator cannot settle the calls this send answers.
+      const answeredToolCallIds = new Set(
+        newMessagesWithId.flatMap((m) =>
+          m.type === "tool" ? [m.tool_call_id] : [],
+        ),
+      );
       const accumulator = new AdkEventAccumulator(
         messagesRef.current.filter((m) => !resentIds.has(m.id)),
+        longRunningToolIdsRef.current.filter(
+          (id) => !answeredToolCallIds.has(id),
+        ),
       );
       for (const event of messagesToEvents(newMessagesWithId)) {
         accumulator.processEvent(event);
       }
       setMessagesImmediate(accumulator.getMessages());
+      setLongRunningToolIds(accumulator.getLongRunningToolIds());
 
+      // Google ADK replaces active runs, while React LangGraph queues sends.
+      abortControllerRef.current?.abort();
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
 
       try {
-        const response = await stream(newMessagesWithId, {
-          ...config,
-          abortSignal: abortController.signal,
-          initialize: async () => {
-            return await aui.threadListItem.initialize();
-          },
-        });
+        const response = await openAbortableIterable(
+          stream(newMessagesWithId, {
+            ...config,
+            abortSignal: abortController.signal,
+            initialize: async () => {
+              return await aui.threadListItem.initialize();
+            },
+          }),
+          abortController.signal,
+        );
+        if (!response) return;
 
-        for await (const event of response) {
+        for await (const event of abortableIterable(
+          response,
+          abortController.signal,
+        )) {
+          if (
+            abortController.signal.aborted ||
+            abortControllerRef.current !== abortController
+          ) {
+            break;
+          }
           const updatedMessages = accumulator.processEvent(event);
           setMessagesImmediate(updatedMessages);
           setStateDelta({
@@ -222,6 +256,7 @@ export const useAdkMessages = ({
       } catch (error) {
         if (
           !abortController.signal.aborted &&
+          abortControllerRef.current === abortController &&
           !(error instanceof Error && error.name === "AbortError")
         ) {
           throw error;
@@ -235,6 +270,7 @@ export const useAdkMessages = ({
     [
       aui,
       setMessagesImmediate,
+      setLongRunningToolIds,
       stream,
       onError,
       onCustomEvent,
@@ -343,7 +379,7 @@ export const messageToEvent = (msg: AdkMessage): AdkEvent => {
             functionResponse: {
               name: msg.name,
               id: msg.tool_call_id,
-              response,
+              response: toAdkFunctionResponse(response, msg.status === "error"),
             },
           },
         ],

package/src/useAdkRuntime.replacement.test.tsx

@@ -0,0 +1,138 @@
+// @vitest-environment jsdom
+
+import { act, render, waitFor } from "@testing-library/react";
+import { type FC } from "react";
+import { describe, expect, it, vi } from "vitest";
+import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
+import type {
+  AssistantRuntime,
+  RemoteThreadListAdapter,
+} from "@assistant-ui/core";
+import { useAdkRuntime } from "./useAdkRuntime";
+import type { AdkEvent } from "./types";
+
+const makeThreadListAdapter = (): RemoteThreadListAdapter => ({
+  list: vi.fn(async () => ({
+    threads: [
+      {
+        status: "regular" as const,
+        remoteId: "adk-1",
+        externalId: "adk-1",
+        title: "ADK session",
+      },
+    ],
+  })),
+  initialize: vi.fn(async () => ({
+    remoteId: "adk-1",
+    externalId: "adk-1",
+  })),
+  rename: vi.fn(async () => {}),
+  archive: vi.fn(async () => {}),
+  unarchive: vi.fn(async () => {}),
+  delete: vi.fn(async () => {}),
+  generateTitle: vi.fn(async () => new ReadableStream() as never),
+  fetch: vi.fn(async () => ({
+    status: "regular" as const,
+    remoteId: "adk-1",
+    externalId: "adk-1",
+    title: "ADK session",
+  })),
+});
+
+const deferred = () => {
+  let resolve!: () => void;
+  const promise = new Promise<void>((r) => {
+    resolve = r;
+  });
+  return { promise, resolve };
+};
+
+describe("useAdkRuntime replacement runs", () => {
+  it.each([
+    { label: "events after cancellation", cancelFirst: true, failFirst: false },
+    {
+      label: "events without cancellation",
+      cancelFirst: false,
+      failFirst: false,
+    },
+    {
+      label: "errors without cancellation",
+      cancelFirst: false,
+      failFirst: true,
+    },
+  ])("ignores superseded run $label", async ({ cancelFirst, failFirst }) => {
+    const gates = [deferred(), deferred()];
+    let calls = 0;
+    const stream = vi.fn(async function* (): AsyncGenerator<AdkEvent> {
+      const call = calls++;
+      await gates[call]!.promise;
+      if (call === 0 && failFirst) throw new Error("stale run failed");
+      yield {
+        id: `event-${call}`,
+        invocationId: `run-${call}`,
+        author: "agent",
+        content: { role: "model", parts: [{ text: `done-${call}` }] },
+      };
+    });
+    const sessionAdapter = makeThreadListAdapter();
+    const capture: { runtime: AssistantRuntime | null } = { runtime: null };
+
+    const Inner: FC = () => {
+      const runtime = useAdkRuntime({
+        stream,
+        sessionAdapter,
+        unstable_allowCancellation: true,
+      });
+      capture.runtime = runtime;
+      return <AssistantRuntimeProvider runtime={runtime} />;
+    };
+
+    await act(async () => {
+      render(<Inner />);
+    });
+    await waitFor(() => expect(capture.runtime).not.toBeNull());
+    await act(async () => {
+      await capture.runtime!.threads.switchToThread("adk-1");
+    });
+
+    let firstSend!: Promise<void>;
+    act(() => {
+      firstSend = capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => expect(stream).toHaveBeenCalledTimes(1));
+
+    let secondSend!: Promise<void>;
+    await act(async () => {
+      if (cancelFirst) await capture.runtime!.thread.cancelRun();
+      secondSend = capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "second" }],
+      });
+    });
+    await waitFor(() => expect(stream).toHaveBeenCalledTimes(2));
+
+    await act(async () => {
+      gates[0]!.resolve();
+      await firstSend;
+    });
+
+    const messagesAfterFirstSettles = JSON.stringify(
+      capture.runtime!.thread.getState().messages,
+    );
+    expect(messagesAfterFirstSettles).toContain("second");
+    expect(messagesAfterFirstSettles).not.toContain("done-0");
+    expect(capture.runtime!.thread.getState().isRunning).toBe(true);
+
+    await act(async () => {
+      gates[1]!.resolve();
+      await secondSend;
+    });
+    expect(
+      JSON.stringify(capture.runtime!.thread.getState().messages),
+    ).toContain("done-1");
+    expect(capture.runtime!.thread.getState().isRunning).toBe(false);
+  });
+});

package/src/useAdkRuntime.ts

@@ -57,6 +57,7 @@ import {
   toAdkToolConfirmationReply,
 } from "./adkToolApproval";
 import { adkExtras } from "./adkExtras";
+import { ADK_SDK } from "./sdkIdentity";
 
 export type UseAdkRuntimeOptions = ExternalStoreSharedOptions & {
   stream: AdkStreamCallback;
@@ -168,16 +169,18 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   useInsertionEffect(() => {
     isRunningRef.current = effectiveIsRunning;
   }, [effectiveIsRunning]);
+  const runGenerationRef = useRef(0);
 
   const handleSendMessage = async (
     msgs: AdkMessage[],
     config: AdkSendMessageConfig,
   ) => {
+    const generation = ++runGenerationRef.current;
     try {
       setIsRunning(true);
       await sendMessage(msgs, config);
     } finally {
-      setIsRunning(false);
+      if (runGenerationRef.current === generation) setIsRunning(false);
     }
   };
 
@@ -502,6 +505,7 @@ export const useAdkRuntime = ({
 }: UseAdkRuntimeOptions) => {
   const aui = useAui();
   const cloudAdapter = useCloudThreadListAdapter({
+    sdk: ADK_SDK,
     cloud,
     create: createCloudThreadListAdapterCreateFallback(
       create,

package/src/useAdkRuntimeApproval.test.tsx

@@ -97,7 +97,6 @@ const approvalPart = () =>
 afterEach(() => {
   mocks.adapters.length = 0;
   mocks.messages = [];
-  vi.clearAllMocks();
 });
 
 describe("useAdkRuntime tool approvals", () => {