@assistant-ui/react-google-adk 0.0.1

package/src/useAdkMessages.ts

@@ -0,0 +1,259 @@
+/// <reference types="@assistant-ui/core/store" />
+import { useState, useCallback, useRef, useMemo } from "react";
+import { v4 as uuidv4 } from "uuid";
+import { useAui } from "@assistant-ui/store";
+import { AdkEventAccumulator } from "./AdkEventAccumulator";
+import type {
+  AdkEvent,
+  AdkEventPart,
+  AdkMessage,
+  AdkMessageContentPart,
+  AdkMessageMetadata,
+  AdkSendMessageConfig,
+  AdkStreamCallback,
+  AdkToolConfirmation,
+  AdkAuthRequest,
+  OnAdkErrorCallback,
+  OnAdkCustomEventCallback,
+  OnAdkAgentTransferCallback,
+} from "./types";
+
+export type UseAdkMessagesOptions = {
+  stream: AdkStreamCallback;
+  eventHandlers?: {
+    onError?: OnAdkErrorCallback;
+    onCustomEvent?: OnAdkCustomEventCallback;
+    onAgentTransfer?: OnAdkAgentTransferCallback;
+  };
+};
+
+export const useAdkMessages = ({
+  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>>(
+    {},
+  );
+  const [toolConfirmations, setToolConfirmations] = useState<
+    AdkToolConfirmation[]
+  >([]);
+  const [authRequests, setAuthRequests] = useState<AdkAuthRequest[]>([]);
+  const [escalated, setEscalated] = useState(false);
+  const [messageMetadata, setMessageMetadata] = useState<
+    Map<string, AdkMessageMetadata>
+  >(new Map());
+  const lastTransferToAgentRef = useRef<string | undefined>(undefined);
+  const messagesRef = useRef(messages);
+  messagesRef.current = messages;
+  const stateDeltaRef = useRef(stateDelta);
+  stateDeltaRef.current = stateDelta;
+  const artifactDeltaRef = useRef(artifactDelta);
+  artifactDeltaRef.current = artifactDelta;
+  const messageMetadataRef = useRef(messageMetadata);
+  messageMetadataRef.current = messageMetadata;
+
+  const setMessagesImmediate = useCallback((msgs: AdkMessage[]) => {
+    messagesRef.current = msgs;
+    _setMessages(msgs);
+  }, []);
+
+  const abortControllerRef = useRef<AbortController | null>(null);
+
+  const { onError, onCustomEvent, onAgentTransfer } = useMemo(
+    () => eventHandlers ?? {},
+    [eventHandlers],
+  );
+
+  const aui = useAui();
+  const sendMessage = useCallback(
+    async (newMessages: AdkMessage[], config: AdkSendMessageConfig) => {
+      const newMessagesWithId = newMessages.map((m) =>
+        m.id ? m : { ...m, id: uuidv4() },
+      ) as AdkMessage[];
+
+      const accumulator = new AdkEventAccumulator(messagesRef.current);
+      for (const msg of newMessagesWithId) {
+        accumulator.processEvent(messageToEvent(msg));
+      }
+      setMessagesImmediate(accumulator.getMessages());
+
+      const abortController = new AbortController();
+      abortControllerRef.current = abortController;
+
+      try {
+        const response = await stream(newMessagesWithId, {
+          ...config,
+          abortSignal: abortController.signal,
+          initialize: async () => {
+            return await aui.threadListItem().initialize();
+          },
+        });
+
+        for await (const event of response) {
+          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]),
+              );
+            }
+          }
+
+          const transfer = accumulator.getLastTransferToAgent();
+          if (transfer && transfer !== lastTransferToAgentRef.current) {
+            lastTransferToAgentRef.current = transfer;
+            onAgentTransfer?.(transfer);
+          }
+
+          // Fire custom event callback for events with customMetadata
+          if (event.customMetadata && onCustomEvent) {
+            for (const [key, value] of Object.entries(event.customMetadata)) {
+              onCustomEvent(key, value);
+            }
+          }
+
+          if (event.errorCode || event.errorMessage) {
+            onError?.(event.errorMessage ?? event.errorCode);
+          }
+        }
+      } catch (error) {
+        if (
+          !abortController.signal.aborted &&
+          !(error instanceof Error && error.name === "AbortError")
+        ) {
+          throw error;
+        }
+      } finally {
+        if (abortControllerRef.current === abortController) {
+          abortControllerRef.current = null;
+        }
+      }
+    },
+    [
+      aui,
+      setMessagesImmediate,
+      stream,
+      onError,
+      onCustomEvent,
+      onAgentTransfer,
+    ],
+  );
+
+  const cancel = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+    }
+  }, []);
+
+  return {
+    messages,
+    stateDelta,
+    agentInfo,
+    longRunningToolIds,
+    artifactDelta,
+    toolConfirmations,
+    authRequests,
+    escalated,
+    messageMetadata,
+    sendMessage,
+    cancel,
+    setMessages: setMessagesImmediate,
+  };
+};
+
+const messageToEvent = (msg: AdkMessage): AdkEvent => {
+  if (msg.type === "human") {
+    return {
+      id: msg.id ?? uuidv4(),
+      author: "user",
+      content: { role: "user", parts: contentToParts(msg.content) },
+    };
+  }
+
+  if (msg.type === "tool") {
+    let response: unknown;
+    try {
+      response = JSON.parse(msg.content);
+    } catch {
+      response = msg.content;
+    }
+    return {
+      id: msg.id ?? uuidv4(),
+      content: {
+        role: "user",
+        parts: [
+          {
+            functionResponse: {
+              name: msg.name,
+              id: msg.tool_call_id,
+              response,
+            },
+          },
+        ],
+      },
+    };
+  }
+
+  const result: AdkEvent = { id: msg.id ?? uuidv4() };
+  if (msg.author != null) result.author = msg.author;
+  result.content = {
+    role: "model",
+    parts: [
+      ...contentToParts(msg.content),
+      ...(msg.tool_calls?.map((tc) => ({
+        functionCall: { name: tc.name, id: tc.id, args: { ...tc.args } },
+      })) ?? []),
+    ],
+  };
+  return result;
+};
+
+const contentToParts = (
+  content: string | AdkMessageContentPart[],
+): AdkEventPart[] => {
+  if (typeof content === "string") return [{ text: content }];
+  return content.map((part) => {
+    switch (part.type) {
+      case "text":
+        return { text: part.text };
+      case "reasoning":
+        return { text: part.text, thought: true };
+      case "image":
+        return { inlineData: { mimeType: part.mimeType, data: part.data } };
+      case "image_url":
+        return { fileData: { fileUri: part.url } };
+      case "code":
+        return {
+          executableCode: { code: part.code, language: part.language },
+        };
+      case "code_result":
+        return {
+          codeExecutionResult: { output: part.output, outcome: part.outcome },
+        };
+      default:
+        return { text: "" };
+    }
+  });
+};

