@sixb/sqlite 0.1.6 → 0.1.7

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. package/dist/agents/rows.d.ts +7 -4
  2. package/dist/agents/rows.d.ts.map +1 -1
  3. package/dist/agents/runs.d.ts +3 -2
  4. package/dist/agents/runs.d.ts.map +1 -1
  5. package/dist/agents/threads.d.ts.map +1 -1
  6. package/dist/ai-cost-groups.d.ts +5 -0
  7. package/dist/ai-cost-groups.d.ts.map +1 -0
  8. package/dist/ai-cost-storage.d.ts +2 -1
  9. package/dist/ai-cost-storage.d.ts.map +1 -1
  10. package/dist/ai-limit-storage.d.ts +35 -0
  11. package/dist/ai-limit-storage.d.ts.map +1 -0
  12. package/dist/ai-usage-storage.d.ts +1 -0
  13. package/dist/ai-usage-storage.d.ts.map +1 -1
  14. package/dist/auth-storage/access-tokens.d.ts +1 -1
  15. package/dist/auth-storage/access-tokens.d.ts.map +1 -1
  16. package/dist/auth-storage/auth-storage.d.ts +4 -1
  17. package/dist/auth-storage/auth-storage.d.ts.map +1 -1
  18. package/dist/auth-storage/device-authorizations.d.ts +29 -0
  19. package/dist/auth-storage/device-authorizations.d.ts.map +1 -0
  20. package/dist/auth-storage/rows.d.ts +19 -1
  21. package/dist/auth-storage/rows.d.ts.map +1 -1
  22. package/dist/auth-storage/shared.d.ts +1 -1
  23. package/dist/auth-storage/shared.d.ts.map +1 -1
  24. package/dist/index.d.ts +2 -1
  25. package/dist/index.d.ts.map +1 -1
  26. package/dist/index.js +1898 -356
  27. package/dist/index.js.map +18 -15
  28. package/dist/migrations.d.ts.map +1 -1
  29. package/package.json +3 -3
  30. package/src/agents/rows.ts +36 -9
  31. package/src/agents/runs.ts +152 -13
  32. package/src/agents/threads.ts +2 -17
  33. package/src/ai-cost-groups.ts +100 -0
  34. package/src/ai-cost-storage.ts +67 -25
  35. package/src/ai-limit-storage.ts +981 -0
  36. package/src/ai-usage-storage.ts +24 -1
  37. package/src/auth-storage/access-tokens.ts +77 -68
  38. package/src/auth-storage/auth-storage.ts +80 -1
  39. package/src/auth-storage/device-authorizations.ts +153 -0
  40. package/src/auth-storage/rows.ts +41 -0
  41. package/src/auth-storage/shared.ts +1 -0
  42. package/src/index.ts +6 -0
  43. package/src/migrations/029-model-accounting.sql +79 -0
  44. package/src/migrations/030-ai-usage-limits.sql +100 -0
  45. package/src/migrations/031-subagent-runs.sql +231 -0
  46. package/src/migrations/032-conversation-run-spec.sql +78 -0
  47. package/src/migrations/033-retire-agent-definitions.sql +89 -0
  48. package/src/migrations/034-device-authorizations.sql +22 -0
  49. package/src/migrations.ts +18 -0
  50. package/src/workflow-run-storage.ts +8 -8
