diff --git a/server/lib/aiBudgetEnforcement.ts b/server/lib/aiBudgetEnforcement.ts index 8f70340fa..2a516695d 100644 --- a/server/lib/aiBudgetEnforcement.ts +++ b/server/lib/aiBudgetEnforcement.ts @@ -10,6 +10,7 @@ import { } from "@server/db"; import { modelKeyMatches } from "@server/lib/aiModelKeyMatch"; import type { AiUsage } from "@server/lib/aiUsageExtraction"; +import { regionalCache as cache } from "#dynamic/lib/cache"; import logger from "@server/logger"; type BudgetPeriod = AiBudget["period"]; @@ -22,6 +23,39 @@ const PERIOD_DURATIONS_MS: Record, number> = { yearly: 365 * 24 * 60 * 60 * 1000 }; +// Budgets are cheap to be a little stale about (enforcement is already +// check-then-act, not transactional). Re-derive each budget's usage sum +// from aiUsageRecords at most this often; in between, completed requests +// just add their own contribution onto the cached sum instead of +// re-querying/re-aggregating from scratch. +const BUDGET_CACHE_REFRESH_MS = 8_000; +// Redis-level TTL is only a safety net for eviction if a budget stops +// seeing traffic - the actual staleness check is the computedAt timestamp +// stored in the cached value, compared against BUDGET_CACHE_REFRESH_MS. +const BUDGET_CACHE_SAFETY_TTL_SEC = 60; + +function applicableBudgetsCacheKey(ctx: BudgetScopeContext): string { + const roleKey = [...ctx.roleIds].sort((a, b) => a - b).join(","); + return [ + "aiBudget:applicable", + ctx.orgId, + ctx.providerId, + ctx.requestedModel, + ctx.resourceId ?? "", + ctx.siteResourceId ?? "", + roleKey + ].join(":"); +} + +function budgetUsageCacheKey(budgetId: number): string { + return `aiBudget:usage:${budgetId}`; +} + +type CachedBudgetUsage = { + sum: number; + computedAt: number; +}; + // Budget periods are trailing windows from "now", not calendar-aligned // (e.g. "daily" = last 24h). "lifetime" has no lower bound. function windowStart(period: BudgetPeriod, now: number): number { @@ -45,10 +79,26 @@ export type BudgetScopeContext = { * Every budget that could apply to this request: the provider itself, any * model on that provider whose (possibly wildcarded) modelKey matches the * requested model, the target resource/site-resource, and any role the - * requesting user holds in the org. + * requesting user holds in the org. Cached for BUDGET_CACHE_REFRESH_MS since + * budget/model config changes are rare and a request-scoped org/provider/ + * model/resource/role combination repeats constantly under real traffic. */ export async function resolveApplicableBudgets( ctx: BudgetScopeContext +): Promise { + const cacheKey = applicableBudgetsCacheKey(ctx); + const cached = await cache.get(cacheKey); + if (cached !== undefined) { + return cached; + } + + const budgets = await fetchApplicableBudgets(ctx); + await cache.set(cacheKey, budgets, BUDGET_CACHE_REFRESH_MS / 1000); + return budgets; +} + +async function fetchApplicableBudgets( + ctx: BudgetScopeContext ): Promise { const providerModels = await db .select({ modelId: aiModels.modelId, modelKey: aiModels.modelKey }) @@ -217,6 +267,74 @@ export async function sumUsageForBudget( return 0; } +/** + * Cached wrapper around sumUsageForBudget. Reuses a per-budget cached sum + * for up to BUDGET_CACHE_REFRESH_MS, and otherwise falls through to the DB + * aggregation and reseeds the cache. Completed requests within that window + * top the cached sum up via applyUsageToBudgetCache below rather than + * forcing a re-aggregation on every request. + */ +async function getBudgetUsage( + budget: AiBudget, + ctx: BudgetScopeContext, + now: number +): Promise { + const cacheKey = budgetUsageCacheKey(budget.budgetId); + const cached = await cache.get(cacheKey); + if (cached && now - cached.computedAt < BUDGET_CACHE_REFRESH_MS) { + return cached.sum; + } + + const sum = await sumUsageForBudget(budget, ctx, now); + await cache.set( + cacheKey, + { sum, computedAt: now } satisfies CachedBudgetUsage, + BUDGET_CACHE_SAFETY_TTL_SEC + ); + return sum; +} + +/** + * Called once a request's actual usage is known, for every budget that was + * resolved as applicable to it (i.e. checkBudgets' returned `budgets`). + * Adds this request's contribution directly onto each budget's cached sum + * so the next request in the same refresh window doesn't need to re-query + * or re-aggregate. If there's no warm cache entry, or it's already due for + * a refresh, this is a no-op - the next reader re-derives from the DB, + * which by then already includes this request's row via recordUsage. + */ +export async function applyUsageToBudgetCache( + budgets: AiBudget[], + usage: { usd: number; tokens: number } +): Promise { + await Promise.all( + budgets.map(async (budget) => { + const delta = budget.unit === "usd" ? usage.usd : usage.tokens; + if (!delta) { + return; + } + + const cacheKey = budgetUsageCacheKey(budget.budgetId); + const cached = await cache.get(cacheKey); + if ( + !cached || + Date.now() - cached.computedAt >= BUDGET_CACHE_REFRESH_MS + ) { + return; + } + + await cache.set( + cacheKey, + { + sum: cached.sum + delta, + computedAt: cached.computedAt + } satisfies CachedBudgetUsage, + BUDGET_CACHE_SAFETY_TTL_SEC + ); + }) + ); +} + // Throttled to one durable event per budget per breach window, so a soft // budget being exceeded doesn't write a row on every subsequent request // while it stays over. @@ -265,6 +383,10 @@ async function recordBreachEventIfNew( export type BudgetCheckResult = { blocked: boolean; blockingBudget?: AiBudget; + // Every budget resolved as applicable to this request, regardless of + // whether it was breached - pass to applyUsageToBudgetCache once this + // request's actual usage is known. + budgets: AiBudget[]; }; export async function checkBudgets( @@ -272,14 +394,14 @@ export async function checkBudgets( ): Promise { const budgets = await resolveApplicableBudgets(ctx); if (budgets.length === 0) { - return { blocked: false }; + return { blocked: false, budgets: [] }; } const now = Date.now(); let blockingBudget: AiBudget | undefined; for (const budget of budgets) { - const usage = await sumUsageForBudget(budget, ctx, now); + const usage = await getBudgetUsage(budget, ctx, now); if (usage < budget.amount) { continue; } @@ -292,8 +414,8 @@ export async function checkBudgets( } return blockingBudget - ? { blocked: true, blockingBudget } - : { blocked: false }; + ? { blocked: true, blockingBudget, budgets } + : { blocked: false, budgets }; } export type UsageRecordInput = { diff --git a/server/lib/cache.ts b/server/lib/cache.ts index 498d0dd3f..0113099e1 100644 --- a/server/lib/cache.ts +++ b/server/lib/cache.ts @@ -168,5 +168,5 @@ class AdaptiveCache { // Export singleton instance export const cache = new AdaptiveCache(); -export const regionalCache = cache; // Alias for compatability with the private version +export const regionalCache = cache; // Alias for compatibility with the private version export default cache; diff --git a/server/routers/aiGateway/pipeline.ts b/server/routers/aiGateway/pipeline.ts index 802f5f453..22274e049 100644 --- a/server/routers/aiGateway/pipeline.ts +++ b/server/routers/aiGateway/pipeline.ts @@ -1,6 +1,7 @@ import { Request, Response } from "express"; import { and, eq, inArray } from "drizzle-orm"; import { + AiBudget, AiProvider, aiModels, aiProviders, @@ -51,7 +52,11 @@ import { } from "@server/lib/aiModelKeyMatch"; import { aiGatewayUpstreamFetch } from "@server/lib/aiGatewayUpstreamFetch"; import { getModelPricing, calculateAiCost } from "@server/lib/aiModelPricing"; -import { checkBudgets, recordUsage } from "@server/lib/aiBudgetEnforcement"; +import { + applyUsageToBudgetCache, + checkBudgets, + recordUsage +} from "@server/lib/aiBudgetEnforcement"; import { extractUsage, estimateUsage, @@ -533,6 +538,7 @@ function logAiUsageAndCost(args: { resourceId: number | null; siteResourceId: number | null; requestUserId: string | null; + budgets: AiBudget[]; }): void { const { capability, @@ -545,7 +551,8 @@ function logAiUsageAndCost(args: { orgId, resourceId, siteResourceId, - requestUserId + requestUserId, + budgets } = args; let usage: AiUsage | null = extractUsage( @@ -588,6 +595,18 @@ function logAiUsageAndCost(args: { usage, costUsd: cost?.totalCost ?? null }); + + if (budgets.length > 0) { + void applyUsageToBudgetCache(budgets, { + usd: cost?.totalCost ?? 0, + tokens: + usage.promptTokens + + usage.cacheReadTokens + + usage.cacheWriteTokens + + usage.completionTokens + + usage.reasoningTokens + }); + } } } @@ -666,6 +685,7 @@ export async function handleAiGatewayProxy( const { provider } = selection; + let appliedBudgets: AiBudget[] = []; if (orgId) { const budgetCheck = await checkBudgets({ orgId, @@ -676,6 +696,7 @@ export async function handleAiGatewayProxy( roleIds: requestUser?.roleIds ?? [], requestUserId: requestUser?.userId ?? null }); + appliedBudgets = budgetCheck.budgets; if (budgetCheck.blocked) { logger.warn("AI gateway request blocked by budget", { @@ -876,7 +897,8 @@ export async function handleAiGatewayProxy( orgId, resourceId, siteResourceId, - requestUserId: requestUser?.userId ?? null + requestUserId: requestUser?.userId ?? null, + budgets: appliedBudgets }); } return; @@ -895,7 +917,8 @@ export async function handleAiGatewayProxy( orgId, resourceId, siteResourceId, - requestUserId: requestUser?.userId ?? null + requestUserId: requestUser?.userId ?? null, + budgets: appliedBudgets }); return res.send(text); } catch (error) {