Skip to content

Commit 8af8bf8

Browse files
waleedlatif1waleedlatif1
committed
fix(billing): refuse charges into a period the terminal settlement already invoiced
A subscription's deletion settles its terminal period at once (claim, overage, final invoice, bookkeeping), but recordCumulativeUsage kept topping up that period's row for a still-running request because the subscription's period never rolls after deletion. That spend was never billed. The terminal claim now advances the close marker to the period's end under its FOR UPDATE lock, and recordCumulativeUsage reads the marker in its existing FOR SHARE read: a charge whose target period ends at or before the marker throws CumulativeUsagePeriodClosedError, which update-cost answers with the existing non-retryable BILLING_PERIOD_ELAPSED outcome, so the worker quarantines the leg for reconciliation instead of the spend disappearing into an invoiced period. No migration: the marker is an existing column, and every reader treats a marker at or past periodStart as current, so v0.9.6 behaves unchanged.
1 parent fa32bc1 commit 8af8bf8

9 files changed

Lines changed: 190 additions & 70 deletions

File tree

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

Lines changed: 87 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
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. Only the internal-key check is stubbed.
4+
* after its invoice, and spend after the payer's terminal settlement is refused. Only the
5+
* internal-key check is stubbed.
56
*/
67
import { db } from '@sim/db'
78
import { subscription, usageLog, user, userStats } from '@sim/db/schema'
89
import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock'
910
import { generateId } from '@sim/utils/id'
10-
import { eq } from 'drizzle-orm'
11+
import { eq, inArray } from 'drizzle-orm'
1112
import { NextRequest } from 'next/server'
1213
import { afterAll, describe, expect, it, vi } from 'vitest'
1314

@@ -25,31 +26,86 @@ import {
2526
BILLING_ACCOUNT_DECISION_HEADER,
2627
serializeAccountBillingDecisionHeader,
2728
} from '@/lib/billing/core/billing-attribution'
29+
import { claimTerminalPeriod } from '@/lib/billing/cycle-close'
2830
import { POST } from '@/app/api/billing/update-cost/route'
2931

3032
const DAY_MS = 24 * 60 * 60 * 1000
31-
const userId = `update-cost-user-${generateId()}`
32-
const subscriptionId = generateId()
33+
const userIds: string[] = []
3334

3435
afterAll(async () => {
35-
await db.delete(usageLog).where(eq(usageLog.userId, userId))
36-
await db.delete(subscription).where(eq(subscription.id, subscriptionId))
37-
await db.delete(userStats).where(eq(userStats.userId, userId))
38-
await db.delete(user).where(eq(user.id, userId))
36+
if (userIds.length === 0) return
37+
await db.delete(usageLog).where(inArray(usageLog.userId, userIds))
38+
await db.delete(subscription).where(inArray(subscription.referenceId, userIds))
39+
await db.delete(userStats).where(inArray(userStats.userId, userIds))
40+
await db.delete(user).where(inArray(user.id, userIds))
3941
})
4042