@@ -0,0 +1,981 @@
1
+ import type { SQLQueryBindings } from "bun:sqlite"
2
+ import {
3
+ type AiLimitAccountingEntry,
4
+ aiLimitAmountKey,
5
+ aiLimitQuantityFromAmount,
6
+ aiLimitReservationBuckets,
7
+ aiLimitReservationRequestKey,
8
+ aiLimitReservationRequestMatches,
9
+ aiLimitSubjectFromParts,
10
+ aiLimitSubjectKey,
11
+ aiLimitSubjectParts,
12
+ aiLimitSubjectsFromAccountingEntry,
13
+ assertAiLimitAccountingMatchesReservation,
14
+ assertNonBlank,
15
+ cloneValidDate,
16
+ type NormalizedAiLimitAmount,
17
+ normalizeAiLimitAmount,
18
+ normalizeAiLimitQuantities,
19
+ normalizeAiModelCallReservationIdentity,
20
+ normalizeCreateAiLimitPolicy,
21
+ normalizeReserveAiModelCall,
22
+ normalizeUpdateAiLimitPolicy,
23
+ resolveAiLimitActual,
24
+ } from "@sixb/core/internal/ai-limit-storage-provider"
25
+ import type {
26
+ AiLimitPeriod,
27
+ AiLimitPolicy,
28
+ AiLimitPolicyStatus,
29
+ AiLimitQuantity,
30
+ AiLimitReservationBucket,
31
+ AiLimitStorage,
32
+ AiLimitSubject,
33
+ AiModelCallReservation,
34
+ CreateAiLimitPolicyInput,
35
+ DeleteAiLimitPolicyInput,
36
+ GetAiLimitPolicyInput,
37
+ ListAiLimitPoliciesInput,
38
+ ListAiLimitPolicyStatusesInput,
39
+ MarkAiModelCallReservationUnknownInput,
40
+ ReconcileAiModelCallInput,
41
+ RecordAiModelCallLimitActualsInput,
42
+ ReserveAiModelCallInput,
43
+ ReserveAiModelCallResult,
44
+ UpdateAiLimitPolicyInput,
45
+ } from "@sixb/core/storage"
46
+ import { AiLimitStorageError, aiLimitCalendarMonth } from "@sixb/core/storage"
47
+ import { installFreshSqliteSchema } from "./migrations"
48
+ import {
49
+ closeSqliteStoreConnection,
50
+ openSqliteStoreConnection,
51
+ runImmediateTransaction,
52
+ type SqliteStoreConnection,
53
+ } from "./transactions"
54
+
55
+ export interface SqliteAiLimitStorageOptions {
56
+ readonly path?: string
57
+ readonly connection?: SqliteStoreConnection
58
+ }
59
+
60
+ /** SQLite-backed editable AI limits and atomic model-call reservations. */
61
+ export class SqliteAiLimitStorage implements AiLimitStorage {
62
+ private readonly connection: SqliteStoreConnection
63
+
64
+ constructor(options: SqliteAiLimitStorageOptions = {}) {
65
+ this.connection = openSqliteStoreConnection(options)
66
+ if (this.connection.installFreshSchema) installFreshSqliteSchema(this.connection.db)
67
+ }
68
+
69
+ async createPolicy(input: CreateAiLimitPolicyInput): Promise<AiLimitPolicy> {
70
+ const policy = normalizeCreateAiLimitPolicy(input)
71
+ return runImmediateTransaction(this.connection.db, () => {
72
+ const subject = aiLimitSubjectParts(policy.subject)
73
+ const limit = normalizeAiLimitAmount(policy.limit)
74
+ try {
75
+ this.connection.db
76
+ .query(
77
+ `
78
+ INSERT INTO ai_usage_limits (
79
+ project_id, id, subject_type, subject_id, meter, currency, limit_amount,
80
+ period_kind, enabled, created_at, updated_at
81
+ ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
82
+ `
83
+ )
84
+ .run(
85
+ policy.projectId,
86
+ policy.id,
87
+ subject.type,
88
+ subject.id,
89
+ limit.meter,
90
+ limit.currency,
91
+ limit.amount.toString(),
92
+ policy.period,
93
+ policy.enabled ? 1 : 0,
94
+ policy.createdAt.toISOString(),
95
+ policy.updatedAt.toISOString()
96
+ )
97
+ } catch (error) {
98
+ if (!isConstraintError(error)) throw error
99
+ throw duplicatePolicy(policy)
100
+ }
101
+ return this.requirePolicy(policy.projectId, policy.id)
102
+ })
103
+ }
104
+
105
+ async updatePolicy(input: UpdateAiLimitPolicyInput): Promise<AiLimitPolicy> {
106
+ const update = normalizeUpdateAiLimitPolicy(input)
107
+ return runImmediateTransaction(this.connection.db, () => {
108
+ const existing = this.requirePolicy(update.projectId, update.id)
109
+ if (update.limit !== undefined) assertSameDimension(existing.limit, update.limit)
110
+ const limit = normalizeAiLimitAmount(update.limit ?? existing.limit)
111
+ this.connection.db
112
+ .query(
113
+ `
114
+ UPDATE ai_usage_limits
115
+ SET limit_amount = ?, enabled = ?, updated_at = ?
116
+ WHERE project_id = ? AND id = ?
117
+ `
118
+ )
119
+ .run(
120
+ limit.amount.toString(),
121
+ (update.enabled ?? existing.enabled) ? 1 : 0,
122
+ update.updatedAt.toISOString(),
123
+ update.projectId,
124
+ update.id
125
+ )
126
+ return this.requirePolicy(update.projectId, update.id)
127
+ })
128
+ }
129
+
130
+ async deletePolicy(input: DeleteAiLimitPolicyInput): Promise<boolean> {
131
+ validatePolicyIdentity(input)
132
+ return runImmediateTransaction(this.connection.db, () => {
133
+ const result = this.connection.db
134
+ .query("DELETE FROM ai_usage_limits WHERE project_id = ? AND id = ?")
135
+ .run(input.projectId, input.id)
136
+ return result.changes > 0
137
+ })
138
+ }
139
+
140
+ async getPolicy(input: GetAiLimitPolicyInput): Promise<AiLimitPolicy | null> {
141
+ validatePolicyIdentity(input)
142
+ const row = this.connection.db
143
+ .query("SELECT * FROM ai_usage_limits WHERE project_id = ? AND id = ?")
144
+ .get(input.projectId, input.id) as AiLimitPolicyRow | null
145
+ return row ? policyFromRow(row) : null
146
+ }
147
+
148
+ async listPolicies(input: ListAiLimitPoliciesInput): Promise<readonly AiLimitPolicy[]> {
149
+ assertNonBlank(input.projectId, "projectId")
150
+ return this.policyRows(input.projectId, input.includeDisabled ?? false).map(policyFromRow)
151
+ }
152
+
153
+ async listPolicyStatuses(
154
+ input: ListAiLimitPolicyStatusesInput
155
+ ): Promise<readonly AiLimitPolicyStatus[]> {
156
+ assertNonBlank(input.projectId, "projectId")
157
+ const period = aiLimitCalendarMonth(input.at ?? new Date())
158
+ const existingGroups =
159
+ input.existingGroupIds === undefined ? undefined : new Set(input.existingGroupIds)
160
+ return runImmediateTransaction(this.connection.db, () =>
161
+ this.policyRows(input.projectId, input.includeDisabled ?? false).map((row) =>
162
+ this.statusForPolicy(policyFromRow(row), period, existingGroups)
163
+ )
164
+ )
165
+ }
166
+
167
+ async reserveModelCall(input: ReserveAiModelCallInput): Promise<ReserveAiModelCallResult> {
168
+ const request = normalizeReserveAiModelCall(input)
169
+ return runImmediateTransaction(this.connection.db, () => {
170
+ const existing = this.findReservation(request.identity)
171
+ if (existing) {
172
+ if (!aiLimitReservationRequestMatches(existing, request)) {
173
+ throw new AiLimitStorageError(
174
+ "reservation_conflict",
175
+ `[SixbSqlite] AI model-call reservation '${input.callId}' was replayed with different subjects, estimates, or period.`
176
+ )
177
+ }
178
+ if (existing.state !== "active") {
179
+ return { status: "terminal", reservation: existing, created: false }
180
+ }
181
+ return { status: "reserved", reservation: existing, created: false }
182
+ }
183
+ this.assertExecutionExists(request.identity.projectId, request.identity.executionId)
184
+
185
+ const subjectKeys = new Set(request.subjects.map(aiLimitSubjectKey))
186
+ const estimates = new Map(
187
+ request.estimates.map((quantity) => {
188
+ const amount = normalizeAiLimitAmount(quantity)
189
+ return [aiLimitAmountKey(amount), amount] as const
190
+ })
191
+ )
192
+ const enabledPolicies = this.policyRows(request.identity.projectId, false).map(policyFromRow)
193
+ const exhaustedPolicies: AiLimitPolicyStatus[] = []
194
+ const unavailablePolicies: AiLimitPolicyStatus[] = []
195
+ const unavailableReasons = new Set<"missingEstimate" | "incompleteAccounting">()
196
+ for (const policy of enabledPolicies) {
197
+ if (!subjectKeys.has(aiLimitSubjectKey(policy.subject))) continue
198
+ const limit = normalizeAiLimitAmount(policy.limit)
199
+ const estimate = estimates.get(aiLimitAmountKey(limit))
200
+ const status = this.statusForPolicy(policy, request.period)
201
+ if (!estimate || status.accountingStatus === "unavailable") {
202
+ unavailablePolicies.push(status)
203
+ if (!estimate) unavailableReasons.add("missingEstimate")
204
+ if (status.accountingStatus === "unavailable") {
205
+ unavailableReasons.add("incompleteAccounting")
206
+ }
207
+ continue
208
+ }
209
+ const consumption = status.consumption
210
+ const consumed =
211
+ normalizeAiLimitAmount(consumption.actual).amount +
212
+ normalizeAiLimitAmount(consumption.reserved).amount +
213
+ normalizeAiLimitAmount(consumption.unknown).amount
214
+ if (consumed + estimate.amount > limit.amount) exhaustedPolicies.push(status)
215
+ }
216
+ if (unavailablePolicies.length > 0) {
217
+ unavailablePolicies.sort((left, right) => left.policy.id.localeCompare(right.policy.id))
218
+ return {
219
+ status: "unavailable",
220
+ unavailablePolicies,
221
+ reasons: [...unavailableReasons].sort(),
222
+ }
223
+ }
224
+ if (exhaustedPolicies.length > 0) {
225
+ exhaustedPolicies.sort((left, right) => left.policy.id.localeCompare(right.policy.id))
226
+ return {
227
+ status: "denied",
228
+ exhaustedPolicies,
229
+ resetAt: new Date(
230
+ Math.min(...exhaustedPolicies.map((status) => status.period.resetAt.getTime()))
231
+ ),
232
+ }
233
+ }
234
+
235
+ const buckets = aiLimitReservationBuckets(
236
+ enabledPolicies,
237
+ request.subjects,
238
+ request.estimates
239
+ )
240
+ if (buckets.length === 0) return { status: "notRequired" }
241
+ for (const bucket of buckets) {
242
+ const amount = normalizeAiLimitAmount(bucket.estimate)
243
+ this.changeCounters(
244
+ request.identity.projectId,
245
+ bucket.subject,
246
+ amount,
247
+ request.period,
248
+ { reserved: amount.amount },
249
+ request.reservedAt
250
+ )
251
+ }
252
+ const reservation: AiModelCallReservation = {
253
+ ...request.identity,
254
+ buckets,
255
+ requestKey: aiLimitReservationRequestKey(request),
256
+ period: request.period,
257
+ state: "active",
258
+ reservedAt: request.reservedAt,
259
+ updatedAt: new Date(request.reservedAt),
260
+ }
261
+ this.insertReservation(reservation)
262
+ return { status: "reserved", reservation: structuredClone(reservation), created: true }
263
+ })
264
+ }
265
+
266
+ async recordModelCallActuals(input: RecordAiModelCallLimitActualsInput): Promise<void> {
267
+ assertNonBlank(input.projectId, "projectId")
268
+ assertNonBlank(input.usageRecordId, "usageRecordId")
269
+ const recordedAt = cloneValidDate(input.recordedAt ?? new Date(), "recordedAt")
270
+ runImmediateTransaction(this.connection.db, () => {
271
+ const entry = this.requireAccountingEntry(input.projectId, input.usageRecordId)
272
+ const period = aiLimitCalendarMonth(entry.occurredAt)
273
+ for (const subjectValue of aiLimitSubjectsFromAccountingEntry(entry)) {
274
+ const subject = aiLimitSubjectParts(subjectValue)
275
+ const rows = this.connection.db
276
+ .query(
277
+ `
278
+ SELECT meter, currency, actual_amount, accounting_status
279
+ FROM ai_usage_limit_periods
280
+ WHERE project_id = ? AND subject_type = ? AND subject_id = ?
281
+ AND period_start = ?
282
+ `
283
+ )
284
+ .all(
285
+ input.projectId,
286
+ subject.type,
287
+ subject.id,
288
+ period.start.toISOString()
289
+ ) as AiLimitActualStateRow[]
290
+ for (const row of rows) {
291
+ const actual = resolveAiLimitActual([entry], row)
292
+ this.connection.db
293
+ .query(
294
+ `
295
+ UPDATE ai_usage_limit_periods
296
+ SET actual_amount = ?, accounting_status = ?, updated_at = ?
297
+ WHERE project_id = ? AND subject_type = ? AND subject_id = ?
298
+ AND meter = ? AND currency = ? AND period_start = ?
299
+ `
300
+ )
301
+ .run(
302
+ (BigInt(row.actual_amount) + actual.amount).toString(),
303
+ row.accounting_status === "unavailable" || actual.accountingStatus === "unavailable"
304
+ ? "unavailable"
305
+ : "complete",
306
+ recordedAt.toISOString(),
307
+ input.projectId,
308
+ subject.type,
309
+ subject.id,
310
+ row.meter,
311
+ row.currency,
312
+ period.start.toISOString()
313
+ )
314
+ }
315
+ }
316
+ })
317
+ }
318
+
319
+ async reconcileModelCall(input: ReconcileAiModelCallInput): Promise<AiModelCallReservation> {
320
+ const identity = normalizeAiModelCallReservationIdentity(input)
321
+ assertNonBlank(input.usageRecordId, "usageRecordId")
322
+ const reconciledAt = cloneValidDate(input.reconciledAt ?? new Date(), "reconciledAt")
323
+ return runImmediateTransaction(this.connection.db, () => {
324
+ const reservation = this.findReservation(identity)
325
+ if (!reservation) throw missingReservation(identity)
326
+ if (reservation.state === "reconciled") {
327
+ if (reservation.usageRecordId === input.usageRecordId) {
328
+ return reservation
329
+ }
330
+ throw reconciliationConflict(identity.callId)
331
+ }
332
+ const accounting = this.requireAccountingEntry(identity.projectId, input.usageRecordId)
333
+ assertAiLimitAccountingMatchesReservation(reservation, accounting)
334
+ this.reconcileEstimateCounters(reservation, accounting, reconciledAt)
335
+ const reconciled: AiModelCallReservation = {
336
+ ...reservation,
337
+ state: "reconciled",
338
+ usageRecordId: input.usageRecordId,
339
+ updatedAt: reconciledAt,
340
+ }
341
+ this.updateReservation(reconciled)
342
+ return structuredClone(reconciled)
343
+ })
344
+ }
345
+
346
+ async markReservationUnknown(
347
+ input: MarkAiModelCallReservationUnknownInput
348
+ ): Promise<AiModelCallReservation> {
349
+ const identity = normalizeAiModelCallReservationIdentity(input)
350
+ const markedAt = cloneValidDate(input.markedAt ?? new Date(), "markedAt")
351
+ return runImmediateTransaction(this.connection.db, () => {
352
+ const reservation = this.findReservation(identity)
353
+ if (!reservation) throw missingReservation(identity)
354
+ if (reservation.state === "unknown") return reservation
355
+ if (reservation.state !== "active") throw invalidState(reservation, "mark unknown")
356
+ this.moveEstimateCounters(reservation, { reserved: -1n, unknown: 1n }, markedAt)
357
+ const unknown = { ...reservation, state: "unknown" as const, updatedAt: markedAt }
358
+ this.updateReservation(unknown)
359
+ return structuredClone(unknown)
360
+ })
361
+ }
362
+
363
+ close(): void {
364
+ closeSqliteStoreConnection(this.connection)
365
+ }
366
+
367
+ private policyRows(projectId: string, includeDisabled: boolean): AiLimitPolicyRow[] {
368
+ return this.connection.db
369
+ .query(
370
+ `
371
+ SELECT * FROM ai_usage_limits
372
+ WHERE project_id = ? AND (? = 1 OR enabled = 1)
373
+ ORDER BY id
374
+ `
375
+ )
376
+ .all(projectId, includeDisabled ? 1 : 0) as AiLimitPolicyRow[]
377
+ }
378
+
379
+ private requirePolicy(projectId: string, id: string): AiLimitPolicy {
380
+ const row = this.connection.db
381
+ .query("SELECT * FROM ai_usage_limits WHERE project_id = ? AND id = ?")
382
+ .get(projectId, id) as AiLimitPolicyRow | null
383
+ if (row) return policyFromRow(row)
384
+ throw new AiLimitStorageError(
385
+ "missing_policy",
386
+ `[SixbSqlite] AI limit policy '${id}' does not exist in project '${projectId}'.`
387
+ )
388
+ }
389
+
390
+ private statusForPolicy(
391
+ policy: AiLimitPolicy,
392
+ period: AiLimitPeriod,
393
+ existingGroups?: ReadonlySet<string>
394
+ ): AiLimitPolicyStatus {
395
+ const subject = aiLimitSubjectParts(policy.subject)
396
+ const limit = normalizeAiLimitAmount(policy.limit)
397
+ let row = this.connection.db
398
+ .query(
399
+ `
400
+ SELECT actual_amount, reserved_amount, unknown_amount, accounting_status
401
+ FROM ai_usage_limit_periods
402
+ WHERE project_id = ? AND subject_type = ? AND subject_id = ?
403
+ AND meter = ? AND currency = ? AND period_start = ?
404
+ `
405
+ )
406
+ .get(
407
+ policy.projectId,
408
+ subject.type,
409
+ subject.id,
410
+ limit.meter,
411
+ limit.currency,
412
+ period.start.toISOString()
413
+ ) as AiLimitPeriodStateRow | null
414
+ // Re-read incomplete accounting inside the transaction; retain in-flight estimates.
415
+ if (!row || row.accounting_status === "unavailable") {
416
+ const actual = resolveAiLimitActual(this.accountingEntries(policy, period), limit)
417
+ row = {
418
+ actual_amount: actual.amount.toString(),
419
+ reserved_amount: row?.reserved_amount ?? "0",
420
+ unknown_amount: row?.unknown_amount ?? "0",
421
+ accounting_status: actual.accountingStatus,
422
+ }
423
+ this.connection.db
424
+ .query(
425
+ `
426
+ INSERT INTO ai_usage_limit_periods (
427
+ project_id, subject_type, subject_id, meter, currency, period_start, period_end,
428
+ actual_amount, reserved_amount, unknown_amount, accounting_status, updated_at
429
+ ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, '0', '0', ?, ?)
430
+ ON CONFLICT (project_id, subject_type, subject_id, meter, currency, period_start)
431
+ DO UPDATE SET
432
+ actual_amount = excluded.actual_amount,
433
+ accounting_status = excluded.accounting_status,
434
+ updated_at = excluded.updated_at
435
+ `
436
+ )
437
+ .run(
438
+ policy.projectId,
439
+ subject.type,
440
+ subject.id,
441
+ limit.meter,
442
+ limit.currency,
443
+ period.start.toISOString(),
444
+ period.end.toISOString(),
445
+ row.actual_amount,
446
+ row.accounting_status,
447
+ new Date().toISOString()
448
+ )
449
+ }
450
+ return policyStatus(policy, period, row, existingGroups)
451
+ }
452
+
453
+ private changeCounters(
454
+ projectId: string,
455
+ subjectValue: AiLimitSubject,
456
+ amount: NormalizedAiLimitAmount,
457
+ period: AiLimitPeriod,
458
+ changes: { readonly reserved?: bigint; readonly unknown?: bigint },
459
+ updatedAt: Date
460
+ ): void {
461
+ const subject = aiLimitSubjectParts(subjectValue)
462
+ const row = this.connection.db
463
+ .query(
464
+ `
465
+ SELECT actual_amount, reserved_amount, unknown_amount, accounting_status
466
+ FROM ai_usage_limit_periods
467
+ WHERE project_id = ? AND subject_type = ? AND subject_id = ?
468
+ AND meter = ? AND currency = ? AND period_start = ?
469
+ `
470
+ )
471
+ .get(
472
+ projectId,
473
+ subject.type,
474
+ subject.id,
475
+ amount.meter,
476
+ amount.currency,
477
+ period.start.toISOString()
478
+ ) as AiLimitPeriodStateRow | null
479
+ const reserved = BigInt(row?.reserved_amount ?? "0") + (changes.reserved ?? 0n)
480
+ const unknown = BigInt(row?.unknown_amount ?? "0") + (changes.unknown ?? 0n)
481
+ if (reserved < 0n || unknown < 0n) {
482
+ throw new Error("[SixbSqlite] AI limit period counter became negative.")
483
+ }
484
+ this.connection.db
485
+ .query(
486
+ `
487
+ INSERT INTO ai_usage_limit_periods (
488
+ project_id, subject_type, subject_id, meter, currency, period_start, period_end,
489
+ actual_amount, reserved_amount, unknown_amount, accounting_status, updated_at
490
+ ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
491
+ ON CONFLICT (project_id, subject_type, subject_id, meter, currency, period_start)
492
+ DO UPDATE SET
493
+ period_end = excluded.period_end,
494
+ actual_amount = excluded.actual_amount,
495
+ reserved_amount = excluded.reserved_amount,
496
+ unknown_amount = excluded.unknown_amount,
497
+ accounting_status = excluded.accounting_status,
498
+ updated_at = excluded.updated_at
499
+ `
500
+ )
501
+ .run(
502
+ projectId,
503
+ subject.type,
504
+ subject.id,
505
+ amount.meter,
506
+ amount.currency,
507
+ period.start.toISOString(),
508
+ period.end.toISOString(),
509
+ row?.actual_amount ?? "0",
510
+ reserved.toString(),
511
+ unknown.toString(),
512
+ row?.accounting_status ?? "complete",
513
+ updatedAt.toISOString()
514
+ )
515
+ }
516
+
517
+ private moveEstimateCounters(
518
+ reservation: AiModelCallReservation,
519
+ directions: { readonly reserved?: bigint; readonly unknown?: bigint },
520
+ updatedAt: Date
521
+ ): void {
522
+ for (const bucket of reservation.buckets) {
523
+ const amount = normalizeAiLimitAmount(bucket.estimate)
524
+ this.changeCounters(
525
+ reservation.projectId,
526
+ bucket.subject,
527
+ amount,
528
+ reservation.period,
529
+ {
530
+ ...(directions.reserved === undefined
531
+ ? {}
532
+ : { reserved: directions.reserved * amount.amount }),
533
+ ...(directions.unknown === undefined
534
+ ? {}
535
+ : { unknown: directions.unknown * amount.amount }),
536
+ },
537
+ updatedAt
538
+ )
539
+ }
540
+ }
541
+
542
+ private reconcileEstimateCounters(
543
+ reservation: AiModelCallReservation,
544
+ entry: AiLimitAccountingEntry,
545
+ updatedAt: Date
546
+ ): void {
547
+ for (const bucket of reservation.buckets) {
548
+ const amount = normalizeAiLimitAmount(bucket.estimate)
549
+ const actual = resolveAiLimitActual([entry], amount)
550
+ if (actual.accountingStatus === "complete") {
551
+ this.changeCounters(
552
+ reservation.projectId,
553
+ bucket.subject,
554
+ amount,
555
+ reservation.period,
556
+ reservation.state === "active"
557
+ ? { reserved: -amount.amount }
558
+ : { unknown: -amount.amount },
559
+ updatedAt
560
+ )
561
+ } else if (reservation.state === "active") {
562
+ this.changeCounters(
563
+ reservation.projectId,
564
+ bucket.subject,
565
+ amount,
566
+ reservation.period,
567
+ { reserved: -amount.amount, unknown: amount.amount },
568
+ updatedAt
569
+ )
570
+ }
571
+ }
572
+ }
573
+
574
+ private insertReservation(reservation: AiModelCallReservation): void {
575
+ this.connection.db
576
+ .query(
577
+ `
578
+ INSERT INTO ai_model_call_reservations (
579
+ project_id, execution_id, attempt, call_id, buckets, request_key,
580
+ period_start, period_end, state, usage_record_id, reserved_at, updated_at
581
+ ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
582
+ `
583
+ )
584
+ .run(...reservationSqlValues(reservation))
585
+ }
586
+
587
+ private updateReservation(reservation: AiModelCallReservation): void {
588
+ this.connection.db
589
+ .query(
590
+ `
591
+ UPDATE ai_model_call_reservations
592
+ SET state = ?, usage_record_id = ?, updated_at = ?
593
+ WHERE project_id = ? AND execution_id = ? AND attempt = ? AND call_id = ?
594
+ `
595
+ )
596
+ .run(
597
+ reservation.state,
598
+ reservation.usageRecordId ?? null,
599
+ reservation.updatedAt.toISOString(),
600
+ reservation.projectId,
601
+ reservation.executionId,
602
+ reservation.attempt,
603
+ reservation.callId
604
+ )
605
+ }
606
+
607
+ private findReservation(input: {
608
+ readonly projectId: string
609
+ readonly executionId: string
610
+ readonly attempt: number
611
+ readonly callId: string
612
+ }): AiModelCallReservation | null {
613
+ const row = this.connection.db
614
+ .query(
615
+ `
616
+ SELECT * FROM ai_model_call_reservations
617
+ WHERE project_id = ? AND execution_id = ? AND attempt = ? AND call_id = ?
618
+ `
619
+ )
620
+ .get(
621
+ input.projectId,
622
+ input.executionId,
623
+ input.attempt,
624
+ input.callId
625
+ ) as AiModelCallReservationRow | null
626
+ return row ? reservationFromRow(row) : null
627
+ }
628
+
629
+ private accountingEntries(
630
+ policy: AiLimitPolicy,
631
+ period: AiLimitPeriod
632
+ ): readonly AiLimitAccountingEntry[] {
633
+ const subject = aiLimitSubjectParts(policy.subject)
634
+ const rows = this.connection.db
635
+ .query(
636
+ `
637
+ SELECT usage.project_id, usage.id AS usage_record_id, usage.execution_id,
638
+ usage.attempt, usage.call_id, usage.occurred_at,
639
+ CAST(usage.total_tokens AS TEXT) AS total_tokens,
640
+ executions.requested_by_user_id,
641
+ executions.requested_by_service_account_id,
642
+ valuations.status AS valuation_status,
643
+ valuations.currency AS valuation_currency,
644
+ CAST(valuations.amount_nanos AS TEXT) AS valuation_amount_nanos
645
+ FROM ai_model_call_usage AS usage
646
+ JOIN executions
647
+ ON executions.project_id = usage.project_id AND executions.id = usage.execution_id
648
+ LEFT JOIN ai_model_call_valuations AS valuations
649
+ ON valuations.project_id = usage.project_id
650
+ AND valuations.usage_record_id = usage.id
651
+ WHERE usage.project_id = ? AND usage.occurred_at >= ? AND usage.occurred_at < ?
652
+ AND (
653
+ ? = 'project'
654
+ OR (
655
+ ? = 'group'
656
+ AND EXISTS (
657
+ SELECT 1 FROM ai_model_call_usage_groups AS usage_groups
658
+ WHERE usage_groups.project_id = usage.project_id
659
+ AND usage_groups.usage_record_id = usage.id
660
+ AND usage_groups.group_id = ?
661
+ )
662
+ )
663
+ OR (? = 'user' AND executions.requested_by_user_id = ?)
664
+ OR (
665
+ ? = 'serviceAccount'
666
+ AND executions.requested_by_service_account_id = ?
667
+ )
668
+ )
669
+ ORDER BY usage.occurred_at, usage.id
670
+ `
671
+ )
672
+ .all(
673
+ policy.projectId,
674
+ period.start.toISOString(),
675
+ period.end.toISOString(),
676
+ subject.type,
677
+ subject.type,
678
+ subject.id,
679
+ subject.type,
680
+ subject.id,
681
+ subject.type,
682
+ subject.id
683
+ ) as AiLimitAccountingRow[]
684
+ return rows.map((row) => accountingEntryFromRow(row, policy.subject))
685
+ }
686
+
687
+ private requireAccountingEntry(projectId: string, usageRecordId: string): AiLimitAccountingEntry {
688
+ const row = this.connection.db
689
+ .query(
690
+ `
691
+ SELECT usage.project_id, usage.id AS usage_record_id, usage.execution_id,
692
+ usage.attempt, usage.call_id, usage.occurred_at,
693
+ CAST(usage.total_tokens AS TEXT) AS total_tokens,
694
+ executions.requested_by_user_id,
695
+ executions.requested_by_service_account_id,
696
+ valuations.status AS valuation_status,
697
+ valuations.currency AS valuation_currency,
698
+ CAST(valuations.amount_nanos AS TEXT) AS valuation_amount_nanos
699
+ FROM ai_model_call_usage AS usage
700
+ JOIN executions
701
+ ON executions.project_id = usage.project_id AND executions.id = usage.execution_id
702
+ LEFT JOIN ai_model_call_valuations AS valuations
703
+ ON valuations.project_id = usage.project_id
704
+ AND valuations.usage_record_id = usage.id
705
+ WHERE usage.project_id = ? AND usage.id = ?
706
+ `
707
+ )
708
+ .get(projectId, usageRecordId) as AiLimitAccountingRow | null
709
+ if (row) {
710
+ const groups = this.connection.db
711
+ .query(
712
+ `
713
+ SELECT group_id FROM ai_model_call_usage_groups
714
+ WHERE project_id = ? AND usage_record_id = ?
715
+ ORDER BY group_id
716
+ `
717
+ )
718
+ .all(projectId, usageRecordId) as { readonly group_id: string }[]
719
+ return accountingEntryFromRow(
720
+ row,
721
+ undefined,
722
+ groups.map((group) => group.group_id)
723
+ )
724
+ }
725
+ throw new AiLimitStorageError(
726
+ "missing_usage_record",
727
+ `[SixbSqlite] AI usage record '${usageRecordId}' does not exist in project '${projectId}'.`
728
+ )
729
+ }
730
+
731
+ private assertExecutionExists(projectId: string, executionId: string): void {
732
+ const row = this.connection.db
733
+ .query("SELECT 1 FROM executions WHERE project_id = ? AND id = ?")
734
+ .get(projectId, executionId)
735
+ if (row) return
736
+ throw new AiLimitStorageError(
737
+ "missing_execution",
738
+ `[SixbSqlite] AI limit execution '${executionId}' does not exist in project '${projectId}'.`
739
+ )
740
+ }
741
+ }
742
+
743
+ interface AiLimitPolicyRow {
744
+ readonly project_id: string
745
+ readonly id: string
746
+ readonly subject_type: AiLimitSubject["type"]
747
+ readonly subject_id: string
748
+ readonly meter: AiLimitQuantity["meter"]
749
+ readonly currency: string
750
+ readonly limit_amount: string
751
+ readonly period_kind: "calendarMonth"
752
+ readonly enabled: number
753
+ readonly created_at: string
754
+ readonly updated_at: string
755
+ }
756
+
757
+ interface AiLimitPeriodStateRow {
758
+ readonly actual_amount: string
759
+ readonly reserved_amount: string
760
+ readonly unknown_amount: string
761
+ readonly accounting_status: "complete" | "unavailable"
762
+ }
763
+
764
+ interface AiLimitActualStateRow {
765
+ readonly meter: AiLimitQuantity["meter"]
766
+ readonly currency: string
767
+ readonly actual_amount: string
768
+ readonly accounting_status: "complete" | "unavailable"
769
+ }
770
+
771
+ interface AiLimitAccountingRow {
772
+ readonly project_id: string
773
+ readonly usage_record_id: string
774
+ readonly execution_id: string
775
+ readonly attempt: number
776
+ readonly call_id: string
777
+ readonly occurred_at: string
778
+ readonly total_tokens: string | null
779
+ readonly requested_by_user_id: string | null
780
+ readonly requested_by_service_account_id: string | null
781
+ readonly valuation_status: "rated" | "unpriceable" | null
782
+ readonly valuation_currency: string | null
783
+ readonly valuation_amount_nanos: string | null
784
+ }
785
+
786
+ interface AiModelCallReservationRow {
787
+ readonly project_id: string
788
+ readonly execution_id: string
789
+ readonly attempt: number
790
+ readonly call_id: string
791
+ readonly buckets: string
792
+ readonly request_key: string
793
+ readonly period_start: string
794
+ readonly period_end: string
795
+ readonly state: AiModelCallReservation["state"]
796
+ readonly usage_record_id: string | null
797
+ readonly reserved_at: string
798
+ readonly updated_at: string
799
+ }
800
+
801
+ function policyFromRow(row: AiLimitPolicyRow): AiLimitPolicy {
802
+ return {
803
+ id: row.id,
804
+ projectId: row.project_id,
805
+ subject: aiLimitSubjectFromParts(row.subject_type, row.subject_id),
806
+ limit: aiLimitQuantityFromAmount({
807
+ meter: row.meter,
808
+ currency: row.currency,
809
+ amount: BigInt(row.limit_amount),
810
+ }),
811
+ period: row.period_kind,
812
+ enabled: row.enabled === 1,
813
+ createdAt: new Date(row.created_at),
814
+ updatedAt: new Date(row.updated_at),
815
+ }
816
+ }
817
+
818
+ function policyStatus(
819
+ policy: AiLimitPolicy,
820
+ period: AiLimitPeriod,
821
+ row: AiLimitPeriodStateRow,
822
+ existingGroups?: ReadonlySet<string>
823
+ ): AiLimitPolicyStatus {
824
+ const limit = normalizeAiLimitAmount(policy.limit)
825
+ const actual = BigInt(row.actual_amount)
826
+ const reserved = BigInt(row.reserved_amount)
827
+ const unknown = BigInt(row.unknown_amount)
828
+ const consumed = actual + reserved + unknown
829
+ const quantity = (amount: bigint) => aiLimitQuantityFromAmount({ ...limit, amount })
830
+ return {
831
+ policy,
832
+ period: structuredClone(period),
833
+ consumption: {
834
+ actual: quantity(actual),
835
+ reserved: quantity(reserved),
836
+ unknown: quantity(unknown),
837
+ remaining: quantity(limit.amount > consumed ? limit.amount - consumed : 0n),
838
+ },
839
+ accountingStatus: row.accounting_status,
840
+ exhausted: consumed >= limit.amount,
841
+ orphaned:
842
+ policy.subject.type === "group" &&
843
+ existingGroups !== undefined &&
844
+ !existingGroups.has(policy.subject.id),
845
+ }
846
+ }
847
+
848
+ function accountingEntryFromRow(
849
+ row: AiLimitAccountingRow,
850
+ subject?: AiLimitSubject,
851
+ requesterGroupIds = subject?.type === "group" ? [subject.id] : []
852
+ ): AiLimitAccountingEntry {
853
+ const totalTokens = row.total_tokens === null ? undefined : Number(row.total_tokens)
854
+ const requester = row.requested_by_user_id
855
+ ? ({ type: "user", id: row.requested_by_user_id } as const)
856
+ : row.requested_by_service_account_id
857
+ ? ({ type: "serviceAccount", id: row.requested_by_service_account_id } as const)
858
+ : undefined
859
+ return {
860
+ projectId: row.project_id,
861
+ usageRecordId: row.usage_record_id,
862
+ executionId: row.execution_id,
863
+ attempt: row.attempt,
864
+ callId: row.call_id,
865
+ occurredAt: new Date(row.occurred_at),
866
+ ...(totalTokens !== undefined && Number.isSafeInteger(totalTokens) ? { totalTokens } : {}),
867
+ requesterGroupIds,
868
+ ...(requester ? { requester } : {}),
869
+ ...(row.valuation_status === "rated" &&
870
+ row.valuation_currency !== null &&
871
+ row.valuation_amount_nanos !== null
872
+ ? {
873
+ valuation: {
874
+ status: "rated" as const,
875
+ money: {
876
+ currency: row.valuation_currency,
877
+ amountNanos: row.valuation_amount_nanos,
878
+ },
879
+ },
880
+ }
881
+ : row.valuation_status === "unpriceable"
882
+ ? { valuation: { status: "unavailable" as const } }
883
+ : {}),
884
+ }
885
+ }
886
+
887
+ function reservationFromRow(row: AiModelCallReservationRow): AiModelCallReservation {
888
+ return {
889
+ projectId: row.project_id,
890
+ executionId: row.execution_id,
891
+ attempt: row.attempt,
892
+ callId: row.call_id,
893
+ buckets: normalizeBucketsJson(row.buckets),
894
+ requestKey: row.request_key,
895
+ period: {
896
+ kind: "calendarMonth",
897
+ start: new Date(row.period_start),
898
+ end: new Date(row.period_end),
899
+ resetAt: new Date(row.period_end),
900
+ },
901
+ state: row.state,
902
+ ...(row.usage_record_id === null ? {} : { usageRecordId: row.usage_record_id }),
903
+ reservedAt: new Date(row.reserved_at),
904
+ updatedAt: new Date(row.updated_at),
905
+ }
906
+ }
907
+
908
+ function normalizeBucketsJson(value: string): readonly AiLimitReservationBucket[] {
909
+ const buckets = JSON.parse(value) as readonly AiLimitReservationBucket[]
910
+ return buckets.map((bucket) => ({
911
+ subject: structuredClone(bucket.subject),
912
+ estimate: normalizeAiLimitQuantities([bucket.estimate], "stored bucket")[0]!,
913
+ }))
914
+ }
915
+
916
+ function reservationSqlValues(reservation: AiModelCallReservation): SQLQueryBindings[] {
917
+ return [
918
+ reservation.projectId,
919
+ reservation.executionId,
920
+ reservation.attempt,
921
+ reservation.callId,
922
+ JSON.stringify(reservation.buckets),
923
+ reservation.requestKey,
924
+ reservation.period.start.toISOString(),
925
+ reservation.period.end.toISOString(),
926
+ reservation.state,
927
+ reservation.usageRecordId ?? null,
928
+ reservation.reservedAt.toISOString(),
929
+ reservation.updatedAt.toISOString(),
930
+ ]
931
+ }
932
+
933
+ function assertSameDimension(previous: AiLimitQuantity, next: AiLimitQuantity): void {
934
+ if (
935
+ aiLimitAmountKey(normalizeAiLimitAmount(previous)) !==
936
+ aiLimitAmountKey(normalizeAiLimitAmount(next))
937
+ ) {
938
+ throw new TypeError("[Sixb] AI limit policy meter and cost currency are immutable.")
939
+ }
940
+ }
941
+
942
+ function validatePolicyIdentity(input: GetAiLimitPolicyInput): void {
943
+ assertNonBlank(input.projectId, "projectId")
944
+ assertNonBlank(input.id, "policy id")
945
+ }
946
+
947
+ function duplicatePolicy(policy: AiLimitPolicy): AiLimitStorageError {
948
+ return new AiLimitStorageError(
949
+ "duplicate_policy",
950
+ `[SixbSqlite] An AI limit policy already exists for policy ID or subject-meter dimension '${policy.id}'.`
951
+ )
952
+ }
953
+
954
+ function missingReservation(input: {
955
+ readonly executionId: string
956
+ readonly attempt: number
957
+ readonly callId: string
958
+ }): AiLimitStorageError {
959
+ return new AiLimitStorageError(
960
+ "missing_reservation",
961
+ `[SixbSqlite] AI model-call reservation '${input.callId}' does not exist for execution '${input.executionId}' attempt ${input.attempt}.`
962
+ )
963
+ }
964
+
965
+ function reconciliationConflict(callId: string): AiLimitStorageError {
966
+ return new AiLimitStorageError(
967
+ "reconciliation_conflict",
968
+ `[SixbSqlite] AI model-call reservation '${callId}' was reconciled with a different usage record.`
969
+ )
970
+ }
971
+
972
+ function invalidState(reservation: AiModelCallReservation, operation: string): AiLimitStorageError {
973
+ return new AiLimitStorageError(
974
+ "invalid_reservation_state",
975
+ `[SixbSqlite] Cannot ${operation} AI model-call reservation '${reservation.callId}' in state '${reservation.state}'.`
976
+ )
977
+ }
978
+
979
+ function isConstraintError(error: unknown): boolean {
980
+ return error instanceof Error && /constraint/i.test(error.message)
981
+ }