@assistant-ui/react-ag-ui 0.0.2

package/src/runtime/event-parser.ts

@@ -0,0 +1,172 @@
+import type { AgUiEvent } from "./types";
+
+const isString = (value: unknown): value is string => typeof value === "string";
+const isNonEmptyString = (value: unknown): value is string =>
+  isString(value) && value.length > 0;
+
+const withOptional = <T extends object>(
+  base: T,
+  optionals: Record<string, unknown>,
+) => {
+  const definedEntries = Object.entries(optionals).filter(
+    ([, value]) => value !== undefined,
+  );
+  return definedEntries.length === 0
+    ? base
+    : ({ ...base, ...Object.fromEntries(definedEntries) } as T);
+};
+
+export const parseAgUiEvent = (event: unknown): AgUiEvent | null => {
+  if (!event || typeof event !== "object") return null;
+  const payload = event as Record<string, unknown>;
+  const typeValue = payload["type"];
+  if (!isString(typeValue)) return null;
+
+  const getString = (key: string) =>
+    isString(payload[key]) ? (payload[key] as string) : undefined;
+
+  switch (typeValue) {
+    case "RUN_STARTED": {
+      const runId = getString("runId");
+      return runId ? { type: "RUN_STARTED", runId } : null;
+    }
+    case "RUN_FINISHED": {
+      const runId = getString("runId");
+      return runId ? { type: "RUN_FINISHED", runId } : null;
+    }
+    case "RUN_CANCELLED": {
+      const runId = getString("runId");
+      return withOptional({ type: "RUN_CANCELLED" as const }, { runId });
+    }
+    case "RUN_ERROR": {
+      return withOptional(
+        { type: "RUN_ERROR" as const },
+        {
+          message: getString("message"),
+          code: getString("code"),
+        },
+      );
+    }
+    case "TEXT_MESSAGE_START":
+      return withOptional(
+        { type: "TEXT_MESSAGE_START" as const },
+        { messageId: getString("messageId") },
+      );
+    case "TEXT_MESSAGE_CONTENT": {
+      const delta = getString("delta");
+      if (!isNonEmptyString(delta)) return null;
+      return withOptional(
+        { type: "TEXT_MESSAGE_CONTENT" as const, delta },
+        { messageId: getString("messageId") },
+      );
+    }
+    case "TEXT_MESSAGE_END":
+      return withOptional(
+        { type: "TEXT_MESSAGE_END" as const },
+        { messageId: getString("messageId") },
+      );
+    case "TEXT_MESSAGE_CHUNK": {
+      const delta = getString("delta") ?? "";
+      return withOptional(
+        { type: "TEXT_MESSAGE_CHUNK" as const, delta },
+        { messageId: getString("messageId") },
+      );
+    }
+    case "THINKING_START":
+      return withOptional(
+        { type: "THINKING_START" as const },
+        { title: getString("title") },
+      );
+    case "THINKING_TEXT_MESSAGE_START":
+      return { type: "THINKING_TEXT_MESSAGE_START" };
+    case "THINKING_TEXT_MESSAGE_CONTENT": {
+      const delta = getString("delta") ?? "";
+      return { type: "THINKING_TEXT_MESSAGE_CONTENT", delta };
+    }
+    case "THINKING_TEXT_MESSAGE_END":
+      return { type: "THINKING_TEXT_MESSAGE_END" };
+    case "THINKING_END":
+      return { type: "THINKING_END" };
+    case "TOOL_CALL_START": {
+      const toolCallId = getString("toolCallId");
+      if (!toolCallId) return null;
+      return withOptional(
+        { type: "TOOL_CALL_START" as const, toolCallId },
+        {
+          toolCallName: getString("toolCallName"),
+          parentMessageId: getString("parentMessageId"),
+        },
+      );
+    }
+    case "TOOL_CALL_ARGS": {
+      const toolCallId = getString("toolCallId");
+      if (!toolCallId) return null;
+      const delta = getString("delta") ?? "";
+      return { type: "TOOL_CALL_ARGS", toolCallId, delta };
+    }
+    case "TOOL_CALL_END": {
+      const toolCallId = getString("toolCallId");
+      return toolCallId ? { type: "TOOL_CALL_END", toolCallId } : null;
+    }
+    case "TOOL_CALL_CHUNK":
+      return withOptional(
+        { type: "TOOL_CALL_CHUNK" as const },
+        {
+          toolCallId: getString("toolCallId"),
+          toolCallName: getString("toolCallName"),
+          parentMessageId: getString("parentMessageId"),
+          delta: getString("delta"),
+        },
+      );
+    case "TOOL_CALL_RESULT": {
+      const toolCallId = getString("toolCallId");
+      if (!toolCallId) return null;
+      return withOptional(
+        {
+          type: "TOOL_CALL_RESULT" as const,
+          toolCallId,
+          content: getString("content") ?? "",
+        },
+        {
+          messageId: getString("messageId"),
+          role: payload["role"] === "tool" ? "tool" : undefined,
+        },
+      );
+    }
+    case "STATE_SNAPSHOT":
+      return { type: "STATE_SNAPSHOT", snapshot: payload["snapshot"] };
+    case "STATE_DELTA":
+      return {
+        type: "STATE_DELTA",
+        delta: Array.isArray(payload["delta"])
+          ? (payload["delta"] as any[])
+          : [],
+      };
+    case "MESSAGES_SNAPSHOT":
+      return {
+        type: "MESSAGES_SNAPSHOT",
+        messages: Array.isArray(payload["messages"])
+          ? (payload["messages"] as any[])
+          : [],
+      };
+    case "RAW":
+      return withOptional(
+        { type: "RAW" as const, event: payload["event"] },
+        { source: getString("source") },
+      );
+    case "CUSTOM": {
+      const name = getString("name");
+      if (!name) return null;
+      return { type: "CUSTOM", name, value: payload["value"] };
+    }
+    default:
+      return withOptional(
+        { type: "RAW" as const, event: payload },
+        {
+          source: isString(payload["type"])
+            ? (payload["type"] as string)
+            : undefined,
+        },
+      );
+  }
+};

