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

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,
@@ -59,49 +57,29 @@ const useAdkMessagesInternal = ({
   eventHandlers,
   onMessages,
 }: UseAdkMessagesInternalOptions) => {
-  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>>(
-    {},
+  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
@@ -110,17 +88,9 @@ const useAdkMessagesInternal = ({
    */
   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.
@@ -129,14 +99,9 @@ const useAdkMessagesInternal = ({
   // 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);
@@ -164,9 +129,10 @@ const useAdkMessagesInternal = ({
           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),
         ),
       );
@@ -176,10 +142,13 @@ const useAdkMessagesInternal = ({
       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.
@@ -222,28 +191,20 @@ const useAdkMessagesInternal = ({
           if (affectedMessages.length > 0) {
             onMessages?.(affectedMessages, config.runConfig);
           }
-          setMessagesImmediate(updatedMessages);
-          setStateDelta({
-            ...stateDeltaRef.current,
-            ...accumulator.getStateDelta(),
-          });
-          setAgentInfo(accumulator.getAgentInfo());
-          setLongRunningToolIds(accumulator.getLongRunningToolIds());
-          setArtifactDelta({
-            ...artifactDeltaRef.current,
-            ...accumulator.getArtifactDelta(),
+          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(),
+            },
           });
-          setToolConfirmations(accumulator.getToolConfirmations());
-          setAuthRequests(accumulator.getAuthRequests());
-          setEscalated(accumulator.isEscalated());
-          {
-            const newMeta = accumulator.getMessageMetadata();
-            if (newMeta.size > 0) {
-              setMessageMetadata(
-                new Map([...messageMetadataRef.current, ...newMeta]),
-              );
-            }
-          }
 
           const transfer = accumulator.getLastTransferToAgent();
           if (transfer && transfer !== lastTransferToAgent) {
@@ -286,12 +247,13 @@ const useAdkMessagesInternal = ({
       } 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",
             );
@@ -318,8 +280,8 @@ const useAdkMessagesInternal = ({
     },
     [
       aui,
+      controller,
       setMessagesImmediate,
-      setLongRunningToolIds,
       stream,
       onError,
       onCustomEvent,
@@ -337,6 +299,7 @@ const useAdkMessagesInternal = ({
   useReplaySafeEffect(() => cancel, []);
 
   return {
+    controller,
     messages,
     stateDelta,
     agentInfo,
@@ -357,11 +320,13 @@ const useAdkMessagesInternal = ({
 export const useAdkMessages = ({
   stream,
   eventHandlers,
-}: UseAdkMessagesOptions) =>
-  useAdkMessagesInternal({
+}: UseAdkMessagesOptions) => {
+  const { controller: _controller, ...result } = useAdkMessagesInternal({
     stream,
     ...(eventHandlers !== undefined && { eventHandlers }),
   });
+  return result;
+};
 
 export { useAdkMessagesInternal };
 

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