feat(api-ai): add user and ip rate limiting

apps/api/src/ai/ai-rate-limit.service.ts

@@ -0,0 +1,123 @@
+import { Injectable } from "@nestjs/common";
+import { ConfigService } from "@nestjs/config";
+
+type AiRateLimitBucket = {
+  count: number;
+  resetAt: number;
+};
+
+export type AiRateLimitResult =
+  | {
+      allowed: true;
+    }
+  | {
+      allowed: false;
+      reason: "USER" | "IP";
+      retryAfterMs: number;
+      limit: number;
+      windowMs: number;
+    };
+
+@Injectable()
+export class AiRateLimitService {
+  private readonly userBuckets = new Map<string, AiRateLimitBucket>();
+  private readonly ipBuckets = new Map<string, AiRateLimitBucket>();
+  private readonly windowMs: number;
+  private readonly userLimit: number;
+  private readonly ipLimit: number;
+
+  constructor(private readonly configService: ConfigService) {
+    this.windowMs = this.readPositiveInt("AI_RATE_LIMIT_WINDOW_MS", 60_000);
+    this.userLimit = this.readPositiveInt("AI_RATE_LIMIT_USER_MAX", 20);
+    this.ipLimit = this.readPositiveInt("AI_RATE_LIMIT_IP_MAX", 60);
+  }
+
+  consume(userId: string, clientIp: string | null): AiRateLimitResult {
+    const now = Date.now();
+    const userBucket = this.getBucket(this.userBuckets, userId, now);
+    if (userBucket.count >= this.userLimit) {
+      return {
+        allowed: false,
+        reason: "USER",
+        retryAfterMs: Math.max(0, userBucket.resetAt - now),
+        limit: this.userLimit,
+        windowMs: this.windowMs
+      };
+    }
+
+    const normalizedIp = this.normalizeIp(clientIp);
+    const ipBucket = normalizedIp ? this.getBucket(this.ipBuckets, normalizedIp, now) : null;
+    if (ipBucket && ipBucket.count >= this.ipLimit) {
+      return {
+        allowed: false,
+        reason: "IP",
+        retryAfterMs: Math.max(0, ipBucket.resetAt - now),
+        limit: this.ipLimit,
+        windowMs: this.windowMs
+      };
+    }
+
+    userBucket.count += 1;
+    if (ipBucket) {
+      ipBucket.count += 1;
+    }
+
+    this.cleanupExpiredBuckets(this.userBuckets, now);
+    this.cleanupExpiredBuckets(this.ipBuckets, now);
+
+    return {
+      allowed: true
+    };
+  }
+
+  private getBucket(
+    buckets: Map<string, AiRateLimitBucket>,
+    key: string,
+    now: number
+  ): AiRateLimitBucket {
+    const currentBucket = buckets.get(key);
+    if (!currentBucket || now >= currentBucket.resetAt) {
+      const nextBucket: AiRateLimitBucket = {
+        count: 0,
+        resetAt: now + this.windowMs
+      };
+      buckets.set(key, nextBucket);
+      return nextBucket;
+    }
+
+    return currentBucket;
+  }
+
+  private cleanupExpiredBuckets(buckets: Map<string, AiRateLimitBucket>, now: number): void {
+    if (buckets.size <= 256) {
+      return;
+    }
+
+    for (const [key, bucket] of buckets.entries()) {
+      if (now >= bucket.resetAt) {
+        buckets.delete(key);
+      }
+    }
+  }
+
+  private normalizeIp(clientIp: string | null): string | null {
+    if (!clientIp) {
+      return null;
+    }
+
+    const normalizedIp = clientIp.trim();
+    return normalizedIp.length > 0 ? normalizedIp : null;
+  }
+
+  private readPositiveInt(key: string, fallbackValue: number): number {
+    const rawValue = this.configService.get<string | number | undefined>(key);
+    const parsedValue =
+      typeof rawValue === "number" ? rawValue : Number.parseInt(String(rawValue ?? ""), 10);
+
+    if (!Number.isFinite(parsedValue) || parsedValue <= 0) {
+      return fallbackValue;
+    }
+
+    return parsedValue;
+  }
+}

apps/api/src/ai/ai.controller.ts

