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

package/src/useAdkRuntime.ts

@@ -1,6 +1,5 @@
 import {
   useCallback,
-  useEffect,
   useInsertionEffect,
   useMemo,
   useRef,
@@ -22,6 +21,7 @@ import {
 import {
   createAbortableThreadLoad,
   createCloudThreadListAdapterCreateFallback,
+  isRecord,
 } from "@assistant-ui/core/internal";
 import {
   useCloudThreadListAdapter,
@@ -30,6 +30,7 @@ import {
   useExternalStoreRuntime,
 } from "@assistant-ui/core/react";
 import { useAui } from "@assistant-ui/store";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 import type { AssistantCloud } from "assistant-cloud";
 import type { RemoteThreadListAdapter } from "@assistant-ui/core";
 import type {
@@ -41,7 +42,7 @@ import type {
   OnAdkCustomEventCallback,
   OnAdkAgentTransferCallback,
 } from "./types";
-import { useAdkMessages } from "./useAdkMessages";
+import { useAdkMessagesInternal } from "./useAdkMessages";
 import {
   convertAdkMessage,
   createAdkMessageConverter,
@@ -123,6 +124,62 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     eventHandlers,
   } = options;
   const aui = useAui();
+  const runConfigByToolCallIdRef = useRef(new Map<string, unknown>());
+
+  const rememberMessageOwnership = useCallback(
+    (newMessages: AdkMessage[], runConfig: unknown) => {
+      const toolOwnership = runConfigByToolCallIdRef.current;
+      for (const message of newMessages) {
+        if (message.type !== "ai") continue;
+        for (const toolCall of message.tool_calls ?? []) {
+          if (!isRecord(toolCall)) continue;
+          if (!toolOwnership.has(toolCall.id)) {
+            toolOwnership.set(toolCall.id, runConfig);
+          }
+        }
+      }
+    },
+    [],
+  );
+
+  const seedMessageOwnership = useCallback((history: AdkMessage[]) => {
+    const currentOwnership = runConfigByToolCallIdRef.current;
+    const nextOwnership = new Map<string, unknown>();
+    for (const message of history) {
+      if (message.type !== "ai") continue;
+      for (const toolCall of message.tool_calls ?? []) {
+        if (!isRecord(toolCall)) continue;
+        // Loaded ids must remain present even without a local owner because
+        // streamed event windows use has() to avoid attributing them later.
+        nextOwnership.set(
+          toolCall.id,
+          currentOwnership.has(toolCall.id)
+            ? currentOwnership.get(toolCall.id)
+            : undefined,
+        );
+      }
+    }
+    runConfigByToolCallIdRef.current = nextOwnership;
+  }, []);
+
+  const pruneMessageOwnership = useCallback((history: AdkMessage[]) => {
+    const toolCallIds = new Set<string>();
+    for (const message of history) {
+      if (message.type !== "ai") continue;
+      for (const toolCall of message.tool_calls ?? []) {
+        if (!isRecord(toolCall)) continue;
+        toolCallIds.add(toolCall.id);
+      }
+    }
+    for (const id of runConfigByToolCallIdRef.current.keys()) {
+      if (!toolCallIds.has(id)) runConfigByToolCallIdRef.current.delete(id);
+    }
+  }, []);
+
+  const getToolRunConfig = useCallback((toolCallId: string) => {
+    return runConfigByToolCallIdRef.current.get(toolCallId);
+  }, []);
+
   const {
     messages,
     stateDelta,
@@ -136,22 +193,37 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     sendMessage,
     cancel,
     setMessages,
-    replaceMessages,
-    applySnapshot,
-  } = useAdkMessages({
+    replaceMessages: replaceAdkMessages,
+    applySnapshot: applyAdkSnapshot,
+  } = useAdkMessagesInternal({
     stream,
     ...(eventHandlers && { eventHandlers }),
+    onMessages: rememberMessageOwnership,
   });
 
   const loadRef = useRef(load);
   useInsertionEffect(() => {
     loadRef.current = load;
   }, [load]);
-  const loadController = useMemo(createAbortableThreadLoad, []);
+  const [loadController] = useState(createAbortableThreadLoad);
   const messagesRef = useRef(messages);
   useInsertionEffect(() => {
     messagesRef.current = messages;
   }, [messages]);
+  const applySnapshot = useCallback(
+    (snapshot: AdkThreadSnapshot) => {
+      seedMessageOwnership(snapshot.messages);
+      applyAdkSnapshot(snapshot);
+    },
+    [applyAdkSnapshot, seedMessageOwnership],
+  );
+  const replaceMessages = useCallback(
+    (nextMessages: AdkMessage[]) => {
+      pruneMessageOwnership(nextMessages);
+      replaceAdkMessages(nextMessages);
+    },
+    [pruneMessageOwnership, replaceAdkMessages],
+  );
   const [isLoadingThread, setIsLoadingThread] = useState(
     () =>
       load !== undefined && aui.threadListItem.getState().externalId != null,
@@ -175,10 +247,20 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     msgs: AdkMessage[],
     config: AdkSendMessageConfig,
   ) => {
+    const isToolContinuation =
+      msgs.length > 0 && msgs.every((msg) => msg.type === "tool");
+    const continuationConfig =
+      isToolContinuation && config.runConfig === undefined
+        ? {
+            ...config,
+            runConfig: getToolRunConfig(msgs[0]!.tool_call_id),
+          }
+        : config;
+
     const generation = ++runGenerationRef.current;
     try {
       setIsRunning(true);
-      await sendMessage(msgs, config);
+      await sendMessage(msgs, continuationConfig);
     } finally {
       if (runGenerationRef.current === generation) setIsRunning(false);
     }
@@ -325,7 +407,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     [threadListItem, loadController, applySnapshot],
   );
 
-  useEffect(() => {
+  useReplaySafeEffect(() => {
     runLoad();
     return () => {
       // Whatever is current, not this effect's own controller: a refetch swaps
@@ -333,7 +415,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       loadController.abort();
       setIsLoadingThread(false);
     };
-  }, [loadController, runLoad]);
+  }, [threadListItem]);
 
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),

package/src/useAdkRuntimeApproval.test.tsx

@@ -1,17 +1,33 @@
 import { act, renderHook } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type {
+  AddToolResultOptions,
+  AppendMessage,
   RespondToToolApprovalOptions,
   ThreadMessage,
   ToolCallMessagePart,
 } from "@assistant-ui/core";
-import type { AdkMessage } from "./types";
+import type { AdkMessage, AdkSendMessageConfig } from "./types";
 
-const mocks = vi.hoisted(() => ({
-  adapters: [] as unknown[],
-  sendMessage: vi.fn().mockResolvedValue(undefined),
-  messages: [] as AdkMessage[],
-}));
+const mocks = vi.hoisted(() => {
+  const threadListItem = {
+    source: null as object | null,
+    externalId: undefined as string | undefined,
+    getState: () => ({
+      externalId: threadListItem.externalId,
+    }),
+    initialize: vi.fn(),
+  };
+  return {
+    adapters: [] as unknown[],
+    sendMessage: vi.fn().mockResolvedValue(undefined),
+    messages: [] as AdkMessage[],
+    messageRunConfig: undefined as unknown,
+    streamedMessages: undefined as AdkMessage[] | undefined,
+    applySnapshot: vi.fn(),
+    threadListItem,
+  };
+});
 
 vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
   ...(await importOriginal<typeof import("@assistant-ui/core/react")>()),
@@ -27,49 +43,88 @@ vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
 vi.mock("@assistant-ui/store", async (importOriginal) => ({
   ...(await importOriginal<typeof import("@assistant-ui/store")>()),
   useAui: () => ({
-    threadListItem: {
-      source: null,
-      getState: () => ({ externalId: undefined }),
-      initialize: vi.fn(),
-    },
+    threadListItem: mocks.threadListItem,
   }),
 }));
 
-vi.mock("./useAdkMessages", async (importOriginal) => ({
-  ...(await importOriginal<typeof import("./useAdkMessages")>()),
-  useAdkMessages: () => ({
-    messages: mocks.messages,
-    stateDelta: {},
-    agentInfo: {},
-    longRunningToolIds: [],
-    artifactDelta: {},
-    // Deliberately empty: the projection must come from the transcript, not
-    // from derived confirmation state that a mid-run accumulator swap drops.
-    toolConfirmations: [],
-    authRequests: [],
-    escalated: false,
-    messageMetadata: new Map(),
-    sendMessage: mocks.sendMessage,
-    cancel: vi.fn(),
-    setMessages: vi.fn(),
-    replaceMessages: vi.fn(),
-    applySnapshot: vi.fn(),
-  }),
-}));
+vi.mock("./useAdkMessages", async (importOriginal) => {
+  const useMockAdkMessages = (options: {
+    onMessages?: (messages: AdkMessage[], runConfig: unknown) => void;
+  }) => {
+    options.onMessages?.(mocks.messages, mocks.messageRunConfig);
+    const sendMessage = async (
+      messages: AdkMessage[],
+      config: AdkSendMessageConfig,
+    ) => {
+      await mocks.sendMessage(messages, config);
+      if (mocks.streamedMessages) {
+        options.onMessages?.(mocks.streamedMessages, config.runConfig);
+      }
+    };
+    return {
+      messages: mocks.messages,
+      stateDelta: {},
+      agentInfo: {},
+      longRunningToolIds: [],
+      artifactDelta: {},
+      // Deliberately empty: the projection must come from the transcript, not
+      // from derived confirmation state that a mid-run accumulator swap drops.
+      toolConfirmations: [],
+      authRequests: [],
+      escalated: false,
+      messageMetadata: new Map(),
+      sendMessage,
+      cancel: vi.fn(),
+      setMessages: vi.fn(),
+      replaceMessages: vi.fn(),
+      applySnapshot: mocks.applySnapshot,
+    };
+  };
+
+  return {
+    ...(await importOriginal<typeof import("./useAdkMessages")>()),
+    useAdkMessages: useMockAdkMessages,
+    useAdkMessagesInternal: useMockAdkMessages,
+  };
+});
 
 import { AdkEventAccumulator } from "./AdkEventAccumulator";
 import { useAdkRuntime } from "./useAdkRuntime";
 
-type ApprovalAdapter = {
+type RuntimeAdapter = {
   messages: readonly ThreadMessage[];
+  extras: {
+    send: (
+      messages: AdkMessage[],
+      config: AdkSendMessageConfig,
+    ) => Promise<void>;
+  };
+  onNew?: (message: AppendMessage) => Promise<void> | void;
+  onEdit?: (message: AppendMessage) => Promise<void> | void;
+  onAddToolResult?: (options: AddToolResultOptions) => Promise<void> | void;
   onRespondToToolApproval?: (
     options: RespondToToolApprovalOptions,
   ) => Promise<void> | void;
+  onRefetchThread?: () => Promise<void> | void;
 };
 
 const CONFIRMATION_CALL = "adk-confirmation-1";
 
-const latestAdapter = () => mocks.adapters.at(-1) as ApprovalAdapter;
+const latestAdapter = () => mocks.adapters.at(-1) as RuntimeAdapter;
+
+const makeUserMessage = (
+  text: string,
+  runConfig: AppendMessage["runConfig"],
+): AppendMessage =>
+  ({
+    role: "user",
+    content: [{ type: "text", text }],
+    attachments: [],
+    parentId: null,
+    sourceId: null,
+    runConfig,
+    metadata: { custom: {} },
+  }) as unknown as AppendMessage;
 
 const makeConfirmationRequest = (): AdkMessage => ({
   id: "ai-1",
@@ -96,20 +151,230 @@ const approvalPart = () =>
 
 afterEach(() => {
   mocks.adapters.length = 0;
+  mocks.sendMessage.mockClear();
   mocks.messages = [];
+  mocks.messageRunConfig = undefined;
+  mocks.streamedMessages = undefined;
+  mocks.applySnapshot.mockReset();
+  mocks.threadListItem.source = null;
+  mocks.threadListItem.externalId = undefined;
 });
 
 describe("useAdkRuntime tool approvals", () => {
+  it("resumes a delayed tool result with its originating run config", async () => {
+    const runConfigA = { custom: { model: "model-a" } };
+    const runConfigB = { custom: { model: "model-b" } };
+    const { rerender } = renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    await act(async () => {
+      await latestAdapter().onNew!(makeUserMessage("first", runConfigA));
+    });
+
+    mocks.messages = [
+      { id: "u-1", type: "human", content: "first" },
+      {
+        id: "ai-1",
+        type: "ai",
+        content: [],
+        tool_calls: [
+          {
+            id: "tool-a",
+            name: "lookup",
+            args: {},
+          },
+        ],
+      },
+    ];
+    mocks.messageRunConfig = runConfigA;
+    rerender();
+
+    await act(async () => {
+      await latestAdapter().onNew!(makeUserMessage("second", runConfigB));
+      await latestAdapter().onAddToolResult!({
+        messageId: "ai-1",
+        toolCallId: "tool-a",
+        toolName: "lookup",
+        result: { value: "done" },
+        isError: false,
+      });
+    });
+
+    expect(
+      mocks.sendMessage.mock.calls.map((call) => call[1].runConfig),
+    ).toEqual([runConfigA, runConfigB, runConfigA]);
+  });
+
+  it("preserves pending tool ownership across a thread refetch", async () => {
+    const runConfig = { custom: { model: "model-a" } };
+    const laterRunConfig = { custom: { model: "model-b" } };
+    const currentMessages: AdkMessage[] = [
+      {
+        id: "ai-1",
+        type: "ai",
+        content: [],
+        tool_calls: [{ id: "tool-a", name: "lookup", args: {} }],
+      },
+    ];
+    const loadedMessages: AdkMessage[] = [
+      ...currentMessages,
+      {
+        id: "ai-2",
+        type: "ai",
+        content: [],
+        tool_calls: [
+          null as never,
+          { id: "tool-loaded", name: "lookup", args: {} },
+        ],
+      },
+    ];
+    const load = vi.fn(async () => ({ messages: loadedMessages }));
+    mocks.threadListItem.source = {};
+
+    const { rerender } = renderHook(() =>
+      useAdkRuntime({ stream: vi.fn(), load }),
+    );
+
+    await act(async () => {
+      await latestAdapter().onNew!(makeUserMessage("first", runConfig));
+    });
+    mocks.messages = currentMessages;
+    mocks.messageRunConfig = runConfig;
+    rerender();
+
+    mocks.threadListItem.externalId = "thread-a";
+    rerender();
+    await act(async () => {
+      await latestAdapter().onRefetchThread!();
+    });
+
+    await act(async () => {
+      mocks.streamedMessages = [loadedMessages.at(-1)!];
+      await latestAdapter().extras.send(
+        [{ id: "u-later", type: "human", content: "later" }],
+        { runConfig: laterRunConfig },
+      );
+      mocks.streamedMessages = undefined;
+      await latestAdapter().onAddToolResult!({
+        messageId: "ai-1",
+        toolCallId: "tool-a",
+        toolName: "lookup",
+        result: { value: "done" },
+        isError: false,
+      });
+      await latestAdapter().onAddToolResult!({
+        messageId: "ai-2",
+        toolCallId: "tool-loaded",
+        toolName: "lookup",
+        result: { value: "loaded" },
+        isError: false,
+      });
+    });
+
+    expect(
+      mocks.sendMessage.mock.calls.slice(-2).map((call) => call[1]),
+    ).toEqual([{ runConfig }, { runConfig: undefined }]);
+  });
+
+  it("tracks and prunes ownership around malformed tool-call entries", async () => {
+    const runConfig = { custom: { model: "model-a" } };
+    mocks.messages = [
+      { id: "u-1", type: "human", content: "first" },
+      {
+        id: "ai-1",
+        type: "ai",
+        content: [],
+        tool_calls: [null as never, { id: "tool-a", name: "lookup", args: {} }],
+      },
+    ];
+    mocks.messageRunConfig = runConfig;
+    renderHook(() =>
+      useAdkRuntime({
+        stream: vi.fn(),
+        getCheckpointId: vi.fn(async () => null),
+      }),
+    );
+
+    await act(async () => {
+      await latestAdapter().onEdit!({
+        ...makeUserMessage("edited", runConfig),
+        parentId: "ai-1",
+      });
+      await latestAdapter().onAddToolResult!({
+        messageId: "ai-1",
+        toolCallId: "tool-a",
+        toolName: "lookup",
+        result: { value: "done" },
+        isError: false,
+      });
+    });
+
+    expect(mocks.sendMessage.mock.calls.at(-1)![1]).toEqual({ runConfig });
+  });
+
+  it("attributes new tool calls to an explicitly configured continuation", async () => {
+    const runConfigA = { custom: { model: "model-a" } };
+    const runConfigB = { custom: { model: "model-b" } };
+    renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    await act(async () => {
+      await latestAdapter().onNew!(makeUserMessage("first", runConfigA));
+      mocks.streamedMessages = [
+        {
+          id: "ai-b",
+          type: "ai",
+          content: [],
+          tool_calls: [{ id: "tool-b", name: "lookup", args: {} }],
+        },
+      ];
+      await latestAdapter().extras.send(
+        [
+          {
+            id: "tool-result-a",
+            type: "tool",
+            tool_call_id: "tool-a",
+            name: "lookup",
+            content: JSON.stringify({ value: "first" }),
+          },
+        ],
+        { runConfig: runConfigB },
+      );
+    });
+
+    await act(async () => {
+      await latestAdapter().onAddToolResult!({
+        messageId: "ai-b",
+        toolCallId: "tool-b",
+        toolName: "lookup",
+        result: { value: "second" },
+        isError: false,
+      });
+    });
+
+    expect(mocks.sendMessage.mock.calls.at(-1)![1]).toEqual({
+      runConfig: runConfigB,
+    });
+  });
+
   it("exposes, answers, and settles the default approval seam across a rerender", async () => {
+    const runConfig = { custom: { model: "model-a" } };
+    const nextRunConfig = { custom: { model: "model-b" } };
     // Retained across the rerender: core caches converted messages by input
     // object, so only a rebuilt converter can surface the settled decision.
     const confirmationRequest = makeConfirmationRequest();
+    const { rerender } = renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+
+    await act(async () => {
+      await latestAdapter().onNew!(
+        makeUserMessage("delete the file", runConfig),
+      );
+    });
+
     mocks.messages = [
       { id: "u-1", type: "human", content: "delete the file" },
       confirmationRequest,
     ];
-
-    const { rerender } = renderHook(() => useAdkRuntime({ stream: vi.fn() }));
+    mocks.messageRunConfig = runConfig;
+    rerender();
 
     expect(latestAdapter().messages.at(-1)!.status).toMatchObject({
       type: "requires-action",
@@ -121,6 +386,9 @@ describe("useAdkRuntime tool approvals", () => {
     });
 
     await act(async () => {
+      await latestAdapter().onNew!(
+        makeUserMessage("start another run", nextRunConfig),
+      );
       await latestAdapter().onRespondToToolApproval!({
         approvalId: CONFIRMATION_CALL,
         approved: false,
@@ -135,6 +403,7 @@ describe("useAdkRuntime tool approvals", () => {
         content: JSON.stringify({ confirmed: false }),
       }),
     ]);
+    expect(mocks.sendMessage.mock.calls.at(-1)![1]).toEqual({ runConfig });
 
     mocks.messages = [
       ...mocks.messages,