feat(quota): wire ai_chat_messages quota into chat controller
All checks were successful
Deploy API Server / build-and-unit (push) Successful in 39s
Deploy API Server / current-integration (push) Successful in 33s
Deploy API Server / backward-compat (push) Successful in 0s
Deploy API Server / deploy (push) Successful in 1m5s

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

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
wangdl 2026-06-26 20:21:49 +08:00
parent 6a5ab31ddd
commit 807488a929
2 changed files with 36 additions and 9 deletions

View File

@ -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 { ApiTags, ApiOperation, ApiBearerAuth } from '@nestjs/swagger';
import type { Response } from 'express'; import type { Response } from 'express';
import { RagChatService } from './rag-chat.service'; import { RagChatService } from './rag-chat.service';
import type { CreateSessionResult } 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 { CurrentUser } from '../../common/decorators/current-user.decorator';
import type { UserPayload } from '../../common/types'; import type { UserPayload } from '../../common/types';
@ -10,7 +11,10 @@ import type { UserPayload } from '../../common/types';
@Controller('rag-chat') @Controller('rag-chat')
@ApiBearerAuth() @ApiBearerAuth()
export class RagChatController { export class RagChatController {
constructor(private readonly svc: RagChatService) {} constructor(
private readonly svc: RagChatService,
private readonly quotaGuard: QuotaGuardService,
) {}
@Post('sessions') @Post('sessions')
@ApiOperation({ summary: '创建/打开对话 (open-or-create)' }) @ApiOperation({ summary: '创建/打开对话 (open-or-create)' })
@ -67,23 +71,45 @@ export class RagChatController {
@Post('sessions/:id/messages') @Post('sessions/:id/messages')
@ApiOperation({ summary: '发送消息(同步)' }) @ApiOperation({ summary: '发送消息(同步)' })
async sendMessage(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }) { 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') @Post('sessions/:id/stream')
@ApiOperation({ summary: '发送消息SSE 流式,支持思考过程)' }) @ApiOperation({ summary: '发送消息SSE 流式,支持思考过程)' })
async sendMessageStream(@CurrentUser() user: UserPayload, @Param('id') id: string, @Body() dto: { content: string }, @Res() res: Response) { 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('Content-Type', 'text/event-stream; charset=utf-8');
res.setHeader('Cache-Control', 'no-cache'); res.setHeader('Cache-Control', 'no-cache');
res.setHeader('Connection', 'keep-alive'); res.setHeader('Connection', 'keep-alive');
res.setHeader('X-Accel-Buffering', 'no'); res.setHeader('X-Accel-Buffering', 'no');
const userId = String(user.id); try {
for await (const chunk of this.svc.sendMessageStream(userId, id, dto.content)) { for await (const chunk of this.svc.sendMessageStream(userId, id, dto.content)) {
if (res.destroyed) break; if (res.destroyed) { this.quotaGuard.cancel(opId).catch(() => {}); break; }
res.write(`data: ${JSON.stringify(chunk)}\n\n`); 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') @Patch('sessions/:id')

View File

@ -4,9 +4,10 @@ import { AdminRagChatController } from './admin-rag-chat.controller';
import { RagChatService } from './rag-chat.service'; import { RagChatService } from './rag-chat.service';
import { PrismaService } from '../../infrastructure/database/prisma.service'; import { PrismaService } from '../../infrastructure/database/prisma.service';
import { AiModule } from '../ai/ai.module'; import { AiModule } from '../ai/ai.module';
import { MembershipModule } from '../membership/membership.module';
@Module({ @Module({
imports: [AiModule], imports: [AiModule, MembershipModule],
controllers: [RagChatController, AdminRagChatController], controllers: [RagChatController, AdminRagChatController],
providers: [RagChatService, PrismaService], providers: [RagChatService, PrismaService],
exports: [RagChatService], exports: [RagChatService],