@assistant-ui/react-ag-ui 0.0.2

package/src/runtime/adapter/conversions.ts

@@ -0,0 +1,165 @@
+"use client";
+
+type ThreadMessageLike = {
+  id: string;
+  role: string;
+  content: any;
+  name?: string;
+};
+
+const generateId = () =>
+  (globalThis.crypto as any)?.randomUUID?.() ??
+  Math.random().toString(36).slice(2);
+
+type AgUiToolCall = {
+  id: string;
+  type: "function";
+  function: { name: string; arguments: string };
+};
+
+export type AgUiMessage =
+  | {
+      id: string;
+      role: string;
+      content: string;
+      name?: string;
+      toolCalls?: AgUiToolCall[];
+    }
+  | {
+      id: string;
+      role: "tool";
+      content: string;
+      toolCallId: string;
+      error?: string;
+    };
+
+type ToolCallPart = {
+  type: "tool-call";
+  toolCallId?: string;
+  toolName?: string;
+  argsText?: string;
+  args?: Record<string, unknown>;
+  result?: unknown;
+  isError?: boolean;
+};
+
+const normaliseToolCall = (part: ToolCallPart) => {
+  const id = part.toolCallId ?? generateId();
+  const argsText =
+    typeof part.argsText === "string"
+      ? part.argsText
+      : JSON.stringify(part.args ?? {});
+  return {
+    id,
+    call: {
+      id,
+      type: "function" as const,
+      function: {
+        name: part.toolName ?? "tool",
+        arguments: argsText,
+      },
+    },
+  };
+};
+
+export const toAgUiMessages = (
+  messages: readonly ThreadMessageLike[],
+): AgUiMessage[] => {
+  const converted: AgUiMessage[] = [];
+
+  for (const message of messages) {
+    const role = message.role;
+    const content = extractText(message.content);
+
+    if (role === "assistant" && Array.isArray(message.content)) {
+      const toolCallParts = message.content.filter(
+        (part): part is ToolCallPart => part?.type === "tool-call",
+      );
+
+      const toolCalls = toolCallParts.map((part) => {
+        const { id, call } = normaliseToolCall(part);
+        return { id, call, part };
+      });
+
+      converted.push({
+        id: message.id,
+        role: "assistant",
+        content,
+        ...(message.name ? { name: message.name } : {}),
+        ...(toolCalls.length > 0
+          ? { toolCalls: toolCalls.map((entry) => entry.call) }
+          : {}),
+      });
+
+      for (const { id: toolCallId, part } of toolCalls) {
+        if (part.result === undefined) continue;
+
+        const resultContent =
+          typeof part.result === "string"
+            ? part.result
+            : JSON.stringify(part.result);
+        converted.push({
+          id: `${toolCallId}:tool`,
+          role: "tool",
+          content: resultContent,
+          toolCallId,
+          ...(part.isError ? { error: resultContent } : {}),
+        });
+      }
+
+      continue;
+    }
+
+    if (role === "tool") {
+      const toolCallId = (message as any).toolCallId ?? generateId();
+      converted.push({
+        id: message.id,
+        role: "tool",
+        content,
+        toolCallId,
+        ...(typeof (message as any).error === "string"
+          ? { error: (message as any).error }
+          : undefined),
+      });
+      continue;
+    }
+
+    converted.push({
+      id: message.id,
+      role,
+      content,
+      ...(message.name ? { name: message.name } : {}),
+    });
+  }
+
+  return converted;
+};
+
+const extractText = (content: any): string => {
+  if (!Array.isArray(content)) {
+    if (typeof content === "string") return content;
+    return "";
+  }
+
+  return content
+    .filter((part) => part?.type === "text" && typeof part?.text === "string")
+    .map((part) => part.text as string)
+    .join("\n");
+};
+
+export const toAgUiTools = (tools: Record<string, any> | undefined) => {
+  if (!tools) return [];
+
+  return Object.entries(tools)
+    .filter(([, tool]) => !tool?.disabled && tool?.type !== "backend")
+    .map(([name, tool]) => ({
+      name,
+      description: tool?.description ?? undefined,
+      parameters:
+        typeof tool?.parameters?.toJSON === "function"
+          ? tool.parameters.toJSON()
+          : typeof tool?.parameters?.toJSONSchema === "function"
+            ? tool.parameters.toJSONSchema()
+            : tool?.parameters,
+    }));
+};

