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({}),
+});