feat: add blob upload support for copilot (#6584)

packages/backend/server/src/core/quota/storage.ts

@@ -7,7 +7,10 @@ import { OneGB } from './constant';
 import { QuotaService } from './service';
 import { formatSize, QuotaQueryType } from './types';
 
-type QuotaBusinessType = QuotaQueryType & { businessBlobLimit: number };
+type QuotaBusinessType = QuotaQueryType & {
+  businessBlobLimit: number;
+  unlimited: boolean;
+};
 
 @Injectable()
 export class QuotaManagementService {
@@ -59,6 +62,52 @@ export class QuotaManagementService {
     }, 0);
   }
 
+  private generateQuotaCalculator(
+    quota: number,
+    blobLimit: number,
+    usedSize: number,
+    unlimited = false
+  ) {
+    const checkExceeded = (recvSize: number) => {
+      const total = usedSize + recvSize;
+      // only skip total storage check if workspace has unlimited feature
+      if (total > quota && !unlimited) {
+        this.logger.log(`storage size limit exceeded: ${total} > ${quota}`);
+        return true;
+      } else if (recvSize > blobLimit) {
+        this.logger.log(`blob size limit exceeded: ${recvSize} > ${blobLimit}`);
+        return true;
+      } else {
+        return false;
+      }
+    };
+    return checkExceeded;
+  }
+
+  async getQuotaCalculator(userId: string) {
+    const quota = await this.getUserQuota(userId);
+    const { storageQuota, businessBlobLimit } = quota;
+    const usedSize = await this.getUserUsage(userId);
+
+    return this.generateQuotaCalculator(
+      storageQuota,
+      businessBlobLimit,
+      usedSize
+    );
+  }
+
+  async getQuotaCalculatorByWorkspace(workspaceId: string) {
+    const { storageQuota, usedSize, businessBlobLimit, unlimited } =
+      await this.getWorkspaceUsage(workspaceId);
+
+    return this.generateQuotaCalculator(
+      storageQuota,
+      businessBlobLimit,
+      usedSize,
+      unlimited
+    );
+  }
+
   // get workspace's owner quota and total size of used
   // quota was apply to owner's account
   async getWorkspaceUsage(workspaceId: string): Promise<QuotaBusinessType> {
@@ -79,6 +128,12 @@ export class QuotaManagementService {
     } = await this.quota.getUserQuota(owner.id);
     // get all workspaces size of owner used
     const usedSize = await this.getUserUsage(owner.id);
+    // relax restrictions if workspace has unlimited feature
+    // todo(@darkskygit): need a mechanism to allow feature as a middleware to edit quota
+    const unlimited = await this.feature.hasWorkspaceFeature(
+      workspaceId,
+      FeatureType.UnlimitedWorkspace
+    );
 
     const quota = {
       name,
@@ -90,15 +145,10 @@ export class QuotaManagementService {
       copilotActionLimit,
       humanReadable,
       usedSize,
+      unlimited,
     };
 
-    // relax restrictions if workspace has unlimited feature
-    // todo(@darkskygit): need a mechanism to allow feature as a middleware to edit quota
-    const unlimited = await this.feature.hasWorkspaceFeature(
-      workspaceId,
-      FeatureType.UnlimitedWorkspace
-    );
-    if (unlimited) {
+    if (quota.unlimited) {
       return this.mergeUnlimitedQuota(quota);
     }
 

packages/backend/server/src/core/workspaces/resolvers/blob.ts

@@ -1,8 +1,4 @@
-import {
-  ForbiddenException,
-  Logger,
-  PayloadTooLargeException,
-} from '@nestjs/common';
+import { Logger, PayloadTooLargeException, UseGuards } from '@nestjs/common';
 import {
   Args,
   Int,
@@ -16,20 +12,23 @@ import { SafeIntResolver } from 'graphql-scalars';
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 
 import type { FileUpload } from '../../../fundamentals';
-import { MakeCache, PreventCache } from '../../../fundamentals';
+import {
+  CloudThrottlerGuard,
+  MakeCache,
+  PreventCache,
+} from '../../../fundamentals';
 import { CurrentUser } from '../../auth';
-import { FeatureManagementService, FeatureType } from '../../features';
 import { QuotaManagementService } from '../../quota';
 import { WorkspaceBlobStorage } from '../../storage';
 import { PermissionService } from '../permission';
 import { Permission, WorkspaceBlobSizes, WorkspaceType } from '../types';
 
+@UseGuards(CloudThrottlerGuard)
 @Resolver(() => WorkspaceType)
 export class WorkspaceBlobResolver {
   logger = new Logger(WorkspaceBlobResolver.name);
   constructor(
     private readonly permissions: PermissionService,
-    private readonly feature: FeatureManagementService,
     private readonly quota: QuotaManagementService,
     private readonly storage: WorkspaceBlobStorage
   ) {}
@@ -124,34 +123,8 @@ export class WorkspaceBlobResolver {
       Permission.Write
     );
 
-    const { storageQuota, usedSize, businessBlobLimit } =
-      await this.quota.getWorkspaceUsage(workspaceId);
-
-    const unlimited = await this.feature.hasWorkspaceFeature(
-      workspaceId,
-      FeatureType.UnlimitedWorkspace
-    );
-
-    const checkExceeded = (recvSize: number) => {
-      if (!storageQuota) {
-        throw new ForbiddenException('Cannot find user quota.');
-      }
-      const total = usedSize + recvSize;
-      // only skip total storage check if workspace has unlimited feature
-      if (total > storageQuota && !unlimited) {
-        this.logger.log(
-          `storage size limit exceeded: ${total} > ${storageQuota}`
-        );
-        return true;
-      } else if (recvSize > businessBlobLimit) {
-        this.logger.log(
-          `blob size limit exceeded: ${recvSize} > ${businessBlobLimit}`
-        );
-        return true;
-      } else {
-        return false;
-      }
-    };
+    const checkExceeded =
+      await this.quota.getQuotaCalculatorByWorkspace(workspaceId);
 
     if (checkExceeded(0)) {
       throw new PayloadTooLargeException(