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

package/src/useAdkRuntime.cancellation.test.tsx

@@ -0,0 +1,544 @@
+// @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 { useAdkLongRunningToolIds, useAdkSend } from "./hooks";
+import { AdkEventAccumulator } from "./AdkEventAccumulator";
+import type { AdkEvent, AdkStreamCallback, AdkThreadSnapshot } 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 makePendingStream = (...events: AdkEvent[]): AdkStreamCallback =>
+  async function* (_messages, { abortSignal }) {
+    yield* events;
+    await new Promise<void>((resolve) => {
+      abortSignal.addEventListener("abort", () => resolve(), { once: true });
+    });
+  };
+
+const renderAdkRuntime = async (
+  stream: AdkStreamCallback,
+  snapshot?: AdkThreadSnapshot,
+) => {
+  const capture: {
+    runtime: AssistantRuntime | null;
+    longRunningToolIds: string[];
+    send: ReturnType<typeof useAdkSend> | null;
+  } = { runtime: null, longRunningToolIds: [], send: null };
+  const sessionAdapter = makeThreadListAdapter();
+
+  const CaptureExtras: FC = () => {
+    capture.longRunningToolIds = useAdkLongRunningToolIds();
+    capture.send = useAdkSend();
+    return null;
+  };
+
+  const Inner: FC = () => {
+    const runtime = useAdkRuntime({
+      stream,
+      sessionAdapter,
+      unstable_allowCancellation: true,
+      ...(snapshot && { load: async () => snapshot }),
+    });
+    capture.runtime = runtime;
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <CaptureExtras />
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  await act(async () => {
+    render(<Inner />);
+  });
+  await waitFor(() => expect(capture.runtime).not.toBeNull());
+  await act(async () => {
+    await capture.runtime!.threads.switchToThread("adk-1");
+  });
+
+  return capture;
+};
+
+describe("useAdkRuntime cancellation", () => {
+  it.each(["none", "send", "stream"] as const)(
+    "preserves inherited confirmation IDs on Stop (answer: %s)",
+    async (answer) => {
+      const seed = new AdkEventAccumulator();
+      seed.processEvent({
+        id: "older-confirmations",
+        author: "agent",
+        longRunningToolIds: ["old-1", "old-2"],
+        content: {
+          role: "model",
+          parts: ["old-1", "old-2"].map((id) => ({
+            functionCall: {
+              id,
+              name: "adk_request_confirmation",
+              args: {
+                originalFunctionCall: {
+                  id: `target-${id}`,
+                  name: "lookup",
+                  args: {},
+                },
+                toolConfirmation: { hint: "Allow lookup?" },
+              },
+            },
+          })),
+        },
+      });
+      const stream = vi.fn(
+        makePendingStream(
+          ...(answer === "stream"
+            ? [
+                {
+                  id: "settled-confirmation",
+                  author: "user",
+                  content: {
+                    role: "user",
+                    parts: [
+                      {
+                        functionResponse: {
+                          id: "old-1",
+                          name: "adk_request_confirmation",
+                          response: { confirmed: true },
+                        },
+                      },
+                    ],
+                  },
+                } satisfies AdkEvent,
+              ]
+            : []),
+          {
+            id: "new-tool-call",
+            author: "agent",
+            longRunningToolIds: ["new-call"],
+            content: {
+              role: "model",
+              parts: [
+                {
+                  functionCall: {
+                    id: "new-call",
+                    name: "lookup",
+                    args: {},
+                  },
+                },
+              ],
+            },
+          },
+        ),
+      );
+      const capture = await renderAdkRuntime(stream, {
+        messages: seed.getMessages(),
+        longRunningToolIds: seed.getLongRunningToolIds(),
+        toolConfirmations: seed.getToolConfirmations(),
+      });
+      await waitFor(() =>
+        expect(capture.longRunningToolIds).toEqual(["old-1", "old-2"]),
+      );
+
+      act(() => {
+        if (answer === "send") {
+          void capture.send!(
+            [
+              {
+                id: "approval-reply",
+                type: "tool",
+                tool_call_id: "old-1",
+                name: "adk_request_confirmation",
+                content: '{"confirmed":true}',
+              },
+            ],
+            {},
+          );
+        } else {
+          capture.runtime!.thread.append({
+            role: "user",
+            content: [{ type: "text", text: "continue" }],
+          });
+        }
+      });
+      const pendingInherited =
+        answer === "none" ? ["old-1", "old-2"] : ["old-2"];
+      await waitFor(() =>
+        expect(capture.longRunningToolIds).toEqual([
+          ...pendingInherited,
+          "new-call",
+        ]),
+      );
+
+      act(() => {
+        capture.runtime!.thread.cancelRun();
+      });
+      await waitFor(() =>
+        expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+      );
+      expect(capture.longRunningToolIds).toEqual(pendingInherited);
+
+      act(() => {
+        capture.runtime!.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "next turn" }],
+        });
+      });
+      await waitFor(() => expect(stream).toHaveBeenCalledTimes(2));
+      const cancellations = stream.mock.calls[1]![0].filter(
+        (m) => m.type === "tool",
+      );
+      expect(cancellations).toMatchObject([
+        { tool_call_id: "new-call", content: '{"cancelled":true}' },
+      ]);
+    },
+  );
+
+  it("preserves a user message staged during a cancelled run", async () => {
+    const capture = await renderAdkRuntime(
+      makePendingStream({
+        id: "partial",
+        author: "agent",
+        partial: true,
+        content: { role: "model", parts: [{ text: "partial answer" }] },
+      }),
+    );
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      });
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().messages.at(-1)).toMatchObject({
+        role: "assistant",
+        content: [{ type: "text", text: "partial answer" }],
+      }),
+    );
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "staged follow-up" }],
+        startRun: false,
+      });
+    });
+    await waitFor(() =>
+      expect(
+        capture
+          .runtime!.thread.getState()
+          .messages.filter((m) => m.role === "user")
+          .at(-1),
+      ).toMatchObject({
+        content: [{ type: "text", text: "staged follow-up" }],
+      }),
+    );
+    const stagedMessage = capture
+      .runtime!.thread.getState()
+      .messages.filter((m) => m.role === "user")
+      .at(-1);
+    expect(capture.runtime!.thread.getState().isRunning).toBe(true);
+
+    act(() => {
+      capture.runtime!.thread.cancelRun();
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+    );
+    expect(
+      capture
+        .runtime!.thread.getState()
+        .messages.find((m) => m.id === stagedMessage!.id),
+    ).toEqual(stagedMessage);
+    expect(
+      capture
+        .runtime!.thread.getState()
+        .messages.find((m) => m.role === "assistant")?.status,
+    ).toEqual({ type: "incomplete", reason: "cancelled" });
+  });
+
+  it.each([false, true])(
+    "cancels a finalized tool-call message (long-running: %s)",
+    async (longRunning) => {
+      const stream = vi.fn(
+        makePendingStream({
+          id: "tool-call",
+          author: "agent",
+          ...(longRunning && { longRunningToolIds: ["call-1"] }),
+          content: {
+            role: "model",
+            parts: [
+              { functionCall: { id: "call-1", name: "lookup", args: {} } },
+            ],
+          },
+        }),
+      );
+      const capture = await renderAdkRuntime(stream);
+
+      act(() => {
+        capture.runtime!.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "hello" }],
+        });
+      });
+      await waitFor(() =>
+        expect(
+          capture.runtime!.thread.getState().messages.at(-1),
+        ).toMatchObject({
+          role: "assistant",
+          content: [{ type: "tool-call", toolCallId: "call-1" }],
+        }),
+      );
+
+      expect(capture.longRunningToolIds).toEqual(longRunning ? ["call-1"] : []);
+
+      act(() => {
+        capture.runtime!.thread.cancelRun();
+      });
+      await waitFor(() =>
+        expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+      );
+
+      expect(
+        capture.runtime!.thread.getState().messages.at(-1)?.status,
+      ).toEqual({
+        type: "incomplete",
+        reason: "cancelled",
+      });
+      expect(capture.longRunningToolIds).toEqual([]);
+
+      act(() => {
+        capture.runtime!.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "next turn" }],
+        });
+      });
+      await waitFor(() => expect(stream).toHaveBeenCalledTimes(2));
+      expect(stream.mock.calls[1]![0]).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({
+            type: "tool",
+            tool_call_id: "call-1",
+            content: '{"cancelled":true}',
+          }),
+        ]),
+      );
+    },
+  );
+
+  it("cancels the assistant message when the last event is a tool response", async () => {
+    const capture = await renderAdkRuntime(
+      makePendingStream(
+        {
+          id: "tool-call",
+          author: "agent",
+          content: {
+            role: "model",
+            parts: [
+              { functionCall: { id: "call-1", name: "lookup", args: {} } },
+            ],
+          },
+        },
+        {
+          id: "tool-response",
+          author: "agent",
+          content: {
+            role: "user",
+            parts: [
+              {
+                functionResponse: {
+                  id: "call-1",
+                  name: "lookup",
+                  response: { result: "found" },
+                },
+              },
+            ],
+          },
+        },
+      ),
+    );
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      });
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().messages.at(-1)).toMatchObject({
+        role: "assistant",
+        content: [{ type: "tool-call", result: '{"result":"found"}' }],
+      }),
+    );
+
+    act(() => {
+      capture.runtime!.thread.cancelRun();
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+    );
+    expect(capture.runtime!.thread.getState().messages.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+  });
+
+  it("keeps an earlier assistant message unchanged when stopped before new output", async () => {
+    let sends = 0;
+    const capture = await renderAdkRuntime(async function* (messages, config) {
+      if (sends++ === 0) {
+        yield {
+          id: "earlier-tool-call",
+          author: "agent",
+          content: {
+            role: "model",
+            parts: [
+              { functionCall: { id: "call-1", name: "lookup", args: {} } },
+            ],
+          },
+        };
+      } else {
+        yield* await makePendingStream()(messages, config);
+      }
+    });
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => {
+      expect(sends).toBe(1);
+      expect(capture.runtime!.thread.getState().isRunning).toBe(false);
+    });
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "second" }],
+      });
+    });
+    await waitFor(() => expect(sends).toBe(2));
+    const earlierMessage = capture
+      .runtime!.thread.getState()
+      .messages.find((m) => m.role === "assistant");
+    expect(earlierMessage).toBeDefined();
+
+    act(() => {
+      capture.runtime!.thread.cancelRun();
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+    );
+    expect(
+      capture
+        .runtime!.thread.getState()
+        .messages.find((m) => m.id === earlierMessage!.id),
+    ).toEqual(earlierMessage);
+  });
+
+  it("marks a partial assistant message as cancelled when Stop is pressed", async () => {
+    const capture = await renderAdkRuntime(
+      makePendingStream({
+        id: "partial",
+        invocationId: "run-1",
+        author: "agent",
+        partial: true,
+        content: { role: "model", parts: [{ text: "partial answer" }] },
+      }),
+    );
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      });
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().messages.at(-1)).toMatchObject({
+        role: "assistant",
+        content: [{ type: "text", text: "partial answer" }],
+      }),
+    );
+
+    act(() => {
+      capture.runtime!.thread.cancelRun();
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+    );
+
+    expect(capture.runtime!.thread.getState().messages.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+  });
+
+  it("keeps a completed assistant message complete when Stop is pressed", async () => {
+    const capture = await renderAdkRuntime(
+      makePendingStream({
+        id: "complete",
+        invocationId: "run-1",
+        author: "agent",
+        content: { role: "model", parts: [{ text: "final answer" }] },
+      }),
+    );
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      });
+    });
+    await waitFor(() =>
+      expect(
+        capture.runtime!.thread.getState().messages.at(-1)?.status,
+      ).toEqual({
+        type: "complete",
+        reason: "stop",
+      }),
+    );
+    expect(capture.runtime!.thread.getState().isRunning).toBe(true);
+
+    act(() => {
+      capture.runtime!.thread.cancelRun();
+    });
+    await waitFor(() =>
+      expect(capture.runtime!.thread.getState().isRunning).toBe(false),
+    );
+
+    expect(capture.runtime!.thread.getState().messages.at(-1)?.status).toEqual({
+      type: "complete",
+      reason: "stop",
+    });
+  });
+});

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

