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

package/src/AdkClient.ts

@@ -243,10 +243,20 @@ async function* parseSSEResponse(response: Response): AsyncGenerator<AdkEvent> {
   const decoder = new TextDecoder();
   const sseDecoder = new SSEEventDecoder({ trailing: "dispatch" });
 
+  let shouldCancel = true;
   try {
     while (true) {
-      const { done, value } = await reader.read();
+      let result: ReadableStreamReadResult<Uint8Array>;
+      try {
+        result = await reader.read();
+      } catch (error) {
+        shouldCancel = false;
+        throw error;
+      }
+
+      const { done, value } = result;
       if (done) {
+        shouldCancel = false;
         for (const event of sseDecoder.push(decoder.decode())) {
           yield JSON.parse(event.data) as AdkEvent;
         }
@@ -263,6 +273,10 @@ async function* parseSSEResponse(response: Response): AsyncGenerator<AdkEvent> {
     const trailing = sseDecoder.flush();
     if (trailing !== null) yield JSON.parse(trailing.data) as AdkEvent;
   } finally {
-    reader.releaseLock();
+    try {
+      if (shouldCancel) await reader.cancel().catch(() => undefined);
+    } finally {
+      reader.releaseLock();
+    }
   }
 }

package/src/AdkEventAccumulator.test.ts

@@ -239,6 +239,133 @@ describe("AdkEventAccumulator - function responses", () => {
       content: JSON.stringify({ results: [] }),
     });
   });
+
+  // A session load replays the stored events through a fresh accumulator.
+  it("gives a tool message the same id on every replay of an event", () => {
+    const event = makeEvent({
+      id: "evt-tool",
+      author: "agent",
+      content: {
+        role: "model",
+        parts: [
+          {
+            functionResponse: {
+              name: "search",
+              id: "tc-1",
+              response: { results: [] },
+            },
+          },
+        ],
+      },
+    });
+
+    const first = new AdkEventAccumulator().processEvent(event);
+    const second = new AdkEventAccumulator().processEvent(event);
+
+    expect(first[0]!.id).toBe(second[0]!.id);
+  });
+
+  it("keeps two tool messages from one event distinct", () => {
+    const msgs = new AdkEventAccumulator().processEvent(
+      makeEvent({
+        id: "evt-tool",
+        author: "agent",
+        content: {
+          role: "model",
+          parts: [
+            {
+              functionResponse: {
+                name: "search",
+                id: "tc-1",
+                response: { a: 1 },
+              },
+            },
+            {
+              functionResponse: {
+                name: "lookup",
+                id: "tc-2",
+                response: { b: 2 },
+              },
+            },
+          ],
+        },
+      }),
+    );
+
+    const toolMessages = msgs.filter((m) => m.type === "tool");
+    expect(toolMessages).toHaveLength(2);
+    expect(toolMessages[0]!.id).not.toBe(toolMessages[1]!.id);
+  });
+
+  it("gives an assistant message the same id on every replay", () => {
+    const events = [
+      makeTextEvent("Hel", true),
+      makeTextEvent("lo", true),
+      makeTextEvent("Hello"),
+    ];
+    const replay = () => {
+      const acc = new AdkEventAccumulator();
+      let msgs: AdkMessage[] = [];
+      for (const event of events) msgs = acc.processEvent(event);
+      return msgs;
+    };
+
+    expect(replay().map((m) => m.id)).toEqual(replay().map((m) => m.id));
+  });
+
+  it("keeps two assistant messages opened by one event distinct and stable", () => {
+    const event = makeEvent({
+      id: "evt-mixed",
+      author: "agent",
+      content: {
+        role: "model",
+        parts: [
+          { text: "before" },
+          {
+            functionResponse: {
+              name: "search",
+              id: "tc-1",
+              response: { ok: true },
+            },
+          },
+          { text: "after" },
+        ],
+      },
+    });
+
+    const first = new AdkEventAccumulator().processEvent(event);
+    const second = new AdkEventAccumulator().processEvent(event);
+
+    const aiIds = first.filter((m) => m.type === "ai").map((m) => m.id);
+    expect(aiIds.length).toBeGreaterThan(1);
+    expect(new Set(aiIds).size).toBe(aiIds.length);
+    expect(first.map((m) => m.id)).toEqual(second.map((m) => m.id));
+  });
+
+  it("keeps two tool messages distinct when the payload omits response ids", () => {
+    const event = makeEvent({
+      id: "evt-tool",
+      author: "agent",
+      content: {
+        role: "model",
+        parts: [
+          { functionResponse: { name: "search", response: { a: 1 } } },
+          { functionResponse: { name: "lookup", response: { b: 2 } } },
+        ],
+      },
+    });
+
+    const first = new AdkEventAccumulator()
+      .processEvent(event)
+      .filter((m) => m.type === "tool");
+    const second = new AdkEventAccumulator()
+      .processEvent(event)
+      .filter((m) => m.type === "tool");
+
+    expect(first).toHaveLength(2);
+    expect(first[0]!.id).not.toBe(first[1]!.id);
+    expect(first.map((m) => m.id)).toEqual(second.map((m) => m.id));
+  });
 });
 
 describe("AdkEventAccumulator - code execution", () => {

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/adkExtras.ts

@@ -1,4 +1,4 @@
-import { createRuntimeExtras } from "@assistant-ui/core/internal";
+import { createRuntimeExtras } from "@assistant-ui/core/react";
 import type { AdkRuntimeExtras } from "./types";
 
 export const adkExtras = createRuntimeExtras<AdkRuntimeExtras>("useAdkRuntime");

package/src/convertAdkMessages.test.ts

@@ -25,6 +25,106 @@ describe("convertAdkMessage - human messages", () => {
       content: [{ type: "text", text: "Hello" }],
     });
   });