package/src/runtime/logger.ts

@@ -0,0 +1,9 @@
+export type Logger = {
+  debug: (...a: any[]) => void;
+  error: (...a: any[]) => void;
+};
+
+export const makeLogger = (l?: Partial<Logger>): Logger => ({
+  debug: (...a) => l?.debug?.(...a),
+  error: (...a) => l?.error?.(...a),
+});

package/src/runtime/types.ts

@@ -0,0 +1,66 @@
+import type {
+  AttachmentAdapter,
+  FeedbackAdapter,
+  SpeechSynthesisAdapter,
+  ThreadHistoryAdapter,
+} from "@assistant-ui/react";
+import type { HttpAgent } from "@ag-ui/client";
+import type { Logger } from "./logger";
+
+export type UseAgUiRuntimeAdapters = {
+  attachments?: AttachmentAdapter;
+  speech?: SpeechSynthesisAdapter;
+  feedback?: FeedbackAdapter;
+  history?: ThreadHistoryAdapter;
+};
+
+export type UseAgUiRuntimeOptions = {
+  agent: HttpAgent;
+  logger?: Partial<Logger>;
+  showThinking?: boolean;
+  onError?: (e: Error) => void;
+  onCancel?: () => void;
+  adapters?: UseAgUiRuntimeAdapters;
+};
+
+export type AgUiEvent =
+  | { type: "RUN_STARTED"; runId: string }
+  | { type: "RUN_FINISHED"; runId: string }
+  | { type: "RUN_CANCELLED"; runId?: string }
+  | { type: "RUN_ERROR"; message?: string; code?: string }
+  | { type: "TEXT_MESSAGE_START"; messageId?: string }
+  | { type: "TEXT_MESSAGE_CONTENT"; messageId?: string; delta: string }
+  | { type: "TEXT_MESSAGE_END"; messageId?: string }
+  | { type: "TEXT_MESSAGE_CHUNK"; delta: string }
+  | { type: "THINKING_START"; title?: string }
+  | { type: "THINKING_TEXT_MESSAGE_START" }
+  | { type: "THINKING_TEXT_MESSAGE_CONTENT"; delta: string }
+  | { type: "THINKING_TEXT_MESSAGE_END" }
+  | { type: "THINKING_END" }
+  | {
+      type: "TOOL_CALL_START";
+      toolCallId: string;
+      toolCallName?: string;
+      parentMessageId?: string;
+    }
+  | { type: "TOOL_CALL_ARGS"; toolCallId: string; delta: string }
+  | { type: "TOOL_CALL_END"; toolCallId: string }
+  | {
+      type: "TOOL_CALL_CHUNK";
+      toolCallId?: string;
+      toolCallName?: string;
+      parentMessageId?: string;
+      delta?: string;
+    }
+  | {
+      type: "TOOL_CALL_RESULT";
+      messageId?: string;
+      toolCallId: string;
+      content: string;
+      role?: "tool";
+    }
+  | { type: "RAW"; event: any; source?: string }
+  | { type: "CUSTOM"; name: string; value: any }
+  | { type: "STATE_SNAPSHOT"; snapshot: any }
+  | { type: "STATE_DELTA"; delta: any[] }
+  | { type: "MESSAGES_SNAPSHOT"; messages: any[] };