package/src/useAdkRuntime.ts

@@ -0,0 +1,398 @@
+/// <reference types="@assistant-ui/core/store" />
+import { useEffect, useRef, useState } from "react";
+import {
+  getExternalStoreMessages,
+  type AttachmentAdapter,
+  type FeedbackAdapter,
+  type SpeechSynthesisAdapter,
+  type AppendMessage,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+import {
+  type ToolExecutionStatus,
+  useCloudThreadListAdapter,
+  useRemoteThreadListRuntime,
+  useExternalMessageConverter,
+  useExternalStoreRuntime,
+  useToolInvocations,
+} from "@assistant-ui/core/react";
+import { useAui } from "@assistant-ui/store";
+import type { AssistantCloud } from "assistant-cloud";
+import type { RemoteThreadListAdapter } from "@assistant-ui/core";
+import type {
+  AdkMessage,
+  AdkSendMessageConfig,
+  AdkStreamCallback,
+  OnAdkErrorCallback,
+  OnAdkCustomEventCallback,
+  OnAdkAgentTransferCallback,
+} from "./types";
+import { useAdkMessages } from "./useAdkMessages";
+import { convertAdkMessage } from "./convertAdkMessages";
+import { symbolAdkRuntimeExtras, type AdkRuntimeExtras } from "./hooks";
+import { v4 as uuidv4 } from "uuid";
+
+const getMessageContent = (msg: AppendMessage) => {
+  const allContent = [
+    ...msg.content,
+    ...(msg.attachments?.flatMap((a) => a.content) ?? []),
+  ];
+  const content = allContent.map((part) => {
+    const type = part.type;
+    switch (type) {
+      case "text":
+        return { type: "text" as const, text: part.text };
+      case "image":
+        return { type: "image_url" as const, url: part.image };
+      case "file":
+        return {
+          type: "text" as const,
+          text: `[File: ${part.filename ?? "file"}]`,
+        };
+
+      case "tool-call":
+        throw new Error("Tool call appends are not supported.");
+
+      default:
+        const _exhaustiveCheck: "reasoning" | "source" | "audio" | "data" =
+          type;
+        throw new Error(
+          `Unsupported append message part type: ${_exhaustiveCheck}`,
+        );
+    }
+  });
+
+  if (content.length === 1 && content[0]?.type === "text") {
+    return content[0].text ?? "";
+  }
+
+  return content;
+};
+
+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) {
+      for (const tc of msg.tool_calls) {
+        pending.set(tc.id, tc);
+      }
+    }
+    if (msg.type === "tool") {
+      pending.delete(msg.tool_call_id);
+    }
+  }
+  return [...pending.values()];
+};
+
+const truncateAdkMessages = (
+  threadMessages: readonly ThreadMessage[],
+  parentId: string | null,
+): AdkMessage[] => {
+  if (parentId === null) return [];
+  const parentIndex = threadMessages.findIndex((m) => m.id === parentId);
+  if (parentIndex === -1) return [];
+  const truncated: AdkMessage[] = [];
+  for (let i = 0; i <= parentIndex && i < threadMessages.length; i++) {
+    truncated.push(...getExternalStoreMessages<AdkMessage>(threadMessages[i]!));
+  }
+  return truncated;
+};
+
+export type UseAdkRuntimeOptions = {
+  stream: AdkStreamCallback;
+  autoCancelPendingToolCalls?: boolean | undefined;
+  unstable_allowCancellation?: boolean | undefined;
+  getCheckpointId?: (
+    threadId: string,
+    parentMessages: AdkMessage[],
+  ) => Promise<string | null>;
+  load?: (threadId: string) => Promise<{ messages: AdkMessage[] }>;
+  create?: () => Promise<{ externalId: string }>;
+  delete?: (threadId: string) => Promise<void>;
+  adapters?:
+    | {
+        attachments?: AttachmentAdapter;
+        speech?: SpeechSynthesisAdapter;
+        feedback?: FeedbackAdapter;
+      }
+    | undefined;
+  eventHandlers?:
+    | {
+        onError?: OnAdkErrorCallback;
+        onCustomEvent?: OnAdkCustomEventCallback;
+        onAgentTransfer?: OnAdkAgentTransferCallback;
+      }
+    | undefined;
+  cloud?: AssistantCloud | undefined;
+  /**
+   * A `RemoteThreadListAdapter` to use instead of the cloud adapter.
+   * Use with `createAdkSessionAdapter` for ADK session-backed persistence.
+   */
+  sessionAdapter?: RemoteThreadListAdapter | undefined;
+};
+
+const useAdkRuntimeImpl = ({
+  autoCancelPendingToolCalls,
+  adapters: { attachments, feedback, speech } = {},
+  unstable_allowCancellation,
+  stream,
+  load,
+  getCheckpointId,
+  eventHandlers,
+}: UseAdkRuntimeOptions) => {
+  const aui = useAui();
+  const {
+    messages,
+    stateDelta,
+    agentInfo,
+    longRunningToolIds,
+    artifactDelta,
+    toolConfirmations,
+    authRequests,
+    escalated,
+    messageMetadata,
+    sendMessage,
+    cancel,
+    setMessages,
+  } = useAdkMessages({
+    stream,
+    ...(eventHandlers && { eventHandlers }),
+  });
+
+  const [isRunning, setIsRunning] = useState(false);
+  const [toolStatuses, setToolStatuses] = useState<
+    Record<string, ToolExecutionStatus>
+  >({});
+  const hasExecutingTools = Object.values(toolStatuses).some(
+    (s) => s?.type === "executing",
+  );
+  const effectiveIsRunning = isRunning || hasExecutingTools;
+
+  const handleSendMessage = async (
+    msgs: AdkMessage[],
+    config: AdkSendMessageConfig,
+  ) => {
+    try {
+      setIsRunning(true);
+      await sendMessage(msgs, config);
+    } finally {
+      setIsRunning(false);
+    }
+  };
+
+  const threadMessages = useExternalMessageConverter({
+    callback: convertAdkMessage,
+    messages,
+    isRunning: effectiveIsRunning,
+  });
+
+  const threadMessagesRef = useRef(threadMessages);
+  threadMessagesRef.current = threadMessages;
+
+  const [runtimeRef] = useState(() => ({
+    get current() {
+      return runtime;
+    },
+  }));
+
+  const toolInvocations = useToolInvocations({
+    state: { messages: threadMessages, isRunning: effectiveIsRunning },
+    getTools: () => runtimeRef.current.thread.getModelContext().tools,
+    onResult: (command) => {
+      if (command.type === "add-tool-result") {
+        void handleSendMessage(
+          [
+            {
+              id: uuidv4(),
+              type: "tool",
+              name: command.toolName,
+              tool_call_id: command.toolCallId,
+              content: JSON.stringify(command.result),
+              status: command.isError ? "error" : "success",
+            },
+          ],
+          {},
+        );
+      }
+    },
+    setToolStatuses,
+  });
+
+  const runtime = useExternalStoreRuntime({
+    isRunning: effectiveIsRunning,
+    messages: threadMessages,
+    adapters: { attachments, feedback, speech },
+    extras: {
+      [symbolAdkRuntimeExtras]: true,
+      agentInfo,
+      stateDelta,
+      artifactDelta,
+      longRunningToolIds,
+      toolConfirmations,
+      authRequests,
+      escalated,
+      messageMetadata,
+      send: handleSendMessage,
+    } satisfies AdkRuntimeExtras,
+    onNew: async (msg) => {
+      await toolInvocations.abort();
+
+      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" },
+            )
+          : [];
+
+      return handleSendMessage(
+        [
+          ...cancellations,
+          {
+            id: uuidv4(),
+            type: "human",
+            content: getMessageContent(msg),
+          },
+        ],
+        { runConfig: msg.runConfig },
+      );
+    },
+    onEdit: getCheckpointId
+      ? async (msg) => {
+          await toolInvocations.abort();
+          const truncated = truncateAdkMessages(
+            threadMessagesRef.current,
+            msg.parentId,
+          );
+          setMessages(truncated);
+          const externalId = aui.threadListItem().getState().externalId;
+          const checkpointId = externalId
+            ? await getCheckpointId(externalId, truncated)
+            : null;
+          return handleSendMessage(
+            [
+              {
+                id: uuidv4(),
+                type: "human",
+                content: getMessageContent(msg),
+              },
+            ],
+            {
+              runConfig: msg.runConfig,
+              ...(checkpointId && { checkpointId }),
+            },
+          );
+        }
+      : undefined,
+    onReload: getCheckpointId
+      ? async (parentId, config) => {
+          await toolInvocations.abort();
+          const truncated = truncateAdkMessages(
+            threadMessagesRef.current,
+            parentId,
+          );
+          setMessages(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,
+      result,
+      isError,
+      artifact,
+    }) => {
+      await handleSendMessage(
+        [
+          {
+            id: uuidv4(),
+            type: "tool",
+            name: toolName,
+            tool_call_id: toolCallId,
+            content: JSON.stringify(result),
+            artifact,
+            status: isError ? "error" : "success",
+          },
+        ],
+        {},
+      );
+    },
+    onResumeToolCall: (options) =>
+      toolInvocations.resume(options.toolCallId, options.payload),
+    onCancel: unstable_allowCancellation
+      ? async () => {
+          cancel();
+          await toolInvocations.abort();
+        }
+      : undefined,
+  });
+
+  {
+    const loadRef = useRef(load);
+    useEffect(() => {
+      loadRef.current = load;
+    });
+
+    useEffect(() => {
+      const loadFn = loadRef.current;
+      if (!loadFn) return;
+
+      const externalId = aui.threadListItem().getState().externalId;
+      if (externalId == null) return;
+
+      loadFn(externalId).then(
+        ({ messages: msgs }) => {
+          setMessages(msgs);
+        },
+        (e) => {
+          console.warn("Failed to load ADK session:", e);
+        },
+      );
+    }, [aui, setMessages]);
+  }
+
+  return runtime;
+};
+
+export const useAdkRuntime = ({
+  cloud,
+  sessionAdapter,
+  create,
+  delete: deleteFn,
+  ...options
+}: UseAdkRuntimeOptions) => {
+  const aui = useAui();
+  const cloudAdapter = useCloudThreadListAdapter({
+    cloud,
+    create: async () => {
+      if (create) return create();
+      if (aui.threadListItem.source) return aui.threadListItem().initialize();
+      return { externalId: undefined };
+    },
+    delete: deleteFn,
+  });
+
+  const adapter = sessionAdapter ?? cloudAdapter;
+
+  return useRemoteThreadListRuntime({
+    runtimeHook: function RuntimeHook() {
+      return useAdkRuntimeImpl(options);
+    },
+    adapter,
+    allowNesting: true,
+  });
+};