@assistant-ui/react-google-adk 0.0.31 → 0.0.32

package/src/hooks.ts

@@ -2,6 +2,7 @@ import { generateId } from "@assistant-ui/core";
 import { useAui } from "@assistant-ui/store";
 import { useShallowSelector } from "@assistant-ui/store/internal";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
+import { toAdkAuthReply } from "./adkAuthRequest";
 import { adkExtras } from "./adkExtras";
 import { toAdkConfirmationReply } from "./adkToolApproval";
 import type {
@@ -74,23 +75,21 @@ export const useAdkConfirmTool = () => {
       .send([toAdkConfirmationReply(toolCallId, confirmed, payload)], {});
 };
 
-/** Returns a function to submit auth credentials for a pending auth request. */
+/**
+ * Returns a function to submit auth credentials for a pending auth request.
+ * The reply carries the request's auth config with the credential as its
+ * `exchangedAuthCredential`, which is the shape ADK resumes the tool on, so the
+ * id must name a request `useAdkAuthRequests` currently lists.
+ */
 export const useAdkSubmitAuth = () => {
   const aui = useAui();
-  return (toolCallId: string, credential: AdkAuthCredential) =>
-    adkExtras.get(aui).send(
-      [
-        {
-          id: generateId(),
-          type: "tool",
-          tool_call_id: toolCallId,
-          name: "adk_request_credential",
-          content: JSON.stringify(credential),
-          status: "success",
-        },
-      ],
+  return (toolCallId: string, credential: AdkAuthCredential) => {
+    const extras = adkExtras.get(aui);
+    return extras.send(
+      [toAdkAuthReply(toolCallId, credential, extras.authRequests)],
       {},
     );
+  };
 };
 
 /** Returns a function to submit the user's answer for a pending `adk_request_input` HITL interrupt. */

package/src/raceWithAbortSignal.test.ts

@@ -0,0 +1,73 @@
+import { describe, expect, it, vi } from "vitest";
+import { raceWithAbortSignal } from "./raceWithAbortSignal";
+
+describe("raceWithAbortSignal", () => {
+  it("invokes the operation synchronously without a signal", async () => {
+    const order: string[] = [];
+
+    const result = raceWithAbortSignal(undefined, () => {
+      order.push("operation");
+      return "done";
+    });
+    order.push("after");
+
+    expect(order).toEqual(["operation", "after"]);
+    await expect(result).resolves.toBe("done");
+  });
+
+  it("converts a synchronous operation error to a rejection", async () => {
+    const error = new Error("failed");
+
+    const result = raceWithAbortSignal(undefined, () => {
+      throw error;
+    });
+
+    await expect(result).rejects.toBe(error);
+  });
+
+  it("rejects a pending operation with the abort reason", async () => {
+    const controller = new AbortController();
+    const reason = new Error("cancelled");
+    let resolveOperation!: (value: string) => void;
+    const operation = new Promise<string>((resolve) => {
+      resolveOperation = resolve;
+    });
+
+    const result = raceWithAbortSignal(controller.signal, () => operation);
+    controller.abort(reason);
+
+    await expect(result).rejects.toBe(reason);
+    resolveOperation("late result");
+  });
+
+  it("rejects before invoking an operation for an already aborted signal", async () => {
+    const controller = new AbortController();
+    const reason = new Error("already cancelled");
+    const operation = vi.fn(() => "done");
+    controller.abort(reason);
+
+    const result = raceWithAbortSignal(controller.signal, operation);
+
+    await expect(result).rejects.toBe(reason);
+    expect(operation).not.toHaveBeenCalled();
+  });
+
+  it("removes the abort listener after the operation settles", async () => {
+    const controller = new AbortController();
+    const removeEventListener = vi.spyOn(
+      controller.signal,
+      "removeEventListener",
+    );
+    const result = raceWithAbortSignal(controller.signal, () => "done");
+
+    await expect(result).resolves.toBe("done");
+    expect(removeEventListener).toHaveBeenCalledOnce();
+    expect(removeEventListener).toHaveBeenCalledWith(
+      "abort",
+      expect.any(Function),
+    );
+
+    controller.abort(new Error("late abort"));
+    await expect(result).resolves.toBe("done");
+  });
+});

package/src/raceWithAbortSignal.ts

@@ -0,0 +1,48 @@
+const getAbortReason = (signal: AbortSignal): unknown => {
+  if (signal.reason !== undefined) return signal.reason;
+  const error = new Error("The operation was aborted");
+  error.name = "AbortError";
+  return error;
+};
+
+export const raceWithAbortSignal = <T>(
+  signal: AbortSignal | undefined,
+  operation: () => T | PromiseLike<T>,
+): Promise<T> => {
+  if (!signal) {
+    try {
+      return Promise.resolve(operation());
+    } catch (error) {
+      return Promise.reject(error);
+    }
+  }
+  if (signal.aborted) return Promise.reject(getAbortReason(signal));
+
+  return new Promise<T>((resolve, reject) => {
+    let settled = false;
+    const cleanup = () => signal.removeEventListener("abort", handleAbort);
+    const resolveOnce = (value: T) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      resolve(value);
+    };
+    const rejectOnce = (error: unknown) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      reject(error);
+    };
+    const handleAbort = () => rejectOnce(getAbortReason(signal));
+
+    signal.addEventListener("abort", handleAbort, { once: true });
+    let result: T | PromiseLike<T>;
+    try {
+      result = operation();
+    } catch (error) {
+      rejectOnce(error);
+      return;
+    }
+    Promise.resolve(result).then(resolveOnce, rejectOnce);
+  });
+};