@@ -0,0 +1,166 @@
+// @vitest-environment jsdom
+
+import { act } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+import type { AdkThreadSnapshot } from "./types";
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+const refreshHarness = vi.hoisted(() => {
+  const state: {
+    renderer: RendererInternals | undefined;
+    roots: Set<unknown>;
+  } = { renderer: undefined, roots: new Set() };
+  vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+  vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+    supportsFiber: true,
+    inject: (internals: RendererInternals) => {
+      state.renderer = internals;
+      return 1;
+    },
+    onScheduleFiberRoot: () => {},
+    onCommitFiberRoot: (_id: number, root: unknown) => state.roots.add(root),
+    onCommitFiberUnmount: () => {},
+  });
+  return state;
+});
+
+const { aui } = vi.hoisted(() => ({
+  aui: {
+    threadListItem: {
+      source: {},
+      getState: () => ({ externalId: "adk-1" }),
+    },
+  },
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => aui,
+}));
+
+vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/core/react")>()),
+  useCloudThreadListAdapter: () => ({}),
+  useRemoteThreadListRuntime: ({
+    runtimeHook,
+  }: {
+    runtimeHook: () => unknown;
+  }) => runtimeHook(),
+  useExternalMessageConverter: ({ messages }: { messages: unknown }) =>
+    messages,
+  useExternalStoreRuntime: (options: unknown) => options,
+}));
+
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+const { useAdkRuntime } = await import("./useAdkRuntime");
+
+afterEach(() => {
+  cleanup();
+  refreshHarness.renderer?.setRefreshHandler(() => undefined);
+  refreshHarness.roots.clear();
+});
+afterAll(() => vi.unstubAllGlobals());
+
+const deferred = <T,>() => {
+  let resolve!: (value: T) => void;
+  const promise = new Promise<T>((res) => {
+    resolve = res;
+  });
+  return { promise, resolve };
+};
+
+const refresh = async (Before: unknown, After: unknown) => {
+  const family: Family = { current: After };
+  refreshHarness.renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const root of refreshHarness.roots) {
+      refreshHarness.renderer!.scheduleRefresh(root, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+};
+
+it("keeps an initial thread load through Fast Refresh and aborts a later load on unmount", async () => {
+  const initial = deferred<AdkThreadSnapshot>();
+  const refetch = deferred<AdkThreadSnapshot>();
+  const signals: AbortSignal[] = [];
+  const load = vi.fn(
+    (_id: string, options?: { signal?: AbortSignal | undefined }) => {
+      signals.push(options!.signal!);
+      return signals.length === 1 ? initial.promise : refetch.promise;
+    },
+  );
+  let runtime:
+    | {
+        messages: unknown;
+        onRefetchThread: () => Promise<void>;
+      }
+    | undefined;
+  let rendered: string | undefined;
+  const host = (name: string) => () => {
+    rendered = name;
+    runtime = useAdkRuntime({
+      stream: async function* () {},
+      load,
+    }) as unknown as typeof runtime;
+    return null;
+  };
+  const Before = host("before");
+  const After = host("after");
+  const view = render(<Before />);
+  await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+  expect(signals[0]!.aborted).toBe(false);
+
+  await refresh(Before, After);
+  expect(rendered).toBe("after");
+  expect(signals[0]!.aborted).toBe(false);
+  expect(load).toHaveBeenCalledTimes(1);
+
+  let refetchSettled = false;
+  act(() => {
+    void runtime!.onRefetchThread().then(() => {
+      refetchSettled = true;
+    });
+  });
+  expect(load).toHaveBeenCalledTimes(1);
+  expect(refetchSettled).toBe(false);
+
+  await act(async () => {
+    initial.resolve({
+      messages: [
+        {
+          id: "history-1",
+          type: "ai",
+          content: [{ type: "text", text: "history landed" }],
+        },
+      ],
+    });
+  });
+  await waitFor(() => expect(refetchSettled).toBe(true));
+  await waitFor(() =>
+    expect(JSON.stringify(runtime!.messages)).toContain("history landed"),
+  );
+
+  act(() => {
+    void runtime!.onRefetchThread();
+  });
+  await waitFor(() => expect(load).toHaveBeenCalledTimes(2));
+  expect(signals[1]!.aborted).toBe(false);
+  view.unmount();
+  await act(async () => {});
+  expect(signals[1]!.aborted).toBe(true);
+  refetch.resolve({ messages: [] });
+});

package/src/useAdkRuntime.refetch.test.tsx

@@ -330,6 +330,7 @@ describe("useAdkRuntime refetch", () => {
     expect(signals[1]?.aborted).toBe(false);
 
     unmount();
+    await act(async () => {});
 
     expect(signals[1]?.aborted).toBe(true);
     pending.resolve({ messages: [] });