41-
function callback(requestKey: string, cost: number, decision: string): NextRequest {
43+
interface Payer {
44+
userId: string
45+
subscriptionId: string
46+
/** The direct-v1 decision of a run admitted in the subscription's period. */
47+
decision: string
48+
}
49+
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> {
52+
const userId = `update-cost-user-${generateId()}`
53+
const subscriptionId = generateId()
54+
userIds.push(userId)
55+
await db.insert(user).values({
56+
id: userId,
57+
name: 'Update Cost Test',
58+
email: `${userId}@update-cost.test`,
59+
emailVerified: true,
60+
createdAt: new Date(),
61+
updatedAt: new Date(),
62+
})
63+
await db.insert(userStats).values({ id: generateId(), userId })
64+
await db.insert(subscription).values({
65+
id: subscriptionId,
66+
plan: 'pro',
67+
referenceId: userId,
68+
status: 'active',
69+
periodStart: period.start,
70+
periodEnd: period.end,
71+
lastClosedPeriodStart: period.start,
72+
})
73+
const decision = serializeAccountBillingDecisionHeader({
74+
userId,
75+
billingEntity: { type: 'user', id: userId },
76+
billingPeriod: {
77+
start: period.start.toISOString(),
78+
end: period.end.toISOString(),
79+
source: 'stripe',
80+
},
81+
payerSubscriptionId: subscriptionId,
82+
})
83+
return { userId, subscriptionId, decision }
84+
}
85+
86+
function requestRows(requestKey: string) {
87+
return db
88+
.select({
89+
eventKey: usageLog.eventKey,
90+
cost: usageLog.cost,
91+
billingPeriodStart: usageLog.billingPeriodStart,
92+
})
93+
.from(usageLog)
94+
.where(inArray(usageLog.eventKey, [`update-cost:${requestKey}`, `update-cost:${requestKey}@1`]))
95+
}
96+
97+
function callback(payer: Payer, requestKey: string, cost: number): NextRequest {
4298
return new NextRequest('http://localhost:3000/api/billing/update-cost', {
4399
method: 'POST',
44100
headers: {
45101
'content-type': 'application/json',
46102
'x-api-key': 'internal',
47103
'x-sim-billing-protocol': 'direct-v1',
48104
'x-sim-billing-request-id': requestKey,
49-
[BILLING_ACCOUNT_DECISION_HEADER]: decision,
105+
[BILLING_ACCOUNT_DECISION_HEADER]: payer.decision,
50106
},
51107
body: JSON.stringify({
52-
userId,
108+
userId: payer.userId,
53109
cost,
54110
model: 'test-model',
55111
source: 'copilot',
@@ -63,50 +119,17 @@ describe('direct-v1 cost callbacks in PostgreSQL', () => {
63119
const now = Date.now()
64120
const admitted = { start: new Date(now - 10 * DAY_MS), end: new Date(now + 20 * DAY_MS) }
65121
const rolled = { start: new Date(now - 60 * 60 * 1000), end: new Date(now + 30 * DAY_MS) }
66-
await db.insert(user).values({
67-
id: userId,
68-
name: 'Update Cost Test',
69-
email: `${userId}@update-cost.test`,
70-
emailVerified: true,
71-
createdAt: new Date(now),
72-
updatedAt: new Date(now),
73-
})
74-
await db.insert(userStats).values({ id: generateId(), userId })
75-
await db.insert(subscription).values({
76-
id: subscriptionId,
77-
plan: 'pro',
78-
referenceId: userId,
79-
status: 'active',
80-
periodStart: admitted.start,
81-
periodEnd: admitted.end,
82-
})
83-
const decision = serializeAccountBillingDecisionHeader({
84-
userId,
85-
billingEntity: { type: 'user', id: userId },
86-
billingPeriod: {
87-
start: admitted.start.toISOString(),
88-
end: admitted.end.toISOString(),
89-
source: 'stripe',
90-
},
91-
payerSubscriptionId: subscriptionId,
92-
})
122+
const payer = await createPayer(admitted)
93123
const requestKey = generateId()
94124

95-
expect((await POST(callback(requestKey, 0.5, decision), {})).status).toBe(200)
125+
expect((await POST(callback(payer, requestKey, 0.5), {})).status).toBe(200)
96126
await db
97127
.update(subscription)
98128
.set({ periodStart: rolled.start, periodEnd: rolled.end })
99-
.where(eq(subscription.id, subscriptionId))
100-
expect((await POST(callback(requestKey, 0.8, decision), {})).status).toBe(200)
129+
.where(eq(subscription.id, payer.subscriptionId))
130+
expect((await POST(callback(payer, requestKey, 0.8), {})).status).toBe(200)
101131

102-
const rows = await db
103-
.select({
104-
eventKey: usageLog.eventKey,
105-
cost: usageLog.cost,
106-
billingPeriodStart: usageLog.billingPeriodStart,
107-
})
108-
.from(usageLog)
109-
.where(eq(usageLog.userId, userId))
132+
const rows = await requestRows(requestKey)
110133
const byKey = new Map(rows.map((row) => [row.eventKey, row]))
111134
expect(rows).toHaveLength(2)
112135
expect(Number(byKey.get(`update-cost:${requestKey}`)?.cost)).toBeCloseTo(0.5)
@@ -118,4 +141,19 @@ describe('direct-v1 cost callbacks in PostgreSQL', () => {
118141
rolled.start.getTime()
119142
)
120143
})
144+
145+
it("refuses spend that lands after the payer's terminal settlement", async () => {
146+
const now = Date.now()
147+
const period = { start: new Date(now - 10 * DAY_MS), end: new Date(now + 20 * DAY_MS) }
148+
const payer = await createPayer(period)
149+
const requestKey = generateId()
150+
151+
expect((await POST(callback(payer, requestKey, 0.5), {})).status).toBe(200)
152+
await claimTerminalPeriod(payer.subscriptionId)
153+
const late = await POST(callback(payer, requestKey, 0.8), {})
154+
155+
expect(late.status).toBe(409)
156+
expect(await late.json()).toMatchObject({ code: 'BILLING_PERIOD_ELAPSED', retryable: false })
157+
expect((await requestRows(requestKey)).map((row) => Number(row.cost))).toEqual([0.5])
158+
})
121159
})

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import {
3030
import {
3131
type CumulativeUsageContextField,
3232
CumulativeUsageContextMismatchError,
33+
CumulativeUsagePeriodClosedError,
3334
recordCumulativeUsage,
3435
} from '@/lib/billing/core/usage-log'
3536
import {
@@ -512,7 +513,8 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
512513
const pgCode = getPostgresErrorCode(error)
513514
const pgConstraint = getPostgresConstraintName(error)
514515
const reconciliationOutcome =
515-
error instanceof ThresholdSettlementError && !error.retryable
516+
(error instanceof ThresholdSettlementError && !error.retryable) ||
517+
error instanceof CumulativeUsagePeriodClosedError
516518
? BILLING_CALLBACK_OUTCOME.billingPeriodElapsed
517519
: pgCode === '23503' && pgConstraint === 'usage_log_user_id_user_id_fk'
518520
? BILLING_CALLBACK_OUTCOME.billingUserNotFound

‎apps/sim/lib/billing/core/usage-log.integration.ts‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({ isOrgScopedSubscription: v
2424

2525
import {
2626
CumulativeUsageContextMismatchError,
27+
CumulativeUsagePeriodClosedError,
2728
getBillingPeriodUsageCost,
2829
getBillingPeriodUsageCostByUser,
2930
getStampedPeriodRangeUsageCostByUser,
@@ -121,7 +122,10 @@ describe('Cumulative billing with PostgreSQL', () => {
121122
CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key)
122123
WHERE event_key IS NOT NULL;
123124
CREATE TABLE driver_probe (id text PRIMARY KEY);
124-
CREATE TABLE subscription (id text PRIMARY KEY, period_start timestamp, period_end timestamp)
125+
CREATE TABLE subscription (
126+
id text PRIMARY KEY, period_start timestamp, period_end timestamp,
127+
last_closed_period_start timestamp
128+
)
125129
`)
126130
transaction.mockImplementation(async (callback: (tx: Transaction) => Promise<unknown>) => {
127131
const pause = nextPause
@@ -362,12 +366,16 @@ describe('Cumulative billing with PostgreSQL', () => {
362366
]
363367
const payer = { type: 'organization', id: 'payer' } as const
364368

369+
/** Moves the subscription to a window whose predecessor the cycle close has settled. */
365370
async function setSubscriptionWindow(start: Date, end: Date) {
366371
await connection`
367-
insert into subscription (id, period_start, period_end)
368-
values ('sub-1', ${start.toISOString()}::timestamptz at time zone 'UTC', ${end.toISOString()}::timestamptz at time zone 'UTC')
372+
insert into subscription (id, period_start, period_end, last_closed_period_start)
373+
values ('sub-1', ${start.toISOString()}::timestamptz at time zone 'UTC', ${end.toISOString()}::timestamptz at time zone 'UTC', ${start.toISOString()}::timestamptz at time zone 'UTC')
369374
on conflict (id) do update
370-
set period_start = excluded.period_start, period_end = excluded.period_end
375+
set period_start = excluded.period_start, period_end = excluded.period_end,
376+
last_closed_period_start = greatest(
377+
subscription.last_closed_period_start, excluded.last_closed_period_start
378+
)
371379
`
372380
}
373381

@@ -528,6 +536,19 @@ describe('Cumulative billing with PostgreSQL', () => {
528536
expect(await ledgerRows()).toEqual([{ event_key: usage(0).eventKey, cost: '0.4' }])
529537
})
530538

539+
it('refuses a charge that would roll into a period the terminal settlement already summed', async () => {
540+
await setSubscriptionPeriod(0)
541+
await charge(0.4)
542+
await setSubscriptionPeriod(1)
543+
await connection`
544+
update subscription
545+
set last_closed_period_start = ${periods[2].toISOString()}::timestamptz at time zone 'UTC'
546+
`
547+
548+
await expect(charge(1)).rejects.toBeInstanceOf(CumulativeUsagePeriodClosedError)
549+
expect(await ledgerRows()).toEqual([{ event_key: usage(0).eventKey, cost: '0.4' }])
550+
})
551+
531552
it('holds an early period-start move until an in-flight top-up commits', async () => {
532553
const start = new Date(Date.now() - 24 * 60 * 60 * 1000)
533554
const end = new Date(Date.now() + 30 * 24 * 60 * 60 * 1000)

‎apps/sim/lib/billing/core/usage-log.ts‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -590,7 +590,9 @@ export interface RecordCumulativeUsageParams {
590590
* arrives after that subscription has moved past the period of the request's latest row is
591591
* recorded in a new row stamped with the subscription's current period, so a request that
592592
* outlives its billing period is invoiced by the period it was spent in rather than topping up
593-
* a period that has already been closed. Omit it for reporting-window and free payers.
593+
* a period that has already been closed. A charge into a period the subscription's close marker
594+
* has already passed (its terminal settlement) throws {@link CumulativeUsagePeriodClosedError}.
595+
* Omit it for reporting-window and free payers.
594596
*
595597
* Mixed versions: code that predates period rows reads only the request key. If such code
596598
* (during a deploy, or after a rollback) handles a later callback for a request that already
@@ -680,6 +682,23 @@ export class CumulativeUsageContextMismatchError extends Error {
680682
}
681683
}
682684

685+
/**
686+
* A cumulative charge whose billing period the payer has already settled: the subscription ended
687+
* and its final invoice summed that period. The charge is refused rather than recorded where no
688+
* invoice will ever read it.
689+
*/
690+
export class CumulativeUsagePeriodClosedError extends Error {
691+
constructor(
692+
readonly eventKey: string,
693+
readonly billingPeriod: { start: Date; end: Date }
694+
) {
695+
super(
696+
`Cumulative usage event "${eventKey}" targets a billing period that has already been settled`
697+
)
698+
this.name = 'CumulativeUsagePeriodClosedError'
699+
}
700+
}
701+
683702
interface CumulativeUsageLedgerBinding {
684703
userId: string
685704
workspaceId: string | null
@@ -882,14 +901,15 @@ export async function recordCumulativeUsage(
882901
return { billed: false, delta: 0, total: recorded, billingPeriod: latestPeriod }
883902
}
884903

885-
// The payer's current period, share-locked so a change to the subscription's period (a
886-
// rollover, or an anchor reset inside the old period) waits for this write to commit, and
887-
// whatever a close later sums for the old period is final.
904+
// The payer's current period and close marker, share-locked so a change to either (a
905+
// rollover, an anchor reset inside the old period, or a terminal settlement) waits for this
906+
// write to commit, and whatever a close later sums for the old period is final.
888907
const [currentPeriod] = payerSubscriptionId
889908
? await tx
890909
.select({
891910
start: subscriptionTable.periodStart,
892911
end: subscriptionTable.periodEnd,
912+
closedThrough: subscriptionTable.lastClosedPeriodStart,
893913
})
894914
.from(subscriptionTable)
895915
.where(eq(subscriptionTable.id, payerSubscriptionId))
@@ -910,6 +930,16 @@ export async function recordCumulativeUsage(
910930
if (rolledPeriod && latest && chain.length >= MAX_CUMULATIVE_PERIOD_ROWS) {
911931
throw new Error(`Cumulative usage event "${eventKey}" spans too many billing periods`)
912932
}
933+
// A marker at or past the target period's end means that period is already settled — a
934+
// terminal settlement marks it whatever the subscription's bounds — so nothing would
935+
// ever invoice this charge.
936+
const targetPeriod = rolledPeriod ?? latestPeriod
937+
if (
938+
currentPeriod?.closedThrough &&
939+
currentPeriod.closedThrough.getTime() >= targetPeriod.end.getTime()
940+
) {
941+
throw new CumulativeUsagePeriodClosedError(eventKey, targetPeriod)
942+
}
913943

914944
enterStage('write')
915945
if (latest && !rolledPeriod) {
@@ -929,7 +959,6 @@ export async function recordCumulativeUsage(
929959
return { billed: true, delta, total: newTotal, billingPeriod: latestPeriod }
930960
}
931961

932-
const targetPeriod = rolledPeriod ?? billingContext.billingPeriod
933962
const rowMetadata = periodUsageMetadata(metadata, chain)
934963
await recordUsage({
935964
userId,

‎apps/sim/lib/billing/cycle-close.ts‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -191,12 +191,16 @@ export async function closeElapsedPeriodBeforeDeletion(subscriptionId: string):
191191
* Claim the terminal period for a subscription that is being deleted, BEFORE
192192
* the deletion handler computes and charges final overage. Reads the
193193
* subscription row fresh (webhook payloads can be stale across a rollover)
194-
* and advances the close marker to its current `periodStart` in one
195-
* transaction, serializing with the sweep on the subscription row: an
196-
* in-flight sweep close then fails its guarded marker claim and rolls back —
197-
* including its outbox invoice — so deletion and sweep can never both bill
198-
* the same period. Call `closeElapsedPeriodBeforeDeletion` first so a lagging
199-
* elapsed period is settled rather than jumped. Returns the fresh period
194+
* and, in one transaction, advances the close marker to the terminal period's end:
195+
* the period is settled from here on, so a cost callback that commits after
196+
* this claim is refused rather than topping up a period the final invoice has
197+
* already summed (`recordCumulativeUsage` reads the marker under a share lock
198+
* on the same row, so every charge either commits before this claim or sees
199+
* the marker). This also serializes with the sweep on the subscription row:
200+
* an in-flight sweep close then fails its guarded marker claim and rolls
201+
* back — including its outbox invoice — so deletion and sweep can never both
202+
* bill the same period. Call `closeElapsedPeriodBeforeDeletion` first so a
203+
* lagging elapsed period is settled rather than jumped. Returns the period
200204
* bounds for the deletion flow to settle against, plus `markerWasCurrent`:
201205
* whether the close marker had already caught up to the terminal period.
202206
* The `billedOverageThisPeriod` tracker only ever holds collections for the
@@ -241,7 +245,10 @@ export async function claimTerminalPeriod(
241245
const markerWasCurrent =
242246
!!row.lastClosedPeriodStart &&
243247
row.lastClosedPeriodStart.getTime() >= row.periodStart.getTime()
244-
if (!markerWasCurrent && options.sealLagging) {
248+
if (!markerWasCurrent && !options.sealLagging) {
249+
return { periodStart: row.periodStart, periodEnd: row.periodEnd, markerWasCurrent }
250+
}
251+
if (!markerWasCurrent) {
245252
logger.error(
246253
'Sealing an unclosed elapsed period at terminal claim; residual overage forgiven',
247254
{
@@ -250,8 +257,8 @@ export async function claimTerminalPeriod(
250257
periodStart: row.periodStart.toISOString(),
251258
}
252259
)
253-
await claimCloseMarker(tx, subscriptionId, row.periodStart)
254260
}
261+
await claimCloseMarker(tx, subscriptionId, row.periodEnd ?? row.periodStart)
255262
return { periodStart: row.periodStart, periodEnd: row.periodEnd, markerWasCurrent }
256263
})
257264
}

0 commit comments

Comments
 (0)