ai-rate-limit.ts

ai-rate-limit.tsapps/api/src/modules/ai/ai-rate-limit.tsraw ← ai2,824 B← Back to the app
1/**
2 * TASK-0123 — the SHARED `ai.calls` per-user rate limiter (spec §Limits:
3 * "30 writing calls / min per user — shares the AI cap from TASK-0122 —
4 * implemented as one shared `ai.calls` limiter").
5 *
6 * Both Ask AI chat (TASK-0122) and the writing tools count against the SAME
7 * Redis sliding window (`ai:calls:rl:<userId>`, 30 calls / 60 s) so a user
8 * cannot bypass the cap by mixing features. Policy is unchanged from the
9 * chat limiter:
10 *
11 *   - sliding zset window (zremrangebyscore → zadd → zcard → pexpire);
12 *   - degrades OPEN when Redis is unavailable (availability over
13 *     strictness — the form limiter's established policy);
14 *   - the counter is INCREMENTED on the checked call (the call itself counts
15 *     against the window, so the 31st within 60 s is the one that 429s).
16 */
17import { Inject, Injectable } from "@nestjs/common";
18 
19import { RedisService } from "../../common/redis.service.js";
20import { logger } from "../../common/logger.js";
21 
22/** 30 AI calls / min per user (chat + writing combined; 429 RATE_LIMITED). */
23export const AI_CALLS_PER_USER_LIMIT = 30;
24export const AI_CALLS_RL_WINDOW_MS = 60_000;
25 
26export interface AiRateLimitResult {
27  allowed: boolean;
28  remaining: number;
29  retryAfterSeconds: number;
30}
31 
32@Injectable()
33export class AiRateLimitService {
34  constructor(@Inject(RedisService) private readonly redis: RedisService) {}
35 
36  private lastWarn = 0;
37 
38  /**
39   * Count one AI call (chat or writing) for the user; returns the window
40   * state. Call BEFORE doing the work (the increment is the reservation).
41   */
42  async check(userId: string, limit: number = AI_CALLS_PER_USER_LIMIT): Promise<AiRateLimitResult> {
43    const key = `ai:calls:rl:${userId}`;
44    try {
45      const now = Date.now();
46      const pipe = this.redis.getRedis().pipeline();
47      pipe.zremrangebyscore(key, 0, now - AI_CALLS_RL_WINDOW_MS);
48      pipe.zadd(key, now, `${now}:${Math.random().toString(36).slice(2)}`);
49      pipe.zcard(key);
50      pipe.pexpire(key, AI_CALLS_RL_WINDOW_MS);
51      const results = await pipe.exec();
52      const card = results?.[2]?.[1];
53      const count = typeof card === "number" ? card : 1;
54      const remaining = Math.max(0, limit - count);
55      return {
56        allowed: count <= limit,
57        remaining,
58        retryAfterSeconds: count <= limit ? 0 : Math.ceil(AI_CALLS_RL_WINDOW_MS / 1000),
59      };
60    } catch (err) {
61      if (Date.now() - this.lastWarn > 60_000) {
62        this.lastWarn = Date.now();
63        logger.warn(`ai.calls rate limiter degraded (redis unavailable): ${(err as Error).message}`);
64      }
65      return { allowed: true, remaining: limit, retryAfterSeconds: 0 };
66    }
67  }
68}
69 
70/** The Redis key shape (integration tests + docs). */
71export function aiCallsRlKey(userId: string): string {
72  return `ai:calls:rl:${userId}`;
73}
74