Skip to content

Commit d4bd072

Browse files
waleedlatif1waleedlatif1
committed
fix(billing): roll a direct-v1 run's spend into the payer's current Stripe period and judge its standing mid-run
- A direct-v1 account decision now carries the payer's subscription from admission, and a cost callback rolls later spend into that subscription's current period exactly as an attributed run does, so a run that outlives its period never tops up a closed one. Only a Stripe period rolls; reporting windows and free payers keep their frozen period. - A direct-v1 cost callback reports the admitted payer's standing, and the direct gate checks the actor and payer for a block before their spend, so a blocked account is never paused with the upgrade card. The account block check moves to billing core and is shared with continuation validation. - The attributed mid-run gate returns early when billing is off. - Tests: a direct run across a rollover against real PostgreSQL, per-period token shares, and the usage card replay asserted on parsed segments.
1 parent e1d1102 commit d4bd072

15 files changed

Lines changed: 365 additions & 74 deletions

File tree

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

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/**
2+
* Cost callbacks against real PostgreSQL: a direct-v1 run that outlives its admitted Stripe period
3+
* 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.
5+
*/
6+
import { db } from '@sim/db'
7+
import { subscription, usageLog, user, userStats } from '@sim/db/schema'
8+
import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock'
9+
import { generateId } from '@sim/utils/id'
10+
import { eq } from 'drizzle-orm'
11+
import { NextRequest } from 'next/server'
12+
import { afterAll, describe, expect, it, vi } from 'vitest'
13+
14+
vi.mock('@/lib/core/config/env-flags', () => ({
15+
...envFlagsMock,
16+
isHosted: true,
17+
isBillingEnabled: true,
18+
}))
19+
vi.mock('@/lib/mothership/request/http', async (importOriginal) => ({
20+
...(await importOriginal<typeof import('@/lib/mothership/request/http')>()),
21+
checkInternalApiKey: () => ({ success: true }),
22+
}))
23+
24+
import {
25+
BILLING_ACCOUNT_DECISION_HEADER,
26+
serializeAccountBillingDecisionHeader,
27+
} from '@/lib/billing/core/billing-attribution'
28+
import { POST } from '@/app/api/billing/update-cost/route'
29+
30+
const DAY_MS = 24 * 60 * 60 * 1000
31+
const userId = `update-cost-user-${generateId()}`
32+
const subscriptionId = generateId()
33+
34+
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))
39+
})
40+
41+
function callback(requestKey: string, cost: number, decision: string): NextRequest {
42+
return new NextRequest('http://localhost:3000/api/billing/update-cost', {
43+
method: 'POST',
44+
headers: {
45+
'content-type': 'application/json',
46+
'x-api-key': 'internal',
47+
'x-sim-billing-protocol': 'direct-v1',
48+
'x-sim-billing-request-id': requestKey,
49+
[BILLING_ACCOUNT_DECISION_HEADER]: decision,
50+
},
51+
body: JSON.stringify({
52+
userId,
53+
cost,
54+
model: 'test-model',
55+
source: 'copilot',
56+
idempotencyKey: requestKey,
57+
}),
58+
})
59+
}
60+
61+
describe('direct-v1 cost callbacks in PostgreSQL', () => {
62+
it('records spend after a Stripe rollover in the payer current period', async () => {
63+
const now = Date.now()
64+
const admitted = { start: new Date(now - 10 * DAY_MS), end: new Date(now + 20 * DAY_MS) }
65+
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+
})
93+
const requestKey = generateId()
94+
95+
expect((await POST(callback(requestKey, 0.5, decision), {})).status).toBe(200)
96+
await db
97+
.update(subscription)
98+
.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)
101+
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))
110+
const byKey = new Map(rows.map((row) => [row.eventKey, row]))
111+
expect(rows).toHaveLength(2)
112+
expect(Number(byKey.get(`update-cost:${requestKey}`)?.cost)).toBeCloseTo(0.5)
113+
expect(byKey.get(`update-cost:${requestKey}`)?.billingPeriodStart?.getTime()).toBe(
114+
admitted.start.getTime()
115+
)
116+
expect(Number(byKey.get(`update-cost:${requestKey}@1`)?.cost)).toBeCloseTo(0.3)
117+
expect(byKey.get(`update-cost:${requestKey}@1`)?.billingPeriodStart?.getTime()).toBe(
118+
rolled.start.getTime()
119+
)
120+
})
121+
})

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,16 @@ import {
44
billingAttributionMock,
55
billingAttributionMockFns,
66
} from '@sim/testing/mocks/billing-attribution.mock'
7+
import { billingCoreMock, billingCoreMockFns } from '@sim/testing/mocks/billing-core.mock'
8+
import { billingPlanMock, billingPlanMockFns } from '@sim/testing/mocks/billing-plan.mock'
79
import {
810
billingUsageLogMock,
911
billingUsageLogMockFns,
1012
} from '@sim/testing/mocks/billing-usage-log.mock'
13+
import {
14+
billingUsageMonitorMock,
15+
billingUsageMonitorMockFns,
16+
} from '@sim/testing/mocks/billing-usage-monitor.mock'
1117
import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing/mocks/copilot-http.mock'
1218
import { mothershipOtelMock } from '@sim/testing/mocks/mothership-otel.mock'
1319
import { sleep } from '@sim/utils/helpers'
@@ -42,6 +48,12 @@ vi.mock('@/lib/billing/core/usage-log', () => billingUsageLogMock)
4248

