@assistant-ui/react-a2a 0.1.0

package/src/appendA2AChunk.ts

@@ -0,0 +1,121 @@
+import { A2AMessage } from "./types";
+import { parsePartialJsonObject } from "assistant-stream/utils";
+
+export const appendA2AChunk = (
+  prev: A2AMessage | undefined,
+  curr: A2AMessage,
+): A2AMessage => {
+  // If no previous message or different message type, return current as-is
+  if (!prev || prev.role !== curr.role || prev.id !== curr.id) {
+    return curr;
+  }
+
+  // For assistant messages, we need to handle streaming content and tool calls
+  if (curr.role === "assistant") {
+    const newContent = Array.isArray(prev.content)
+      ? [...prev.content]
+      : typeof prev.content === "string"
+        ? [{ type: "text" as const, text: prev.content }]
+        : [];
+
+    // Append new content chunks
+    if (typeof curr.content === "string") {
+      const lastIndex = newContent.length - 1;
+      const lastPart = newContent[lastIndex];
+
+      if (lastPart?.type === "text") {
+        // Append to existing text part
+        (lastPart as { type: "text"; text: string }).text += curr.content;
+      } else {
+        // Create new text part
+        newContent.push({ type: "text", text: curr.content });
+      }
+    } else if (Array.isArray(curr.content)) {
+      for (const contentPart of curr.content) {
+        const lastIndex = newContent.length - 1;
+        const lastPart = newContent[lastIndex];
+
+        if (contentPart.type === "text" && lastPart?.type === "text") {
+          // Append to existing text part
+          (lastPart as { type: "text"; text: string }).text += contentPart.text;
+        } else {
+          // Add new content part
+          newContent.push(contentPart);
+        }
+      }
+    }
+
+    // Merge tool calls - A2A typically sends complete tool calls rather than chunks
+    const newToolCalls = [...(prev.tool_calls ?? [])];
+    if (curr.tool_calls) {
+      for (const toolCall of curr.tool_calls) {
+        const existingIndex = newToolCalls.findIndex(
+          (tc) => tc.id === toolCall.id,
+        );
+        if (existingIndex >= 0) {
+          // Update existing tool call (merge args if needed)
+          const existing = newToolCalls[existingIndex]!;
+          newToolCalls[existingIndex] = {
+            ...existing,
+            ...toolCall,
+            // If argsText is provided in chunks, concatenate it
+            argsText: (existing.argsText || "") + (toolCall.argsText || ""),
+            // Try to parse merged args, fallback to existing or new args
+            args:
+              parsePartialJsonObject(
+                (existing.argsText || "") + (toolCall.argsText || ""),
+              ) ||
+              toolCall.args ||
+              existing.args,
+          };
+        } else {
+          // Add new tool call
+          newToolCalls.push(toolCall);
+        }
+      }
+    }
+
+    // Merge artifacts
+    const newArtifacts = [...(prev.artifacts ?? [])];
+    if (curr.artifacts) {
+      for (const artifact of curr.artifacts) {
+        const existingIndex = newArtifacts.findIndex(
+          (a) => a.name === artifact.name,
+        );
+        if (existingIndex >= 0) {
+          // Merge artifact parts
+          const existingArtifact = newArtifacts[existingIndex]!;
+          newArtifacts[existingIndex] = {
+            name: existingArtifact.name,
+            parts: [...existingArtifact.parts, ...artifact.parts],
+          };
+        } else {
+          // Add new artifact
+          newArtifacts.push(artifact);
+        }
+      }
+    }
+
+    const result: A2AMessage = {
+      ...prev,
+      content: newContent,
+    };
+    const newStatus = curr.status || prev.status;
+    if (newStatus) result.status = newStatus;
+    if (newToolCalls.length > 0) result.tool_calls = newToolCalls;
+    if (newArtifacts.length > 0) result.artifacts = newArtifacts;
+    return result;
+  }
+
+  // For other message types (user, system, tool), just return the current message
+  // as they typically don't stream in chunks
+  const result: A2AMessage = {
+    ...prev,
+    ...curr,
+  };
+  // Preserve any existing artifacts and merge with new ones
+  if (curr.artifacts || prev.artifacts) {
+    result.artifacts = [...(prev.artifacts ?? []), ...(curr.artifacts ?? [])];
+  }
+  return result;
+};

