diff --git a/server/cleanup.ts b/server/cleanup.ts index 10e9f4cc3..3e7a335b4 100644 --- a/server/cleanup.ts +++ b/server/cleanup.ts @@ -3,12 +3,14 @@ import { flushConnectionLogToDb } from "#dynamic/routers/newt"; import { flushSiteBandwidthToDb } from "@server/routers/gerbil/receiveBandwidth"; import { stopPingAccumulator } from "@server/routers/newt/pingAccumulator"; import { cleanup as wsCleanup } from "#dynamic/routers/ws"; +import { shutdownUsageRecorder } from "@server/lib/aiBudgetEnforcement"; async function cleanup() { await stopPingAccumulator(); await flushBandwidthToDb(); await flushConnectionLogToDb(); await flushSiteBandwidthToDb(); + await shutdownUsageRecorder(); await wsCleanup(); process.exit(0); diff --git a/server/lib/aiBudgetEnforcement.ts b/server/lib/aiBudgetEnforcement.ts index 2a516695d..4177af9b7 100644 --- a/server/lib/aiBudgetEnforcement.ts +++ b/server/lib/aiBudgetEnforcement.ts @@ -1,4 +1,4 @@ -import { and, eq, gte, inArray, isNull, or, sql, SQL } from "drizzle-orm"; +import { and, eq, gte, inArray, isNull, or, sql, SQL, type InferInsertModel } from "drizzle-orm"; import { AiBudget, aiBudgetBreachEvents, @@ -430,6 +430,91 @@ export type UsageRecordInput = { createdAt?: number; }; +type AiUsageRecordInsert = InferInsertModel; + +// In-memory buffer for batching AI usage record inserts, mirroring the +// approach in server/routers/badger/logRequestAudit.ts. Usage rows are read +// back on every budget-cache miss (see getBudgetUsage above), which happens +// at least every BUDGET_CACHE_REFRESH_MS, so this buffer is flushed much +// more aggressively than the request audit log to keep the table from +// lagging behind what budget enforcement needs. Unlike the audit log, there +// is no retention/cleanup job for this table - usage history is kept +// indefinitely for billing and historical reporting. +const usageRecordBuffer: AiUsageRecordInsert[] = []; + +const USAGE_BATCH_SIZE = 20; // Write to DB every 20 records +const USAGE_BATCH_INTERVAL_MS = 1000; // Or every 1 second, whichever comes first +const USAGE_MAX_BUFFER_SIZE = 5000; // Prevent unbounded memory growth +let usageFlushTimer: NodeJS.Timeout | null = null; +let isUsageFlushInProgress = false; + +async function flushUsageRecords() { + if (usageRecordBuffer.length === 0 || isUsageFlushInProgress) { + return; + } + + isUsageFlushInProgress = true; + + const recordsToWrite = usageRecordBuffer.splice(0, usageRecordBuffer.length); + + try { + // Use a transaction to ensure all inserts succeed or fail together + await db.transaction(async (tx) => { + // Batch insert in groups to avoid overwhelming the database + const DB_BATCH_SIZE = 25; + for (let i = 0; i < recordsToWrite.length; i += DB_BATCH_SIZE) { + const batch = recordsToWrite.slice(i, i + DB_BATCH_SIZE); + await tx.insert(aiUsageRecords).values(batch); + } + }); + logger.debug(`Flushed ${recordsToWrite.length} AI usage records to database`); + } catch (error) { + logger.error("Error flushing AI usage records:", error); + // On transaction error, put records back at the front of the buffer + // to retry, but only if the buffer isn't too large + if (usageRecordBuffer.length < USAGE_MAX_BUFFER_SIZE - recordsToWrite.length) { + usageRecordBuffer.unshift(...recordsToWrite); + logger.info(`Re-queued ${recordsToWrite.length} AI usage records for retry`); + } else { + logger.error(`Buffer full, dropped ${recordsToWrite.length} AI usage records`); + } + } finally { + isUsageFlushInProgress = false; + // If buffer filled up while we were flushing, flush again + if (usageRecordBuffer.length >= USAGE_BATCH_SIZE) { + flushUsageRecords().catch((err) => + logger.error("Error in follow-up AI usage flush:", err) + ); + } + } +} + +function scheduleUsageFlush() { + if (usageFlushTimer === null) { + usageFlushTimer = setTimeout(() => { + usageFlushTimer = null; + flushUsageRecords().catch((err) => + logger.error("Error in scheduled AI usage flush:", err) + ); + }, USAGE_BATCH_INTERVAL_MS); + } +} + +/** + * Gracefully flush all pending AI usage records (call this on shutdown). + */ +export async function shutdownUsageRecorder() { + if (usageFlushTimer) { + clearTimeout(usageFlushTimer); + usageFlushTimer = null; + } + // Force flush even if one is in progress by waiting and retrying + while (isUsageFlushInProgress) { + await new Promise((resolve) => setTimeout(resolve, 100)); + } + await flushUsageRecords(); +} + export async function recordUsage(input: UsageRecordInput): Promise { try { const { usage } = input; @@ -440,7 +525,15 @@ export async function recordUsage(input: UsageRecordInput): Promise { usage.completionTokens + usage.reasoningTokens; - await db.insert(aiUsageRecords).values({ + // Prevent unbounded buffer growth - drop oldest entries if buffer is too large + if (usageRecordBuffer.length >= USAGE_MAX_BUFFER_SIZE) { + const dropped = usageRecordBuffer.splice(0, USAGE_BATCH_SIZE); + logger.warn( + `AI usage record buffer exceeded max size (${USAGE_MAX_BUFFER_SIZE}), dropped ${dropped.length} oldest entries` + ); + } + + usageRecordBuffer.push({ orgId: input.orgId, providerId: input.providerId, resourceId: input.resourceId, @@ -457,6 +550,15 @@ export async function recordUsage(input: UsageRecordInput): Promise { estimated: usage.estimated, createdAt: input.createdAt ?? Date.now() }); + + // Flush immediately if buffer is full, otherwise schedule a flush + if (usageRecordBuffer.length >= USAGE_BATCH_SIZE) { + flushUsageRecords().catch((err) => + logger.error("Error flushing AI usage records:", err) + ); + } else { + scheduleUsageFlush(); + } } catch (error) { logger.error("Failed to record AI usage", { error }); } diff --git a/server/private/cleanup.ts b/server/private/cleanup.ts index 6c934cb0b..a1ad24976 100644 --- a/server/private/cleanup.ts +++ b/server/private/cleanup.ts @@ -18,12 +18,14 @@ import { flushBandwidthToDb } from "@server/routers/newt/handleReceiveBandwidthM import { flushConnectionLogToDb } from "#private/routers/newt"; import { flushSiteBandwidthToDb } from "@server/routers/gerbil/receiveBandwidth"; import { stopPingAccumulator } from "@server/routers/newt/pingAccumulator"; +import { shutdownUsageRecorder } from "@server/lib/aiBudgetEnforcement"; async function cleanup() { await stopPingAccumulator(); await flushBandwidthToDb(); await flushConnectionLogToDb(); await flushSiteBandwidthToDb(); + await shutdownUsageRecorder(); await rateLimitService.cleanup(); await wsCleanup(); await logStreamingManager.shutdown();