package/src/runtime/adapter/run-aggregator.ts

@@ -0,0 +1,376 @@
+"use client";
+
+import type {
+  ChatModelRunResult,
+  ThreadAssistantMessagePart,
+  ToolCallMessagePart,
+} from "@assistant-ui/react";
+import type { AgUiEvent } from "../types";
+import type { Logger } from "../logger";
+
+type Emit = (update: ChatModelRunResult) => void;
+
+type ToolCallState = {
+  toolCallId: string;
+  toolCallName: string;
+  argsText: string;
+  parsedArgs: Record<string, unknown> | undefined;
+  result: unknown;
+  isError: boolean | undefined;
+  parentMessageId?: string;
+};
+
+export type RunAggregatorOptions = {
+  showThinking: boolean;
+  logger: Logger;
+  emit: Emit;
+};
+
+/**
+ * Collects AG-UI events into assistant-ui run snapshots that can be yielded from a ChatModelAdapter.
+ *
+ * The aggregator keeps a single assistant message worth of parts. Each incoming event updates the parts and
+ * emits a fresh snapshot through the provided `emit` callback.
+ */
+export class RunAggregator {
+  private readonly emitUpdate: Emit;
+  private readonly showThinking: boolean;
+  private readonly logger: Logger;
+
+  private status: ChatModelRunResult["status"] | undefined;
+  private readonly textParts = new Map<
+    string,
+    { buffer: string; touched: boolean }
+  >();
+  private activeTextMessageId: string | undefined;
+  private reasoningBuffer = "";
+  private reasoningActive = false;
+  private readonly toolCalls = new Map<string, ToolCallState>();
+  private readonly partOrder: (
+    | { kind: "text"; key: string }
+    | { kind: "reasoning" }
+    | { kind: "tool-call"; toolCallId: string }
+  )[] = [];
+  private hasReasoningPart = false;
+  private textPartCounter = 0;
+
+  constructor(options: RunAggregatorOptions) {
+    this.emitUpdate = options.emit;
+    this.showThinking = options.showThinking;
+    this.logger = options.logger;
+  }
+
+  handle(event: AgUiEvent): void {
+    switch (event.type) {
+      case "RUN_STARTED": {
+        this.clearTextParts();
+        this.reasoningBuffer = "";
+        this.reasoningActive = false;
+        this.toolCalls.clear();
+        this.partOrder.length = 0;
+        this.hasReasoningPart = false;
+        this.textPartCounter = 0;
+        this.activeTextMessageId = undefined;
+        this.status = { type: "running" };
+        this.emit();
+        break;
+      }
+      case "RUN_FINISHED": {
+        this.status = { type: "complete", reason: "unknown" };
+        this.emit();
+        break;
+      }
+      case "RUN_ERROR": {
+        this.status = {
+          type: "incomplete",
+          reason: "error",
+          ...(event.message !== undefined ? { error: event.message } : {}),
+        };
+        this.emit();
+        break;
+      }
+      case "RUN_CANCELLED": {
+        this.status = { type: "incomplete", reason: "cancelled" };
+        this.emit();
+        break;
+      }
+
+      case "TEXT_MESSAGE_START": {
+        const id = this.startTextMessage(event.messageId);
+        if (id) {
+          this.markTextPartTouched(id);
+        }
+        this.emit();
+        break;
+      }
+      case "TEXT_MESSAGE_CONTENT":
+      case "TEXT_MESSAGE_CHUNK": {
+        if (!event.delta) break;
+        const id = this.resolveTextMessageId(
+          "messageId" in event ? event.messageId : undefined,
+        );
+        this.appendText(id, event.delta);
+        this.emit();
+        break;
+      }
+      case "TEXT_MESSAGE_END": {
+        if (event.messageId && this.activeTextMessageId === event.messageId) {
+          this.activeTextMessageId = undefined;
+        }
+        this.emit();
+        break;
+      }
+
+      case "THINKING_START":
+      case "THINKING_TEXT_MESSAGE_START": {
+        if (!this.showThinking) break;
+        this.reasoningActive = true;
+        if (!this.reasoningBuffer) this.reasoningBuffer = "";
+        this.ensureReasoningPart();
+        this.emit();
+        break;
+      }
+      case "THINKING_TEXT_MESSAGE_CONTENT": {
+        if (!this.showThinking || !event.delta) break;
+        this.reasoningBuffer += event.delta;
+        this.ensureReasoningPart();
+        this.emit();
+        break;
+      }
+      case "THINKING_TEXT_MESSAGE_END":
+      case "THINKING_END": {
+        if (!this.showThinking) break;
+        this.emit();
+        break;
+      }
+
+      case "TOOL_CALL_START": {
+        this.startToolCall(
+          event.toolCallId,
+          event.toolCallName,
+          event.parentMessageId,
+        );
+        this.emit();
+        break;
+      }
+      case "TOOL_CALL_ARGS":
+      case "TOOL_CALL_CHUNK": {
+        if (!event.delta) break;
+        this.appendToolArgs(event.toolCallId, event.delta);
+        this.emit();
+        break;
+      }
+      case "TOOL_CALL_END": {
+        this.emit();
+        break;
+      }
+      case "TOOL_CALL_RESULT": {
+        this.finishToolCall(
+          event.toolCallId,
+          event.content ?? "",
+          event.role === "tool" ? false : undefined,
+        );
+        this.emit();
+        break;
+      }
+
+      default: {
+        this.logger.debug?.("[agui] aggregator ignored event", event);
+      }
+    }
+  }
+
+  private clearTextParts(): void {
+    this.textParts.clear();
+  }
+
+  private generateTextKey(): string {
+    this.textPartCounter += 1;
+    return `text-${this.textPartCounter}`;
+  }
+
+  private startTextMessage(messageId?: string): string {
+    const id = messageId ?? this.generateTextKey();
+    this.ensureTextPart(id);
+    this.activeTextMessageId = id;
+    return id;
+  }
+
+  private resolveTextMessageId(messageId?: string): string {
+    if (messageId) {
+      this.ensureTextPart(messageId);
+      this.activeTextMessageId = messageId;
+      return messageId;
+    }
+
+    if (this.activeTextMessageId) {
+      return this.activeTextMessageId;
+    }
+
+    const generated = this.generateTextKey();
+    this.ensureTextPart(generated);
+    this.activeTextMessageId = generated;
+    return generated;
+  }
+
+  private ensureTextPart(id: string): void {
+    if (!this.textParts.has(id)) {
+      this.textParts.set(id, { buffer: "", touched: false });
+      if (
+        !this.partOrder.some((part) => part.kind === "text" && part.key === id)
+      ) {
+        this.partOrder.push({ kind: "text", key: id });
+      }
+    }
+  }
+
+  private markTextPartTouched(id: string): void {
+    const entry = this.textParts.get(id);
+    if (!entry) return;
+    entry.touched = true;
+  }
+
+  private appendText(id: string, delta: string): void {
+    this.ensureTextPart(id);
+    const entry = this.textParts.get(id);
+    if (!entry) return;
+    entry.buffer += delta;
+    entry.touched = true;
+  }
+
+  private startToolCall(
+    id: string | undefined,
+    name?: string,
+    parentMessageId?: string,
+  ) {
+    if (!id) return;
+    if (
+      !this.partOrder.some(
+        (part) => part.kind === "tool-call" && part.toolCallId === id,
+      )
+    ) {
+      this.partOrder.push({ kind: "tool-call", toolCallId: id });
+    }
+    const state: ToolCallState = {
+      toolCallId: id,
+      toolCallName: name ?? "tool",
+      argsText: "",
+      parsedArgs: undefined,
+      result: undefined,
+      isError: undefined,
+    };
+    if (parentMessageId) {
+      state.parentMessageId = parentMessageId;
+    }
+    this.toolCalls.set(id, state);
+  }
+
+  private appendToolArgs(id: string | undefined, delta: string) {
+    const entry = id ? this.toolCalls.get(id) : undefined;
+    if (!entry) return;
+    entry.argsText += delta;
+    try {
+      const parsed = JSON.parse(entry.argsText);
+      if (parsed && typeof parsed === "object") {
+        entry.parsedArgs = parsed as Record<string, unknown>;
+      } else {
+        entry.parsedArgs = undefined;
+      }
+    } catch {
+      entry.parsedArgs = undefined;
+    }
+  }
+
+  private finishToolCall(id: string, content: string, isError?: boolean) {
+    if (!id) return;
+    let entry = this.toolCalls.get(id);
+    if (!entry) {
+      entry = {
+        toolCallId: id,
+        toolCallName: "tool",
+        argsText: "",
+        parsedArgs: undefined,
+        result: undefined,
+        isError: undefined,
+      };
+      this.toolCalls.set(id, entry);
+    }
+    if (
+      !this.partOrder.some(
+        (part) => part.kind === "tool-call" && part.toolCallId === id,
+      )
+    ) {
+      this.partOrder.push({ kind: "tool-call", toolCallId: id });
+    }
+    entry.result = this.tryParseJSON(content);
+    entry.isError = isError;
+  }
+
+  private tryParseJSON(value: string): unknown {
+    if (!value) return value;
+    try {
+      return JSON.parse(value);
+    } catch {
+      return value;
+    }
+  }
+
+  private emit(): void {
+    const snapshot: ThreadAssistantMessagePart[] = [];
+
+    for (const part of this.partOrder) {
+      if (part.kind === "reasoning") {
+        if (
+          this.showThinking &&
+          (this.reasoningActive || this.reasoningBuffer.length > 0)
+        ) {
+          snapshot.push({
+            type: "reasoning",
+            text: this.reasoningBuffer,
+          } as const);
+        }
+        continue;
+      }
+
+      if (part.kind === "text") {
+        const entry = this.textParts.get(part.key);
+        if (entry?.touched) {
+          snapshot.push({ type: "text", text: entry.buffer } as const);
+        }
+        continue;
+      }
+
+      const entry = this.toolCalls.get(part.toolCallId);
+      if (!entry) continue;
+      const toolPart: ToolCallMessagePart = {
+        type: "tool-call",
+        toolCallId: entry.toolCallId,
+        toolName: entry.toolCallName,
+        args: (entry.parsedArgs ?? {}) as any,
+        argsText: entry.argsText,
+        ...(entry.result !== undefined ? { result: entry.result } : {}),
+        ...(entry.isError !== undefined ? { isError: entry.isError } : {}),
+        ...(entry.parentMessageId ? { parentId: entry.parentMessageId } : {}),
+      } as ToolCallMessagePart;
+      snapshot.push(toolPart);
+    }
+
+    const result: ChatModelRunResult = {
+      content: snapshot,
+      ...(this.status ? { status: this.status } : undefined),
+    };
+    this.emitUpdate(result);
+  }
+
+  private ensureReasoningPart(): void {
+    if (this.hasReasoningPart) return;
+    // ensure reasoning appears before the first text segment if possible
+    const textIndex = this.partOrder.findIndex((part) => part.kind === "text");
+    if (textIndex === -1) {
+      this.partOrder.push({ kind: "reasoning" });
+    } else {
+      this.partOrder.splice(textIndex, 0, { kind: "reasoning" });
+    }
+    this.hasReasoningPart = true;
+  }
+}

