refactor(server): use feature model (#9932)

packages/backend/server/src/plugins/copilot/session.ts

@@ -12,8 +12,8 @@ import {
   CopilotSessionNotFound,
   PrismaTransaction,
 } from '../../base';
-import { FeatureManagementService } from '../../core/features';
 import { QuotaService } from '../../core/quota';
+import { Models } from '../../models';
 import { ChatMessageCache } from './message';
 import { PromptService } from './prompt';
 import {
@@ -195,10 +195,10 @@ export class ChatSessionService {
 
   constructor(
     private readonly db: PrismaClient,
-    private readonly feature: FeatureManagementService,
     private readonly quota: QuotaService,
     private readonly messageCache: ChatMessageCache,
-    private readonly prompt: PromptService
+    private readonly prompt: PromptService,
+    private readonly models: Models
   ) {}
 
   private async haveSession(
@@ -545,12 +545,15 @@ export class ChatSessionService {
   }
 
   async getQuota(userId: string) {
-    const isCopilotUser = await this.feature.isCopilotUser(userId);
+    const isCopilotUser = await this.models.userFeature.has(
+      userId,
+      'unlimited_copilot'
+    );
 
     let limit: number | undefined;
     if (!isCopilotUser) {
       const quota = await this.quota.getUserQuota(userId);
-      limit = quota.feature.copilotActionLimit;
+      limit = quota.copilotActionLimit;
     }
 
     const used = await this.countUserMessages(userId);

packages/backend/server/src/plugins/copilot/storage.ts

@@ -12,7 +12,7 @@ import {
   StorageProviderFactory,
   URLHelper,
 } from '../../base';
-import { QuotaManagementService } from '../../core/quota';
+import { QuotaService } from '../../core/quota';
 
 @Injectable()
 export class CopilotStorage {
@@ -22,7 +22,7 @@ export class CopilotStorage {
     private readonly config: Config,
     private readonly url: URLHelper,
     private readonly storageFactory: StorageProviderFactory,
-    private readonly quota: QuotaManagementService
+    private readonly quota: QuotaService
   ) {
     this.provider = this.storageFactory.create(
       this.config.plugins.copilot.storage
@@ -57,7 +57,7 @@ export class CopilotStorage {
 
   @CallMetric('ai', 'blob_upload')
   async handleUpload(userId: string, blob: FileUpload) {
-    const checkExceeded = await this.quota.getQuotaCalculator(userId);
+    const checkExceeded = await this.quota.getUserQuotaCalculator(userId);
 
     if (checkExceeded(0)) {
       throw new BlobQuotaExceeded();

packages/backend/server/src/plugins/license/service.ts

@@ -12,7 +12,7 @@ import {
   WorkspaceLicenseAlreadyExists,
 } from '../../base';
 import { PermissionService } from '../../core/permission';
-import { QuotaManagementService, QuotaType } from '../../core/quota';
+import { Models } from '../../models';
 import { SubscriptionPlan, SubscriptionRecurring } from '../payment/types';
 
 interface License {
@@ -29,9 +29,9 @@ export class LicenseService {
   constructor(
     private readonly config: Config,
     private readonly db: PrismaClient,
-    private readonly quota: QuotaManagementService,
     private readonly event: EventBus,
-    private readonly permission: PermissionService
+    private readonly permission: PermissionService,
+    private readonly models: Models
   ) {}
 
   async getLicense(workspaceId: string) {
@@ -316,14 +316,13 @@ export class LicenseService {
   }: Events['workspace.subscription.activated']) {
     switch (plan) {
       case SubscriptionPlan.SelfHostedTeam:
-        await this.quota.addTeamWorkspace(
+        await this.models.workspaceFeature.add(
           workspaceId,
-          `${recurring} team subscription activated`
-        );
-        await this.quota.updateWorkspaceConfig(
-          workspaceId,
-          QuotaType.TeamPlanV1,
-          { memberLimit: quantity }
+          'team_plan_v1',
+          `${recurring} team subscription activated`,
+          {
+            memberLimit: quantity,
+          }
         );
         await this.permission.refreshSeatStatus(workspaceId, quantity);
         break;
@@ -339,7 +338,7 @@ export class LicenseService {
   }: Events['workspace.subscription.canceled']) {
     switch (plan) {
       case SubscriptionPlan.SelfHostedTeam:
-        await this.quota.removeTeamWorkspace(workspaceId);
+        await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1');
         break;
       default:
         break;

packages/backend/server/src/plugins/payment/manager/user.ts

@@ -15,10 +15,7 @@ import {
   TooManyRequest,
   URLHelper,
 } from '../../../base';
-import {
-  EarlyAccessType,
-  FeatureManagementService,
-} from '../../../core/features';
+import { EarlyAccessType, FeatureService } from '../../../core/features';
 import {
   CouponType,
   KnownStripeInvoice,
@@ -59,7 +56,7 @@ export class UserSubscriptionManager extends SubscriptionManager {
     stripe: Stripe,
     db: PrismaClient,
     private readonly runtime: Runtime,
-    private readonly feature: FeatureManagementService,
+    private readonly feature: FeatureService,
     private readonly event: EventBus,
     private readonly url: URLHelper,
     private readonly mutex: Mutex

packages/backend/server/src/plugins/payment/quota.ts

@@ -1,25 +1,17 @@
 import { Injectable } from '@nestjs/common';
 import { OnEvent } from '@nestjs/event-emitter';
 
-import { FeatureManagementService } from '../../core/features';
 import { PermissionService } from '../../core/permission';
-import {
-  QuotaManagementService,
-  QuotaService,
-  QuotaType,
-} from '../../core/quota';
 import { WorkspaceService } from '../../core/workspaces/resolvers';
+import { Models } from '../../models';
 import { SubscriptionPlan } from './types';
 
 @Injectable()
 export class QuotaOverride {
   constructor(
-    private readonly quota: QuotaService,
-    private readonly manager: QuotaManagementService,
     private readonly permission: PermissionService,
     private readonly workspace: WorkspaceService,
-    private readonly feature: FeatureManagementService,
-    private readonly quotaService: QuotaService
+    private readonly models: Models
   ) {}
 
   @OnEvent('workspace.subscription.activated')
@@ -31,21 +23,17 @@ export class QuotaOverride {
   }: Events['workspace.subscription.activated']) {
     switch (plan) {
       case 'team': {
-        const hasTeamWorkspace = await this.quota.hasWorkspaceQuota(
+        const isTeam = await this.workspace.isTeamWorkspace(workspaceId);
+        await this.models.workspaceFeature.add(
           workspaceId,
-          QuotaType.TeamPlanV1
-        );
-        await this.manager.addTeamWorkspace(
-          workspaceId,
-          `${recurring} team subscription activated`
-        );
-        await this.quota.updateWorkspaceConfig(
-          workspaceId,
-          QuotaType.TeamPlanV1,
-          { memberLimit: quantity }
+          'team_plan_v1',
+          `${recurring} team subscription activated`,
+          {
+            memberLimit: quantity,
+          }
         );
         await this.permission.refreshSeatStatus(workspaceId, quantity);
-        if (!hasTeamWorkspace) {
+        if (!isTeam) {
           // this event will triggered when subscription is activated or changed
           // we only send emails when the team workspace is activated
           await this.workspace.sendTeamWorkspaceUpgradedEmail(workspaceId);
@@ -64,7 +52,7 @@ export class QuotaOverride {
   }: Events['workspace.subscription.canceled']) {
     switch (plan) {
       case SubscriptionPlan.Team:
-        await this.manager.removeTeamWorkspace(workspaceId);
+        await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1');
         break;
       default:
         break;
@@ -79,14 +67,16 @@ export class QuotaOverride {
   }: Events['user.subscription.activated']) {
     switch (plan) {
       case SubscriptionPlan.AI:
-        await this.feature.addCopilot(userId, 'subscription activated');
+        await this.models.userFeature.add(
+          userId,
+          'unlimited_copilot',
+          'subscription activated'
+        );
         break;
       case SubscriptionPlan.Pro:
-        await this.quotaService.switchUserQuota(
+        await this.models.userFeature.add(
           userId,
-          recurring === 'lifetime'
-            ? QuotaType.LifetimeProPlanV1
-            : QuotaType.ProPlanV1,
+          recurring === 'lifetime' ? 'lifetime_pro_plan_v1' : 'pro_plan_v1',
           'subscription activated'
         );
         break;
@@ -102,16 +92,20 @@ export class QuotaOverride {
   }: Events['user.subscription.canceled']) {
     switch (plan) {
       case SubscriptionPlan.AI:
-        await this.feature.removeCopilot(userId);
+        await this.models.userFeature.remove(userId, 'unlimited_copilot');
         break;
       case SubscriptionPlan.Pro: {
         // edge case: when user switch from recurring Pro plan to `Lifetime` plan,
         // a subscription canceled event will be triggered because `Lifetime` plan is not subscription based
-        const quota = await this.quotaService.getUserQuota(userId);
-        if (quota.feature.name !== QuotaType.LifetimeProPlanV1) {
-          await this.quotaService.switchUserQuota(
+        const isLifetimeUser = await this.models.userFeature.has(
+          userId,
+          'lifetime_pro_plan_v1'
+        );
+
+        if (!isLifetimeUser) {
+          await this.models.userFeature.switchQuota(
             userId,
-            QuotaType.FreePlanV1,
+            'free_plan_v1',
             'subscription canceled'
           );
         }

packages/backend/server/src/plugins/payment/service.ts

@@ -26,7 +26,7 @@ import {
   UserNotFound,
 } from '../../base';
 import { CurrentUser } from '../../core/auth';
-import { FeatureManagementService } from '../../core/features';
+import { FeatureService } from '../../core/features';
 import { Models } from '../../models';
 import {
   CheckoutParams,
@@ -83,7 +83,7 @@ export class SubscriptionService implements OnApplicationBootstrap {
     private readonly config: Config,
     private readonly stripe: Stripe,
     private readonly db: PrismaClient,
-    private readonly feature: FeatureManagementService,
+    private readonly feature: FeatureService,
     private readonly models: Models,
     private readonly userManager: UserSubscriptionManager,
     private readonly workspaceManager: WorkspaceSubscriptionManager,