From 807488a929afe68aaa0b7fc5ba6bb21bff0d49ea Mon Sep 17 00:00:00 2001
From: wangdl <wangdl@longde.cloud>
Date: Fri, 26 Jun 2026 20:21:49 +0800
Subject: [PATCH] feat(quota): wire ai_chat_messages quota into chat controller
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Add QuotaGuard check to sendMessage and sendMessageStream
- Stream disconnect → cancel reservation
- Import MembershipModule in RagChatModule

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/modules/rag-chat/rag-chat.controller.ts | 42 +++++++++++++++++----
 src/modules/rag-chat/rag-chat.module.ts     |  3 +-
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/src/modules/rag-chat/rag-chat.controller.ts b/src/modules/rag-chat/rag-chat.controller.ts
index 732b6c4..e3bd9d3 100644
--- a/src/modules/rag-chat/rag-chat.controller.ts
+++ b/src/modules/rag-chat/rag-chat.controller.ts
@@ -1,8 +1,9 @@
-import { Controller, Get, Post, Patch, Delete, Body, Param, Res, Query, HttpCode } from '@nestjs/common';
+import { Controller, Get, Post, Patch, Delete, Body, Param, Res, Query, HttpCode, BadRequestException } from '@nestjs/common';
 import { ApiTags, ApiOperation, ApiBearerAuth } from '@nestjs/swagger';
 import type { Response } from 'express';
 import { RagChatService } from './rag-chat.service';
 import type { CreateSessionResult } from './rag-chat.service';
+import { QuotaGuardService } from '../membership/quota-guard.service';
 import { CurrentUser } from '../../common/decorators/current-user.decorator';
 import type { UserPayload } from '../../common/types';
 
@@ -10,7 +11,10 @@ import type { UserPayload } from '../../common/types';
 @Controller('rag-chat')
 @ApiBearerAuth()
 export class RagChatController {
-  constructor(private readonly svc: RagChatService) {}
+  constructor(
+    private readonly svc: RagChatService,
+    private readonly quotaGuard: QuotaGuardService,
+  ) {}
 
   @Post('sessions')
   @ApiOperation({ summary: '创建/打开对话 (open-or-create)' })
@@ -67,23 +71,45 @@ export class RagChatController {
   @Post('sessions/:id/messages')
   @ApiOperation({ summary: '发送消息（同步）' })
   async sendMessage(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }) {
-    return this.svc.sendMessage(String(user.id), id, dto.content);
+    const userId = String(user.id);
+    const opId = `chat:${id}:${Date.now()}`;
+    const check = await this.quotaGuard.check(userId, 'ai_chat_messages', opId);
+    if (!check.allowed) throw new BadRequestException(this.quotaGuard.buildError(check));
+    try {
+      const result = await this.svc.sendMessage(userId, id, dto.content);
+      this.quotaGuard.confirm(opId).catch(() => {});
+      return result;
+    } catch {
+      this.quotaGuard.cancel(opId).catch(() => {});
+      throw new BadRequestException('AI Chat failed');
+    }
   }
 
   @Post('sessions/:id/stream')
   @ApiOperation({ summary: '发送消息（SSE 流式，支持思考过程）' })
   async sendMessageStream(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }, @Res() res: Response) {
+    const userId = String(user.id);
+    const opId = `chat_stream:${id}:${Date.now()}`;
+    const check = await this.quotaGuard.check(userId, 'ai_chat_messages', opId);
+    if (!check.allowed) {
+      res.status(429).json(this.quotaGuard.buildError(check));
+      return;
+    }
     res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
     res.setHeader('Cache-Control', 'no-cache');
     res.setHeader('Connection', 'keep-alive');
     res.setHeader('X-Accel-Buffering', 'no');
 
-    const userId = String(user.id);
-    for await (const chunk of this.svc.sendMessageStream(userId, id, dto.content)) {
-      if (res.destroyed) break;
-      res.write(`data: ${JSON.stringify(chunk)}\n\n`);
+    try {
+      for await (const chunk of this.svc.sendMessageStream(userId, id, dto.content)) {
+        if (res.destroyed) { this.quotaGuard.cancel(opId).catch(() => {}); break; }
+        res.write(`data: ${JSON.stringify(chunk)}\n\n`);
+      }
+      if (!res.destroyed) { res.end(); this.quotaGuard.confirm(opId).catch(() => {}); }
+    } catch {
+      if (!res.destroyed) res.end();
+      this.quotaGuard.cancel(opId).catch(() => {});
     }
-    res.end();
   }
 
   @Patch('sessions/:id')
diff --git a/src/modules/rag-chat/rag-chat.module.ts b/src/modules/rag-chat/rag-chat.module.ts
index 1c1bbd0..6628469 100644
--- a/src/modules/rag-chat/rag-chat.module.ts
+++ b/src/modules/rag-chat/rag-chat.module.ts
@@ -4,9 +4,10 @@ import { AdminRagChatController } from './admin-rag-chat.controller';
 import { RagChatService } from './rag-chat.service';
 import { PrismaService } from '../../infrastructure/database/prisma.service';
 import { AiModule } from '../ai/ai.module';
+import { MembershipModule } from '../membership/membership.module';
 
 @Module({
-  imports: [AiModule],
+  imports: [AiModule, MembershipModule],
   controllers: [RagChatController, AdminRagChatController],
   providers: [RagChatService, PrismaService],
   exports: [RagChatService],