feat(api): add support for custom plan quota override for existing users

- Introduce `overridePlan` field in Subscription model to allow custom quota settings - Create PlanQuota interface to define custom quota structure - Implement logic to parse and use custom plan quota when available - Modify token and storage usage meter creation to support custom quotas
refly-ai · Mar 10, 2025 · 8d19c5b · 8d19c5b
1 parent 2d6e077
commit 8d19c5b
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 24 deletions.
diff --git a/apps/api/prisma/schema.prisma b/apps/api/prisma/schema.prisma
@@ -699,6 +699,8 @@ model Subscription {
   status         String    @map("status")
   /// Whether this is a trial subscription
   isTrial        Boolean   @default(false) @map("is_trial")
+  /// Override plan quota (JSON string)
+  overridePlan   String?   @map("override_plan")
   /// Cancel timestamp
   cancelAt       DateTime? @map("cancel_at") @db.Timestamptz()
   /// Create timestamp

diff --git a/apps/api/src/subscription/subscription.dto.ts b/apps/api/src/subscription/subscription.dto.ts
@@ -15,6 +15,12 @@ import {
 } from '@prisma/client';
 import { pick } from '@/utils';
 
+export interface PlanQuota {
+  t1CountQuota: number;
+  t2CountQuota: number;
+  fileCountQuota: number;
+}
+
 export interface CreateSubscriptionParam {
   subscriptionId: string;
   customerId: string;

diff --git a/apps/api/src/subscription/subscription.service.ts b/apps/api/src/subscription/subscription.service.ts
@@ -11,7 +11,7 @@ import {
   SubscriptionUsageData,
   User,
 } from '@refly-packages/openapi-schema';
-import { genTokenUsageMeterID, genStorageUsageMeterID } from '@refly-packages/utils';
+import { genTokenUsageMeterID, genStorageUsageMeterID, safeParseJSON } from '@refly-packages/utils';
 import {
   CreateSubscriptionParam,
   SyncTokenUsageJobData,
@@ -22,6 +22,7 @@ import {
   CheckStorageUsageResult,
   SyncRequestUsageJobData,
   CheckFileParseUsageResult,
+  PlanQuota,
 } from '@/subscription/subscription.dto';
 import { pick } from '@/utils';
 import {
@@ -273,12 +274,8 @@ export class SubscriptionService implements OnModuleInit {
           endAt,
           t1CountQuota: plan?.t1CountQuota ?? this.config.get('quota.request.t1'),
           t1CountUsed: 0,
-          t1TokenQuota: plan?.t1TokenQuota ?? this.config.get('quota.token.t1'),
-          t1TokenUsed: 0,
           t2CountQuota: plan?.t2CountQuota ?? this.config.get('quota.request.t2'),
           t2CountUsed: 0,
-          t2TokenQuota: plan?.t2TokenQuota ?? this.config.get('quota.token.t2'),
-          t2TokenUsed: 0,
         },
       });
 
@@ -292,8 +289,6 @@ export class SubscriptionService implements OnModuleInit {
         data: {
           subscriptionId: sub.subscriptionId,
           fileCountQuota: plan?.fileCountQuota ?? this.config.get('quota.storage.file'),
-          objectStorageQuota: plan?.objectStorageQuota ?? this.config.get('quota.storage.object'),
-          vectorStorageQuota: plan?.vectorStorageQuota ?? this.config.get('quota.storage.vector'),
         },
       });
 
@@ -355,10 +350,6 @@ export class SubscriptionService implements OnModuleInit {
         data: {
           subscriptionId: null,
           fileCountQuota: freePlan?.fileCountQuota ?? this.config.get('quota.storage.file'),
-          objectStorageQuota:
-            freePlan?.objectStorageQuota ?? this.config.get('quota.storage.object'),
-          vectorStorageQuota:
-            freePlan?.vectorStorageQuota ?? this.config.get('quota.storage.vector'),
         },
       });
     });
@@ -704,9 +695,15 @@ export class SubscriptionService implements OnModuleInit {
           ? new Date(startAt.getFullYear(), startAt.getMonth(), startAt.getDate() + 1)
           : new Date(startAt.getFullYear(), startAt.getMonth() + 1, startAt.getDate());
 
-      const plan = await this.prisma.subscriptionPlan.findFirstOrThrow({
-        where: { planType },
-      });
+      let plan: PlanQuota | null = null;
+      if (sub?.overridePlan) {
+        plan = safeParseJSON(sub.overridePlan) as PlanQuota;
+      }
+      if (!plan) {
+        plan = await this.prisma.subscriptionPlan.findFirstOrThrow({
+          where: { planType },
+        });
+      }
 
       return prisma.tokenUsageMeter.create({
         data: {
@@ -717,12 +714,8 @@ export class SubscriptionService implements OnModuleInit {
           endAt,
           t1CountQuota: plan?.t1CountQuota ?? this.config.get('quota.request.t1'),
           t1CountUsed: 0,
-          t1TokenQuota: plan?.t1TokenQuota ?? this.config.get('quota.token.t1'),
-          t1TokenUsed: 0,
           t2CountQuota: plan?.t2CountQuota ?? this.config.get('quota.request.t2'),
           t2CountUsed: 0,
-          t2TokenQuota: plan?.t2TokenQuota ?? this.config.get('quota.token.t2'),
-          t2TokenUsed: 0,
         },
       });
     });
@@ -757,13 +750,15 @@ export class SubscriptionService implements OnModuleInit {
       });
 
       // Find the storage quota for the plan
-      const planType = sub?.planType || 'free';
-      const plan = await this.prisma.subscriptionPlan.findFirstOrThrow({
-        where: { planType },
-      });
-
+      let plan: PlanQuota | null = null;
+      if (sub?.overridePlan) {
+        plan = safeParseJSON(sub.overridePlan) as PlanQuota;
+      }
       if (!plan) {
-        throw new Error(`No subscription plan found for type ${planType}`);
+        const planType = sub?.planType || 'free';
+        plan = await this.prisma.subscriptionPlan.findFirstOrThrow({
+          where: { planType },
+        });
       }
 
       if (activeMeter) {