Skip to content

Commit 9e2afe0

Browse files
authored
fix(billing): hold free plans to their allowance mid-run (#9462)
* fix(billing): hold free plans to their allowance mid-run * fix(billing): read free payers fresh in the provider check and judge free by the admitted plan * fix(billing): bound free payers' fresh usage reads and shrink the test-types baseline
1 parent 584848a commit 9e2afe0

14 files changed

Lines changed: 468 additions & 110 deletions

File tree

‎apps/docs/content/docs/platform/costs.mdx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,7 @@ Execution rate buckets can burst above the sustained per-minute rate. General AP
249249

250250
| Plan | Concurrent runs |
251251
| --- | --- |
252-
| Free | 10 |
252+
| Free | 2 |
253253
| Pro | 50 |
254254
| Max | 200 |
255255
| Enterprise default | 1,000; customizable |

‎apps/sim/app/api/billing/update-cost/route.integration.ts‎

Lines changed: 128 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
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
*/
78
import { 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'
910
import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock'
1011
import { generateId } from '@sim/utils/id'
1112
import { eq, inArray } from 'drizzle-orm'
@@ -24,8 +25,12 @@ vi.mock('@/lib/mothership/request/http', async (importOriginal) => ({
2425

2526
import {
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'
2934
import { claimTerminalPeriod } from '@/lib/billing/cycle-close'
3035
import { POST } from '@/app/api/billing/update-cost/route'
3136

@@ -35,6 +40,8 @@ const userIds: string[] = []
3540
afterAll(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+
97125
function 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+
117204
describe('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+
})

‎apps/sim/app/api/billing/update-cost/route.test.ts‎

Lines changed: 83 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -931,6 +931,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
931931
source: 'stripe' as const,
932932
},
933933
}
934+
const PAID_SUBSCRIPTION = { id: 'sub-1', plan: 'team', status: 'active', seats: 4 }
935+
const PAID_ATTRIBUTION = { ...CURRENT_ATTRIBUTION, payerSubscription: PAID_SUBSCRIPTION }
934936

935937
function attributedCallback() {
936938
callbackSequence += 1
@@ -1005,7 +1007,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
10051007
})
10061008
})
10071009

1008-
it('serves a cached admission to every step and re-reads a refusal', async () => {
1010+
it('serves a paid payer a cached admission to every step and re-reads a refusal', async () => {
1011+
mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
10091012
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false })
10101013

10111014
expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(false)
@@ -1022,6 +1025,17 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
10221025
expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(false)
10231026
})
10241027

1028+
it('reads a free payer standing past its cached admission after every charge', async () => {
1029+
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false })
1030+
expect((await (await POST(attributedCallback(), {})).json()).usageExceeded).toBe(false)
1031+
1032+
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
1033+
await expect((await POST(attributedCallback(), {})).json()).resolves.toMatchObject({
1034+
usageExceeded: true,
1035+
usageUpgrade: { action: 'upgrade_plan' },
1036+
})
1037+
})
1038+
10251039
it('answers a duplicate retry with the verdict its lost first answer carried', async () => {
10261040
mockRecordCumulativeUsage.mockResolvedValue({ billed: false, delta: 0, total: 0.5 })
10271041
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
@@ -1209,7 +1223,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
12091223
it.each([
12101224
['an unreadable ledger', { isExceeded: true, reason: 'usage_unavailable' }],
12111225
['a blocked account', { isExceeded: true, reason: 'billing_blocked', scope: 'payer' }],
1212-
])('does not pause a run for %s', async (_case, verdict) => {
1226+
])('does not pause a paid run for %s', async (_case, verdict) => {
1227+
mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
12131228
mockCheckAttributedUsageLimits.mockResolvedValue(verdict)
12141229

12151230
const body = await (await POST(attributedCallback())).json()
@@ -1218,6 +1233,42 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
12181233
expect(body).not.toHaveProperty('usageUpgrade')
12191234
})
12201235

1236+
it.each([
1237+
[
1238+
'an unreadable ledger',
1239+
() =>
1240+
mockCheckAttributedUsageLimits.mockResolvedValue({
1241+
isExceeded: true,
1242+
reason: 'usage_unavailable',
1243+
}),
1244+
],
1245+
[
1246+
'a failed gate read',
1247+
() => mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out')),
1248+
],
1249+
])('pauses a free run for %s, keeping the recorded charge', async (_case, answer) => {
1250+
answer()
1251+
1252+
const res = await POST(attributedCallback(), {})
1253+
1254+
expect(res.status).toBe(200)
1255+
await expect(res.json()).resolves.toMatchObject({
1256+
success: true,
1257+
usageExceeded: true,
1258+
usageUpgrade: { action: 'upgrade_plan' },
1259+
})
1260+
})
1261+
1262+
it('does not pause a free run for a blocked account', async () => {
1263+
mockCheckAttributedUsageLimits.mockResolvedValue({
1264+
isExceeded: true,
1265+
reason: 'billing_blocked',
1266+
scope: 'payer',
1267+
})
1268+
1269+
expect((await (await POST(attributedCallback(), {})).json()).usageExceeded).toBe(false)
1270+
})
1271+
12211272
it('tells a member over the cap their organization set who can raise it', async () => {
12221273
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'member' })
12231274

@@ -1284,8 +1335,11 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
12841335
expect(body.usageExceeded).toBe(true)
12851336
})
12861337

