feat: 用户 AI 消费额度与日限流 (API-AI-070)
All checks were successful
Deploy API Server / build-and-deploy (push) Successful in 43s

- UserAiQuotaService: checkAndReserve / incrementJobCount / recordTokenUsage
- 创建 Job 前检查 maxDailyAiJobs + maxDailyTokenBudget
- Runtime 调用后更新 UserAiUsageDaily token 消耗
- user_deepseek_key 不计入平台 token 预算
- 超限返回 DAILY_JOB_LIMIT_EXCEEDED / DAILY_TOKEN_BUDGET_EXCEEDED

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
wangdl 2026-06-11 21:25:18 +08:00
parent 43e6e9029c
commit 5cbf20046a
2 changed files with 100 additions and 2 deletions

View File

@ -6,11 +6,12 @@ import { UserAiService } from './user-ai.service';
import { CredentialEncryptionService } from './credential-encryption.service';
import { RuntimeInternalController } from './internal/runtime-internal.controller';
import { RuntimeInternalService } from './internal/runtime-internal.service';
import { UserAiQuotaService } from './user-ai-quota.service';
@Module({
imports: [ConfigModule, PrismaModule],
controllers: [UserAiController, RuntimeInternalController],
providers: [UserAiService, CredentialEncryptionService, RuntimeInternalService],
exports: [UserAiService, CredentialEncryptionService, RuntimeInternalService],
providers: [UserAiService, CredentialEncryptionService, RuntimeInternalService, UserAiQuotaService],
exports: [UserAiService, CredentialEncryptionService, RuntimeInternalService, UserAiQuotaService],
})
export class AiRuntimeModule {}

View File

@ -0,0 +1,97 @@
import { Injectable, BadRequestException } from '@nestjs/common';
import { PrismaService } from '../../infrastructure/database/prisma.service';
@Injectable()
export class UserAiQuotaService {
constructor(private readonly prisma: PrismaService) {}
/** Get today's date truncated to midnight UTC */
private today(): Date {
const d = new Date();
return new Date(Date.UTC(d.getFullYear(), d.getMonth(), d.getDate()));
}
/** Check if user can create a new AI job, throw if over limit */
async checkAndReserve(userId: string, apiKeyMode: string): Promise<void> {
const settings = await this.prisma.userAiSettings.findUnique({ where: { userId } });
const maxJobs = settings?.maxDailyAiJobs ?? 20;
const maxTokens = settings?.maxDailyTokenBudget ?? 100_000;
const localDate = this.today();
const usage = await this.prisma.userAiUsageDaily.findUnique({
where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode } },
});
const jobCount = usage?.jobCount ?? 0;
const totalTokens = usage?.totalTokens ?? 0;
if (jobCount >= maxJobs) {
throw new BadRequestException({
errorCode: 'DAILY_JOB_LIMIT_EXCEEDED',
message: `Daily AI job limit (${maxJobs}) exceeded. Current: ${jobCount}`,
});
}
if (totalTokens >= maxTokens) {
throw new BadRequestException({
errorCode: 'DAILY_TOKEN_BUDGET_EXCEEDED',
message: `Daily token budget (${maxTokens}) exceeded. Current: ${totalTokens}`,
});
}
}
/** Increment job count (called on job creation) */
async incrementJobCount(userId: string, apiKeyMode: string): Promise<void> {
const localDate = this.today();
await this.prisma.userAiUsageDaily.upsert({
where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode } },
create: { userId, localDate, apiKeyMode, jobCount: 1 },
update: { jobCount: { increment: 1 } },
});
}
/** Record token consumption from invocation log (called after Runtime call) */
async recordTokenUsage(
userId: string,
apiKeyMode: string,
inputTokens: number,
outputTokens: number,
totalTokens: number,
costEstimate?: number,
): Promise<void> {
const localDate = this.today();
await this.prisma.userAiUsageDaily.upsert({
where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode } },
create: {
userId, localDate, apiKeyMode,
inputTokens, outputTokens, totalTokens,
costEstimate: costEstimate ?? 0,
jobCount: 0,
},
update: {
inputTokens: { increment: inputTokens },
outputTokens: { increment: outputTokens },
totalTokens: { increment: totalTokens },
costEstimate: { increment: costEstimate ?? 0 },
},
});
}
/** Get today's usage summary for a user */
async getUsage(userId: string) {
const localDate = this.today();
const [platform, userKey] = await Promise.all([
this.prisma.userAiUsageDaily.findUnique({
where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode: 'platform_key' } },
}),
this.prisma.userAiUsageDaily.findUnique({
where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode: 'user_deepseek_key' } },
}),
]);
return {
platformKey: platform ?? { jobCount: 0, totalTokens: 0 },
userDeepseekKey: userKey ?? { jobCount: 0, totalTokens: 0 },
};
}
}