@assistant-ui/react-google-adk 0.0.4 → 0.0.6

package/src/useAdkMessages.ts

@@ -64,6 +64,22 @@ export const useAdkMessages = ({
     _setMessages(msgs);
   }, []);
 
+  // Replace the message list AND reset derived per-turn HITL state.
+  // Used by truncation paths (edit, reload, load) so that stale interrupt
+  // markers and per-message metadata from the removed messages don't leak
+  // into the next turn.
+  const replaceMessages = useCallback(
+    (msgs: AdkMessage[]) => {
+      setMessagesImmediate(msgs);
+      setLongRunningToolIds([]);
+      setToolConfirmations([]);
+      setAuthRequests([]);
+      setEscalated(false);
+      setMessageMetadata(new Map());
+    },
+    [setMessagesImmediate],
+  );
+
   const abortControllerRef = useRef<AbortController | null>(null);
 
   const { onError, onCustomEvent, onAgentTransfer } = useMemo(
@@ -180,6 +196,7 @@ export const useAdkMessages = ({
     sendMessage,
     cancel,
     setMessages: setMessagesImmediate,
+    replaceMessages,
   };
 };
 

package/src/useAdkRuntime.test.ts

@@ -0,0 +1,138 @@
+import { describe, it, expect } from "vitest";
+import { getPendingCancellations, getPendingToolCalls } from "./useAdkRuntime";
+import type { AdkMessage } from "./types";
+
+const aiWithToolCalls = (
+  id: string,
+  toolCalls: Array<{ id: string; name: string }>,
+): AdkMessage => ({
+  id,
+  type: "ai",
+  content: [],
+  tool_calls: toolCalls.map((tc) => ({
+    id: tc.id,
+    name: tc.name,
+    args: {},
+    argsText: "{}",
+  })),
+});
+
+const toolResponse = (
+  id: string,
+  toolCallId: string,
+  name: string,
+): AdkMessage => ({
+  id,
+  type: "tool",
+  tool_call_id: toolCallId,
+  name,
+  content: JSON.stringify({ ok: true }),
+  status: "success",
+});
+
+describe("getPendingToolCalls", () => {
+  it("returns tool calls without matching tool responses", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [
+        { id: "tc-1", name: "tool_a" },
+        { id: "tc-2", name: "tool_b" },
+      ]),
+    ];
+    expect(getPendingToolCalls(messages)).toEqual([
+      { id: "tc-1", name: "tool_a", args: {}, argsText: "{}" },
+      { id: "tc-2", name: "tool_b", args: {}, argsText: "{}" },
+    ]);
+  });
+
+  it("excludes tool calls that have a matching tool response", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [
+        { id: "tc-1", name: "tool_a" },
+        { id: "tc-2", name: "tool_b" },
+      ]),
+      toolResponse("t-1", "tc-1", "tool_a"),
+    ];
+    expect(getPendingToolCalls(messages)).toEqual([
+      { id: "tc-2", name: "tool_b", args: {}, argsText: "{}" },
+    ]);
+  });
+
+  it("returns empty for threads with no ai messages", () => {
+    expect(getPendingToolCalls([])).toEqual([]);
+  });
+});
+
+describe("getPendingCancellations", () => {
+  it("emits a {cancelled:true} tool message for every pending tool call", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [{ id: "tc-1", name: "tool_a" }]),
+    ];
+    const result = getPendingCancellations(messages, []);
+    expect(result).toHaveLength(1);
+    expect(result[0]).toMatchObject({
+      type: "tool",
+      name: "tool_a",
+      tool_call_id: "tc-1",
+      content: JSON.stringify({ cancelled: true }),
+      status: "error",
+    });
+  });
+
+  it("skips tool calls whose id is in longRunningToolIds", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [
+        { id: "lrt-1", name: "adk_request_input" },
+        { id: "tc-2", name: "regular_tool" },
+      ]),
+    ];
+    const result = getPendingCancellations(messages, ["lrt-1"]);
+    expect(result).toHaveLength(1);
+    expect(result[0]).toMatchObject({
+      type: "tool",
+      name: "regular_tool",
+      tool_call_id: "tc-2",
+    });
+  });
+
+  it("skips all three HITL interrupt types", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [
+        { id: "lrt-1", name: "adk_request_input" },
+        { id: "lrt-2", name: "adk_request_confirmation" },
+        { id: "lrt-3", name: "adk_request_credential" },
+      ]),
+    ];
+    const result = getPendingCancellations(messages, [
+      "lrt-1",
+      "lrt-2",
+      "lrt-3",
+    ]);
+    expect(result).toEqual([]);
+  });
+
+  it("returns empty when no tool calls are pending", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [{ id: "tc-1", name: "tool_a" }]),
+      toolResponse("t-1", "tc-1", "tool_a"),
+    ];
+    expect(getPendingCancellations(messages, [])).toEqual([]);
+  });
+
+  it("cancels regular tool calls even when HITL interrupts are resolved", () => {
+    const messages: AdkMessage[] = [
+      aiWithToolCalls("ai-1", [
+        { id: "lrt-1", name: "adk_request_input" },
+        { id: "tc-2", name: "regular_tool" },
+      ]),
+      toolResponse("t-1", "lrt-1", "adk_request_input"),
+    ];
+    // lrt-1 lingering in longRunningToolIds is harmless because
+    // getPendingToolCalls already excluded it via the tool response.
+    const result = getPendingCancellations(messages, ["lrt-1"]);
+    expect(result).toHaveLength(1);
+    expect(result[0]).toMatchObject({
+      tool_call_id: "tc-2",
+      name: "regular_tool",
+    });
+  });
+});

