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

package/src/AdkEventAccumulator.ts

@@ -14,6 +14,23 @@ import type { ReadonlyJSONObject } from "assistant-stream/utils";
 
 type InProgressMessage = AdkMessage & { type: "ai" };
 
+/**
+ * A session load replays the stored events through a fresh accumulator, so a
+ * message needs an id derived from the event that carries it rather than one
+ * minted per replay. An event with no id of its own has never been through the
+ * session and has nothing stable to derive from, so it keeps a generated one.
+ *
+ * A human message keeps the bare event id it has always had. The other kinds
+ * take a suffixed namespace, since one event can carry several of them: a tool
+ * message by the index of its part, an assistant message by how many this
+ * event has already opened.
+ */
+const toolMessageId = (event: AdkEvent, partIndex: number): string =>
+  event.id ? `${event.id}:${partIndex}` : uuidv4();
+
+const aiMessageId = (event: AdkEvent, ordinal: number): string =>
+  event.id ? `${event.id}:ai${ordinal === 0 ? "" : ordinal}` : uuidv4();
+
 const ADK_REQUEST_CONFIRMATION = "adk_request_confirmation";
 const ADK_REQUEST_CREDENTIAL = "adk_request_credential";
 
@@ -187,6 +204,9 @@ export class AdkEventAccumulator {
   private authRequests: AdkAuthRequest[] = [];
   private escalated = false;
   private messageMetadataMap = new Map<string, AdkMessageMetadata>();
+  // How many assistant messages each event has opened, so a replay of that
+  // event opens them with the same ids.
+  private aiMessageOrdinals = new Map<string, number>();
   constructor(initialMessages?: AdkMessage[]) {
     if (initialMessages) {
       for (const msg of initialMessages) {
@@ -344,8 +364,8 @@ export class AdkEventAccumulator {
       }
     }
 
-    for (const part of parts) {
-      this.processPart(part, event);
+    for (const [index, part] of parts.entries()) {
+      this.processPart(part, event, index);
     }
 
     // Track per-message metadata (grounding, citation, usage)
@@ -379,7 +399,11 @@ export class AdkEventAccumulator {
     return this.getMessages();
   }
 
-  private processPart(part: AdkEventPart, event: AdkEvent): void {
+  private processPart(
+    part: AdkEventPart,
+    event: AdkEvent,
+    partIndex: number,
+  ): void {
     // Detect special ADK function calls
     if (part.functionCall && !event.partial) {
       const name = part.functionCall.name;
@@ -475,7 +499,7 @@ export class AdkEventAccumulator {
     if (part.functionResponse) {
       this.finalizeCurrentMessage();
       const toolMsg: AdkMessage = {
-        id: uuidv4(),
+        id: toolMessageId(event, partIndex),
         type: "tool",
         tool_call_id: part.functionResponse.id ?? "",
         name: part.functionResponse.name,
@@ -559,7 +583,9 @@ export class AdkEventAccumulator {
       }
     }
 
-    const id = uuidv4();
+    const ordinal = this.aiMessageOrdinals.get(event.id ?? "") ?? 0;
+    this.aiMessageOrdinals.set(event.id ?? "", ordinal + 1);
+    const id = aiMessageId(event, ordinal);
     const msg: InProgressMessage = {
       id,
       type: "ai",

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/server/adkEventStream.test.ts

@@ -73,4 +73,42 @@ describe("adkEventStream", () => {
     await readSSE(response);
     expect(onError).toHaveBeenCalledTimes(1);
   });
+
+  it.each([
+    [
+      "throws",
+      () => {
+        throw new Error("callback failed");
+      },
+    ],
+    [
+      "rejects",
+      async () => {
+        throw new Error("callback failed");
+      },
+    ],
+  ])(
+    "still emits the stream error when onError %s",
+    async (_behavior, onError) => {
+      const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
+      async function* throwingGen() {
+        throw new Error("stream failed");
+      }
+
+      try {
+        const response = adkEventStream(throwingGen() as any, { onError });
+        const text = await readSSE(response);
+        await vi.waitFor(() => expect(errorSpy).toHaveBeenCalledTimes(1));
+
+        const lines = text.split("\n\n").filter((l) => l.startsWith("data: "));
+        expect(lines).toHaveLength(1);
+        expect(JSON.parse(lines[0]!.slice(6))).toMatchObject({
+          errorCode: "STREAM_ERROR",
+          errorMessage: "stream failed",
+        });
+      } finally {
+        errorSpy.mockRestore();
+      }
+    },
+  );
 });

package/src/server/adkEventStream.ts

@@ -101,6 +101,21 @@ export type AdkEventStreamOptions = {
   onError?: (error: unknown) => void;
 };
 
+const reportCallbackError = (error: unknown) => {
+  console.error("[react-google-adk] onError callback threw an error", error);
+};
+
+const notifyError = (
+  callback: AdkEventStreamOptions["onError"],
+  error: unknown,
+) => {
+  try {
+    void Promise.resolve(callback?.(error)).catch(reportCallbackError);
+  } catch (callbackError) {
+    reportCallbackError(callbackError);
+  }
+};
+
 /**
  * Converts an AsyncGenerator of ADK SDK Events into an SSE Response.
  *
@@ -131,7 +146,7 @@ export const adkEventStream = (
         }
       } catch (e) {
         if (!cancelled) {
-          options?.onError?.(e);
+          notifyError(options?.onError, e);
           const errorEvent: AdkEvent = {
             id: "",
             errorCode: "STREAM_ERROR",

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,
   };
 };