4349
vi.mock('@/lib/billing/core/billing-attribution', () => billingAttributionMock)
4450

51+
vi.mock('@/lib/billing/core/billing', () => billingCoreMock)
52+
53+
vi.mock('@/lib/billing/core/plan', () => billingPlanMock)
54+
55+
vi.mock('@/lib/billing/calculations/usage-monitor', () => billingUsageMonitorMock)
56+
4557
vi.mock('@/lib/billing/threshold-billing', () => ({
4658
checkAndBillOverageThreshold: mockCheckAndBillOverageThreshold,
4759
checkAndBillPayerOverageThreshold: mockCheckAndBillPayerOverageThreshold,
@@ -977,6 +989,62 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
977989
})
978990
})
979991

992+
describe('a direct-v1 run', () => {
993+
function directCallback() {
994+
callbackSequence += 1
995+
const billingRequestId = `0190c03f-9f7d-4b79-8b58-${String(callbackSequence).padStart(12, '0')}`
996+
return createMockRequest(
997+
'POST',
998+
{
999+
userId: 'user-1',
1000+
cost: 0.5 * callbackSequence,
1001+
model: 'claude-opus-4.8',
1002+
source: 'workspace-chat',
1003+
idempotencyKey: billingRequestId,
1004+
},
1005+
{
1006+
'x-api-key': 'internal',
1007+
'x-sim-billing-protocol': 'direct-v1',
1008+
'x-sim-billing-request-id': billingRequestId,
1009+
'x-sim-billing-account-decision': 'serialized-account-decision',
1010+
}
1011+
)
1012+
}
1013+
1014+
beforeEach(() => {
1015+
mockRequireAccountBillingDecisionHeader.mockReturnValue(ACCOUNT_BILLING_DECISION)
1016+
billingCoreMockFns.mockGetOrganizationSubscription.mockResolvedValue(null)
1017+
billingPlanMockFns.mockGetHighestPrioritySubscription.mockResolvedValue(null)
1018+
billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({ blocked: false })
1019+
billingUsageMonitorMockFns.mockCheckUsageStatus.mockResolvedValue({
1020+
isExceeded: true,
1021+
currentUsage: 12,
1022+
limit: 10,
1023+
})
1024+
})
1025+
1026+
it('tells the worker when its admitted payer has crossed its usage limit', async () => {
1027+
const body = await (await POST(directCallback())).json()
1028+
1029+
expect(body).toMatchObject({
1030+
usageExceeded: true,
1031+
usageUpgrade: { reason: 'usage_limit' },
1032+
})
1033+
})
1034+
1035+
it('never pauses a blocked payer with the usage card', async () => {
1036+
billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({
1037+
blocked: true,
1038+
scope: 'payer',
1039+
})
1040+
1041+
const body = await (await POST(directCallback())).json()
1042+
1043+
expect(body.usageExceeded).toBe(false)
1044+
expect(billingUsageMonitorMockFns.mockCheckUsageStatus).not.toHaveBeenCalled()
1045+
})
1046+
})
1047+
9801048
describe('a run that outlives its billing period', () => {
9811049
const PAYER_SUBSCRIPTION = {
9821050
id: 'sub-1',
@@ -995,6 +1063,14 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
9951063
...CURRENT_ATTRIBUTION,
9961064
payerSubscription: PAYER_SUBSCRIPTION,
9971065
})
1066+
mockToBillingContext.mockReturnValue({
1067+
billingEntity: { type: 'organization', id: 'org-1' },
1068+
billingPeriod: {
1069+
start: new Date('2026-07-01T00:00:00.000Z'),
1070+
end: new Date('2026-08-01T00:00:00.000Z'),
1071+
source: 'stripe',
1072+
},
1073+
})
9981074
})
9991075