package/src/useAgUiRuntime.ts

@@ -0,0 +1,101 @@
+"use client";
+
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+  useExternalStoreRuntime,
+  useRuntimeAdapters,
+} from "@assistant-ui/react";
+import type {
+  AssistantRuntime,
+  AppendMessage,
+  ExternalStoreAdapter,
+  ThreadMessage,
+} from "@assistant-ui/react";
+import type { ReadonlyJSONValue } from "assistant-stream/utils";
+import { makeLogger } from "./runtime/logger";
+import type { UseAgUiRuntimeOptions } from "./runtime/types";
+import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
+
+export function useAgUiRuntime(
+  options: UseAgUiRuntimeOptions,
+): AssistantRuntime {
+  const logger = useMemo(() => makeLogger(options.logger), [options.logger]);
+  const [version, setVersion] = useState(0);
+  const notifyUpdate = useCallback(() => setVersion((v) => v + 1), []);
+  const coreRef = useRef<AgUiThreadRuntimeCore | null>(null);
+  const runtimeAdapters = useRuntimeAdapters();
+
+  const attachmentsAdapter =
+    options.adapters?.attachments ?? runtimeAdapters?.attachments;
+  const historyAdapter = options.adapters?.history ?? runtimeAdapters?.history;
+  const speechAdapter = options.adapters?.speech;
+  const feedbackAdapter = options.adapters?.feedback;
+
+  if (!coreRef.current) {
+    coreRef.current = new AgUiThreadRuntimeCore({
+      agent: options.agent,
+      logger,
+      showThinking: options.showThinking ?? true,
+      ...(options.onError ? { onError: options.onError } : {}),
+      ...(options.onCancel ? { onCancel: options.onCancel } : {}),
+      ...(historyAdapter ? { history: historyAdapter } : {}),
+      notifyUpdate,
+    });
+  }
+
+  const core = coreRef.current;
+  core.updateOptions({
+    agent: options.agent,
+    logger,
+    showThinking: options.showThinking ?? true,
+    ...(options.onError ? { onError: options.onError } : {}),
+    ...(options.onCancel ? { onCancel: options.onCancel } : {}),
+    ...(historyAdapter ? { history: historyAdapter } : {}),
+  });
+
+  const adapterAdapters = useMemo(() => {
+    const value: NonNullable<ExternalStoreAdapter<ThreadMessage>["adapters"]> =
+      {};
+    if (attachmentsAdapter) value.attachments = attachmentsAdapter;
+    if (speechAdapter) value.speech = speechAdapter;
+    if (feedbackAdapter) value.feedback = feedbackAdapter;
+    return Object.keys(value).length ? value : undefined;
+  }, [attachmentsAdapter, speechAdapter, feedbackAdapter]);
+
+  const store = useMemo(
+    () =>
+      ({
+        messages: core.getMessages(),
+        state: core.getState(),
+        isRunning: core.isRunning(),
+        onNew: (message: AppendMessage) => core.append(message),
+        onEdit: (message: AppendMessage) => core.edit(message),
+        onReload: (parentId: string | null, config: { runConfig?: any }) =>
+          core.reload(parentId, config),
+        onCancel: () => core.cancel(),
+        onAddToolResult: (options) => core.addToolResult(options),
+        onResume: (config) => core.resume(config),
+        setMessages: (messages: readonly ThreadMessage[]) =>
+          core.applyExternalMessages(messages),
+        onImport: (messages: readonly ThreadMessage[]) =>
+          core.applyExternalMessages(messages),
+        onLoadExternalState: (state: ReadonlyJSONValue) =>
+          core.loadExternalState(state),
+        adapters: adapterAdapters,
+      }) satisfies ExternalStoreAdapter<ThreadMessage>,
+    // version is intentionally included to trigger re-computation when core state changes via notifyUpdate
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [adapterAdapters, core, version],
+  );
+
+  const runtime = useExternalStoreRuntime(store);
+
+  useEffect(() => {
+    core.attachRuntime(runtime);
+    return () => {
+      core.detachRuntime();
+    };
+  }, [core, runtime]);
+
+  return runtime;
+}