@zing-protocol/zing-sdk 0.0.1

package/src/mutations/walrus.ts

@@ -0,0 +1,862 @@
+import { signAndExecuteGasPoolTransaction } from "./signer.js";
+import {
+  certifyArticleBlob,
+  finalizePublishArticle,
+  startPublishArticle,
+} from "../_generated/zing_studio/app.js";
+import {
+  articleBlobAddMetadata,
+  articleBlobInsertOrUpdateMetadataPair,
+  articleBlobRemoveMetadata,
+} from "../_generated/zing_studio/article.js";
+import { _new as newMetadata } from "../_generated/zing_studio/deps/walrus/metadata.js";
+import {
+  encodedBlobLength,
+  encodeQuiltPatchId,
+  getDerivedStorageID,
+  getDerivedStudioID,
+  getPublishedArticleAndBlobId,
+  signersToBitmap,
+} from "../getters.js";
+import { encryptData } from "../mutations/seal.js";
+import { bcs } from "@mysten/sui/bcs";
+import { Transaction } from "@mysten/sui/transactions";
+import { blobIdToInt } from "@mysten/walrus";
+import type { PublishArticleEvent } from "../_generated/zing_studio/article.js";
+import type { WalrusUploadFlowStage } from "../config/common.js";
+import type { SharedObjectRef, WalrusSystemState } from "../config/types.js";
+import type { ZKLoginSignerState } from "../types.js";
+import type { SuiGrpcClient } from "@mysten/sui/grpc";
+import type { SuiJsonRpcClient } from "@mysten/sui/jsonRpc";
+import type {
+  CertifyBlobOptions,
+  EncodingType,
+  WalrusClient,
+  WalrusFile,
+  WriteFilesFlowOptions,
+  WriteFilesFlowRegisterOptions,
+  WriteFilesFlowUploadOptions,
+} from "@mysten/walrus";
+
+const REDSTUFF_CODING_TYPE = 1;
+// ZING_PREFIX
+const FIXED_FILE_IV = new Uint8Array([
+  4, 122, 105, 110, 103, 0, 0, 0, 0, 0, 0, 0,
+]);
+const BYTES_PER_UNIT_SIZE = 1024 * 1024;
+
+type ZingWriteFlowConfig = {
+  zingStudioPackageId: string;
+  zingStudioV0PackageId: string;
+  walrusPackageId: string;
+  zingStudioConfigSharedObjectRef: SharedObjectRef;
+  walrusSystemSharedObjectRef: SharedObjectRef;
+  zingStorageTreasurySharedObjectRef: SharedObjectRef;
+  blobRegisteredEventType: string;
+  publishArticleEventType: string;
+};
+
+export function buildPublishArticleTransaction(
+  config: ZingWriteFlowConfig,
+  walrusSystem: WalrusSystemState,
+  walTreasury: string,
+  storageByEpoch: Record<number, number>,
+  fileSize: number,
+  owner: string,
+  subscriptionLevel: null | number,
+  blobId: string,
+  rootHash: Uint8Array,
+  articleMetadata: {
+    identifier: string;
+    tags: {
+      "content-type": string;
+    };
+    size: number;
+  }[],
+  attributes: Record<string, string | null>,
+  existingAttributes: Record<string, string> | null,
+) {
+  const {
+    zingStudioPackageId,
+    walrusPackageId,
+    zingStudioConfigSharedObjectRef,
+    walrusSystemSharedObjectRef,
+    zingStorageTreasurySharedObjectRef,
+  } = config;
+  return (tx: Transaction) => {
+    const { encodedSize, writeCost } = storageCost(walrusSystem, fileSize, 2);
+
+    console.log({ encodedSize });
+    const currentEpoch = walrusSystem.committee.epoch;
+
+    if (!storageByEpoch?.[currentEpoch]) {
+      throw new Error(`No Epochs Found for ${currentEpoch}`);
+    }
+
+    if (!storageByEpoch?.[currentEpoch + 1]) {
+      throw new Error(`No Epochs Found for ${currentEpoch + 1}`);
+    }
+
+    if (storageByEpoch[currentEpoch] < encodedSize) {
+      throw new Error(
+        `Insufficient Storages size ${storageByEpoch[currentEpoch]} for Epochs ${currentEpoch} with required size ${encodedSize}`,
+      );
+    }
+
+    if (storageByEpoch[currentEpoch + 1] < encodedSize) {
+      throw new Error(
+        `Insufficient Storages size ${storageByEpoch[currentEpoch + 1]} for Epochs ${currentEpoch + 1} with required size ${encodedSize}`,
+      );
+    }
+
+    if (BigInt(walTreasury) < writeCost) {
+      throw new Error(
+        `Insufficient walTreasury Balance ${walTreasury} for writeCost: ${writeCost}`,
+      );
+    }
+
+    console.log({ fileSize, encodedSize });
+    console.log({ articleMetadata });
+    const [article, publishReceipt] = tx.add(
+      startPublishArticle({
+        package: zingStudioPackageId,
+        arguments: {
+          config: tx.sharedObjectRef(zingStudioConfigSharedObjectRef),
+          studio: tx.object(
+            getDerivedStudioID(zingStudioConfigSharedObjectRef.objectId, owner),
+          ),
+          walrusSystem: tx.sharedObjectRef(walrusSystemSharedObjectRef),
+          storageTreasury: tx.sharedObjectRef(
+            zingStorageTreasurySharedObjectRef,
+          ),
+          storageSpace: tx.object(
+            getDerivedStorageID(
+              zingStudioConfigSharedObjectRef.objectId,
+              owner,
+            ),
+          ),
+          subscriptionLevel: tx.pure.option("u8", subscriptionLevel),
+          blobId: blobIdToInt(blobId),
+          rootHash: BigInt(bcs.u256().parse(rootHash)),
+          unencodedBlobSize: fileSize,
+          encodingType: REDSTUFF_CODING_TYPE,
+          identifiers: tx.pure.vector(
+            "string",
+            articleMetadata.map((m) => m.identifier),
+          ),
+          blobIndexes: tx.pure.vector(
+            "u64",
+            articleMetadata.map(() => 0),
+          ),
+          mimeTypes: tx.pure.vector(
+            "string",
+            articleMetadata.map((m) => m.tags["content-type"]),
+          ),
+          metadataSizes: tx.pure.vector(
+            "u64",
+            articleMetadata.map((m) => m.size),
+          ),
+        },
+      }),
+    );
+
+    if (!existingAttributes) {
+      tx.add(
+        articleBlobAddMetadata({
+          package: zingStudioPackageId,
+          arguments: {
+            config: tx.sharedObjectRef(zingStudioConfigSharedObjectRef),
+            self: article,
+            blobIndex: tx.pure.u64(0),
+            metadata: newMetadata({
+              package: walrusPackageId,
+            }),
+          },
+        }),
+      );
+    }
+
+    Object.keys(attributes).forEach((key) => {
+      const value = attributes[key];
+
+      if (value === null) {
+        if (existingAttributes && key in existingAttributes) {
+          tx.add(
+            articleBlobRemoveMetadata({
+              package: zingStudioPackageId,
+              arguments: {
+                config: tx.sharedObjectRef(zingStudioConfigSharedObjectRef),
+                self: article,
+                blobIndex: tx.pure.u64(0),
+                key,
+              },
+            }),
+          );
+        }
+      } else {
+        tx.add(
+          articleBlobInsertOrUpdateMetadataPair({
+            package: zingStudioPackageId,
+            arguments: {
+              config: tx.sharedObjectRef(zingStudioConfigSharedObjectRef),
+              self: article,
+              blobIndex: tx.pure.u64(0),
+              key,
+              value,
+            },
+          }),
+        );
+      }
+    });
+
+    // put article to studio
+    tx.add(
+      finalizePublishArticle({
+        package: zingStudioPackageId,
+        arguments: {
+          config: tx.sharedObjectRef(zingStudioConfigSharedObjectRef),
+          studio: tx.object(
+            getDerivedStudioID(zingStudioConfigSharedObjectRef.objectId, owner),
+          ),
+          receipt: publishReceipt,
+          article: article,
+        },
+      }),
+    );
+  };
+}
+
+export function buildCertifyArticleBlobTransaction(
+  config: {
+    zingStudioPackageId: string;
+    zingStudioConfigSharedObjectRef: SharedObjectRef;
+    walrusSystemSharedObjectRef: SharedObjectRef;
+  },
+  walrusClient: WalrusClient,
+  owner: string,
+  articleId: string,
+  {
+    blobId,
+    blobObjectId,
+    confirmations,
+    certificate,
+    deletable,
+  }: CertifyBlobOptions,
+) {
+  const {
+    zingStudioPackageId,
+    zingStudioConfigSharedObjectRef,
+    walrusSystemSharedObjectRef,
+  } = config;
+  return async (tx: Transaction) => {
+    const systemState = await walrusClient.systemState();
+    const combinedSignature =
+      certificate ??
+      (await walrusClient.certificateFromConfirmations({
+        confirmations,
+        blobId,
+        deletable,
+        blobObjectId,
+      }));
+    tx.add(
+      certifyArticleBlob({
+        package: zingStudioPackageId,
+        arguments: {
+          config: tx.sharedObjectRef(zingStudioConfigSharedObjectRef),
+          studio: tx.object(
+            getDerivedStudioID(zingStudioConfigSharedObjectRef.objectId, owner),
+          ),
+          walrusSystem: tx.sharedObjectRef(walrusSystemSharedObjectRef),
+          articleId: tx.object(articleId),
+          blobIndex: tx.pure.u64(0), // TODO
+          signature: tx.pure.vector("u8", combinedSignature.signature),
+          signersBitmap: tx.pure.vector(
+            "u8",
+            signersToBitmap(
+              combinedSignature.signers,
+              systemState.committee.members.length,
+            ),
+          ),
+          message: tx.pure.vector("u8", combinedSignature.serializedMessage),
+        },
+      }),
+    );
+  };
+}
+
+export const zingWriteFlow = (
+  config: ZingWriteFlowConfig,
+  suiGrpcClient: SuiGrpcClient,
+  suiJsonRpcClient: SuiJsonRpcClient,
+  walrusClient: WalrusClient,
+  walrusSystem: WalrusSystemState,
+  walTreasury: string,
+  storageByEpoch: Record<number, number>,
+  fileKey: CryptoKey,
+  studioObjectId: string,
+  owner: string,
+  subscriptionLevel: null | number,
+  articleMetadata: {
+    identifier: string;
+    tags: {
+      "content-type": string;
+    };
+    size: number;
+  }[],
+  options: WriteFilesFlowOptions,
+  withUploadRelayClient: boolean,
+  blobRegisteredEventType: string,
+  publishArticleEventType: string,
+) => {
+  const { files } = options;
+
+  const encode = async () => {
+    const { quilt, index } = await walrusClient.encodeQuilt({
+      blobs: await Promise.all(
+        files.map(async (file, i) => {
+          const contents = await file.bytes();
+          const bytes = await encryptData(fileKey, contents, FIXED_FILE_IV);
+          return {
+            contents: bytes,
+            identifier: (await file.getIdentifier()) ?? `file-${i}`,
+            tags: (await file.getTags()) ?? {},
+          };
+        }),
+      ),
+    });
+
+    console.log({ walrus_system: await walrusClient.systemState() });
+    console.log({ quiltlength: quilt.length });
+
+    // store encryptedBytes
+    const metadata = withUploadRelayClient
+      ? await walrusClient.computeBlobMetadata({
+          bytes: quilt,
+        })
+      : await walrusClient.encodeBlob(quilt);
+
+    return {
+      metadata,
+      size: quilt.length,
+      data: withUploadRelayClient ? quilt : undefined,
+      index,
+    };
+  };
+
+  const register = (
+    { data, metadata, index, size }: Awaited<ReturnType<typeof encode>>,
+    { deletable, owner, attributes }: WriteFilesFlowRegisterOptions,
+  ) => {
+    const transaction = new Transaction();
+    // transaction.setSenderIfNotSet(owner);
+
+    if (withUploadRelayClient) {
+      const meta = metadata as Awaited<
+        ReturnType<typeof walrusClient.computeBlobMetadata>
+      >;
+      transaction.add(
+        walrusClient.sendUploadRelayTip({
+          size,
+          blobDigest: meta.blobDigest,
+          nonce: meta.nonce,
+        }),
+      );
+    }
+
+    // 1. rent storage space and publish article
+    transaction.add(
+      buildPublishArticleTransaction(
+        config,
+        walrusSystem,
+        walTreasury,
+        storageByEpoch,
+        size,
+        owner,
+        subscriptionLevel,
+        metadata.blobId,
+        metadata.rootHash,
+        articleMetadata,
+        { _walrusBlobType: "quilt", ...attributes },
+        null,
+      ),
+    );
+
+    console.log({ registerTransaction: transaction });
+
+    return {
+      registerTransaction: transaction,
+      index,
+      data,
+      metadata,
+      deletable,
+    };
+  };
+
+  const upload = async (
+    { index, data, metadata, deletable }: Awaited<ReturnType<typeof register>>,
+    { digest }: WriteFilesFlowUploadOptions,
+  ) => {
+    const { parsedPublishArticleEvent, blobObjectId } =
+      await getPublishedArticleAndBlobId(
+        suiGrpcClient,
+        suiJsonRpcClient,
+        walrusClient,
+        digest,
+        blobRegisteredEventType,
+        publishArticleEventType,
+      );
+    const articleId = parsedPublishArticleEvent.article_id;
+
+    if (withUploadRelayClient) {
+      const meta = metadata as Awaited<
+        ReturnType<typeof walrusClient.computeBlobMetadata>
+      >;
+      return {
+        index,
+        parsedPublishArticleEvent,
+        articleId,
+        blobObjectId,
+        metadata,
+        deletable,
+        certificate: (
+          await walrusClient.writeBlobToUploadRelay({
+            blobId: metadata.blobId,
+            blob: data!,
+            nonce: meta.nonce,
+            txDigest: digest,
+            blobObjectId,
+            deletable,
+            encodingType: meta.metadata.encodingType as EncodingType,
+          })
+        ).certificate,
+      };
+    }
+
+    const meta = metadata as Awaited<
+      ReturnType<typeof walrusClient.encodeBlob>
+    >;
+
+    return {
+      index,
+      blobObjectId,
+      articleId,
+      parsedPublishArticleEvent,
+      metadata,
+      deletable,
+      confirmations: await walrusClient.writeEncodedBlobToNodes({
+        blobId: metadata.blobId,
+        objectId: blobObjectId,
+        metadata: meta.metadata,
+        sliversByNode: meta.sliversByNode,
+        deletable,
+      }),
+    };
+  };
+
+  const certify = ({
+    index,
+    metadata,
+    confirmations,
+    certificate,
+    articleId,
+    blobObjectId,
+    deletable,
+  }: Awaited<ReturnType<typeof upload>>) => {
+    console.log({ confirmations });
+    const tx = new Transaction();
+
+    if (confirmations) {
+      tx.add(
+        buildCertifyArticleBlobTransaction(
+          {
+            zingStudioPackageId: config.zingStudioPackageId,
+            zingStudioConfigSharedObjectRef:
+              config.zingStudioConfigSharedObjectRef,
+            walrusSystemSharedObjectRef: config.walrusSystemSharedObjectRef,
+          },
+          walrusClient,
+          owner,
+          articleId,
+          {
+            certificate,
+            blobId: metadata.blobId,
+            blobObjectId,
+            confirmations,
+            deletable,
+          },
+        ),
+      );
+    } else {
+      tx.add(
+        buildCertifyArticleBlobTransaction(
+          {
+            zingStudioPackageId: config.zingStudioPackageId,
+            zingStudioConfigSharedObjectRef:
+              config.zingStudioConfigSharedObjectRef,
+            walrusSystemSharedObjectRef: config.walrusSystemSharedObjectRef,
+          },
+          walrusClient,
+          owner,
+          articleId,
+          {
+            certificate,
+            blobId: metadata.blobId,
+            blobObjectId,
+            deletable,
+          },
+        ),
+      );
+    }
+    return {
+      index,
+      blobObjectId,
+      metadata,
+      transaction: tx,
+    };
+  };
+
+  async function listFiles({
+    index,
+    blobObjectId,
+    metadata,
+  }: Awaited<ReturnType<typeof certify>>) {
+    return index.patches.map((patch) => ({
+      id: encodeQuiltPatchId({
+        quiltId: metadata.blobId,
+        patchId: {
+          version: 1,
+          startIndex: patch.startIndex,
+          endIndex: patch.endIndex,
+        },
+      }),
+      blobId: metadata.blobId,
+      blobObjectId,
+    }));
+  }
+
+  const stepResults: {
+    encode?: Awaited<ReturnType<typeof encode>>;
+    register?: Awaited<ReturnType<typeof register>>;
+    upload?: Awaited<ReturnType<typeof upload>>;
+    certify?: Awaited<ReturnType<typeof certify>>;
+    listFiles?: never;
+  } = {};
+
+  function getResults<T extends keyof typeof stepResults>(
+    step: T,
+    current: keyof typeof stepResults,
+  ): NonNullable<(typeof stepResults)[T]> {
+    if (!stepResults[step]) {
+      throw new Error(`${step} must be executed before calling ${current}`);
+    }
+    return stepResults[step];
+  }
+
+  return {
+    encode: async () => {
+      if (!stepResults.encode) {
+        stepResults.encode = await encode();
+      }
+    },
+    register: (options: WriteFilesFlowRegisterOptions) => {
+      stepResults.register = register(
+        getResults("encode", "register"),
+        options,
+      );
+      return stepResults.register.registerTransaction;
+    },
+    upload: async (options: WriteFilesFlowUploadOptions) => {
+      stepResults.upload = await upload(
+        getResults("register", "upload"),
+        options,
+      );
+      return stepResults.upload.parsedPublishArticleEvent;
+    },
+    certify: () => {
+      stepResults.certify = certify(getResults("upload", "certify"));
+      return stepResults.certify.transaction;
+    },
+    listFiles: async () => listFiles(getResults("certify", "listFiles")),
+  };
+};
+
+export async function uploadFileToWalrus(
+  zingWriteFlowConfig: ZingWriteFlowConfig,
+  suiGrpcClient: SuiGrpcClient,
+  suiJsonRpcClient: SuiJsonRpcClient,
+  walrusClient: WalrusClient,
+  walrusSystem: WalrusSystemState,
+  walTreasury: string,
+  storageByEpoch: Record<number, number>,
+  fileKey: CryptoKey,
+  walrusFiles: WalrusFile[],
+  subscriptionLevel: null | number,
+  metadata: {
+    identifier: string;
+    tags: {
+      "content-type": string;
+    };
+    size: number;
+  }[],
+  signer: ZKLoginSignerState,
+  gasPoolBearerToken: string,
+  onProgress?: (stage: WalrusUploadFlowStage) => void,
+): Promise<typeof PublishArticleEvent.$inferType> {
+  const report = (stage: WalrusUploadFlowStage) => {
+    if (onProgress) onProgress(stage);
+  };
+  const owner = signer.suiAddress;
+
+  if (!owner) throw new Error("no signer");
+  // Step 1: Create and encode the flow (can be done immediately when file is selected)
+  const flow = zingWriteFlow(
+    zingWriteFlowConfig,
+    suiGrpcClient,
+    suiJsonRpcClient,
+    walrusClient,
+    walrusSystem,
+    walTreasury,
+    storageByEpoch,
+    fileKey,
+    getDerivedStudioID(
+      zingWriteFlowConfig.zingStudioConfigSharedObjectRef.objectId,
+      owner,
+    ),
+    owner,
+    subscriptionLevel,
+    metadata,
+    {
+      files: walrusFiles,
+    },
+    true,
+    zingWriteFlowConfig.blobRegisteredEventType,
+    zingWriteFlowConfig.publishArticleEventType,
+  );
+
+  report("encode");
+  await flow.encode();
+
+  // Step 2: Register the blob (triggered by user clicking a register button after the encode step)
+  async function handleRegister() {
+    if (!owner) throw new Error("suiaddress not found in zkloginsignerstate");
+
+    report("register_blob");
+    const registerTx = flow.register({
+      epochs: 1,
+      owner,
+      deletable: true,
+    });
+
+    const registerTxEffects = await signAndExecuteGasPoolTransaction(
+      signer,
+      suiJsonRpcClient,
+      registerTx,
+    );
+
+    if (!registerTxEffects?.transactionDigest)
+      throw new Error("fail to execute register Blob transaction");
+
+    const finalizedTransaction = await suiJsonRpcClient.waitForTransaction({
+      digest: registerTxEffects.transactionDigest,
+    });
+    // Step 3: Upload the data to storage nodes
+    // This can be done immediately after the register step, or as a separate step the user initiates
+    report("upload");
+    return await flow.upload({
+      digest: finalizedTransaction.digest,
+    });
+  }
+
+  const parsedPublishArticleEvent = await handleRegister();
+
+  // Step 4: Certify the blob (triggered by user clicking a certify button after the blob is uploaded)
+  async function handleCertify() {
+    report("certify_blob");
+    const certifyTx = flow.certify();
+
+    const certifyTxEffects = await signAndExecuteGasPoolTransaction(
+      signer,
+      suiJsonRpcClient,
+      certifyTx,
+    );
+
+    if (!certifyTxEffects?.transactionDigest)
+      throw new Error("fail to execute certify Blob transaction");
+
+    await suiJsonRpcClient.waitForTransaction({
+      digest: certifyTxEffects.transactionDigest,
+    });
+
+    report("publish");
+  }
+
+  await handleCertify();
+
+  return parsedPublishArticleEvent;
+}
+
+export async function downloadImagesFromMarkdown<
+  T extends "blob" | "uint8array",
+>(
+  markdown: string,
+  as: T = "blob" as T,
+): Promise<{ url: string; data: T extends "blob" ? Blob : Uint8Array }[]> {
+  // Match ![alt](url) or [![alt](url)](...)
+  const imageRegex = /!\[[^\]]*\]\(([^)]+)\)/g;
+  const urls = new Set<string>();
+  let match;
+
+  while ((match = imageRegex.exec(markdown)) !== null) {
+    let url = match[1].trim();
+
+    // Remove optional Markdown title e.g. (url "title")
+    const spaceIndex = url.indexOf(" ");
+    if (spaceIndex !== -1 && url[0] !== '"') {
+      url = url.slice(0, spaceIndex);
+    }
+
+    // Include both HTTP URLs and data URLs (base64 encoded images)
+    if (
+      url.startsWith("http") ||
+      url.startsWith("https") ||
+      url.startsWith("data:")
+    ) {
+      urls.add(url);
+    }
+  }
+
+  const images: { url: string; data: T extends "blob" ? Blob : Uint8Array }[] =
+    [];
+
+  for (const url of urls) {
+    try {
+      if (url.startsWith("data:")) {
+        // Handle base64 encoded data URLs
+        const response = await fetch(url);
+        if (!response.ok) throw new Error(`Failed to process data URL`);
+
+        if (as === "blob") {
+          const blob = await response.blob();
+          images.push({
+            url,
+            data: blob as T extends "blob" ? Blob : Uint8Array,
+          });
+        } else {
+          const buffer = await response.arrayBuffer();
+          images.push({
+            url,
+            data: new Uint8Array(buffer) as T extends "blob"
+              ? Blob
+              : Uint8Array,
+          });
+        }
+      } else {
+        // Handle regular HTTP/HTTPS URLs
+        const response = await fetch(url);
+        if (!response.ok) throw new Error(`Failed to fetch ${url}`);
+
+        if (as === "blob") {
+          const blob = await response.blob();
+          images.push({
+            url,
+            data: blob as T extends "blob" ? Blob : Uint8Array,
+          });
+        } else {
+          const buffer = await response.arrayBuffer();
+          images.push({
+            url,
+            data: new Uint8Array(buffer) as T extends "blob"
+              ? Blob
+              : Uint8Array,
+          });
+        }
+      }
+    } catch (err) {
+      console.error(
+        `Error processing image from ${url.startsWith("data:") ? "data URL" : url}:`,
+        err,
+      );
+    }
+  }
+
+  return images;
+}
+/**
+ * Calculate the cost of storing a blob for a given a size and number of epochs.
+ */
+export function storageCost(
+  systemState: WalrusSystemState,
+  size: number,
+  epochs: number,
+) {
+  const encodedSize = encodedBlobLength(size, systemState.committee.n_shards);
+  const storageUnits = storageUnitsFromSize(encodedSize);
+  const storageCost =
+    BigInt(storageUnits) *
+    BigInt(systemState.storage_price_per_unit_size) *
+    BigInt(epochs);
+  BigInt(epochs);
+
+  const writeCost =
+    BigInt(storageUnits) * BigInt(systemState.write_price_per_unit_size);
+
+  return {
+    encodedSize,
+    storageCost,
+    writeCost,
+    totalCost: storageCost + writeCost,
+  };
+}
+
+export function storageUnitsFromSize(size: number): number {
+  return Math.ceil(size / BYTES_PER_UNIT_SIZE);
+}
+
+/**
+ * Replace all base64 image URLs in markdown with asset identifiers (e.g. assets-0.png).
+ *
+ * @param markdown - the original markdown string containing data URLs
+ * @param assets - the array returned from downloadImagesFromMarkdown()
+ * @returns a new markdown string with data URLs replaced by identifiers
+ */
+export function replaceBase64ImagesWithIdentifiers(
+  markdown: string,
+  assets: { url: string; data: Blob | Uint8Array }[],
+): string {
+  let replacedMarkdown = markdown;
+
+  // Regex to match ![alt](url) markdown image syntax
+  const imageRegex = /!\[([^\]]*)\]\(([^)]+)\)/g;
+
+  replacedMarkdown = replacedMarkdown.replace(imageRegex, (match, alt, url) => {
+    url = url.trim();
+
+    // Find the matching asset by comparing the original URL
+    const assetIndex = assets.findIndex((asset) => asset.url === url);
+
+    if (assetIndex === -1) {
+      // No matching asset found, keep original
+      return match;
+    }
+
+    const asset = assets[assetIndex];
+
+    // Try to infer extension from mime type
+    let extension = ".bin";
+    if (asset.data instanceof Blob) {
+      const type = asset.data.type;
+      if (type.includes("png")) extension = ".png";
+      else if (type.includes("jpeg") || type.includes("jpg"))
+        extension = ".jpg";
+      else if (type.includes("gif")) extension = ".gif";
+      else if (type.includes("webp")) extension = ".webp";
+      else if (type.includes("svg")) extension = ".svg";
+    }
+
+    // Replace with new identifier reference
+    const newIdentifier = `assets-${assetIndex}${extension}`;
+    return `![${alt}](${newIdentifier})`;
+  });
+
+  return replacedMarkdown;
+}