From fef069584547b88bca0ae32f1d822b333f1f4ba0 Mon Sep 17 00:00:00 2001
From: wangdl <wangdl@longde.cloud>
Date: Fri, 26 Jun 2026 20:28:29 +0800
Subject: [PATCH] feat(quota): wire knowledge_base_count quota into KB
 controller

---
 .../knowledge-base.controller.ts               | 18 +++++++++++++++---
 .../knowledge-base/knowledge-base.module.ts    |  2 ++
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/modules/knowledge-base/knowledge-base.controller.ts b/src/modules/knowledge-base/knowledge-base.controller.ts
index 298c716..f4ecd23 100644
--- a/src/modules/knowledge-base/knowledge-base.controller.ts
+++ b/src/modules/knowledge-base/knowledge-base.controller.ts
@@ -1,6 +1,7 @@
-import { Controller, Get, Post, Patch, Delete, Body, Param, Query } from '@nestjs/common';
+import { Controller, Get, Post, Patch, Delete, Body, Param, Query, BadRequestException } from '@nestjs/common';
 import { ApiTags, ApiOperation, ApiResponse } from '@nestjs/swagger';
 import { KnowledgeBaseService } from './knowledge-base.service';
+import { QuotaGuardService } from '../membership/quota-guard.service';
 import { CurrentUser } from '../../common/decorators/current-user.decorator';
 import { PaginationDto } from '../../common/dto/pagination.dto';
 import type { UserPayload } from '../../common/types';
@@ -8,12 +9,23 @@ import type { UserPayload } from '../../common/types';
 @ApiTags('knowledge-base')
 @Controller('knowledge-bases')
 export class KnowledgeBaseController {
-  constructor(private readonly service: KnowledgeBaseService) {}
+  constructor(
+    private readonly service: KnowledgeBaseService,
+    private readonly quotaGuard: QuotaGuardService,
+  ) {}
 
   @Post()
   @ApiOperation({ summary: '创建知识库' })
   async create(@CurrentUser() user: UserPayload, @Body() dto: any) {
-    return this.service.create(String(user?.id || 'anonymous'), dto);
+    const userId = String(user?.id || 'anonymous');
+    const opId = `kb_create:${userId}:${Date.now()}`;
+    const check = await this.quotaGuard.check(userId, 'knowledge_base_count', opId);
+    if (!check.allowed) throw new BadRequestException(this.quotaGuard.buildError(check));
+    try {
+      const result = await this.service.create(userId, dto);
+      this.quotaGuard.confirm(opId).catch(() => {});
+      return result;
+    } catch { this.quotaGuard.cancel(opId).catch(() => {}); throw new BadRequestException('Knowledge base creation failed'); }
   }
 
   @Get()
diff --git a/src/modules/knowledge-base/knowledge-base.module.ts b/src/modules/knowledge-base/knowledge-base.module.ts
index a615287..f395d6e 100644
--- a/src/modules/knowledge-base/knowledge-base.module.ts
+++ b/src/modules/knowledge-base/knowledge-base.module.ts
@@ -4,8 +4,10 @@ import { KnowledgeBaseService } from './knowledge-base.service';
 import { KnowledgeBaseRepository } from './knowledge-base.repository';
 import { SystemKnowledgeBaseSeed } from './system-kb.seed';
 import { PrismaService } from '../../infrastructure/database/prisma.service';
+import { MembershipModule } from '../membership/membership.module';
 
 @Module({
+  imports: [MembershipModule],
   controllers: [KnowledgeBaseController],
   providers: [KnowledgeBaseService, KnowledgeBaseRepository, SystemKnowledgeBaseSeed, PrismaService],
   exports: [KnowledgeBaseService],