@assistant-ui/core 0.1.0

package/src/runtime/utils/message-repository.ts

@@ -0,0 +1,390 @@
+import type { RunConfig, ThreadMessage } from "../../types";
+import { generateId, generateOptimisticId } from "../../utils/id";
+import type { ThreadMessageLike } from "./thread-message-like";
+import { getAutoStatus } from "./auto-status";
+import { fromThreadMessageLike } from "./thread-message-like";
+
+export type ExportedMessageRepositoryItem = {
+  message: ThreadMessage;
+  parentId: string | null;
+  runConfig?: RunConfig;
+};
+
+export type ExportedMessageRepository = {
+  headId?: string | null;
+  messages: Array<{
+    message: ThreadMessage;
+    parentId: string | null;
+    runConfig?: RunConfig;
+  }>;
+};
+
+export const ExportedMessageRepository = {
+  fromArray: (
+    messages: readonly ThreadMessageLike[],
+  ): ExportedMessageRepository => {
+    const conv = messages.map((m) =>
+      fromThreadMessageLike(
+        m,
+        generateId(),
+        getAutoStatus(false, false, false, false, undefined),
+      ),
+    );
+
+    return {
+      messages: conv.map((m, idx) => ({
+        parentId: idx > 0 ? conv[idx - 1]!.id : null,
+        message: m,
+      })),
+    };
+  },
+};
+
+type RepositoryParent = {
+  children: string[];
+  next: RepositoryMessage | null;
+};
+
+type RepositoryMessage = RepositoryParent & {
+  prev: RepositoryMessage | null;
+  current: ThreadMessage;
+  level: number;
+};
+
+const findHead = (
+  message: RepositoryMessage | RepositoryParent,
+): RepositoryMessage | null => {
+  if (message.next) return findHead(message.next);
+  if ("current" in message) return message;
+  return null;
+};
+
+class CachedValue<T> {
+  private _value: T | null = null;
+
+  constructor(private func: () => T) {}
+
+  get value() {
+    if (this._value === null) {
+      this._value = this.func();
+    }
+    return this._value;
+  }
+
+  dirty() {
+    this._value = null;
+  }
+}
+
+export class MessageRepository {
+  private messages = new Map<string, RepositoryMessage>();
+  private head: RepositoryMessage | null = null;
+  private root: RepositoryParent = {
+    children: [],
+    next: null,
+  };
+
+  private updateLevels(message: RepositoryMessage, newLevel: number) {
+    message.level = newLevel;
+    for (const childId of message.children) {
+      const childMessage = this.messages.get(childId);
+      if (childMessage) {
+        this.updateLevels(childMessage, newLevel + 1);
+      }
+    }
+  }
+
+  private performOp(
+    newParent: RepositoryMessage | null,
+    child: RepositoryMessage,
+    operation: "cut" | "link" | "relink",
+  ) {
+    const parentOrRoot = child.prev ?? this.root;
+    const newParentOrRoot = newParent ?? this.root;
+
+    if (operation === "relink" && parentOrRoot === newParentOrRoot) return;
+
+    if (operation !== "link") {
+      parentOrRoot.children = parentOrRoot.children.filter(
+        (m) => m !== child.current.id,
+      );
+
+      if (parentOrRoot.next === child) {
+        const fallbackId = parentOrRoot.children.at(-1);
+        const fallback = fallbackId ? this.messages.get(fallbackId) : null;
+        if (fallback === undefined) {
+          throw new Error(
+            "MessageRepository(performOp/cut): Fallback sibling message not found. This is likely an internal bug in assistant-ui.",
+          );
+        }
+        parentOrRoot.next = fallback;
+      }
+    }
+
+    if (operation !== "cut") {
+      for (
+        let current: RepositoryMessage | null = newParent;
+        current;
+        current = current.prev
+      ) {
+        if (current.current.id === child.current.id) {
+          throw new Error(
+            "MessageRepository(performOp/link): A message with the same id already exists in the parent tree. This error occurs if the same message id is found multiple times. This is likely an internal bug in assistant-ui.",
+          );
+        }
+      }
+
+      newParentOrRoot.children = [
+        ...newParentOrRoot.children,
+        child.current.id,
+      ];
+
+      if (findHead(child) === this.head || newParentOrRoot.next === null) {
+        newParentOrRoot.next = child;
+      }
+
+      child.prev = newParent;
+
+      const newLevel = newParent ? newParent.level + 1 : 0;
+      this.updateLevels(child, newLevel);
+    }
+  }
+
+  private _messages = new CachedValue<readonly ThreadMessage[]>(() => {
+    const messages = new Array<ThreadMessage>((this.head?.level ?? -1) + 1);
+    for (let current = this.head; current; current = current.prev) {
+      messages[current.level] = current.current;
+    }
+    return messages;
+  });
+
+  get headId() {
+    return this.head?.current.id ?? null;
+  }
+
+  getMessages(headId?: string) {
+    if (headId === undefined || headId === this.head?.current.id) {
+      return this._messages.value;
+    }
+
+    const headMessage = this.messages.get(headId);
+    if (!headMessage) {
+      throw new Error(
+        "MessageRepository(getMessages): Head message not found. This is likely an internal bug in assistant-ui.",
+      );
+    }
+
+    const messages = new Array<ThreadMessage>(headMessage.level + 1);
+    for (
+      let current: RepositoryMessage | null = headMessage;
+      current;
+      current = current.prev
+    ) {
+      messages[current.level] = current.current;
+    }
+    return messages;
+  }
+
+  addOrUpdateMessage(parentId: string | null, message: ThreadMessage) {
+    const existingItem = this.messages.get(message.id);
+    const prev = parentId ? this.messages.get(parentId) : null;
+    if (prev === undefined)
+      throw new Error(
+        "MessageRepository(addOrUpdateMessage): Parent message not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    if (existingItem) {
+      existingItem.current = message;
+      this.performOp(prev, existingItem, "relink");
+      this._messages.dirty();
+      return;
+    }
+
+    const newItem: RepositoryMessage = {
+      prev,
+      current: message,
+      next: null,
+      children: [],
+      level: prev ? prev.level + 1 : 0,
+    };
+
+    this.messages.set(message.id, newItem);
+    this.performOp(prev, newItem, "link");
+
+    if (this.head === prev) {
+      this.head = newItem;
+    }
+
+    this._messages.dirty();
+  }
+
+  getMessage(messageId: string) {
+    const message = this.messages.get(messageId);
+    if (!message)
+      throw new Error(
+        "MessageRepository(updateMessage): Message not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    return {
+      parentId: message.prev?.current.id ?? null,
+      message: message.current,
+      index: message.level,
+    };
+  }
+
+  appendOptimisticMessage(parentId: string | null, message: ThreadMessageLike) {
+    let optimisticId: string;
+    do {
+      optimisticId = generateOptimisticId();
+    } while (this.messages.has(optimisticId));
+
+    this.addOrUpdateMessage(
+      parentId,
+      fromThreadMessageLike(message, optimisticId, { type: "running" }),
+    );
+
+    return optimisticId;
+  }
+
+  deleteMessage(messageId: string, replacementId?: string | null | undefined) {
+    const message = this.messages.get(messageId);
+
+    if (!message)
+      throw new Error(
+        "MessageRepository(deleteMessage): Message not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    const replacement =
+      replacementId === undefined
+        ? message.prev
+        : replacementId === null
+          ? null
+          : this.messages.get(replacementId);
+    if (replacement === undefined)
+      throw new Error(
+        "MessageRepository(deleteMessage): Replacement not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    for (const child of message.children) {
+      const childMessage = this.messages.get(child);
+      if (!childMessage)
+        throw new Error(
+          "MessageRepository(deleteMessage): Child message not found. This is likely an internal bug in assistant-ui.",
+        );
+      this.performOp(replacement, childMessage, "relink");
+    }
+
+    this.performOp(null, message, "cut");
+    this.messages.delete(messageId);
+
+    if (this.head === message) {
+      this.head = findHead(replacement ?? this.root);
+    }
+
+    this._messages.dirty();
+  }
+
+  getBranches(messageId: string) {
+    const message = this.messages.get(messageId);
+    if (!message)
+      throw new Error(
+        "MessageRepository(getBranches): Message not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    const { children } = message.prev ?? this.root;
+    return children;
+  }
+
+  switchToBranch(messageId: string) {
+    const message = this.messages.get(messageId);
+    if (!message)
+      throw new Error(
+        "MessageRepository(switchToBranch): Branch not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    const prevOrRoot = message.prev ?? this.root;
+    prevOrRoot.next = message;
+
+    this.head = findHead(message);
+
+    this._messages.dirty();
+  }
+
+  resetHead(messageId: string | null) {
+    if (messageId === null) {
+      this.clear();
+      return;
+    }
+
+    const message = this.messages.get(messageId);
+    if (!message)
+      throw new Error(
+        "MessageRepository(resetHead): Branch not found. This is likely an internal bug in assistant-ui.",
+      );
+
+    if (message.children.length > 0) {
+      const deleteDescendants = (msg: RepositoryMessage) => {
+        for (const childId of msg.children) {
+          const childMessage = this.messages.get(childId);
+          if (childMessage) {
+            deleteDescendants(childMessage);
+            this.messages.delete(childId);
+          }
+        }
+      };
+      deleteDescendants(message);
+
+      message.children = [];
+      message.next = null;
+    }
+
+    this.head = message;
+    for (
+      let current: RepositoryMessage | null = message;
+      current;
+      current = current.prev
+    ) {
+      if (current.prev) {
+        current.prev.next = current;
+      } else {
+        this.root.next = current;
+      }
+    }
+
+    this._messages.dirty();
+  }
+
+  clear(): void {
+    this.messages.clear();
+    this.head = null;
+    this.root = {
+      children: [],
+      next: null,
+    };
+    this._messages.dirty();
+  }
+
+  export(): ExportedMessageRepository {
+    const exportItems: ExportedMessageRepository["messages"] = [];
+
+    for (const [, message] of this.messages) {
+      exportItems.push({
+        message: message.current,
+        parentId: message.prev?.current.id ?? null,
+      });
+    }
+
+    return {
+      headId: this.head?.current.id ?? null,
+      messages: exportItems,
+    };
+  }
+
+  import({ headId, messages }: ExportedMessageRepository) {
+    for (const { message, parentId } of messages) {
+      this.addOrUpdateMessage(parentId, message);
+    }
+
+    this.resetHead(headId ?? messages.at(-1)?.message.id ?? null);
+  }
+}

package/src/runtime/utils/thread-message-like.ts

@@ -0,0 +1,225 @@
+import { parsePartialJsonObject } from "assistant-stream/utils";
+import { generateId } from "../../utils/id";
+import type {
+  MessageTiming,
+  ReasoningMessagePart,
+  SourceMessagePart,
+  ThreadStep,
+  MessageStatus,
+  TextMessagePart,
+  ImageMessagePart,
+  ThreadMessage,
+  ThreadAssistantMessagePart,
+  ThreadAssistantMessage,
+  ThreadUserMessagePart,
+  ThreadUserMessage,
+  ThreadSystemMessage,
+  CompleteAttachment,
+  FileMessagePart,
+  DataMessagePart,
+  Unstable_AudioMessagePart,
+} from "../../types";
+import { ReadonlyJSONObject, ReadonlyJSONValue } from "assistant-stream/utils";
+
+export type ThreadMessageLike = {
+  readonly role: "assistant" | "user" | "system";
+  readonly content:
+    | string
+    | readonly (
+        | TextMessagePart
+        | ReasoningMessagePart
+        | SourceMessagePart
+        | ImageMessagePart
+        | FileMessagePart
+        | DataMessagePart
+        | Unstable_AudioMessagePart
+        | {
+            readonly type: "tool-call";
+            readonly toolCallId?: string;
+            readonly toolName: string;
+            readonly args?: ReadonlyJSONObject;
+            readonly argsText?: string;
+            readonly artifact?: any;
+            readonly result?: any | undefined;
+            readonly isError?: boolean | undefined;
+            readonly parentId?: string | undefined;
+            readonly messages?: readonly ThreadMessage[] | undefined;
+          }
+      )[];
+  readonly id?: string | undefined;
+  readonly createdAt?: Date | undefined;
+  readonly status?: MessageStatus | undefined;
+  readonly attachments?: readonly CompleteAttachment[] | undefined;
+  readonly metadata?:
+    | {
+        readonly unstable_state?: ReadonlyJSONValue;
+        readonly unstable_annotations?:
+          | readonly ReadonlyJSONValue[]
+          | undefined;
+        readonly unstable_data?: readonly ReadonlyJSONValue[] | undefined;
+        readonly steps?: readonly ThreadStep[] | undefined;
+        readonly timing?: MessageTiming | undefined;
+        readonly submittedFeedback?: { readonly type: "positive" | "negative" };
+        readonly custom?: Record<string, unknown> | undefined;
+      }
+    | undefined;
+};
+
+export const fromThreadMessageLike = (
+  like: ThreadMessageLike,
+  fallbackId: string,
+  fallbackStatus: MessageStatus,
+): ThreadMessage => {
+  const { role, id, createdAt, attachments, status, metadata } = like;
+  const common = {
+    id: id ?? fallbackId,
+    createdAt: createdAt ?? new Date(),
+  };
+
+  const content =
+    typeof like.content === "string"
+      ? [{ type: "text" as const, text: like.content }]
+      : like.content;
+
+  const sanitizeImageContent = ({
+    image,
+    ...rest
+  }: ImageMessagePart): ImageMessagePart | null => {
+    const match = image.match(
+      /^data:image\/(png|jpeg|jpg|gif|webp);base64,(.*)$/,
+    );
+    if (match) {
+      return { ...rest, image };
+    }
+    console.warn(`Invalid image data format detected`);
+    return null;
+  };
+
+  if (role !== "user" && attachments?.length)
+    throw new Error("attachments are only supported for user messages");
+
+  if (role !== "assistant" && status)
+    throw new Error("status is only supported for assistant messages");
+
+  if (role !== "assistant" && metadata?.steps)
+    throw new Error("metadata.steps is only supported for assistant messages");
+
+  switch (role) {
+    case "assistant":
+      return {
+        ...common,
+        role,
+        content: content
+          .map((part): ThreadAssistantMessagePart | null => {
+            const type = part.type;
+            switch (type) {
+              case "text":
+              case "reasoning":
+                if (part.text.trim().length === 0) return null;
+                return part;
+
+              case "file":
+              case "source":
+                return part;
+
+              case "image":
+                return sanitizeImageContent(part);
+
+              case "data":
+                return part;
+
+              case "tool-call": {
+                const { parentId, messages, ...basePart } = part;
+                const commonProps = {
+                  ...basePart,
+                  toolCallId: part.toolCallId ?? `tool-${generateId()}`,
+                  ...(parentId !== undefined && { parentId }),
+                  ...(messages !== undefined && { messages }),
+                };
+
+                if (part.args) {
+                  return {
+                    ...commonProps,
+                    args: part.args,
+                    argsText: part.argsText ?? JSON.stringify(part.args),
+                  };
+                }
+                return {
+                  ...commonProps,
+                  args: parsePartialJsonObject(part.argsText ?? "") ?? {},
+                  argsText: part.argsText ?? "",
+                };
+              }
+
+              default: {
+                const unhandledType: "audio" = type;
+                throw new Error(
+                  `Unsupported assistant message part type: ${unhandledType}`,
+                );
+              }
+            }
+          })
+          .filter((c) => !!c),
+        status: status ?? fallbackStatus,
+        metadata: {
+          unstable_state: metadata?.unstable_state ?? null,
+          unstable_annotations: metadata?.unstable_annotations ?? [],
+          unstable_data: metadata?.unstable_data ?? [],
+          custom: metadata?.custom ?? {},
+          steps: metadata?.steps ?? [],
+          ...(metadata?.timing && { timing: metadata.timing }),
+          ...(metadata?.submittedFeedback && {
+            submittedFeedback: metadata.submittedFeedback,
+          }),
+        },
+      } satisfies ThreadAssistantMessage;
+
+    case "user":
+      return {
+        ...common,
+        role,
+        content: content.map((part): ThreadUserMessagePart => {
+          const type = part.type;
+          switch (type) {
+            case "text":
+            case "image":
+            case "audio":
+            case "file":
+            case "data":
+              return part;
+
+            default: {
+              const unhandledType: "tool-call" | "reasoning" | "source" = type;
+              throw new Error(
+                `Unsupported user message part type: ${unhandledType}`,
+              );
+            }
+          }
+        }),
+        attachments: attachments ?? [],
+        metadata: {
+          custom: metadata?.custom ?? {},
+        },
+      } satisfies ThreadUserMessage;
+
+    case "system":
+      if (content.length !== 1 || content[0]!.type !== "text")
+        throw new Error(
+          "System messages must have exactly one text message part.",
+        );
+
+      return {
+        ...common,
+        role,
+        content: content as [TextMessagePart],
+        metadata: {
+          custom: metadata?.custom ?? {},
+        },
+      } satisfies ThreadSystemMessage;
+
+    default: {
+      const unsupportedRole: never = role;
+      throw new Error(`Unknown message role: ${unsupportedRole}`);
+    }
+  }
+};

package/src/runtimes/assistant-transport/utils.ts

@@ -0,0 +1,37 @@
+import { Tool, toToolsJSONSchema, type ToolJSONSchema } from "assistant-stream";
+
+/**
+ * @deprecated Use `toToolsJSONSchema` from `assistant-stream` instead.
+ */
+export function toAISDKTools(
+  tools: Record<string, Tool>,
+): Record<string, ToolJSONSchema> {
+  return toToolsJSONSchema(tools, { filter: () => true });
+}
+
+/**
+ * @deprecated Use `toToolsJSONSchema` from `assistant-stream` instead, which includes filtering by default.
+ */
+export function getEnabledTools(
+  tools: Record<string, Tool>,
+): Record<string, Tool> {
+  return Object.fromEntries(
+    Object.entries(tools).filter(
+      ([, tool]) => !tool.disabled && tool.type !== "backend",
+    ),
+  );
+}
+
+export async function createRequestHeaders(
+  headersValue:
+    | Record<string, string>
+    | Headers
+    | (() => Promise<Record<string, string> | Headers>),
+): Promise<Headers> {
+  const resolvedHeaders =
+    typeof headersValue === "function" ? await headersValue() : headersValue;
+
+  const headers = new Headers(resolvedHeaders);
+  headers.set("Content-Type", "application/json");
+  return headers;
+}