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

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react-google-adk",
-  "version": "0.0.34",
+  "version": "0.0.36",
   "description": "Google ADK adapter for assistant-ui",
   "keywords": [
     "google-adk",
@@ -34,10 +34,10 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@assistant-ui/core": "^0.3.23",
-    "@assistant-ui/store": "^0.3.17",
+    "@assistant-ui/core": "^0.3.25",
+    "@assistant-ui/store": "^0.3.18",
     "assistant-cloud": "*",
-    "assistant-stream": "^0.3.47"
+    "assistant-stream": "^0.3.49"
   },
   "peerDependencies": {
     "@google/adk": ">=0.5.0",
@@ -58,7 +58,7 @@
     "@testing-library/dom": "^10.4.2",
     "@testing-library/react": "^16.3.3",
     "@types/react": "^19.3.0",
-    "jsdom": "^30.0.1",
+    "jsdom": "^30.1.2",
     "react": "^19.3.0",
     "vitest": "^5.0.3"
   },
@@ -78,6 +78,7 @@
   "scripts": {
     "build": "aui-build",
     "test": "vitest run",
+    "test:coverage": "vitest run --coverage --coverage.reporter=text --coverage.reporter=html --coverage.reporter=json-summary",
     "typecheck": "tsc --noEmit",
     "test:peer-react18": "vitest run --config vitest.peer-react18.config.ts",
     "test:types:peer-react18": "aui-react18-types",

package/src/AdkClient.test.ts

@@ -189,7 +189,8 @@ describe("createAdkStream - proxy mode", () => {
     const messages: AdkMessage[] = [
       { id: "m1", type: "human", content: "Hello" },
     ];
-    const gen = await stream(messages, makeConfig());
+    const config = makeConfig();
+    const gen = await stream(messages, config);
     // drain
     for await (const _ of gen) {
       /* noop */
@@ -200,7 +201,27 @@ describe("createAdkStream - proxy mode", () => {
     expect(url).toBe("/api/adk");
     expect(init?.method).toBe("POST");
     const body = JSON.parse(init?.body as string);
-    expect(body).toMatchObject({ message: "Hello" });
+    expect(body).toMatchObject({ message: "Hello", sessionId: "session-1" });
+    expect(config.initialize).toHaveBeenCalledOnce();
+  });
+
+  it("falls back to the remote thread ID for proxy sessions", async () => {
+    mockFetch.mockResolvedValueOnce(sseResponse(sseBody("")));
+    const initialize = vi
+      .fn()
+      .mockResolvedValue({ remoteId: "remote-1", externalId: undefined });
+
+    const stream = createAdkStream({ api: "/api/adk" });
+    const gen = await stream(
+      [{ id: "m1", type: "human", content: "Hello" }],
+      makeConfig({ initialize }),
+    );
+    for await (const _ of gen) {
+      /* noop */
+    }
+
+    const body = JSON.parse(mockFetch.mock.calls[0]![1]?.body as string);
+    expect(body.sessionId).toBe("remote-1");
   });
 
   it("sends runConfig and checkpointId in proxy body", async () => {
@@ -603,6 +624,61 @@ describe("createAdkStream - SSE parsing", () => {
     expect(collected[1]!.id).toBe("e2");
   });
 
+  it("forwards configurable SSE line and event limits", async () => {
+    const event: AdkEvent = {
+      id: "e1",
+      content: { parts: [{ text: "x".repeat(64) }] },
+    };
+    const text = `data: ${JSON.stringify(event)}\n\n`;
+    const consume = async (options: {
+      maxStreamLineLength: number;
+      maxStreamEventLength: number;
+    }) => {
+      mockFetch.mockResolvedValueOnce(sseResponse(sseBody(text)));
+      const stream = createAdkStream({ api: "/api/adk", ...options });
+      const gen = await stream(
+        [{ id: "m1", type: "human", content: "Hi" }],
+        makeConfig(),
+      );
+      const events: AdkEvent[] = [];
+      for await (const parsedEvent of gen) events.push(parsedEvent);
+      return events;
+    };
+
+    await expect(
+      consume({ maxStreamLineLength: 32, maxStreamEventLength: 1_024 }),
+    ).rejects.toThrow("SSE line exceeds maxLineLength");
+    await expect(
+      consume({ maxStreamLineLength: 1_024, maxStreamEventLength: 32 }),
+    ).rejects.toThrow("SSE event exceeds maxEventLength");
+    await expect(
+      consume({ maxStreamLineLength: 1_024, maxStreamEventLength: 1_024 }),
+    ).resolves.toEqual([event]);
+  });
+
+  it("releases the response reader when decoder limits are invalid", async () => {
+    const cancel = vi.fn();
+    const body = new ReadableStream<Uint8Array>({ cancel });
+    mockFetch.mockResolvedValueOnce(sseResponse(body));
+    const stream = createAdkStream({
+      api: "/api/adk",
+      maxStreamLineLength: 0,
+    });
+    const consume = async () => {
+      const gen = await stream(
+        [{ id: "m1", type: "human", content: "Hi" }],
+        makeConfig(),
+      );
+      for await (const _event of gen) void _event;
+    };
+
+    await expect(consume()).rejects.toThrow(
+      "maxLineLength must be a positive safe integer",
+    );
+    expect(cancel).toHaveBeenCalledOnce();
+    expect(body.locked).toBe(false);
+  });
+
   it.each(["{}", "null", "[]", '["event"]', '"event"'])(
     "rejects empty or non-object stream events: %s",
     async (payload) => {

package/src/AdkClient.ts

@@ -1,7 +1,7 @@
 import { SSEEventDecoder } from "assistant-stream/utils";
+import { raceWithAbortSignal } from "@assistant-ui/core/internal";
 import { contentToParts } from "./contentToParts";
 import { parseAdkEventValue } from "./parseAdkEvent";
-import { raceWithAbortSignal } from "./raceWithAbortSignal";
 import { toAdkFunctionResponse } from "./toAdkFunctionResponse";
 import { trimTrailingSlashes } from "./trimTrailingSlashes";
 import type {
@@ -41,6 +41,12 @@ export type CreateAdkStreamOptions = {
     | Record<string, string>
     | (() => Record<string, string> | Promise<Record<string, string>>)
     | undefined;
+
+  /** Maximum UTF-16 code units accepted in one SSE line. Defaults to 16 MiB. */
+  maxStreamLineLength?: number | undefined;
+
+  /** Maximum UTF-16 code units retained across one SSE event. Defaults to 16 MiB. */
+  maxStreamEventLength?: number | undefined;
 };
 
 /**
@@ -97,7 +103,8 @@ export function createAdkStream(
     } else {
       // Proxy mode: POST in parseAdkRequest-compatible format
       url = options.api;
-      body = messagesToProxyBody(messages, config);
+      const { remoteId, externalId } = await config.initialize();
+      body = messagesToProxyBody(messages, config, externalId ?? remoteId);
     }
 
     const response = await fetch(url, {
@@ -114,7 +121,7 @@ export function createAdkStream(
     }
 
     validateEventStreamContentType(response);
-    yield* parseSSEResponse(response);
+    yield* parseSSEResponse(response, options);
   };
 }
 
@@ -209,8 +216,9 @@ function messagesToProxyBody(
     checkpointId?: string | undefined;
     stateDelta?: Record<string, unknown> | undefined;
   },
+  sessionId: string,
 ): Record<string, unknown> {
-  const body: Record<string, unknown> = {};
+  const body: Record<string, unknown> = { sessionId };
 
   if (config.runConfig != null) body.runConfig = config.runConfig;
   if (config.checkpointId != null) body.checkpointId = config.checkpointId;
@@ -255,16 +263,26 @@ function messagesToProxyBody(
   return body;
 }
 
-async function* parseSSEResponse(response: Response): AsyncGenerator<AdkEvent> {
+async function* parseSSEResponse(
+  response: Response,
+  options: Pick<
+    CreateAdkStreamOptions,
+    "maxStreamLineLength" | "maxStreamEventLength"
+  >,
+): AsyncGenerator<AdkEvent> {
   if (!response.body) {
     throw new Error("Expected ADK stream response body, received no body");
   }
   const reader = response.body.getReader();
   const decoder = new TextDecoder();
-  const sseDecoder = new SSEEventDecoder({ trailing: "dispatch" });
 
   let shouldCancel = true;
   try {
+    const sseDecoder = new SSEEventDecoder({
+      trailing: "dispatch",
+      maxLineLength: options.maxStreamLineLength,
+      maxEventLength: options.maxStreamEventLength,
+    });
     while (true) {
       let result: ReadableStreamReadResult<Uint8Array>;
       try {

package/src/AdkSessionAdapter.ts

@@ -5,12 +5,12 @@ import type {
   RemoteThreadListResponse,
   RemoteThreadMetadata,
 } from "@assistant-ui/core";
+import { raceWithAbortSignal } from "@assistant-ui/core/internal";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
 import { normalizeAdkPart } from "./normalizeAdkPart";
 import { parseAdkEventValue } from "./parseAdkEvent";
 import type { AdkMessage, AdkThreadSnapshot } from "./types";
 import { trimTrailingSlashes } from "./trimTrailingSlashes";
-import { raceWithAbortSignal } from "./raceWithAbortSignal";
 
 export type AdkSessionAdapterOptions = {
   /**

package/src/AdkThreadController.test.ts

@@ -0,0 +1,90 @@
+import { describe, expect, it, vi } from "vitest";
+import { AdkThreadController } from "./AdkThreadController";
+
+describe("AdkThreadController", () => {
+  it("selects staged messages through a parent in transcript order", () => {
+    const controller = new AdkThreadController();
+    const first = { id: "first", type: "human" as const, content: "one" };
+    const second = { id: "second", type: "human" as const, content: "two" };
+    const later = { id: "later", type: "human" as const, content: "three" };
+    controller.dispatch({
+      type: "staged.stage",
+      entry: { message: first, runConfig: { custom: { source: "first" } } },
+    });
+    controller.dispatch({
+      type: "staged.stage",
+      entry: { message: second, runConfig: { custom: { source: "second" } } },
+    });
+    controller.dispatch({
+      type: "staged.stage",
+      entry: { message: later, runConfig: undefined },
+    });
+
+    expect(controller.getStagedMessageCount()).toBe(3);
+    expect(
+      controller.getStagedRun("second", [
+        second,
+        { id: "canonical", type: "ai", content: "reply" },
+        first,
+        later,
+      ]),
+    ).toEqual({
+      messages: [second],
+      runConfig: { custom: { source: "second" } },
+    });
+    expect(
+      controller.getStagedRun("later", [
+        second,
+        { id: "canonical", type: "ai", content: "reply" },
+        first,
+        later,
+      ])?.messages,
+    ).toEqual([second, first, later]);
+    expect(controller.getStagedRun("canonical", [first, second])).toBeNull();
+
+    controller.dispatch({ type: "staged.unstage", ids: ["first", "second"] });
+    expect(controller.getStagedMessageCount()).toBe(1);
+  });
+
+  it("notifies for staging but not for unstaging an absent id", () => {
+    const controller = new AdkThreadController();
+    const listener = vi.fn();
+    controller.subscribe(listener);
+
+    controller.dispatch({
+      type: "staged.stage",
+      entry: {
+        message: { id: "first", type: "human", content: "one" },
+        runConfig: undefined,
+      },
+    });
+    expect(listener).toHaveBeenCalledTimes(1);
+
+    controller.dispatch({ type: "staged.unstage", ids: ["missing"] });
+    expect(listener).toHaveBeenCalledTimes(1);
+  });
+
+  it("publishes the reduced state synchronously to subscribers", () => {
+    const controller = new AdkThreadController();
+    const first = vi.fn(() => controller.getState());
+    const second = vi.fn();
+    const unsubscribeFirst = controller.subscribe(first);
+    controller.subscribe(second);
+    const initial = controller.getState();
+    const messages = [
+      { id: "message-1", type: "human" as const, content: "hi" },
+    ];
+
+    controller.dispatch({ type: "messages.set", messages });
+
+    expect(controller.getState()).not.toBe(initial);
+    expect(controller.getState().messages).toBe(messages);
+    expect(first).toHaveReturnedWith(controller.getState());
+    expect(second).toHaveBeenCalledTimes(1);
+
+    unsubscribeFirst();
+    controller.dispatch({ type: "messages.replaced", messages: [] });
+    expect(first).toHaveBeenCalledTimes(1);
+    expect(second).toHaveBeenCalledTimes(2);
+  });
+});

package/src/AdkThreadController.ts

@@ -0,0 +1,45 @@
+import {
+  createAdkThreadState,
+  reduceAdkThreadState,
+  type AdkThreadAction,
+} from "./adkThreadState";
+import type { AdkMessage } from "./types";
+
+export class AdkThreadController {
+  private state = createAdkThreadState();
+  private readonly listeners = new Set<() => void>();
+
+  public getState = () => this.state;
+
+  public getStagedMessageCount = () => this.state.stagedEntries.size;
+
+  public getStagedRun = (
+    parentId: string | null,
+    messages: AdkMessage[] = this.state.messages,
+  ) => {
+    const entries = this.state.stagedEntries;
+    if (!parentId || !entries.has(parentId)) return null;
+
+    const staged: AdkMessage[] = [];
+    for (const message of messages) {
+      if (message.id && entries.has(message.id)) {
+        staged.push(entries.get(message.id)!.message);
+      }
+      if (message.id === parentId) break;
+    }
+
+    return { messages: staged, runConfig: entries.get(parentId)!.runConfig };
+  };
+
+  public subscribe = (listener: () => void) => {
+    this.listeners.add(listener);
+    return () => this.listeners.delete(listener);
+  };
+
+  public dispatch = (action: AdkThreadAction) => {
+    const nextState = reduceAdkThreadState(this.state, action);
+    if (nextState === this.state) return;
+    this.state = nextState;
+    for (const listener of this.listeners) listener();
+  };
+}

package/src/adkThreadState.test.ts

@@ -0,0 +1,207 @@
+import { describe, expect, it } from "vitest";
+import { createAdkThreadState, reduceAdkThreadState } from "./adkThreadState";
+import type { AdkMessage } from "./types";
+
+const message: AdkMessage = {
+  id: "message-1",
+  type: "human",
+  content: "hello",
+};
+
+describe("reduceAdkThreadState", () => {
+  it("stages messages by id without changing the prior state", () => {
+    const initial = createAdkThreadState();
+    const first = {
+      message: { ...message },
+      runConfig: { custom: { model: "first" } },
+    };
+    const second = {
+      message: { ...message, id: "message-2" },
+      runConfig: { custom: { model: "second" } },
+    };
+
+    const staged = reduceAdkThreadState(initial, {
+      type: "staged.stage",
+      entry: first,
+    });
+    const next = reduceAdkThreadState(staged, {
+      type: "staged.stage",
+      entry: second,
+    });
+
+    expect(initial.stagedEntries.size).toBe(0);
+    expect(staged.stagedEntries?.size).toBe(1);
+    expect([...next.stagedEntries!]).toEqual([
+      ["message-1", first],
+      ["message-2", second],
+    ]);
+    expect(staged.stagedEntries?.size).toBe(1);
+  });
+
+  it("unstages a promoted run and preserves remaining entries through snapshots", () => {
+    const entries = ["message-1", "message-2", "message-3"].map((id) => ({
+      message: { ...message, id },
+      runConfig: undefined,
+    }));
+    const staged = entries.reduce(
+      (state, entry) =>
+        reduceAdkThreadState(state, { type: "staged.stage", entry }),
+      createAdkThreadState(),
+    );
+    const promoted = reduceAdkThreadState(staged, {
+      type: "staged.unstage",
+      ids: ["message-1", "message-2"],
+    });
+    const snapshotted = reduceAdkThreadState(promoted, {
+      type: "snapshot.applied",
+      snapshot: { messages: [message] },
+    });
+
+    expect([...promoted.stagedEntries!.keys()]).toEqual(["message-3"]);
+    expect([...snapshotted.stagedEntries!.keys()]).toEqual(["message-3"]);
+    expect(staged.stagedEntries?.size).toBe(3);
+    expect(
+      reduceAdkThreadState(promoted, {
+        type: "staged.unstage",
+        ids: ["missing"],
+      }),
+    ).toBe(promoted);
+  });
+
+  it("merges event deltas and metadata with previous state", () => {
+    const previous = {
+      ...createAdkThreadState(),
+      stagedEntries: new Map([
+        ["message-1", { message, runConfig: undefined }],
+      ]),
+      stateDelta: { retained: 1, replaced: "old" },
+      artifactDelta: { retained: 2, replaced: 3 },
+      messageMetadata: new Map([
+        ["retained", { groundingMetadata: "old" }],
+        ["replaced", { citationMetadata: "old" }],
+      ]),
+    };
+    const published = {
+      ...createAdkThreadState(),
+      messages: [message],
+      stateDelta: { replaced: "new" },
+      artifactDelta: { replaced: 4 },
+      agentInfo: { name: "agent" },
+      longRunningToolIds: ["tool-1"],
+      escalated: true,
+      messageMetadata: new Map([["replaced", { citationMetadata: "new" }]]),
+    };
+    const { stagedEntries: _stagedEntries, ...publishedState } = published;
+
+    const next = reduceAdkThreadState(previous, {
+      type: "event.published",
+      state: publishedState,
+    });
+
+    expect(next).toMatchObject({
+      messages: [message],
+      stateDelta: { retained: 1, replaced: "new" },
+      artifactDelta: { retained: 2, replaced: 4 },
+      agentInfo: { name: "agent" },
+      longRunningToolIds: ["tool-1"],
+      escalated: true,
+    });
+    expect([...next.messageMetadata]).toEqual([
+      ["retained", { groundingMetadata: "old" }],
+      ["replaced", { citationMetadata: "new" }],
+    ]);
+    expect(next.stagedEntries).toBe(previous.stagedEntries);
+    expect(previous.stateDelta).toEqual({ retained: 1, replaced: "old" });
+    expect(previous.messageMetadata.get("replaced")).toEqual({
+      citationMetadata: "old",
+    });
+
+    const withoutMetadata = reduceAdkThreadState(next, {
+      type: "event.published",
+      state: { ...publishedState, messageMetadata: new Map() },
+    });
+    expect(withoutMetadata.messageMetadata).toBe(next.messageMetadata);
+  });
+
+  it("swaps a snapshot and clears omitted fields", () => {
+    const previous = {
+      ...createAdkThreadState(),
+      stateDelta: { stale: true },
+      artifactDelta: { stale: 1 },
+      agentInfo: { name: "stale-agent", branch: "stale" },
+      longRunningToolIds: ["stale"],
+      toolConfirmations: [
+        {
+          toolCallId: "stale-tool",
+          toolName: "approve",
+          args: {},
+          hint: "stale",
+          confirmed: false,
+        },
+      ],
+      authRequests: [{ toolCallId: "stale-auth", authConfig: {} }],
+      escalated: true,
+      messageMetadata: new Map([["stale", { groundingMetadata: "stale" }]]),
+    };
+    const snapshot = {
+      messages: [message],
+      stateDelta: { loaded: true },
+      messageMetadata: new Map([["message-1", { usageMetadata: 1 }]]),
+    };
+
+    const next = reduceAdkThreadState(previous, {
+      type: "snapshot.applied",
+      snapshot,
+    });
+
+    expect(next).toEqual({
+      messages: snapshot.messages,
+      stateDelta: snapshot.stateDelta,
+      agentInfo: {},
+      longRunningToolIds: [],
+      artifactDelta: {},
+      toolConfirmations: [],
+      authRequests: [],
+      escalated: false,
+      messageMetadata: snapshot.messageMetadata,
+      stagedEntries: previous.stagedEntries,
+    });
+  });
+
+  it("replaces messages and clears per-turn state while retaining thread deltas", () => {
+    const previous = {
+      ...createAdkThreadState(),
+      stateDelta: { retained: true },
+      artifactDelta: { file: 1 },
+      agentInfo: { name: "agent" },
+      longRunningToolIds: ["tool-1"],
+      toolConfirmations: [
+        {
+          toolCallId: "tool-1",
+          toolName: "search",
+          args: {},
+          hint: "approve",
+          confirmed: false,
+        },
+      ],
+      authRequests: [{ toolCallId: "tool-2", authConfig: {} }],
+      escalated: true,
+      messageMetadata: new Map([["old", { usageMetadata: 1 }]]),
+    };
+
+    const next = reduceAdkThreadState(previous, {
+      type: "messages.replaced",
+      messages: [message],
+    });
+
+    expect(next.messages).toEqual([message]);
+    expect(next.longRunningToolIds).toEqual([]);
+    expect(next.toolConfirmations).toEqual([]);
+    expect(next.authRequests).toEqual([]);
+    expect(next.escalated).toBe(false);
+    expect(next.messageMetadata.size).toBe(0);
+    expect(next.stateDelta).toBe(previous.stateDelta);
+    expect(next.artifactDelta).toBe(previous.artifactDelta);
+    expect(next.agentInfo).toBe(previous.agentInfo);
+  });
+});