package/src/runtime/adapter/subscriber.ts

@@ -0,0 +1,134 @@
+"use client";
+
+import type { AgUiEvent } from "../types";
+import { parseAgUiEvent } from "../event-parser";
+
+type Dispatch = (event: AgUiEvent) => void;
+
+type Subscriber = {
+  onEvent?: (payload: { event: unknown }) => void;
+  onTextMessageStartEvent?: (payload: { event: unknown }) => void;
+  onTextMessageContentEvent?: (payload: { event: unknown }) => void;
+  onTextMessageEndEvent?: (payload: { event: unknown }) => void;
+  onTextMessageChunkEvent?: (payload: { event: unknown }) => void;
+  onThinkingStartEvent?: (payload: { event: unknown }) => void;
+  onThinkingEndEvent?: (payload: { event: unknown }) => void;
+  onThinkingTextMessageStartEvent?: (payload: { event: unknown }) => void;
+  onThinkingTextMessageContentEvent?: (payload: { event: unknown }) => void;
+  onThinkingTextMessageEndEvent?: (payload: { event: unknown }) => void;
+  onToolCallStartEvent?: (payload: { event: unknown }) => void;
+  onToolCallArgsEvent?: (payload: { event: unknown }) => void;
+  onToolCallEndEvent?: (payload: { event: unknown }) => void;
+  onToolCallChunkEvent?: (payload: { event: unknown }) => void;
+  onToolCallResultEvent?: (payload: { event: unknown }) => void;
+  onStateSnapshotEvent?: (payload: { event: unknown }) => void;
+  onStateDeltaEvent?: (payload: { event: unknown }) => void;
+  onMessagesSnapshotEvent?: (payload: { event: unknown }) => void;
+  onCustomEvent?: (payload: { event: unknown }) => void;
+  onRawEvent?: (payload: { event: unknown }) => void;
+  onRunFinalized?: () => void;
+  onRunFailed?: (payload: { error: Error }) => void;
+};
+
+const ensureEvent = (
+  raw: unknown,
+  type: AgUiEvent["type"],
+): AgUiEvent | null => {
+  if (raw && typeof raw === "object") {
+    const payload = raw as Record<string, unknown>;
+    if (typeof payload["type"] === "string") {
+      return parseAgUiEvent(payload);
+    }
+    return parseAgUiEvent({ type, ...payload });
+  }
+  return parseAgUiEvent({ type });
+};
+
+const dispatchIfValid = (
+  dispatch: Dispatch,
+  raw: unknown,
+  type: AgUiEvent["type"],
+) => {
+  const event = ensureEvent(raw, type);
+  if (!event) return;
+  dispatch(event);
+};
+
+type SubscriberOptions = {
+  dispatch: Dispatch;
+  runId: string;
+  onRunFailed?: (error: Error) => void;
+};
+
+export const createAgUiSubscriber = (
+  options: SubscriberOptions,
+): Subscriber => {
+  const { dispatch, runId, onRunFailed } = options;
+  return {
+    onEvent: ({ event }) => {
+      const typeCandidate =
+        event && typeof event === "object"
+          ? (event as Record<string, unknown>)["type"]
+          : undefined;
+      if (typeof typeCandidate === "string") {
+        // Typed handlers will receive this via the discriminated callbacks; avoid duplicates.
+        return;
+      }
+      const parsed = parseAgUiEvent(event);
+      if (parsed) dispatch(parsed);
+    },
+    onTextMessageStartEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TEXT_MESSAGE_START"),
+    onTextMessageContentEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TEXT_MESSAGE_CONTENT"),
+    onTextMessageEndEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TEXT_MESSAGE_END"),
+    onTextMessageChunkEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TEXT_MESSAGE_CHUNK"),
+    onThinkingStartEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "THINKING_START"),
+    onThinkingEndEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "THINKING_END"),
+    onThinkingTextMessageStartEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "THINKING_TEXT_MESSAGE_START"),
+    onThinkingTextMessageContentEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "THINKING_TEXT_MESSAGE_CONTENT"),
+    onThinkingTextMessageEndEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "THINKING_TEXT_MESSAGE_END"),
+    onToolCallStartEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TOOL_CALL_START"),
+    onToolCallArgsEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TOOL_CALL_ARGS"),
+    onToolCallEndEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TOOL_CALL_END"),
+    onToolCallChunkEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TOOL_CALL_CHUNK"),
+    onToolCallResultEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "TOOL_CALL_RESULT"),
+    onStateSnapshotEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "STATE_SNAPSHOT"),
+    onStateDeltaEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "STATE_DELTA"),
+    onMessagesSnapshotEvent: ({ event }) =>
+      dispatchIfValid(dispatch, event, "MESSAGES_SNAPSHOT"),
+    onCustomEvent: ({ event }) => dispatchIfValid(dispatch, event, "CUSTOM"),
+    onRawEvent: ({ event }) => dispatchIfValid(dispatch, event, "RAW"),
+    onRunFinalized: () => dispatch({ type: "RUN_FINISHED", runId }),
+    onRunFailed: ({ error }) => {
+      onRunFailed?.(error);
+      const message =
+        typeof error.message === "string" ? error.message : "Run failed";
+      const code =
+        typeof (error as any)?.code === "string"
+          ? (error as any).code
+          : undefined;
+      dispatch({
+        type: "RUN_ERROR" as const,
+        ...(message !== undefined ? { message } : {}),
+        ...(code !== undefined ? { code } : {}),
+      } satisfies AgUiEvent);
+    },
+  };
+};
+
+export type AgUiSubscriber = ReturnType<typeof createAgUiSubscriber>;