budget enforcement logic

This commit is contained in:
Owen
2026-08-10 17:17:24 -04:00
parent 88515a7dad
commit 8b8e7913dc
4 changed files with 607 additions and 6 deletions
+98
View File
@@ -1744,6 +1744,102 @@ export const aiBudgets = sqliteTable(
]
);
export const aiUsageRecords = sqliteTable(
"aiUsageRecords",
{
id: integer("id").primaryKey({ autoIncrement: true }),
orgId: text("orgId")
.notNull()
.references(() => orgs.orgId, { onDelete: "cascade" }),
providerId: integer("providerId")
.notNull()
.references(() => aiProviders.providerId, { onDelete: "cascade" }),
resourceId: integer("resourceId").references(
() => resources.resourceId,
{ onDelete: "cascade" }
),
siteResourceId: integer("siteResourceId").references(
() => siteResources.siteResourceId,
{ onDelete: "cascade" }
),
userId: text("userId").references(() => users.userId, {
onDelete: "set null"
}),
requestedModel: text("requestedModel").notNull(),
promptTokens: integer("promptTokens").notNull().default(0),
cacheReadTokens: integer("cacheReadTokens").notNull().default(0),
cacheWriteTokens: integer("cacheWriteTokens").notNull().default(0),
completionTokens: integer("completionTokens").notNull().default(0),
reasoningTokens: integer("reasoningTokens").notNull().default(0),
totalTokens: integer("totalTokens").notNull().default(0),
costUsd: real("costUsd"),
estimated: integer("estimated", { mode: "boolean" })
.notNull()
.default(false),
createdAt: integer("createdAt").notNull()
},
(t) => [
index("idx_ai_usage_records_org_provider_created").on(
t.orgId,
t.providerId,
t.createdAt
),
index("idx_ai_usage_records_org_resource_created").on(
t.orgId,
t.resourceId,
t.createdAt
),
index("idx_ai_usage_records_org_site_resource_created").on(
t.orgId,
t.siteResourceId,
t.createdAt
),
index("idx_ai_usage_records_org_user_created").on(
t.orgId,
t.userId,
t.createdAt
)
]
);
export const aiBudgetBreachEvents = sqliteTable(
"aiBudgetBreachEvents",
{
id: integer("id").primaryKey({ autoIncrement: true }),
orgId: text("orgId")
.notNull()
.references(() => orgs.orgId, { onDelete: "cascade" }),
budgetId: integer("budgetId")
.notNull()
.references(() => aiBudgets.budgetId, { onDelete: "cascade" }),
enforcement: text("enforcement").$type<"hard" | "soft">().notNull(),
unit: text("unit").$type<"usd" | "tokens">().notNull(),
period: text("period")
.$type<
| "monthly"
| "yearly"
| "lifetime"
| "daily"
| "hourly"
| "weekly"
>()
.notNull(),
amount: real("amount").notNull(),
usageAmount: real("usageAmount").notNull(),
blocked: integer("blocked", { mode: "boolean" }).notNull(),
requestUserId: text("requestUserId").references(() => users.userId, {
onDelete: "set null"
}),
createdAt: integer("createdAt").notNull()
},
(t) => [
index("idx_ai_budget_breach_events_budget_created").on(
t.budgetId,
t.createdAt
)
]
);
export type Org = InferSelectModel<typeof orgs>;
export type User = InferSelectModel<typeof users>;
export type Site = InferSelectModel<typeof sites>;
@@ -1829,6 +1925,8 @@ export type UserPolicy = InferSelectModel<typeof userPolicies>;
export type AiProvider = InferSelectModel<typeof aiProviders>;
export type AiModel = InferSelectModel<typeof aiModels>;
export type AiBudget = InferSelectModel<typeof aiBudgets>;
export type AiUsageRecord = InferSelectModel<typeof aiUsageRecords>;
export type AiBudgetBreachEvent = InferSelectModel<typeof aiBudgetBreachEvents>;
export type ResourceAiProvider = InferSelectModel<typeof resourceAiProviders>;
export type SiteResourceAiProvider = InferSelectModel<
typeof siteResourceAiProviders