feat: check quota correctly (#6561)

packages/backend/server/src/core/features/feature.ts

@@ -54,10 +54,23 @@ export class UnlimitedWorkspaceFeatureConfig extends FeatureConfig {
   }
 }
 
+export class UnlimitedCopilotFeatureConfig extends FeatureConfig {
+  override config!: Feature & { feature: FeatureType.UnlimitedCopilot };
+
+  constructor(data: any) {
+    super(data);
+
+    if (this.config.feature !== FeatureType.UnlimitedCopilot) {
+      throw new Error('Invalid feature config: type is not UnlimitedWorkspace');
+    }
+  }
+}
+
 const FeatureConfigMap = {
   [FeatureType.Copilot]: CopilotFeatureConfig,
   [FeatureType.EarlyAccess]: EarlyAccessFeatureConfig,
   [FeatureType.UnlimitedWorkspace]: UnlimitedWorkspaceFeatureConfig,
+  [FeatureType.UnlimitedCopilot]: UnlimitedCopilotFeatureConfig,
 };
 
 export type FeatureConfigType<F extends FeatureType> = InstanceType<

packages/backend/server/src/core/features/management.ts

@@ -35,7 +35,6 @@ export class FeatureManagementService {
     return this.feature.addUserFeature(
       userId,
       FeatureType.EarlyAccess,
-      2,
       'Early access user'
     );
   }
@@ -116,9 +115,9 @@ export class FeatureManagementService {
     return this.feature.listFeatureWorkspaces(feature);
   }
 
-  async getUserFeatures(userId: string): Promise<FeatureType[]> {
-    return (await this.feature.getUserFeatures(userId)).map(
-      f => f.feature.name
-    );
+  // ======== User Feature ========
+  async getActivatedUserFeatures(userId: string): Promise<FeatureType[]> {
+    const features = await this.feature.getActivatedUserFeatures(userId);
+    return features.map(f => f.feature.name);
   }
 }

packages/backend/server/src/core/features/service.ts