package/src/server/adkEventStream.test.ts

@@ -20,6 +20,153 @@ async function readSSE(response: Response): Promise<string> {
 }
 
 describe("adkEventStream", () => {
+  it("does not advance the generator before the response is read", async () => {
+    let produced = 0;
+    async function* events() {
+      for (let index = 0; index < 1000; index++) {
+        produced++;
+        yield { id: String(index) };
+      }
+    }
+    const response = adkEventStream(events());
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    try {
+      expect(produced).toBe(0);
+    } finally {
+      await response.body!.cancel();
+    }
+  });
+
+  it("buffers at most one event ahead of a slow reader", async () => {
+    let produced = 0;
+    async function* events() {
+      for (let index = 0; index < 1000; index++) {
+        produced++;
+        yield { id: String(index) };
+      }
+    }
+    const reader = adkEventStream(events()).body!.getReader();
+    try {
+      expect(new TextDecoder().decode((await reader.read()).value)).toBe(
+        ":ok\n\n",
+      );
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(produced).toBe(1);
+      expect(new TextDecoder().decode((await reader.read()).value)).toContain(
+        '"id":"0"',
+      );
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(produced).toBe(2);
+    } finally {
+      await reader.cancel();
+      reader.releaseLock();
+    }
+  });
+
+  it("finalizes the generator when the reader cancels between events", async () => {
+    const finalized = vi.fn();
+    const onError = vi.fn();
+    async function* events() {
+      try {
+        while (true) yield { id: "event" };
+      } finally {
+        finalized();
+      }
+    }
+    const reader = adkEventStream(events(), { onError }).body!.getReader();
+    await reader.read();
+    await reader.read();
+    await reader.cancel();
+    expect(finalized).toHaveBeenCalledTimes(1);
+    expect(onError).not.toHaveBeenCalled();
+    expect(await reader.read()).toEqual({ done: true, value: undefined });
+    reader.releaseLock();
+  });
+
+  it("delivers every event in order when the response is fully consumed", async () => {
+    async function* events() {
+      for (let index = 0; index < 1000; index++) yield { id: String(index) };
+    }
+    const text = await readSSE(adkEventStream(events()));
+    const ids = text
+      .split("\n\n")
+      .filter((line) => line.startsWith("data: "))
+      .map((line) => JSON.parse(line.slice(6)).id);
+    expect(ids).toEqual(
+      Array.from({ length: 1000 }, (_, index) => String(index)),
+    );
+  });
+
+  it("allows onError to cancel the reader without writing to the closed stream", async () => {
+    let reader: ReadableStreamDefaultReader<Uint8Array>;
+    let cancelled: Promise<void> | undefined;
+    async function* events() {
+      throw new Error("stream failed");
+    }
+    const response = adkEventStream(events(), {
+      onError: () => {
+        cancelled = reader.cancel();
+      },
+    });
+    reader = response.body!.getReader();
+    await reader.read();
+    expect(await reader.read()).toEqual({ done: true, value: undefined });
+    await cancelled;
+    reader.releaseLock();
+  });
+
+  it.each(["yield", "finish", "reject"])(
+    "does not write after cancellation while next() later %s",
+    async (outcome) => {
+      let release!: () => void;
+      const waiting = new Promise<void>((resolve) => {
+        release = resolve;
+      });
+      let started!: () => void;
+      const entered = new Promise<void>((resolve) => {
+        started = resolve;
+      });
+      const finalized = vi.fn();
+      const onError = vi.fn();
+      async function* events() {
+        try {
+          started();
+          await waiting;
+          if (outcome === "reject") throw new Error("late error");
+          if (outcome === "yield") yield { id: "late" };
+        } finally {
+          finalized();
+        }
+      }
+      const reader = adkEventStream(events(), { onError }).body!.getReader();
+      await reader.read();
+      await entered;
+      const cancelled = reader.cancel();
+      release();
+      await cancelled;
+      expect(finalized).toHaveBeenCalledTimes(1);
+      expect(onError).not.toHaveBeenCalled();
+      expect(await reader.read()).toEqual({ done: true, value: undefined });
+      reader.releaseLock();
+    },
+  );
+
+  it("finalizes the source and emits an error frame if event serialization fails", async () => {
+    const finalized = vi.fn();
+    const onError = vi.fn();
+    async function* events() {
+      try {
+        yield { id: "invalid", customMetadata: { unsupported: 1n } };
+      } finally {
+        finalized();
+      }
+    }
+    const text = await readSSE(adkEventStream(events(), { onError }));
+    expect(text).toContain('"errorCode":"STREAM_ERROR"');
+    expect(finalized).toHaveBeenCalledTimes(1);
+    expect(onError).toHaveBeenCalledTimes(1);
+  });
+
   it("returns a Response with text/event-stream content-type", () => {
     const response = adkEventStream(yieldEvents([]));
     expect(response.headers.get("Content-Type")).toBe("text/event-stream");

package/src/server/adkEventStream.ts

@@ -133,31 +133,38 @@ export const adkEventStream = (
   const encoder = new TextEncoder();
   let cancelled = false;
   const stream = new ReadableStream({
-    async start(controller) {
+    start(controller) {
       // Initial SSE comment to keep connection alive through proxies
       controller.enqueue(encoder.encode(":ok\n\n"));
+    },
+    async pull(controller) {
       try {
-        for await (const event of events) {
-          if (cancelled) break;
-          const wireEvent = convertSdkEvent(event);
-          controller.enqueue(
-            encoder.encode(`data: ${JSON.stringify(wireEvent)}\n\n`),
-          );
+        const { done, value } = await events.next();
+        if (cancelled) return;
+        if (done) {
+          controller.close();
+          return;
         }
+        const wireEvent = convertSdkEvent(value);
+        controller.enqueue(
+          encoder.encode(`data: ${JSON.stringify(wireEvent)}\n\n`),
+        );
       } catch (e) {
-        if (!cancelled) {
-          notifyError(options?.onError, e);
-          const errorEvent: AdkEvent = {
-            id: "",
-            errorCode: "STREAM_ERROR",
-            errorMessage:
-              e instanceof Error ? e.message : "Unknown stream error",
-          };
-          controller.enqueue(
-            encoder.encode(`data: ${JSON.stringify(errorEvent)}\n\n`),
-          );
-        }
-      } finally {
+        if (cancelled) return;
+        try {
+          await events.return?.(undefined as any);
+        } catch {}
+        if (cancelled) return;
+        notifyError(options?.onError, e);
+        if (cancelled) return;
+        const errorEvent: AdkEvent = {
+          id: "",
+          errorCode: "STREAM_ERROR",
+          errorMessage: e instanceof Error ? e.message : "Unknown stream error",
+        };
+        controller.enqueue(
+          encoder.encode(`data: ${JSON.stringify(errorEvent)}\n\n`),
+        );
         controller.close();
       }
     },