package/src/convertA2AMessages.ts

@@ -0,0 +1,108 @@
+"use client";
+
+import { useExternalMessageConverter } from "@assistant-ui/react";
+import { A2AMessage } from "./types";
+import { ToolCallMessagePart } from "@assistant-ui/react";
+import { ThreadUserMessage } from "@assistant-ui/react";
+
+const contentToParts = (content: A2AMessage["content"]) => {
+  if (typeof content === "string")
+    return [{ type: "text" as const, text: content }];
+  return content
+    .map((part): ThreadUserMessage["content"][number] | null => {
+      const type = part.type;
+      switch (type) {
+        case "text":
+          return { type: "text", text: part.text };
+        case "image_url":
+          if (typeof part.image_url === "string") {
+            return { type: "image", image: part.image_url };
+          } else {
+            return {
+              type: "image",
+              image: part.image_url.url,
+            };
+          }
+        case "data":
+          // Convert data parts to text representation for display
+          return {
+            type: "text",
+            text: `[Data: ${JSON.stringify(part.data)}]`,
+          };
+        default:
+          return null;
+      }
+    })
+    .filter((part): part is NonNullable<typeof part> => part !== null);
+};
+
+export const convertA2AMessage = (message: A2AMessage) => {
+  const role = message.role;
+  switch (role) {
+    case "system":
+      return {
+        id: message.id,
+        role: "system" as const,
+        content: [{ type: "text" as const, text: message.content as string }],
+      };
+
+    case "user":
+      return {
+        id: message.id,
+        role: "user" as const,
+        content: contentToParts(message.content),
+      };
+
+    case "assistant": {
+      const toolCallParts: ToolCallMessagePart[] =
+        message.tool_calls?.map((toolCall) => ({
+          type: "tool-call",
+          toolCallId: toolCall.id,
+          toolName: toolCall.name,
+          args: toolCall.args,
+          argsText: toolCall.argsText ?? JSON.stringify(toolCall.args),
+        })) ?? [];
+
+      return {
+        id: message.id,
+        role: "assistant" as const,
+        content: [...contentToParts(message.content), ...toolCallParts],
+        status: message.status,
+      };
+    }
+
+    case "tool":
+      return {
+        id: message.id,
+        role: "user" as const,
+        content: [
+          {
+            type: "tool-call" as const,
+            toolCallId: message.tool_call_id!,
+            toolName: "", // A2A doesn't store tool name in tool messages
+            result: JSON.parse(message.content as string),
+            isError:
+              message.status?.type === "incomplete" &&
+              message.status?.reason === "error",
+          },
+        ],
+      };
+
+    default:
+      const _exhaustiveCheck: never = role;
+      throw new Error(`Unknown message role: ${_exhaustiveCheck}`);
+  }
+};
+
+export const convertA2AMessages = (messages: A2AMessage[]) =>
+  messages.map(convertA2AMessage);
+
+export const useA2AMessageConverter = (
+  messages: A2AMessage[],
+  isRunning: boolean,
+) =>
+  useExternalMessageConverter({
+    callback: convertA2AMessage,
+    messages,
+    isRunning,
+  });

package/src/index.ts

@@ -0,0 +1,6 @@
+export * from "./useA2ARuntime";
+export * from "./useA2AMessages";
+export * from "./convertA2AMessages";
+export * from "./types";
+export * from "./A2AMessageAccumulator";
+export * from "./appendA2AChunk";

package/src/testUtils.ts

@@ -0,0 +1,11 @@
+import { A2AMessage } from "./types";
+import { A2AMessagesEvent } from "./useA2AMessages";
+
+export const mockStreamCallbackFactory = (
+  events: Array<A2AMessagesEvent<A2AMessage>>,
+) =>
+  async function* () {
+    for (const event of events) {
+      yield event;
+    }
+  };

package/src/types.ts