@@ -1,4 +1,13 @@
-import { Body, Controller, Get, Headers, Post, Query, UnauthorizedException } from "@nestjs/common";
+import {
+  Body,
+  Controller,
+  Get,
+  Headers,
+  Ip,
+  Post,
+  Query,
+  UnauthorizedException
+} from "@nestjs/common";
 import { AiChatDto } from "./dto/ai-chat.dto";
 import { ListAiUsageLogsQueryDto } from "./dto/list-ai-usage-logs-query.dto";
 import { UpsertAiProviderBindingDto } from "./dto/upsert-ai-provider-binding.dto";
@@ -39,9 +48,10 @@ export class AiController {
   @Post("chat")
   async chat(
     @Headers("x-user-id") userIdHeader: string | string[] | undefined,
+    @Ip() clientIp: string,
     @Body() body: AiChatDto
   ): Promise<AiChatResponse> {
-    return this.aiService.chat(this.resolveUserId(userIdHeader), body);
+    return this.aiService.chat(this.resolveUserId(userIdHeader), body, clientIp);
   }
 
   private resolveUserId(userIdHeader: string | string[] | undefined): string {

apps/api/src/ai/ai.module.ts

@@ -1,5 +1,6 @@
 import { Module } from "@nestjs/common";
 import { PrismaModule } from "../prisma/prisma.module";
+import { AiRateLimitService } from "./ai-rate-limit.service";
 import { AiController } from "./ai.controller";
 import { AiProviderRegistryService } from "./ai-provider-registry.service";
 import { AiService } from "./ai.service";
@@ -9,6 +10,12 @@ import { OpenAiCompatibleProvider } from "./providers/openai-compatible.provider
 @Module({
   imports: [PrismaModule],
   controllers: [AiController],
-  providers: [AiService, AiProviderRegistryService, OpenAiCompatibleProvider, AstrbotProvider]
+  providers: [
+    AiService,
+    AiRateLimitService,
+    AiProviderRegistryService,
+    OpenAiCompatibleProvider,
+    AstrbotProvider
+  ]
 })
 export class AiModule {}

apps/api/src/ai/ai.service.ts

@@ -1,4 +1,11 @@
-import { BadGatewayException, BadRequestException, Injectable, Logger } from "@nestjs/common";
+import {
+  BadGatewayException,
+  BadRequestException,
+  HttpException,
+  HttpStatus,
+  Injectable,
+  Logger
+} from "@nestjs/common";
 import {
   AiChannel,
   AiUsageLog,
@@ -10,6 +17,7 @@ import {
 } from "../../generated/prisma/client";
 import { PrismaService } from "../prisma/prisma.service";
 import { DataEncryptionService } from "../security/data-encryption.service";
+import { AiRateLimitService } from "./ai-rate-limit.service";
 import { AiProviderRegistryService } from "./ai-provider-registry.service";
 import { AiChatDto } from "./dto/ai-chat.dto";
 import { ListAiUsageLogsQueryDto } from "./dto/list-ai-usage-logs-query.dto";
@@ -105,7 +113,8 @@ export class AiService {
   constructor(
     private readonly prismaService: PrismaService,
     private readonly aiProviderRegistryService: AiProviderRegistryService,
-    private readonly dataEncryptionService: DataEncryptionService
+    private readonly dataEncryptionService: DataEncryptionService,
+    private readonly aiRateLimitService: AiRateLimitService
   ) {}
 
   async listBindings(userId: string): Promise<ListAiBindingsResponse> {
@@ -242,7 +251,26 @@ export class AiService {
     return this.serializeBinding(result);
   }
 
-  async chat(userId: string, dto: AiChatDto): Promise<AiChatResponse> {
+  async chat(
+    userId: string,
+    dto: AiChatDto,
+    clientIp: string | null = null
+  ): Promise<AiChatResponse> {
+    const rateLimitResult = this.aiRateLimitService.consume(userId, clientIp);
+    if (!rateLimitResult.allowed) {
+      throw new HttpException(
+        {
+          message: "AI 请求过于频繁，请稍后再试",
+          code: "AI_RATE_LIMITED",
+          dimension: rateLimitResult.reason === "USER" ? "user" : "ip",
+          retryAfterMs: rateLimitResult.retryAfterMs,
+          limit: rateLimitResult.limit,
+          windowMs: rateLimitResult.windowMs
+        },
+        HttpStatus.TOO_MANY_REQUESTS
+      );
+    }
+
     const attempts: AiRouteAttempt[] = [];
     const plan = await this.buildRoutePlan(userId, dto.channel ?? null);
     const promptMessage = await this.buildPromptMessage(userId, dto.message, dto.localTasks ?? []);