11/**
22 * Cost callbacks against real PostgreSQL: a direct-v1 run that outlives its admitted Stripe period
33 * records its later spend in the payer's current period, so the closed period is never topped up
4- * after its invoice, and spend after the payer's terminal settlement is refused. Only the
5- * internal-key check is stubbed.
4+ * after its invoice, and spend after the payer's terminal settlement is refused. A free payer's
5+ * callback that takes it past its limit pauses the run even while an earlier admission is
6+ * cached, and a paid payer keeps the cached admission. Only the internal-key check is stubbed.
67 */
78import { db } from '@sim/db'
8- import { subscription , usageLog , user , userStats } from '@sim/db/schema'
9+ import { project , subscription , usageLog , user , userStats , workspace } from '@sim/db/schema'
910import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock'
1011import { generateId } from '@sim/utils/id'
1112import { eq , inArray } from 'drizzle-orm'
@@ -24,8 +25,12 @@ vi.mock('@/lib/mothership/request/http', async (importOriginal) => ({
2425
2526import {
2627 BILLING_ACCOUNT_DECISION_HEADER ,
28+ BILLING_ATTRIBUTION_HEADER ,
29+ type BillingAttributionSnapshot ,
2730 serializeAccountBillingDecisionHeader ,
31+ serializeBillingAttributionHeader ,
2832} from '@/lib/billing/core/attribution'
33+ import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
2934import { claimTerminalPeriod } from '@/lib/billing/cycle-close'
3035import { POST } from '@/app/api/billing/update-cost/route'
3136
@@ -35,6 +40,8 @@ const userIds: string[] = []
3540afterAll ( async ( ) => {
3641 if ( userIds . length === 0 ) return
3742 await db . delete ( usageLog ) . where ( inArray ( usageLog . userId , userIds ) )
43+ await db . delete ( workspace ) . where ( inArray ( workspace . ownerId , userIds ) )
44+ await db . delete ( project ) . where ( inArray ( project . ownerId , userIds ) )
3845 await db . delete ( subscription ) . where ( inArray ( subscription . referenceId , userIds ) )
3946 await db . delete ( userStats ) . where ( inArray ( userStats . userId , userIds ) )
4047 await db . delete ( user ) . where ( inArray ( user . id , userIds ) )
@@ -47,10 +54,9 @@ interface Payer {
4754 decision : string
4855}
4956
50- /** A user on a pro subscription for `period`, whose close marker has caught up to it . */
51- async function createPayer ( period : { start : Date ; end : Date } ) : Promise < Payer > {
57+ /** A user with the default $5 usage limit and no subscription . */
58+ async function createUser ( ) : Promise < string > {
5259 const userId = `update-cost-user-${ generateId ( ) } `
53- const subscriptionId = generateId ( )
5460 userIds . push ( userId )
5561 await db . insert ( user ) . values ( {
5662 id : userId ,
@@ -61,6 +67,13 @@ async function createPayer(period: { start: Date; end: Date }): Promise<Payer> {
6167 updatedAt : new Date ( ) ,
6268 } )
6369 await db . insert ( userStats ) . values ( { id : generateId ( ) , userId } )
70+ return userId
71+ }
72+
73+ /** A user on a pro subscription for `period`, whose close marker has caught up to it. */
74+ async function createPayer ( period : { start : Date ; end : Date } ) : Promise < Payer > {
75+ const userId = await createUser ( )
76+ const subscriptionId = generateId ( )
6477 await db . insert ( subscription ) . values ( {
6578 id : subscriptionId ,
6679 plan : 'pro' ,
@@ -94,6 +107,21 @@ function requestRows(requestKey: string) {
94107 . where ( inArray ( usageLog . eventKey , [ `update-cost:${ requestKey } ` , `update-cost:${ requestKey } @1` ] ) )
95108}
96109
110+ /** A free user's direct-v1 decision, admitted in the open default period as admission mints it. */
111+ async function createFreePayer ( ) : Promise < Payer > {
112+ const userId = await createUser ( )
113+ const decision = serializeAccountBillingDecisionHeader ( {
114+ userId,
115+ billingEntity : { type : 'user' , id : userId } ,
116+ billingPeriod : {
117+ start : defaultBillingPeriod ( ) . start . toISOString ( ) ,
118+ end : defaultBillingPeriod ( ) . end . toISOString ( ) ,
119+ source : 'default' ,
120+ } ,
121+ } )
122+ return { userId, subscriptionId : '' , decision }
123+ }
124+
97125function callback ( payer : Payer , requestKey : string , cost : number ) : NextRequest {
98126 return new NextRequest ( 'http://localhost:3000/api/billing/update-cost' , {
99127 method : 'POST' ,
@@ -114,6 +142,65 @@ function callback(payer: Payer, requestKey: string, cost: number): NextRequest {
114142 } )
115143}
116144
145+ /** A free user's personal workspace and the attribution its runs are admitted with. */
146+ async function createFreeWorkspacePayer ( ) : Promise < BillingAttributionSnapshot > {
147+ const userId = await createUser ( )
148+ const projectId = generateId ( )
149+ const workspaceId = generateId ( )
150+ await db . insert ( project ) . values ( { id : projectId , name : 'Update Cost Test' , ownerId : userId } )
151+ await db . insert ( workspace ) . values ( {
152+ id : workspaceId ,
153+ projectId,
154+ name : 'Update Cost Test' ,
155+ ownerId : userId ,
156+ billedAccountUserId : userId ,
157+ } )
158+ return {
159+ actorUserId : userId ,
160+ workspaceId,
161+ organizationId : null ,
162+ billedAccountUserId : userId ,
163+ billingEntity : { type : 'user' , id : userId } ,
164+ billingPeriod : {
165+ start : defaultBillingPeriod ( ) . start . toISOString ( ) ,
166+ end : defaultBillingPeriod ( ) . end . toISOString ( ) ,
167+ source : 'default' ,
168+ } ,
169+ payerSubscription : null ,
170+ }
171+ }
172+
173+ function attributedCallback (
174+ attribution : BillingAttributionSnapshot ,
175+ requestKey : string ,
176+ cost : number
177+ ) : NextRequest {
178+ return new NextRequest ( 'http://localhost:3000/api/billing/update-cost' , {
179+ method : 'POST' ,
180+ headers : {
181+ 'content-type' : 'application/json' ,
182+ 'x-api-key' : 'internal' ,
183+ 'x-sim-billing-protocol' : 'attribution-v1' ,
184+ 'x-sim-billing-request-id' : requestKey ,
185+ [ BILLING_ATTRIBUTION_HEADER ] : serializeBillingAttributionHeader ( attribution ) ,
186+ } ,
187+ body : JSON . stringify ( {
188+ userId : attribution . actorUserId ,
189+ workspaceId : attribution . workspaceId ,
190+ cost,
191+ model : 'test-model' ,
192+ source : 'copilot' ,
193+ idempotencyKey : requestKey ,
194+ } ) ,
195+ } )
196+ }
197+
198+ async function usageExceeded ( response : ReturnType < typeof POST > ) : Promise < boolean > {
199+ const res = await response
200+ expect ( res . status ) . toBe ( 200 )
201+ return ( await res . json ( ) ) . usageExceeded
202+ }
203+
117204describe ( 'direct-v1 cost callbacks in PostgreSQL' , ( ) => {
118205 it ( 'records spend after a Stripe rollover in the payer current period' , async ( ) => {
119206 const now = Date . now ( )
@@ -157,3 +244,38 @@ describe('direct-v1 cost callbacks in PostgreSQL', () => {
157244 expect ( ( await requestRows ( requestKey ) ) . map ( ( row ) => Number ( row . cost ) ) ) . toEqual ( [ 0.5 ] )
158245 } )
159246} )
247+
248+ /**
249+ * Each payer's first callback caches an admission ($1 of a $5 limit), and its second, inside the
250+ * near-limit cache window, records spend that takes it to $6.
251+ */
252+ describe ( 'usage standing after a charge that crosses the limit' , ( ) => {
253+ it ( 'pauses a free direct-v1 run although its admission is cached' , async ( ) => {
254+ const payer = await createFreePayer ( )
255+
256+ expect ( await usageExceeded ( POST ( callback ( payer , generateId ( ) , 1 ) , { } ) ) ) . toBe ( false )
257+ expect ( await usageExceeded ( POST ( callback ( payer , generateId ( ) , 5 ) , { } ) ) ) . toBe ( true )
258+ } )
259+
260+ it ( 'pauses a free attributed run although its admission is cached' , async ( ) => {
261+ const attribution = await createFreeWorkspacePayer ( )
262+
263+ expect ( await usageExceeded ( POST ( attributedCallback ( attribution , generateId ( ) , 1 ) , { } ) ) ) . toBe (
264+ false
265+ )
266+ expect ( await usageExceeded ( POST ( attributedCallback ( attribution , generateId ( ) , 5 ) , { } ) ) ) . toBe (
267+ true
268+ )
269+ } )
270+
271+ it ( "keeps serving a paid payer's cached admission" , async ( ) => {
272+ const now = Date . now ( )
273+ const payer = await createPayer ( {
274+ start : new Date ( now - 10 * DAY_MS ) ,
275+ end : new Date ( now + 20 * DAY_MS ) ,
276+ } )
277+
278+ expect ( await usageExceeded ( POST ( callback ( payer , generateId ( ) , 1 ) , { } ) ) ) . toBe ( false )
279+ expect ( await usageExceeded ( POST ( callback ( payer , generateId ( ) , 5 ) , { } ) ) ) . toBe ( false )
280+ } )
281+ } )
0 commit comments