feat: set custom upload fn, uploads in progress count toward available slots, preserve original attachment id

Author: MartinCupela

src/messageComposer/attachmentManager.ts

@@ -1,3 +1,4 @@
+import { DEFAULT_ATTACHMENT_MANAGER_CONFIG } from './configuration';
 import { isLocalImageAttachment, isUploadedAttachment } from './attachmentIdentity';
 import {
   createFileFromBlobs,
@@ -11,6 +12,7 @@ import {
 } from './fileUtils';
 import { StateStore } from '../store';
 import { generateUUIDv4 } from '../utils';
+import { mergeWith } from '../utils/mergeWith';
 import { DEFAULT_UPLOAD_SIZE_LIMIT_BYTES } from '../constants';
 import type { Channel } from '../channel';
 import type {
@@ -25,15 +27,12 @@ import type {
   LocalVoiceRecordingAttachment,
   UploadPermissionCheckResult,
 } from './types';
+import type { ChannelResponse, DraftMessage, LocalMessage } from '../types';
 import type {
-  ChannelResponse,
-  DraftMessage,
-  LocalMessage,
-  SendFileAPIResponse,
-} from '../types';
-import { mergeWith } from '../utils/mergeWith';
-import { DEFAULT_ATTACHMENT_MANAGER_CONFIG } from './configuration/configuration';
-import type { AttachmentManagerConfig } from './configuration/types';
+  AttachmentManagerConfig,
+  MinimumUploadRequestResult,
+  UploadRequestFn,
+} from './configuration';
 
 type LocalNotImageAttachment =
   | LocalFileAttachment
@@ -155,7 +154,11 @@ export class AttachmentManager {
   }
 
   get availableUploadSlots() {
-    return this.config.maxNumberOfFilesPerMessage - this.successfulUploadsCount;
+    return (
+      this.config.maxNumberOfFilesPerMessage -
+      this.successfulUploadsCount -
+      this.uploadsInProgressCount
+    );
   }
 
   getUploadsByState(state: AttachmentLoadingState) {
@@ -336,18 +339,24 @@ export class AttachmentManager {
     // the following is substitute for: if (noFiles && !isImage) return att
     if (!this.fileUploadFilter(attachment)) return;
 
-    return await this.fileToLocalUploadAttachment(attachment.localMetadata.file);
+    const newAttachment = await this.fileToLocalUploadAttachment(
+      attachment.localMetadata.file,
+    );
+    if (attachment.localMetadata.id) {
+      newAttachment.localMetadata.id = attachment.localMetadata.id;
+    }
+    return newAttachment;
+  };
+
+  setCustomUploadFn = (doUploadRequest: UploadRequestFn) => {
+    this.configState.partialNext({ doUploadRequest });
   };
 
   /**
-   * todo: docs how to customize the image and file upload by overriding do
+   * Method to perform the default upload behavior without checking for custom upload functions
+   * to prevent recursive calls
    */
-
-  doUploadRequest = (fileLike: FileReference | FileLike) => {
-    if (this.config.doUploadRequest) {
-      return this.config.doUploadRequest(fileLike);
-    }
-
+  doDefaultUploadRequest = async (fileLike: FileReference | FileLike) => {
     if (isFileReference(fileLike)) {
       return this.channel[isImageFile(fileLike) ? 'sendImage' : 'sendFile'](
         fileLike.uri,
@@ -364,7 +373,23 @@ export class AttachmentManager {
           mimeType: fileLike.type,
         });
 
-    return this.channel[isImageFile(fileLike) ? 'sendImage' : 'sendFile'](file);
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    const { duration, ...result } =
+      await this.channel[isImageFile(fileLike) ? 'sendImage' : 'sendFile'](file);
+    return result;
+  };
+
+  /**
+   * todo: docs how to customize the image and file upload by overriding do
+   */
+
+  doUploadRequest = async (fileLike: FileReference | FileLike) => {
+    const customUploadFn = this.config.doUploadRequest;
+    if (customUploadFn) {
+      return await customUploadFn(fileLike);
+    }
+
+    return this.doDefaultUploadRequest(fileLike);
   };
 
   uploadAttachment = async (attachment: LocalUploadAttachment) => {
@@ -376,7 +401,11 @@ export class AttachmentManager {
 
     if (localAttachment.localMetadata.uploadState === 'blocked') {
       this.upsertAttachments([localAttachment]);
-      return;
+      this.client.notifications.addError({
+        message: 'Error uploading attachment',
+        origin: { emitter: 'AttachmentManager', context: { attachment } },
+      });
+      return localAttachment;
     }
 
     this.upsertAttachments([
@@ -389,7 +418,7 @@ export class AttachmentManager {
       },
     ]);
 
-    let response: SendFileAPIResponse;
+    let response: MinimumUploadRequestResult;
     try {
       response = await this.doUploadRequest(localAttachment.localMetadata.file);
     } catch (error) {
@@ -417,8 +446,7 @@ export class AttachmentManager {
       };
 
       this.upsertAttachments([failedAttachment]);
-
-      throw finalError;
+      return failedAttachment;
     }
 
     if (!response) {

test/unit/MessageComposer/attachmentManager.test.ts

@@ -325,6 +325,39 @@ describe('AttachmentManager', () => {
       // Should have max slots minus the number of attachments in the message
       expect(manager.availableUploadSlots).toBe(API_MAX_FILES_ALLOWED_PER_MESSAGE - 2);
     });
+
+    it('should take into consideration uploads in progress', () => {
+      const { attachmentManager } = setup();
+
+      // Set up state with successful uploads and uploads in progress
+      attachmentManager.state.next({
+        attachments: [
+          {
+            type: 'image',
+            image_url: 'test-image-url',
+            localMetadata: {
+              id: 'test-uuid-1',
+              uploadState: 'finished',
+              file: new File([''], 'test1.jpg', { type: 'image/jpeg' }),
+            },
+          },
+          {
+            type: 'image',
+            image_url: 'test-image-url',
+            localMetadata: {
+              id: 'test-uuid-2',
+              uploadState: 'uploading',
+              file: new File([''], 'test2.jpg', { type: 'image/jpeg' }),
+            },
+          },
+        ],
+      });
+
+      // Should have max slots minus successful uploads (1) minus uploads in progress (1)
+      expect(attachmentManager.availableUploadSlots).toBe(
+        API_MAX_FILES_ALLOWED_PER_MESSAGE - 2,
+      );
+    });
   });
 
   describe('initState', () => {
@@ -691,7 +724,23 @@ describe('AttachmentManager', () => {
       mockChannel.sendImage.mockRejectedValueOnce(new Error('Upload failed'));
       const file = new File([''], 'test.jpg', { type: 'image/jpeg' });
 
-      await expect(attachmentManager.uploadFiles([file])).rejects.toThrowError();
+      await expect(attachmentManager.uploadFiles([file])).resolves.toEqual([
+        {
+          fallback: 'test.jpg',
+          file_size: 0,
+          localMetadata: {
+            id: 'test-uuid',
+            file,
+            uploadState: 'failed',
+            previewUri: expect.any(String),
+            uploadPermissionCheck: {
+              uploadBlocked: false,
+            },
+          },
+          mime_type: 'image/jpeg',
+          type: 'image',
+        },
+      ]);
 
       expect(attachmentManager.failedUploadsCount).toBe(1);
       expect(mockClient.notifications.addError).toHaveBeenCalledWith({
@@ -703,6 +752,70 @@ describe('AttachmentManager', () => {
       });
     });
 
+    it('should register notification for blocked file', async () => {
+      const { attachmentManager, mockClient } = setup();
+
+      // Create a blocked attachment
+      const blockedAttachment = {
+        type: 'image',
+        localMetadata: {
+          id: 'test-id',
+          file: new File([''], 'test.jpg', { type: 'image/jpeg' }),
+        },
+      };
+
+      // Mock getUploadConfigCheck to return blocked
+      vi.spyOn(attachmentManager, 'getUploadConfigCheck').mockResolvedValue({
+        uploadBlocked: true,
+        reason: 'size_limit',
+      });
+
+      await attachmentManager.uploadAttachment(blockedAttachment);
+
+      // Verify notification was added
+      expect(mockClient.notifications.addError).toHaveBeenCalledWith({
+        message: 'Error uploading attachment',
+        origin: {
+          emitter: 'AttachmentManager',
+          context: { attachment: blockedAttachment },
+        },
+      });
+    });
+
+    it('should use custom upload function when provided', async () => {
+      const { attachmentManager, mockChannel } = setup();
+
+      // Create a custom upload function
+      const customUploadFn = vi.fn().mockResolvedValue({ file: 'custom-upload-url' });
+
+      // Set the custom upload function
+      attachmentManager.setCustomUploadFn(customUploadFn);
+
+      // Create a file to upload
+      const file = new File([''], 'test.jpg', { type: 'image/jpeg' });
+
+      // Mock fileToLocalUploadAttachment to return a valid attachment
+      const attachment = {
+        type: 'image',
+        localMetadata: {
+          id: 'test-id',
+          file,
+          uploadState: 'pending',
+        },
+      };
+
+      vi.spyOn(attachmentManager, 'ensureLocalUploadAttachment').mockResolvedValue(
+        attachment,
+      );
+
+      // Upload the attachment
+      await attachmentManager.uploadAttachment(attachment);
+
+      // Verify the custom upload function was called
+      expect(customUploadFn).toHaveBeenCalledWith(file);
+      expect(mockChannel.sendImage).not.toHaveBeenCalled();
+    });
+
     it('should respect maxNumberOfFilesPerMessage', async () => {
       const { attachmentManager } = setup();
       const files = Array(API_MAX_FILES_ALLOWED_PER_MESSAGE + 1)
@@ -848,5 +961,44 @@ describe('AttachmentManager', () => {
 
       expect(result).toEqual(expectedAttachment);
     });
+
+    it('should preserve original ID if it exists', async () => {
+      const { attachmentManager } = setup();
+      const ensureLocalUploadAttachment = (attachmentManager as any)
+        .ensureLocalUploadAttachment;
+
+      // Set a fileUploadFilter that allows all files
+      attachmentManager.fileUploadFilter = () => true;
+
+      // Create an attachment with an ID
+      const originalId = 'original-test-id';
+      const file = new File([''], 'test.jpg', { type: 'image/jpeg' });
+
+      // Mock fileToLocalUploadAttachment to return a new attachment
+      const newAttachment = {
+        type: 'image',
+        image_url: 'test-url',
+        localMetadata: {
+          id: 'new-test-id', // Different ID
+          file: new File([''], 'test.jpg', { type: 'image/jpeg' }),
+          uploadState: 'finished',
+        },
+      };
+
+      vi.spyOn(attachmentManager, 'fileToLocalUploadAttachment').mockResolvedValue(
+        newAttachment,
+      );
+
+      // Call with original ID
+      const result = await ensureLocalUploadAttachment({
+        localMetadata: {
+          id: originalId,
+          file,
+        },
+      });
+
+      // Verify the original ID was preserved
+      expect(result.localMetadata.id).toBe(originalId);
+    });
   });
 });