@assistant-ui/react-google-adk 0.0.20 → 0.0.21

package/src/AdkSessionAdapter.test.ts

@@ -293,6 +293,49 @@ describe("createAdkSessionAdapter - fetch", () => {
 // ── load() ──
 
 describe("createAdkSessionAdapter - load", () => {
+  it("returns the per-turn state the events imply, not just the messages", async () => {
+    const session = {
+      id: "s1",
+      events: [
+        {
+          id: "e1",
+          author: "agent",
+          longRunningToolIds: ["tc-1"],
+          actions: { stateDelta: { step: 2 }, escalate: true },
+          content: {
+            role: "model",
+            parts: [{ functionCall: { name: "search", id: "tc-1", args: {} } }],
+          },
+        },
+      ],
+    };
+    mockFetch.mockResolvedValueOnce(
+      new Response(JSON.stringify(session), { status: 200 }),
+    );
+
+    const { load } = createAdkSessionAdapter(baseOptions);
+    const result = await load("s1");
+
+    expect(result.longRunningToolIds).toEqual(["tc-1"]);
+    expect(result.stateDelta).toEqual({ step: 2 });
+    expect(result.messageMetadata).toBeInstanceOf(Map);
+    expect(result.toolConfirmations).toEqual([]);
+    expect(result.authRequests).toEqual([]);
+  });
+
+  it("passes an abort signal through to the request", async () => {
+    mockFetch.mockResolvedValueOnce(
+      new Response(JSON.stringify({ id: "s1", events: [] }), { status: 200 }),
+    );
+    const controller = new AbortController();
+
+    const { load } = createAdkSessionAdapter(baseOptions);
+    await load("s1", { signal: controller.signal });
+
+    const init = mockFetch.mock.calls[0]![1] as RequestInit;
+    expect(init.signal).toBe(controller.signal);
+  });
+
   it("reconstructs messages from session events", async () => {
     const session = {
       id: "s1",

package/src/AdkSessionAdapter.ts

@@ -6,7 +6,7 @@ import type {
   RemoteThreadMetadata,
 } from "@assistant-ui/core";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
-import type { AdkEvent, AdkMessage } from "./types";
+import type { AdkEvent, AdkMessage, AdkThreadSnapshot } from "./types";
 import { trimTrailingSlashes } from "./trimTrailingSlashes";
 
 export type AdkSessionAdapterOptions = {
@@ -41,7 +41,10 @@ export type AdkArtifactData = {
 
 type AdkSessionAdapterResult = {
   adapter: RemoteThreadListAdapter;
-  load: (sessionId: string) => Promise<{ messages: AdkMessage[] }>;
+  load: (
+    sessionId: string,
+    options?: { signal?: AbortSignal | undefined },
+  ) => Promise<AdkThreadSnapshot>;
   artifacts: {
     list: (sessionId: string) => Promise<string[]>;
     load: (
@@ -302,10 +305,12 @@ export function createAdkSessionAdapter(
 
   const load = async (
     sessionId: string,
-  ): Promise<{ messages: AdkMessage[] }> => {
+    options?: { signal?: AbortSignal | undefined },
+  ): Promise<AdkThreadSnapshot> => {
     const headers = await getHeaders();
     const res = await fetch(`${baseUrl}/${encodeURIComponent(sessionId)}`, {
       headers,
+      ...(options?.signal ? { signal: options.signal } : {}),
     });
     if (!res.ok) {
       throw new Error(`Failed to load session: ${res.status}`);
@@ -329,7 +334,19 @@ export function createAdkSessionAdapter(
     for (const event of events) {
       messages = accumulator.processEvent(event);
     }
-    return { messages };
+    // The per-turn state rides along, so a refetch can swap the thread over in
+    // one commit instead of reconstructing it from the messages alone.
+    return {
+      messages,
+      longRunningToolIds: accumulator.getLongRunningToolIds(),
+      toolConfirmations: accumulator.getToolConfirmations(),
+      authRequests: accumulator.getAuthRequests(),
+      escalated: accumulator.isEscalated(),
+      messageMetadata: accumulator.getMessageMetadata(),
+      stateDelta: accumulator.getStateDelta(),
+      artifactDelta: accumulator.getArtifactDelta(),
+      agentInfo: accumulator.getAgentInfo(),
+    };
   };
 
   const artifactBaseUrl = (sessionId: string) =>

package/src/convertAdkMessages.test.ts

@@ -26,32 +26,66 @@ describe("convertAdkMessage - human messages", () => {
     });
   });
 
-  it("restores an audio/mp3 file part as an audio message part", () => {
+  it("restores a file_url part as a file part stamped with sourceType url", () => {
     const msg: AdkMessage = {
       id: "m1",
       type: "human",
-      content: [{ type: "file", mimeType: "audio/mp3", data: "QUJD" }],
+      content: [
+        {
+          type: "file_url",
+          url: "gs://bucket/report.pdf",
+          mimeType: "application/pdf",
+        },
+      ],
     };
     const result = convertAdkMessage(msg, {});
     expect(result).toMatchObject({
       role: "user",
-      content: [{ type: "audio", audio: { data: "QUJD", format: "mp3" } }],
+      content: [
+        {
+          type: "file",
+          data: "gs://bucket/report.pdf",
+          mimeType: "application/pdf",
+          sourceType: "url",
+        },
+      ],
     });
   });
 
-  it("restores an audio/wav file part as an audio message part", () => {
+  it("falls back to application/octet-stream for file_url parts without mimeType", () => {
     const msg: AdkMessage = {
       id: "m1",
       type: "human",
-      content: [{ type: "file", mimeType: "audio/wav", data: "QUJD" }],
+      content: [{ type: "file_url", url: "gs://bucket/blob" }],
     };
     const result = convertAdkMessage(msg, {});
     expect(result).toMatchObject({
       role: "user",
-      content: [{ type: "audio", audio: { data: "QUJD", format: "wav" } }],
+      content: [
+        {
+          type: "file",
+          data: "gs://bucket/blob",
+          mimeType: "application/octet-stream",
+          sourceType: "url",
+        },
+      ],
     });
   });
 
+  it("keeps an audio file part as a file part", () => {
+    for (const mimeType of ["audio/mp3", "audio/wav"]) {
+      const msg: AdkMessage = {
+        id: "m1",
+        type: "human",
+        content: [{ type: "file", mimeType, data: "QUJD" }],
+      };
+      expect(convertAdkMessage(msg, {})).toMatchObject({
+        role: "user",
+        content: [{ type: "file", data: "QUJD", mimeType }],
+      });
+    }
+  });
+
   it("keeps attachment-derived audio file parts (with filename) as file parts", () => {
     const msg: AdkMessage = {
       id: "m1",

package/src/convertAdkMessages.ts

@@ -8,8 +8,13 @@ type ContentPart =
   | { type: "text"; text: string }
   | { type: "reasoning"; text: string }
   | { type: "image"; image: string }
-  | { type: "file"; data: string; mimeType: string; filename?: string }
-  | { type: "audio"; audio: { data: string; format: "mp3" | "wav" } }
+  | {
+      type: "file";
+      data: string;
+      mimeType: string;
+      filename?: string;
+      sourceType?: "url";
+    }
   | { type: "data"; name: string; data: unknown };
 
 const contentToParts = (
@@ -33,26 +38,22 @@ const contentToParts = (
           };
         case "image_url":
           return { type: "image", image: part.url };
-        case "file": {
-          const format =
-            role === "user" && part.filename == null
-              ? part.mimeType === "audio/wav"
-                ? ("wav" as const)
-                : part.mimeType === "audio/mp3"
-                  ? ("mp3" as const)
-                  : null
-              : null;
-          if (format) {
-            return { type: "audio", audio: { data: part.data, format } };
-          }
+        case "file":
           return {
             type: "file",
             data: part.data,
             mimeType: part.mimeType,
             ...(part.filename != null && { filename: part.filename }),
           };
-        }
         case "file_url":
+          if (role === "user") {
+            return {
+              type: "file",
+              data: part.url,
+              mimeType: part.mimeType ?? "application/octet-stream",
+              sourceType: "url",
+            };
+          }
           return {
             type: "data",
             name: "file_url",

package/src/index.ts

@@ -9,6 +9,7 @@ export type {
   AdkEventPart,
   AdkEventActions,
   AdkMessage,
+  AdkThreadSnapshot,
   AdkMessageContentPart,
   AdkToolCall,
   AdkToolConfirmation,

package/src/types.ts

@@ -287,3 +287,24 @@ export type AdkRuntimeExtras = {
   escalated: boolean;
   messageMetadata: Map<string, AdkMessageMetadata>;
 };
+
+/**
+ * What a session load reconstructs. The messages are the whole thread; the
+ * rest is the per-turn state the events imply, which a caller that replays
+ * them through {@link AdkEventAccumulator} already has. Omitting a field
+ * clears it, since a snapshot that cannot describe the state is not evidence
+ * that the state survived.
+ */
+export type AdkThreadSnapshot = {
+  messages: AdkMessage[];
+  longRunningToolIds?: string[] | undefined;
+  toolConfirmations?: AdkToolConfirmation[] | undefined;
+  authRequests?: AdkAuthRequest[] | undefined;
+  escalated?: boolean | undefined;
+  messageMetadata?: Map<string, AdkMessageMetadata> | undefined;
+  stateDelta?: Record<string, unknown> | undefined;
+  artifactDelta?: Record<string, number> | undefined;
+  agentInfo?:
+    | { name?: string | undefined; branch?: string | undefined }
+    | undefined;
+};

package/src/useAdkMessages.ts

@@ -11,6 +11,7 @@ import type {
   AdkStreamCallback,
   AdkToolConfirmation,
   AdkAuthRequest,
+  AdkThreadSnapshot,
   OnAdkErrorCallback,
   OnAdkCustomEventCallback,
   OnAdkAgentTransferCallback,
@@ -62,8 +63,29 @@ export const useAdkMessages = ({
     _setMessages(msgs);
   }, []);
 
+  /**
+   * Swap the thread over to a loaded snapshot in one commit. Unlike
+   * {@link replaceMessages} this never passes through a cleared state, so a
+   * refetch that lands while a confirmation is on screen replaces it rather
+   * than blanking it first.
+   */
+  const applySnapshot = useCallback(
+    (snapshot: AdkThreadSnapshot) => {
+      setMessagesImmediate(snapshot.messages);
+      setLongRunningToolIds(snapshot.longRunningToolIds ?? []);
+      setToolConfirmations(snapshot.toolConfirmations ?? []);
+      setAuthRequests(snapshot.authRequests ?? []);
+      setEscalated(snapshot.escalated ?? false);
+      setMessageMetadata(snapshot.messageMetadata ?? new Map());
+      setStateDelta(snapshot.stateDelta ?? {});
+      setArtifactDelta(snapshot.artifactDelta ?? {});
+      setAgentInfo(snapshot.agentInfo ?? {});
+    },
+    [setMessagesImmediate],
+  );
+
   // Replace the message list AND reset derived per-turn HITL state.
-  // Used by truncation paths (edit, reload, load) so that stale interrupt
+  // Used by truncation paths (edit, reload) so that stale interrupt
   // markers and per-message metadata from the removed messages don't leak
   // into the next turn.
   const replaceMessages = useCallback(
@@ -195,6 +217,7 @@ export const useAdkMessages = ({
     cancel,
     setMessages: setMessagesImmediate,
     replaceMessages,
+    applySnapshot,
   };
 };
 

package/src/useAdkRuntime.refetch.test.tsx

@@ -0,0 +1,336 @@
+// @vitest-environment jsdom
+
+import { act, render, renderHook, waitFor } from "@testing-library/react";
+import { type FC, type ReactNode } 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 { useAui } from "@assistant-ui/store";
+import { useAdkRuntime } from "./useAdkRuntime";
+import type { AdkMessage, AdkThreadSnapshot } from "./types";
+
+const deferred = <T,>() => {
+  let resolve!: (value: T) => void;
+  let reject!: (reason: unknown) => void;
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+  return { promise, resolve, reject };
+};
+
+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 wrapperFactory = (runtime: AssistantRuntime) => {
+  const Wrapper = ({ children }: { children: ReactNode }) => (
+    <AssistantRuntimeProvider runtime={runtime}>
+      {children}
+    </AssistantRuntimeProvider>
+  );
+  Wrapper.displayName = "TestWrapper";
+  return Wrapper;
+};
+
+const aiMessage = (id: string, text: string): AdkMessage => ({
+  id,
+  type: "ai",
+  content: [{ type: "text", text }],
+});
+
+const renderAdk = async (
+  load?: (
+    threadId: string,
+    options?: { signal?: AbortSignal | undefined },
+  ) => Promise<AdkThreadSnapshot>,
+) => {
+  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
+  // to be in the same tree as the hook for a thread switch to settle
+  const Inner: FC = () => {
+    const runtime = useAdkRuntime({
+      stream: streamMock as never,
+      ...(load ? { load } : {}),
+      sessionAdapter: makeThreadListAdapter(),
+    });
+    capture.runtime = runtime;
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        {null}
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  let unmount!: () => void;
+  await act(async () => {
+    ({ unmount } = render(<Inner />));
+  });
+  await waitFor(() => expect(capture.runtime).not.toBeNull());
+
+  await act(async () => {
+    await capture.runtime!.threads.switchToThread("adk-1");
+  });
+
+  return { capture, streamMock, unmount };
+};
+
+describe("useAdkRuntime refetch", () => {
+  it("declares the refetch capability only when a load is supplied", async () => {
+    const withLoad = await renderAdk(async () => ({ messages: [] }));
+    expect(
+      withLoad.capture.runtime!.thread.getState().capabilities.refetchThread,
+    ).toBe(true);
+
+    const withoutLoad = await renderAdk();
+    expect(
+      withoutLoad.capture.runtime!.thread.getState().capabilities.refetchThread,
+    ).toBe(false);
+  });
+
+  it("refetches in place, keeping the composer draft and the runtime", async () => {
+    let call = 0;
+    const load = vi.fn(async () => {
+      call++;
+      return {
+        messages: [aiMessage(`m-${call}`, `load ${call}`)],
+      } satisfies AdkThreadSnapshot;
+    });
+
+    const { capture } = await renderAdk(load);
+    const wrapper = wrapperFactory(capture.runtime!);
+    const { result: auiResult } = renderHook(() => useAui(), { wrapper });
+
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+    await waitFor(() =>
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("load 1"),
+    );
+
+    await act(async () => {
+      auiResult.current.composer.setText("draft that must survive");
+    });
+
+    await act(async () => {
+      await capture.runtime!.threads.reloadMainThread();
+    });
+
+    expect(load).toHaveBeenCalledTimes(2);
+    await waitFor(() =>
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("load 2"),
+    );
+    expect(auiResult.current.composer.getState().text).toBe(
+      "draft that must survive",
+    );
+  });
+
+  it("swaps the per-turn state over with the messages", async () => {
+    let call = 0;
+    const load = vi.fn(async () => {
+      call++;
+      return call === 1
+        ? {
+            messages: [aiMessage("m-1", "first")],
+            toolConfirmations: [
+              { toolCallId: "tc-1", toolName: "search", hint: "before" },
+            ],
+          }
+        : {
+            messages: [aiMessage("m-2", "second")],
+            toolConfirmations: [
+              { toolCallId: "tc-2", toolName: "search", hint: "after" },
+            ],
+          };
+    });
+
+    const { capture } = await renderAdk(load as never);
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+
+    await act(async () => {
+      await capture.runtime!.threads.reloadMainThread();
+    });
+
+    await waitFor(() =>
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("second"),
+    );
+    // the confirmation from the fresh snapshot replaced the earlier one rather
+    // than the thread passing through a cleared state
+    const extras = capture.runtime!.thread.getState().extras as {
+      toolConfirmations: { hint: string }[];
+    };
+    expect(extras.toolConfirmations).toHaveLength(1);
+    expect(extras.toolConfirmations[0]!.hint).toBe("after");
+  });
+
+  it("leaves the thread untouched when the refetch fails, and rejects", async () => {
+    let call = 0;
+    const load = vi.fn(async () => {
+      call++;
+      if (call === 1) return { messages: [aiMessage("m-1", "first")] };
+      throw new Error("refetch failed");
+    });
+
+    const { capture } = await renderAdk(load as never);
+    await waitFor(() =>
+      expect(
+        JSON.stringify(capture.runtime!.thread.getState().messages),
+      ).toContain("first"),
+    );
+
+    await act(async () => {
+      await expect(capture.runtime!.threads.reloadMainThread()).rejects.toThrow(
+        "refetch failed",
+      );
+    });
+
+    expect(
+      JSON.stringify(capture.runtime!.thread.getState().messages),
+    ).toContain("first");
+  });
+
+  it("defers to an initial load still in flight rather than taking it over", async () => {
+    const pending = deferred<AdkThreadSnapshot>();
+    let call = 0;
+    const load = vi.fn(async () => {
+      call++;
+      return call === 1
+        ? pending.promise
+        : { messages: [aiMessage("m-2", "second")] };
+    });
+
+    const streamMock = vi.fn(async function* () {});
+    const capture: { runtime: AssistantRuntime | null } = { runtime: null };
+    const Inner: FC = () => {
+      const runtime = useAdkRuntime({
+        stream: streamMock as never,
+        load: load as never,
+        sessionAdapter: makeThreadListAdapter(),
+      });
+      capture.runtime = runtime;
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          {null}
+        </AssistantRuntimeProvider>
+      );
+    };
+    await act(async () => {
+      render(<Inner />);
+    });
+    await act(async () => {
+      await capture.runtime!.threads.switchToThread("adk-1");
+    });
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+
+    let settled = false;
+    act(() => {
+      capture.runtime!.threads.reloadMainThread().then(() => {
+        settled = true;
+      });
+    });
+
+    // the refetch waits on the initial load instead of starting a second one
+    expect(load).toHaveBeenCalledTimes(1);
+    expect(settled).toBe(false);
+
+    await act(async () => {
+      pending.resolve({ messages: [aiMessage("m-1", "first")] });
+    });
+    await waitFor(() => expect(settled).toBe(true));
+    expect(
+      JSON.stringify(capture.runtime!.thread.getState().messages),
+    ).toContain("first");
+  });
+
+  it("does not report an aborted load as a failure", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const pending = deferred<AdkThreadSnapshot>();
+    let call = 0;
+    const load = vi.fn(
+      async (_id: string, options?: { signal?: AbortSignal }) => {
+        call++;
+        if (call === 1) {
+          options?.signal?.addEventListener("abort", () =>
+            pending.reject(
+              Object.assign(new Error("aborted"), { name: "AbortError" }),
+            ),
+          );
+          return pending.promise;
+        }
+        return { messages: [] };
+      },
+    );
+
+    const { unmount } = await renderAdk(load as never);
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+
+    unmount();
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(warn).not.toHaveBeenCalled();
+    warn.mockRestore();
+  });
+
+  it("aborts a refetch still in flight when the thread unmounts", async () => {
+    const pending = deferred<AdkThreadSnapshot>();
+    let call = 0;
+    const signals: (AbortSignal | undefined)[] = [];
+    const load = vi.fn(
+      async (_id: string, options?: { signal?: AbortSignal }) => {
+        call++;
+        signals.push(options?.signal);
+        if (call === 1) return { messages: [] };
+        return pending.promise;
+      },
+    );
+
+    const { capture, unmount } = await renderAdk(load as never);
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(1));
+
+    act(() => {
+      capture.runtime!.threads.reloadMainThread().catch(() => {});
+    });
+    await waitFor(() => expect(load).toHaveBeenCalledTimes(2));
+    expect(signals[1]?.aborted).toBe(false);
+
+    unmount();
+
+    expect(signals[1]?.aborted).toBe(true);
+    pending.resolve({ messages: [] });
+  });
+});