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

package/src/useAdkMessages.ts

@@ -1,7 +1,18 @@
-import { useState, useCallback, useEffect, useRef, useMemo } from "react";
+import {
+  useState,
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useRef,
+  useMemo,
+} 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 type {
@@ -60,14 +71,21 @@ 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.
   const messagesRef = useRef(messages);
-  messagesRef.current = messages;
   const stateDeltaRef = useRef(stateDelta);
-  stateDeltaRef.current = stateDelta;
+  useInsertionEffect(() => {
+    stateDeltaRef.current = stateDelta;
+  }, [stateDelta]);
   const artifactDeltaRef = useRef(artifactDelta);
-  artifactDeltaRef.current = artifactDelta;
+  useInsertionEffect(() => {
+    artifactDeltaRef.current = artifactDelta;
+  }, [artifactDelta]);
   const messageMetadataRef = useRef(messageMetadata);
-  messageMetadataRef.current = messageMetadata;
+  useInsertionEffect(() => {
+    messageMetadataRef.current = messageMetadata;
+  }, [messageMetadata]);
 
   const setMessagesImmediate = useCallback((msgs: AdkMessage[]) => {
     messagesRef.current = msgs;
@@ -138,19 +156,34 @@ export const useAdkMessages = ({
       }
       setMessagesImmediate(accumulator.getMessages());
 
+      // 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({
@@ -208,6 +241,7 @@ export const useAdkMessages = ({
       } catch (error) {
         if (
           !abortController.signal.aborted &&
+          abortControllerRef.current === abortController &&
           !(error instanceof Error && error.name === "AbortError")
         ) {
           throw 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.toolOwnership.test.tsx

@@ -0,0 +1,236 @@
+// @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,
+  useAssistantTool,
+} from "@assistant-ui/core/react";
+import type {
+  AssistantRuntime,
+  RemoteThreadListAdapter,
+  ToolCallMessagePart,
+} 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: "Existing 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: "Existing ADK session",
+  })),
+});
+
+const toolCalls = (runtime: AssistantRuntime): ToolCallMessagePart[] =>
+  runtime.thread
+    .getState()
+    .messages.flatMap((message) => message.content as readonly unknown[])
+    .filter(
+      (part): part is ToolCallMessagePart =>
+        (part as ToolCallMessagePart).type === "tool-call",
+    );
+
+const GATED_CALL = {
+  id: "adk-original-1",
+  name: "delete_file",
+  args: { path: "/tmp/a" },
+};
+
+const callEvent = (longRunningToolIds?: string[]): AdkEvent => ({
+  id: "ev-1",
+  invocationId: "inv-1",
+  author: "agent",
+  content: { role: "model", parts: [{ functionCall: GATED_CALL }] },
+  ...(longRunningToolIds && { longRunningToolIds }),
+});
+
+const confirmationEvent = (): AdkEvent => ({
+  id: "ev-2",
+  invocationId: "inv-1",
+  author: "agent",
+  content: {
+    role: "model",
+    parts: [
+      {
+        functionCall: {
+          id: "adk-confirmation-1",
+          name: "adk_request_confirmation",
+          args: {
+            originalFunctionCall: GATED_CALL,
+            toolConfirmation: { hint: "Delete /tmp/a?" },
+          },
+        },
+      },
+    ],
+  },
+  longRunningToolIds: ["adk-confirmation-1"],
+});
+
+const serverResultEvent = (): AdkEvent => ({
+  id: "ev-2",
+  invocationId: "inv-1",
+  author: "agent",
+  content: {
+    role: "user",
+    parts: [
+      {
+        functionResponse: {
+          id: "adk-original-1",
+          name: "delete_file",
+          response: { server: "deleted" },
+        },
+      },
+    ],
+  },
+});
+
+/**
+ * Streams `first`, hands control back so the window between the two events is
+ * observable, then streams `second` once the returned `release` is called.
+ */
+const renderStreamingAdk = async (first: AdkEvent, second: AdkEvent) => {
+  let release!: () => void;
+  const held = new Promise<void>((r) => (release = r));
+  let firstDelivered!: () => void;
+  const firstDeliveredP = new Promise<void>((r) => (firstDelivered = r));
+
+  const stream = vi.fn(async function* (): AsyncGenerator<AdkEvent> {
+    yield first;
+    firstDelivered();
+    await held;
+    yield second;
+  });
+
+  const execute = vi.fn(async () => ({ deleted: true }));
+  const capture: { runtime: AssistantRuntime | null } = { runtime: null };
+
+  const DeleteFileTool = () => {
+    useAssistantTool({
+      toolName: "delete_file",
+      description: "delete a file",
+      parameters: {
+        type: "object" as const,
+        properties: { path: { type: "string" } },
+      },
+      execute,
+    });
+    return null;
+  };
+
+  const Inner: FC = () => {
+    const runtime = useAdkRuntime({
+      stream: stream as never,
+      sessionAdapter: makeThreadListAdapter(),
+    });
+    capture.runtime = runtime;
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <DeleteFileTool />
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  await act(async () => {
+    render(<Inner />);
+  });
+  await waitFor(() => expect(capture.runtime).not.toBeNull());
+  await act(async () => {
+    await capture.runtime!.threads.switchToThread("adk-1");
+  });
+  await act(async () => {
+    await new Promise((r) => setTimeout(r, 0));
+  });
+
+  await act(async () => {
+    void capture.runtime!.thread.append({
+      role: "user",
+      content: [{ type: "text", text: "delete it" }],
+    });
+    await firstDeliveredP;
+  });
+  await act(async () => {
+    await new Promise((r) => setTimeout(r, 0));
+  });
+
+  const settle = async () => {
+    await act(async () => {
+      release();
+      await new Promise((r) => setTimeout(r, 0));
+    });
+    await act(async () => {
+      await new Promise((r) => setTimeout(r, 0));
+    });
+  };
+
+  return { capture, execute, settle };
+};
+
+describe("useAdkRuntime tool ownership", () => {
+  it("does not run a frontend execute on a call ADK is about to gate", async () => {
+    const { capture, execute, settle } = await renderStreamingAdk(
+      callEvent(),
+      confirmationEvent(),
+    );
+
+    const inWindow = toolCalls(capture.runtime!)[0]!;
+    expect(inWindow.argsText).toBe('{"path":"/tmp/a"}');
+    expect(inWindow.approval).toBeUndefined();
+    expect(inWindow.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    const gated = toolCalls(capture.runtime!).find(
+      (part) => part.toolCallId === "adk-original-1",
+    )!;
+    expect(gated.approval).toEqual({ id: "adk-confirmation-1" });
+    expect(gated.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("does not run a frontend execute on a call ADK resolves itself", async () => {
+    const { capture, execute, settle } = await renderStreamingAdk(
+      callEvent(),
+      serverResultEvent(),
+    );
+
+    expect(toolCalls(capture.runtime!)[0]!.result).toBeUndefined();
+    expect(execute).not.toHaveBeenCalled();
+
+    await settle();
+
+    expect(toolCalls(capture.runtime!)[0]!.result).toBe('{"server":"deleted"}');
+    expect(execute).not.toHaveBeenCalled();
+  });
+
+  it("runs a frontend execute on a call ADK marks long-running", async () => {
+    const { capture, execute } = await renderStreamingAdk(
+      callEvent(["adk-original-1"]),
+      serverResultEvent(),
+    );
+
+    await waitFor(() => expect(execute).toHaveBeenCalledTimes(1));
+    await waitFor(() =>
+      expect(toolCalls(capture.runtime!)[0]!.result).toBeDefined(),
+    );
+  });
+});

package/src/useAdkRuntime.ts

@@ -1,4 +1,11 @@
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 import {
   pickExternalStoreSharedOptions,
   type AttachmentAdapter,
@@ -8,6 +15,7 @@ import {
   type RealtimeVoiceAdapter,
   type SpeechSynthesisAdapter,
   type AppendMessage,
+  type ToolCallMessagePart,
   type ToolExecutionStatus,
   generateId,
 } from "@assistant-ui/core";
@@ -135,10 +143,14 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   });
 
   const loadRef = useRef(load);
-  loadRef.current = load;
+  useInsertionEffect(() => {
+    loadRef.current = load;
+  }, [load]);
   const loadController = useMemo(createAbortableThreadLoad, []);
   const messagesRef = useRef(messages);
-  messagesRef.current = messages;
+  useInsertionEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
   const [isLoadingThread, setIsLoadingThread] = useState(
     () =>
       load !== undefined && aui.threadListItem.getState().externalId != null,
@@ -153,25 +165,46 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   );
   const effectiveIsRunning = isRunning || hasExecutingTools;
   const isRunningRef = useRef(effectiveIsRunning);
-  isRunningRef.current = effectiveIsRunning;
+  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);
     }
   };
 
   const { approvals: toolApprovals, key: toolApprovalsKey } =
     projectAdkToolApprovals(messages);
+  // The messageConverter memo below reads this during render, where the ref
+  // must carry the same render's approvals; a commit-scoped write would feed
+  // the memo the previous commit's approvals whenever the key changes. No
+  // callback reads it — approval replies project from the committed messages.
   const toolApprovalsRef = useRef(toolApprovals);
   toolApprovalsRef.current = toolApprovals;
 
+  const longRunningToolIdsRef = useRef(longRunningToolIds);
+  useInsertionEffect(() => {
+    longRunningToolIdsRef.current = longRunningToolIds;
+  }, [longRunningToolIds]);
+  // ADK resolves every call it did not mark long-running itself, and yields
+  // that call to the client one or more events before its own response, so
+  // only a long-running call is the client's to execute.
+  const isClientToolCall = useCallback(
+    (toolCall: ToolCallMessagePart) =>
+      longRunningToolIdsRef.current.includes(toolCall.toolCallId),
+    [],
+  );
+
   const messageConverter = useMemo(
     () =>
       toolApprovalsKey === ""
@@ -187,10 +220,17 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   });
 
   const threadMessagesRef = useRef(threadMessages);
-  threadMessagesRef.current = threadMessages;
+  useInsertionEffect(() => {
+    threadMessagesRef.current = threadMessages;
+  }, [threadMessages]);
 
+  // Staging assigns adkMessagesRef.current directly, so the effect must key on
+  // the committed messages alone; a dep-less publication would clobber the
+  // optimistic value on any unrelated commit.
   const adkMessagesRef = useRef(messages);
-  adkMessagesRef.current = messages;
+  useInsertionEffect(() => {
+    adkMessagesRef.current = messages;
+  }, [messages]);
 
   const stagedMessagesRef = useRef(
     new Map<
@@ -300,6 +340,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     isLoading: isLoadingThread,
     messages: threadMessages,
     unstable_enableToolInvocations: true,
+    unstable_isClientToolCall: isClientToolCall,
     setToolStatuses,
     adapters: { attachments, dictation, feedback, speech, voice },
     extras: adkExtras.provide({
@@ -431,7 +472,12 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     },
     onRespondToToolApproval: async (options) => {
       await handleSendMessage(
-        [toAdkToolConfirmationReply(options, toolApprovalsRef.current)],
+        [
+          toAdkToolConfirmationReply(
+            options,
+            projectAdkToolApprovals(adkMessagesRef.current).approvals,
+          ),
+        ],
         {},
       );
     },