refactor(bot): share image attachment uploads

apps/web/src/lib/bot/images.ts

@@ -1,13 +1,10 @@
-import { PutObjectCommand } from '@aws-sdk/client-s3';
 import {
   CLOUD_AGENT_IMAGE_ALLOWED_TYPES,
   CLOUD_AGENT_IMAGE_MAX_COUNT,
-  CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES,
-  CLOUD_AGENT_IMAGE_MIME_TO_EXTENSION,
   type CloudAgentImageAllowedType,
 } from '@/lib/cloud-agent/constants';
 import type { Images } from '@/lib/images-schema';
-import { r2Client, r2CloudAgentAttachmentsBucketName } from '@/lib/r2/client';
+import { assertImageAttachmentSize, uploadImageAttachment } from '@/lib/r2/cloud-agent-attachments';
 import { captureException } from '@sentry/nextjs';
 import type { Attachment, Message } from 'chat';
 import { randomUUID } from 'crypto';
@@ -53,37 +50,26 @@ export async function extractAndUploadImages(
   for (const attachment of toProcess) {
     try {
       const imageId = randomUUID();
-      const ext = CLOUD_AGENT_IMAGE_MIME_TO_EXTENSION[attachment.mimeType];
-      const filename = `${imageId}.${ext}`;
-      const r2Key = `${userId}/cloud-agent/${messageUuid}/${filename}`;
 
-      if (
-        typeof attachment.size === 'number' &&
-        attachment.size > CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES
-      ) {
-        throw new Error(
-          `Image ${attachment.name ?? filename} exceeds ${CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES / (1024 * 1024)}MB limit (${(attachment.size / (1024 * 1024)).toFixed(1)}MB)`
-        );
+      if (typeof attachment.size === 'number') {
+        assertImageAttachmentSize({
+          service: 'cloud-agent',
+          contentLength: attachment.size,
+          name: attachment.name ?? imageId,
+        });
       }
 
       const data = await attachment.fetchData();
 
-      if (data.byteLength > CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES) {
-        throw new Error(
-          `Image ${attachment.name ?? filename} exceeds ${CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES / (1024 * 1024)}MB limit (${(data.byteLength / (1024 * 1024)).toFixed(1)}MB)`
-        );
-      }
-
-      await r2Client.send(
-        new PutObjectCommand({
-          Bucket: r2CloudAgentAttachmentsBucketName,
-          Key: r2Key,
-          Body: data,
-          ContentType: attachment.mimeType,
-          ContentLength: data.byteLength,
-          Metadata: { userId, messageUuid, imageId },
-        })
-      );
+      const { filename } = await uploadImageAttachment({
+        service: 'cloud-agent',
+        userId,
+        messageUuid,
+        imageId,
+        contentType: attachment.mimeType,
+        body: data,
+        name: attachment.name,
+      });
 
       filenames.push(filename);
     } catch (error) {

apps/web/src/lib/r2/cloud-agent-attachments.ts

@@ -1,7 +1,9 @@
 import { PutObjectCommand } from '@aws-sdk/client-s3';
 import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
+import { APP_BUILDER_IMAGE_MAX_SIZE_BYTES } from '@/lib/app-builder/constants';
 import {
   CLOUD_AGENT_IMAGE_MIME_TO_EXTENSION,
+  CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES,
   CLOUD_AGENT_IMAGE_PRESIGNED_URL_EXPIRY_SECONDS,
 } from '@/lib/cloud-agent/constants';
 import type { CloudAgentImageAllowedType } from '@/lib/cloud-agent/constants';
@@ -13,15 +15,57 @@ function getExtensionFromContentType(contentType: CloudAgentImageAllowedType): s
   return CLOUD_AGENT_IMAGE_MIME_TO_EXTENSION[contentType];
 }
 
-function getImageKey(
-  service: Service,
-  userId: string,
-  messageUuid: string,
-  imageId: string,
-  contentType: CloudAgentImageAllowedType
-): string {
+function getMaxSizeBytes(service: Service): number {
+  return service === 'app-builder'
+    ? APP_BUILDER_IMAGE_MAX_SIZE_BYTES
+    : CLOUD_AGENT_IMAGE_MAX_SIZE_BYTES;
+}
+
+export function assertImageAttachmentSize({
+  service,
+  contentLength,
+  name,
+}: {
+  service: Service;
+  contentLength: number;
+  name?: string;
+}): void {
+  const maxSizeBytes = getMaxSizeBytes(service);
+
+  if (contentLength <= maxSizeBytes) return;
+
+  throw new Error(
+    `Image ${name ?? 'attachment'} exceeds ${maxSizeBytes / (1024 * 1024)}MB limit (${(contentLength / (1024 * 1024)).toFixed(1)}MB)`
+  );
+}
+
+export type GetImageAttachmentPathParams = {
+  service: Service;
+  userId: string;
+  messageUuid: string;
+  imageId: string;
+  contentType: CloudAgentImageAllowedType;
+};
+
+export type GetImageAttachmentPathResult = {
+  filename: string;
+  key: string;
+};
+
+export function getImageAttachmentPath({
+  service,
+  userId,
+  messageUuid,
+  imageId,
+  contentType,
+}: GetImageAttachmentPathParams): GetImageAttachmentPathResult {
   const ext = getExtensionFromContentType(contentType);
-  return `${userId}/${service}/${messageUuid}/${imageId}.${ext}`;
+  const filename = `${imageId}.${ext}`;
+
+  return {
+    filename,
+    key: `${userId}/${service}/${messageUuid}/${filename}`,
+  };
 }
 
 export type GenerateImageUploadUrlParams = {
@@ -39,17 +83,33 @@ export type GenerateImageUploadUrlResult = {
   expiresAt: string;
 };
 
-export async function generateImageUploadUrl({
-  service,
+export type UploadImageAttachmentParams = GetImageAttachmentPathParams & {
+  body: Uint8Array;
+  name?: string;
+};
+
+export type UploadImageAttachmentResult = GetImageAttachmentPathResult;
+
+function getImageAttachmentPutObjectCommand({
+  key,
   userId,
   messageUuid,
   imageId,
   contentType,
   contentLength,
-}: GenerateImageUploadUrlParams): Promise<GenerateImageUploadUrlResult> {
-  const key = getImageKey(service, userId, messageUuid, imageId, contentType);
+  body,
+}: {
+  key: string;
+  userId: string;
+  messageUuid: string;
+  imageId: string;
+  contentType: CloudAgentImageAllowedType;
+  contentLength: number;
+  body?: Uint8Array;
+}): PutObjectCommand {
+  const bodyInput = body === undefined ? {} : { Body: body };
 
-  const command = new PutObjectCommand({
+  return new PutObjectCommand({
     Bucket: r2CloudAgentAttachmentsBucketName,
     Key: key,
     ContentType: contentType,
@@ -59,6 +119,77 @@ export async function generateImageUploadUrl({
       messageUuid,
       imageId,
     },
+    ...bodyInput,
+  });
+}
+
+export async function uploadImageAttachment({
+  service,
+  userId,
+  messageUuid,
+  imageId,
+  contentType,
+  body,
+  name,
+}: UploadImageAttachmentParams): Promise<UploadImageAttachmentResult> {
+  const { filename, key } = getImageAttachmentPath({
+    service,
+    userId,
+    messageUuid,
+    imageId,
+    contentType,
+  });
+
+  assertImageAttachmentSize({
+    service,
+    contentLength: body.byteLength,
+    name: name ?? filename,
+  });
+
+  await r2Client.send(
+    getImageAttachmentPutObjectCommand({
+      key,
+      userId,
+      messageUuid,
+      imageId,
+      contentType,
+      contentLength: body.byteLength,
+      body,
+    })
+  );
+
+  return { filename, key };
+}
+
+export async function generateImageUploadUrl({
+  service,
+  userId,
+  messageUuid,
+  imageId,
+  contentType,
+  contentLength,
+}: GenerateImageUploadUrlParams): Promise<GenerateImageUploadUrlResult> {
+  const { filename, key } = getImageAttachmentPath({
+    service,
+    userId,
+    messageUuid,
+    imageId,
+    contentType,
+  });
+
+  assertImageAttachmentSize({
+    service,
+    contentLength,
+    name: filename,
+  });
+
+  const command = getImageAttachmentPutObjectCommand({
+    key,
+    userId,
+    messageUuid,
+    imageId,
+    contentType,
+    contentLength,
   });
 
   const signedUrl = await getSignedUrl(r2Client, command, {