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

package/src/useAdkMessages.fast-refresh.test.tsx

@@ -0,0 +1,114 @@
+// @vitest-environment jsdom
+
+import { act, StrictMode, useEffect } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+let renderer: RendererInternals | undefined;
+const fiberRoots = new Set<unknown>();
+vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+  supportsFiber: true,
+  inject: (internals: RendererInternals) => {
+    renderer = internals;
+    return 1;
+  },
+  onScheduleFiberRoot: () => {},
+  onCommitFiberRoot: (_id: number, root: unknown) => fiberRoots.add(root),
+  onCommitFiberUnmount: () => {},
+});
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+const { useAdkMessages } = await import("./useAdkMessages");
+
+afterEach(() => {
+  cleanup();
+  renderer?.setRefreshHandler(() => undefined);
+  fiberRoots.clear();
+});
+afterAll(() => vi.unstubAllGlobals());
+
+const refresh = async (Before: unknown, After: unknown) => {
+  const family: Family = { current: After };
+  renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const root of fiberRoots) {
+      renderer!.scheduleRefresh(root, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+};
+
+it("keeps a direct useAdkMessages stream through Fast Refresh and cancels on unmount", async () => {
+  let signal: AbortSignal | undefined;
+  let send: (() => void) | undefined;
+  let rendered: string | undefined;
+  const stream = vi.fn((_messages, config: { abortSignal: AbortSignal }) => {
+    signal = config.abortSignal;
+    return new Promise<never>(() => {});
+  });
+  const host = (name: string) => () => {
+    rendered = name;
+    const { sendMessage } = useAdkMessages({ stream: stream as never });
+    send = () => {
+      void sendMessage([{ id: "user", type: "human", content: "hello" }], {});
+    };
+    return null;
+  };
+  const Before = host("before");
+  const After = host("after");
+  const view = render(<Before />);
+
+  act(() => send!());
+  await waitFor(() => expect(stream).toHaveBeenCalledTimes(1));
+  expect(signal!.aborted).toBe(false);
+
+  await refresh(Before, After);
+  expect(rendered).toBe("after");
+  expect(signal!.aborted).toBe(false);
+  expect(stream).toHaveBeenCalledTimes(1);
+
+  view.unmount();
+  await act(async () => {});
+  expect(signal!.aborted).toBe(true);
+});
+
+it("keeps a stream started during StrictMode's effect replay", async () => {
+  let signal: AbortSignal | undefined;
+  let launched = false;
+  const stream = vi.fn((_messages, config: { abortSignal: AbortSignal }) => {
+    signal = config.abortSignal;
+    return new Promise<never>(() => {});
+  });
+  const Host = () => {
+    const { sendMessage } = useAdkMessages({ stream: stream as never });
+    useEffect(() => {
+      if (launched) return;
+      launched = true;
+      void sendMessage([{ id: "user", type: "human", content: "hello" }], {});
+    }, [sendMessage]);
+    return null;
+  };
+  const view = render(
+    <StrictMode>
+      <Host />
+    </StrictMode>,
+  );
+  await waitFor(() => expect(stream).toHaveBeenCalledTimes(1));
+  expect(signal!.aborted).toBe(false);
+  view.unmount();
+  await act(async () => {});
+  expect(signal!.aborted).toBe(true);
+});

package/src/useAdkMessages.test.ts

@@ -77,6 +77,40 @@ describe("optimistic tool outcomes", () => {
 });
 
 describe("ADK runtime callbacks", () => {
+  it("reports the same agent transfer again in a later run", async () => {
+    const onAgentTransfer = vi.fn();
+    const stream: AdkStreamCallback = async function* () {
+      yield {
+        id: "transfer",
+        actions: { transferToAgent: "researcher" },
+      };
+      yield {
+        id: "transfer-duplicate",
+        actions: { transferToAgent: "researcher" },
+      };
+    };
+    const { result } = renderHook(() =>
+      useAdkMessages({ stream, eventHandlers: { onAgentTransfer } }),
+    );
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "first" }],
+        {},
+      );
+      expect(onAgentTransfer).toHaveBeenCalledTimes(1);
+
+      await result.current.sendMessage(
+        [{ id: "user-2", type: "human", content: "second" }],
+        {},
+      );
+    });
+
+    expect(onAgentTransfer).toHaveBeenCalledTimes(2);
+    expect(onAgentTransfer).toHaveBeenNthCalledWith(1, "researcher");
+    expect(onAgentTransfer).toHaveBeenNthCalledWith(2, "researcher");
+  });
+
   it.each(["onAgentTransfer", "onCustomEvent", "onError"] as const)(
     "continues streaming when %s throws",
     async (callbackName) => {
@@ -305,6 +339,7 @@ describe("ADK stream lifecycle", () => {
     await started;
 
     unmount();
+    await act(async () => {});
 
     expect(runSignal?.aborted).toBe(true);
     await expect(sendPromise).resolves.toBeUndefined();

package/src/useAdkMessages.ts

@@ -1,12 +1,12 @@
 import {
   useState,
   useCallback,
-  useEffect,
   useInsertionEffect,
   useRef,
   useMemo,
 } from "react";
 import { generateId } from "@assistant-ui/core";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 import { useAui } from "@assistant-ui/store";
 import {
   abortableIterable,
@@ -71,7 +71,6 @@ export const useAdkMessages = ({
   const [messageMetadata, setMessageMetadata] = useState<
     Map<string, AdkMessageMetadata>
   >(new Map());
-  const lastTransferToAgentRef = useRef<string | undefined>(undefined);
   // 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);
@@ -172,6 +171,7 @@ export const useAdkMessages = ({
       setLongRunningToolIds(accumulator.getLongRunningToolIds());
       setToolConfirmations(accumulator.getToolConfirmations());
       setAuthRequests(accumulator.getAuthRequests());
+      let lastTransferToAgent: string | undefined;
 
       // Google ADK replaces active runs, while React LangGraph queues sends.
       abortControllerRef.current?.abort();
@@ -226,8 +226,8 @@ export const useAdkMessages = ({
           }
 
           const transfer = accumulator.getLastTransferToAgent();
-          if (transfer && transfer !== lastTransferToAgentRef.current) {
-            lastTransferToAgentRef.current = transfer;
+          if (transfer && transfer !== lastTransferToAgent) {
+            lastTransferToAgent = transfer;
             invokeAdkRuntimeCallback(
               "onAgentTransfer",
               onAgentTransfer,
@@ -286,7 +286,7 @@ export const useAdkMessages = ({
     }
   }, []);
 
-  useEffect(() => cancel, [cancel]);
+  useReplaySafeEffect(() => cancel, []);
 
   return {
     messages,

package/src/useAdkRuntime.refetch.test.tsx

@@ -11,6 +11,7 @@ import type {
 import { useAui } from "@assistant-ui/store";
 import { useAdkRuntime } from "./useAdkRuntime";
 import type { AdkMessage, AdkThreadSnapshot } from "./types";
+import { settleOutsideAct } from "./tests/settleOutsideAct";
 
 const deferred = <T,>() => {
   let resolve!: (value: T) => void;
@@ -97,9 +98,9 @@ const renderAdk = async (
   });
   await waitFor(() => expect(capture.runtime).not.toBeNull());
 
-  await act(async () => {
-    await capture.runtime!.threads.switchToThread("adk-1");
-  });
+  await settleOutsideAct(() =>
+    capture.runtime!.threads.switchToThread("adk-1"),
+  );
 
   return { capture, streamMock, unmount };
 };
@@ -250,9 +251,9 @@ describe("useAdkRuntime refetch", () => {
     await act(async () => {
       render(<Inner />);
     });
-    await act(async () => {
-      await capture.runtime!.threads.switchToThread("adk-1");
-    });
+    await settleOutsideAct(() =>
+      capture.runtime!.threads.switchToThread("adk-1"),
+    );
     await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
 
     let settled = false;

package/src/useAdkRuntime.replacement.test.tsx

@@ -10,6 +10,7 @@ import type {
 } from "@assistant-ui/core";
 import { useAdkRuntime } from "./useAdkRuntime";
 import type { AdkEvent } from "./types";
+import { settleOutsideAct } from "./tests/settleOutsideAct";
 
 const makeThreadListAdapter = (): RemoteThreadListAdapter => ({
   list: vi.fn(async () => ({
@@ -97,9 +98,9 @@ describe("useAdkRuntime replacement runs", () => {
       render(<Inner />);
     });
     await waitFor(() => expect(capture.runtime).not.toBeNull());
-    await act(async () => {
-      await capture.runtime!.threads.switchToThread("adk-1");
-    });
+    await settleOutsideAct(() =>
+      capture.runtime!.threads.switchToThread("adk-1"),
+    );
 
     act(() => {
       capture.runtime!.thread.append({

package/src/useAdkRuntime.toolOwnership.test.tsx

@@ -14,6 +14,7 @@ import type {
 } from "@assistant-ui/core";
 import { useAdkRuntime } from "./useAdkRuntime";
 import type { AdkEvent } from "./types";
+import { settleOutsideAct } from "./tests/settleOutsideAct";
 
 const makeThreadListAdapter = (): RemoteThreadListAdapter => ({
   list: vi.fn(async () => ({
@@ -153,9 +154,9 @@ const renderStreamingAdk = async (first: AdkEvent, second: AdkEvent) => {
     render(<Inner />);
   });
   await waitFor(() => expect(capture.runtime).not.toBeNull());
-  await act(async () => {
-    await capture.runtime!.threads.switchToThread("adk-1");
-  });
+  await settleOutsideAct(() =>
+    capture.runtime!.threads.switchToThread("adk-1"),
+  );
   await act(async () => {
     await new Promise((r) => setTimeout(r, 0));
   });