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

package/src/useAdkRuntime.refetch.test.tsx

@@ -1,7 +1,7 @@
 // @vitest-environment jsdom
 
 import { act, render, renderHook, waitFor } from "@testing-library/react";
-import { type FC, type ReactNode } from "react";
+import { StrictMode, type FC, type ReactNode } from "react";
 import { describe, expect, it, vi } from "vitest";
 import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
 import type {
@@ -10,7 +10,7 @@ import type {
 } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
 import { useAdkRuntime } from "./useAdkRuntime";
-import type { AdkMessage, AdkThreadSnapshot } from "./types";
+import type { AdkEvent, AdkMessage, AdkThreadSnapshot } from "./types";
 import { settleOutsideAct } from "./tests/settleOutsideAct";
 
 const deferred = <T,>() => {
@@ -72,8 +72,16 @@ const renderAdk = async (
     threadId: string,
     options?: { signal?: AbortSignal | undefined },
   ) => Promise<AdkThreadSnapshot>,
+  getCheckpointId?: (
+    threadId: string,
+    parentMessages: AdkMessage[],
+  ) => Promise<string | null>,
+  streamMock = vi.fn(async function* (
+    _messages: unknown,
+    _config: { checkpointId?: string },
+  ): AsyncGenerator<AdkEvent> {}),
+  strictMode = false,
 ) => {
-  const streamMock = vi.fn(async function* () {});
   const capture: { runtime: AssistantRuntime | null } = { runtime: null };
 
   // the runtime hook's binder mounts inside the provider, so the provider has
@@ -82,6 +90,7 @@ const renderAdk = async (
     const runtime = useAdkRuntime({
       stream: streamMock as never,
       ...(load ? { load } : {}),
+      ...(getCheckpointId ? { getCheckpointId } : {}),
       sessionAdapter: makeThreadListAdapter(),
     });
     capture.runtime = runtime;
@@ -94,7 +103,10 @@ const renderAdk = async (
 
   let unmount!: () => void;
   await act(async () => {
-    ({ unmount } = render(<Inner />));
+    const element = <Inner />;
+    ({ unmount } = render(
+      strictMode ? <StrictMode>{element}</StrictMode> : element,
+    ));
   });
   await waitFor(() => expect(capture.runtime).not.toBeNull());
 
@@ -106,6 +118,83 @@ const renderAdk = async (
 };
 
 describe("useAdkRuntime refetch", () => {
+  it("cancels pending calls from history when sending during the initial load", async () => {
+    const loaded = deferred<AdkThreadSnapshot>();
+    const streamMock = vi.fn(async function* (
+      _messages: unknown,
+      _config: { checkpointId?: string },
+    ): AsyncGenerator<AdkEvent> {});
+    const { capture } = await renderAdk(
+      () => loaded.promise,
+      undefined,
+      streamMock,
+    );
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "next question" }],
+      });
+    });
+    expect(streamMock).not.toHaveBeenCalled();
+
+    await act(async () => {
+      loaded.resolve({
+        messages: [
+          { id: "h1", type: "human", content: "earlier question" },
+          {
+            id: "a1",
+            type: "ai",
+            content: [],
+            tool_calls: [
+              { id: "cancel-me", name: "lookup", args: {} },
+              { id: "keep-me", name: "wait_for_user", args: {} },
+            ],
+          },
+        ],
+        longRunningToolIds: ["keep-me"],
+      });
+    });
+    await waitFor(() => expect(streamMock).toHaveBeenCalledOnce());
+    expect(streamMock.mock.calls[0]![0]).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          type: "tool",
+          tool_call_id: "cancel-me",
+          content: '{"cancelled":true}',
+        }),
+      ]),
+    );
+    expect(streamMock.mock.calls[0]![0]).not.toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({ tool_call_id: "keep-me" }),
+      ]),
+    );
+  });
+
+  it("drops a queued send when switching threads during the initial load", async () => {
+    const loaded = deferred<AdkThreadSnapshot>();
+    const load = vi.fn(() => loaded.promise);
+    const { capture, streamMock } = await renderAdk(load);
+    await waitFor(() => expect(load).toHaveBeenCalledOnce());
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "old thread question" }],
+      });
+    });
+    await settleOutsideAct(() => capture.runtime!.threads.switchToNewThread());
+    await act(async () => {
+      loaded.resolve({ messages: [aiMessage("old-answer", "old answer")] });
+    });
+
+    expect(streamMock).not.toHaveBeenCalled();
+    expect(
+      JSON.stringify(capture.runtime!.thread.getState().messages),
+    ).not.toContain("old thread question");
+  });
+
   it("declares the refetch capability only when a load is supplied", async () => {
     const withLoad = await renderAdk(async () => ({ messages: [] }));
     expect(
@@ -157,6 +246,99 @@ describe("useAdkRuntime refetch", () => {
     );
   });
 
+  it("waits for the initial snapshot before sending a new turn", async () => {
+    const loaded = deferred<AdkThreadSnapshot>();
+    const streamStarted = deferred<void>();
+    const finishStream = deferred<void>();
+    const load = vi.fn(() => loaded.promise);
+    const streamMock = vi.fn(async function* (): AsyncGenerator<AdkEvent> {
+      streamStarted.resolve();
+      await finishStream.promise;
+      yield {
+        id: "event-1",
+        invocationId: "run-1",
+        author: "agent",
+        content: { role: "model", parts: [{ text: "new answer" }] },
+      };
+    });
+    const { capture } = await renderAdk(load, undefined, streamMock);
+    await waitFor(() => expect(load).toHaveBeenCalledOnce());
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "new question" }],
+      });
+    });
+    expect(streamMock).not.toHaveBeenCalled();
+
+    await act(async () => {
+      loaded.resolve({
+        messages: [{ id: "h1", type: "human", content: "earlier question" }],
+      });
+      await streamStarted.promise;
+    });
+    finishStream.resolve();
+    await waitFor(() =>
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("new answer"),
+    );
+
+    const messages = JSON.stringify(
+      capture.runtime!.thread.getState().messages,
+    );
+    expect(messages).toContain("earlier question");
+    expect(messages).toContain("new question");
+    expect(messages).toContain("new answer");
+  });
+
+  it("waits for the replay-safe initial load in StrictMode", async () => {
+    const loads: Array<{
+      threadId: string;
+      result: ReturnType<typeof deferred<AdkThreadSnapshot>>;
+    }> = [];
+    const load = vi.fn((threadId: string) => {
+      const result = deferred<AdkThreadSnapshot>();
+      loads.push({ threadId, result });
+      return result.promise;
+    });
+    const streamMock = vi.fn(async function* (): AsyncGenerator<AdkEvent> {
+      yield {
+        id: "event-1",
+        invocationId: "run-1",
+        author: "agent",
+        content: { role: "model", parts: [{ text: "new answer" }] },
+      };
+    });
+    const { capture } = await renderAdk(load, undefined, streamMock, true);
+    await waitFor(() => expect(loads).toHaveLength(1));
+    expect(loads[0]?.threadId).toBe("adk-1");
+
+    act(() => {
+      capture.runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "new question" }],
+      });
+    });
+    expect(streamMock).not.toHaveBeenCalled();
+
+    loads[0]!.result.resolve({
+      messages: [{ id: "h1", type: "human", content: "earlier question" }],
+    });
+    await waitFor(() =>
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("new answer"),
+    );
+
+    const messages = JSON.stringify(
+      capture.runtime!.thread.getState().messages,
+    );
+    expect(messages).toContain("earlier question");
+    expect(messages).toContain("new question");
+  });
+
   it("swaps the per-turn state over with the messages", async () => {
     let call = 0;
     const load = vi.fn(async () => {
@@ -223,6 +405,368 @@ describe("useAdkRuntime refetch", () => {
     ).toContain("first");
   });
 
+  it("hides the answer during lookup and restores it when the lookup fails", async () => {
+    const checkpoint = deferred<string | null>();
+    const getCheckpointId = vi.fn(() => checkpoint.promise);
+    const checkpointError = new Error("checkpoint lookup failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    try {
+      const { capture } = await renderAdk(
+        async () => ({
+          messages: [
+            { id: "q-1", type: "human", content: "question" },
+            aiMessage("a-1", "answer"),
+          ],
+        }),
+        getCheckpointId,
+      );
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+      const answer = capture
+        .runtime!.thread.getState()
+        .messages.find((message) => message.role === "assistant")!;
+
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledOnce());
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("answer");
+      await act(async () => {
+        checkpoint.reject(checkpointError);
+      });
+      await waitFor(() => expect(consoleError).toHaveBeenCalledOnce());
+
+      expect(consoleError).toHaveBeenCalledExactlyOnceWith(
+        "[assistant-ui] Message reload failed",
+        checkpointError,
+      );
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("answer");
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("restores a failed reload but does not overwrite a newer send", async () => {
+    const checkpoints = [deferred<string | null>(), deferred<string | null>()];
+    const getCheckpointId = vi.fn(
+      () => checkpoints[getCheckpointId.mock.calls.length - 1]!.promise,
+    );
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    try {
+      const { capture, streamMock } = await renderAdk(
+        async () => ({
+          messages: [
+            { id: "q-1", type: "human", content: "question" },
+            aiMessage("a-1", "answer"),
+          ],
+        }),
+        getCheckpointId,
+      );
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+      const answer = capture
+        .runtime!.thread.getState()
+        .messages.find((message) => message.role === "assistant")!;
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledOnce());
+      await act(async () => {
+        checkpoints[0]!.reject(new Error("first lookup failed"));
+      });
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledTimes(2));
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("answer");
+
+      act(() => {
+        capture.runtime!.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "follow-up" }],
+        });
+      });
+      await waitFor(() => expect(streamMock).toHaveBeenCalledOnce());
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("follow-up");
+
+      await act(async () => {
+        checkpoints[1]!.reject(new Error("second lookup failed"));
+      });
+
+      expect(streamMock).toHaveBeenCalledOnce();
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("answer");
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("follow-up");
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("restores a failed reload but does not overwrite a newer load", async () => {
+    const checkpoints = [deferred<string | null>(), deferred<string | null>()];
+    const pendingLoad = deferred<AdkThreadSnapshot>();
+    const getCheckpointId = vi.fn(
+      () => checkpoints[getCheckpointId.mock.calls.length - 1]!.promise,
+    );
+    let loadCount = 0;
+    const load = vi.fn(() => {
+      loadCount++;
+      return loadCount === 1
+        ? Promise.resolve({
+            messages: [
+              { id: "q-1", type: "human" as const, content: "question" },
+              aiMessage("a-1", "answer"),
+            ],
+          })
+        : loadCount === 2
+          ? pendingLoad.promise
+          : Promise.resolve({ messages: [aiMessage("loaded", "new history")] });
+    });
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    try {
+      const { capture } = await renderAdk(load, getCheckpointId);
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+      const answer = capture
+        .runtime!.thread.getState()
+        .messages.find((m) => m.role === "assistant")!;
+
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledOnce());
+      await act(async () => {
+        checkpoints[0]!.reject(new Error("first lookup failed"));
+      });
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledTimes(2));
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("answer");
+
+      let refetch!: Promise<void>;
+      act(() => {
+        refetch = capture.runtime!.threads.reloadMainThread();
+      });
+      await waitFor(() => expect(load).toHaveBeenCalledTimes(2));
+      await act(async () => {
+        checkpoints[1]!.reject(new Error("second lookup failed"));
+      });
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("answer");
+
+      await act(async () => {
+        pendingLoad.resolve({ messages: [aiMessage("loaded", "new history")] });
+        await refetch;
+      });
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("answer");
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("new history");
+
+      await act(async () => {
+        await capture.runtime!.threads.reloadMainThread();
+      });
+      expect(load).toHaveBeenCalledTimes(3);
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("new history");
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("answer");
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it.each(["resolves", "rejects"] as const)(
+    "restores the answer when a refetch %s during a failed lookup",
+    async (outcome) => {
+      const checkpoint = deferred<string | null>();
+      const pendingLoad = deferred<AdkThreadSnapshot>();
+      const getCheckpointId = vi.fn(() => checkpoint.promise);
+      let loadCount = 0;
+      const load = vi.fn(async () => {
+        loadCount++;
+        return loadCount === 1
+          ? {
+              messages: [
+                { id: "q-1", type: "human" as const, content: "question" },
+                aiMessage("a-1", "answer"),
+              ],
+            }
+          : pendingLoad.promise;
+      });
+      const consoleError = vi
+        .spyOn(console, "error")
+        .mockImplementation(() => {});
+      try {
+        const { capture } = await renderAdk(load, getCheckpointId);
+        await waitFor(() =>
+          expect(
+            JSON.stringify(capture.runtime!.thread.getState().messages),
+          ).toContain("answer"),
+        );
+        const answer = capture
+          .runtime!.thread.getState()
+          .messages.find((message) => message.role === "assistant")!;
+
+        act(() => {
+          capture.runtime!.thread.getMessageById(answer.id).reload();
+        });
+        await waitFor(() => expect(getCheckpointId).toHaveBeenCalledOnce());
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).not.toContain("answer");
+
+        let refetch!: Promise<void>;
+        act(() => {
+          refetch = capture.runtime!.threads.reloadMainThread();
+        });
+        await waitFor(() => expect(load).toHaveBeenCalledTimes(2));
+        await act(async () => {
+          if (outcome === "resolves") {
+            pendingLoad.resolve({
+              messages: [aiMessage("loaded", "new history")],
+            });
+            await refetch;
+          } else {
+            pendingLoad.reject(new Error("refetch failed"));
+            await expect(refetch).rejects.toThrow("refetch failed");
+          }
+        });
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).not.toContain("answer");
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).not.toContain("new history");
+
+        await act(async () => {
+          checkpoint.reject(new Error("checkpoint lookup failed"));
+        });
+        await waitFor(() => expect(consoleError).toHaveBeenCalledOnce());
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer");
+      } finally {
+        consoleError.mockRestore();
+      }
+    },
+  );
+
+  it("restores a failed reload but does not overwrite a newer reload", async () => {
+    const checkpoints = [
+      deferred<string | null>(),
+      deferred<string | null>(),
+      deferred<string | null>(),
+    ];
+    const getCheckpointId = vi.fn(
+      () => checkpoints[getCheckpointId.mock.calls.length - 1]!.promise,
+    );
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    try {
+      const { capture, streamMock } = await renderAdk(
+        async () => ({
+          messages: [
+            { id: "q-1", type: "human", content: "question" },
+            aiMessage("a-1", "answer"),
+          ],
+        }),
+        getCheckpointId,
+      );
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+      const answer = capture
+        .runtime!.thread.getState()
+        .messages.find((message) => message.role === "assistant")!;
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledOnce());
+      await act(async () => {
+        checkpoints[0]!.reject(new Error("first lookup failed"));
+      });
+      await waitFor(() =>
+        expect(
+          JSON.stringify(capture.runtime!.thread.getState().messages),
+        ).toContain("answer"),
+      );
+
+      act(() => {
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+        capture.runtime!.thread.getMessageById(answer.id).reload();
+      });
+      await waitFor(() => expect(getCheckpointId).toHaveBeenCalledTimes(3));
+
+      await act(async () => {
+        checkpoints[1]!.reject(new Error("older lookup failed"));
+      });
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).not.toContain("answer");
+
+      await act(async () => {
+        checkpoints[2]!.reject(new Error("newer lookup failed"));
+      });
+
+      expect(streamMock).not.toHaveBeenCalled();
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("answer");
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
   it("defers to an initial load still in flight rather than taking it over", async () => {
     const pending = deferred<AdkThreadSnapshot>();
     let call = 0;