diff --git a/src/controllers/template.controller.ts b/src/controllers/template.controller.ts index f7ee689..0e8a8a3 100644 --- a/src/controllers/template.controller.ts +++ b/src/controllers/template.controller.ts @@ -34,7 +34,11 @@ import { TemplateListDto, BatchExecuteDto, } from '../dto/template.dto'; -import { TemplateExecutionEntity, ExecutionStatus, ExecutionType } from '../entities/template-execution.entity'; +import { + TemplateExecutionEntity, + ExecutionStatus, + ExecutionType, +} from '../entities/template-execution.entity'; import { InjectRepository } from '@nestjs/typeorm'; import { Repository } from 'typeorm'; import { ApiCommonResponses } from '../decorators/api-common-responses.decorator'; @@ -54,7 +58,7 @@ export class TemplateController { @InjectRepository(TemplateExecutionEntity) private readonly executionRepository: Repository, private readonly unifiedContentService: UnifiedContentService, - ) { } + ) {} @Post(':templateId/execute') @ApiOperation({ @@ -141,10 +145,22 @@ export class TemplateController { }, }, }) + @SwaggerApiResponse({ + status: 429, + description: '任务数量限制', + schema: { + example: { + code: 429, + message: + '当前账号有3个任务正在进行中，且距开始时间不足5分钟，请稍后再试', + data: null, + }, + }, + }) async executeTemplateByCode( @Param('code') code: string, @Body() body: { imageUrl: string }, - @Request() req + @Request() req, ): Promise> { try { const { imageUrl } = body; @@ -153,13 +169,17 @@ export class TemplateController { throw new HttpException('imageUrl is required', HttpStatus.BAD_REQUEST); } + const userId = req.user.userId; + + // 检查用户当前的任务限制 + await this.checkUserTaskLimit(userId); + // 首先获取模板配置以确定模板类型 const templateConfig = await this.templateFactory.getTemplateByCode(code); if (!templateConfig) { throw new HttpException('Template not found', HttpStatus.NOT_FOUND); } - const userId = req.user.userId; const platform = req.user.platform; // 1. 先进行图片内容审核 @@ -191,13 +211,16 @@ export class TemplateController { // 3. 审核通过，执行模板 const template = await this.templateFactory.createTemplateByCode(code); const taskId = await template.execute(imageUrl); - + // 4. 将任务保存到 TemplateExecutionEntity const execution = this.executionRepository.create({ templateId: templateConfig.id, userId, platform: req.user.platform, - type: templateConfig.templateType === TemplateType.VIDEO ? ExecutionType.VIDEO : ExecutionType.IMAGE, + type: + templateConfig.templateType === TemplateType.VIDEO + ? ExecutionType.VIDEO + : ExecutionType.IMAGE, prompt: '', // 可以从请求参数中获取，如果有的话 inputImageUrl: imageUrl, taskId: taskId, // 保存外部系统返回的任务ID，用于回调匹配 @@ -206,13 +229,13 @@ export class TemplateController { creditCost: templateConfig.creditCost, startedAt: new Date(), }); - + const savedExecution = await this.executionRepository.save(execution); // 返回任务id (执行记录的ID) 以及审核信息 return ResponseUtil.success(savedExecution.id, '模板执行已启动'); } catch (error) { - console.error(error) + console.error(error); throw new HttpException( error.message || 'Template execution failed', error instanceof HttpException ? error.getStatus() : HttpStatus.INTERNAL_SERVER_ERROR, @@ -220,6 +243,48 @@ export class TemplateController { } } + /** + * 检查用户当前任务限制 + * @param userId 用户ID + * @throws HttpException 如果超出任务限制 + */ + private async checkUserTaskLimit(userId: string): Promise { + // 查询用户当前进行中的任务 + const processingTasks = await this.executionRepository.find({ + where: { + userId, + status: ExecutionStatus.PROCESSING, + }, + order: { + startedAt: 'ASC', + }, + }); + + // 如果当前没有进行中的任务，允许执行 + if (processingTasks.length === 0) { + return; + } + + // 如果当前进行中任务数量已达到3个，需要检查时间限制 + if (processingTasks.length >= 3) { + const now = new Date(); + const fiveMinutesAgo = new Date(now.getTime() - 5 * 60 * 1000); + + // 检查是否有任务在5分钟内开始 + const recentTasks = processingTasks.filter( + (task) => + task.startedAt && task.startedAt.getTime() > fiveMinutesAgo.getTime(), + ); + + if (recentTasks.length > 3) { + throw new HttpException( + `有${recentTasks.length}个任务正在进行中，请稍后再试`, + HttpStatus.TOO_MANY_REQUESTS, + ); + } + } + } + @Post(':templateId/batch-execute') @ApiOperation({ summary: '批量执行模板', @@ -474,7 +539,9 @@ export class TemplateController { * @returns 执行进度信息 */ @Get('execution/:taskId/progress') - async getExecutionProgress(@Param('taskId', ParseIntPipe) taskId: number): Promise> { + async getExecutionProgress( + @Param('taskId', ParseIntPipe) taskId: number, + ): Promise> { try { const execution = await this.executionRepository.findOne({ where: { id: taskId }, @@ -488,27 +555,30 @@ export class TemplateController { ); } - return ResponseUtil.success({ - taskId: execution.id, - templateId: execution.templateId, - templateName: execution.template?.name, - userId: execution.userId, - platform: execution.platform, - type: execution.type, - status: execution.status, - progress: execution.progress, - inputImageUrl: execution.inputImageUrl, - outputUrl: execution.outputUrl, - thumbnailUrl: execution.thumbnailUrl, - errorMessage: execution.errorMessage, - creditCost: execution.creditCost, - startedAt: execution.startedAt, - completedAt: execution.completedAt, - executionDuration: execution.executionDuration, - createdAt: execution.createdAt, - updatedAt: execution.updatedAt, - executionResult: execution.executionResult - }, '获取执行进度成功'); + return ResponseUtil.success( + { + taskId: execution.id, + templateId: execution.templateId, + templateName: execution.template?.name, + userId: execution.userId, + platform: execution.platform, + type: execution.type, + status: execution.status, + progress: execution.progress, + inputImageUrl: execution.inputImageUrl, + outputUrl: execution.outputUrl, + thumbnailUrl: execution.thumbnailUrl, + errorMessage: execution.errorMessage, + creditCost: execution.creditCost, + startedAt: execution.startedAt, + completedAt: execution.completedAt, + executionDuration: execution.executionDuration, + createdAt: execution.createdAt, + updatedAt: execution.updatedAt, + executionResult: execution.executionResult, + }, + '获取执行进度成功', + ); } catch (error) { if (error instanceof HttpException) { throw error;