From 8d19c5bf6eb479d881ab043a65b8c6af19095688 Mon Sep 17 00:00:00 2001
From: mrcfps <mrc@powerformer.com>
Date: Mon, 10 Mar 2025 15:50:30 +0800
Subject: [PATCH] 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
---
 apps/api/prisma/schema.prisma                 |  2 +
 apps/api/src/subscription/subscription.dto.ts |  6 +++
 .../src/subscription/subscription.service.ts  | 43 ++++++++-----------
 3 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/apps/api/prisma/schema.prisma b/apps/api/prisma/schema.prisma
index c85cc6c16..ab2fe40fb 100644
--- 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
index 09204a62a..faf72449a 100644
--- 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
index 19b543457..594982c00 100644
--- 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) {