Skip to content

Commit 2091953

Browse files
authored
fix(billing): build the mid-run usage card from the admitted payer and model the empty new-turn 402 (#8502)
* fix(billing): build the mid-run usage card from the admitted payer and model the empty new-turn 402 A direct-v1 run's mid-run verdict reads the payer saved in its account decision, but the upgrade card was resolved from the actor's current subscription, so a payer/actor mismatch picked the wrong action and copy. The exceeded account verdict now carries the payer and subscription it already read, and update-cost and the validate continuation pass it to resolveUsageUpgradePayload instead of a second lookup. The validate contract declared every 402 as a JSON refusal, while a new turn's 402 has no body; the 402 schema now allows the empty body. The wire is unchanged. * fix(billing): drop the unreachable card-read deadline from the usage upgrade card Every exceeded verdict that reaches update-cost now carries its payer, so the deadline-bounded actor subscription lookup could no longer run. Remove the parameter, its call-site argument, the stale TSDoc, and the test that passed without exercising it. * test(copilot): pin the new-turn usage refusal to an empty 402
1 parent 7e8d3ff commit 2091953

7 files changed

Lines changed: 100 additions & 46 deletions

File tree

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

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1040,32 +1040,38 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
10401040
})
10411041
})
10421042

1043-
it('never pauses a blocked payer with the usage card', async () => {
1044-
billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({
1045-
blocked: true,
1046-
scope: 'payer',
1043+
it("offers the card for its admitted payer's plan, not the actor's current one", async () => {
1044+
billingCoreMockFns.mockGetOrganizationSubscription.mockResolvedValue({
1045+
id: 'sub-account-org',
1046+
referenceId: 'account-org',
1047+
plan: 'team',
1048+
status: 'active',
1049+
seats: 4,
1050+
})
1051+
billingPlanMockFns.mockGetHighestPrioritySubscription.mockResolvedValue({
1052+
id: 'sub-personal',
1053+
referenceId: 'user-1',
1054+
plan: 'pro',
1055+
status: 'active',
10471056
})
10481057

10491058
const body = await (await POST(directCallback())).json()
10501059

1051-
expect(body.usageExceeded).toBe(false)
1060+
expect(body.usageUpgrade).toMatchObject({
1061+
action: 'increase_limit',
1062+
message: expect.stringContaining("organization's usage limit"),
1063+
})
10521064
})
10531065

1054-
it('keeps the exceeded verdict with the plan-upgrade card when the card read outlasts the callback budget', async () => {
1055-
billingPlanMockFns.mockGetHighestPrioritySubscription.mockImplementation(async () => {
1056-
await sleep(1500)
1057-
return { plan: 'pro' }
1066+
it('never pauses a blocked payer with the usage card', async () => {
1067+
billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({
1068+
blocked: true,
1069+
scope: 'payer',
10581070
})
1059-
const startedAt = Date.now()
10601071

10611072
const body = await (await POST(directCallback())).json()
10621073

1063-
expect(body).toMatchObject({
1064-
success: true,
1065-
usageExceeded: true,
1066-
usageUpgrade: { action: 'upgrade_plan' },
1067-
})
1068-
expect(Date.now() - startedAt).toBeLessThan(1400)
1074+
expect(body.usageExceeded).toBe(false)
10691075
})
10701076
})
10711077

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp
8686
* already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the
8787
* refusal to the next step or re-check rather than ending a paying run on a database blip,
8888
* and so does a verdict read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}. An exceeded
89-
* verdict always pauses the run; a card read past that budget falls back to the plan-upgrade card.
89+
* verdict always pauses the run.
9090
*/
9191
async function readUsageStanding(
9292
userId: string,
@@ -99,10 +99,9 @@ async function readUsageStanding(
9999
? () => readMidRunAccountUsageVerdict(accountDecision)
100100
: null
101101
if (!isHosted || !readVerdict) return { usageExceeded: false }
102-
const deadlineAt = Date.now() + USAGE_STANDING_TIMEOUT_MS
103102
let verdict: MidRunUsageVerdict
104103
try {
105-
verdict = await withinDeadline(readVerdict, deadlineAt)
104+
verdict = await withinDeadline(readVerdict, Date.now() + USAGE_STANDING_TIMEOUT_MS)
106105
} catch {
107106
logger.warn('Usage standing read outlasted the callback budget; answering not exceeded')
108107
return { usageExceeded: false }
@@ -114,9 +113,8 @@ async function readUsageStanding(
114113
usageExceeded: true,
115114
usageUpgrade: await resolveUsageUpgradePayload(
116115
userId,
117-
billingAttribution,
118-
verdict.scope,
119-
deadlineAt
116+
billingAttribution ?? verdict.payer,
117+
verdict.scope
120118
),
121119
}
122120
}

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

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,10 @@ vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock)
140140

141141
vi.mock('@/lib/workspaces/utils', () => workspacesUtilsMock)
142142

143-
import { validateCopilotApiKeyBodySchema } from '@/lib/api/contracts/copilot'
143+
import {
144+
validateCopilotApiKeyBodySchema,
145+
validateCopilotApiKeyContract,
146+
} from '@/lib/api/contracts/copilot'
144147
import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage'
145148
import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
146149
import { POST } from '@/app/api/copilot/api-keys/validate/route'
@@ -248,6 +251,17 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => {
248251
expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled()
249252
})
250253

254+
it("sends a new turn's usage refusal as the empty 402 its contract declares", async () => {
255+
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
256+
257+
const res = await POST(request(SELF_HOSTED_VALIDATE_BODY))
258+
259+
expect(res.status).toBe(402)
260+
expect(await res.text()).toBe('')
261+
const refusalSchema = validateCopilotApiKeyContract.response.statusSchemas?.[402]
262+
expect(refusalSchema?.safeParse(undefined).success).toBe(true)
263+
})
264+
251265
it('preserves the actor member cap for markerless self-hosted admission', async () => {
252266
mockCheckAttributedUsageLimits.mockResolvedValue({
253267
isExceeded: true,
@@ -593,6 +607,28 @@ describe('validation lifecycle purposes', () => {
593607
})
594608
})
595609

610+
it("refuses a direct-v1 continuation with the card for its admitted payer's plan", async () => {
611+
mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 })
612+
mockGetOrganizationSubscription.mockResolvedValue({
613+
id: 'sub-account-org',
614+
referenceId: 'account-org',
615+
plan: 'team',
616+
status: 'active',
617+
seats: 4,
618+
})
619+
mockGetHighestPrioritySubscription.mockResolvedValue(null)
620+
621+
const response = await POST(request(body, directHeaders))
622+
623+
expect(response.status).toBe(402)
624+
await expect(response.json()).resolves.toMatchObject({
625+
usageUpgrade: {
626+
action: 'increase_limit',
627+
message: expect.stringContaining("organization's usage limit"),
628+
},
629+
})
630+
})
631+
596632
it('admits a direct-v1 continuation whose usage cannot be read', async () => {
597633
mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, unavailable: true })
598634
expect((await POST(request(body, directHeaders))).status).toBe(200)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -453,7 +453,7 @@ export const POST = withRouteHandler((req: NextRequest) =>
453453
span.setAttribute(TraceAttr.HttpStatusCode, 402)
454454
const usageUpgrade = await resolveUsageUpgradePayload(
455455
userId,
456-
billing?.kind === 'attributed' ? billing.attribution : undefined,
456+
billing?.kind === 'attributed' ? billing.attribution : verdict.payer,
457457
verdict.scope
458458
)
459459
return NextResponse.json<ValidateCopilotApiKeyUsageExceeded>(

‎apps/sim/lib/api/contracts/copilot.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -433,7 +433,7 @@ export const validateCopilotApiKeyContract = defineRouteContract({
433433
status: [200, 402],
434434
statusSchemas: {
435435
200: validateCopilotApiKeyResponseSchema,
436-
402: validateCopilotApiKeyRefusalSchema,
436+
402: validateCopilotApiKeyRefusalSchema.optional(),
437437
},
438438
},
439439
error: validateCopilotApiKeyErrorSchema,

‎apps/sim/lib/billing/core/mid-run-usage.ts‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import {
1818
USAGE_GATE_SETTLE_TIMEOUT_MS,
1919
USAGE_GATE_TTL_MS,
2020
} from '@/lib/billing/core/usage-gate-cache'
21+
import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade'
2122
import { coalesceLocally } from '@/lib/concurrency/singleflight'
2223
import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags'
2324

@@ -26,7 +27,8 @@ const logger = createLogger('MidRunUsage')
2627
/**
2728
* A run's standing while it is under way, read through the execution usage gate:
2829
* - `exceeded`: the payer (or the actor's member cap) spent its limit; the run pauses with the
29-
* upgrade card.
30+
* upgrade card. A direct-v1 verdict carries the payer it read, so the card names that payer's
31+
* plan rather than the actor's.
3032
* - `blocked`: the account is blocked (payment failed, dispute); the run is refused as a blocked
3133
* account, never with the upgrade card.
3234
* - `unknown`: usage, or the payer's current period, could not be read. Admission fails closed
@@ -35,7 +37,11 @@ const logger = createLogger('MidRunUsage')
3537
*/
3638
export type MidRunUsageVerdict =
3739
| { status: 'within' }
38-
| { status: 'exceeded'; scope?: AttributedUsageLimitsResult['scope'] }
40+
| {
41+
status: 'exceeded'
42+
scope?: AttributedUsageLimitsResult['scope']
43+
payer?: UsageUpgradePayer
44+
}
3945
| { status: 'blocked'; message?: string }
4046
| { status: 'unknown' }
4147

@@ -213,7 +219,13 @@ export async function readMidRunAccountUsageVerdict(
213219
USAGE_GATE_SETTLE_TIMEOUT_MS
214220
)
215221
if (usage.unavailable) return { status: 'unknown' }
216-
if (usage.isExceeded) return { status: 'exceeded', scope: 'payer' }
222+
if (usage.isExceeded) {
223+
return {
224+
status: 'exceeded',
225+
scope: 'payer',
226+
payer: { billingEntity: payer, payerSubscription: subscription },
227+
}
228+
}
217229
const within: MidRunUsageVerdict = { status: 'within' }
218230
accountVerdictCache.set(key, within)
219231
return within

‎apps/sim/lib/billing/usage-upgrade.ts‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
11
import { createLogger } from '@sim/logger'
22
import { getErrorMessage } from '@sim/utils/errors'
33
import type { UsageUpgradePayload } from '@/lib/api/contracts/subscription'
4-
import type {
5-
AttributedUsageLimitsResult,
6-
BillingAttributionSnapshot,
7-
} from '@/lib/billing/core/billing-attribution'
4+
import type { AttributedUsageLimitsResult } from '@/lib/billing/core/billing-attribution'
85
import { getHighestPrioritySubscription } from '@/lib/billing/core/plan'
6+
import type { BillingEntity } from '@/lib/billing/core/usage-log'
97
import { isEnterprise, isPaid } from '@/lib/billing/plan-helpers'
108
import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
11-
import { withinDeadline } from '@/lib/core/utils/deadline'
129

1310
const logger = createLogger('UsageUpgrade')
1411

@@ -18,33 +15,38 @@ const UPGRADE_PLAN_MESSAGE =
1815
const MEMBER_CAP_MESSAGE =
1916
"You've reached the usage limit your organization set for you this billing period. Only an organization owner or admin can raise it — please ask them to continue."
2017

18+
/**
19+
* The payer a run is billed to, as the upgrade card needs it: its billing entity and its
20+
* subscription's plan. An attribution snapshot is one; a direct-v1 run's mid-run verdict carries one.
21+
*/
22+
export interface UsageUpgradePayer {
23+
readonly billingEntity: Readonly<BillingEntity>
24+
readonly payerSubscription: { readonly plan: string } | null
25+
}
26+
2127
/**
2228
* The upgrade card for a payer over its usage limit: a plan upgrade for a free payer, a limit
2329
* increase for a paid one, with copy naming who can raise an organization's limit. A member
24-
* over the cap their organization set gets copy naming who can raise that cap. An attributed
25-
* run reads the plan from its admission snapshot without a query; otherwise the actor's current
26-
* subscription decides, and a lookup that fails or outlasts `deadlineAt` falls back to the
27-
* plan-upgrade card.
30+
* over the cap their organization set gets copy naming who can raise that cap. A known payer
31+
* decides the card without a query; otherwise the actor's current subscription decides, and a
32+
* lookup that fails falls back to the plan-upgrade card.
2833
*/
2934
export async function resolveUsageUpgradePayload(
3035
userId: string,
31-
billingAttribution?: BillingAttributionSnapshot,
32-
scope?: AttributedUsageLimitsResult['scope'],
33-
deadlineAt?: number
36+
payer?: UsageUpgradePayer,
37+
scope?: AttributedUsageLimitsResult['scope']
3438
): Promise<UsageUpgradePayload> {
3539
if (scope === 'member') {
3640
return { reason: 'usage_limit', action: 'increase_limit', message: MEMBER_CAP_MESSAGE }
3741
}
3842
let plan: string | undefined
3943
let orgScoped = false
4044
try {
41-
if (billingAttribution) {
42-
plan = billingAttribution.payerSubscription?.plan
43-
orgScoped = billingAttribution.billingEntity.type === 'organization'
45+
if (payer) {
46+
plan = payer.payerSubscription?.plan
47+
orgScoped = payer.billingEntity.type === 'organization'
4448
} else {
45-
const subscription = await (deadlineAt === undefined
46-
? getHighestPrioritySubscription(userId)
47-
: withinDeadline(() => getHighestPrioritySubscription(userId), deadlineAt))
49+
const subscription = await getHighestPrioritySubscription(userId)
4850
plan = subscription?.plan
4951
orgScoped = isOrgScopedSubscription(subscription, userId)
5052
}

0 commit comments

Comments
 (0)