@assistant-ui/react-google-adk 0.0.15 → 0.0.17

package/src/hooks.ts

@@ -1,6 +1,7 @@
-import { useAui, useAuiState } from "@assistant-ui/store";
+import { useAui } from "@assistant-ui/store";
 import { v4 as uuidv4 } from "uuid";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
+import { adkExtras } from "./adkExtras";
 import type {
   AdkMessage,
   AdkSendMessageConfig,
@@ -10,115 +11,50 @@ import type {
   AdkMessageMetadata,
 } from "./types";
 
-export const symbolAdkRuntimeExtras = Symbol("adk-runtime-extras");
-
-export type AdkRuntimeExtras = {
-  [symbolAdkRuntimeExtras]: true;
-  send: (messages: AdkMessage[], config: AdkSendMessageConfig) => Promise<void>;
-  agentInfo: { name?: string | undefined; branch?: string | undefined };
-  stateDelta: Record<string, unknown>;
-  artifactDelta: Record<string, number>;
-  longRunningToolIds: string[];
-  toolConfirmations: AdkToolConfirmation[];
-  authRequests: AdkAuthRequest[];
-  escalated: boolean;
-  messageMetadata: Map<string, AdkMessageMetadata>;
-};
-
-const asAdkRuntimeExtras = (extras: unknown): AdkRuntimeExtras => {
-  if (
-    typeof extras !== "object" ||
-    extras == null ||
-    !(symbolAdkRuntimeExtras in extras)
-  )
-    throw new Error(
-      "This method can only be called when you are using useAdkRuntime",
-    );
-
-  return extras as AdkRuntimeExtras;
-};
+const EMPTY_STATE_DELTA: Record<string, unknown> = {};
+const EMPTY_ARTIFACT_DELTA: Record<string, number> = {};
+const EMPTY_LONG_RUNNING_TOOL_IDS: string[] = [];
+const EMPTY_TOOL_CONFIRMATIONS: AdkToolConfirmation[] = [];
+const EMPTY_AUTH_REQUESTS: AdkAuthRequest[] = [];
+const EMPTY_MESSAGE_METADATA = new Map<string, AdkMessageMetadata>();
 
 /** Returns the name and branch of the currently active ADK agent. */
-export const useAdkAgentInfo = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return undefined;
-    return asAdkRuntimeExtras(extras).agentInfo;
-  });
-};
+export const useAdkAgentInfo = () =>
+  adkExtras.use((e) => e.agentInfo, undefined);
 
 /** Returns the accumulated session state delta from ADK events. */
-export const useAdkSessionState = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return {};
-    return asAdkRuntimeExtras(extras).stateDelta;
-  });
-};
+export const useAdkSessionState = () =>
+  adkExtras.use((e) => e.stateDelta, EMPTY_STATE_DELTA);
 
 /** Returns a function to send raw ADK messages. */
 export const useAdkSend = () => {
   const aui = useAui();
-  return (messages: AdkMessage[], config: AdkSendMessageConfig) => {
-    const extras = aui.thread().getState().extras;
-    const { send } = asAdkRuntimeExtras(extras);
-    return send(messages, config);
-  };
+  return (messages: AdkMessage[], config: AdkSendMessageConfig) =>
+    adkExtras.get(aui).send(messages, config);
 };
 
 /** Returns the IDs of long-running tools awaiting external input. */
-export const useAdkLongRunningToolIds = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return [];
-    return asAdkRuntimeExtras(extras).longRunningToolIds;
-  });
-};
+export const useAdkLongRunningToolIds = () =>
+  adkExtras.use((e) => e.longRunningToolIds, EMPTY_LONG_RUNNING_TOOL_IDS);
 
 /** Returns pending tool confirmation requests (from SecurityPlugin etc). */
-export const useAdkToolConfirmations = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return [];
-    return asAdkRuntimeExtras(extras).toolConfirmations;
-  });
-};
+export const useAdkToolConfirmations = () =>
+  adkExtras.use((e) => e.toolConfirmations, EMPTY_TOOL_CONFIRMATIONS);
 
 /** Returns pending auth credential requests from tools. */
-export const useAdkAuthRequests = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return [];
-    return asAdkRuntimeExtras(extras).authRequests;
-  });
-};
+export const useAdkAuthRequests = () =>
+  adkExtras.use((e) => e.authRequests, EMPTY_AUTH_REQUESTS);
 
 /** Returns the accumulated artifact delta (filename → version). */
-export const useAdkArtifacts = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return {};
-    return asAdkRuntimeExtras(extras).artifactDelta;
-  });
-};
+export const useAdkArtifacts = () =>
+  adkExtras.use((e) => e.artifactDelta, EMPTY_ARTIFACT_DELTA);
 
 /** Returns whether any agent has escalated (requested human handoff). */
