budget-guard.service.ts

budget-guard.service.tsapps/api/src/modules/ai/budget-guard.service.tsraw ← ai3,760 B← Back to the app
1/**
2 * AI module — budget guard service (Phase 12, TASK-0121).
3 *
4 * Wraps {@link createAiBudgetGuard} in an @Injectable so the gateway factory
5 * and the AI services share ONE guard instance (one rate-window state, one
6 * notification hook):
7 *
8 *   - `guard.check(workspaceId)` — awaited by EVERY AI endpoint before the
9 *     gateway call (429 rate window / 402 AI_BUDGET_EXCEEDED, Postgres-
10 *     backed month-to-date, replica-consistent via Redis);
11 *   - the exhaustion hook — in-app notification to every workspace
12 *     OWNER/ADMIN (the `notification.created` stream event carries the email
13 *     channel via the worker's email decision) + the
14 *     `ai.budget_exceeded` audit event (actor = system; abuse indicator);
15 *     throttled to one per workspace per 24 h (the Redis NX slot inside the
16 *     guard, NOT here — a duplicate call within the window is a no-op).
17 */
18import { Inject, Injectable, OnModuleInit } from "@nestjs/common";
19import type { Request } from "express";
20import type Redis from "ioredis";
21import type { PrismaClient } from "@prisma/client";
22import { PrismaService } from "../../common/prisma.service.js";
23import { RedisService } from "../../common/redis.service.js";
24import { AuditService } from "../../common/audit.service.js";
25import { logger } from "../../common/logger.js";
26import { NotificationsService } from "../notifications/notifications.service.js";
27import { createAiBudgetGuard, type AiBudgetGuard } from "./gateway-provider.js";
28 
29@Injectable()
30export class AiBudgetGuardService implements OnModuleInit {
31  private _guard: AiBudgetGuard | null = null;
32 
33  constructor(
34    @Inject(PrismaService) private readonly prisma: PrismaService,
35    @Inject(RedisService) private readonly redis: RedisService,
36    @Inject(AuditService) private readonly audit: AuditService,
37    @Inject(NotificationsService) private readonly notifications: NotificationsService
38  ) {}
39 
40  onModuleInit(): void {
41    this._guard = createAiBudgetGuard({
42      prisma: this.prisma as unknown as PrismaClient,
43      redis: this.redis.getRedis() as Redis,
44      onBudgetExceeded: (workspaceId, used, limit, req) => this.onBudgetExceeded(workspaceId, used, limit, req)
45    });
46  }
47 
48  /** The guard (available after module init — Nest boots the module before
49   *  routing, so request-path access is always safe). */
50  get guard(): AiBudgetGuard {
51    if (!this._guard) throw new Error("AiBudgetGuardService not initialized");
52    return this._guard;
53  }
54 
55  /** The exhaustion side effect (see module doc). */
56  async onBudgetExceeded(workspaceId: string, used: number, limit: number, req?: Request): Promise<void> {
57    const prisma = this.prisma as unknown as PrismaClient;
58    const admins = await prisma.workspaceMember.findMany({
59      where: { workspaceId, role: { in: ["OWNER", "ADMIN"] } },
60      select: { userId: true }
61    });
62 
63    const body =
64      `The workspace AI token budget was reached (${used} used of ${limit} monthly). ` +
65      `New AI calls are blocked until the limit is raised (workspace settings → usageLimits.aiMonthlyTokens).`;
66 
67    for (const admin of admins) {
68      try {
69        await this.notifications.notify({
70          workspaceId,
71          userId: admin.userId,
72          type: "AUTOMATION",
73          title: "AI budget exceeded",
74          body,
75          link: { type: "workspace", id: workspaceId }
76        });
77      } catch (err) {
78        logger.warn(`ai budget notification failed for user ${admin.userId}: ${(err as Error).message}`);
79      }
80    }
81 
82    await this.audit.audit(req, undefined, {
83      action: "ai.budget_exceeded",
84      resourceType: "workspace",
85      resourceId: workspaceId,
86      workspaceId,
87      metadata: { used, limit, adminsNotified: admins.length }
88    });
89  }
90}
91