assistant-cloud/src/FormattedCloudPersistence.ts

Cloud integration for assistant-ui

import type { ReadonlyJSONObject } from "assistant-stream/utils";
import type { CloudMessageRunPhase } from "./CloudMessagePersistence";

/**
 * Format adapter shape — structurally identical to the MessageFormatAdapter
 * in @assistant-ui/react, but defined here to avoid cross-package type moves.
 * TypeScript's structural typing ensures these are interchangeable.
 */
export type MessageFormatAdapter<TMessage, TStorageFormat> = {
  format: string;
  encode(item: { parentId: string | null; message: TMessage }): TStorageFormat;
  decode(stored: {
    id: string;
    parent_id: string | null;
    format: string;
    content: TStorageFormat;
  }): { parentId: string | null; message: TMessage };
  getId(message: TMessage): string;
};

/**
 * Wraps a CloudMessagePersistence with a MessageFormatAdapter's encode and
 * decode. The persistence parameter is typed structurally, so a caller does
 * not need to import the class.
 *
 * `append` and `update` take an optional run phase that is passed through to
 * the persistence: `running` opens (on append) or renews (on update) the
 * message's run lease, and an end reason makes the write the lease's final
 * write and ends it, the same as for `aui/v0` messages.
 */
export const createFormattedPersistence = <TMessage, TStorageFormat>(
  persistence: {
    append: (
      threadId: string,
      messageId: string,
      parentId: string | null,
      format: string,
      content: ReadonlyJSONObject,
      run?: CloudMessageRunPhase,
    ) => Promise<void>;
    load: (threadId: string, format?: string) => Promise<any[]>;
    isPersisted: (messageId: string) => boolean;
    update?: (
      threadId: string,
      messageId: string,
      format: string,
      content: ReadonlyJSONObject,
      run?: CloudMessageRunPhase,
    ) => Promise<void>;
    delete?: (threadId: string, messageIds: string[]) => Promise<void>;
  },
  adapter: MessageFormatAdapter<TMessage, TStorageFormat>,
) => ({
  append: async (
    threadId: string,
    item: { parentId: string | null; message: TMessage },
    run?: CloudMessageRunPhase,
  ): Promise<void> => {
    const messageId = adapter.getId(item.message);
    const encoded = adapter.encode(item);
    return persistence.append(
      threadId,
      messageId,
      item.parentId,
      adapter.format,
      encoded as ReadonlyJSONObject,
      ...(run !== undefined ? [run] : []),
    );
  },
  update: persistence.update
    ? async (
        threadId: string,
        item: { parentId: string | null; message: TMessage },
        messageId: string,
        run?: CloudMessageRunPhase,
      ): Promise<void> => {
        const encoded = adapter.encode(item);
        return persistence.update!(
          threadId,
          messageId,
          adapter.format,
          encoded as ReadonlyJSONObject,
          ...(run !== undefined ? [run] : []),
        );
      }
    : undefined,
  delete: persistence.delete
    ? async (
        threadId: string,
        items: { parentId: string | null; message: TMessage }[],
      ): Promise<void> =>
        persistence.delete!(
          threadId,
          items.map((item) => adapter.getId(item.message)),
        )
    : undefined,
  load: async (threadId: string) => {
    const messages = await persistence.load(threadId, adapter.format);
    return {
      messages: messages
        .filter((m) => m.format === adapter.format)
        .map((m) =>
          adapter.decode({
            id: m.id,
            parent_id: m.parent_id,
            format: m.format,
            content: m.content as TStorageFormat,
          }),
        )
        .reverse(),
    };
  },
  isPersisted: (messageId: string) => persistence.isPersisted(messageId),
});