diff --git a/prisma/migrations/20260523192858_add_task_log/migration.sql b/prisma/migrations/20260523192858_add_task_log/migration.sql deleted file mode 100644 index 9a23ab0..0000000 --- a/prisma/migrations/20260523192858_add_task_log/migration.sql +++ /dev/null @@ -1,9 +0,0 @@ -CREATE TABLE TaskLog ( - id VARCHAR(191) NOT NULL, queueName VARCHAR(64) NOT NULL, jobId VARCHAR(100) NOT NULL, - status VARCHAR(16) NOT NULL DEFAULT 'enqueued', payload JSON, - error TEXT, attempts INT NOT NULL DEFAULT 0, - createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), - updatedAt DATETIME(3) NOT NULL, - INDEX TaskLog_queueName_idx(queueName), INDEX TaskLog_status_idx(status), - INDEX TaskLog_createdAt_idx(createdAt), PRIMARY KEY (id) -) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; diff --git a/prisma/migrations/20260523201015_add_quota_billing/migration.sql b/prisma/migrations/20260523201015_add_quota_billing/migration.sql new file mode 100644 index 0000000..75b2018 --- /dev/null +++ b/prisma/migrations/20260523201015_add_quota_billing/migration.sql @@ -0,0 +1,38 @@ +CREATE TABLE MembershipPlan ( + id VARCHAR(191) NOT NULL, name VARCHAR(100) NOT NULL, description VARCHAR(500), + price DOUBLE NOT NULL DEFAULT 0, currency VARCHAR(8) NOT NULL DEFAULT 'CNY', + aiCallLimit INT NOT NULL DEFAULT 100, storageLimit INT NOT NULL DEFAULT 0, + ocrLimit INT NOT NULL DEFAULT 0, visionLimit INT NOT NULL DEFAULT 0, + embeddingLimit INT NOT NULL DEFAULT 0, active BOOLEAN NOT NULL DEFAULT true, + createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + updatedAt DATETIME(3) NOT NULL, PRIMARY KEY (id) +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; + +CREATE TABLE UserMembership ( + id VARCHAR(191) NOT NULL, userId VARCHAR(191) NOT NULL, planId VARCHAR(191) NOT NULL, + startedAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), expiresAt DATETIME(3), + autoRenew BOOLEAN NOT NULL DEFAULT false, active BOOLEAN NOT NULL DEFAULT true, + createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + updatedAt DATETIME(3) NOT NULL, + INDEX UserMembership_userId_idx(userId), + PRIMARY KEY (id), + CONSTRAINT UserMembership_planId_fkey FOREIGN KEY (planId) REFERENCES MembershipPlan(id) ON DELETE RESTRICT ON UPDATE CASCADE, + CONSTRAINT UserMembership_userId_fkey FOREIGN KEY (userId) REFERENCES User(id) ON DELETE RESTRICT ON UPDATE CASCADE +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; + +CREATE TABLE QuotaUsage ( + id VARCHAR(191) NOT NULL, userId VARCHAR(191) NOT NULL, + quotaType VARCHAR(32) NOT NULL, amount INT NOT NULL, resource VARCHAR(255), + createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + INDEX QuotaUsage_userId_quotaType_idx(userId, quotaType), + INDEX QuotaUsage_createdAt_idx(createdAt), PRIMARY KEY (id) +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; + +CREATE TABLE CostDailySummary ( + id VARCHAR(191) NOT NULL, date DATETIME(3) NOT NULL, + provider VARCHAR(32) NOT NULL, model VARCHAR(100), + calls INT NOT NULL DEFAULT 0, tokens INT NOT NULL DEFAULT 0, cost DOUBLE NOT NULL DEFAULT 0, + createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + UNIQUE INDEX CostDailySummary_date_provider_model_key(date, provider, model), + PRIMARY KEY (id) +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; diff --git a/prisma/migrations/20260523201212_add_user_membership_quota/migration.sql b/prisma/migrations/20260523201212_add_user_membership_quota/migration.sql new file mode 100644 index 0000000..ea768aa --- /dev/null +++ b/prisma/migrations/20260523201212_add_user_membership_quota/migration.sql @@ -0,0 +1,21 @@ +ALTER TABLE MembershipPlan DROP COLUMN IF EXISTS maxKnowledgeBases, DROP COLUMN IF EXISTS maxStorageBytes, DROP COLUMN IF EXISTS maxFileSizeBytes, DROP COLUMN IF EXISTS monthlyOcrPages, DROP COLUMN IF EXISTS monthlyVisionPages, DROP COLUMN IF EXISTS monthlyChatCount, DROP COLUMN IF EXISTS monthlyAiAnalysisCount, DROP COLUMN IF EXISTS monthlyRecallCount, DROP COLUMN IF EXISTS monthlyCardGenCount; +CREATE TABLE IF NOT EXISTS UserMembership ( + id VARCHAR(191) NOT NULL, userId VARCHAR(191) NOT NULL, planId VARCHAR(191) NOT NULL, + startedAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), expiresAt DATETIME(3), + active BOOLEAN NOT NULL DEFAULT true, createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + updatedAt DATETIME(3) NOT NULL, + INDEX UserMembership_userId_idx(userId), PRIMARY KEY (id), + CONSTRAINT UserMembership_planId_fkey FOREIGN KEY (planId) REFERENCES MembershipPlan(id) ON DELETE RESTRICT ON UPDATE CASCADE +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; +CREATE TABLE IF NOT EXISTS QuotaUsage ( + id VARCHAR(191) NOT NULL, userId VARCHAR(191) NOT NULL, quotaType VARCHAR(32) NOT NULL, + amount INT NOT NULL, resource VARCHAR(255), + createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + INDEX QuotaUsage_userId_quotaType_idx(userId, quotaType), INDEX QuotaUsage_createdAt_idx(createdAt), PRIMARY KEY (id) +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; +CREATE TABLE IF NOT EXISTS CostDailySummary ( + id VARCHAR(191) NOT NULL, date DATETIME(3) NOT NULL, provider VARCHAR(32) NOT NULL, + model VARCHAR(100), calls INT NOT NULL DEFAULT 0, tokens INT NOT NULL DEFAULT 0, cost DOUBLE NOT NULL DEFAULT 0, + createdAt DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), + UNIQUE INDEX CostDailySummary_date_provider_model_key(date, provider, model), PRIMARY KEY (id) +) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index ea494a0..8e635e5 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -23,6 +23,7 @@ model User { authAccounts AuthAccount[] refreshTokens RefreshToken[] + memberships UserMembership[] profile UserProfile? preferences UserPreference? consents UserConsent[] @@ -793,6 +794,7 @@ model MembershipPlan { monthlyRecallCount Int @default(0) monthlyCardGenCount Int @default(0) isActive Boolean @default(true) + memberships UserMembership[] createdAt DateTime @default(now()) updatedAt DateTime @updatedAt } @@ -978,3 +980,44 @@ model TaskLog { @@index([status]) @@index([createdAt]) } + +model UserMembership { + id String @id @default(cuid()) + userId String + planId String + startedAt DateTime @default(now()) + expiresAt DateTime? + active Boolean @default(true) + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt + + user User @relation(fields: [userId], references: [id]) + plan MembershipPlan @relation(fields: [planId], references: [id]) + + @@index([userId]) +} + +model QuotaUsage { + id String @id @default(cuid()) + userId String + quotaType String @db.VarChar(32) + amount Int + resource String? @db.VarChar(255) + createdAt DateTime @default(now()) + + @@index([userId, quotaType]) + @@index([createdAt]) +} + +model CostDailySummary { + id String @id @default(cuid()) + date DateTime + provider String @db.VarChar(32) + model String? @db.VarChar(100) + calls Int @default(0) + tokens Int @default(0) + cost Float @default(0) + createdAt DateTime @default(now()) + + @@unique([date, provider, model]) +} diff --git a/src/app.module.ts b/src/app.module.ts index afa6941..c20247c 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -17,6 +17,7 @@ import { AdminAuthModule } from './modules/admin-auth/admin-auth.module'; import { AdminDashboardModule } from './modules/admin-dashboard/admin-dashboard.module'; import { AdminUsersModule } from './modules/admin-users/admin-users.module'; import { ContentSafetyModule } from './modules/content-safety/content-safety.module'; +import { QuotaModule } from './modules/quota/quota.module'; import { AdminMetricsModule } from './modules/admin-metrics/admin-metrics.module'; import { AdminThrottleModule } from './modules/admin-throttle/admin-throttle.module'; import { AppConfigModule } from './modules/config/config.module'; @@ -107,6 +108,7 @@ import appleConfig from './config/apple.config'; AdminDashboardModule, AdminUsersModule, ContentSafetyModule, + QuotaModule, AdminMetricsModule, AdminThrottleModule, AppConfigModule, diff --git a/src/modules/quota/quota.controller.ts b/src/modules/quota/quota.controller.ts new file mode 100644 index 0000000..39b762d --- /dev/null +++ b/src/modules/quota/quota.controller.ts @@ -0,0 +1,43 @@ +import { Controller, Get, Post, Param, Body, UseGuards } from '@nestjs/common'; +import { ApiTags, ApiBearerAuth, ApiOperation } from '@nestjs/swagger'; +import { PrismaService } from '../../infrastructure/database/prisma.service'; +import { QuotaService } from './quota.service'; +import { AdminAuthGuard } from '../../common/guards/admin-auth.guard'; +import { AdminRolesGuard } from '../../common/guards/admin-roles.guard'; +import { AdminRoles } from '../../common/decorators/admin-roles.decorator'; +import type { AdminRole } from '../../common/types/admin-role.enum'; + +@ApiTags('admin-quota') +@Controller('admin-api/quota') +@UseGuards(AdminAuthGuard, AdminRolesGuard) +@ApiBearerAuth() +export class QuotaController { + constructor(private readonly prisma: PrismaService, private readonly quota: QuotaService) {} + + @Get('plans') + @AdminRoles('SUPER_ADMIN' as AdminRole) + async plans() { return this.prisma.membershipPlan.findMany() } + + @Post('plans') + @AdminRoles('SUPER_ADMIN' as AdminRole) + async createPlan(@Body() d: any) { return this.prisma.membershipPlan.create({ data: d }) } + + @Get('memberships') + @AdminRoles('SUPER_ADMIN' as AdminRole) + async memberships() { return this.prisma.userMembership.findMany({ include: { user: { select: { email: true } }, plan: true }, orderBy: { createdAt: 'desc' }, take: 50 }) } + + @Get('usage/:userId') + @AdminRoles('SUPER_ADMIN' as AdminRole) + async userUsage(@Param('userId') userId: string) { + const types = ['ai_call', 'ocr', 'storage', 'vision', 'embedding']; + const result: Record = {}; + for (const t of types) { + result[t] = { used: await this.quota.getUsed(userId, t), limit: await this.quota.getRemaining(userId, t) + await this.quota.getUsed(userId, t) }; + } + return result; + } + + @Get('costs') + @AdminRoles('SUPER_ADMIN' as AdminRole) + async costs() { return this.prisma.costDailySummary.findMany({ orderBy: { date: 'desc' }, take: 30 }) } +} diff --git a/src/modules/quota/quota.module.ts b/src/modules/quota/quota.module.ts new file mode 100644 index 0000000..b6c1570 --- /dev/null +++ b/src/modules/quota/quota.module.ts @@ -0,0 +1,14 @@ +import { Module } from '@nestjs/common'; +import { QuotaController } from './quota.controller'; +import { QuotaService } from './quota.service'; +import { PrismaService } from '../../infrastructure/database/prisma.service'; +import { RedisService } from '../../infrastructure/redis/redis.service'; +import { AdminAuthGuard } from '../../common/guards/admin-auth.guard'; +import { AdminRolesGuard } from '../../common/guards/admin-roles.guard'; + +@Module({ + controllers: [QuotaController], + providers: [QuotaService, PrismaService, RedisService, AdminAuthGuard, AdminRolesGuard], + exports: [QuotaService], +}) +export class QuotaModule {} diff --git a/src/modules/quota/quota.service.ts b/src/modules/quota/quota.service.ts new file mode 100644 index 0000000..cbcc66c --- /dev/null +++ b/src/modules/quota/quota.service.ts @@ -0,0 +1,79 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { PrismaService } from '../../infrastructure/database/prisma.service'; +import { RedisService } from '../../infrastructure/redis/redis.service'; + +const QUOTA_PREFIX = 'quota:'; +const RESET_DAILY = 86400; + +@Injectable() +export class QuotaService { + private readonly logger = new Logger(QuotaService.name); + + constructor( + private readonly prisma: PrismaService, + private readonly redis: RedisService, + ) {} + + async getRemaining(userId: string, quotaType: string): Promise { + const planLimit = await this.getPlanLimit(userId, quotaType); + const used = await this.getUsed(userId, quotaType); + return Math.max(0, planLimit - used); + } + + async checkAndDeduct(userId: string, quotaType: string, amount: number): Promise { + const remaining = await this.getRemaining(userId, quotaType); + if (remaining < amount) { + this.logger.warn(`Quota exhausted: ${userId} ${quotaType} need=${amount} have=${remaining}`); + return false; + } + + const key = `${QUOTA_PREFIX}${userId}:${quotaType}`; + try { + await this.redis.incr(key); + await this.redis.expire(key, RESET_DAILY); + } catch {} + + await this.prisma.quotaUsage.create({ + data: { userId, quotaType, amount, resource: quotaType }, + }); + + return true; + } + + async getUsed(userId: string, quotaType: string): Promise { + try { + const key = `${QUOTA_PREFIX}${userId}:${quotaType}`; + const val = await this.redis.get(key); + if (val) return parseInt(val); + } catch {} + + const result = await this.prisma.quotaUsage.aggregate({ + _sum: { amount: true }, + where: { userId, quotaType, createdAt: { gte: new Date(Date.now() - RESET_DAILY * 1000) } }, + }); + return result._sum.amount || 0; + } + + private async getPlanLimit(userId: string, quotaType: string): Promise { + const membership = await this.prisma.userMembership.findFirst({ + where: { userId, active: true }, + include: { plan: true }, + }); + + if (!membership?.plan) { + // Free tier defaults + const defaults: Record = { ai_call: 10, ocr: 5, storage: 0, vision: 3, embedding: 50 }; + return defaults[quotaType] || 10; + } + + const plan = membership.plan; + const limits: Record = { + ai_call: plan.monthlyChatCount, + ocr: plan.monthlyOcrPages, + storage: Number(plan.maxStorageBytes), + vision: plan.monthlyVisionPages, + embedding: plan.monthlyRecallCount, + }; + return limits[quotaType] || 10; + } +}