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

package/src/server/parseAdkRequest.test.ts

@@ -28,11 +28,18 @@ describe("parseAdkRequest", () => {
     });
   });
 
-  it("parses a message request with stateDelta", async () => {
+  it("parses a message request with sessionId and stateDelta", async () => {
     const result = await parseAdkRequest(
-      makeRequest({ message: "Hello", stateDelta: { count: 1 } }),
+      makeRequest({
+        message: "Hello",
+        sessionId: "session-1",
+        stateDelta: { count: 1 },
+      }),
     );
-    expect(result).toMatchObject({ stateDelta: { count: 1 } });
+    expect(result).toMatchObject({
+      sessionId: "session-1",
+      stateDelta: { count: 1 },
+    });
   });
 
   it("parses a message request with parts (multimodal)", async () => {
@@ -130,6 +137,7 @@ describe("parseAdkRequest", () => {
     [{ parts: [null] }, 'field "parts"'],
     [{ type: "unknown", message: "hello" }, 'field "type"'],
     [{ message: "hello", checkpointId: 42 }, 'field "checkpointId"'],
+    [{ message: "hello", sessionId: 42 }, 'field "sessionId"'],
     [{ message: "hello", stateDelta: [] }, 'field "stateDelta"'],
   ])("rejects malformed message requests %#", async (body, error) => {
     await expect(parseAdkRequest(makeRequest(body))).rejects.toThrow(error);

package/src/server/parseAdkRequest.ts

@@ -7,6 +7,7 @@ type ParsedAdkRequest =
       type: "message";
       text: string;
       parts?: Array<Record<string, unknown>> | undefined;
+      sessionId?: string | undefined;
       config: AdkSendMessageConfig;
       stateDelta?: Record<string, unknown> | undefined;
     }
@@ -16,6 +17,7 @@ type ParsedAdkRequest =
       toolName: string;
       result: unknown;
       isError: boolean;
+      sessionId?: string | undefined;
       config: AdkSendMessageConfig;
       stateDelta?: Record<string, unknown> | undefined;
     };
@@ -162,6 +164,7 @@ export const parseAdkRequest = async (
   if (body.runConfig !== undefined) config.runConfig = body.runConfig;
   const checkpointId = readOptionalString(body, "checkpointId");
   if (checkpointId !== undefined) config.checkpointId = checkpointId;
+  const sessionId = readOptionalString(body, "sessionId");
 
   const stateDelta = body.stateDelta;
   if (stateDelta !== undefined && !isRecord(stateDelta)) {
@@ -181,6 +184,7 @@ export const parseAdkRequest = async (
       toolName: readString(body, "toolName"),
       result: body.result,
       isError: body.isError ?? false,
+      ...(sessionId !== undefined && { sessionId }),
       config,
       ...(stateDelta != null && { stateDelta }),
     };
@@ -213,6 +217,7 @@ export const parseAdkRequest = async (
     type: "message",
     text: text ?? "",
     ...(parts !== undefined && { parts }),
+    ...(sessionId !== undefined && { sessionId }),
     config,
     ...(stateDelta != null && { stateDelta }),
   };
@@ -226,7 +231,8 @@ export const parseAdkRequest = async (
  * ```ts
  * const parsed = await parseAdkRequest(req);
  * const newMessage = toAdkContent(parsed);
- * const events = runner.runAsync({ userId, sessionId, newMessage, stateDelta: parsed.stateDelta });
+ * const stateDelta = validateSessionState(parsed.stateDelta);
+ * const events = runner.runAsync({ userId, sessionId, newMessage, stateDelta });
  * return adkEventStream(events);
  * ```
  */

package/src/useAdkMessages.test.ts

@@ -1,3 +1,4 @@
+/** @vitest-environment jsdom */
 import { act, renderHook } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 
@@ -17,6 +18,7 @@ import {
   messageToEvent,
   messagesToEvents,
   useAdkMessages,
+  useAdkMessagesInternal,
 } from "./useAdkMessages";
 import { projectAdkToolApprovals } from "./adkToolApproval";
 import { createAdkStream } from "./AdkClient";
@@ -212,6 +214,47 @@ describe("ADK runtime callbacks", () => {
 });
 
 describe("ADK stream lifecycle", () => {
+  it("reports streamed tool calls with the run config that produced them", async () => {
+    const runConfig = { custom: { model: "model-a" } };
+    const onMessages = vi.fn();
+    const stream: AdkStreamCallback = async function* () {
+      yield {
+        id: "event-1",
+        author: "agent",
+        content: {
+          role: "model",
+          parts: [
+            { functionCall: { id: "tool-1", name: "lookup", args: {} } },
+            { functionCall: { id: "tool-2", name: "search", args: {} } },
+          ],
+        },
+      };
+    };
+    const { result } = renderHook(() =>
+      useAdkMessagesInternal({ stream, onMessages }),
+    );
+
+    await act(async () => {
+      await result.current.sendMessage(
+        [{ id: "user-1", type: "human", content: "look it up" }],
+        { runConfig },
+      );
+    });
+
+    expect(onMessages).toHaveBeenLastCalledWith(
+      expect.arrayContaining([
+        expect.objectContaining({
+          type: "ai",
+          tool_calls: [
+            expect.objectContaining({ id: "tool-1" }),
+            expect.objectContaining({ id: "tool-2" }),
+          ],
+        }),
+      ]),
+      runConfig,
+    );
+  });
+
   it("settles a superseded send while its stream is still opening", async () => {
     const signals: AbortSignal[] = [];
     const parked = new Promise<AsyncGenerator<AdkEvent>>(() => {});

package/src/useAdkMessages.ts

@@ -2,9 +2,9 @@ import { isRecord } from "@assistant-ui/core/internal";
 import {
   useState,
   useCallback,
-  useInsertionEffect,
   useRef,
   useMemo,
+  useSyncExternalStore,
 } from "react";
 import { generateId } from "@assistant-ui/core";
 import { useReplaySafeEffect } from "@assistant-ui/store/internal";
@@ -15,16 +15,14 @@ import {
   openAbortableIterable,
 } from "@assistant-ui/core/internal";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
+import { AdkThreadController } from "./AdkThreadController";
 import { contentToParts } from "./contentToParts";
 import { toAdkFunctionResponse } from "./toAdkFunctionResponse";
 import type {
   AdkEvent,
   AdkMessage,
-  AdkMessageMetadata,
   AdkSendMessageConfig,
   AdkStreamCallback,
-  AdkToolConfirmation,
-  AdkAuthRequest,
   AdkThreadSnapshot,
   OnAdkErrorCallback,
   OnAdkCustomEventCallback,
@@ -40,6 +38,10 @@ export type UseAdkMessagesOptions = {
   };
 };
 
+type UseAdkMessagesInternalOptions = UseAdkMessagesOptions & {
+  onMessages?: (messages: AdkMessage[], runConfig: unknown) => void;
+};
+
 type AdkRuntimeCallbackName = "onError" | "onCustomEvent" | "onAgentTransfer";
 
 const invokeAdkRuntimeCallback = <TArgs extends readonly unknown[]>(
@@ -50,53 +52,34 @@ const invokeAdkRuntimeCallback = <TArgs extends readonly unknown[]>(
   void invokeUserCallback("react-google-adk", name, callback, ...args);
 };
 
-export const useAdkMessages = ({
+const useAdkMessagesInternal = ({
   stream,
   eventHandlers,
-}: UseAdkMessagesOptions) => {
-  const [messages, _setMessages] = useState<AdkMessage[]>([]);
-  const [stateDelta, setStateDelta] = useState<Record<string, unknown>>({});
-  const [agentInfo, setAgentInfo] = useState<{
-    name?: string | undefined;
-    branch?: string | undefined;
-  }>({});
-  const [longRunningToolIds, _setLongRunningToolIds] = useState<string[]>([]);
-  const [artifactDelta, setArtifactDelta] = useState<Record<string, number>>(
-    {},
+  onMessages,
+}: UseAdkMessagesInternalOptions) => {
+  const [controller] = useState(() => new AdkThreadController());
+  const {
+    messages,
+    stateDelta,
+    agentInfo,
+    longRunningToolIds,
+    artifactDelta,
+    toolConfirmations,
+    authRequests,
+    escalated,
+    messageMetadata,
+  } = useSyncExternalStore(
+    controller.subscribe,
+    controller.getState,
+    controller.getState,
   );
-  const [toolConfirmations, setToolConfirmations] = useState<
-    AdkToolConfirmation[]
-  >([]);
-  const [authRequests, setAuthRequests] = useState<AdkAuthRequest[]>([]);
-  const [escalated, setEscalated] = useState(false);
-  const [messageMetadata, setMessageMetadata] = useState<
-    Map<string, AdkMessageMetadata>
-  >(new Map());
-  // setMessagesImmediate and setLongRunningToolIds are the only writers of their state and publish these refs with it, so neither ref trails a commit.
-  const messagesRef = useRef(messages);
-  const longRunningToolIdsRef = useRef(longRunningToolIds);
-  const stateDeltaRef = useRef(stateDelta);
-  useInsertionEffect(() => {
-    stateDeltaRef.current = stateDelta;
-  }, [stateDelta]);
-  const artifactDeltaRef = useRef(artifactDelta);
-  useInsertionEffect(() => {
-    artifactDeltaRef.current = artifactDelta;
-  }, [artifactDelta]);
-  const messageMetadataRef = useRef(messageMetadata);
-  useInsertionEffect(() => {
-    messageMetadataRef.current = messageMetadata;
-  }, [messageMetadata]);
-
-  const setMessagesImmediate = useCallback((msgs: AdkMessage[]) => {
-    messagesRef.current = msgs;
-    _setMessages(msgs);
-  }, []);
-  const setLongRunningToolIds = useCallback((ids: string[]) => {
-    longRunningToolIdsRef.current = ids;
-    _setLongRunningToolIds(ids);
-  }, []);
 
+  const setMessagesImmediate = useCallback(
+    (msgs: AdkMessage[]) => {
+      controller.dispatch({ type: "messages.set", messages: msgs });
+    },
+    [controller],
+  );
   /**
    * Swap the thread over to a loaded snapshot in one commit. Unlike
    * {@link replaceMessages} this never passes through a cleared state, so a
@@ -105,17 +88,9 @@ export const useAdkMessages = ({
    */
   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 ?? {});
+      controller.dispatch({ type: "snapshot.applied", snapshot });
     },
-    [setLongRunningToolIds, setMessagesImmediate],
+    [controller],
   );
 
   // Replace the message list AND reset derived per-turn HITL state.
@@ -124,14 +99,9 @@ export const useAdkMessages = ({
   // into the next turn.
   const replaceMessages = useCallback(
     (msgs: AdkMessage[]) => {
-      setMessagesImmediate(msgs);
-      setLongRunningToolIds([]);
-      setToolConfirmations([]);
-      setAuthRequests([]);
-      setEscalated(false);
-      setMessageMetadata(new Map());
+      controller.dispatch({ type: "messages.replaced", messages: msgs });
     },
-    [setLongRunningToolIds, setMessagesImmediate],
+    [controller],
   );
 
   const abortControllerRef = useRef<AbortController | null>(null);
@@ -159,9 +129,10 @@ export const useAdkMessages = ({
           m.type === "tool" ? [m.tool_call_id] : [],
         ),
       );
+      const currentState = controller.getState();
       const accumulator = new AdkEventAccumulator(
-        messagesRef.current.filter((m) => !resentIds.has(m.id)),
-        longRunningToolIdsRef.current.filter(
+        currentState.messages.filter((m) => !resentIds.has(m.id)),
+        currentState.longRunningToolIds.filter(
           (id) => !answeredToolCallIds.has(id),
         ),
       );
@@ -171,10 +142,13 @@ export const useAdkMessages = ({
       const initialMessages = accumulator.getMessages();
       const initialMessageIds = new Set(initialMessages.map((m) => m.id));
       const initialLongRunningToolIds = accumulator.getLongRunningToolIds();
-      setMessagesImmediate(initialMessages);
-      setLongRunningToolIds(initialLongRunningToolIds);
-      setToolConfirmations(accumulator.getToolConfirmations());
-      setAuthRequests(accumulator.getAuthRequests());
+      controller.dispatch({
+        type: "run.started",
+        messages: initialMessages,
+        longRunningToolIds: initialLongRunningToolIds,
+        toolConfirmations: accumulator.getToolConfirmations(),
+        authRequests: accumulator.getAuthRequests(),
+      });
       let lastTransferToAgent: string | undefined;
 
       // Google ADK replaces active runs, while React LangGraph queues sends.
@@ -206,28 +180,31 @@ export const useAdkMessages = ({
             break;
           }
           const updatedMessages = accumulator.processEvent(event);
-          setMessagesImmediate(updatedMessages);
-          setStateDelta({
-            ...stateDeltaRef.current,
-            ...accumulator.getStateDelta(),
-          });
-          setAgentInfo(accumulator.getAgentInfo());
-          setLongRunningToolIds(accumulator.getLongRunningToolIds());
-          setArtifactDelta({
-            ...artifactDeltaRef.current,
-            ...accumulator.getArtifactDelta(),
-          });
-          setToolConfirmations(accumulator.getToolConfirmations());
-          setAuthRequests(accumulator.getAuthRequests());
-          setEscalated(accumulator.isEscalated());
-          {
-            const newMeta = accumulator.getMessageMetadata();
-            if (newMeta.size > 0) {
-              setMessageMetadata(
-                new Map([...messageMetadataRef.current, ...newMeta]),
-              );
-            }
+          // Each event part can append at most one message, and a function call
+          // stays on the current assistant message until a later part finalizes
+          // it, so every message touched by this event is within this tail.
+          const affectedMessageCount = Math.max(
+            event.content?.parts?.length ?? 0,
+            1,
+          );
+          const affectedMessages = updatedMessages.slice(-affectedMessageCount);
+          if (affectedMessages.length > 0) {
+            onMessages?.(affectedMessages, config.runConfig);
           }
+          controller.dispatch({
+            type: "event.published",
+            state: {
+              messages: updatedMessages,
+              stateDelta: accumulator.getStateDelta(),
+              agentInfo: accumulator.getAgentInfo(),
+              longRunningToolIds: accumulator.getLongRunningToolIds(),
+              artifactDelta: accumulator.getArtifactDelta(),
+              toolConfirmations: accumulator.getToolConfirmations(),
+              authRequests: accumulator.getAuthRequests(),
+              escalated: accumulator.isEscalated(),
+              messageMetadata: accumulator.getMessageMetadata(),
+            },
+          });
 
           const transfer = accumulator.getLastTransferToAgent();
           if (transfer && transfer !== lastTransferToAgent) {
@@ -270,12 +247,13 @@ export const useAdkMessages = ({
       } finally {
         if (abortControllerRef.current === abortController) {
           if (abortController.signal.aborted) {
-            setLongRunningToolIds(
-              accumulator
+            controller.dispatch({
+              type: "longRunningToolIds.set",
+              ids: accumulator
                 .getLongRunningToolIds()
                 .filter((id) => initialLongRunningToolIds.includes(id)),
-            );
-            const updatedMessages = messagesRef.current;
+            });
+            const updatedMessages = controller.getState().messages;
             const lastAssistantMessage = updatedMessages.findLast(
               (m) => m.type === "ai",
             );
@@ -302,12 +280,13 @@ export const useAdkMessages = ({
     },
     [
       aui,
+      controller,
       setMessagesImmediate,
-      setLongRunningToolIds,
       stream,
       onError,
       onCustomEvent,
       onAgentTransfer,
+      onMessages,
     ],
   );
 
@@ -320,6 +299,7 @@ export const useAdkMessages = ({
   useReplaySafeEffect(() => cancel, []);
 
   return {
+    controller,
     messages,
     stateDelta,
     agentInfo,
@@ -337,6 +317,19 @@ export const useAdkMessages = ({
   };
 };
 
+export const useAdkMessages = ({
+  stream,
+  eventHandlers,
+}: UseAdkMessagesOptions) => {
+  const { controller: _controller, ...result } = useAdkMessagesInternal({
+    stream,
+    ...(eventHandlers !== undefined && { eventHandlers }),
+  });
+  return result;
+};
+
+export { useAdkMessagesInternal };
+
 /**
  * Transport sends every human and tool message of one `send` call as a single
  * ADK `Content`, and ADK parses that event's function responses before running

package/src/useAdkRuntime.cancellation.test.tsx

@@ -11,6 +11,7 @@ import type {
 import { useAdkRuntime } from "./useAdkRuntime";
 import { useAdkLongRunningToolIds, useAdkSend } from "./hooks";
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
+import { settleOutsideAct } from "./tests/settleOutsideAct";
 import type { AdkEvent, AdkStreamCallback, AdkThreadSnapshot } from "./types";
 
 const makeThreadListAdapter = (): RemoteThreadListAdapter => ({
@@ -85,9 +86,9 @@ const renderAdkRuntime = async (
     render(<Inner />);
   });
   await waitFor(() => expect(capture.runtime).not.toBeNull());
-  await act(async () => {
-    await capture.runtime!.threads.switchToThread("adk-1");
-  });
+  await settleOutsideAct(() =>
+    capture.runtime!.threads.switchToThread("adk-1"),
+  );
 
   return capture;
 };

package/src/useAdkRuntime.cloud-options.test.tsx

@@ -0,0 +1,59 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import type { AssistantCloud } from "assistant-cloud";
+import type { AssistantRuntime } from "@assistant-ui/core";
+
+const mocks = vi.hoisted(() => ({
+  cloudAdapter: {},
+  runtime: {},
+  useCloudThreadListAdapter: vi.fn(() => ({})),
+  useRemoteThreadListRuntime: vi.fn(() => ({})),
+}));
+
+vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/core/react")>()),
+  useCloudThreadListAdapter: mocks.useCloudThreadListAdapter,
+  useRemoteThreadListRuntime: mocks.useRemoteThreadListRuntime,
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => ({
+    threadListItem: {
+      source: null,
+      getState: () => ({ externalId: undefined }),
+      initialize: vi.fn(),
+    },
+  }),
+}));
+
+import { ADK_SDK } from "./sdkIdentity";
+import { useAdkRuntime } from "./useAdkRuntime";
+
+describe("useAdkRuntime Cloud options", () => {
+  it("forwards the Cloud scope to the thread-list adapter", () => {
+    const cloud = {} as AssistantCloud;
+    mocks.useCloudThreadListAdapter.mockReturnValue(mocks.cloudAdapter);
+    mocks.useRemoteThreadListRuntime.mockReturnValue(
+      mocks.runtime as AssistantRuntime,
+    );
+
+    renderHook(() =>
+      useAdkRuntime({
+        stream: vi.fn(),
+        cloud,
+        scopeId: "workspace-1",
+      }),
+    );
+
+    expect(mocks.useCloudThreadListAdapter).toHaveBeenCalledWith(
+      expect.objectContaining({
+        cloud,
+        scopeId: "workspace-1",
+        sdk: ADK_SDK,
+      }),
+    );
+  });
+});