1
26
27import { ApiError, ErrorCodes } from "@factory/shared";
28
29import { DEFAULT_BUDGETS } from "./config.js";
30
31
32export const AiErrorCodes = {
33 AI_BUDGET_EXCEEDED: "AI_BUDGET_EXCEEDED",
34 SERVICE_UNAVAILABLE: "SERVICE_UNAVAILABLE",
35} as const;
36
37export type AiErrorCode = (typeof AiErrorCodes)[keyof typeof AiErrorCodes];
38
39export interface UsageRecord {
40 id: string;
41 workspaceId: string;
42 userId?: string;
43 capability?: string;
44 kind: "chat" | "embed" | "generate";
45 model: string;
46 provider: string;
47 tokensIn: number;
48 tokensOut: number;
49 latencyMs: number;
50
51 costUsd: number;
52 requestKind: "stream" | "batch";
53 status: "ok" | "error" | "timeout";
54 createdAt: number;
55}
56
57export interface UsageSink {
58 (record: UsageRecord): void | Promise<void>;
59}
60
61let counter = 0;
62function nextId(): string {
63 counter += 1;
64 return `usage_${counter.toString(36)}_${Date.now().toString(36)}`;
65}
66
67
72export class UsageLedger {
73 readonly records: UsageRecord[] = [];
74 private readonly sink: UsageSink | undefined;
75 private readonly now: () => number;
76
77 constructor(options: { sink?: UsageSink; now?: () => number } = {}) {
78 this.sink = options.sink;
79 this.now = options.now ?? Date.now;
80 }
81
82 async record(input: Omit<UsageRecord, "id" | "createdAt"> & { id?: string }): Promise<UsageRecord> {
83 const record: UsageRecord = { id: input.id ?? nextId(), createdAt: this.now(), ...input };
84 this.records.push(record);
85 if (this.sink) await this.sink(record);
86 return record;
87 }
88
89
90 tokensSince(workspaceId: string, sinceMs: number): number {
91 let sum = 0;
92 for (const r of this.records) {
93 if (r.workspaceId === workspaceId && r.createdAt >= sinceMs) sum += r.tokensIn + r.tokensOut;
94 }
95 return sum;
96 }
97
98
99 costSince(workspaceId: string, sinceMs: number): number {
100 let sum = 0;
101 for (const r of this.records) {
102 if (r.workspaceId === workspaceId && r.createdAt >= sinceMs) sum += r.costUsd;
103 }
104 return sum;
105 }
106
107
108 callsSince(workspaceId: string, sinceMs: number): number {
109 let n = 0;
110 for (const r of this.records) {
111 if (r.workspaceId === workspaceId && r.createdAt >= sinceMs) n += 1;
112 }
113 return n;
114 }
115
116
117 monthlyRollup(workspaceId: string, monthStartMs: number): { tokens: number; costUsd: number; calls: number } {
118 return {
119 tokens: this.tokensSince(workspaceId, monthStartMs),
120 costUsd: this.costSince(workspaceId, monthStartMs),
121 calls: this.callsSince(workspaceId, monthStartMs),
122 };
123 }
124}
125
126
127export interface Budgets {
128
129 tokensPerMonth: number;
130
131 requestsPerMinute: number;
132
133 maxOutputTokens: number;
134
135 dailySpendUsd: number;
136}
137
138
139export function budgetExceededError(
140 reason: "monthly_token_budget_exceeded" | "daily_spend_budget_exceeded",
141 used: number,
142 limit: number,
143): ApiError {
144 return new ApiError(ErrorCodes.AI_ERROR, {
145 status: 402,
146 message:
147 reason === "monthly_token_budget_exceeded"
148 ? "Monthly AI token budget reached. Ask an admin."
149 : "Daily AI spend budget reached. Ask an admin.",
150 details: { code: AiErrorCodes.AI_BUDGET_EXCEEDED, reason, used, limit },
151 });
152}
153
154
160export class BudgetService {
161 private readonly ledger: UsageLedger;
162 private readonly now: () => number;
163
164 private readonly windows = new Map<string, number[]>();
165
166 private readonly lastNotification = new Map<string, number>();
167
168 constructor(ledger: UsageLedger, now?: () => number) {
169 this.ledger = ledger;
170 this.now = now ?? Date.now;
171 }
172
173
174 private monthStart(): number {
175 const d = new Date(this.now());
176 return new Date(d.getFullYear(), d.getMonth(), 1).getTime();
177 }
178
179
180 private dayStart(): number {
181 const d = new Date(this.now());
182 return new Date(d.getFullYear(), d.getMonth(), d.getDate()).getTime();
183 }
184
185
190 check(workspaceId: string, budgets: Budgets, estimatedTokensOut = 0): void {
191 const b = { ...DEFAULT_BUDGETS, ...budgets };
192
193
194 const nowMs = this.now();
195 const windowMs = 60_000;
196 let stamps = this.windows.get(workspaceId) ?? [];
197 stamps = stamps.filter((t) => nowMs - t < windowMs);
198 if (stamps.length >= b.requestsPerMinute) {
199 this.windows.set(workspaceId, stamps);
200 throw new ApiError(ErrorCodes.RATE_LIMITED, {
201 message: "AI rate limit reached. Try again shortly.",
202 details: { code: "RATE_LIMITED", limit: b.requestsPerMinute, windowSeconds: 60 },
203 });
204 }
205 stamps.push(nowMs);
206 this.windows.set(workspaceId, stamps);
207
208
209 if (b.tokensPerMonth > 0) {
210 const used = this.ledger.tokensSince(workspaceId, this.monthStart());
211 if (used + estimatedTokensOut >= b.tokensPerMonth) {
212 this.notifyBudgetExceeded(workspaceId);
213 throw budgetExceededError("monthly_token_budget_exceeded", used, b.tokensPerMonth);
214 }
215 }
216
217
218 if (b.dailySpendUsd > 0) {
219 const used = this.ledger.costSince(workspaceId, this.dayStart());
220 if (used >= b.dailySpendUsd) {
221 this.notifyBudgetExceeded(workspaceId);
222 throw budgetExceededError("daily_spend_budget_exceeded", used, b.dailySpendUsd);
223 }
224 }
225 }
226
227
228 notifyBudgetExceeded(workspaceId: string, emit?: (workspaceId: string) => void): boolean {
229 const nowMs = this.now();
230 const last = this.lastNotification.get(workspaceId) ?? 0;
231 if (nowMs - last < 24 * 3600 * 1000) return false;
232 this.lastNotification.set(workspaceId, nowMs);
233 emit?.(workspaceId);
234 return true;
235 }
236
237
238 clampMaxTokens(requested: number | undefined, budgets: Budgets): number {
239 const cap = { ...DEFAULT_BUDGETS, ...budgets }.maxOutputTokens;
240 if (requested === undefined) return cap;
241 return Math.min(requested, cap);
242 }
243}
244