-export const useAdkEscalation = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return false;
-    return asAdkRuntimeExtras(extras).escalated;
-  });
-};
+export const useAdkEscalation = () => adkExtras.use((e) => e.escalated, false);
 
 /** Returns per-message metadata (grounding, citation, usage). Keyed by message ID. */
-export const useAdkMessageMetadata = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return new Map<string, AdkMessageMetadata>();
-    return asAdkRuntimeExtras(extras).messageMetadata;
-  });
-};
+export const useAdkMessageMetadata = () =>
+  adkExtras.use((e) => e.messageMetadata, EMPTY_MESSAGE_METADATA);
 
 // ── Convenience helpers for interactive flows ──
 
@@ -129,10 +65,8 @@ export const useAdkConfirmTool = () => {
     toolCallId: string,
     confirmed: boolean,
     payload?: ReadonlyJSONValue,
-  ) => {
-    const extras = aui.thread().getState().extras;
-    const { send } = asAdkRuntimeExtras(extras);
-    return send(
+  ) =>
+    adkExtras.get(aui).send(
       [
         {
           id: uuidv4(),
@@ -148,16 +82,13 @@ export const useAdkConfirmTool = () => {
       ],
       {},
     );
-  };
 };
 
 /** Returns a function to submit auth credentials for a pending auth request. */
 export const useAdkSubmitAuth = () => {
   const aui = useAui();
-  return (toolCallId: string, credential: AdkAuthCredential) => {
-    const extras = aui.thread().getState().extras;
-    const { send } = asAdkRuntimeExtras(extras);
-    return send(
+  return (toolCallId: string, credential: AdkAuthCredential) =>
+    adkExtras.get(aui).send(
       [
         {
           id: uuidv4(),
@@ -170,16 +101,13 @@ export const useAdkSubmitAuth = () => {
       ],
       {},
     );
-  };
 };
 
 /** Returns a function to submit the user's answer for a pending `adk_request_input` HITL interrupt. */
 export const useAdkSubmitInput = () => {
   const aui = useAui();
-  return (toolCallId: string, result: ReadonlyJSONValue) => {
-    const extras = aui.thread().getState().extras;
-    const { send } = asAdkRuntimeExtras(extras);
-    return send(
+  return (toolCallId: string, result: ReadonlyJSONValue) =>
+    adkExtras.get(aui).send(
       [
         {
           id: uuidv4(),
@@ -192,7 +120,6 @@ export const useAdkSubmitInput = () => {
       ],
       {},
     );
-  };
 };
 
 // ── State prefix helpers ──
@@ -215,28 +142,22 @@ const filterByPrefix = (
 };
 
 /** Returns app-level state (keys prefixed with `app:`, prefix stripped). */
-export const useAdkAppState = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return {};
-    return filterByPrefix(asAdkRuntimeExtras(extras).stateDelta, APP_PREFIX);
-  });
-};
+export const useAdkAppState = () =>
+  adkExtras.use(
+    (e) => filterByPrefix(e.stateDelta, APP_PREFIX),
+    EMPTY_STATE_DELTA,
+  );
 
 /** Returns user-level state (keys prefixed with `user:`, prefix stripped). */
-export const useAdkUserState = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return {};
-    return filterByPrefix(asAdkRuntimeExtras(extras).stateDelta, USER_PREFIX);
-  });
-};
+export const useAdkUserState = () =>
+  adkExtras.use(
+    (e) => filterByPrefix(e.stateDelta, USER_PREFIX),
+    EMPTY_STATE_DELTA,
+  );
 
 /** Returns temp state (keys prefixed with `temp:`, prefix stripped). Not persisted. */
-export const useAdkTempState = () => {
-  return useAuiState((s) => {
-    const extras = s.thread.extras;
-    if (!extras) return {};
-    return filterByPrefix(asAdkRuntimeExtras(extras).stateDelta, TEMP_PREFIX);
-  });
-};
+export const useAdkTempState = () =>
+  adkExtras.use(
+    (e) => filterByPrefix(e.stateDelta, TEMP_PREFIX),
+    EMPTY_STATE_DELTA,
+  );

package/src/server/parseAdkRequest.test.ts

@@ -85,12 +85,14 @@ describe("parseAdkRequest", () => {
       method: "POST",
       body: "not json",
     });
-    await expect(parseAdkRequest(req)).rejects.toThrow();
+    await expect(parseAdkRequest(req)).rejects.toThrow(
+      'Invalid JSON in Google ADK proxy request body. Expected a JSON object like {"message":"Hello"} or {"type":"tool-result",...}.',
+    );
   });
 
   it("throws on non-object body", async () => {
     await expect(parseAdkRequest(makeRequest([1, 2, 3]))).rejects.toThrow(
-      "Request body must be a JSON object",
+      "Google ADK proxy request body must be a JSON object",
     );
   });
 });