1287-
it('keeps a run going when its current period cannot be read', async () => {
1288-
mockRequireBillingAttributionHeader.mockReturnValue(STRIPE_ATTRIBUTION)
1338+
it('keeps a paid run going when its current period cannot be read', async () => {
1339+
mockRequireBillingAttributionHeader.mockReturnValue({
1340+
...STRIPE_ATTRIBUTION,
1341+
payerSubscription: PAID_SUBSCRIPTION,
1342+
})
12891343
mockRefreshAttributionPeriod.mockRejectedValue(new Error('subscription read timed out'))
12901344
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
12911345

@@ -1304,31 +1358,39 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
13041358
vi.useRealTimers()
13051359
})
13061360

1307-
it('answers not exceeded when the standing read outlasts the callback budget', async () => {
1308-
mockCheckAttributedUsageLimits.mockImplementation(async () => {
1309-
await sleep(1500)
1310-
return { isExceeded: true, scope: 'payer' }
1311-
})
1361+
it.each([
1362+
['a paid', PAID_ATTRIBUTION, false],
1363+
['a free', CURRENT_ATTRIBUTION, true],
1364+
])(
1365+
'answers %s payer whose standing read outlasts the callback budget',
1366+
async (_case, attribution, exceeded) => {
1367+
mockRequireBillingAttributionHeader.mockReturnValue(attribution)
1368+
mockCheckAttributedUsageLimits.mockImplementation(async () => {
1369+
await sleep(1500)
1370+
return { isExceeded: false }
1371+
})
13121372

1313-
const pending = POST(attributedCallback())
1314-
await vi.advanceTimersByTimeAsync(1000)
1315-
const res = await pending
1373+
const pending = POST(attributedCallback(), {})
1374+
await vi.advanceTimersByTimeAsync(1000)
1375+
const res = await pending
13161376

1317-
expect(res.status).toBe(200)
1318-
await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: false })
1319-
// Settle the abandoned read so no later test coalesces onto it.
1320-
await vi.advanceTimersByTimeAsync(500)
1321-
})
1377+
expect(res.status).toBe(200)
1378+
await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: exceeded })
1379+
// Settle the abandoned read so no later test coalesces onto it.
1380+
await vi.advanceTimersByTimeAsync(500)
1381+
}
1382+
)
13221383

1323-
it('does not pause a run on a verdict read across the end of its period', async () => {
1384+
it('does not pause a paid run on a verdict read across the end of its period', async () => {
13241385
const straddling = {
1325-
...CURRENT_ATTRIBUTION,
1386+
...PAID_ATTRIBUTION,
13261387
billingPeriod: {
13271388
start: '2026-07-01T00:00:00.000Z',
13281389
end: new Date(Date.now() + 40).toISOString(),
13291390
source: 'stripe' as const,
13301391
},
13311392
}
1393+
mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
13321394
mockRefreshAttributionPeriod.mockResolvedValue(straddling)
13331395
mockCheckAttributedUsageLimits.mockImplementation(async () => {
13341396
await sleep(80)
@@ -1363,7 +1425,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
13631425
})
13641426
})
13651427

1366-
it('keeps a recorded charge successful when the gate read fails', async () => {
1428+
it('keeps a paid recorded charge successful and unpaused when the gate read fails', async () => {
1429+
mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
13671430
mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out'))
13681431

13691432
const res = await POST(attributedCallback())

0 commit comments

Comments
 (0)