@@ -10,27 +10,47 @@ import * as schema from '@sim/db/schema'
1010import { readTestDatabaseUrl , readTestRedisUrl } from '@sim/db/testing/test-infrastructure'
1111import { redisConfigMock , redisConfigMockFns } from '@sim/testing/mocks/redis-config.mock'
1212import { generateId } from '@sim/utils/id'
13+ import { randomInt , randomItem } from '@sim/utils/random'
1314import { drizzle } from 'drizzle-orm/postgres-js'
1415import Redis from 'ioredis'
1516import postgres from 'postgres'
1617import { afterAll , afterEach , beforeAll , beforeEach , describe , expect , it , vi } from 'vitest'
1718
18- const { transaction } = vi . hoisted ( ( ) => ( { transaction : vi . fn ( ) } ) )
19+ const { transaction, summedPeriods } = vi . hoisted ( ( ) => ( {
20+ transaction : vi . fn ( ) ,
21+ summedPeriods : [ ] as UsageQueryPeriod [ ] ,
22+ } ) )
1923const databaseUrl = readTestDatabaseUrl ( )
2024const redisUrl = readTestRedisUrl ( )
2125
2226vi . mock ( '@sim/db' , ( ) => ( { db : { transaction } , dbReplica : { } } ) )
2327vi . mock ( '@/lib/core/config/redis' , ( ) => redisConfigMock )
28+ vi . mock ( '@/lib/billing/core/usage-log' , async ( importOriginal ) => {
29+ const actual = await importOriginal < typeof import ( '@/lib/billing/core/usage-log' ) > ( )
30+ return {
31+ ...actual ,
32+ getBillingPeriodUsageCost : ( ...args : Parameters < typeof actual . getBillingPeriodUsageCost > ) => {
33+ summedPeriods . push ( args [ 1 ] )
34+ return actual . getBillingPeriodUsageCost ( ...args )
35+ } ,
36+ }
37+ } )
2438
2539import {
2640 REPORTING_USAGE_CACHE_TTL_MS ,
2741 readSoftGateUsageCost,
2842} from '@/lib/billing/core/reporting-usage-cache'
29- import type { BillingEntity , UsageQueryPeriod } from '@/lib/billing/core/usage-log'
43+ import { invalidateReportingUsagePrefixes } from '@/lib/billing/core/reporting-usage-epoch'
44+ import {
45+ type BillingEntity ,
46+ type RecordCumulativeUsageParams ,
47+ recordCumulativeUsage ,
48+ type UsageQueryPeriod ,
49+ } from '@/lib/billing/core/usage-log'
3050
3151const schemaName = `reporting_usage_${ generateId ( ) . replaceAll ( '-' , '' ) } `
3252const connection = postgres ( databaseUrl , {
33- max : 2 ,
53+ max : 4 ,
3454 prepare : false ,
3555 connection : { search_path : schemaName } ,
3656 onnotice : ( ) => undefined ,
@@ -43,28 +63,128 @@ const REPORTING: UsageQueryPeriod = {
4363 source : 'reporting' ,
4464}
4565
46- function sharedKey ( payer : BillingEntity ) : string {
47- return `usage:reporting:v1:${ payer . type } :${ payer . id } :reporting:${ REPORTING . start . toISOString ( ) } :${ REPORTING . end . toISOString ( ) } `
66+ const DAY_MS = 86_400_000
67+ const MINUTE_MS = 60_000
68+
69+ /** A year-long window that started 100 days ago, so the clock is always inside it. */
70+ function openReportingWindow ( ) : UsageQueryPeriod {
71+ const start = Math . floor ( ( Date . now ( ) - 100 * DAY_MS ) / DAY_MS ) * DAY_MS
72+ return { start : new Date ( start ) , end : new Date ( start + 365 * DAY_MS ) , source : 'reporting' }
73+ }
74+
75+ function windowKey ( payer : BillingEntity , period : UsageQueryPeriod ) : string {
76+ return `${ payer . type } :${ payer . id } :reporting:${ period . start . toISOString ( ) } :${ period . end . toISOString ( ) } `
77+ }
78+
79+ function sharedKey ( payer : BillingEntity , period : UsageQueryPeriod = REPORTING ) : string {
80+ return `usage:reporting:v1:${ windowKey ( payer , period ) } `
81+ }
82+
83+ function prefixKey ( payer : BillingEntity , period : UsageQueryPeriod ) : string {
84+ return `usage:reporting:v1:prefix:${ windowKey ( payer , period ) } `
85+ }
86+
87+ function prefixLeaseKey ( payer : BillingEntity , period : UsageQueryPeriod ) : string {
88+ return `usage:reporting:v1:prefix-lease:${ windowKey ( payer , period ) } `
89+ }
90+
91+ function payerEpochKey ( payer : BillingEntity ) : string {
92+ return `usage:reporting:v1:prefix-epoch:${ payer . type } :${ payer . id } `
93+ }
94+
95+ type ReportingUsageCacheModule = typeof import ( '@/lib/billing/core/reporting-usage-cache' )
96+
97+ /** Fresh module instances, so each in-process cache is cold, as in a new process. */
98+ async function newProcesses ( count : number ) : Promise < ReportingUsageCacheModule [ ] > {
99+ const processes : ReportingUsageCacheModule [ ] = [ ]
100+ for ( let index = 0 ; index < count ; index ++ ) {
101+ vi . resetModules ( )
102+ processes . push ( await import ( '@/lib/billing/core/reporting-usage-cache' ) )
103+ }
104+ return processes
105+ }
106+
107+ async function readInNewProcess ( payer : BillingEntity , period : UsageQueryPeriod ) : Promise < number > {
108+ const [ cache ] = await newProcesses ( 1 )
109+ return cache . readSoftGateUsageCost ( payer , period )
110+ }
111+
112+ /** Sums that start at the window's start: a whole-window sum or a prefix, never a tail. */
113+ function windowStartSums ( period : UsageQueryPeriod ) : UsageQueryPeriod [ ] {
114+ return summedPeriods . filter ( ( summed ) => summed . start . getTime ( ) === period . start . getTime ( ) )
115+ }
116+
117+ async function insertUsage ( payer : BillingEntity , cost : number , createdAt : Date ) : Promise < string > {
118+ const id = generateId ( )
119+ await connection `INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at)
120+ VALUES (${ id } , ${ cost } , ${ payer . type } , ${ payer . id } , ${ createdAt . toISOString ( ) } )`
121+ return id
122+ }
123+
124+ function cumulativeCharge (
125+ payer : BillingEntity ,
126+ eventKey : string ,
127+ cost : number
128+ ) : RecordCumulativeUsageParams {
129+ return {
130+ userId : 'actor' ,
131+ workspaceId : 'workspace' ,
132+ billingEntity : payer ,
133+ billingPeriod : { start : REPORTING . start , end : REPORTING . end } ,
134+ source : 'workspace-chat' ,
135+ model : 'test-model' ,
136+ eventKey,
137+ cost,
138+ }
139+ }
140+
141+ /** A cumulative charge whose row was first written at `createdAt`. */
142+ async function startCumulativeCharge (
143+ payer : BillingEntity ,
144+ cost : number ,
145+ createdAt : Date
146+ ) : Promise < string > {
147+ const eventKey = `update-cost:${ generateId ( ) } `
148+ await recordCumulativeUsage ( cumulativeCharge ( payer , eventKey , cost ) )
149+ await connection `UPDATE usage_log SET created_at = ${ createdAt . toISOString ( ) }
150+ WHERE event_key = ${ eventKey } `
151+ return eventKey
152+ }
153+
154+ async function exactSum ( payer : BillingEntity , period : UsageQueryPeriod ) : Promise < number > {
155+ const [ row ] = await connection < { cost : string } [ ] > `SELECT COALESCE(SUM(cost), 0) AS cost
156+ FROM usage_log WHERE billing_entity_type = ${ payer . type } AND billing_entity_id = ${ payer . id }
157+ AND created_at >= ${ period . start . toISOString ( ) } AND created_at < ${ period . end . toISOString ( ) } `
158+ return Number ( row ?. cost ?? 0 )
48159}
49160
50161describe . runIf ( Boolean ( redisUrl ) ) ( 'shared reporting usage read' , ( ) => {
51162 let redis : Redis
52163 let payer : BillingEntity
53164
165+ async function prefixStored ( payer : BillingEntity , period : UsageQueryPeriod ) : Promise < void > {
166+ await vi . waitFor ( async ( ) => expect ( await redis . exists ( prefixKey ( payer , period ) ) ) . toBe ( 1 ) )
167+ }
168+
54169 beforeAll ( async ( ) => {
55170 redis = new Redis ( redisUrl ! , { lazyConnect : true , maxRetriesPerRequest : 0 } )
56171 await redis . connect ( )
57172 await connection . unsafe ( `CREATE SCHEMA "${ schemaName } "` )
58173 await connection . unsafe ( `CREATE TABLE usage_log (
59- id text PRIMARY KEY, cost numeric NOT NULL, billing_entity_type text,
174+ id text PRIMARY KEY, user_id text, category text, source text, description text,
175+ metadata jsonb, cost numeric NOT NULL, event_key text, billing_entity_type text,
60176 billing_entity_id text, billing_period_start timestamp, billing_period_end timestamp,
61- created_at timestamp NOT NULL
62- )` )
177+ workspace_id text, workflow_id text, execution_id text,
178+ created_at timestamp NOT NULL DEFAULT now()
179+ );
180+ CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key)
181+ WHERE event_key IS NOT NULL` )
63182 transaction . mockImplementation ( ( callback ) => database . transaction ( callback ) )
64183 } )
65184
66185 beforeEach ( async ( ) => {
67186 transaction . mockClear ( )
187+ summedPeriods . length = 0
68188 redisConfigMockFns . mockGetRedisClient . mockReturnValue ( redis )
69189 payer = { type : 'organization' , id : generateId ( ) }
70190 await connection `INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at)
@@ -75,7 +195,15 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
75195
76196 afterEach ( async ( ) => {
77197 vi . useRealTimers ( )
78- await redis . del ( sharedKey ( payer ) )
198+ await redis . del (
199+ ...[ REPORTING , openReportingWindow ( ) ] . flatMap ( ( period ) => [
200+ sharedKey ( payer , period ) ,
201+ prefixKey ( payer , period ) ,
202+ prefixLeaseKey ( payer , period ) ,
203+ ] ) ,
204+ payerEpochKey ( payer ) ,
205+ 'usage:reporting:v1:prefix-epoch'
206+ )
79207 } )
80208
81209 afterAll ( async ( ) => {
@@ -149,6 +277,170 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
149277 expect ( await redis . get ( sharedKey ( payer ) ) ) . toBe ( '9.99' )
150278 } )
151279
280+ it ( 'sums the whole window in one query when Redis is absent' , async ( ) => {
281+ redisConfigMockFns . mockGetRedisClient . mockReturnValue ( null )
282+ const period = openReportingWindow ( )
283+ await insertUsage ( payer , 2 , new Date ( period . start . getTime ( ) + DAY_MS ) )
284+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
285+
286+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 5 )
287+ expect ( summedPeriods ) . toEqual ( [ period ] )
288+ } )
289+
290+ it ( 'leases a due refresh to one of the processes missing together' , async ( ) => {
291+ const period = openReportingWindow ( )
292+ await insertUsage ( payer , 2 , new Date ( period . start . getTime ( ) + DAY_MS ) )
293+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
294+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 5 )
295+ await insertUsage ( payer , 4 , new Date ( period . start . getTime ( ) + 2 * DAY_MS ) )
296+
297+ const now = Date . now ( )
298+ vi . useFakeTimers ( { toFake : [ 'Date' ] } )
299+ vi . setSystemTime ( now + 4 * MINUTE_MS )
300+ await redis . del ( sharedKey ( payer , period ) , prefixLeaseKey ( payer , period ) )
301+ const processes = await newProcesses ( 5 )
302+ const costs = await Promise . all (
303+ processes . map ( ( cache ) => cache . readSoftGateUsageCost ( payer , period ) )
304+ )
305+
306+ expect ( windowStartSums ( period ) ) . toHaveLength ( 2 )
307+ expect ( costs ) . toContain ( 9 )
308+ for ( const cost of costs ) expect ( [ 5 , 9 ] ) . toContain ( cost )
309+ } )
310+
311+ it ( 'sums the whole window at once when no prefix is stored and another process is summing one' , async ( ) => {
312+ const period = openReportingWindow ( )
313+ await insertUsage ( payer , 2 , new Date ( period . start . getTime ( ) + DAY_MS ) )
314+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
315+ await redis . set ( prefixLeaseKey ( payer , period ) , '1' , 'PX' , MINUTE_MS )
316+ const [ cache ] = await newProcesses ( 1 )
317+
318+ const startedAt = performance . now ( )
319+ await expect ( cache . readSoftGateUsageCost ( payer , period ) ) . resolves . toBe ( 5 )
320+ expect ( performance . now ( ) - startedAt ) . toBeLessThan ( 1_000 )
321+ expect ( summedPeriods ) . toEqual ( [ period ] )
322+ } )
323+
324+ it ( 'never reads above the ledger as it grows, and matches it once the prefix is refreshed' , async ( ) => {
325+ const period = openReportingWindow ( )
326+ const oldRows = [
327+ await insertUsage ( payer , 1.25 , new Date ( period . start . getTime ( ) + DAY_MS ) ) ,
328+ await insertUsage ( payer , 0.5 , new Date ( Date . now ( ) - 10 * DAY_MS ) ) ,
329+ ]
330+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 1.75 )
331+
332+ for ( let round = 0 ; round < 8 ; round ++ ) {
333+ const recent = new Date ( Date . now ( ) - randomInt ( 0 , 4 * MINUTE_MS ) )
334+ await insertUsage ( payer , randomInt ( 1 , 500 ) / 100 , recent )
335+ if ( round % 2 === 0 ) {
336+ const late = new Date ( period . start . getTime ( ) + randomInt ( 1 , 90 ) * DAY_MS )
337+ await insertUsage ( payer , randomInt ( 1 , 500 ) / 100 , late )
338+ } else {
339+ await connection `UPDATE usage_log SET cost = cost + ${ randomInt ( 1 , 500 ) / 100 }
340+ WHERE id = ${ randomItem ( oldRows ) } `
341+ }
342+ await redis . del ( sharedKey ( payer , period ) )
343+
344+ const served = await readInNewProcess ( payer , period )
345+ expect ( served ) . toBeLessThanOrEqual ( ( await exactSum ( payer , period ) ) + 1e-9 )
346+ }
347+ expect ( windowStartSums ( period ) ) . toHaveLength ( 1 )
348+
349+ const now = Date . now ( )
350+ vi . useFakeTimers ( { toFake : [ 'Date' ] } )
351+ vi . setSystemTime ( now + 4 * MINUTE_MS )
352+ await redis . del ( sharedKey ( payer , period ) , prefixLeaseKey ( payer , period ) )
353+
354+ const refreshed = await readInNewProcess ( payer , period )
355+ expect ( refreshed ) . toBeCloseTo ( await exactSum ( payer , period ) , 9 )
356+ expect ( windowStartSums ( period ) ) . toHaveLength ( 2 )
357+ } )
358+
359+ it ( 'never reads above the ledger after ledger rows are deleted and prefixes are invalidated' , async ( ) => {
360+ const period = openReportingWindow ( )
361+ await insertUsage ( payer , 2 , new Date ( period . start . getTime ( ) + DAY_MS ) )
362+ const departedRow = await insertUsage ( payer , 6 , new Date ( period . start . getTime ( ) + 2 * DAY_MS ) )
363+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
364+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 11 )
365+ await prefixStored ( payer , period )
366+
367+ await connection `DELETE FROM usage_log WHERE id = ${ departedRow } `
368+ await invalidateReportingUsagePrefixes ( )
369+ await redis . del ( sharedKey ( payer , period ) )
370+
371+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( await exactSum ( payer , period ) )
372+ } )
373+
374+ it ( 'sees a top-up to a row inside the stored prefix on the next read' , async ( ) => {
375+ const period = openReportingWindow ( )
376+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
377+ const eventKey = await startCumulativeCharge (
378+ payer ,
379+ 2 ,
380+ new Date ( period . start . getTime ( ) + DAY_MS )
381+ )
382+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 5 )
383+ await prefixStored ( payer , period )
384+
385+ await recordCumulativeUsage ( cumulativeCharge ( payer , eventKey , 6 ) )
386+ await vi . waitFor ( async ( ) => expect ( await redis . exists ( payerEpochKey ( payer ) ) ) . toBe ( 1 ) )
387+ await redis . del ( sharedKey ( payer , period ) )
388+
389+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 9 )
390+ expect ( await exactSum ( payer , period ) ) . toBe ( 9 )
391+ } )
392+
393+ it ( 'leaves the prefix in place for a top-up to a row still in the tail, and reads it exactly' , async ( ) => {
394+ const period = openReportingWindow ( )
395+ await insertUsage ( payer , 2 , new Date ( period . start . getTime ( ) + DAY_MS ) )
396+ const eventKey = await startCumulativeCharge ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
397+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 5 )
398+ await prefixStored ( payer , period )
399+
400+ await recordCumulativeUsage ( cumulativeCharge ( payer , eventKey , 7 ) )
401+ await redis . del ( sharedKey ( payer , period ) )
402+
403+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 9 )
404+ expect ( await redis . exists ( payerEpochKey ( payer ) ) ) . toBe ( 0 )
405+ expect ( windowStartSums ( period ) ) . toHaveLength ( 1 )
406+ } )
407+
408+ it ( 'never over-counts a top-up that commits while a prefix refresh sums' , async ( ) => {
409+ const period = openReportingWindow ( )
410+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
411+ const eventKey = await startCumulativeCharge (
412+ payer ,
413+ 2 ,
414+ new Date ( period . start . getTime ( ) + DAY_MS )
415+ )
416+ const [ refreshing ] = await newProcesses ( 1 )
417+ transaction . mockImplementationOnce ( async ( callback ) => {
418+ const result = await database . transaction ( callback )
419+ await recordCumulativeUsage ( cumulativeCharge ( payer , eventKey , 6 ) )
420+ return result
421+ } )
422+
423+ const served = await refreshing . readSoftGateUsageCost ( payer , period )
424+ expect ( served ) . toBeLessThanOrEqual ( await exactSum ( payer , period ) )
425+ await prefixStored ( payer , period )
426+ await vi . waitFor ( async ( ) => expect ( await redis . exists ( payerEpochKey ( payer ) ) ) . toBe ( 1 ) )
427+ await redis . del ( sharedKey ( payer , period ) )
428+
429+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( await exactSum ( payer , period ) )
430+ } )
431+
432+ it ( 'serves a stale prefix while another process holds the refresh lease' , async ( ) => {
433+ const period = openReportingWindow ( )
434+ await insertUsage ( payer , 2 , new Date ( period . start . getTime ( ) + DAY_MS ) )
435+ await insertUsage ( payer , 3 , new Date ( Date . now ( ) - MINUTE_MS ) )
436+ const staleCutoff = Date . now ( ) - 10 * MINUTE_MS
437+ await redis . set ( prefixKey ( payer , period ) , `::${ staleCutoff } :7` , 'PX' , MINUTE_MS )
438+ await redis . set ( prefixLeaseKey ( payer , period ) , '1' , 'PX' , MINUTE_MS )
439+
440+ await expect ( readInNewProcess ( payer , period ) ) . resolves . toBe ( 10 )
441+ expect ( windowStartSums ( period ) ) . toHaveLength ( 0 )
442+ } )
443+
152444 it ( 'sums the ledger when the Redis client cannot be built' , async ( ) => {
153445 redisConfigMockFns . mockGetRedisClient . mockImplementation ( ( ) => {
154446 throw new Error ( 'Invalid Redis configuration' )
0 commit comments