package/src/server/parseAdkRequest.ts

@@ -46,11 +46,13 @@ export const parseAdkRequest = async (
   try {
     body = (await request.json()) as Record<string, unknown>;
   } catch {
-    throw new Error("Invalid JSON in request body");
+    throw new Error(
+      'Invalid JSON in Google ADK proxy request body. Expected a JSON object like {"message":"Hello"} or {"type":"tool-result",...}.',
+    );
   }
 
   if (!body || typeof body !== "object" || Array.isArray(body)) {
-    throw new Error("Request body must be a JSON object");
+    throw new Error("Google ADK proxy request body must be a JSON object");
   }
 
   const config: AdkSendMessageConfig = {};

package/src/trimTrailingSlashes.ts

@@ -0,0 +1,5 @@
+export const trimTrailingSlashes = (value: string) => {
+  let end = value.length;
+  while (end > 0 && value[end - 1] === "/") end -= 1;
+  return value.slice(0, end);
+};

package/src/types.ts

@@ -274,3 +274,16 @@ export type OnAdkCustomEventCallback = (
 export type OnAdkAgentTransferCallback = (
   toAgent: string,
 ) => void | Promise<void>;
+
+/** Private state and actions `useAdkRuntime` exposes through `thread.extras`. */
+export type AdkRuntimeExtras = {
+  send: (messages: AdkMessage[], config: AdkSendMessageConfig) => Promise<void>;
+  agentInfo: { name?: string | undefined; branch?: string | undefined };
+  stateDelta: Record<string, unknown>;
+  artifactDelta: Record<string, number>;
+  longRunningToolIds: string[];
+  toolConfirmations: AdkToolConfirmation[];
+  authRequests: AdkAuthRequest[];
+  escalated: boolean;
+  messageMetadata: Map<string, AdkMessageMetadata>;
+};

package/src/useAdkRuntime.ts

@@ -11,6 +11,7 @@ import {
   type AppendMessage,
   type ThreadMessage,
   type ToolExecutionStatus,
+  generateId,
 } from "@assistant-ui/core";
 import {
   useCloudThreadListAdapter,
@@ -31,7 +32,7 @@ import type {
 } from "./types";
 import { useAdkMessages } from "./useAdkMessages";
 import { convertAdkMessage } from "./convertAdkMessages";
-import { symbolAdkRuntimeExtras, type AdkRuntimeExtras } from "./hooks";
+import { adkExtras } from "./adkExtras";
 import { v4 as uuidv4 } from "uuid";
 
 /** @internal — exported for unit tests. */
@@ -138,8 +139,26 @@ const truncateAdkMessages = (
   return truncated;
 };
 
+const toAdkUserMessage = (
+  msg: AppendMessage,
+  id = generateId(),
+): AdkMessage & { type: "human"; id: string } => ({
+  id,
+  type: "human",
+  content: getMessageContent(msg),
+});
+
 export type UseAdkRuntimeOptions = ExternalStoreSharedOptions & {
   stream: AdkStreamCallback;
+  /**
+   * Called whenever the active thread's canonical (remote) ID changes, so the
+   * value can be treated as a managed/controlled variable (e.g. synced to a URL
+   * query param). Only the settled remote ID is emitted: while a freshly created
+   * thread is still optimistic the value is `undefined`, and the real ID is
+   * emitted once the thread is initialized; the transient local ID is never
+   * surfaced.
+   */
+  onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
   autoCancelPendingToolCalls?: boolean | undefined;
   unstable_allowCancellation?: boolean | undefined;
   getCheckpointId?: (
@@ -196,6 +215,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     messageMetadata,
     sendMessage,
     cancel,
+    setMessages,
     replaceMessages,
   } = useAdkMessages({
     stream,
@@ -232,6 +252,50 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
   const threadMessagesRef = useRef(threadMessages);
   threadMessagesRef.current = threadMessages;
 
+  const adkMessagesRef = useRef(messages);
+  adkMessagesRef.current = messages;
+
+  const stagedMessagesRef = useRef(
+    new Map<
+      string,
+      {
+        message: AdkMessage & { id: string };
+        runConfig: AppendMessage["runConfig"];
+      }
+    >(),
+  );
+  const [stagedMessageCount, setStagedMessageCount] = useState(0);
+  const hasStagedMessages = stagedMessageCount > 0;
+
+  const getStagedRun = (parentId: string | null) => {
+    if (!parentId || !stagedMessagesRef.current.has(parentId)) return null;
+
+    const staged: AdkMessage[] = [];
+    for (const message of adkMessagesRef.current) {
+      if (message.id && stagedMessagesRef.current.has(message.id)) {
+        staged.push(stagedMessagesRef.current.get(message.id)!.message);
+      }
+      if (message.id === parentId) break;
+    }
+
+    return {
+      messages: staged,
+      runConfig: stagedMessagesRef.current.get(parentId)!.runConfig,
+    };
+  };
+
+  const stageUserMessage = (msg: AppendMessage) => {
+    const stagedMessage = toAdkUserMessage(msg);
+    stagedMessagesRef.current.set(stagedMessage.id, {
+      message: stagedMessage,
+      runConfig: msg.runConfig,
+    });
+    setStagedMessageCount(stagedMessagesRef.current.size);
+    const nextMessages = [...adkMessagesRef.current, stagedMessage];
+    adkMessagesRef.current = nextMessages;
+    setMessages(nextMessages);
+  };
+
   const runtime = useExternalStoreRuntime({
     ...pickExternalStoreSharedOptions(options),
     isRunning: effectiveIsRunning,
@@ -239,8 +303,7 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
     unstable_enableToolInvocations: true,
     setToolStatuses,
     adapters: { attachments, dictation, feedback, speech, voice },
-    extras: {
-      [symbolAdkRuntimeExtras]: true,
+    extras: adkExtras.provide({
       agentInfo,
       stateDelta,
       artifactDelta,
@@ -250,8 +313,13 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
       escalated,
       messageMetadata,
       send: handleSendMessage,
-    } satisfies AdkRuntimeExtras,
+    }),
     onNew: async (msg) => {
+      if (!(msg.startRun ?? msg.role === "user")) {
+        stageUserMessage(msg);
+        return;
+      }
+
       const cancellations =
         autoCancelPendingToolCalls !== false
           ? getPendingCancellations(messages, longRunningToolIds)
@@ -276,6 +344,18 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
             msg.parentId,
           );
           replaceMessages(truncated);
+          if (!(msg.startRun ?? msg.role === "user")) {
+            const stagedMessage = toAdkUserMessage(msg);
+            stagedMessagesRef.current.set(stagedMessage.id, {
+              message: stagedMessage,
+              runConfig: msg.runConfig,
+            });
+            setStagedMessageCount(stagedMessagesRef.current.size);
+            const nextMessages = [...truncated, stagedMessage];
+            adkMessagesRef.current = nextMessages;
+            setMessages(nextMessages);
+            return;
+          }
           const externalId = aui.threadListItem().getState().externalId;
           const checkpointId = externalId
             ? await getCheckpointId(externalId, truncated)
@@ -295,23 +375,39 @@ const useAdkRuntimeImpl = (options: UseAdkRuntimeOptions) => {
           );
         }
       : undefined,
-    onReload: getCheckpointId
-      ? async (parentId, config) => {
-          const truncated = truncateAdkMessages(
-            threadMessagesRef.current,
-            parentId,
-          );
-          replaceMessages(truncated);
-          const externalId = aui.threadListItem().getState().externalId;
-          const checkpointId = externalId
-            ? await getCheckpointId(externalId, truncated)
-            : null;
-          return handleSendMessage([], {
-            runConfig: config.runConfig,
-            ...(checkpointId && { checkpointId }),
-          });
+    ...(getCheckpointId || hasStagedMessages
+      ? {
+          onReload: async (parentId, config) => {
+            const stagedRun = getStagedRun(parentId);
+            if (stagedRun) {
+              for (const message of stagedRun.messages) {
+                stagedMessagesRef.current.delete(message.id);
+              }
+              setStagedMessageCount(stagedMessagesRef.current.size);
+              return handleSendMessage(stagedRun.messages, {
+                runConfig: config.runConfig ?? stagedRun.runConfig,
+              });
+            }
+
+            if (!getCheckpointId)
+              throw new Error("Runtime does not support reloading messages.");
+
+            const truncated = truncateAdkMessages(
+              threadMessagesRef.current,
+              parentId,
+            );
+            replaceMessages(truncated);
+            const externalId = aui.threadListItem().getState().externalId;
+            const checkpointId = externalId
+              ? await getCheckpointId(externalId, truncated)
+              : null;
+            return handleSendMessage([], {
+              runConfig: config.runConfig,
+              ...(checkpointId && { checkpointId }),
+            });
+          },
         }
-      : undefined,
+      : {}),
     onAddToolResult: async ({
       toolCallId,
       toolName,
@@ -373,6 +469,7 @@ export const useAdkRuntime = ({
   sessionAdapter,
   create,
   delete: deleteFn,
+  onThreadIdChange,
   ...options
 }: UseAdkRuntimeOptions) => {
   const aui = useAui();
@@ -394,5 +491,6 @@ export const useAdkRuntime = ({
     },
     adapter,
     allowNesting: true,
+    onThreadIdChange,
   });
 };