@@ -0,0 +1,114 @@
+import { MessageStatus } from "@assistant-ui/react";
+import { ReadonlyJSONObject } from "assistant-stream/utils";
+
+// A2A Message Types
+export type A2AMessage = {
+  id?: string;
+  role: "user" | "assistant" | "system" | "tool";
+  content: string | A2AMessageContent[];
+  tool_calls?: A2AToolCall[];
+  tool_call_id?: string;
+  artifacts?: A2AArtifact[];
+  status?: MessageStatus;
+};
+
+export type A2AMessageContent =
+  | { type: "text"; text: string }
+  | { type: "image_url"; image_url: string | { url: string } }
+  | { type: "data"; data: any };
+
+export type A2AToolCall = {
+  id: string;
+  name: string;
+  args: ReadonlyJSONObject;
+  argsText?: string;
+};
+
+export type A2AArtifact = {
+  name: string;
+  parts: A2AArtifactPart[];
+};
+
+export type A2AArtifactPart = {
+  kind: "text" | "data" | "file";
+  data?: any;
+  text?: string;
+  metadata?: Record<string, any>;
+};
+
+// A2A Events (similar to LangGraph events)
+export type A2AEvent = {
+  event: A2AKnownEventTypes | string;
+  data: any;
+};
+
+export enum A2AKnownEventTypes {
+  TaskUpdate = "task-update",
+  TaskComplete = "task-complete",
+  TaskFailed = "task-failed",
+  Artifacts = "artifacts",
+  StateUpdate = "state-update",
+  Error = "error",
+}
+
+// A2A Task State
+export type A2ATaskState = {
+  id: string;
+  state: "pending" | "working" | "completed" | "failed";
+  progress?: number;
+  message?: string;
+};
+
+// A2A Task Result
+export type A2ATaskResult = {
+  id: string;
+  status: {
+    state: "pending" | "working" | "completed" | "failed";
+    message?: string;
+  };
+  artifacts?: A2AArtifact[];
+  history?: Array<{
+    messageId: string;
+    role: string;
+    parts?: Array<{ kind: string; text?: string; data?: any }>;
+  }>;
+};
+
+// A2A Configuration
+export type A2AConfig = {
+  contextId?: string;
+  runConfig?: Record<string, any>;
+};
+
+// A2A Send Message Configuration
+export type A2ASendMessageConfig = A2AConfig & {
+  command?: A2ACommand;
+};
+
+// A2A Commands (for interrupts/resume)
+export type A2ACommand = {
+  resume?: string;
+  interrupt?: string;
+};
+
+// A2A Stream Callback
+export type A2AStreamCallback<TMessage> = (
+  messages: TMessage[],
+  config: A2ASendMessageConfig & { abortSignal: AbortSignal },
+) => Promise<AsyncGenerator<A2AEvent>> | AsyncGenerator<A2AEvent>;
+
+// Event handler callback types
+export type OnTaskUpdateEventCallback = (
+  data: A2ATaskState,
+) => void | Promise<void>;
+export type OnArtifactsEventCallback = (
+  artifacts: A2AArtifact[],
+) => void | Promise<void>;
+export type OnErrorEventCallback = (error: unknown) => void | Promise<void>;
+export type OnStateUpdateEventCallback = (
+  state: unknown,
+) => void | Promise<void>;
+export type OnCustomEventCallback = (
+  type: string,
+  data: unknown,
+) => void | Promise<void>;

package/src/useA2AMessages.ts

