Merge pull request #411 from refly-ai/fix/subscription-webhook

Fix/subscription webhook
refly-ai · Jan 24, 2025 · 377cec4 · 377cec4
2 parents b181538 + 1a5c3d4
commit 377cec4
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/apps/api/src/app.module.ts b/apps/api/src/app.module.ts
@@ -1,8 +1,10 @@
 import { ExecutionContext, Module } from '@nestjs/common';
 import { BullModule } from '@nestjs/bullmq';
+import { APP_GUARD } from '@nestjs/core';
 import { ConfigModule, ConfigService } from '@nestjs/config';
 import { LoggerModule } from 'nestjs-pino';
 import { SkipThrottle, ThrottlerGuard, ThrottlerModule, seconds } from '@nestjs/throttler';
+import { ThrottlerStorageRedisService } from '@nest-lab/throttler-storage-redis';
 import api from '@opentelemetry/api';
 
 import { CommonModule } from '@/common/common.module';
@@ -25,7 +27,6 @@ import { CanvasModule } from './canvas/canvas.module';
 import { CollabModule } from './collab/collab.module';
 import { ActionModule } from './action/action.module';
 import { RedisService } from '@/common/redis.service';
-import { APP_GUARD } from '@nestjs/core';
 
 class CustomThrottlerGuard extends ThrottlerGuard {
   protected async shouldSkip(context: ExecutionContext): Promise<boolean> {
@@ -60,7 +61,7 @@ class CustomThrottlerGuard extends ThrottlerGuard {
           },
         ],
         getTracker: (req) => (req.ips?.length ? req.ips[0] : req.ip),
-        // storage: new ThrottlerStorageRedisService(redis),
+        storage: new ThrottlerStorageRedisService(redis),
       }),
     }),
     LoggerModule.forRoot({

diff --git a/apps/api/src/subscription/subscription.service.ts b/apps/api/src/subscription/subscription.service.ts
@@ -220,8 +220,12 @@ export class SubscriptionService implements OnModuleInit {
           subscriptionId: sub.subscriptionId,
           startAt: startOfDay(now),
           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,
         },
@@ -236,6 +240,7 @@ 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'),
         },
@@ -296,6 +301,7 @@ 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:
@@ -504,8 +510,13 @@ export class SubscriptionService implements OnModuleInit {
       }
 
       // Otherwise, create a new meter
-      const startAt = lastMeter?.endAt ?? startOfDay(now);
+      let startAt: Date;
       const planType = sub?.planType || 'free';
+      if (planType === 'free') {
+        startAt = startOfDay(now);
+      } else {
+        startAt = lastMeter?.endAt ?? startOfDay(now);
+      }
 
       // For free plan, the meter ends at the next day
       // For paid plan, the meter ends at the next month