package/src/useAdkRuntime.ts

@@ -69,7 +69,8 @@ const getMessageContent = (msg: AppendMessage) => {
   return content;
 };
 
-const getPendingToolCalls = (messages: AdkMessage[]) => {
+/** @internal — exported for unit tests. */
+export const getPendingToolCalls = (messages: AdkMessage[]) => {
   const pending = new Map<string, { id: string; name: string }>();
   for (const msg of messages) {
     if (msg.type === "ai" && msg.tool_calls) {
@@ -84,6 +85,35 @@ const getPendingToolCalls = (messages: AdkMessage[]) => {
   return [...pending.values()];
 };
 
+/**
+ * @internal — exported for unit tests.
+ *
+ * Returns `{cancelled: true}` tool responses for pending tool calls when the
+ * user sends a new turn, EXCEPT for HITL interrupts marked via
+ * `long_running_tool_ids` (`adk_request_input`, `adk_request_confirmation`,
+ * `adk_request_credential`). Those must be answered through a dedicated tool
+ * UI + submit helper, not auto-cancelled.
+ */
+export const getPendingCancellations = (
+  messages: AdkMessage[],
+  longRunningToolIds: readonly string[],
+): Array<AdkMessage & { type: "tool" }> => {
+  const longRunningSet = new Set(longRunningToolIds);
+  return getPendingToolCalls(messages)
+    .filter((t) => !longRunningSet.has(t.id))
+    .map(
+      (t) =>
+        ({
+          id: uuidv4(),
+          type: "tool",
+          name: t.name,
+          tool_call_id: t.id,
+          content: JSON.stringify({ cancelled: true }),
+          status: "error",
+        }) satisfies AdkMessage & { type: "tool" },
+    );
+};
+
 const truncateAdkMessages = (
   threadMessages: readonly ThreadMessage[],
   parentId: string | null,
@@ -153,7 +183,7 @@ const useAdkRuntimeImpl = ({
     messageMetadata,
     sendMessage,
     cancel,
-    setMessages,
+    replaceMessages,
   } = useAdkMessages({
     stream,
     ...(eventHandlers && { eventHandlers }),
@@ -239,17 +269,7 @@ const useAdkRuntimeImpl = ({
 
       const cancellations =
         autoCancelPendingToolCalls !== false
-          ? getPendingToolCalls(messages).map(
-              (t) =>
-                ({
-                  id: uuidv4(),
-                  type: "tool",
-                  name: t.name,
-                  tool_call_id: t.id,
-                  content: JSON.stringify({ cancelled: true }),
-                  status: "error",
-                }) satisfies AdkMessage & { type: "tool" },
-            )
+          ? getPendingCancellations(messages, longRunningToolIds)
           : [];
 
       return handleSendMessage(
@@ -271,7 +291,7 @@ const useAdkRuntimeImpl = ({
             threadMessagesRef.current,
             msg.parentId,
           );
-          setMessages(truncated);
+          replaceMessages(truncated);
           const externalId = aui.threadListItem().getState().externalId;
           const checkpointId = externalId
             ? await getCheckpointId(externalId, truncated)
@@ -298,7 +318,7 @@ const useAdkRuntimeImpl = ({
             threadMessagesRef.current,
             parentId,
           );
-          setMessages(truncated);
+          replaceMessages(truncated);
           const externalId = aui.threadListItem().getState().externalId;
           const checkpointId = externalId
             ? await getCheckpointId(externalId, truncated)
@@ -356,13 +376,13 @@ const useAdkRuntimeImpl = ({
 
       loadFn(externalId).then(
         ({ messages: msgs }) => {
-          setMessages(msgs);
+          replaceMessages(msgs);
         },
         (e) => {
           console.warn("Failed to load ADK session:", e);
         },
       );
-    }, [aui, setMessages]);
+    }, [aui, replaceMessages]);
   }
 
   return runtime;