+
+  it("restores a file_url part as a file part stamped with sourceType url", () => {
+    const msg: AdkMessage = {
+      id: "m1",
+      type: "human",
+      content: [
+        {
+          type: "file_url",
+          url: "gs://bucket/report.pdf",
+          mimeType: "application/pdf",
+        },
+      ],
+    };
+    const result = convertAdkMessage(msg, {});
+    expect(result).toMatchObject({
+      role: "user",
+      content: [
+        {
+          type: "file",
+          data: "gs://bucket/report.pdf",
+          mimeType: "application/pdf",
+          sourceType: "url",
+        },
+      ],
+    });
+  });
+
+  it("falls back to application/octet-stream for file_url parts without mimeType", () => {
+    const msg: AdkMessage = {
+      id: "m1",
+      type: "human",
+      content: [{ type: "file_url", url: "gs://bucket/blob" }],
+    };
+    const result = convertAdkMessage(msg, {});
+    expect(result).toMatchObject({
+      role: "user",
+      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",
+      type: "human",
+      content: [
+        {
+          type: "file",
+          mimeType: "audio/wav",
+          data: "QUJD",
+          filename: "memo.wav",
+        },
+      ],
+    };
+    const result = convertAdkMessage(msg, {});
+    expect(result).toMatchObject({
+      role: "user",
+      content: [
+        {
+          type: "file",
+          mimeType: "audio/wav",
+          data: "QUJD",
+          filename: "memo.wav",
+        },
+      ],
+    });
+  });
+
+  it("keeps file parts with other audio mime types as file parts", () => {
+    const msg: AdkMessage = {
+      id: "m1",
+      type: "human",
+      content: [{ type: "file", mimeType: "audio/ogg", data: "QUJD" }],
+    };
+    const result = convertAdkMessage(msg, {});
+    expect(result).toMatchObject({
+      role: "user",
+      content: [{ type: "file", mimeType: "audio/ogg", data: "QUJD" }],
+    });
+  });
 });
 
 describe("convertAdkMessage - ai messages", () => {
@@ -78,6 +178,19 @@ describe("convertAdkMessage - ai messages", () => {
     });
   });
 
+  it("keeps audio/mp3 file parts as file parts on assistant messages", () => {
+    const msg: AdkMessage = {
+      id: "m1",
+      type: "ai",
+      content: [{ type: "file", mimeType: "audio/mp3", data: "QUJD" }],
+    };
+    const result = convertAdkMessage(msg, {});
+    expect(result).toMatchObject({
+      role: "assistant",
+      content: [{ type: "file", mimeType: "audio/mp3", data: "QUJD" }],
+    });
+  });
+
   it("converts file content parts to file message parts", () => {
     const msg: AdkMessage = {
       id: "m1",

package/src/convertAdkMessages.ts

@@ -8,10 +8,19 @@ type ContentPart =
   | { type: "text"; text: string }
   | { type: "reasoning"; text: string }
   | { type: "image"; image: string }
-  | { type: "file"; data: string; mimeType: string; filename?: string }
+  | {
+      type: "file";
+      data: string;
+      mimeType: string;
+      filename?: string;
+      sourceType?: "url";
+    }
   | { type: "data"; name: string; data: unknown };
 
-const contentToParts = (content: AdkMessage["content"]): ContentPart[] => {
+const contentToParts = (
+  content: AdkMessage["content"],
+  role: "user" | "assistant",
+): ContentPart[] => {
   if (typeof content === "string")
     return [{ type: "text" as const, text: content }];
 
@@ -37,6 +46,14 @@ const contentToParts = (content: AdkMessage["content"]): ContentPart[] => {
             ...(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",
@@ -72,7 +89,7 @@ export const convertAdkMessage: useExternalMessageConverter.Callback<
       return {
         role: "user",
         id: message.id,
-        content: contentToParts(message.content),
+        content: contentToParts(message.content, "user"),
       };
 
     case "ai": {
@@ -88,7 +105,10 @@ export const convertAdkMessage: useExternalMessageConverter.Callback<
       return {
         role: "assistant",
         id: message.id,
-        content: [...contentToParts(message.content), ...toolCallParts],
+        content: [
+          ...contentToParts(message.content, "assistant"),
+          ...toolCallParts,
+        ],
         ...(message.status && { status: message.status }),
         ...(message.author && {
           metadata: {

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