@@ -59,11 +59,17 @@ export class FeatureService {
   async addUserFeature(
     userId: string,
     feature: FeatureType,
-    version: number,
     reason: string,
     expiredAt?: Date | string
   ) {
     return this.prisma.$transaction(async tx => {
+      const latestVersion = await tx.features
+        .aggregate({
+          where: { feature },
+          _max: { version: true },
+        })
+        .then(r => r._max.version || 1);
+
       const latestFlag = await tx.userFeatures.findFirst({
         where: {
           userId,
@@ -95,7 +101,7 @@ export class FeatureService {
                 connect: {
                   feature_version: {
                     feature,
-                    version,
+                    version: latestVersion,
                   },
                   type: FeatureKind.Feature,
                 },
@@ -157,6 +163,33 @@ export class FeatureService {
     return configs.filter(feature => !!feature.feature);
   }
 
+  async getActivatedUserFeatures(userId: string) {
+    const features = await this.prisma.userFeatures.findMany({
+      where: {
+        user: { id: userId },
+        feature: { type: FeatureKind.Feature },
+        activated: true,
+        OR: [{ expiredAt: null }, { expiredAt: { gt: new Date() } }],
+      },
+      select: {
+        activated: true,
+        reason: true,
+        createdAt: true,
+        expiredAt: true,
+        featureId: true,
+      },
+    });
+
+    const configs = await Promise.all(
+      features.map(async feature => ({
+        ...feature,
+        feature: await getFeature(this.prisma, feature.featureId),
+      }))
+    );
+
+    return configs.filter(feature => !!feature.feature);
+  }
+
   async listFeatureUsers(feature: FeatureType) {
     return this.prisma.userFeatures
       .findMany({

packages/backend/server/src/core/features/types/common.ts

@@ -1,8 +1,11 @@
 import { registerEnumType } from '@nestjs/graphql';
 
 export enum FeatureType {
-  Copilot = 'copilot',
+  // user feature
   EarlyAccess = 'early_access',
+  UnlimitedCopilot = 'unlimited_copilot',
+  // workspace feature
+  Copilot = 'copilot',
   UnlimitedWorkspace = 'unlimited_workspace',
 }
 

packages/backend/server/src/core/features/types/index.ts

@@ -3,6 +3,7 @@ import { z } from 'zod';
 import { FeatureType } from './common';
 import { featureCopilot } from './copilot';
 import { featureEarlyAccess } from './early-access';
+import { featureUnlimitedCopilot } from './unlimited-copilot';
 import { featureUnlimitedWorkspace } from './unlimited-workspace';
 
 /// ======== common schema ========
@@ -52,6 +53,12 @@ export const Features: Feature[] = [
     version: 1,
     configs: {},
   },
+  {
+    feature: FeatureType.UnlimitedCopilot,
+    type: FeatureKind.Feature,
+    version: 1,
+    configs: {},
+  },
 ];
 
 /// ======== schema infer ========
@@ -65,6 +72,7 @@ export const FeatureSchema = commonFeatureSchema
       featureCopilot,
       featureEarlyAccess,
       featureUnlimitedWorkspace,
+      featureUnlimitedCopilot,
     ])
   );
 

packages/backend/server/src/core/features/types/unlimited-copilot.ts

@@ -0,0 +1,8 @@
+import { z } from 'zod';
+
+import { FeatureType } from './common';
+
+export const featureUnlimitedCopilot = z.object({
+  feature: z.literal(FeatureType.UnlimitedCopilot),
+  configs: z.object({}),
+});

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

@@ -20,5 +20,5 @@ import { QuotaManagementService } from './storage';
 export class QuotaModule {}
 
 export { QuotaManagementService, QuotaService };
-export { Quota_FreePlanV1_1, Quota_ProPlanV1, Quotas } from './schema';
+export { Quota_FreePlanV1_1, Quota_ProPlanV1 } from './schema';
 export { QuotaQueryType, QuotaType } from './types';

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

@@ -117,14 +117,61 @@ export const Quotas: Quota[] = [
       copilotActionLimit: 10,
     },
   },
+  {
+    feature: QuotaType.ProPlanV1,
+    type: FeatureKind.Quota,
+    version: 2,
+    configs: {
+      // quota name
+      name: 'Pro',
+      // single blob limit 100MB
+      blobLimit: 100 * OneMB,
+      // total blob limit 100GB
+      storageQuota: 100 * OneGB,
+      // history period of validity 30 days
+      historyPeriod: 30 * OneDay,
+      // member limit 10
+      memberLimit: 10,
+      // copilot action limit 10
+      copilotActionLimit: 10,
+    },
+  },
+  {
+    feature: QuotaType.RestrictedPlanV1,
+    type: FeatureKind.Quota,
+    version: 2,
+    configs: {
+      // quota name
+      name: 'Restricted',
+      // single blob limit 10MB
+      blobLimit: OneMB,
+      // total blob limit 1GB
+      storageQuota: 10 * OneMB,
+      // history period of validity 30 days
+      historyPeriod: 30 * OneDay,
+      // member limit 10
+      memberLimit: 10,
+      // copilot action limit 10
+      copilotActionLimit: 10,
+    },
+  },
 ];
 
+export function getLatestQuota(type: QuotaType) {
+  const quota = Quotas.filter(f => f.feature === type);
+  quota.sort((a, b) => b.version - a.version);
+  return quota[0];
+}
+
+export const FreePlan = getLatestQuota(QuotaType.FreePlanV1);
+export const ProPlan = getLatestQuota(QuotaType.ProPlanV1);
+
 export const Quota_FreePlanV1_1 = {
   feature: Quotas[5].feature,
   version: Quotas[5].version,
 };
 
 export const Quota_ProPlanV1 = {
-  feature: Quotas[1].feature,
-  version: Quotas[1].version,
+  feature: Quotas[6].feature,
+  version: Quotas[6].version,
 };

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

@@ -3,13 +3,17 @@ import { PrismaClient } from '@prisma/client';
 
 import type { EventPayload } from '../../fundamentals';
 import { OnEvent, PrismaTransaction } from '../../fundamentals';
-import { FeatureKind } from '../features';
+import { SubscriptionPlan } from '../../plugins/payment/types';
+import { FeatureKind, FeatureService, FeatureType } from '../features';
 import { QuotaConfig } from './quota';
 import { QuotaType } from './types';
 
 @Injectable()
 export class QuotaService {
-  constructor(private readonly prisma: PrismaClient) {}
+  constructor(
+    private readonly prisma: PrismaClient,
+    private readonly feature: FeatureService
+  ) {}
 
   // get activated user quota
   async getUserQuota(userId: string) {
@@ -159,22 +163,49 @@ export class QuotaService {
   @OnEvent('user.subscription.activated')
   async onSubscriptionUpdated({
     userId,
+    plan,
   }: EventPayload<'user.subscription.activated'>) {
-    await this.switchUserQuota(
-      userId,
-      QuotaType.ProPlanV1,
-      'subscription activated'
-    );
+    switch (plan) {
+      case SubscriptionPlan.AI:
+        await this.feature.addUserFeature(
+          userId,
+          FeatureType.UnlimitedCopilot,
+          'subscription activated'
+        );
+        break;
+      case SubscriptionPlan.Pro:
+        await this.switchUserQuota(
+          userId,
+          QuotaType.ProPlanV1,
+          'subscription activated'
+        );
+        break;
+      default:
+        break;
+    }
   }
 
   @OnEvent('user.subscription.canceled')
-  async onSubscriptionCanceled(
-    userId: EventPayload<'user.subscription.canceled'>
-  ) {
-    await this.switchUserQuota(
-      userId,
-      QuotaType.FreePlanV1,
-      'subscription canceled'
-    );
+  async onSubscriptionCanceled({
+    userId,
+    plan,
+  }: EventPayload<'user.subscription.canceled'>) {
+    switch (plan) {
+      case SubscriptionPlan.AI:
+        await this.feature.removeUserFeature(
+          userId,
+          FeatureType.UnlimitedCopilot
+        );
+        break;
+      case SubscriptionPlan.Pro:
+        await this.switchUserQuota(
+          userId,
+          QuotaType.FreePlanV1,
+          'subscription canceled'
+        );
+        break;
+      default:
+        break;
+    }
   }
 }

packages/backend/server/src/core/user/resolver.ts

@@ -115,7 +115,7 @@ export class UserResolver {
     description: 'Enabled features of a user',
   })
   async userFeatures(@CurrentUser() user: CurrentUser) {
-    return this.feature.getUserFeatures(user.id);
+    return this.feature.getActivatedUserFeatures(user.id);
   }
 
   @Throttle({

packages/backend/server/src/core/workspaces/management.ts

@@ -117,12 +117,7 @@ export class WorkspaceManagementResolver {
   async availableFeatures(
     @CurrentUser() user: CurrentUser
   ): Promise<FeatureType[]> {
-    const isEarlyAccessUser = await this.feature.isEarlyAccessUser(user.email);
-    if (isEarlyAccessUser) {
-      return [FeatureType.Copilot];
-    } else {
-      return [];
-    }
+    return await this.feature.getActivatedUserFeatures(user.id);
   }
 
   @ResolveField(() => [FeatureType], {