ai-usage-budget.test.ts

ai-usage-budget.test.tsapps/worker/src/__tests__/ai-usage-budget.test.tsraw ← __tests__16,027 B← Back to the app
1/**
2 * TASK-0121 — usage.rollup + budget.notify unit tests (Docker-free, fake
3 * prisma/redis — the `as unknown as PrismaClient` convention of
4 * ai-embedding.test.ts).
5 *
6 * Covers the DoD acceptance:
7 *   - the rollup sums ONLY the current month (rows before month-start are
8 *     excluded — asserted on the createdAt filter + the sum)
9 *   - the rollup sums ONLY the requested workspace's rows (scoped to its
10 *     own conversation ids; another workspace's ledger is untouched)
11 *   - the budget sweep enqueues budget.notify only when crossed AND the 24 h
12 *     slot is free (slot set → no enqueue)
13 *   - budget.notify: claims the slot (NX), notifies every OWNER/ADMIN,
14 *     publishes notification.created per admin, writes the audit row — and
15 *     SKIPS when the slot is already held (the 24 h throttle)
16 */
17import { describe, expect, it, vi } from "vitest";
18import type { PrismaClient } from "@prisma/client";
19import type { Redis } from "ioredis";
20import {
21  aiJobSchema,
22  aiMonthlyBudget,
23  budgetNotifySlot,
24  runBudgetNotify,
25  runUsageRollup
26} from "../processors/ai.processor.js";
27 
28// ---------------------------------------------------------------------------
29// Fakes
30// ---------------------------------------------------------------------------
31 
32interface MessageRow {
33  conversationId: string;
34  createdAt: Date;
35  tokensIn: number | null;
36  tokensOut: number | number;
37}
38 
39function fakePrisma(opts: {
40  workspaces?: { id: string; settings: unknown; usageLimits: unknown }[];
41  conversations?: { id: string; workspaceId: string }[];
42  messages?: MessageRow[];
43  members?: { userId: string; workspaceId: string; role: string }[];
44}) {
45  const calls = {
46    aggregate: vi.fn(),
47    findManyWs: vi.fn(),
48    findManyConv: vi.fn(),
49    findManyMembers: vi.fn(),
50    notificationCreate: vi.fn(),
51    auditCreate: vi.fn()
52  };
53 
54  calls.findManyWs.mockImplementation(async (args: { where: { id?: string } }) => {
55    const all = opts.workspaces ?? [];
56    if (args.where.id) return all.filter((w) => w.id === args.where.id);
57    return all;
58  });
59  calls.findManyConv.mockImplementation(async (args: { where: { workspaceId: string } }) =>
60    (opts.conversations ?? []).filter((c) => c.workspaceId === args.where.workspaceId).map((c) => ({ id: c.id }))
61  );
62  calls.aggregate.mockImplementation(async (args: {
63    where: { conversationId: { in: string[] }; createdAt: { gte: Date } };
64  }) => {
65    const inConv = new Set(args.where.conversationId.in);
66    const rows = (opts.messages ?? []).filter(
67      (m) => inConv.has(m.conversationId) && m.createdAt >= args.where.createdAt.gte
68    );
69    const ti = rows.reduce((s, m) => s + (m.tokensIn ?? 0), 0);
70    const to = rows.reduce((s, m) => s + (m.tokensOut ?? 0), 0);
71    return { _sum: { tokensIn: ti, tokensOut: to } };
72  });
73  // Mirrors the DB: both the workspace AND the `role: {in}` where clause
74  // apply (the real Prisma filters in Postgres; the fake must too).
75  calls.findManyMembers.mockImplementation(async (args: {
76    where: { workspaceId: string; role: { in: string[] } };
77  }) =>
78    (opts.members ?? [])
79      .filter((m) => m.workspaceId === args.where.workspaceId && args.where.role.in.includes(m.role))
80      .map((m) => ({ userId: m.userId }))
81  );
82  let notifSeq = 0;
83  calls.notificationCreate.mockImplementation(async () => ({ id: `ntf-${++notifSeq}` }));
84  calls.auditCreate.mockImplementation(async () => ({ id: "aud-1" }));
85 
86  const prisma = {
87    workspace: { findMany: calls.findManyWs },
88    aiConversation: { findMany: calls.findManyConv },
89    aiMessage: { aggregate: calls.aggregate },
90    workspaceMember: { findMany: calls.findManyMembers },
91    notification: { create: calls.notificationCreate },
92    auditLog: { create: calls.auditCreate }
93  } as unknown as PrismaClient;
94 
95  return { prisma, calls };
96}
97 
98function fakeRedis(opts: { slotValue?: string | null; setReturns?: string | "OK" } = {}) {
99  const calls = {
100    get: vi.fn(async (key: string) => (key === budgetNotifySlot("ws-1") ? (opts.slotValue ?? null) : null)),
101    set: vi.fn(async () => opts.setReturns ?? "OK"),
102    // xadd(key, id, "event", <name>, "payload", <json>, "ts", <iso>) — the
103    // real ioredis signature is variadic, so the mock declares it as such
104    // (keeps the call-shape assertions below typeable).
105    xadd: vi.fn(async (..._args: string[]) => "1234-0")
106  };
107  return { redis: calls as unknown as Redis, calls };
108}
109 
110function monthStart(): Date {
111  const d = new Date();
112  d.setUTCDate(1);
113  d.setUTCHours(0, 0, 0, 0);
114  return d;
115}
116 
117// ---------------------------------------------------------------------------
118// payload validation
119// ---------------------------------------------------------------------------
120 
121describe("aiJobSchema (TASK-0121 job types)", () => {
122  it("accepts usage.rollup (workspaceId optional) and budget.notify", () => {
123    expect(aiJobSchema.parse({ type: "usage.rollup" }).type).toBe("usage.rollup");
124    // (the parse() result is the job union — access the variant field via the
125    //  narrow cast now that the union has more members than TASK-0121 added)
126    expect((aiJobSchema.parse({ type: "usage.rollup", workspaceId: "ws-1" }) as { workspaceId?: string }).workspaceId).toBe("ws-1");
127    expect(
128      aiJobSchema.parse({ type: "budget.notify", workspaceId: "ws-1", used: 10, limit: 100 }).type
129    ).toBe("budget.notify");
130  });
131 
132  it("rejects a budget.notify without the numeric budget fields", () => {
133    expect(() => aiJobSchema.parse({ type: "budget.notify", workspaceId: "ws-1" })).toThrow();
134    expect(() => aiJobSchema.parse({ type: "budget.notify", workspaceId: "ws-1", used: -1, limit: 10 })).toThrow();
135    expect(() => aiJobSchema.parse({ type: "budget.notify", workspaceId: "ws-1", used: 10, limit: 0 })).toThrow();
136  });
137});
138 
139// ---------------------------------------------------------------------------
140// aiMonthlyBudget
141// ---------------------------------------------------------------------------
142 
143describe("aiMonthlyBudget (admin override wins over settings)", () => {
144  it("uses max(usageLimits, settings.budgets.tokensPerMonth)", () => {
145    expect(aiMonthlyBudget({ usageLimits: { aiMonthlyTokens: 500 }, settings: { ai: { budgets: { tokensPerMonth: 900 } } } })).toBe(900);
146    expect(aiMonthlyBudget({ usageLimits: { aiMonthlyTokens: 1500 }, settings: {} })).toBe(1500);
147  });
148 
149  it("unset or 0 = unlimited (0)", () => {
150    expect(aiMonthlyBudget({ usageLimits: {}, settings: {} })).toBe(0);
151    expect(aiMonthlyBudget({ usageLimits: { aiMonthlyTokens: 0 }, settings: {} })).toBe(0);
152    expect(aiMonthlyBudget({ usageLimits: "not-json", settings: null as unknown })).toBe(0);
153  });
154});
155 
156// ---------------------------------------------------------------------------
157// runUsageRollup
158// ---------------------------------------------------------------------------
159 
160describe("runUsageRollup (TASK-0121 nightly reconcile)", () => {
161  it("sums ONLY the current month and ONLY the workspace's own rows", async () => {
162    const ms = monthStart();
163    const before = new Date(ms.getTime() - 60_000); // last month
164    const inMonth = new Date(ms.getTime() + 60_000);
165    const { prisma, calls } = fakePrisma({
166      workspaces: [{ id: "ws-1", usageLimits: { aiMonthlyTokens: 10_000 }, settings: {} }],
167      conversations: [
168        { id: "conv-a", workspaceId: "ws-1" },
169        { id: "conv-b", workspaceId: "ws-1" }
170      ],
171      messages: [
172        { conversationId: "conv-a", createdAt: before, tokensIn: 500, tokensOut: 500 }, // previous month — excluded
173        { conversationId: "conv-a", createdAt: inMonth, tokensIn: 100, tokensOut: 50 },
174        { conversationId: "conv-b", createdAt: inMonth, tokensIn: null, tokensOut: 150 }
175      ]
176    });
177    const { redis } = fakeRedis();
178    const enqueueBudgetNotify = vi.fn(async () => ({}));
179 
180    const res = await runUsageRollup(
181      { prisma, redis, enqueueBudgetNotify },
182      { type: "usage.rollup", workspaceId: "ws-1" },
183      "job-1"
184    );
185 
186    expect(res).toMatchObject({ reconciled: 1, budgetsExceeded: 0, enqueued: 0 });
187    // The aggregate was scoped to ws-1's conversation ids AND month start.
188    const aggArg = calls.aggregate.mock.calls[0]![0] as {
189      where: { conversationId: { in: string[] }; createdAt: { gte: Date } };
190    };
191    expect(aggArg.where.conversationId.in.sort()).toEqual(["conv-a", "conv-b"]);
192    expect(aggArg.where.createdAt.gte.getTime()).toBe(ms.getTime());
193    // used = 100+50+150 = 300 (the 1000 previous-month tokens are excluded)
194    expect(enqueueBudgetNotify).not.toHaveBeenCalled();
195  });
196 
197  it("scopes to ONE workspace: another workspace's ledger never enters the sum", async () => {
198    const inMonth = new Date();
199    const { prisma } = fakePrisma({
200      // only ws-1 is in scope (single-workspace run); ws-2 has a full ledger
201      workspaces: [{ id: "ws-1", usageLimits: {}, settings: {} }],
202      conversations: [{ id: "conv-1", workspaceId: "ws-1" }],
203      messages: [
204        { conversationId: "conv-1", createdAt: inMonth, tokensIn: 10, tokensOut: 10 }
205      ]
206    });
207    const { redis } = fakeRedis();
208    const res = await runUsageRollup(
209      { prisma, redis, enqueueBudgetNotify: vi.fn(async () => ({})) },
210      { type: "usage.rollup", workspaceId: "ws-1" },
211      null
212    );
213    expect(res.reconciled).toBe(1);
214  });
215 
216  it("enqueues budget.notify when the budget was crossed and the slot is free", async () => {
217    const inMonth = new Date();
218    const { prisma } = fakePrisma({
219      workspaces: [{ id: "ws-1", usageLimits: { aiMonthlyTokens: 100 }, settings: {} }],
220      conversations: [{ id: "conv-1", workspaceId: "ws-1" }],
221      messages: [{ conversationId: "conv-1", createdAt: inMonth, tokensIn: 60, tokensOut: 60 }]
222    });
223    const { redis } = fakeRedis({ slotValue: null });
224    const enqueueBudgetNotify = vi.fn(async () => ({}));
225 
226    const res = await runUsageRollup(
227      { prisma, redis, enqueueBudgetNotify },
228      { type: "usage.rollup", workspaceId: "ws-1" },
229      null
230    );
231 
232    expect(res).toMatchObject({ reconciled: 1, budgetsExceeded: 1, enqueued: 1 });
233    expect(enqueueBudgetNotify).toHaveBeenCalledWith({
234      type: "budget.notify",
235      workspaceId: "ws-1",
236      used: 120,
237      limit: 100
238    });
239  });
240 
241  it("does NOT re-enqueue when the 24 h slot is already set", async () => {
242    const inMonth = new Date();
243    const { prisma } = fakePrisma({
244      workspaces: [{ id: "ws-1", usageLimits: { aiMonthlyTokens: 100 }, settings: {} }],
245      conversations: [{ id: "conv-1", workspaceId: "ws-1" }],
246      messages: [{ conversationId: "conv-1", createdAt: inMonth, tokensIn: 60, tokensOut: 60 }]
247    });
248    const { redis } = fakeRedis({ slotValue: "1" });
249    const enqueueBudgetNotify = vi.fn(async () => ({}));
250 
251    const res = await runUsageRollup(
252      { prisma, redis, enqueueBudgetNotify },
253      { type: "usage.rollup", workspaceId: "ws-1" },
254      null
255    );
256 
257    expect(res).toMatchObject({ budgetsExceeded: 1, enqueued: 0 });
258    expect(enqueueBudgetNotify).not.toHaveBeenCalled();
259  });
260 
261  it("unlimited budget (0) never counts as exceeded", async () => {
262    const inMonth = new Date();
263    const { prisma } = fakePrisma({
264      workspaces: [{ id: "ws-1", usageLimits: {}, settings: {} }],
265      conversations: [{ id: "conv-1", workspaceId: "ws-1" }],
266      messages: [{ conversationId: "conv-1", createdAt: inMonth, tokensIn: 1e9, tokensOut: 1 }]
267    });
268    const { redis } = fakeRedis();
269    const enqueueBudgetNotify = vi.fn(async () => ({}));
270    const res = await runUsageRollup(
271      { prisma, redis, enqueueBudgetNotify },
272      { type: "usage.rollup", workspaceId: "ws-1" },
273      null
274    );
275    expect(res.budgetsExceeded).toBe(0);
276    expect(enqueueBudgetNotify).not.toHaveBeenCalled();
277  });
278});
279 
280// ---------------------------------------------------------------------------
281// runBudgetNotify
282// ---------------------------------------------------------------------------
283 
284describe("runBudgetNotify (TASK-0121 admin notification)", () => {
285  const JOB = { type: "budget.notify" as const, workspaceId: "ws-1", used: 150, limit: 100 };
286 
287  it("notifies every OWNER/ADMIN + publishes notification.created + audits", async () => {
288    const { prisma, calls } = fakePrisma({
289      members: [
290        { userId: "owner-1", workspaceId: "ws-1", role: "OWNER" },
291        { userId: "admin-1", workspaceId: "ws-1", role: "ADMIN" },
292        { userId: "member-1", workspaceId: "ws-1", role: "MEMBER" } // not an admin
293      ]
294    });
295    const { redis, calls: redisCalls } = fakeRedis({ setReturns: "OK" });
296 
297    const res = await runBudgetNotify({ prisma, redis: redis as unknown as Redis }, JOB, "job-9");
298 
299    expect(res).toEqual({ notified: 2 });
300    // slot claimed with NX + 24 h TTL
301    expect(redisCalls.set).toHaveBeenCalledWith(budgetNotifySlot("ws-1"), "1", "EX", 24 * 3600, "NX");
302    expect(calls.notificationCreate).toHaveBeenCalledTimes(2);
303    const first = calls.notificationCreate.mock.calls[0]![0] as { data: Record<string, unknown> };
304    expect(first.data).toMatchObject({
305      workspaceId: "ws-1",
306      userId: "owner-1",
307      type: "AUTOMATION",
308      channel: "IN_APP",
309      title: "AI budget exceeded"
310    });
311    expect(String(first.data.body)).toContain("150 used of 100");
312    // one stream event per admin notification (the email channel's input)
313    expect(redisCalls.xadd).toHaveBeenCalledTimes(2);
314    // xadd(key, id, "event", <name>, "payload", <json>, "ts", <iso>)
315    const call = redisCalls.xadd.mock.calls[0]!;
316    expect(call[0]).toBe("factory:events");
317    expect(call[3]).toBe("notification.created");
318    expect(call[4]).toBe("payload");
319    const payload = JSON.parse(call[5] as string);
320    expect(payload).toMatchObject({ notificationId: "ntf-1", type: "AUTOMATION", workspaceId: "ws-1" });
321    expect(payload.deepLink).toEqual({ type: "workspace", id: "ws-1" });
322    // the audit row (system actor)
323    expect(calls.auditCreate).toHaveBeenCalledTimes(1);
324    const audit = calls.auditCreate.mock.calls[0]![0] as { data: Record<string, unknown> };
325    expect(audit.data).toMatchObject({
326      actorType: "SYSTEM",
327      action: "ai.budget_exceeded",
328      workspaceId: "ws-1"
329    });
330    expect(audit.data.metadata).toMatchObject({ used: 150, limit: 100, adminsNotified: 2 });
331  });
332 
333  it("skips when the slot is already held (24 h throttle — the API hook won the race)", async () => {
334    const { prisma, calls } = fakePrisma({ members: [{ userId: "owner-1", workspaceId: "ws-1", role: "OWNER" }] });
335    const { redis } = fakeRedis({ setReturns: "PENDING" }); // NX lost → not "OK"
336 
337    const res = await runBudgetNotify({ prisma, redis: redis as unknown as Redis }, JOB, null);
338 
339    expect(res).toEqual({ skipped: true, reason: "already notified" });
340    expect(calls.notificationCreate).not.toHaveBeenCalled();
341    expect(calls.auditCreate).not.toHaveBeenCalled();
342  });
343 
344  it("no OWNER/ADMIN → no rows, no crash (the slot stays set for 24 h)", async () => {
345    const { prisma, calls } = fakePrisma({ members: [{ userId: "member-1", workspaceId: "ws-1", role: "MEMBER" }] });
346    const { redis } = fakeRedis({ setReturns: "OK" });
347 
348    const res = await runBudgetNotify({ prisma, redis: redis as unknown as Redis }, JOB, null);
349 
350    expect(res).toEqual({ notified: 0 });
351    expect(calls.notificationCreate).not.toHaveBeenCalled();
352    expect(calls.auditCreate).not.toHaveBeenCalled();
353  });
354 
355  it("a redis slot failure fails OPEN (notifies) — a duplicate is cheaper than a miss", async () => {
356    const { prisma, calls } = fakePrisma({ members: [{ userId: "owner-1", workspaceId: "ws-1", role: "OWNER" }] });
357    const redisCalls = {
358      get: vi.fn(async () => null),
359      set: vi.fn(async () => {
360        throw new Error("redis down");
361      }),
362      xadd: vi.fn(async () => "1-0")
363    };
364    const res = await runBudgetNotify({ prisma, redis: redisCalls as unknown as Redis }, JOB, null);
365    expect(res).toEqual({ notified: 1 });
366    expect(calls.notificationCreate).toHaveBeenCalledTimes(1);
367  });
368});
369