@@ -0,0 +1,180 @@
+import { useState, useCallback, useRef, useMemo } from "react";
+import { v4 as uuidv4 } from "uuid";
+import { A2AMessageAccumulator } from "./A2AMessageAccumulator";
+import {
+  A2AEvent,
+  A2AKnownEventTypes,
+  A2ATaskState,
+  A2AArtifact,
+  A2ASendMessageConfig,
+  A2AStreamCallback,
+  OnTaskUpdateEventCallback,
+  OnArtifactsEventCallback,
+  OnErrorEventCallback,
+  OnStateUpdateEventCallback,
+  OnCustomEventCallback,
+} from "./types";
+
+export type A2AMessagesEvent<_TMessage> = A2AEvent;
+
+const DEFAULT_APPEND_MESSAGE = <TMessage>(
+  _: TMessage | undefined,
+  curr: TMessage,
+) => curr;
+
+export const useA2AMessages = <TMessage extends { id?: string }>({
+  stream,
+  appendMessage = DEFAULT_APPEND_MESSAGE,
+  eventHandlers,
+}: {
+  stream: A2AStreamCallback<TMessage>;
+  appendMessage?: (prev: TMessage | undefined, curr: TMessage) => TMessage;
+  eventHandlers?: {
+    onTaskUpdate?: OnTaskUpdateEventCallback;
+    onArtifacts?: OnArtifactsEventCallback;
+    onError?: OnErrorEventCallback;
+    onStateUpdate?: OnStateUpdateEventCallback;
+    onCustomEvent?: OnCustomEventCallback;
+  };
+}) => {
+  const [messages, setMessages] = useState<TMessage[]>([]);
+  const [taskState, setTaskState] = useState<A2ATaskState | undefined>();
+  const [artifacts, setArtifacts] = useState<A2AArtifact[]>([]);
+  const abortControllerRef = useRef<AbortController | null>(null);
+
+  const { onTaskUpdate, onArtifacts, onError, onStateUpdate, onCustomEvent } =
+    useMemo(() => eventHandlers ?? {}, [eventHandlers]);
+
+  const sendMessage = useCallback(
+    async (newMessages: TMessage[], config: A2ASendMessageConfig) => {
+      // ensure all messages have an ID
+      const newMessagesWithId = newMessages.map((m) =>
+        m.id ? m : { ...m, id: uuidv4() },
+      );
+
+      const accumulator = new A2AMessageAccumulator({
+        initialMessages: messages,
+        appendMessage,
+      });
+      setMessages(accumulator.addMessages(newMessagesWithId));
+
+      const abortController = new AbortController();
+      abortControllerRef.current = abortController;
+      const response = await stream(newMessagesWithId, {
+        ...config,
+        abortSignal: abortController.signal,
+      });
+
+      for await (const event of response) {
+        switch (event.event) {
+          case A2AKnownEventTypes.TaskUpdate:
+            const taskData = event.data as A2ATaskState;
+            setTaskState(taskData);
+            onTaskUpdate?.(taskData);
+            break;
+
+          case A2AKnownEventTypes.TaskComplete:
+            // Extract messages and artifacts from completed task
+            const { messages: taskMessages, artifacts: taskArtifacts } =
+              event.data;
+            if (taskMessages) {
+              setMessages(accumulator.addMessages(taskMessages));
+            }
+            if (taskArtifacts) {
+              setArtifacts(taskArtifacts);
+              onArtifacts?.(taskArtifacts);
+            }
+            // Clear task state on completion
+            setTaskState(undefined);
+            break;
+
+          case A2AKnownEventTypes.TaskFailed:
+            onError?.(event.data);
+            // Update task state to failed
+            if (taskState) {
+              setTaskState({
+                ...taskState,
+                state: "failed",
+                message: event.data?.message,
+              });
+            }
+            break;
+
+          case A2AKnownEventTypes.Artifacts:
+            const artifactData = event.data as A2AArtifact[];
+            setArtifacts(artifactData);
+            onArtifacts?.(artifactData);
+            break;
+
+          case A2AKnownEventTypes.StateUpdate:
+            onStateUpdate?.(event.data);
+            break;
+
+          case A2AKnownEventTypes.Error:
+            onError?.(event.data);
+            // Update the last assistant message with error status if available
+            const messages = accumulator.getMessages();
+            const lastAssistantMessage = messages.findLast(
+              (m): m is TMessage & { role: string; id: string } =>
+                m != null &&
+                "role" in m &&
+                m.role === "assistant" &&
+                m.id != null,
+            );
+            if (lastAssistantMessage) {
+              const errorMessage = {
+                ...lastAssistantMessage,
+                status: {
+                  type: "incomplete" as const,
+                  reason: "error" as const,
+                  error: event.data,
+                },
+              };
+              setMessages(accumulator.addMessages([errorMessage]));
+            }
+            break;
+
+          default:
+            if (onCustomEvent) {
+              onCustomEvent(event.event, event.data);
+            } else {
+              console.warn(
+                "Unhandled A2A event received:",
+                event.event,
+                event.data,
+              );
+            }
+            break;
+        }
+      }
+    },
+    [
+      messages,
+      appendMessage,
+      stream,
+      onTaskUpdate,
+      onArtifacts,
+      onError,
+      onStateUpdate,
+      onCustomEvent,
+      taskState,
+    ],
+  );
+
+  const cancel = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+    }
+  }, [abortControllerRef]);
+
+  return {
+    messages,
+    artifacts,
+    taskState,
+    sendMessage,
+    cancel,
+    setMessages,
+    setArtifacts,
+    setTaskState,
+  };
+};