10001076
it('settles the threshold against the period the charge was stamped into', async () => {

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

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,11 @@ import {
2222
resolveLegacyV0BillingAttribution,
2323
toBillingContext,
2424
} from '@/lib/billing/core/billing-attribution'
25-
import { type MidRunUsageVerdict, readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage'
25+
import {
26+
type MidRunUsageVerdict,
27+
readMidRunAccountUsageVerdict,
28+
readMidRunUsageVerdict,
29+
} from '@/lib/billing/core/mid-run-usage'
2630
import {
2731
type CumulativeUsageContextField,
2832
CumulativeUsageContextMismatchError,
@@ -75,24 +79,28 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp
7579
/**
7680
* Reads the run payer's standing after a cost callback, so a long run stops at its next step
7781
* once it crosses the limit instead of at its next admission, with the card the worker writes
78-
* to its log. A duplicate callback answers too: it is often a retry whose first answer was lost.
79-
* Served from the execution usage gate: an admission is cached per payer and actor for the gate
80-
* TTL and a refusal is always re-read, so steady-state steps cost no ledger read. The charge is
82+
* to its log. The payer is the attributed run's, or the one a direct-v1 run was admitted with.
83+
* A duplicate callback answers too: it is often a retry whose first answer was lost. An
84+
* admission is cached per payer and actor for the gate TTL and a refusal is always re-read, so
85+
* steady-state steps cost no ledger read. The charge is
8186
* already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the
8287
* refusal to the next step or re-check rather than ending a paying run on a database blip,
8388
* and so does a read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}.
8489
*/
8590
async function readUsageStanding(
8691
userId: string,
87-
billingAttribution: BillingAttributionSnapshot | undefined
92+
billingAttribution: BillingAttributionSnapshot | undefined,
93+
accountDecision: AccountBillingDecision | undefined
8894
): Promise<BillingUsageVerdict> {
89-
if (!isHosted || !billingAttribution) return { usageExceeded: false }
95+
const readVerdict = billingAttribution
96+
? () => readMidRunUsageVerdict(billingAttribution)
97+
: accountDecision
98+
? () => readMidRunAccountUsageVerdict(accountDecision)
99+
: null
100+
if (!isHosted || !readVerdict) return { usageExceeded: false }
90101
let verdict: MidRunUsageVerdict
91102
try {
92-
verdict = await withinDeadline(
93-
() => readMidRunUsageVerdict(billingAttribution),
94-
Date.now() + USAGE_STANDING_TIMEOUT_MS
95-
)
103+
verdict = await withinDeadline(readVerdict, Date.now() + USAGE_STANDING_TIMEOUT_MS)
96104
} catch {
97105
logger.warn('Usage standing read outlasted the callback budget; answering not exceeded')
98106
return { usageExceeded: false }
@@ -370,8 +378,8 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
370378
// period (see `payerSubscriptionId`), so a closed period is never topped up. Reporting-window
371379
// payers are closed from live anchors, and free payers have no close to miss.
372380
const rolloverSubscriptionId =
373-
billingAttribution && billingContext?.billingPeriod.source !== 'reporting'
374-
? (billingAttribution.payerSubscription?.id ?? undefined)
381+
billingContext?.billingPeriod.source === 'stripe'
382+
? (billingAttribution?.payerSubscription?.id ?? accountDecision?.payerSubscriptionId)
375383
: undefined
376384
const usageStartedAt = Date.now()
377385
const result = await recordCumulativeUsage({
@@ -426,7 +434,7 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
426434
await checkAndBillOverageThreshold(userId, undefined, { onError: 'throw' })
427435
}
428436

429-
const usageVerdict = await readUsageStanding(userId, billingAttribution)
437+
const usageVerdict = await readUsageStanding(userId, billingAttribution, accountDecision)
430438
const duration = Date.now() - startTime
431439

432440
// Same-or-lower cumulative than already recorded: nothing new to bill.

‎apps/sim/app/api/copilot/api-keys/validate/route.test.ts‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,12 @@ const { mockAuthorizeOrganizationChatDelegation: mockAuthorizeOrganizationChat }
152152
const { mockDeriveBillingContext } = billingUsageLogMockFns
153153
const { mockGetHighestPrioritySubscription } = billingPlanMockFns
154154
const { mockGetOrganizationSubscription } = billingCoreMockFns
155-
const { mockCheckServerSideUsageLimits, mockCheckUsageStatus } = billingUsageMonitorMockFns
155+
const {
156+
mockCheckBillingBlocked,
157+
mockCheckBillingEntityBlocked,
158+
mockCheckServerSideUsageLimits,
159+
mockCheckUsageStatus,
160+
} = billingUsageMonitorMockFns
156161

157162
const mockIsEnterprisePlan = billingSubscriptionMockFns.mockIsEnterprisePlan
158163
const mockGetUserEntityPermissions = permissionsMockFns.mockGetUserEntityPermissions
@@ -421,7 +426,10 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => {
421426
expect(mockResolveBillingAttribution).not.toHaveBeenCalled()
422427
expect(mockGetUserEntityPermissions).not.toHaveBeenCalled()
423428
expect(mockGetWorkspaceBillingSettings).not.toHaveBeenCalled()
424-
expect(mockSerializeAccountBillingDecisionHeader).toHaveBeenCalledWith(ACCOUNT_BILLING_DECISION)
429+
expect(mockSerializeAccountBillingDecisionHeader).toHaveBeenCalledWith({
430+
...ACCOUNT_BILLING_DECISION,
431+
payerSubscriptionId: ACCOUNT_SUBSCRIPTION.id,
432+
})
425433
expect(res.headers.get('x-sim-billing-account-decision')).toBe('serialized-account-decision')
426434
})
427435

@@ -514,6 +522,8 @@ describe('validation lifecycle purposes', () => {
514522
mockCheckContinuationBilling.mockReset().mockResolvedValue({ blocked: false })
515523
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false })
516524
mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 1, limit: 10 })
525+
mockCheckBillingBlocked.mockResolvedValue({ blocked: false })
526+
mockCheckBillingEntityBlocked.mockResolvedValue({ blocked: false })
517527
mockIsEnterprisePlan.mockResolvedValue(false)
518528
mockGetOrganizationSubscription.mockResolvedValue({
519529
id: 'sub-org-1',
@@ -810,6 +820,14 @@ describe('validation lifecycle purposes', () => {
810820
expect(mockCheckUsageStatus).toHaveBeenCalledTimes(3)
811821
})
812822

823+
it('never reads the usage gate for an attributed continuation when billing is off', async () => {
824+
setEnvFlags({ isHosted: false, isBillingEnabled: false })
825+
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
826+
827+
expect((await POST(request(body, attributedHeaders))).status).toBe(200)
828+
expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled()
829+
})
830+
813831
it('never reads the ledger for a direct-v1 continuation when billing is off', async () => {
814832
setEnvFlags({ isHosted: false, isBillingEnabled: false })
815833
mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 })

‎apps/sim/app/api/copilot/api-keys/validate/route.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,7 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise
319319
? { source: billingContext.billingPeriod.source }
320320
: {}),
321321
},
322+
...(subscription ? { payerSubscriptionId: subscription.id } : {}),
322323
},
323324
}
324325
}

‎apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -888,7 +888,9 @@ describe('usage card written to a worker log', () => {
888888

889889
const { segments } = parseSpecialTags(replayed, false)
890890

891-
expect(segments).toContainEqual({ type: 'usage_upgrade', data: usageUpgrade })
892-
expect(renderedText(segments)).toBe('Finished the first report.')
891+
expect(segments).toEqual([
892+
{ type: 'text', content: 'Finished the first report.' },
893+
{ type: 'usage_upgrade', data: usageUpgrade },
894+
])
893895
})
894896
})

0 commit comments

Comments
 (0)