From 62eb19d5910d62f686762770991e273c93ce6c92 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 02:07:46 -0700 Subject: [PATCH 01/13] fix(billing): enforce the plan usage limit mid-run and bill runs that outlive their period Unbounded Chat runs need spend enforced inside a run, not only at its edges. - update-cost answers every callback (200 and duplicate 409) with a top-level usageExceeded verdict read through the cached execution usage gate, plus the usageUpgrade card payload when exceeded. - Continuation validation and the lifecycle's continuation admission read the original payer's spend through the same gate. A refused validation returns 402 { code: USAGE_LIMIT_EXCEEDED, error, usageUpgrade }; a blocked account returns 402 { code: BILLING_BLOCKED, error } and never gets the usage card. A refused lifecycle continuation renders the upgrade card and stops the worker run so the next message after an upgrade is not refused as busy. - Mid-run paths treat an unreadable ledger as unknown and keep the run going; admission before a run still fails closed. - The card's action and copy come from one resolver shared by Sim's synthetic card and the verdicts the worker echoes into its durable log, with member-cap copy for a member over the cap their organization set. - A cumulative charge that outlives its Stripe billing period records later spend in one row per later period, stamped with the payer's current period under a share lock on the subscription row, so a closed period is never topped up and the whole run is invoiced exactly once. Threshold settlement follows the stamped period. - Sync the worker's billing contract (usageExceeded, UsageUpgrade, UsageLimitRefusal). --- .../app/api/billing/update-cost/route.test.ts | 233 ++++++++++++++++++ apps/sim/app/api/billing/update-cost/route.ts | 75 +++++- .../copilot/api-keys/validate/route.test.ts | 78 +++++- .../api/copilot/api-keys/validate/route.ts | 57 ++++- .../special-tags/special-tags.test.ts | 17 ++ apps/sim/lib/api/contracts/copilot.ts | 46 +++- apps/sim/lib/api/contracts/subscription.ts | 30 ++- .../calculations/usage-monitor.test.ts | 15 ++ .../lib/billing/calculations/usage-monitor.ts | 6 + .../billing/core/billing-attribution.test.ts | 23 ++ .../lib/billing/core/billing-attribution.ts | 7 + apps/sim/lib/billing/core/mid-run-usage.ts | 48 ++++ apps/sim/lib/billing/core/usage-analytics.ts | 4 + .../lib/billing/core/usage-log.integration.ts | 114 ++++++++- apps/sim/lib/billing/core/usage-log.test.ts | 4 +- apps/sim/lib/billing/core/usage-log.ts | 207 +++++++++++++--- apps/sim/lib/billing/usage-upgrade.ts | 67 +++++ apps/sim/lib/mothership/generated/billing.ts | 18 ++ .../lib/mothership/request/go/stream.test.ts | 32 +++ apps/sim/lib/mothership/request/go/stream.ts | 20 ++ .../request/lifecycle/admission.test.ts | 52 ++++ .../mothership/request/lifecycle/admission.ts | 14 +- .../mothership/request/lifecycle/run.test.ts | 72 +++++- .../lib/mothership/request/lifecycle/run.ts | 51 +++- .../mothership/request/tools/billing.test.ts | 11 + .../lib/mothership/request/tools/billing.ts | 59 +---- 26 files changed, 1239 insertions(+), 121 deletions(-) create mode 100644 apps/sim/lib/billing/core/mid-run-usage.ts create mode 100644 apps/sim/lib/billing/usage-upgrade.ts diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index dd89be51224..4cbc6f66fb8 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -47,6 +47,7 @@ vi.mock('@/lib/billing/threshold-billing', () => ({ ThresholdSettlementError: MockThresholdSettlementError, })) +import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { BillingCallbackBody, BillingCallbackHeaders, @@ -69,6 +70,7 @@ const mockRequireBillingAttributionHeader = const mockResolveLegacyV0BillingAttribution = billingAttributionMockFns.mockResolveLegacyV0BillingAttribution const mockToBillingContext = billingAttributionMockFns.mockToBillingContext +const mockCheckAttributedUsageLimits = billingAttributionMockFns.mockCheckAttributedUsageLimits afterAll(resetEnvFlagsMock) @@ -847,3 +849,234 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => { expect(mockRecordCumulativeUsage).not.toHaveBeenCalled() }) }) + +describe('POST /api/billing/update-cost — mid-run usage gate', () => { + let callbackSequence = 0 + + function attributedCallback() { + callbackSequence += 1 + const billingRequestId = `0190c03f-9f7d-4b79-8b58-${String(callbackSequence).padStart(12, '0')}` + return createMockRequest( + 'POST', + { + userId: 'user-1', + cost: 0.5 * callbackSequence, + model: 'claude-opus-4.8', + source: 'workspace-chat', + workspaceId: 'ws-1', + idempotencyKey: billingRequestId, + }, + { + 'x-api-key': 'internal', + 'x-sim-billing-protocol': 'attribution-v1', + 'x-sim-billing-request-id': billingRequestId, + 'x-sim-billing-attribution': 'serialized-attribution', + } + ) + } + + beforeEach(() => { + resetUsageGateCache() + setEnvFlags({ isBillingEnabled: true, isHosted: true }) + mockCheckInternalApiKey.mockReturnValue({ success: true }) + mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 }) + mockCheckAndBillPayerOverageThreshold.mockResolvedValue(undefined) + mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockToBillingContext.mockReturnValue({ + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { + start: new Date('2026-07-01T00:00:00.000Z'), + end: new Date('2026-08-01T00:00:00.000Z'), + }, + }) + }) + + it('tells the worker when the run payer has crossed its usage limit', async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const res = await POST(attributedCallback()) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toMatchObject({ + success: true, + usageExceeded: true, + usageUpgrade: { + reason: 'usage_limit', + action: 'upgrade_plan', + message: expect.stringContaining('usage limit'), + }, + }) + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(ATTRIBUTION) + }) + + it('offers a paid organization payer the increase-limit card', async () => { + mockRequireBillingAttributionHeader.mockReturnValue({ + ...ATTRIBUTION, + payerSubscription: { id: 'sub-1', plan: 'team', status: 'active', seats: 4 }, + }) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const res = await POST(attributedCallback()) + + const body = await res.json() + expect(body.usageUpgrade).toMatchObject({ + action: 'increase_limit', + message: expect.stringContaining('organization'), + }) + }) + + it('serves a cached admission to every step and re-reads a refusal', async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + + for (let step = 0; step < 5; step++) { + const body = await (await POST(attributedCallback())).json() + expect(body.usageExceeded).toBe(false) + expect(body).not.toHaveProperty('usageUpgrade') + } + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(1) + + resetUsageGateCache() + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + await POST(attributedCallback()) + await POST(attributedCallback()) + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(3) + }) + + it('answers a duplicate retry with the verdict its lost first answer carried', async () => { + mockRecordCumulativeUsage.mockResolvedValue({ billed: false, delta: 0, total: 0.5 }) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const res = await POST(attributedCallback()) + + expect(res.status).toBe(409) + await expect(res.json()).resolves.toMatchObject({ + code: 'DUPLICATE_BILLING_EVENT', + usageExceeded: true, + usageUpgrade: { action: 'upgrade_plan' }, + }) + }) + + describe('a run that outlives its billing period', () => { + const PAYER_SUBSCRIPTION = { + id: 'sub-1', + plan: 'team', + status: 'active', + seats: 4, + } + const CURRENT_PERIOD = { + start: new Date('2026-08-01T00:00:00.000Z'), + end: new Date('2026-09-01T00:00:00.000Z'), + } + + beforeEach(() => { + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + mockRequireBillingAttributionHeader.mockReturnValue({ + ...ATTRIBUTION, + payerSubscription: PAYER_SUBSCRIPTION, + }) + }) + + it('settles the threshold against the period the charge was stamped into', async () => { + mockRecordCumulativeUsage.mockResolvedValue({ + billed: true, + delta: 0.5, + total: 1.5, + billingPeriod: CURRENT_PERIOD, + }) + + const res = await POST(attributedCallback()) + + expect(res.status).toBe(200) + expect(mockRecordCumulativeUsage).toHaveBeenCalledWith( + expect.objectContaining({ payerSubscriptionId: 'sub-1' }) + ) + expect(mockCheckAndBillPayerOverageThreshold).toHaveBeenCalledWith( + { type: 'organization', id: 'org-1' }, + expect.objectContaining({ expectedBillingPeriod: CURRENT_PERIOD }) + ) + }) + + it('leaves a period that closed under a recorded charge to the cycle close', async () => { + mockRecordCumulativeUsage.mockResolvedValue({ + billed: true, + delta: 0.5, + total: 1.5, + billingPeriod: { + start: new Date('2026-07-01T00:00:00.000Z'), + end: new Date('2026-08-01T00:00:00.000Z'), + }, + }) + mockCheckAndBillPayerOverageThreshold.mockRejectedValue( + new MockThresholdSettlementError('billing_period_elapsed') + ) + + const res = await POST(attributedCallback()) + + expect(res.status).toBe(200) + }) + + it('keeps reporting-window payers on their frozen period', async () => { + mockToBillingContext.mockReturnValue({ + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { + start: new Date('2026-07-01T00:00:00.000Z'), + end: new Date('2026-08-01T00:00:00.000Z'), + source: 'reporting', + }, + }) + mockRecordCumulativeUsage.mockResolvedValue({ + billed: true, + delta: 0.5, + total: 0.5, + billingPeriod: CURRENT_PERIOD, + }) + + await POST(attributedCallback()) + + expect(mockRecordCumulativeUsage).toHaveBeenCalledWith( + expect.not.objectContaining({ payerSubscriptionId: expect.anything() }) + ) + }) + }) + + it.each([ + ['an unreadable ledger', { isExceeded: true, reason: 'usage_unavailable' }], + ['a blocked account', { isExceeded: true, reason: 'billing_blocked', scope: 'payer' }], + ])('does not pause a run for %s', async (_case, verdict) => { + mockCheckAttributedUsageLimits.mockResolvedValue(verdict) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(false) + expect(body).not.toHaveProperty('usageUpgrade') + }) + + it('tells a member over the cap their organization set who can raise it', async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'member' }) + + const body = await (await POST(attributedCallback())).json() + + expect(body).toMatchObject({ + usageUpgrade: { message: expect.stringMatching(/limit your organization set for you/) }, + }) + }) + + it('keeps a recorded charge successful when the gate read fails', async () => { + mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out')) + + const res = await POST(attributedCallback()) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toMatchObject({ usageExceeded: false }) + expect(mockRecordCumulativeUsage).toHaveBeenCalledOnce() + }) + + it('reports no exceeded usage when billing is disabled', async () => { + setEnvFlags({ isBillingEnabled: false, isHosted: true }) + + const res = await POST(attributedCallback()) + + await expect(res.json()).resolves.toMatchObject({ usageExceeded: false }) + expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index 0f789f87c99..4c0fe4789ae 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -2,7 +2,11 @@ import type { Span } from '@opentelemetry/api' import { createLogger } from '@sim/logger' import { getPostgresConstraintName, getPostgresErrorCode, toError } from '@sim/utils/errors' import { type NextRequest, NextResponse } from 'next/server' -import { billingUpdateCostContract } from '@/lib/api/contracts/subscription' +import { + type BillingUpdateCostResponse, + type BillingUsageVerdict, + billingUpdateCostContract, +} from '@/lib/api/contracts/subscription' import { parseRequest } from '@/lib/api/server' import { type AccountBillingDecision, @@ -18,6 +22,7 @@ import { resolveLegacyV0BillingAttribution, toBillingContext, } from '@/lib/billing/core/billing-attribution' +import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' import { type CumulativeUsageContextField, CumulativeUsageContextMismatchError, @@ -28,6 +33,7 @@ import { checkAndBillPayerOverageThreshold, ThresholdSettlementError, } from '@/lib/billing/threshold-billing' +import { resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' import { generateRequestId } from '@/lib/core/utils/request' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' @@ -57,6 +63,30 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp ) } +/** + * Reads the run payer's standing after a cost callback, so a long run stops at its next step + * once it crosses the limit instead of at its next admission, with the card the worker writes + * to its log. A duplicate callback answers too: it is often a retry whose first answer was lost. + * Served from the execution usage gate: an admission is cached per payer and actor for the gate + * TTL and a refusal is always re-read, so steady-state steps cost no ledger read. The charge is + * already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the + * refusal to the next step or re-check rather than ending a paying run on a database blip. + */ +async function readUsageStanding( + userId: string, + billingAttribution: BillingAttributionSnapshot | undefined +): Promise { + if (!isHosted || !billingAttribution) return { usageExceeded: false } + const verdict = await readMidRunUsageVerdict(billingAttribution) + // Only a spent limit pauses the run. A blocked account is refused at the run's next + // continuation or re-check, with blocked-account copy rather than the upgrade card. + if (verdict.status !== 'exceeded') return { usageExceeded: false } + return { + usageExceeded: true, + usageUpgrade: await resolveUsageUpgradePayload(userId, billingAttribution, verdict.scope), + } +} + function getBillingResolution( isMarkerlessLegacy: boolean, billingAttribution: BillingAttributionSnapshot | undefined @@ -112,9 +142,10 @@ async function updateCostInner(req: NextRequest, span: Span): Promise({ success: true, message: 'Billing disabled, cost update skipped', + usageExceeded: false, data: { billingEnabled: false, processedAt: new Date().toISOString(), @@ -312,6 +343,13 @@ async function updateCostInner(req: NextRequest, span: Span): Promise({ success: true, + ...usageVerdict, data: { userId, cost, diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 48118441140..a10d2e1de21 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -138,6 +138,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) vi.mock('@/lib/workspaces/utils', () => workspacesUtilsMock) import { validateCopilotApiKeyBodySchema } from '@/lib/api/contracts/copilot' +import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { POST } from '@/app/api/copilot/api-keys/validate/route' const { mockGetWorkspaceBillingSettings } = workspacesUtilsMockFns @@ -506,7 +507,9 @@ describe('validation lifecycle purposes', () => { mockCheckInternalApiKey.mockReturnValue({ success: true }) mockAuthorizeCallback.mockReset().mockResolvedValue(undefined) mockCheckContinuationBilling.mockReset().mockResolvedValue({ blocked: false }) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockIsEnterprisePlan.mockResolvedValue(false) + resetUsageGateCache() }) it('defaults older callers to full admission and rejects unknown purposes', () => { @@ -516,7 +519,7 @@ describe('validation lifecycle purposes', () => { ).toBe(false) }) - it('checks original payer and current scope without repeating spend admission', async () => { + it('checks original payer, current scope, and the original payer spend', async () => { const response = await POST(request(body, attributedHeaders)) expect(response.status).toBe(200) expect(mockAuthorizeCallback).toHaveBeenCalledWith({ ...body, delegationId: requestId }) @@ -527,7 +530,7 @@ describe('validation lifecycle purposes', () => { expect(mockAuthorizeCallback.mock.invocationCallOrder[0]).toBeLessThan( mockCheckContinuationBilling.mock.invocationCallOrder[0] ) - expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(ATTRIBUTION) expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() expect(mockResolveLegacyV0BillingAttribution).not.toHaveBeenCalled() expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled() @@ -638,11 +641,76 @@ describe('validation lifecycle purposes', () => { }) it.each(['actor', 'payer'])('refuses a newly blocked %s on continuation', async (scope) => { - mockCheckContinuationBilling.mockResolvedValueOnce({ blocked: true, scope }) - expect((await POST(request(body, attributedHeaders))).status).toBe(402) + mockCheckContinuationBilling.mockResolvedValueOnce({ + blocked: true, + scope, + message: 'Billing account frozen.', + }) + const response = await POST(request(body, attributedHeaders)) + expect(response.status).toBe(402) + await expect(response.json()).resolves.toEqual({ + code: 'BILLING_BLOCKED', + error: 'Billing account frozen.', + }) expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() }) + it('refuses a payer the usage gate finds blocked as blocked, without the usage card', async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ + isExceeded: true, + reason: 'billing_blocked', + message: 'Organization billing issue.', + scope: 'payer', + }) + const response = await POST(request(body, attributedHeaders)) + expect(response.status).toBe(402) + await expect(response.json()).resolves.toEqual({ + code: 'BILLING_BLOCKED', + error: 'Organization billing issue.', + }) + }) + + it('admits a polled continuation whose spend cannot be read', async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ + isExceeded: true, + reason: 'usage_unavailable', + }) + expect((await POST(request(body, attributedHeaders))).status).toBe(200) + mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out')) + expect((await POST(request(body, attributedHeaders))).status).toBe(200) + }) + + it('refuses a continuation over its usage limit with the card the worker writes', async () => { + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const response = await POST(request(body, attributedHeaders)) + + expect(response.status).toBe(402) + await expect(response.json()).resolves.toEqual({ + code: 'USAGE_LIMIT_EXCEEDED', + error: expect.stringContaining('usage limit'), + usageUpgrade: { + reason: 'usage_limit', + action: 'upgrade_plan', + message: expect.stringContaining('usage limit'), + }, + }) + }) + + it('answers a polled re-check from the cached admission and always re-reads a refusal', async () => { + for (let call = 0; call < 2; call++) queueTableRows(schemaMock.user, [{ id: 'user-1' }]) + for (let poll = 0; poll < 3; poll++) { + expect((await POST(request(body, attributedHeaders))).status).toBe(200) + } + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(1) + + resetUsageGateCache() + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + expect((await POST(request(body, attributedHeaders))).status).toBe(402) + expect((await POST(request(body, attributedHeaders))).status).toBe(402) + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(3) + }) + it('allows cancellation without billing material or spending/standing/plan checks', async () => { const response = await POST( request({ ...body, purpose: 'cancellation' }, { 'x-sim-billing-protocol': 'attribution-v1' }) @@ -752,7 +820,7 @@ describe('validation lifecycle purposes', () => { expect((await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders))).status).toBe( 402 ) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(1) + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(2) expect((await POST(request({ ...body, purpose: 'new-turn' }, directHeaders))).status).toBe(400) expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() }) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 0b431ed12d4..d21004d3e46 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -4,7 +4,13 @@ import { createLogger } from '@sim/logger' import { generateId } from '@sim/utils/id' import { eq } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' -import { validateCopilotApiKeyContract } from '@/lib/api/contracts/copilot' +import { + COPILOT_BILLING_BLOCKED_CODE, + COPILOT_USAGE_LIMIT_EXCEEDED_CODE, + type ValidateCopilotApiKeyBillingBlocked, + type ValidateCopilotApiKeyUsageExceeded, + validateCopilotApiKeyContract, +} from '@/lib/api/contracts/copilot' import { parseRequest, validationErrorResponse } from '@/lib/api/server' import { checkServerSideUsageLimits } from '@/lib/billing/calculations/usage-monitor' import { @@ -20,9 +26,11 @@ import { serializeAccountBillingDecisionHeader, serializeBillingAttributionHeader, } from '@/lib/billing/core/billing-attribution' +import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { isEnterprisePlan } from '@/lib/billing/core/subscription' import { deriveBillingContext } from '@/lib/billing/core/usage-log' +import { resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' import { asOrchestrationError } from '@/lib/core/orchestration/types' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' @@ -53,6 +61,8 @@ import { withIncomingGoSpan } from '@/lib/mothership/request/otel' const logger = createLogger('CopilotApiKeysValidate') +const CONTINUATION_BLOCKED_MESSAGE = 'Continuation billing account is blocked' + function invalidBillingProtocolResponse(): NextResponse { return NextResponse.json({ error: 'Invalid billing attribution protocol' }, { status: 400 }) } @@ -402,13 +412,54 @@ export const POST = withRouteHandler((req: NextRequest) => blocked: blocked?.blocked ?? false, elapsedMs: Math.round(performance.now() - startedAt), }) - if (blocked?.blocked) { + // A continuation, and a worker's periodic re-check of a long run, also reads the + // original payer's spend through the cached execution usage gate. A read that fails + // admits: the run is already under way, and the next re-check reads again. + const verdict = + !blocked?.blocked && + purpose === COPILOT_VALIDATION_PURPOSE.continuation && + billing?.kind === 'attributed' + ? await readMidRunUsageVerdict(billing.attribution) + : null + if (blocked?.blocked || verdict?.status === 'blocked') { span.setAttribute( TraceAttr.CopilotValidateOutcome, CopilotValidateOutcome.UsageExceeded ) span.setAttribute(TraceAttr.HttpStatusCode, 402) - return new NextResponse(null, { status: 402 }) + return NextResponse.json( + { + code: COPILOT_BILLING_BLOCKED_CODE, + error: + (blocked?.blocked + ? blocked.message + : verdict?.status === 'blocked' + ? verdict.message + : undefined) ?? CONTINUATION_BLOCKED_MESSAGE, + }, + { status: 402 } + ) + } + if (verdict?.status === 'exceeded' && billing?.kind === 'attributed') { + logger.info('[API VALIDATION] Continuation usage exceeded', { userId }) + span.setAttribute( + TraceAttr.CopilotValidateOutcome, + CopilotValidateOutcome.UsageExceeded + ) + span.setAttribute(TraceAttr.HttpStatusCode, 402) + const usageUpgrade = await resolveUsageUpgradePayload( + userId, + billing.attribution, + verdict.scope + ) + return NextResponse.json( + { + code: COPILOT_USAGE_LIMIT_EXCEEDED_CODE, + error: usageUpgrade.message, + usageUpgrade, + }, + { status: 402 } + ) } const isEnterprise = purpose === COPILOT_VALIDATION_PURPOSE.cancellation diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts index 5e8d4273732..ed3296e6f14 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts @@ -10,6 +10,7 @@ import { describe, expect, it, vi } from 'vitest' */ vi.mock('@/lib/auth/auth-client', () => authClientMock) +import { formatUsageUpgradeTag } from '@/lib/billing/usage-upgrade' import type { ContentSegment, CredentialItemData, @@ -875,3 +876,19 @@ describe('source tag', () => { } }) }) + +describe('usage card written to a worker log', () => { + it('renders the card Sim hands the worker when the text is replayed after a reload', () => { + const usageUpgrade = { + reason: 'usage_limit', + action: 'increase_limit', + message: "You've reached your usage limit for this billing period.", + } as const + const replayed = `Finished the first report.${formatUsageUpgradeTag(usageUpgrade)}` + + const { segments } = parseSpecialTags(replayed, false) + + expect(segments).toContainEqual({ type: 'usage_upgrade', data: usageUpgrade }) + expect(renderedText(segments)).toBe('Finished the first report.') + }) +}) diff --git a/apps/sim/lib/api/contracts/copilot.ts b/apps/sim/lib/api/contracts/copilot.ts index d9fcd2ff2a5..37316242a31 100644 --- a/apps/sim/lib/api/contracts/copilot.ts +++ b/apps/sim/lib/api/contracts/copilot.ts @@ -3,6 +3,7 @@ import { persistedContentBlockSchema } from '@/lib/api/contracts/copilot-message import { workspaceSearchFiltersSchema } from '@/lib/api/contracts/knowledge/search' import { mothershipResourceSchema } from '@/lib/api/contracts/mothership-resources' import { requiredFieldSchema, workspaceIdSchema } from '@/lib/api/contracts/primitives' +import { usageUpgradePayloadSchema } from '@/lib/api/contracts/subscription' import { type ContractJsonResponse, defineRouteContract } from '@/lib/api/contracts/types' import { ASYNC_TOOL_CONFIRMATION_STATUS, @@ -270,6 +271,41 @@ export const validateCopilotApiKeyResponseSchema = z.object({ }) export type ValidateCopilotApiKeyResponse = z.output +export const COPILOT_USAGE_LIMIT_EXCEEDED_CODE = 'USAGE_LIMIT_EXCEEDED' +export const COPILOT_BILLING_BLOCKED_CODE = 'BILLING_BLOCKED' + +/** + * A continuation refused because the run's original payer is over its usage limit. A worker + * polling continuation validation mid-run writes `usageUpgrade` as a `` tag into + * its log and pauses the run for the limit. + */ +export const validateCopilotApiKeyUsageExceededSchema = z.object({ + code: z.literal(COPILOT_USAGE_LIMIT_EXCEEDED_CODE), + error: z.string(), + usageUpgrade: usageUpgradePayloadSchema, +}) +export type ValidateCopilotApiKeyUsageExceeded = z.output< + typeof validateCopilotApiKeyUsageExceededSchema +> + +/** A continuation refused because the actor or payer account is blocked (payment, dispute). */ +export const validateCopilotApiKeyBillingBlockedSchema = z.object({ + code: z.literal(COPILOT_BILLING_BLOCKED_CODE), + error: z.string(), +}) +export type ValidateCopilotApiKeyBillingBlocked = z.output< + typeof validateCopilotApiKeyBillingBlockedSchema +> + +/** + * A 402 from continuation validation. New-turn admission still answers 402 with no body, which + * older callers treat as a refusal of either kind. + */ +export const validateCopilotApiKeyRefusalSchema = z.union([ + validateCopilotApiKeyUsageExceededSchema, + validateCopilotApiKeyBillingBlockedSchema, +]) + export const listCopilotApiKeysContract = defineRouteContract({ method: 'GET', path: '/api/copilot/api-keys', @@ -392,7 +428,15 @@ export const validateCopilotApiKeyContract = defineRouteContract({ path: '/api/copilot/api-keys/validate', headers: validateCopilotApiKeyHeadersSchema, body: validateCopilotApiKeyBodySchema, - response: { mode: 'json', schema: validateCopilotApiKeyResponseSchema }, + response: { + mode: 'json', + schema: validateCopilotApiKeyResponseSchema, + status: [200, 402], + statusSchemas: { + 200: validateCopilotApiKeyResponseSchema, + 402: validateCopilotApiKeyRefusalSchema, + }, + }, error: validateCopilotApiKeyErrorSchema, }) diff --git a/apps/sim/lib/api/contracts/subscription.ts b/apps/sim/lib/api/contracts/subscription.ts index 4d1efa6dd63..1d765f5c615 100644 --- a/apps/sim/lib/api/contracts/subscription.ts +++ b/apps/sim/lib/api/contracts/subscription.ts @@ -325,7 +325,34 @@ export const billingSwitchPlanResponseSchema = z.object({ message: z.string().optional(), }) -export const billingUpdateCostResponseSchema = z.object({ +/** + * The upgrade card's payload: the JSON body of a `` tag in assistant text, which + * the chat renders as the usage card wherever that text appears (live, replayed, or reloaded). + * Sim decides the action and copy from the payer's plan; a worker ending a run at the usage + * limit writes the tag with this payload verbatim into its durable log. + */ +export const usageUpgradePayloadSchema = z.object({ + reason: z.literal('usage_limit'), + action: z.enum(['upgrade_plan', 'increase_limit']), + message: z.string(), +}) +export type UsageUpgradePayload = z.infer + +/** + * The payer's standing after a cost callback, read through the cached execution usage gate. It + * sits at the top level of the body, beside `success`, where the worker's shared + * `BillingCallbackResult` reads it, on a 200 and on a duplicate 409 alike. A worker that + * predates the fields ignores them. + */ +export const billingUsageVerdictSchema = z.object({ + /** The payer is over its usage limit; the worker pauses the run at its next step boundary. */ + usageExceeded: z.boolean(), + /** Present exactly when `usageExceeded` is true: the card the worker writes to its log. */ + usageUpgrade: usageUpgradePayloadSchema.optional(), +}) +export type BillingUsageVerdict = z.infer + +export const billingUpdateCostResponseSchema = billingUsageVerdictSchema.extend({ success: z.literal(true), message: z.string().optional(), data: z.object({ @@ -336,6 +363,7 @@ export const billingUpdateCostResponseSchema = z.object({ requestId: z.string(), }), }) +export type BillingUpdateCostResponse = z.infer export const billingSwitchPlanContract = defineRouteContract({ method: 'POST', diff --git a/apps/sim/lib/billing/calculations/usage-monitor.test.ts b/apps/sim/lib/billing/calculations/usage-monitor.test.ts index 7571686ee5b..4bd444b3b90 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.test.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.test.ts @@ -130,6 +130,21 @@ describe('checkUsageStatus', () => { }) }) + it('refuses on a ledger read failure but marks the answer unavailable', async () => { + mockGetBillingPeriodUsageCost.mockRejectedValueOnce(new Error('canceling statement')) + + await expect( + checkUsageStatus('user-1', { + referenceId: 'user-1', + plan: 'free', + status: 'active', + seats: 1, + periodStart: new Date('2026-06-01T00:00:00.000Z'), + periodEnd: new Date('2026-07-01T00:00:00.000Z'), + }) + ).resolves.toMatchObject({ isExceeded: true, unavailable: true }) + }) + it('preserves negative ledger-only personal usage', async () => { const periodStart = new Date('2026-06-01T00:00:00.000Z') const periodEnd = new Date('2026-07-01T00:00:00.000Z') diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index f4ebe5bc66b..318ecc366d2 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.ts @@ -43,6 +43,11 @@ interface UsageData { scope: 'user' | 'organization' /** Present only when `scope === 'organization'`. */ organizationId: string | null + /** + * The ledger could not be read, so `isExceeded` is a fail-closed refusal rather than a + * measured one. Admission refuses on it; a run already under way treats it as unknown. + */ + unavailable?: true } /** @@ -183,6 +188,7 @@ export async function checkUsageStatus( limit: 0, scope: 'user', organizationId: null, + unavailable: true, } } } diff --git a/apps/sim/lib/billing/core/billing-attribution.test.ts b/apps/sim/lib/billing/core/billing-attribution.test.ts index 5cec0fa1918..bf4f3ddb521 100644 --- a/apps/sim/lib/billing/core/billing-attribution.test.ts +++ b/apps/sim/lib/billing/core/billing-attribution.test.ts @@ -432,6 +432,29 @@ describe('checkAttributedUsageLimits', () => { expect(mockCheckOrganizationMemberUsageLimit).not.toHaveBeenCalled() }) + it('names a billing block and an unreadable ledger apart from a spent limit', async () => { + mockCheckBillingBlocked.mockResolvedValueOnce({ blocked: true, message: 'Frozen.' }) + await expect(checkAttributedUsageLimits(attribution)).resolves.toMatchObject({ + isExceeded: true, + reason: 'billing_blocked', + }) + + mockCheckUsageStatus.mockResolvedValueOnce({ + currentUsage: 0, + isExceeded: true, + limit: 0, + organizationId: null, + percentUsed: 100, + isWarning: false, + scope: 'user', + unavailable: true, + }) + await expect(checkAttributedUsageLimits(attribution)).resolves.toMatchObject({ + isExceeded: true, + reason: 'usage_unavailable', + }) + }) + it('returns payer exhaustion before checking the actor member cap', async () => { mockCheckUsageStatus.mockResolvedValue({ currentUsage: 100, diff --git a/apps/sim/lib/billing/core/billing-attribution.ts b/apps/sim/lib/billing/core/billing-attribution.ts index 6c205ca33bd..f8211c0eb07 100644 --- a/apps/sim/lib/billing/core/billing-attribution.ts +++ b/apps/sim/lib/billing/core/billing-attribution.ts @@ -111,6 +111,11 @@ export interface AttributedUsageLimitsResult { isExceeded: boolean message?: string scope?: 'actor' | 'payer' | 'member' + /** + * Why an `isExceeded` refusal is not a spent limit: the account is blocked (payment failed, + * dispute), or the payer's usage could not be read and the gate failed closed. + */ + reason?: 'billing_blocked' | 'usage_unavailable' payerUsage?: { currentUsage: number limit: number @@ -919,6 +924,7 @@ export async function checkAttributedUsageLimits( isExceeded: true, message: billingBlock.message, scope: billingBlock.scope, + reason: 'billing_blocked', } } @@ -944,6 +950,7 @@ export async function checkAttributedUsageLimits( message, scope: 'payer', payerUsage: payerSnapshot, + ...(payerUsage.unavailable ? { reason: 'usage_unavailable' as const } : {}), } } diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts new file mode 100644 index 00000000000..e6d39178f44 --- /dev/null +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -0,0 +1,48 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import type { + AttributedUsageLimitsResult, + BillingAttributionSnapshot, +} from '@/lib/billing/core/billing-attribution' +import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache' + +const logger = createLogger('MidRunUsage') + +/** + * A run's standing while it is under way, read through the execution usage gate: + * - `exceeded`: the payer (or the actor's member cap) spent its limit; the run pauses with the + * upgrade card. + * - `blocked`: the account is blocked (payment failed, dispute); the run is refused as a blocked + * account, never with the upgrade card. + * - `unknown`: the gate could not read usage. Admission fails closed on this, but a run already + * under way continues: a database blip must not end a paying user's long run, and the next + * step or re-check reads again. + */ +export type MidRunUsageVerdict = + | { status: 'within' } + | { status: 'exceeded'; scope?: AttributedUsageLimitsResult['scope'] } + | { status: 'blocked'; message?: string } + | { status: 'unknown' } + +export async function readMidRunUsageVerdict( + attribution: BillingAttributionSnapshot +): Promise { + let usage: AttributedUsageLimitsResult + try { + usage = await checkExecutionUsageLimits(attribution) + } catch (error) { + logger.warn('Mid-run usage read failed; continuing the run', { + error: getErrorMessage(error), + }) + return { status: 'unknown' } + } + if (!usage.isExceeded) return { status: 'within' } + if (usage.reason === 'billing_blocked') { + return { status: 'blocked', ...(usage.message ? { message: usage.message } : {}) } + } + if (usage.reason === 'usage_unavailable') { + logger.warn('Mid-run usage could not be read; continuing the run') + return { status: 'unknown' } + } + return { status: 'exceeded', ...(usage.scope ? { scope: usage.scope } : {}) } +} diff --git a/apps/sim/lib/billing/core/usage-analytics.ts b/apps/sim/lib/billing/core/usage-analytics.ts index 92e6b50f9c7..60d5ff1a1e9 100644 --- a/apps/sim/lib/billing/core/usage-analytics.ts +++ b/apps/sim/lib/billing/core/usage-analytics.ts @@ -513,6 +513,10 @@ export function usageBucketTimestamps( * cost in place for as long as its stream runs — which {@link STREAM_TIMEOUT_MS} * caps — plus the retry flushes that follow it. Past the cap and this margin a day or * hour can no longer change and is treated as settled. + * + * Without a run deadline a Chat turn can top up its row for longer than that, so a + * settled hour's cached aggregate can under-report that turn's later spend. This is + * display only: invoices, threshold billing, and the usage gate read live ledger sums. */ export const USAGE_SETTLE_MS = STREAM_TIMEOUT_MS + 2 * 60 * 60 * 1000 diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index 7b8da8fb23f..c34361ff622 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -26,6 +26,7 @@ import { CumulativeUsageContextMismatchError, getBillingPeriodUsageCost, getBillingPeriodUsageCostByUser, + getStampedPeriodRangeUsageCostByUser, type RecordCumulativeUsageParams, recordCumulativeUsage, } from '@/lib/billing/core/usage-log' @@ -119,7 +120,8 @@ describe('Cumulative billing with PostgreSQL', () => { ); CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key) WHERE event_key IS NOT NULL; - CREATE TABLE driver_probe (id text PRIMARY KEY) + CREATE TABLE driver_probe (id text PRIMARY KEY); + CREATE TABLE subscription (id text PRIMARY KEY, period_start timestamp, period_end timestamp) `) transaction.mockImplementation(async (callback: (tx: Transaction) => Promise) => { const pause = nextPause @@ -142,6 +144,7 @@ describe('Cumulative billing with PostgreSQL', () => { beforeEach(async () => { nextPause = undefined await connection`truncate usage_log` + await connection`truncate subscription` }) it.each([ @@ -216,12 +219,12 @@ describe('Cumulative billing with PostgreSQL', () => { expect(recovered.billed).toBe(true) expect(recovered.delta).toBeCloseTo(0.8 - initial, 9) expect(recovered.total).toBe(0.8) - expect(await recordCumulativeUsage(usage(0.8))).toEqual({ + expect(await recordCumulativeUsage(usage(0.8))).toMatchObject({ billed: false, delta: 0, total: 0.8, }) - expect(await recordCumulativeUsage(usage(0.3))).toEqual({ + expect(await recordCumulativeUsage(usage(0.3))).toMatchObject({ billed: false, delta: 0, total: 0.8, @@ -245,7 +248,11 @@ describe('Cumulative billing with PostgreSQL', () => { ) ) expect(await ledgerRows()).toHaveLength(33) - expect(await recordCumulativeUsage(usage(0.8))).toEqual({ billed: false, delta: 0, total: 0.8 }) + expect(await recordCumulativeUsage(usage(0.8))).toMatchObject({ + billed: false, + delta: 0, + total: 0.8, + }) }) it('reads committed pooled and member charges freshly after concurrent executions', async () => { @@ -305,4 +312,103 @@ describe('Cumulative billing with PostgreSQL', () => { expect(await ledgerRows()).toEqual([{ event_key: usage(0).eventKey, cost: '0.8' }]) } ) + + describe('a request that outlives its billing period', () => { + const periods = [ + new Date('2026-09-01T00:00:00.000Z'), + new Date('2026-10-01T00:00:00.000Z'), + new Date('2026-11-01T00:00:00.000Z'), + new Date('2026-12-01T00:00:00.000Z'), + ] + const payer = { type: 'organization', id: 'payer' } as const + + async function setSubscriptionPeriod(index: number) { + await connection` + insert into subscription (id, period_start, period_end) + values ('sub-1', ${periods[index].toISOString()}::timestamptz at time zone 'UTC', ${periods[index + 1].toISOString()}::timestamptz at time zone 'UTC') + on conflict (id) do update + set period_start = excluded.period_start, period_end = excluded.period_end + ` + } + + function charge(cost: number) { + return recordCumulativeUsage({ ...usage(cost), payerSubscriptionId: 'sub-1' }) + } + + /** What the cycle close invoices for one period: the ledger rows stamped with it. */ + async function stampedTotal(index: number) { + const byUser = await getStampedPeriodRangeUsageCostByUser( + payer, + { from: periods[index], to: periods[index + 1] }, + undefined, + database + ) + return [...byUser.values()].reduce((total, cost) => total + cost, 0) + } + + it('invoices a charge that spans a period close exactly once in total', async () => { + await setSubscriptionPeriod(0) + expect(await charge(0.4)).toMatchObject({ billed: true, total: 0.4 }) + + await setSubscriptionPeriod(1) + const closedTotal = await stampedTotal(0) + expect(closedTotal).toBeCloseTo(0.4, 9) + + const afterClose = await charge(1) + expect(afterClose).toMatchObject({ billed: true, total: 1 }) + expect(afterClose.billingPeriod).toEqual({ start: periods[1], end: periods[2] }) + expect(await charge(0.9)).toMatchObject({ billed: false, total: 1 }) + expect(await charge(1.3)).toMatchObject({ billed: true, total: 1.3 }) + expect(await charge(1.3)).toMatchObject({ billed: false, total: 1.3 }) + + await setSubscriptionPeriod(2) + expect(await charge(1.5)).toMatchObject({ billed: true, total: 1.5 }) + + expect(await stampedTotal(0)).toBeCloseTo(closedTotal, 9) + expect(await stampedTotal(1)).toBeCloseTo(0.9, 9) + expect(await stampedTotal(2)).toBeCloseTo(0.2, 9) + const invoiced = (await stampedTotal(0)) + (await stampedTotal(1)) + (await stampedTotal(2)) + expect(invoiced).toBeCloseTo(1.5, 9) + }) + + it('never stamps a charge into a period earlier than its latest row', async () => { + await setSubscriptionPeriod(0) + await charge(0.4) + await setSubscriptionPeriod(1) + await charge(1) + await setSubscriptionPeriod(0) + expect(await charge(1.2)).toMatchObject({ billed: true, total: 1.2 }) + expect(await stampedTotal(0)).toBeCloseTo(0.4, 9) + expect(await stampedTotal(1)).toBeCloseTo(0.8, 9) + }) + + it('stamps a first charge that lands after the close into the current period', async () => { + await setSubscriptionPeriod(1) + expect(await charge(0.7)).toMatchObject({ billed: true, total: 0.7 }) + expect(await charge(0.9)).toMatchObject({ billed: true, total: 0.9 }) + expect(await stampedTotal(0)).toBe(0) + expect(await stampedTotal(1)).toBeCloseTo(0.9, 9) + }) + + it('holds the period advance until an in-flight top-up of the old period commits', async () => { + await setSubscriptionPeriod(0) + await charge(0.4) + const pause = pauseNextTransaction() + const inFlight = charge(0.6) + try { + await pause.reached.promise + const advance = await connection + .begin(async (tx) => { + await tx`select set_config('lock_timeout', '300ms', true)` + await tx`update subscription set period_start = ${periods[1].toISOString()}::timestamptz at time zone 'UTC' where id = 'sub-1'` + }) + .catch((error: unknown) => error) + expect(getPostgresErrorCode(advance)).toBe('55P03') + } finally { + pause.release.resolve() + await inFlight + } + expect(await stampedTotal(0)).toBeCloseTo(0.6, 9) + }) + }) }) diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index f3ccfba347e..3c87c10b9fa 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -280,7 +280,7 @@ describe('recordCumulativeUsage', () => { eventKey: 'update-cost:msg-1-billing', metadata: { inputTokens: 100, outputTokens: 5 }, }) - expect(result).toEqual({ billed: true, delta: 0.3474447, total: 0.3474447 }) + expect(result).toMatchObject({ billed: true, delta: 0.3474447, total: 0.3474447 }) expect(mockInsert).toHaveBeenCalledTimes(1) expect(mockUpdate).not.toHaveBeenCalled() expect(mockValues.mock.calls[0][0][0]).toMatchObject({ @@ -315,7 +315,7 @@ describe('recordCumulativeUsage', () => { cost: 0.4662453, eventKey: 'update-cost:msg-1-billing', }) - expect(result).toEqual({ billed: false, delta: 0, total: 0.4662453 }) + expect(result).toMatchObject({ billed: false, delta: 0, total: 0.4662453 }) expect(updateSet).not.toHaveBeenCalled() expect(mockInsert).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 5a845acad67..28af4201e69 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -1,9 +1,11 @@ import { createHash } from 'node:crypto' import { db, dbReplica } from '@sim/db' -import { usageLog, workflow } from '@sim/db/schema' +import { subscription as subscriptionTable, usageLog, workflow } from '@sim/db/schema' import { createLogger } from '@sim/logger' +import { toNumberOrNull } from '@sim/utils/coerce' import { getPostgresErrorCode, toError } from '@sim/utils/errors' import { generateId } from '@sim/utils/id' +import { toRecordOrNull } from '@sim/utils/object' import { and, desc, eq, gte, inArray, lt, lte, notInArray, or, sql } from 'drizzle-orm' import { type CursorKey, @@ -583,6 +585,14 @@ export interface RecordCumulativeUsageParams { /** Stable per-request key; the single ledger row is keyed on this. */ eventKey: string metadata?: UsageLogMetadata + /** + * The Stripe-period subscription that pays for this request. When given, a top-up that + * arrives after that subscription has moved past the period of the request's latest row is + * recorded in a new row stamped with the subscription's current period, so a request that + * outlives its billing period is invoiced by the period it was spent in rather than topping up + * a period that has already been closed. Omit it for reporting-window and free payers. + */ + payerSubscriptionId?: string } export interface RecordCumulativeUsageResult { @@ -592,6 +602,53 @@ export interface RecordCumulativeUsageResult { delta: number /** The request's recorded cumulative cost after this flush. */ total: number + /** The billing period of the row this flush wrote to, or of the request's latest row. */ + billingPeriod: { start: Date; end: Date } +} + +/** + * The most period rows one request may span: its first row plus one per later billing period. + * A request still billing twelve periods after it started is refused rather than scanned. + */ +const MAX_CUMULATIVE_PERIOD_ROWS = 12 + +/** The ledger key of the `index`-th period a cumulative request rolled into; 0 is the request key. */ +function cumulativePeriodEventKey(eventKey: string, index: number): string { + return index === 0 ? eventKey : `${eventKey}@${index}` +} + +/** Decimal places kept when period row costs are summed or subtracted as floats. */ +const PERIOD_COST_DECIMALS = 12 + +function sumLedgerCost(rows: readonly { cost: string }[]): number { + if (rows.length <= 1) return rows[0] ? Number.parseFloat(rows[0].cost) : 0 + const total = rows.reduce((sum, row) => sum + Number.parseFloat(row.cost), 0) + return Number(total.toFixed(PERIOD_COST_DECIMALS)) +} + +const CUMULATIVE_TOKEN_FIELDS = ['inputTokens', 'outputTokens'] as const + +/** + * A period row's share of a cumulative callback's token counts: the cumulative counts minus what + * the request's other rows already hold, so summing the rows never counts a token twice. + */ +function periodUsageMetadata( + metadata: UsageLogMetadata | undefined, + otherRows: readonly { metadata: unknown }[] +): UsageLogMetadata | undefined { + const cumulative = toRecordOrNull(metadata) + if (!cumulative || otherRows.length === 0) return metadata + const share: Record = { ...cumulative } + for (const field of CUMULATIVE_TOKEN_FIELDS) { + const total = toNumberOrNull(cumulative[field]) + if (total === null) continue + const recorded = otherRows.reduce( + (sum, row) => sum + (toNumberOrNull(toRecordOrNull(row.metadata)?.[field]) ?? 0), + 0 + ) + share[field] = Math.max(0, total - recorded) + } + return share } export type CumulativeUsageContextField = @@ -628,6 +685,8 @@ function assertCumulativeUsageLedgerBinding( workspaceId?: string billingContext: BillingContext eventKey: string + /** A request whose first charge landed after its period closed is stamped with a later one. */ + allowLaterPeriod?: boolean } ): void { const mismatchedFields: CumulativeUsageContextField[] = [] @@ -643,11 +702,15 @@ function assertCumulativeUsageLedgerBinding( ) { mismatchedFields.push('billing entity') } - if ( - existing.billingPeriodStart?.getTime() !== - expected.billingContext.billingPeriod.start.getTime() || - existing.billingPeriodEnd?.getTime() !== expected.billingContext.billingPeriod.end.getTime() - ) { + const frozenPeriod = expected.billingContext.billingPeriod + const samePeriod = + existing.billingPeriodStart?.getTime() === frozenPeriod.start.getTime() && + existing.billingPeriodEnd?.getTime() === frozenPeriod.end.getTime() + const laterPeriod = + expected.allowLaterPeriod === true && + existing.billingPeriodStart !== null && + existing.billingPeriodStart.getTime() >= frozenPeriod.end.getTime() + if (!samePeriod && !laterPeriod) { mismatchedFields.push('billing period') } @@ -703,6 +766,7 @@ export async function recordCumulativeUsage( cost, eventKey, metadata, + payerSubscriptionId, } = params if (workspaceId && (!billingEntity || !billingPeriod)) { @@ -744,10 +808,26 @@ export async function recordCumulativeUsage( await acquireAdvisoryXactLock(tx, 'usage_log_event', eventKey) enterStage('read') - const [existing] = await tx + // The payer's current period, share-locked so a rollover of the subscription waits for + // this top-up to commit: whatever a close later sums for the old period is final. + const [currentPeriod] = payerSubscriptionId + ? await tx + .select({ + start: subscriptionTable.periodStart, + end: subscriptionTable.periodEnd, + }) + .from(subscriptionTable) + .where(eq(subscriptionTable.id, payerSubscriptionId)) + .for('share') + .limit(1) + : [] + + const rows = await tx .select({ id: usageLog.id, + eventKey: usageLog.eventKey, cost: usageLog.cost, + metadata: usageLog.metadata, userId: usageLog.userId, workspaceId: usageLog.workspaceId, billingEntityType: usageLog.billingEntityType, @@ -756,55 +836,108 @@ export async function recordCumulativeUsage( billingPeriodEnd: usageLog.billingPeriodEnd, }) .from(usageLog) - .where(eq(usageLog.eventKey, eventKey)) - .limit(1) - - if (existing) { - assertCumulativeUsageLedgerBinding(existing, { + .where( + payerSubscriptionId + ? inArray( + usageLog.eventKey, + Array.from({ length: MAX_CUMULATIVE_PERIOD_ROWS }, (_, index) => + cumulativePeriodEventKey(eventKey, index) + ) + ) + : eq(usageLog.eventKey, eventKey) + ) + .limit(MAX_CUMULATIVE_PERIOD_ROWS) + + // Period rows are written in order under this lock, so they are the keys 0..n-1. + const chain = payerSubscriptionId + ? Array.from({ length: rows.length }, (_, index) => + rows.find((row) => row.eventKey === cumulativePeriodEventKey(eventKey, index)) + ).filter((row) => row !== undefined) + : rows.slice(0, 1) + if (payerSubscriptionId && chain.length !== rows.length) { + throw new Error(`Cumulative usage event "${eventKey}" has a gap in its period rows`) + } + const [anchor] = chain + if (anchor) { + assertCumulativeUsageLedgerBinding(anchor, { userId, workspaceId, billingContext, eventKey, + allowLaterPeriod: Boolean(payerSubscriptionId), }) } - const recorded = existing ? Number.parseFloat(existing.cost) : 0 + const latest = chain.at(-1) + const latestPeriod = + latest?.billingPeriodStart && latest.billingPeriodEnd + ? { start: latest.billingPeriodStart, end: latest.billingPeriodEnd } + : billingContext.billingPeriod + const recorded = sumLedgerCost(chain) const { shouldBill, delta, newTotal } = resolveCumulativeTopUp(recorded, cost) if (!shouldBill) { enterStage('commit') - return { billed: false, delta: 0, total: recorded } + return { billed: false, delta: 0, total: recorded, billingPeriod: latestPeriod } + } + + // Only ever forward: a subscription period that is not past the latest row keeps topping + // up that row, whatever the wall clock or a replayed webhook says. + const rolledPeriod = + currentPeriod?.start && + currentPeriod.end && + currentPeriod.start.getTime() >= latestPeriod.end.getTime() + ? { start: currentPeriod.start, end: currentPeriod.end } + : null + if (rolledPeriod && latest && chain.length >= MAX_CUMULATIVE_PERIOD_ROWS) { + throw new Error(`Cumulative usage event "${eventKey}" spans too many billing periods`) } enterStage('write') - if (existing) { + if (latest && !rolledPeriod) { + const otherRows = chain.slice(0, -1) + const latestCost = + otherRows.length === 0 + ? newTotal + : Number((newTotal - sumLedgerCost(otherRows)).toFixed(PERIOD_COST_DECIMALS)) await tx .update(usageLog) - .set({ cost: newTotal.toString(), metadata: metadata ?? null }) - .where(eq(usageLog.id, existing.id)) - } else { - await recordUsage({ - userId, - workspaceId, - tx, - billingEntity: billingContext.billingEntity, - billingPeriod: billingContext.billingPeriod, - entries: [ - { - category: 'model', - source, - description: model, - cost: newTotal, - eventKey, - sourceReference: eventKey, - ...(metadata ? { metadata } : {}), - }, - ], - }) + .set({ + cost: latestCost.toString(), + metadata: periodUsageMetadata(metadata, otherRows) ?? null, + }) + .where(eq(usageLog.id, latest.id)) + enterStage('commit') + return { billed: true, delta, total: newTotal, billingPeriod: latestPeriod } } + const targetPeriod = rolledPeriod ?? billingContext.billingPeriod + const rowMetadata = periodUsageMetadata(metadata, chain) + await recordUsage({ + userId, + workspaceId, + tx, + billingEntity: billingContext.billingEntity, + billingPeriod: targetPeriod, + entries: [ + { + category: 'model', + source, + description: model, + cost: chain.length === 0 ? newTotal : Number(delta.toFixed(PERIOD_COST_DECIMALS)), + eventKey: cumulativePeriodEventKey(eventKey, chain.length), + sourceReference: eventKey, + ...(rowMetadata ? { metadata: rowMetadata } : {}), + }, + ], + }) enterStage('commit') - return { billed: true, delta, total: newTotal } + return { + billed: true, + delta, + total: newTotal, + billingPeriod: { start: targetPeriod.start, end: targetPeriod.end }, + } }) succeeded = true return result diff --git a/apps/sim/lib/billing/usage-upgrade.ts b/apps/sim/lib/billing/usage-upgrade.ts new file mode 100644 index 00000000000..e72713a8d8c --- /dev/null +++ b/apps/sim/lib/billing/usage-upgrade.ts @@ -0,0 +1,67 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import type { UsageUpgradePayload } from '@/lib/api/contracts/subscription' +import type { + AttributedUsageLimitsResult, + BillingAttributionSnapshot, +} from '@/lib/billing/core/billing-attribution' +import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' +import { isEnterprise, isPaid } from '@/lib/billing/plan-helpers' +import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' + +const logger = createLogger('UsageUpgrade') + +const UPGRADE_PLAN_MESSAGE = + "You've reached your usage limit. Please upgrade your plan to continue." + +const MEMBER_CAP_MESSAGE = + "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." + +/** + * The upgrade card for a payer over its usage limit: a plan upgrade for a free payer, a limit + * increase for a paid one, with copy naming who can raise an organization's limit. A member + * over the cap their organization set gets copy naming who can raise that cap. An attributed + * run reads the plan from its admission snapshot without a query; otherwise the actor's current + * subscription decides, and a failed lookup falls back to the plan-upgrade card. + */ +export async function resolveUsageUpgradePayload( + userId: string, + billingAttribution?: BillingAttributionSnapshot, + scope?: AttributedUsageLimitsResult['scope'] +): Promise { + if (scope === 'member') { + return { reason: 'usage_limit', action: 'increase_limit', message: MEMBER_CAP_MESSAGE } + } + let plan: string | undefined + let orgScoped = false + try { + if (billingAttribution) { + plan = billingAttribution.payerSubscription?.plan + orgScoped = billingAttribution.billingEntity.type === 'organization' + } else { + const subscription = await getHighestPrioritySubscription(userId) + plan = subscription?.plan + orgScoped = isOrgScopedSubscription(subscription, userId) + } + } catch (error) { + logger.warn('Failed to determine subscription plan, defaulting to upgrade_plan', { + error: getErrorMessage(error), + }) + } + + if (!plan || !isPaid(plan)) { + return { reason: 'usage_limit', action: 'upgrade_plan', message: UPGRADE_PLAN_MESSAGE } + } + // Paid plans get `increase_limit`; the copy says who can raise it when the user cannot. + const message = !orgScoped + ? "You've reached your usage limit for this billing period. Please increase your usage limit from billing settings to continue." + : isEnterprise(plan) + ? "You've reached your organization's usage limit for this billing period. Only an organization admin or Sim support can raise an enterprise limit — reach out to them to continue." + : "You've reached your organization's usage limit for this billing period. Only an organization owner or admin can raise the limit — please ask them to update it from the team billing settings." + return { reason: 'usage_limit', action: 'increase_limit', message } +} + +/** The assistant text that renders {@link payload} as the usage card. */ +export function formatUsageUpgradeTag(payload: UsageUpgradePayload): string { + return `${JSON.stringify(payload)}` +} diff --git a/apps/sim/lib/mothership/generated/billing.ts b/apps/sim/lib/mothership/generated/billing.ts index 12d4ba14b82..34072228ebf 100644 --- a/apps/sim/lib/mothership/generated/billing.ts +++ b/apps/sim/lib/mothership/generated/billing.ts @@ -64,9 +64,27 @@ export const BillingCallbackHeaders = z context.addIssue({ code: "custom", message: "Incomplete or conflicting billing protocol headers" }); }); +/** Sim's plan-aware usage card: the JSON body of the `` tag its chat renders. */ +export const UsageUpgrade = z.object({ + reason: z.literal("usage_limit"), + action: z.enum(["upgrade_plan", "increase_limit"]), + message: z.string().min(1).max(1_000), +}); +export type UsageUpgrade = z.infer; + export const BillingCallbackResult = z.object({ success: z.boolean(), code: z.string().optional(), + /** The payer is over its plan usage limit after this charge. Absent (older Sim) means not over. */ + usageExceeded: z.boolean().optional(), + /** The card for an over-limit payer; a malformed card never turns a settled charge into a retry. */ + usageUpgrade: UsageUpgrade.optional().catch(undefined), +}); + +/** A continuation refused for the usage limit. A body-less 402 is a blocked account instead. */ +export const UsageLimitRefusal = z.object({ + code: z.literal("USAGE_LIMIT_EXCEEDED"), + usageUpgrade: UsageUpgrade, }); export const BillingDuplicateCode = "DUPLICATE_BILLING_EVENT"; diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index 51b07fb5ef7..cb5583f6a8f 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -79,6 +79,7 @@ vi.mock('@/lib/mothership/tools/server/files/file-preview', async () => { }) import { + BillingLimitError, buildPreviewContentUpdate, CopilotBackendError, decodeJsonStringPrefix, @@ -888,6 +889,37 @@ describe('copilot go stream helpers', () => { expect(fetch).toHaveBeenCalledTimes(1) }) + it('keeps a blocked account off the usage card when the worker refuses a leg', async () => { + vi.mocked(fetch).mockResolvedValueOnce( + new Response(JSON.stringify({ code: 'BILLING_BLOCKED', error: 'Billing account frozen.' }), { + status: 402, + }) + ) + const blocked = await runStreamLoop( + 'https://example.com/api/tools/resume', + {}, + createStreamingContext(), + turnScopedExecContext(), + { timeout: 1000 } + ).then( + () => undefined, + (error: unknown) => error + ) + expect(blocked).not.toBeInstanceOf(BillingLimitError) + expect(blocked).toMatchObject({ code: 'forbidden', message: 'Billing account frozen.' }) + + vi.mocked(fetch).mockResolvedValueOnce(new Response(null, { status: 402 })) + await expect( + runStreamLoop( + 'https://example.com/api/tools/resume', + {}, + createStreamingContext(), + turnScopedExecContext(), + { timeout: 1000 } + ) + ).rejects.toBeInstanceOf(BillingLimitError) + }) + it('fails closed when the shared stream ends before a terminal event', async () => { const textEvent = createEvent({ streamId: 'stream-1', diff --git a/apps/sim/lib/mothership/request/go/stream.ts b/apps/sim/lib/mothership/request/go/stream.ts index b55b623458d..a5b61e7521e 100644 --- a/apps/sim/lib/mothership/request/go/stream.ts +++ b/apps/sim/lib/mothership/request/go/stream.ts @@ -2,6 +2,8 @@ import { type Context, SpanStatusCode } from '@opentelemetry/api' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { toRecordOrNull } from '@sim/utils/object' +import { COPILOT_BILLING_BLOCKED_CODE } from '@/lib/api/contracts/copilot' +import { OrchestrationError } from '@/lib/core/orchestration/types' import { ORCHESTRATION_TIMEOUT_MS } from '@/lib/mothership/constants' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { CopilotSseCloseReason } from '@/lib/mothership/generated/trace-attribute-values-v1' @@ -106,6 +108,16 @@ function userFacingRejection(value: unknown): string | undefined { return message } +const BILLING_BLOCKED_MESSAGE = 'Billing account is blocked' + +function parseJsonRecord(body: string): Record | null { + try { + return toRecordOrNull(JSON.parse(body)) + } catch { + return null + } +} + /** * What the user is told about a failed backend response. A 5xx or a gateway page * is upstream detail and stays on the error for logs; a 4xx may carry the @@ -233,6 +245,14 @@ export async function runStreamLoop( const errorText = await response.text().catch(() => '') if (response.status === 402) { + // A blocked account is refused as blocked, never with the usage card. + const refusal = parseJsonRecord(errorText) + if (refusal?.code === COPILOT_BILLING_BLOCKED_CODE) { + throw new OrchestrationError( + 'forbidden', + userFacingRejection(refusal.error) ?? BILLING_BLOCKED_MESSAGE + ) + } throw new BillingLimitError(execContext.userId) } diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index 1e21943fc8b..f4e11eec090 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -1,6 +1,12 @@ import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { + billingUsageGateCacheMock, + billingUsageGateCacheMockFns, +} from '@sim/testing/mocks/billing-usage-gate-cache.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { createAttributedBillingRequestEnvelope } from '@/lib/billing/core/billing-attribution' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { BillingLimitError } from '@/lib/mothership/request/go/stream' import { authorizeLifecycleContinuation, restoreBillingAdmission } from './admission' const mocks = vi.hoisted(() => ({ authorize: vi.fn(), standing: vi.fn() })) @@ -8,6 +14,8 @@ vi.mock('@/lib/mothership/application/authorize-chat-callback', () => ({ authorizeCopilotChatCallback: mocks.authorize, checkCopilotContinuationBilling: mocks.standing, })) +vi.mock('@/lib/billing/core/usage-gate-cache', () => billingUsageGateCacheMock) +const { mockCheckExecutionUsageLimits } = billingUsageGateCacheMockFns const attribution = { actorUserId: 'actor', workspaceId: 'workspace', @@ -27,6 +35,7 @@ const context = { beforeEach(() => { setEnvFlags({ isHosted: true }) mocks.standing.mockResolvedValue({ blocked: false }) + mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: false }) }) afterEach(resetEnvFlagsMock) @@ -69,4 +78,47 @@ describe('continuation admission', () => { authorizeLifecycleContinuation({ ...context, billingAttribution: undefined }) ).rejects.toThrow('missing') }) + it('refuses a continuation whose original payer has crossed its usage limit', async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const refusal = authorizeLifecycleContinuation(context) + + await expect(refusal).rejects.toBeInstanceOf(BillingLimitError) + await expect(refusal).rejects.toMatchObject({ userId: 'actor' }) + expect(mockCheckExecutionUsageLimits).toHaveBeenCalledWith(attribution) + }) + it('keeps a blocked account a forbidden refusal without reading spend', async () => { + mocks.standing.mockResolvedValue({ blocked: true }) + + await expect(authorizeLifecycleContinuation(context)).rejects.toThrow('blocked') + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + it('does not read spend for a self-hosted continuation', async () => { + setEnvFlags({ isHosted: false }) + + await authorizeLifecycleContinuation(context) + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + it('continues a leg when spend cannot be read', async () => { + mockCheckExecutionUsageLimits.mockRejectedValueOnce(new Error('ledger read timed out')) + await expect(authorizeLifecycleContinuation(context)).resolves.toBeUndefined() + + mockCheckExecutionUsageLimits.mockResolvedValueOnce({ + isExceeded: true, + reason: 'usage_unavailable', + }) + await expect(authorizeLifecycleContinuation(context)).resolves.toBeUndefined() + }) + it('refuses a payer the gate finds blocked as a blocked account, not with the usage card', async () => { + mockCheckExecutionUsageLimits.mockResolvedValueOnce({ + isExceeded: true, + reason: 'billing_blocked', + scope: 'payer', + }) + + const refusal = authorizeLifecycleContinuation(context) + + await expect(refusal).rejects.toBeInstanceOf(OrchestrationError) + await expect(refusal).rejects.toThrow('blocked') + }) }) diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.ts b/apps/sim/lib/mothership/request/lifecycle/admission.ts index b685b7bf4d8..8ec1b7e7f40 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.ts @@ -6,12 +6,14 @@ import { COPILOT_BILLING_PROTOCOL_HEADER, requireBillingCallbackAttribution, } from '@/lib/billing/core/billing-attribution' +import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' import { isHosted } from '@/lib/core/config/env-flags' import { OrchestrationError } from '@/lib/core/orchestration/types' import { authorizeCopilotChatCallback, checkCopilotContinuationBilling, } from '@/lib/mothership/application/authorize-chat-callback' +import { BillingLimitError } from '@/lib/mothership/request/go/stream' import { BillingAdmissionSchema } from '@/lib/mothership/request/lifecycle/recovery-config' import type { ExecutionContext } from '@/lib/mothership/request/types' @@ -35,7 +37,13 @@ export function restoreBillingAdmission( return { attribution, envelope } } -/** Every resumed model leg rechecks authority and account standing without reading spend. */ +/** + * Every resumed model leg rechecks authority, account standing, and the original payer's spend. + * Spend is read through the execution usage gate, so a run under its limit pays no ledger read on + * most legs while an over-limit payer is re-read and refused. A spent limit is a + * {@link BillingLimitError}, which the lifecycle turns into the same upgrade card as a refused + * dispatch; a blocked account is refused as blocked; a usage read that fails lets the leg run. + */ export async function authorizeLifecycleContinuation( context: Pick< ExecutionContext, @@ -72,5 +80,9 @@ export async function authorizeLifecycleContinuation( }) if (standing.blocked) throw new OrchestrationError('forbidden', 'Continuation billing account is blocked') + const usage = await readMidRunUsageVerdict(context.billingAttribution) + if (usage.status === 'blocked') + throw new OrchestrationError('forbidden', 'Continuation billing account is blocked') + if (usage.status === 'exceeded') throw new BillingLimitError(context.userId) } } diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 963aeb78ee0..60d59896864 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -199,6 +199,11 @@ vi.mock('@/lib/mothership/request/tools/billing', () => ({ handleBillingLimitResponse: vi.fn(), })) +const mockRequestExplicitStreamAbort = vi.hoisted(() => vi.fn()) +vi.mock('@/lib/mothership/request/session/explicit-abort', () => ({ + requestExplicitStreamAbort: mockRequestExplicitStreamAbort, +})) + vi.mock('@/lib/mothership/request/tools/executor', () => ({ executeToolAndReport: vi.fn(), failPendingToolCall: mockForceFailHungToolCall, @@ -212,6 +217,7 @@ vi.mock('@/lib/mothership/request/enterprise-byok', () => ({ resolveEnterpriseByokKey: mockResolveEnterpriseByokKey, })) +import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { buildPersistedAssistantMessage } from '@/lib/mothership/chat/persisted-message' import { MothershipStreamV1CompletionStatus, @@ -2330,7 +2336,7 @@ describe('runCopilotLifecycle', () => { } }) - it('cold recovery preserves billing identity and does not read spend again', async () => { + it('cold recovery reads spend only against the original billing identity', async () => { const attribution = { actorUserId: 'user-1', workspaceId: 'ws-1', @@ -2341,7 +2347,7 @@ describe('runCopilotLifecycle', () => { payerSubscription: null, } setEnvFlags({ isHosted: true }) - mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true }) + resetUsageGateCache() const billingRequestId = generateId() const onBillingAdmission = vi.fn() await runCopilotLifecycle( @@ -2363,7 +2369,8 @@ describe('runCopilotLifecycle', () => { onBillingAdmission, } ) - expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledOnce() + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(attribution) expect(onBillingAdmission).not.toHaveBeenCalled() expect(continuationAuth).toHaveBeenCalled() expect(mockRunStreamLoop).toHaveBeenCalledOnce() @@ -2416,6 +2423,65 @@ describe('runCopilotLifecycle', () => { expect(result.cancelled).not.toBe(true) }) + it('shows the usage card instead of resuming a run whose payer crossed its limit', async () => { + const billingAttribution = { + actorUserId: 'user-1', + workspaceId: 'ws-1', + organizationId: 'org-1', + billedAccountUserId: 'user-1', + billingEntity: { type: 'organization' as const, id: 'org-1' }, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + payerSubscription: null, + } + setEnvFlags({ isHosted: true }) + resetUsageGateCache() + mockCheckAttributedUsageLimits + .mockResolvedValueOnce({ isExceeded: false }) + .mockResolvedValue({ isExceeded: true, scope: 'payer' }) + mockRunStreamLoop.mockImplementationOnce( + async (_url: string, _init: RequestInit, context: StreamingContext) => { + context.toolCalls.set('tool-1', { + id: 'tool-1', + name: 'read', + status: MothershipStreamV1ToolOutcome.success, + result: { success: true, output: { content: 'file contents' } }, + }) + context.awaitingAsyncContinuation = { + checkpointId: 'ckpt-1', + pendingToolCallIds: ['tool-1'], + } + } + ) + + mockRequestExplicitStreamAbort.mockResolvedValue({ settled: true }) + const result = await runCopilotLifecycle( + { message: 'hello', messageId: 'message-1' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'execution-1', + runId: 'run-1', + billingAttribution, + } + ) + + expect(mockRunStreamLoop).toHaveBeenCalledOnce() + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(2) + expect(handleBillingLimitResponse).toHaveBeenCalledOnce() + expect(handleBillingLimitResponse).toHaveBeenCalledWith( + 'user-1', + expect.anything(), + expect.anything(), + expect.anything() + ) + expect(result.cancelled).not.toBe(true) + expect(result.error).toBeUndefined() + expect(mockRequestExplicitStreamAbort).toHaveBeenCalledWith( + expect.objectContaining({ streamId: 'message-1', userId: 'user-1', chatId: 'chat-1' }) + ) + }) + it('preserves a resume tool name that collides with a configured secret', async () => { const registry = new ResolvedSecretTraceRegistry([ { name: 'TOKEN', plaintext: 'unsafe-tool', encryptedValue: 'ciphertext' }, diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 9a9ac2cd667..c95be22607f 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -518,14 +518,23 @@ export async function runCopilotLifecycle( requestPayload, lifecycleOptions.workspaceId ) - await runCheckpointLoop( - modelSafeRequestPayload, - context, - execContext, - lifecycleOptions, - goRoute, - hostedBillingRequest - ) + try { + await runCheckpointLoop( + modelSafeRequestPayload, + context, + execContext, + lifecycleOptions, + goRoute, + hostedBillingRequest + ) + } catch (error) { + // A continuation refused on spend, or a worker 402 on a child leg of a subagent + // fan-out, ends the turn with the same card as a refused dispatch. + if (!(error instanceof BillingLimitError)) throw error + context.awaitingAsyncContinuation = undefined + await handleBillingLimitResponse(error.userId, context, execContext, lifecycleOptions) + await stopWorkerRunAfterUsageRefusal(context.messageId, execContext) + } } // The backend's terminal `complete` is the turn's verdict. A failure it @@ -1691,6 +1700,32 @@ function isAborted(options: CopilotLifecycleOptions, context: StreamingContext): return !!(options.abortSignal?.aborted || context.wasAborted) } +/** + * A refused continuation leaves the worker run parked on its checkpoint, and a parked run holds + * the chat: the next message would be refused as busy until the sweeper expires it. Stopping it + * frees the chat, so the message sent after an upgrade continues the conversation. + */ +async function stopWorkerRunAfterUsageRefusal( + streamId: string, + execContext: Pick +): Promise { + try { + const { requestExplicitStreamAbort } = await import( + '@/lib/mothership/request/session/explicit-abort' + ) + await requestExplicitStreamAbort({ + streamId, + userId: execContext.userId, + chatId: execContext.chatId, + }) + } catch (error) { + logger.warn('Worker stop after a usage-limit refusal was not delivered', { + streamId, + error: getErrorMessage(error), + }) + } +} + function cancelPendingTools(context: StreamingContext): void { for (const [, toolCall] of context.toolCalls) { if ( diff --git a/apps/sim/lib/mothership/request/tools/billing.test.ts b/apps/sim/lib/mothership/request/tools/billing.test.ts index d1c74216697..9755b299073 100644 --- a/apps/sim/lib/mothership/request/tools/billing.test.ts +++ b/apps/sim/lib/mothership/request/tools/billing.test.ts @@ -102,4 +102,15 @@ describe('handleBillingLimitResponse', () => { payload: { text: expect.stringContaining('"action":"increase_limit"') }, }) }) + + it('terminates the turn with the card after a leg that already ended at a checkpoint', async () => { + const onEvent = vi.fn() + const pausedContext = { streamComplete: true } as StreamingContext + + await handleBillingLimitResponse('actor-1', pausedContext, createExecutionContext(), { + onEvent, + } as OrchestratorOptions) + + expect(onEvent.mock.calls.map(([event]) => event.type)).toEqual(['text', 'complete']) + }) }) diff --git a/apps/sim/lib/mothership/request/tools/billing.ts b/apps/sim/lib/mothership/request/tools/billing.ts index fed35fd4156..09a27517d63 100644 --- a/apps/sim/lib/mothership/request/tools/billing.ts +++ b/apps/sim/lib/mothership/request/tools/billing.ts @@ -1,8 +1,5 @@ import { createLogger } from '@sim/logger' -import { getErrorMessage } from '@sim/utils/errors' -import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' -import { isEnterprise, isPaid } from '@/lib/billing/plan-helpers' -import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' +import { formatUsageUpgradeTag, resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' import { MothershipStreamV1CompletionStatus, MothershipStreamV1EventType, @@ -19,11 +16,12 @@ import type { const logger = createLogger('CopilotBillingEffect') /** - * Handle a 402 billing-limit response from the Go backend. + * Ends the turn with the usage card: a refused dispatch or continuation, a worker 402, or a + * worker usage-limit terminal that arrived without a card of its own. * - * Determines whether the user needs a plan upgrade or a limit increase, - * then dispatches synthetic text + complete events through the handler chain - * so the client renders the upgrade prompt. + * Dispatches synthetic text + complete events through the handler chain so the client renders + * the upgrade prompt and the turn finishes as complete, so the next message after an upgrade + * starts normally. */ export async function handleBillingLimitResponse( userId: string, @@ -31,47 +29,10 @@ export async function handleBillingLimitResponse( execContext: ExecutionContext, options: OrchestratorOptions ): Promise { - let action: 'upgrade_plan' | 'increase_limit' = 'upgrade_plan' - let message = "You've reached your usage limit. Please upgrade your plan to continue." - try { - let plan: string | undefined - let orgScoped = false - if (execContext.billingAttribution) { - plan = execContext.billingAttribution.payerSubscription?.plan - orgScoped = execContext.billingAttribution.billingEntity.type === 'organization' - } else { - const sub = await getHighestPrioritySubscription(userId) - plan = sub?.plan - orgScoped = isOrgScopedSubscription(sub, userId) - } - - if (plan && isPaid(plan)) { - // Paid subs use the existing `increase_limit` action so the UI - // (`UsageUpgradeDisplay`) renders its standard button. The message - // text does the work of clarifying the action when the user can't - // actually self-serve the limit change. - action = 'increase_limit' - if (orgScoped) { - message = isEnterprise(plan) - ? "You've reached your organization's usage limit for this billing period. Only an organization admin or Sim support can raise an enterprise limit — reach out to them to continue." - : "You've reached your organization's usage limit for this billing period. Only an organization owner or admin can raise the limit — please ask them to update it from the team billing settings." - } else { - message = - "You've reached your usage limit for this billing period. Please increase your usage limit from billing settings to continue." - } - } - } catch (error) { - logger.warn('Failed to determine subscription plan, defaulting to upgrade_plan', { - error: getErrorMessage(error), - }) - } - - const upgradePayload = JSON.stringify({ - reason: 'usage_limit', - action, - message, - }) - const syntheticContent = `${upgradePayload}` + const payload = await resolveUsageUpgradePayload(userId, execContext.billingAttribution) + const syntheticContent = formatUsageUpgradeTag(payload) + // The card is this turn's terminal even when the refused leg follows one that already ended. + context.streamComplete = false const syntheticEvents: StreamEvent[] = [ { From 16f34f1728125af070a44878b9800188d28676a8 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 03:03:36 -0700 Subject: [PATCH 02/13] fix(billing): roll forward-moved periods, lock only past period end, skip ended admissions - Roll a cumulative charge into the payer's current period whenever that period starts after the latest row's, so an anchor reset or resync inside the old period never tops up a closed period. - Share-lock the payer's subscription row only once the latest row's period has ended; before that no close can be due, and the lock would starve the rollover update for a busy payer. - Mid-run usage checks report unknown (continue) once the run's admitted period has ended, instead of judging the old period against its allowance. - Refuse request keys containing "@" when period rows are in play, so they cannot collide with another request's period rows. - Document the mixed-version and rollback window: code that predates period rows can double-count a run's post-rollover spend only between a rollover and that period's close (at least an hour); the exposure is cents to dollars. --- .../app/api/billing/update-cost/route.test.ts | 22 +++++- .../copilot/api-keys/validate/route.test.ts | 2 +- apps/sim/lib/billing/core/mid-run-usage.ts | 12 ++- .../lib/billing/core/usage-log.integration.ts | 79 ++++++++++++++++--- apps/sim/lib/billing/core/usage-log.ts | 57 ++++++++----- .../request/lifecycle/admission.test.ts | 16 +++- .../mothership/request/lifecycle/run.test.ts | 4 +- 7 files changed, 153 insertions(+), 39 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index 4cbc6f66fb8..c29b31d48ff 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -852,6 +852,10 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => { describe('POST /api/billing/update-cost — mid-run usage gate', () => { let callbackSequence = 0 + const CURRENT_ATTRIBUTION = { + ...ATTRIBUTION, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, + } function attributedCallback() { callbackSequence += 1 @@ -881,7 +885,7 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { mockCheckInternalApiKey.mockReturnValue({ success: true }) mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 }) mockCheckAndBillPayerOverageThreshold.mockResolvedValue(undefined) - mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockRequireBillingAttributionHeader.mockReturnValue(CURRENT_ATTRIBUTION) mockToBillingContext.mockReturnValue({ billingEntity: { type: 'organization', id: 'org-1' }, billingPeriod: { @@ -906,12 +910,12 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { message: expect.stringContaining('usage limit'), }, }) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(ATTRIBUTION) + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(CURRENT_ATTRIBUTION) }) it('offers a paid organization payer the increase-limit card', async () => { mockRequireBillingAttributionHeader.mockReturnValue({ - ...ATTRIBUTION, + ...CURRENT_ATTRIBUTION, payerSubscription: { id: 'sub-1', plan: 'team', status: 'active', seats: 4 }, }) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) @@ -971,7 +975,7 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { beforeEach(() => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockRequireBillingAttributionHeader.mockReturnValue({ - ...ATTRIBUTION, + ...CURRENT_ATTRIBUTION, payerSubscription: PAYER_SUBSCRIPTION, }) }) @@ -1061,6 +1065,16 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { }) }) + it('does not judge a run against an admitted period that has already ended', async () => { + mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(false) + expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + }) + it('keeps a recorded charge successful when the gate read fails', async () => { mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out')) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index a10d2e1de21..0a9c5831f95 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -63,7 +63,7 @@ const ATTRIBUTION = { billingEntity: { type: 'organization' as const, id: 'org-1' }, billingPeriod: { start: '2026-07-01T00:00:00.000Z', - end: '2026-08-01T00:00:00.000Z', + end: '2099-01-01T00:00:00.000Z', source: 'reporting' as const, }, payerSubscription: null, diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index e6d39178f44..b7102111b62 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -14,9 +14,9 @@ const logger = createLogger('MidRunUsage') * upgrade card. * - `blocked`: the account is blocked (payment failed, dispute); the run is refused as a blocked * account, never with the upgrade card. - * - `unknown`: the gate could not read usage. Admission fails closed on this, but a run already - * under way continues: a database blip must not end a paying user's long run, and the next - * step or re-check reads again. + * - `unknown`: the gate could not read usage, or the run's admitted period has ended. Admission + * fails closed on an unreadable ledger, but a run already under way continues: a database + * blip must not end a paying user's long run, and the next step or re-check reads again. */ export type MidRunUsageVerdict = | { status: 'within' } @@ -27,6 +27,12 @@ export type MidRunUsageVerdict = export async function readMidRunUsageVerdict( attribution: BillingAttributionSnapshot ): Promise { + // The gate judges the admitted snapshot's period. Once that period has ended it would keep + // counting the old period against the old allowance, so a run just past a reset is not judged + // until its next admission reads the new one. + if (Date.now() >= new Date(attribution.billingPeriod.end).getTime()) { + return { status: 'unknown' } + } let usage: AttributedUsageLimitsResult try { usage = await checkExecutionUsageLimits(attribution) diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index c34361ff622..4132cca2238 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -314,38 +314,52 @@ describe('Cumulative billing with PostgreSQL', () => { ) describe('a request that outlives its billing period', () => { + // Past periods: the old period's row is written under the subscription lock only once + // that period has ended. const periods = [ - new Date('2026-09-01T00:00:00.000Z'), - new Date('2026-10-01T00:00:00.000Z'), - new Date('2026-11-01T00:00:00.000Z'), - new Date('2026-12-01T00:00:00.000Z'), + new Date('2025-09-01T00:00:00.000Z'), + new Date('2025-10-01T00:00:00.000Z'), + new Date('2025-11-01T00:00:00.000Z'), + new Date('2025-12-01T00:00:00.000Z'), ] const payer = { type: 'organization', id: 'payer' } as const - async function setSubscriptionPeriod(index: number) { + async function setSubscriptionWindow(start: Date, end: Date) { await connection` insert into subscription (id, period_start, period_end) - values ('sub-1', ${periods[index].toISOString()}::timestamptz at time zone 'UTC', ${periods[index + 1].toISOString()}::timestamptz at time zone 'UTC') + values ('sub-1', ${start.toISOString()}::timestamptz at time zone 'UTC', ${end.toISOString()}::timestamptz at time zone 'UTC') on conflict (id) do update set period_start = excluded.period_start, period_end = excluded.period_end ` } - function charge(cost: number) { - return recordCumulativeUsage({ ...usage(cost), payerSubscriptionId: 'sub-1' }) + async function setSubscriptionPeriod(index: number) { + await setSubscriptionWindow(periods[index], periods[index + 1]) + } + + function charge(cost: number, frozen = { start: periods[0], end: periods[1] }) { + return recordCumulativeUsage({ + ...usage(cost), + billingPeriod: frozen, + payerSubscriptionId: 'sub-1', + }) } /** What the cycle close invoices for one period: the ledger rows stamped with it. */ - async function stampedTotal(index: number) { + async function stampedWindowTotal(from: Date, to: Date) { const byUser = await getStampedPeriodRangeUsageCostByUser( payer, - { from: periods[index], to: periods[index + 1] }, + { from, to }, undefined, database ) return [...byUser.values()].reduce((total, cost) => total + cost, 0) } + function stampedTotal(index: number) { + return stampedWindowTotal(periods[index], periods[index + 1]) + } + it('invoices a charge that spans a period close exactly once in total', async () => { await setSubscriptionPeriod(0) expect(await charge(0.4)).toMatchObject({ billed: true, total: 0.4 }) @@ -410,5 +424,50 @@ describe('Cumulative billing with PostgreSQL', () => { } expect(await stampedTotal(0)).toBeCloseTo(0.6, 9) }) + + it('rolls into a period whose start moved forward before the old period ended', async () => { + await setSubscriptionPeriod(0) + await charge(0.4) + const resetStart = new Date('2025-09-15T00:00:00.000Z') + const resetEnd = new Date('2025-10-15T00:00:00.000Z') + await setSubscriptionWindow(resetStart, resetEnd) + + expect(await charge(1)).toMatchObject({ + billed: true, + billingPeriod: { start: resetStart, end: resetEnd }, + }) + expect(await stampedTotal(0)).toBeCloseTo(0.4, 9) + expect(await stampedWindowTotal(resetStart, resetEnd)).toBeCloseTo(0.6, 9) + }) + + it('does not hold the rollover while the latest period is still running', async () => { + const start = new Date(Date.now() - 24 * 60 * 60 * 1000) + const end = new Date(Date.now() + 30 * 24 * 60 * 60 * 1000) + await setSubscriptionWindow(start, end) + await charge(0.4, { start, end }) + const pause = pauseNextTransaction() + const inFlight = charge(0.6, { start, end }) + try { + await pause.reached.promise + await connection.begin(async (tx) => { + await tx`select set_config('lock_timeout', '300ms', true)` + await tx`update subscription set period_end = period_end where id = 'sub-1'` + }) + } finally { + pause.release.resolve() + await inFlight + } + }) + + it('refuses a request key that could collide with its period rows', async () => { + await setSubscriptionPeriod(0) + await expect( + recordCumulativeUsage({ + ...usage(0.4, 'update-cost:request@1'), + payerSubscriptionId: 'sub-1', + }) + ).rejects.toThrow('@') + expect(await ledgerRows()).toEqual([]) + }) }) }) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 28af4201e69..c05e916f46b 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -591,6 +591,13 @@ export interface RecordCumulativeUsageParams { * recorded in a new row stamped with the subscription's current period, so a request that * outlives its billing period is invoiced by the period it was spent in rather than topping up * a period that has already been closed. Omit it for reporting-window and free payers. + * + * Mixed versions: code that predates period rows reads only the request key. If such code + * (during a deploy, or after a rollback) handles a later callback for a request that already + * has period rows, it re-adds those rows' amount to the first row. That only double-counts + * when it lands between the rollover and that period's close, which waits at least an hour, + * and only for runs spanning a rollover; the exposure is one run's post-rollover spend, cents + * to dollars. */ payerSubscriptionId?: string } @@ -612,7 +619,10 @@ export interface RecordCumulativeUsageResult { */ const MAX_CUMULATIVE_PERIOD_ROWS = 12 -/** The ledger key of the `index`-th period a cumulative request rolled into; 0 is the request key. */ +/** + * The ledger key of the `index`-th period a cumulative request rolled into; 0 is the request key. + * A request key may not itself contain `@`, so these keys never collide with another request's. + */ function cumulativePeriodEventKey(eventKey: string, index: number): string { return index === 0 ? eventKey : `${eventKey}@${index}` } @@ -774,6 +784,9 @@ export async function recordCumulativeUsage( } const billingContext = await resolveBillingContext(userId, billingEntity, billingPeriod) + if (payerSubscriptionId && eventKey.includes('@')) { + throw new Error(`Cumulative usage event "${eventKey}" must not contain "@"`) + } const startedAt = Date.now() let stage: CumulativeUsageStage = 'pool' @@ -808,20 +821,6 @@ export async function recordCumulativeUsage( await acquireAdvisoryXactLock(tx, 'usage_log_event', eventKey) enterStage('read') - // The payer's current period, share-locked so a rollover of the subscription waits for - // this top-up to commit: whatever a close later sums for the old period is final. - const [currentPeriod] = payerSubscriptionId - ? await tx - .select({ - start: subscriptionTable.periodStart, - end: subscriptionTable.periodEnd, - }) - .from(subscriptionTable) - .where(eq(subscriptionTable.id, payerSubscriptionId)) - .for('share') - .limit(1) - : [] - const rows = await tx .select({ id: usageLog.id, @@ -881,12 +880,34 @@ export async function recordCumulativeUsage( return { billed: false, delta: 0, total: recorded, billingPeriod: latestPeriod } } - // Only ever forward: a subscription period that is not past the latest row keeps topping - // up that row, whatever the wall clock or a replayed webhook says. + // The payer's current period. Once the latest row's period has ended, the read is + // share-locked so a rollover of the subscription waits for this top-up to commit, and + // whatever the close later sums for the old period is final. Before that no close can be + // due, and locking every callback would starve the rollover UPDATE for a busy payer. + const periodQuery = payerSubscriptionId + ? tx + .select({ + start: subscriptionTable.periodStart, + end: subscriptionTable.periodEnd, + }) + .from(subscriptionTable) + .where(eq(subscriptionTable.id, payerSubscriptionId)) + .limit(1) + : null + const [currentPeriod] = periodQuery + ? Date.now() >= latestPeriod.end.getTime() + ? await periodQuery.for('share') + : await periodQuery + : [] + + // Only ever forward: a subscription period that does not start after the latest row's + // keeps topping up that row, whatever the wall clock or a replayed webhook says. A start + // that moved forward inside the old period (anchor reset, resync) still rolls, so the old + // period's close is never topped up after the fact. const rolledPeriod = currentPeriod?.start && currentPeriod.end && - currentPeriod.start.getTime() >= latestPeriod.end.getTime() + currentPeriod.start.getTime() > latestPeriod.start.getTime() ? { start: currentPeriod.start, end: currentPeriod.end } : null if (rolledPeriod && latest && chain.length >= MAX_CUMULATIVE_PERIOD_ROWS) { diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index f4e11eec090..7497348b142 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -22,7 +22,7 @@ const attribution = { organizationId: 'original-org', billedAccountUserId: 'original-owner', billingEntity: { type: 'organization' as const, id: 'original-org' }, - billingPeriod: { start: '2026-09-01T00:00:00.000Z', end: '2026-10-01T00:00:00.000Z' }, + billingPeriod: { start: '2026-09-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, payerSubscription: null, } const context = { @@ -121,4 +121,18 @@ describe('continuation admission', () => { await expect(refusal).rejects.toBeInstanceOf(OrchestrationError) await expect(refusal).rejects.toThrow('blocked') }) + it('lets a leg run once its admitted period has ended instead of judging the old period', async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + await expect( + authorizeLifecycleContinuation({ + ...context, + billingAttribution: { + ...attribution, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + }, + }) + ).resolves.toBeUndefined() + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) }) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 60d59896864..de13f3c7f3d 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -2343,7 +2343,7 @@ describe('runCopilotLifecycle', () => { organizationId: 'org-1', billedAccountUserId: 'original-owner', billingEntity: { type: 'organization' as const, id: 'org-1' }, - billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, payerSubscription: null, } setEnvFlags({ isHosted: true }) @@ -2430,7 +2430,7 @@ describe('runCopilotLifecycle', () => { organizationId: 'org-1', billedAccountUserId: 'user-1', billingEntity: { type: 'organization' as const, id: 'org-1' }, - billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, payerSubscription: null, } setEnvFlags({ isHosted: true }) From a2c45c0003e64f9b38ccce4232980c04e5c7154b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 04:00:54 -0700 Subject: [PATCH 03/13] fix(billing): judge long runs against the current period and close review gaps - Mid-run usage checks judge a run that outlived its admitted period against the same payer's current period, re-read after a gate read that straddles the period end, and continue only when the current period cannot be read. Direct-v1 continuations read account usage through the same mid-run rules. - Share-lock the payer's subscription on every period-aware write, so an early period-start move waits for an in-flight top-up. - Reserve "@" in every cumulative idempotency key; update-cost rejects it. - Carry a member cap through BillingLimitError to the member card, and send every usage-limit refusal, including a worker 402 on the first leg, through one handler that also stops the worker run. - Every validate 402 now carries a declared body: USAGE_LIMIT_EXCEEDED, BILLING_BLOCKED, or USAGE_UNAVAILABLE for a new turn refused because usage could not be read. - The update-cost verdict schema ties usageUpgrade to usageExceeded. --- .../app/api/billing/update-cost/route.test.ts | 84 ++++++++++++++++- apps/sim/app/api/billing/update-cost/route.ts | 4 + .../copilot/api-keys/validate/route.test.ts | 47 ++++++++-- .../api/copilot/api-keys/validate/route.ts | 65 +++++++++++--- apps/sim/lib/api/contracts/copilot.ts | 17 +++- apps/sim/lib/api/contracts/subscription.ts | 43 +++++---- .../lib/billing/calculations/usage-monitor.ts | 12 ++- .../lib/billing/core/billing-attribution.ts | 30 +++++++ apps/sim/lib/billing/core/mid-run-usage.ts | 89 ++++++++++++++++--- .../lib/billing/core/usage-log.integration.ts | 37 ++++---- apps/sim/lib/billing/core/usage-log.ts | 21 ++--- .../lib/mothership/request/go/stream.test.ts | 16 ++++ apps/sim/lib/mothership/request/go/stream.ts | 19 +++- .../request/lifecycle/admission.test.ts | 44 ++++++--- .../mothership/request/lifecycle/admission.ts | 2 +- .../mothership/request/lifecycle/run.test.ts | 24 ++++- .../lib/mothership/request/lifecycle/run.ts | 16 ++-- .../mothership/request/tools/billing.test.ts | 16 ++++ .../lib/mothership/request/tools/billing.ts | 5 +- .../src/mocks/billing-attribution.mock.ts | 2 + 20 files changed, 486 insertions(+), 107 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index c29b31d48ff..778bc8256bc 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -10,6 +10,7 @@ import { } from '@sim/testing/mocks/billing-usage-log.mock' import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing/mocks/copilot-http.mock' import { mothershipOtelMock } from '@sim/testing/mocks/mothership-otel.mock' +import { sleep } from '@sim/utils/helpers' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { @@ -47,6 +48,7 @@ vi.mock('@/lib/billing/threshold-billing', () => ({ ThresholdSettlementError: MockThresholdSettlementError, })) +import { billingUpdateCostResponseSchema } from '@/lib/api/contracts/subscription' import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { BillingCallbackBody, @@ -71,6 +73,7 @@ const mockResolveLegacyV0BillingAttribution = billingAttributionMockFns.mockResolveLegacyV0BillingAttribution const mockToBillingContext = billingAttributionMockFns.mockToBillingContext const mockCheckAttributedUsageLimits = billingAttributionMockFns.mockCheckAttributedUsageLimits +const mockRefreshAttributionPeriod = billingAttributionMockFns.mockRefreshAttributionPeriod afterAll(resetEnvFlagsMock) @@ -238,6 +241,18 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => { expect(mockRecordCumulativeUsage).not.toHaveBeenCalled() }) + it('rejects an idempotency key that could collide with a period row key', async () => { + const res = await POST( + createMockRequest( + 'POST', + { ...SELF_HOSTED_UPDATE_COST_BODY, idempotencyKey: 'old-go-key@1' }, + { 'x-api-key': 'internal' } + ) + ) + + expect(res.status).toBe(400) + }) + it('rejects billing-enabled callbacks without a stable idempotency key', async () => { const res = await POST( createMockRequest('POST', KEYLESS_UPDATE_COST_BODY, { 'x-api-key': 'internal' }) @@ -910,7 +925,6 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { message: expect.stringContaining('usage limit'), }, }) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(CURRENT_ATTRIBUTION) }) it('offers a paid organization payer the increase-limit card', async () => { @@ -1065,14 +1079,77 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { }) }) - it('does not judge a run against an admitted period that has already ended', async () => { + /** The gate refuses only when it judges the payer's current period. */ + function refuseOnlyCurrentPeriod() { + mockCheckAttributedUsageLimits.mockImplementation( + async (attribution: typeof CURRENT_ATTRIBUTION) => ({ + isExceeded: attribution.billingPeriod.end === CURRENT_ATTRIBUTION.billingPeriod.end, + scope: 'payer', + }) + ) + } + + it('judges a run past its admitted period against the payer current period', async () => { mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockRefreshAttributionPeriod.mockResolvedValue(CURRENT_ATTRIBUTION) + refuseOnlyCurrentPeriod() + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(true) + }) + + it('keeps a run going when its current period cannot be read', async () => { + mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockRefreshAttributionPeriod.mockRejectedValue(new Error('subscription read timed out')) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) const body = await (await POST(attributedCallback())).json() expect(body.usageExceeded).toBe(false) - expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + }) + + it('rejudges a read that straddles the end of the admitted period', async () => { + const straddling = { + ...CURRENT_ATTRIBUTION, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: new Date(Date.now() + 40).toISOString(), + }, + } + mockRequireBillingAttributionHeader.mockReturnValue(straddling) + mockRefreshAttributionPeriod.mockResolvedValue(CURRENT_ATTRIBUTION) + mockCheckAttributedUsageLimits.mockImplementation( + async (attribution: typeof CURRENT_ATTRIBUTION) => { + if (attribution.billingPeriod.end !== straddling.billingPeriod.end) { + return { isExceeded: false } + } + await sleep(80) + return { isExceeded: true, scope: 'payer' } + } + ) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(false) + }) + + it('never answers a verdict whose card and flag disagree', () => { + const base = { + success: true, + data: { processedAt: new Date(0).toISOString(), requestId: 'request-1' }, + } + const card = { reason: 'usage_limit', action: 'upgrade_plan', message: 'Limit reached.' } + expect( + billingUpdateCostResponseSchema.safeParse({ ...base, usageExceeded: true }).success + ).toBe(false) + expect( + billingUpdateCostResponseSchema.safeParse({ + ...base, + usageExceeded: false, + usageUpgrade: card, + }).success + ).toBe(false) }) it('keeps a recorded charge successful when the gate read fails', async () => { @@ -1091,6 +1168,5 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { const res = await POST(attributedCallback()) await expect(res.json()).resolves.toMatchObject({ usageExceeded: false }) - expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index 4c0fe4789ae..526e8fab89e 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -233,6 +233,10 @@ async function updateCostInner(req: NextRequest, span: Span): Promise@`). + if (idempotencyKey?.includes('@')) { + return invalidBillingProtocolResponse(requestId, span) + } const isMcp = source === 'mcp_copilot' span.setAttributes({ diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 0a9c5831f95..0fe4c86abfc 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -75,7 +75,7 @@ const ACCOUNT_BILLING_DECISION = { billingEntity: { type: 'organization' as const, id: 'account-org' }, billingPeriod: { start: '2026-07-01T00:00:00.000Z', - end: '2026-08-01T00:00:00.000Z', + end: '2099-01-01T00:00:00.000Z', source: 'reporting' as const, }, } @@ -147,7 +147,7 @@ const { mockAuthorizeOrganizationChatDelegation: mockAuthorizeOrganizationChat } mothershipOrganizationChatsMockFns const { mockDeriveBillingContext } = billingUsageLogMockFns const { mockGetHighestPrioritySubscription } = billingPlanMockFns -const { mockCheckServerSideUsageLimits } = billingUsageMonitorMockFns +const { mockCheckServerSideUsageLimits, mockCheckUsageStatus } = billingUsageMonitorMockFns const mockIsEnterprisePlan = billingSubscriptionMockFns.mockIsEnterprisePlan const mockGetUserEntityPermissions = permissionsMockFns.mockGetUserEntityPermissions @@ -508,6 +508,7 @@ describe('validation lifecycle purposes', () => { mockAuthorizeCallback.mockReset().mockResolvedValue(undefined) mockCheckContinuationBilling.mockReset().mockResolvedValue({ blocked: false }) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 1, limit: 10 }) mockIsEnterprisePlan.mockResolvedValue(false) resetUsageGateCache() }) @@ -556,6 +557,23 @@ describe('validation lifecycle purposes', () => { expect(response.headers.get('x-sim-billing-account-decision')).toBeNull() }) + it('refuses a direct-v1 continuation whose account is over its usage limit', async () => { + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) + + const response = await POST(request(body, directHeaders)) + + expect(response.status).toBe(402) + await expect(response.json()).resolves.toMatchObject({ + code: 'USAGE_LIMIT_EXCEEDED', + usageUpgrade: { reason: 'usage_limit' }, + }) + }) + + it('admits a direct-v1 continuation whose usage cannot be read', async () => { + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, unavailable: true }) + expect((await POST(request(body, directHeaders))).status).toBe(200) + }) + it.each([ ['missing attribution', { ...attributedHeaders, 'x-sim-billing-attribution': '' }], [ @@ -711,6 +729,21 @@ describe('validation lifecycle purposes', () => { expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(3) }) + it('refuses a blocked new turn with the blocked body the contract declares', async () => { + mockCheckAttributedUsageLimits.mockResolvedValueOnce({ + isExceeded: true, + reason: 'billing_blocked', + message: 'Billing account frozen.', + scope: 'payer', + }) + const refused = await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders)) + expect(refused.status).toBe(402) + await expect(refused.json()).resolves.toEqual({ + code: 'BILLING_BLOCKED', + error: 'Billing account frozen.', + }) + }) + it('allows cancellation without billing material or spending/standing/plan checks', async () => { const response = await POST( request({ ...body, purpose: 'cancellation' }, { 'x-sim-billing-protocol': 'attribution-v1' }) @@ -817,10 +850,12 @@ describe('validation lifecycle purposes', () => { isExceeded: true, payerUsage: { currentUsage: 120, limit: 100 }, }) - expect((await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders))).status).toBe( - 402 - ) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(2) + const refused = await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders)) + expect(refused.status).toBe(402) + await expect(refused.json()).resolves.toMatchObject({ + code: 'USAGE_LIMIT_EXCEEDED', + usageUpgrade: { reason: 'usage_limit' }, + }) expect((await POST(request({ ...body, purpose: 'new-turn' }, directHeaders))).status).toBe(400) expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() }) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index d21004d3e46..93fdeeb3f42 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -7,7 +7,9 @@ import { type NextRequest, NextResponse } from 'next/server' import { COPILOT_BILLING_BLOCKED_CODE, COPILOT_USAGE_LIMIT_EXCEEDED_CODE, + COPILOT_USAGE_UNAVAILABLE_CODE, type ValidateCopilotApiKeyBillingBlocked, + type ValidateCopilotApiKeyRefusal, type ValidateCopilotApiKeyUsageExceeded, validateCopilotApiKeyContract, } from '@/lib/api/contracts/copilot' @@ -26,7 +28,10 @@ import { serializeAccountBillingDecisionHeader, serializeBillingAttributionHeader, } from '@/lib/billing/core/billing-attribution' -import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' +import { + readMidRunAccountUsageVerdict, + readMidRunUsageVerdict, +} from '@/lib/billing/core/mid-run-usage' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { isEnterprisePlan } from '@/lib/billing/core/subscription' import { deriveBillingContext } from '@/lib/billing/core/usage-log' @@ -62,6 +67,8 @@ import { withIncomingGoSpan } from '@/lib/mothership/request/otel' const logger = createLogger('CopilotApiKeysValidate') const CONTINUATION_BLOCKED_MESSAGE = 'Continuation billing account is blocked' +const BILLING_BLOCKED_MESSAGE = 'Billing account is blocked' +const USAGE_UNAVAILABLE_MESSAGE = 'Usage could not be verified. Please try again.' function invalidBillingProtocolResponse(): NextResponse { return NextResponse.json({ error: 'Invalid billing attribution protocol' }, { status: 400 }) @@ -242,13 +249,42 @@ function resolveContinuationBilling( } } -async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise<{ +interface AdmissionUsage { isExceeded: boolean currentUsage: number limit: number scope: string + message?: string + reason?: 'billing_blocked' | 'usage_unavailable' accountBillingDecision?: AccountBillingDecision -}> { +} + +/** The 402 body for a refused new turn: the same bodies continuation validation answers with. */ +async function admissionRefusal( + userId: string, + admission: AdmissionBillingDecision, + usage: AdmissionUsage +): Promise { + if (usage.reason === 'billing_blocked') { + return { code: COPILOT_BILLING_BLOCKED_CODE, error: usage.message ?? BILLING_BLOCKED_MESSAGE } + } + if (usage.reason === 'usage_unavailable') { + return { + code: COPILOT_USAGE_UNAVAILABLE_CODE, + error: usage.message ?? USAGE_UNAVAILABLE_MESSAGE, + } + } + const usageUpgrade = await resolveUsageUpgradePayload( + userId, + admission.kind === 'attributed' || admission.kind === 'legacy-scoped' + ? admission.attribution + : undefined, + usage.scope === 'member' ? 'member' : undefined + ) + return { code: COPILOT_USAGE_LIMIT_EXCEEDED_CODE, error: usageUpgrade.message, usageUpgrade } +} + +async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise { if (admission.kind === 'attributed' || admission.kind === 'legacy-scoped') { const usage = await checkAttributedUsageLimits(admission.attribution) const enforcedUsage = @@ -258,6 +294,8 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: enforcedUsage?.currentUsage ?? 0, limit: enforcedUsage?.limit ?? 0, scope: usage.scope ?? 'payer', + ...(usage.message ? { message: usage.message } : {}), + ...(usage.reason ? { reason: usage.reason } : {}), } } @@ -272,6 +310,8 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: usage.currentUsage, limit: usage.limit, scope: 'account', + ...(usage.message ? { message: usage.message } : {}), + ...(usage.reason ? { reason: usage.reason } : {}), accountBillingDecision: { userId: admission.userId, billingEntity: billingContext.billingEntity, @@ -292,6 +332,8 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: usage.currentUsage, limit: usage.limit, scope: 'account', + ...(usage.message ? { message: usage.message } : {}), + ...(usage.reason ? { reason: usage.reason } : {}), } } @@ -416,10 +458,10 @@ export const POST = withRouteHandler((req: NextRequest) => // original payer's spend through the cached execution usage gate. A read that fails // admits: the run is already under way, and the next re-check reads again. const verdict = - !blocked?.blocked && - purpose === COPILOT_VALIDATION_PURPOSE.continuation && - billing?.kind === 'attributed' - ? await readMidRunUsageVerdict(billing.attribution) + !blocked?.blocked && purpose === COPILOT_VALIDATION_PURPOSE.continuation && billing + ? billing.kind === 'attributed' + ? await readMidRunUsageVerdict(billing.attribution) + : await readMidRunAccountUsageVerdict(billing.decision) : null if (blocked?.blocked || verdict?.status === 'blocked') { span.setAttribute( @@ -440,7 +482,7 @@ export const POST = withRouteHandler((req: NextRequest) => { status: 402 } ) } - if (verdict?.status === 'exceeded' && billing?.kind === 'attributed') { + if (verdict?.status === 'exceeded') { logger.info('[API VALIDATION] Continuation usage exceeded', { userId }) span.setAttribute( TraceAttr.CopilotValidateOutcome, @@ -449,7 +491,7 @@ export const POST = withRouteHandler((req: NextRequest) => span.setAttribute(TraceAttr.HttpStatusCode, 402) const usageUpgrade = await resolveUsageUpgradePayload( userId, - billing.attribution, + billing?.kind === 'attributed' ? billing.attribution : undefined, verdict.scope ) return NextResponse.json( @@ -518,7 +560,10 @@ export const POST = withRouteHandler((req: NextRequest) => }) span.setAttribute(TraceAttr.CopilotValidateOutcome, CopilotValidateOutcome.UsageExceeded) span.setAttribute(TraceAttr.HttpStatusCode, 402) - return new NextResponse(null, { status: 402 }) + return NextResponse.json( + await admissionRefusal(userId, admission, usage), + { status: 402 } + ) } const responseHeaders: Record = {} diff --git a/apps/sim/lib/api/contracts/copilot.ts b/apps/sim/lib/api/contracts/copilot.ts index 37316242a31..3427cfd2ed5 100644 --- a/apps/sim/lib/api/contracts/copilot.ts +++ b/apps/sim/lib/api/contracts/copilot.ts @@ -297,14 +297,27 @@ export type ValidateCopilotApiKeyBillingBlocked = z.output< typeof validateCopilotApiKeyBillingBlockedSchema > +export const COPILOT_USAGE_UNAVAILABLE_CODE = 'USAGE_UNAVAILABLE' + /** - * A 402 from continuation validation. New-turn admission still answers 402 with no body, which - * older callers treat as a refusal of either kind. + * A new turn refused because usage could not be read. Admission fails closed on it; a run + * already under way is never refused for it. */ +export const validateCopilotApiKeyUsageUnavailableSchema = z.object({ + code: z.literal(COPILOT_USAGE_UNAVAILABLE_CODE), + error: z.string(), +}) +export type ValidateCopilotApiKeyUsageUnavailable = z.output< + typeof validateCopilotApiKeyUsageUnavailableSchema +> + +/** Every 402 from validation, new turn and continuation alike, carries one of these bodies. */ export const validateCopilotApiKeyRefusalSchema = z.union([ validateCopilotApiKeyUsageExceededSchema, validateCopilotApiKeyBillingBlockedSchema, + validateCopilotApiKeyUsageUnavailableSchema, ]) +export type ValidateCopilotApiKeyRefusal = z.output export const listCopilotApiKeysContract = defineRouteContract({ method: 'GET', diff --git a/apps/sim/lib/api/contracts/subscription.ts b/apps/sim/lib/api/contracts/subscription.ts index 1d765f5c615..424f17ceb10 100644 --- a/apps/sim/lib/api/contracts/subscription.ts +++ b/apps/sim/lib/api/contracts/subscription.ts @@ -344,25 +344,34 @@ export type UsageUpgradePayload = z.infer * `BillingCallbackResult` reads it, on a 200 and on a duplicate 409 alike. A worker that * predates the fields ignores them. */ -export const billingUsageVerdictSchema = z.object({ - /** The payer is over its usage limit; the worker pauses the run at its next step boundary. */ - usageExceeded: z.boolean(), - /** Present exactly when `usageExceeded` is true: the card the worker writes to its log. */ - usageUpgrade: usageUpgradePayloadSchema.optional(), -}) +export const billingUsageVerdictSchema = z.discriminatedUnion('usageExceeded', [ + z.object({ + /** The payer is within its usage limit, or its standing could not be read. */ + usageExceeded: z.literal(false), + usageUpgrade: z.never().optional(), + }), + z.object({ + /** The payer is over its usage limit; the worker pauses the run at its next step boundary. */ + usageExceeded: z.literal(true), + /** The card the worker writes to its log. */ + usageUpgrade: usageUpgradePayloadSchema, + }), +]) export type BillingUsageVerdict = z.infer -export const billingUpdateCostResponseSchema = billingUsageVerdictSchema.extend({ - success: z.literal(true), - message: z.string().optional(), - data: z.object({ - userId: z.string().optional(), - cost: z.number().optional(), - billingEnabled: z.boolean().optional(), - processedAt: z.string(), - requestId: z.string(), - }), -}) +export const billingUpdateCostResponseSchema = z + .object({ + success: z.literal(true), + message: z.string().optional(), + data: z.object({ + userId: z.string().optional(), + cost: z.number().optional(), + billingEnabled: z.boolean().optional(), + processedAt: z.string(), + requestId: z.string(), + }), + }) + .and(billingUsageVerdictSchema) export type BillingUpdateCostResponse = z.infer export const billingSwitchPlanContract = defineRouteContract({ diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index 318ecc366d2..31dcdbb0550 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.ts @@ -314,6 +314,8 @@ export async function checkServerSideUsageLimits( currentUsage: number limit: number message?: string + /** Why a refusal is not a spent limit: a blocked account, or usage that could not be read. */ + reason?: 'billing_blocked' | 'usage_unavailable' }> { try { if (!isBillingEnabled) { @@ -342,7 +344,13 @@ export async function checkServerSideUsageLimits( const billingPeriod = preloadedBillingContext?.billingPeriod ?? resolveSubscriptionUsagePeriod(sub) ?? { ...defaultBillingPeriod(), source: 'default' } const currentUsage = await getBillingPeriodUsageCost(billingEntity, billingPeriod) - return { isExceeded: true, currentUsage, limit: 0, message: blocked.message } + return { + isExceeded: true, + currentUsage, + limit: 0, + message: blocked.message, + reason: 'billing_blocked', + } } const usageData = await checkUsageStatus(userId, preloadedSubscription, preloadedBillingContext) @@ -359,6 +367,7 @@ export async function checkServerSideUsageLimits( currentUsage: usageData.currentUsage, limit: usageData.limit, message: usageData.isExceeded ? exceededMessage : undefined, + ...(usageData.unavailable ? { reason: 'usage_unavailable' as const } : {}), } } catch (error) { logger.error('Error in server-side usage limit check', { @@ -375,6 +384,7 @@ export async function checkServerSideUsageLimits( isExceeded: true, currentUsage: 0, limit: 0, + reason: 'usage_unavailable', message: error instanceof Error && error.message.includes('No user stats record found') ? 'User account not properly initialized. Please contact support.' diff --git a/apps/sim/lib/billing/core/billing-attribution.ts b/apps/sim/lib/billing/core/billing-attribution.ts index f8211c0eb07..e9112e1c52a 100644 --- a/apps/sim/lib/billing/core/billing-attribution.ts +++ b/apps/sim/lib/billing/core/billing-attribution.ts @@ -725,6 +725,36 @@ function buildBillingAttributionSnapshot(params: { }) } +/** + * The same payer's attribution for its current usage period, for a run that outlived the period + * it was admitted in. The actor, workspace and payer are kept; only the payer's subscription, + * and so its period, is read again, and a subscription that no longer belongs to the payer is + * refused rather than re-selected. + */ +export async function refreshAttributionPeriod( + attribution: BillingAttributionSnapshot +): Promise { + const validated = assertBillingAttributionSnapshot(attribution) + const payerSubscription = validated.organizationId + ? await getOrganizationSubscription(validated.organizationId, { onError: 'throw' }) + : await getHighestPriorityPersonalSubscription(validated.billedAccountUserId, { + onError: 'throw', + }) + const expectedReferenceId = validated.organizationId ?? validated.billedAccountUserId + if (payerSubscription && payerSubscription.referenceId !== expectedReferenceId) { + throw new Error( + `Resolved subscription ${payerSubscription.id} does not belong to payer ${expectedReferenceId}` + ) + } + return buildBillingAttributionSnapshot({ + actorUserId: validated.actorUserId, + workspaceId: validated.workspaceId, + billedAccountUserId: validated.billedAccountUserId, + organizationId: validated.organizationId, + payerSubscription, + }) +} + /** * Resolves the payer from the workspace without consulting the actor's * subscriptions or organization memberships. diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index b7102111b62..7cf04bbb65d 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -1,8 +1,11 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' -import type { - AttributedUsageLimitsResult, - BillingAttributionSnapshot, +import { checkUsageStatus } from '@/lib/billing/calculations/usage-monitor' +import { + type AccountBillingDecision, + type AttributedUsageLimitsResult, + type BillingAttributionSnapshot, + refreshAttributionPeriod, } from '@/lib/billing/core/billing-attribution' import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache' @@ -14,9 +17,9 @@ const logger = createLogger('MidRunUsage') * upgrade card. * - `blocked`: the account is blocked (payment failed, dispute); the run is refused as a blocked * account, never with the upgrade card. - * - `unknown`: the gate could not read usage, or the run's admitted period has ended. Admission - * fails closed on an unreadable ledger, but a run already under way continues: a database - * blip must not end a paying user's long run, and the next step or re-check reads again. + * - `unknown`: usage, or the payer's current period, could not be read. Admission fails closed + * on an unreadable ledger, but a run already under way continues: a database blip must not + * end a paying user's long run, and the next step or re-check reads again. */ export type MidRunUsageVerdict = | { status: 'within' } @@ -24,15 +27,13 @@ export type MidRunUsageVerdict = | { status: 'blocked'; message?: string } | { status: 'unknown' } -export async function readMidRunUsageVerdict( +function periodHasEnded(attribution: BillingAttributionSnapshot): boolean { + return Date.now() >= new Date(attribution.billingPeriod.end).getTime() +} + +async function readGateVerdict( attribution: BillingAttributionSnapshot ): Promise { - // The gate judges the admitted snapshot's period. Once that period has ended it would keep - // counting the old period against the old allowance, so a run just past a reset is not judged - // until its next admission reads the new one. - if (Date.now() >= new Date(attribution.billingPeriod.end).getTime()) { - return { status: 'unknown' } - } let usage: AttributedUsageLimitsResult try { usage = await checkExecutionUsageLimits(attribution) @@ -52,3 +53,65 @@ export async function readMidRunUsageVerdict( } return { status: 'exceeded', ...(usage.scope ? { scope: usage.scope } : {}) } } + +/** + * The gate judges the snapshot's period, so a run that outlived its admitted period is judged + * against the same payer's current period instead, and a read that straddled the period's end + * is judged again against the new one. If the current period cannot be read the verdict is + * unknown, and the run continues. + */ +export async function readMidRunUsageVerdict( + attribution: BillingAttributionSnapshot +): Promise { + let judged = attribution + for (let attempt = 0; attempt < 2; attempt++) { + if (periodHasEnded(judged)) { + try { + judged = await refreshAttributionPeriod(judged) + if (periodHasEnded(judged)) return { status: 'unknown' } + } catch (error) { + logger.warn('Current billing period could not be read; continuing the run', { + error: getErrorMessage(error), + }) + return { status: 'unknown' } + } + } + const verdict = await readGateVerdict(judged) + if (!periodHasEnded(judged)) return verdict + } + return { status: 'unknown' } +} + +/** + * The same verdict for a direct-v1 run billed to an account decision rather than an attributed + * payer, read through the account usage check. A decision whose period has ended is judged + * against the account's current period. + */ +export async function readMidRunAccountUsageVerdict( + decision: AccountBillingDecision +): Promise { + try { + const ended = Date.now() >= new Date(decision.billingPeriod.end).getTime() + const usage = await checkUsageStatus( + decision.userId, + undefined, + ended + ? undefined + : { + billingEntity: decision.billingEntity, + billingPeriod: { + start: new Date(decision.billingPeriod.start), + end: new Date(decision.billingPeriod.end), + ...(decision.billingPeriod.source ? { source: decision.billingPeriod.source } : {}), + }, + } + ) + if (usage.unavailable) return { status: 'unknown' } + return usage.isExceeded ? { status: 'exceeded', scope: 'payer' } : { status: 'within' } + } catch (error) { + logger.warn('Mid-run account usage read failed; continuing the run', { + error: getErrorMessage(error), + }) + return { status: 'unknown' } + } +} diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index 4132cca2238..83862ade045 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -440,7 +440,7 @@ describe('Cumulative billing with PostgreSQL', () => { expect(await stampedWindowTotal(resetStart, resetEnd)).toBeCloseTo(0.6, 9) }) - it('does not hold the rollover while the latest period is still running', async () => { + it('holds an early period-start move until an in-flight top-up commits', async () => { const start = new Date(Date.now() - 24 * 60 * 60 * 1000) const end = new Date(Date.now() + 30 * 24 * 60 * 60 * 1000) await setSubscriptionWindow(start, end) @@ -449,25 +449,32 @@ describe('Cumulative billing with PostgreSQL', () => { const inFlight = charge(0.6, { start, end }) try { await pause.reached.promise - await connection.begin(async (tx) => { - await tx`select set_config('lock_timeout', '300ms', true)` - await tx`update subscription set period_end = period_end where id = 'sub-1'` - }) + const reset = await connection + .begin(async (tx) => { + await tx`select set_config('lock_timeout', '300ms', true)` + await tx`update subscription set period_start = now() at time zone 'UTC' where id = 'sub-1'` + }) + .catch((error: unknown) => error) + expect(getPostgresErrorCode(reset)).toBe('55P03') } finally { pause.release.resolve() await inFlight } + expect(await stampedWindowTotal(start, end)).toBeCloseTo(0.6, 9) }) - it('refuses a request key that could collide with its period rows', async () => { - await setSubscriptionPeriod(0) - await expect( - recordCumulativeUsage({ - ...usage(0.4, 'update-cost:request@1'), - payerSubscriptionId: 'sub-1', - }) - ).rejects.toThrow('@') - expect(await ledgerRows()).toEqual([]) - }) + it.each([ + ['with', { payerSubscriptionId: 'sub-1' }], + ['without', {}], + ])( + 'refuses a request key that could collide with period rows %s a payer subscription', + async (_case, extra) => { + await setSubscriptionPeriod(0) + await expect( + recordCumulativeUsage({ ...usage(0.4, 'update-cost:request@1'), ...extra }) + ).rejects.toThrow('@') + expect(await ledgerRows()).toEqual([]) + } + ) }) }) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index c05e916f46b..a780a2371d5 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -621,7 +621,7 @@ const MAX_CUMULATIVE_PERIOD_ROWS = 12 /** * The ledger key of the `index`-th period a cumulative request rolled into; 0 is the request key. - * A request key may not itself contain `@`, so these keys never collide with another request's. + * No cumulative request key may contain `@`, so these keys never collide with another request's. */ function cumulativePeriodEventKey(eventKey: string, index: number): string { return index === 0 ? eventKey : `${eventKey}@${index}` @@ -784,7 +784,7 @@ export async function recordCumulativeUsage( } const billingContext = await resolveBillingContext(userId, billingEntity, billingPeriod) - if (payerSubscriptionId && eventKey.includes('@')) { + if (eventKey.includes('@')) { throw new Error(`Cumulative usage event "${eventKey}" must not contain "@"`) } @@ -880,24 +880,19 @@ export async function recordCumulativeUsage( return { billed: false, delta: 0, total: recorded, billingPeriod: latestPeriod } } - // The payer's current period. Once the latest row's period has ended, the read is - // share-locked so a rollover of the subscription waits for this top-up to commit, and - // whatever the close later sums for the old period is final. Before that no close can be - // due, and locking every callback would starve the rollover UPDATE for a busy payer. - const periodQuery = payerSubscriptionId - ? tx + // The payer's current period, share-locked so a change to the subscription's period (a + // rollover, or an anchor reset inside the old period) waits for this write to commit, and + // whatever a close later sums for the old period is final. + const [currentPeriod] = payerSubscriptionId + ? await tx .select({ start: subscriptionTable.periodStart, end: subscriptionTable.periodEnd, }) .from(subscriptionTable) .where(eq(subscriptionTable.id, payerSubscriptionId)) + .for('share') .limit(1) - : null - const [currentPeriod] = periodQuery - ? Date.now() >= latestPeriod.end.getTime() - ? await periodQuery.for('share') - : await periodQuery : [] // Only ever forward: a subscription period that does not start after the latest row's diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index cb5583f6a8f..1aecd534a2b 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -908,6 +908,22 @@ describe('copilot go stream helpers', () => { expect(blocked).not.toBeInstanceOf(BillingLimitError) expect(blocked).toMatchObject({ code: 'forbidden', message: 'Billing account frozen.' }) + vi.mocked(fetch).mockResolvedValueOnce( + new Response( + JSON.stringify({ code: 'USAGE_UNAVAILABLE', error: 'Usage could not be verified.' }), + { status: 402 } + ) + ) + await expect( + runStreamLoop( + 'https://example.com/api/tools/resume', + {}, + createStreamingContext(), + turnScopedExecContext(), + { timeout: 1000 } + ) + ).rejects.toMatchObject({ code: 'forbidden' }) + vi.mocked(fetch).mockResolvedValueOnce(new Response(null, { status: 402 })) await expect( runStreamLoop( diff --git a/apps/sim/lib/mothership/request/go/stream.ts b/apps/sim/lib/mothership/request/go/stream.ts index a5b61e7521e..6cd7dbe9692 100644 --- a/apps/sim/lib/mothership/request/go/stream.ts +++ b/apps/sim/lib/mothership/request/go/stream.ts @@ -2,7 +2,10 @@ import { type Context, SpanStatusCode } from '@opentelemetry/api' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { toRecordOrNull } from '@sim/utils/object' -import { COPILOT_BILLING_BLOCKED_CODE } from '@/lib/api/contracts/copilot' +import { + COPILOT_BILLING_BLOCKED_CODE, + COPILOT_USAGE_UNAVAILABLE_CODE, +} from '@/lib/api/contracts/copilot' import { OrchestrationError } from '@/lib/core/orchestration/types' import { ORCHESTRATION_TIMEOUT_MS } from '@/lib/mothership/constants' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' @@ -138,7 +141,11 @@ function backendErrorMessage(status: number, body: string): string { } export class BillingLimitError extends Error { - constructor(public readonly userId: string) { + /** `member` when the actor hit the cap their organization set, so the card names who can raise it. */ + constructor( + public readonly userId: string, + public readonly scope?: 'actor' | 'payer' | 'member' + ) { super('Usage limit reached') this.name = 'BillingLimitError' } @@ -245,9 +252,13 @@ export async function runStreamLoop( const errorText = await response.text().catch(() => '') if (response.status === 402) { - // A blocked account is refused as blocked, never with the usage card. + // A blocked account, or usage that could not be read, is refused as such, never with the + // usage card. const refusal = parseJsonRecord(errorText) - if (refusal?.code === COPILOT_BILLING_BLOCKED_CODE) { + if ( + refusal?.code === COPILOT_BILLING_BLOCKED_CODE || + refusal?.code === COPILOT_USAGE_UNAVAILABLE_CODE + ) { throw new OrchestrationError( 'forbidden', userFacingRejection(refusal.error) ?? BILLING_BLOCKED_MESSAGE diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index 7497348b142..49e91053cdb 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -1,4 +1,5 @@ import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { billingCoreMock, billingCoreMockFns } from '@sim/testing/mocks/billing-core.mock' import { billingUsageGateCacheMock, billingUsageGateCacheMockFns, @@ -15,6 +16,8 @@ vi.mock('@/lib/mothership/application/authorize-chat-callback', () => ({ checkCopilotContinuationBilling: mocks.standing, })) vi.mock('@/lib/billing/core/usage-gate-cache', () => billingUsageGateCacheMock) +vi.mock('@/lib/billing/core/billing', () => billingCoreMock) +const { mockGetOrganizationSubscription } = billingCoreMockFns const { mockCheckExecutionUsageLimits } = billingUsageGateCacheMockFns const attribution = { actorUserId: 'actor', @@ -121,18 +124,39 @@ describe('continuation admission', () => { await expect(refusal).rejects.toBeInstanceOf(OrchestrationError) await expect(refusal).rejects.toThrow('blocked') }) - it('lets a leg run once its admitted period has ended instead of judging the old period', async () => { - mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + it('judges a leg past its admitted period against the payer current period', async () => { + const ended = { + ...attribution, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + } + mockGetOrganizationSubscription.mockResolvedValue({ + id: 'sub-org', + referenceId: 'original-org', + plan: 'team', + status: 'active', + seats: 4, + periodStart: new Date(attribution.billingPeriod.start), + periodEnd: new Date(attribution.billingPeriod.end), + }) + mockCheckExecutionUsageLimits.mockImplementation(async (judged: typeof attribution) => ({ + isExceeded: judged.billingPeriod.end === attribution.billingPeriod.end, + scope: 'payer', + })) + + await expect( + authorizeLifecycleContinuation({ ...context, billingAttribution: ended }) + ).rejects.toBeInstanceOf(BillingLimitError) + mockGetOrganizationSubscription.mockRejectedValue(new Error('subscription read failed')) await expect( - authorizeLifecycleContinuation({ - ...context, - billingAttribution: { - ...attribution, - billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, - }, - }) + authorizeLifecycleContinuation({ ...context, billingAttribution: ended }) ).resolves.toBeUndefined() - expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + it('carries a member cap into the refusal so the card names who can raise it', async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'member' }) + + await expect(authorizeLifecycleContinuation(context)).rejects.toMatchObject({ + scope: 'member', + }) }) }) diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.ts b/apps/sim/lib/mothership/request/lifecycle/admission.ts index 8ec1b7e7f40..efdcba71e38 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.ts @@ -83,6 +83,6 @@ export async function authorizeLifecycleContinuation( const usage = await readMidRunUsageVerdict(context.billingAttribution) if (usage.status === 'blocked') throw new OrchestrationError('forbidden', 'Continuation billing account is blocked') - if (usage.status === 'exceeded') throw new BillingLimitError(context.userId) + if (usage.status === 'exceeded') throw new BillingLimitError(context.userId, usage.scope) } } diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index de13f3c7f3d..303ca50e481 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -100,11 +100,13 @@ vi.mock('@/lib/mothership/request/go/stream', () => { class BillingLimitError extends Error { userId: string + scope?: string - constructor(userId: string) { + constructor(userId: string, scope?: string) { super('Usage limit reached') this.name = 'BillingLimitError' this.userId = userId + this.scope = scope } } @@ -224,6 +226,7 @@ import { MothershipStreamV1ToolOutcome, } from '@/lib/mothership/generated/mothership-stream-v1' import { + BillingLimitError, CopilotBackendError, STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE, StreamEndedWithoutTerminalError, @@ -2423,6 +2426,22 @@ describe('runCopilotLifecycle', () => { expect(result.cancelled).not.toBe(true) }) + it('stops the worker run when the worker itself refuses a leg at the usage limit', async () => { + mockRequestExplicitStreamAbort.mockResolvedValue({ settled: true }) + mockRunStreamLoop.mockRejectedValueOnce(new BillingLimitError('user-1')) + + const result = await runCopilotLifecycle( + { message: 'hello', messageId: 'message-1' }, + { userId: 'user-1', workspaceId: 'ws-1', chatId: 'chat-1', runId: 'run-1' } + ) + + expect(handleBillingLimitResponse).toHaveBeenCalledOnce() + expect(mockRequestExplicitStreamAbort).toHaveBeenCalledWith( + expect.objectContaining({ streamId: 'message-1', chatId: 'chat-1' }) + ) + expect(result.error).toBeUndefined() + }) + it('shows the usage card instead of resuming a run whose payer crossed its limit', async () => { const billingAttribution = { actorUserId: 'user-1', @@ -2473,7 +2492,8 @@ describe('runCopilotLifecycle', () => { 'user-1', expect.anything(), expect.anything(), - expect.anything() + expect.anything(), + 'payer' ) expect(result.cancelled).not.toBe(true) expect(result.error).toBeUndefined() diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index c95be22607f..ce42c2cc84f 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -528,11 +528,17 @@ export async function runCopilotLifecycle( hostedBillingRequest ) } catch (error) { - // A continuation refused on spend, or a worker 402 on a child leg of a subagent - // fan-out, ends the turn with the same card as a refused dispatch. + // A continuation refused on spend, or a worker 402 on any leg, ends the turn with the + // same card as a refused dispatch and stops the worker run. if (!(error instanceof BillingLimitError)) throw error context.awaitingAsyncContinuation = undefined - await handleBillingLimitResponse(error.userId, context, execContext, lifecycleOptions) + await handleBillingLimitResponse( + error.userId, + context, + execContext, + lifecycleOptions, + error.scope + ) await stopWorkerRunAfterUsageRefusal(context.messageId, execContext) } } @@ -1237,10 +1243,6 @@ async function runCheckpointLoop( } catch (streamError) { context.trace.endSpan(streamSpan, RequestTraceV1SpanStatus.error) context.trace.setActiveSpan(undefined) - if (streamError instanceof BillingLimitError) { - await handleBillingLimitResponse(streamError.userId, context, execContext, options) - break - } const backoff = retry?.nextDelay(streamError, options.abortSignal) ?? null if (backoff !== null) { /** A recovered connection must not finalize with an earlier transport failure. */ diff --git a/apps/sim/lib/mothership/request/tools/billing.test.ts b/apps/sim/lib/mothership/request/tools/billing.test.ts index 9755b299073..2154b06e8b3 100644 --- a/apps/sim/lib/mothership/request/tools/billing.test.ts +++ b/apps/sim/lib/mothership/request/tools/billing.test.ts @@ -113,4 +113,20 @@ describe('handleBillingLimitResponse', () => { expect(onEvent.mock.calls.map(([event]) => event.type)).toEqual(['text', 'complete']) }) + + it('names who can raise the cap for a member over the limit their organization set', async () => { + const onEvent = vi.fn() + + await handleBillingLimitResponse( + 'actor-1', + { streamComplete: false } as StreamingContext, + createExecutionContext(), + { onEvent } as OrchestratorOptions, + 'member' + ) + + expect(onEvent.mock.calls[0]?.[0]).toMatchObject({ + payload: { text: expect.stringContaining('limit your organization set for you') }, + }) + }) }) diff --git a/apps/sim/lib/mothership/request/tools/billing.ts b/apps/sim/lib/mothership/request/tools/billing.ts index 09a27517d63..295bf69ddcc 100644 --- a/apps/sim/lib/mothership/request/tools/billing.ts +++ b/apps/sim/lib/mothership/request/tools/billing.ts @@ -27,9 +27,10 @@ export async function handleBillingLimitResponse( userId: string, context: StreamingContext, execContext: ExecutionContext, - options: OrchestratorOptions + options: OrchestratorOptions, + scope?: 'actor' | 'payer' | 'member' ): Promise { - const payload = await resolveUsageUpgradePayload(userId, execContext.billingAttribution) + const payload = await resolveUsageUpgradePayload(userId, execContext.billingAttribution, scope) const syntheticContent = formatUsageUpgradeTag(payload) // The card is this turn's terminal even when the refused leg follows one that already ended. context.streamComplete = false diff --git a/packages/testing/src/mocks/billing-attribution.mock.ts b/packages/testing/src/mocks/billing-attribution.mock.ts index 3ad64d75b51..ed26d3c271b 100644 --- a/packages/testing/src/mocks/billing-attribution.mock.ts +++ b/packages/testing/src/mocks/billing-attribution.mock.ts @@ -162,6 +162,7 @@ export const billingAttributionMockFns = { mockToBillingContext: vi.fn(toBillingContext), mockCheckAttributedBillingBlocks: vi.fn(), mockCheckAttributedUsageLimits: vi.fn(), + mockRefreshAttributionPeriod: vi.fn(), } /** @@ -212,4 +213,5 @@ export const billingAttributionMock = { toBillingContext: billingAttributionMockFns.mockToBillingContext, checkAttributedBillingBlocks: billingAttributionMockFns.mockCheckAttributedBillingBlocks, checkAttributedUsageLimits: billingAttributionMockFns.mockCheckAttributedUsageLimits, + refreshAttributionPeriod: billingAttributionMockFns.mockRefreshAttributionPeriod, } From a7b74b85a3bfcd449d212eef1fe652682ef0fe4b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 04:34:30 -0700 Subject: [PATCH 04/13] fix(billing): judge mid-run usage against the payer's current period - Mid-run checks always judge the admitted payer's current subscription period (cached for a minute), so an early anchor move or a rollover is judged against the period charges now land in; a straddling read is judged again against the next period. - A direct-v1 continuation checks the payer saved in its account decision, against that payer's current period, never a payer chosen from the actor's current memberships. - An unreadable usage read no longer reports a spent-limit message; new turns refused for it get neutral copy. - Dispatch-time refusals pass the verdict scope, so a member over the cap their organization set gets the member card. --- .../app/api/billing/update-cost/route.test.ts | 19 ++++ .../copilot/api-keys/validate/route.test.ts | 69 +++++++++++- .../api/copilot/api-keys/validate/route.ts | 7 +- .../calculations/usage-monitor.test.ts | 17 +++ .../lib/billing/calculations/usage-monitor.ts | 7 +- apps/sim/lib/billing/constants.ts | 4 + .../billing/core/billing-attribution.test.ts | 7 +- .../lib/billing/core/billing-attribution.ts | 6 +- apps/sim/lib/billing/core/mid-run-usage.ts | 105 ++++++++++++------ .../request/lifecycle/admission.test.ts | 13 ++- .../mothership/request/lifecycle/run.test.ts | 10 +- .../lib/mothership/request/lifecycle/run.ts | 8 +- 12 files changed, 223 insertions(+), 49 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index 778bc8256bc..a87edf49c11 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -49,6 +49,7 @@ vi.mock('@/lib/billing/threshold-billing', () => ({ })) import { billingUpdateCostResponseSchema } from '@/lib/api/contracts/subscription' +import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage' import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { BillingCallbackBody, @@ -896,11 +897,13 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { beforeEach(() => { resetUsageGateCache() + resetMidRunPeriodCache() setEnvFlags({ isBillingEnabled: true, isHosted: true }) mockCheckInternalApiKey.mockReturnValue({ success: true }) mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 }) mockCheckAndBillPayerOverageThreshold.mockResolvedValue(undefined) mockRequireBillingAttributionHeader.mockReturnValue(CURRENT_ATTRIBUTION) + mockRefreshAttributionPeriod.mockImplementation(async (attribution: unknown) => attribution) mockToBillingContext.mockReturnValue({ billingEntity: { type: 'organization', id: 'org-1' }, billingPeriod: { @@ -1099,6 +1102,22 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(body.usageExceeded).toBe(true) }) + it('judges a run whose payer period moved early against the moved period', async () => { + const moved = { + ...CURRENT_ATTRIBUTION, + billingPeriod: { start: '2026-07-15T00:00:00.000Z', end: '2099-02-01T00:00:00.000Z' }, + } + mockRefreshAttributionPeriod.mockResolvedValue(moved) + mockCheckAttributedUsageLimits.mockImplementation(async (attribution: typeof moved) => ({ + isExceeded: attribution.billingPeriod.start === moved.billingPeriod.start, + scope: 'payer', + })) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(true) + }) + it('keeps a run going when its current period cannot be read', async () => { mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) mockRefreshAttributionPeriod.mockRejectedValue(new Error('subscription read timed out')) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 0fe4c86abfc..218b9397d37 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -6,6 +6,7 @@ import { schemaMock, setEnvFlags, } from '@sim/testing' +import { billingCoreMock, billingCoreMockFns } from '@sim/testing/mocks/billing-core.mock' import { billingPlanMock, billingPlanMockFns } from '@sim/testing/mocks/billing-plan.mock' import { billingSubscriptionMock, @@ -118,6 +119,8 @@ vi.mock('@/lib/billing/calculations/usage-monitor', () => billingUsageMonitorMoc vi.mock('@/lib/billing/core/plan', () => billingPlanMock) +vi.mock('@/lib/billing/core/billing', () => billingCoreMock) + vi.mock('@/lib/billing/core/subscription', () => billingSubscriptionMock) vi.mock('@/lib/billing/core/usage-log', () => billingUsageLogMock) @@ -138,6 +141,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) vi.mock('@/lib/workspaces/utils', () => workspacesUtilsMock) import { validateCopilotApiKeyBodySchema } from '@/lib/api/contracts/copilot' +import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage' import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { POST } from '@/app/api/copilot/api-keys/validate/route' @@ -147,6 +151,7 @@ const { mockAuthorizeOrganizationChatDelegation: mockAuthorizeOrganizationChat } mothershipOrganizationChatsMockFns const { mockDeriveBillingContext } = billingUsageLogMockFns const { mockGetHighestPrioritySubscription } = billingPlanMockFns +const { mockGetOrganizationSubscription } = billingCoreMockFns const { mockCheckServerSideUsageLimits, mockCheckUsageStatus } = billingUsageMonitorMockFns const mockIsEnterprisePlan = billingSubscriptionMockFns.mockIsEnterprisePlan @@ -510,7 +515,16 @@ describe('validation lifecycle purposes', () => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 1, limit: 10 }) mockIsEnterprisePlan.mockResolvedValue(false) + mockGetOrganizationSubscription.mockResolvedValue({ + id: 'sub-org-1', + referenceId: 'org-1', + plan: 'enterprise', + status: 'active', + periodStart: new Date(ATTRIBUTION.billingPeriod.start), + periodEnd: new Date(ATTRIBUTION.billingPeriod.end), + }) resetUsageGateCache() + resetMidRunPeriodCache() }) it('defaults older callers to full admission and rejects unknown purposes', () => { @@ -531,7 +545,6 @@ describe('validation lifecycle purposes', () => { expect(mockAuthorizeCallback.mock.invocationCallOrder[0]).toBeLessThan( mockCheckContinuationBilling.mock.invocationCallOrder[0] ) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(ATTRIBUTION) expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() expect(mockResolveLegacyV0BillingAttribution).not.toHaveBeenCalled() expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled() @@ -729,6 +742,60 @@ describe('validation lifecycle purposes', () => { expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(3) }) + it('refuses a new turn whose usage cannot be read with neutral copy', async () => { + mockCheckAttributedUsageLimits.mockResolvedValueOnce({ + isExceeded: true, + reason: 'usage_unavailable', + message: 'Usage limit exceeded: $0.00 used of $0.00 limit.', + scope: 'payer', + }) + const refused = await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders)) + expect(refused.status).toBe(402) + const refusal = await refused.json() + expect(refusal.code).toBe('USAGE_UNAVAILABLE') + expect(refusal.error).not.toMatch(/\$/) + }) + + it('checks the payer saved at admission for a direct-v1 run whose actor changed orgs', async () => { + const endedDecision = { + ...ACCOUNT_BILLING_DECISION, + billingPeriod: { start: '2026-06-01T00:00:00.000Z', end: '2026-07-01T00:00:00.000Z' }, + } + mockGetHighestPrioritySubscription.mockResolvedValue({ + id: 'sub-new-org', + referenceId: 'new-org', + plan: 'team', + status: 'active', + periodStart: new Date('2026-07-01T00:00:00.000Z'), + periodEnd: new Date('2099-01-01T00:00:00.000Z'), + }) + mockGetOrganizationSubscription.mockResolvedValue({ + id: 'sub-account-org', + referenceId: 'account-org', + plan: 'team', + status: 'active', + periodStart: new Date('2026-07-01T00:00:00.000Z'), + periodEnd: new Date('2099-01-01T00:00:00.000Z'), + }) + mockCheckUsageStatus.mockImplementation( + async ( + _userId: string, + _subscription: unknown, + context?: { billingEntity: { id: string } } + ) => ({ + isExceeded: context?.billingEntity.id === 'account-org', + currentUsage: 12, + limit: 10, + }) + ) + + const response = await POST( + request(body, { ...directHeaders, 'x-sim-billing-account-decision': encode(endedDecision) }) + ) + + expect(response.status).toBe(402) + }) + it('refuses a blocked new turn with the blocked body the contract declares', async () => { mockCheckAttributedUsageLimits.mockResolvedValueOnce({ isExceeded: true, diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 93fdeeb3f42..094a2c31df8 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -15,6 +15,7 @@ import { } from '@/lib/api/contracts/copilot' import { parseRequest, validationErrorResponse } from '@/lib/api/server' import { checkServerSideUsageLimits } from '@/lib/billing/calculations/usage-monitor' +import { USAGE_UNAVAILABLE_MESSAGE } from '@/lib/billing/constants' import { type AccountBillingDecision, type BillingAttributionSnapshot, @@ -68,7 +69,6 @@ const logger = createLogger('CopilotApiKeysValidate') const CONTINUATION_BLOCKED_MESSAGE = 'Continuation billing account is blocked' const BILLING_BLOCKED_MESSAGE = 'Billing account is blocked' -const USAGE_UNAVAILABLE_MESSAGE = 'Usage could not be verified. Please try again.' function invalidBillingProtocolResponse(): NextResponse { return NextResponse.json({ error: 'Invalid billing attribution protocol' }, { status: 400 }) @@ -269,10 +269,7 @@ async function admissionRefusal( return { code: COPILOT_BILLING_BLOCKED_CODE, error: usage.message ?? BILLING_BLOCKED_MESSAGE } } if (usage.reason === 'usage_unavailable') { - return { - code: COPILOT_USAGE_UNAVAILABLE_CODE, - error: usage.message ?? USAGE_UNAVAILABLE_MESSAGE, - } + return { code: COPILOT_USAGE_UNAVAILABLE_CODE, error: USAGE_UNAVAILABLE_MESSAGE } } const usageUpgrade = await resolveUsageUpgradePayload( userId, diff --git a/apps/sim/lib/billing/calculations/usage-monitor.test.ts b/apps/sim/lib/billing/calculations/usage-monitor.test.ts index 4bd444b3b90..a571dbfe070 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.test.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.test.ts @@ -213,6 +213,23 @@ describe('checkServerSideUsageLimits', () => { mockGetBillingPeriodUsageCost.mockResolvedValue(125) }) + it('does not describe an unreadable ledger as a spent limit', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ blocked: false }]) + mockGetBillingPeriodUsageCost.mockRejectedValueOnce(new Error('canceling statement')) + + const result = await checkServerSideUsageLimits('user-1', { + referenceId: 'user-1', + plan: 'free', + status: 'active', + seats: 1, + periodStart: new Date('2026-06-01T00:00:00.000Z'), + periodEnd: new Date('2026-07-01T00:00:00.000Z'), + }) + + expect(result).toMatchObject({ isExceeded: true, reason: 'usage_unavailable' }) + expect(result.message ?? '').not.toMatch(/\$/) + }) + it('keeps blocked accounts blocked while reporting their real ledger usage', async () => { dbChainMockFns.limit.mockResolvedValueOnce([{ blocked: true, blockedReason: 'payment_failed' }]) const subscription = { diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index 31dcdbb0550..811b812f10a 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.ts @@ -3,6 +3,7 @@ import { userStats } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { eq } from 'drizzle-orm' +import { USAGE_UNAVAILABLE_MESSAGE } from '@/lib/billing/constants' import { isOrganizationBillingBlocked } from '@/lib/billing/core/access' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' @@ -366,7 +367,11 @@ export async function checkServerSideUsageLimits( isExceeded: usageData.isExceeded, currentUsage: usageData.currentUsage, limit: usageData.limit, - message: usageData.isExceeded ? exceededMessage : undefined, + message: usageData.unavailable + ? USAGE_UNAVAILABLE_MESSAGE + : usageData.isExceeded + ? exceededMessage + : undefined, ...(usageData.unavailable ? { reason: 'usage_unavailable' as const } : {}), } } catch (error) { diff --git a/apps/sim/lib/billing/constants.ts b/apps/sim/lib/billing/constants.ts index c1fd884a293..a3db9fe9acf 100644 --- a/apps/sim/lib/billing/constants.ts +++ b/apps/sim/lib/billing/constants.ts @@ -103,3 +103,7 @@ export const ANNUAL_DISCOUNT_RATE = 0.15 * Effectively unlimited — any limit >= this threshold is treated as uncapped. */ export const ON_DEMAND_UNLIMITED = 999999 + +/** Shown when usage could not be read, instead of a limit the read never measured. */ +export const USAGE_UNAVAILABLE_MESSAGE = + 'Usage could not be verified right now. Please try again in a moment.' diff --git a/apps/sim/lib/billing/core/billing-attribution.test.ts b/apps/sim/lib/billing/core/billing-attribution.test.ts index bf4f3ddb521..3bf33ee85be 100644 --- a/apps/sim/lib/billing/core/billing-attribution.test.ts +++ b/apps/sim/lib/billing/core/billing-attribution.test.ts @@ -449,10 +449,9 @@ describe('checkAttributedUsageLimits', () => { scope: 'user', unavailable: true, }) - await expect(checkAttributedUsageLimits(attribution)).resolves.toMatchObject({ - isExceeded: true, - reason: 'usage_unavailable', - }) + const unavailable = await checkAttributedUsageLimits(attribution) + expect(unavailable).toMatchObject({ isExceeded: true, reason: 'usage_unavailable' }) + expect(unavailable.message ?? '').not.toMatch(/\$/) }) it('returns payer exhaustion before checking the actor member cap', async () => { diff --git a/apps/sim/lib/billing/core/billing-attribution.ts b/apps/sim/lib/billing/core/billing-attribution.ts index e9112e1c52a..938357f5654 100644 --- a/apps/sim/lib/billing/core/billing-attribution.ts +++ b/apps/sim/lib/billing/core/billing-attribution.ts @@ -10,6 +10,7 @@ import { checkUsageStatus, } from '@/lib/billing/calculations/usage-monitor' import { parseBillingConcurrencyLimit } from '@/lib/billing/concurrency-defaults' +import { USAGE_UNAVAILABLE_MESSAGE } from '@/lib/billing/constants' import { getOrganizationSubscription } from '@/lib/billing/core/billing' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/plan' @@ -970,8 +971,9 @@ export async function checkAttributedUsageLimits( if (payerUsage.isExceeded) { const formattedUsage = payerUsage.currentUsage.toFixed(2) const formattedLimit = payerUsage.limit.toFixed(2) - const message = - validatedAttribution.billingEntity.type === 'organization' + const message = payerUsage.unavailable + ? USAGE_UNAVAILABLE_MESSAGE + : validatedAttribution.billingEntity.type === 'organization' ? `Organization usage limit exceeded: $${formattedUsage} pooled of $${formattedLimit} organization limit. Ask a team admin to raise the organization usage limit to continue.` : `Usage limit exceeded: $${formattedUsage} used of $${formattedLimit} limit. Please upgrade your plan or raise your usage limit to continue.` diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 7cf04bbb65d..5ea9552674d 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -1,12 +1,17 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' +import { LRUCache } from 'lru-cache' import { checkUsageStatus } from '@/lib/billing/calculations/usage-monitor' +import { getOrganizationSubscription } from '@/lib/billing/core/billing' import { type AccountBillingDecision, type AttributedUsageLimitsResult, type BillingAttributionSnapshot, refreshAttributionPeriod, } from '@/lib/billing/core/billing-attribution' +import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' +import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/plan' +import { resolveSubscriptionUsagePeriod } from '@/lib/billing/core/reporting-period' import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache' const logger = createLogger('MidRunUsage') @@ -27,6 +32,40 @@ export type MidRunUsageVerdict = | { status: 'blocked'; message?: string } | { status: 'unknown' } +/** + * How long a payer's current period stays cached for mid-run checks. Every model step settles a + * cost callback that reads it; a subscription period change (rollover, anchor reset) reaches the + * check within this long. + */ +const CURRENT_PERIOD_TTL_MS = 60 * 1000 + +const currentPeriodCache = new LRUCache({ + max: 10_000, + ttl: CURRENT_PERIOD_TTL_MS, +}) + +function currentPeriodKey(attribution: BillingAttributionSnapshot): string { + return [ + attribution.actorUserId, + attribution.workspaceId ?? '', + attribution.organizationId ?? '', + attribution.billedAccountUserId, + ].join(':') +} + +/** The admitted payer's attribution for its current subscription period. */ +async function currentAttribution( + attribution: BillingAttributionSnapshot, + fresh: boolean +): Promise { + const key = currentPeriodKey(attribution) + const cached = fresh ? undefined : currentPeriodCache.get(key) + if (cached) return cached + const current = await refreshAttributionPeriod(attribution) + currentPeriodCache.set(key, current) + return current +} + function periodHasEnded(attribution: BillingAttributionSnapshot): boolean { return Date.now() >= new Date(attribution.billingPeriod.end).getTime() } @@ -55,26 +94,25 @@ async function readGateVerdict( } /** - * The gate judges the snapshot's period, so a run that outlived its admitted period is judged - * against the same payer's current period instead, and a read that straddled the period's end - * is judged again against the new one. If the current period cannot be read the verdict is - * unknown, and the run continues. + * Judges a run against its admitted payer's CURRENT subscription period, never the period it was + * admitted in: charges land in whatever period the subscription is in now (a rollover or an + * early anchor reset included), so that is the allowance they count against. A read that + * straddles the end of that period is judged again against the next one. If the current period + * cannot be read the verdict is unknown, and the run continues. */ export async function readMidRunUsageVerdict( attribution: BillingAttributionSnapshot ): Promise { - let judged = attribution for (let attempt = 0; attempt < 2; attempt++) { - if (periodHasEnded(judged)) { - try { - judged = await refreshAttributionPeriod(judged) - if (periodHasEnded(judged)) return { status: 'unknown' } - } catch (error) { - logger.warn('Current billing period could not be read; continuing the run', { - error: getErrorMessage(error), - }) - return { status: 'unknown' } - } + let judged: BillingAttributionSnapshot + try { + judged = await currentAttribution(attribution, attempt > 0) + if (periodHasEnded(judged)) return { status: 'unknown' } + } catch (error) { + logger.warn('Current billing period could not be read; continuing the run', { + error: getErrorMessage(error), + }) + return { status: 'unknown' } } const verdict = await readGateVerdict(judged) if (!periodHasEnded(judged)) return verdict @@ -84,28 +122,26 @@ export async function readMidRunUsageVerdict( /** * The same verdict for a direct-v1 run billed to an account decision rather than an attributed - * payer, read through the account usage check. A decision whose period has ended is judged - * against the account's current period. + * payer. The payer is the one saved in the decision at admission, never re-selected from the + * actor's current memberships, judged against that payer's current subscription period. */ export async function readMidRunAccountUsageVerdict( decision: AccountBillingDecision ): Promise { try { - const ended = Date.now() >= new Date(decision.billingPeriod.end).getTime() - const usage = await checkUsageStatus( - decision.userId, - undefined, - ended - ? undefined - : { - billingEntity: decision.billingEntity, - billingPeriod: { - start: new Date(decision.billingPeriod.start), - end: new Date(decision.billingPeriod.end), - ...(decision.billingPeriod.source ? { source: decision.billingPeriod.source } : {}), - }, - } - ) + const payer = decision.billingEntity + const subscription = + payer.type === 'organization' + ? await getOrganizationSubscription(payer.id, { onError: 'throw' }) + : await getHighestPriorityPersonalSubscription(payer.id, { onError: 'throw' }) + const billingPeriod = resolveSubscriptionUsagePeriod(subscription) ?? { + ...defaultBillingPeriod(), + source: 'default' as const, + } + const usage = await checkUsageStatus(decision.userId, subscription, { + billingEntity: payer, + billingPeriod, + }) if (usage.unavailable) return { status: 'unknown' } return usage.isExceeded ? { status: 'exceeded', scope: 'payer' } : { status: 'within' } } catch (error) { @@ -115,3 +151,8 @@ export async function readMidRunAccountUsageVerdict( return { status: 'unknown' } } } + +/** Drops every cached current period. Test seam; never called in production code. */ +export function resetMidRunPeriodCache(): void { + currentPeriodCache.clear() +} diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index 49e91053cdb..471a3d9d1a1 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -6,6 +6,7 @@ import { } from '@sim/testing/mocks/billing-usage-gate-cache.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { createAttributedBillingRequestEnvelope } from '@/lib/billing/core/billing-attribution' +import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage' import { OrchestrationError } from '@/lib/core/orchestration/types' import { BillingLimitError } from '@/lib/mothership/request/go/stream' import { authorizeLifecycleContinuation, restoreBillingAdmission } from './admission' @@ -39,6 +40,16 @@ beforeEach(() => { setEnvFlags({ isHosted: true }) mocks.standing.mockResolvedValue({ blocked: false }) mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: false }) + mockGetOrganizationSubscription.mockResolvedValue({ + id: 'sub-org', + referenceId: 'original-org', + plan: 'team', + status: 'active', + seats: 4, + periodStart: new Date(attribution.billingPeriod.start), + periodEnd: new Date(attribution.billingPeriod.end), + }) + resetMidRunPeriodCache() }) afterEach(resetEnvFlagsMock) @@ -88,7 +99,6 @@ describe('continuation admission', () => { await expect(refusal).rejects.toBeInstanceOf(BillingLimitError) await expect(refusal).rejects.toMatchObject({ userId: 'actor' }) - expect(mockCheckExecutionUsageLimits).toHaveBeenCalledWith(attribution) }) it('keeps a blocked account a forbidden refusal without reading spend', async () => { mocks.standing.mockResolvedValue({ blocked: true }) @@ -147,6 +157,7 @@ describe('continuation admission', () => { authorizeLifecycleContinuation({ ...context, billingAttribution: ended }) ).rejects.toBeInstanceOf(BillingLimitError) + resetMidRunPeriodCache() mockGetOrganizationSubscription.mockRejectedValue(new Error('subscription read failed')) await expect( authorizeLifecycleContinuation({ ...context, billingAttribution: ended }) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 303ca50e481..0bbcb68ce37 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -2373,7 +2373,12 @@ describe('runCopilotLifecycle', () => { } ) expect(mockCheckAttributedUsageLimits).toHaveBeenCalledOnce() - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(attribution) + expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith( + expect.objectContaining({ + billedAccountUserId: 'original-owner', + billingEntity: attribution.billingEntity, + }) + ) expect(onBillingAdmission).not.toHaveBeenCalled() expect(continuationAuth).toHaveBeenCalled() expect(mockRunStreamLoop).toHaveBeenCalledOnce() @@ -2404,7 +2409,7 @@ describe('runCopilotLifecycle', () => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, message: 'limit reached', - scope: 'payer', + scope: 'member', }) const result = await runCopilotLifecycle( @@ -2422,6 +2427,7 @@ describe('runCopilotLifecycle', () => { expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(billingAttribution) expect(handleBillingLimitResponse).toHaveBeenCalledTimes(1) + expect(vi.mocked(handleBillingLimitResponse).mock.calls[0][4]).toBe('member') expect(mockRunStreamLoop).not.toHaveBeenCalled() expect(result.cancelled).not.toBe(true) }) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index ce42c2cc84f..699a1e6d437 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -509,7 +509,13 @@ export async function runCopilotLifecycle( // The worker terminal was already delivered before the relay died. // Rebuild persistence from that receipt without charging its usage twice. } else if (admission.isExceeded) { - await handleBillingLimitResponse(execContext.userId, context, execContext, lifecycleOptions) + await handleBillingLimitResponse( + execContext.userId, + context, + execContext, + lifecycleOptions, + 'scope' in admission ? admission.scope : undefined + ) } else { if (!isContinuation && hostedBillingRequest) await lifecycleOptions.onBillingAdmission?.(hostedBillingRequest) From 51725ac1b106accc9fb994717ceecd88915a142c Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 04:42:19 -0700 Subject: [PATCH 05/13] chore(billing): sync the worker usage refusal contract --- apps/sim/lib/mothership/generated/billing.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/apps/sim/lib/mothership/generated/billing.ts b/apps/sim/lib/mothership/generated/billing.ts index 34072228ebf..abdec1159ad 100644 --- a/apps/sim/lib/mothership/generated/billing.ts +++ b/apps/sim/lib/mothership/generated/billing.ts @@ -81,10 +81,13 @@ export const BillingCallbackResult = z.object({ usageUpgrade: UsageUpgrade.optional().catch(undefined), }); -/** A continuation refused for the usage limit. A body-less 402 is a blocked account instead. */ +/** + * A continuation refused for the usage limit. A body-less 402 is a blocked account instead. + * The code decides; a malformed card falls back to the default one. + */ export const UsageLimitRefusal = z.object({ code: z.literal("USAGE_LIMIT_EXCEEDED"), - usageUpgrade: UsageUpgrade, + usageUpgrade: UsageUpgrade.optional().catch(undefined), }); export const BillingDuplicateCode = "DUPLICATE_BILLING_EVENT"; From 3a7134d0f280d9731fde973aad4b310b2a469be7 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 04:50:31 -0700 Subject: [PATCH 06/13] fix(billing): reload an ended cached period, keep org payers org-scoped, keep blocked accounts blocked --- .../app/api/billing/update-cost/route.test.ts | 20 +++++++ .../copilot/api-keys/validate/route.test.ts | 13 +++++ .../calculations/usage-monitor.test.ts | 13 +++++ .../lib/billing/calculations/usage-monitor.ts | 53 +++++++++++++------ apps/sim/lib/billing/core/mid-run-usage.ts | 19 ++++++- 5 files changed, 101 insertions(+), 17 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index a87edf49c11..db604e7273b 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -1153,6 +1153,26 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(body.usageExceeded).toBe(false) }) + it('reloads a cached current period once it has ended', async () => { + const ending = { + ...CURRENT_ATTRIBUTION, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: new Date(Date.now() + 50).toISOString(), + }, + } + mockRefreshAttributionPeriod + .mockResolvedValueOnce(ending) + .mockResolvedValue(CURRENT_ATTRIBUTION) + refuseOnlyCurrentPeriod() + await POST(attributedCallback()) + await sleep(100) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(true) + }) + it('never answers a verdict whose card and flag disagree', () => { const base = { success: true, diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 218b9397d37..859768fad1f 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -796,6 +796,19 @@ describe('validation lifecycle purposes', () => { expect(response.status).toBe(402) }) + it('judges a direct-v1 organization payer without a subscription as that organization', async () => { + mockGetOrganizationSubscription.mockResolvedValue(null) + mockCheckUsageStatus.mockImplementation( + async (_userId: string, subscription: { referenceId?: string } | null) => ({ + isExceeded: subscription?.referenceId === 'account-org', + currentUsage: 12, + limit: 10, + }) + ) + + expect((await POST(request(body, directHeaders))).status).toBe(402) + }) + it('refuses a blocked new turn with the blocked body the contract declares', async () => { mockCheckAttributedUsageLimits.mockResolvedValueOnce({ isExceeded: true, diff --git a/apps/sim/lib/billing/calculations/usage-monitor.test.ts b/apps/sim/lib/billing/calculations/usage-monitor.test.ts index a571dbfe070..c45e971f687 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.test.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.test.ts @@ -230,6 +230,19 @@ describe('checkServerSideUsageLimits', () => { expect(result.message ?? '').not.toMatch(/\$/) }) + it('keeps a blocked account blocked when its ledger cannot be read', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ blocked: true, blockedReason: 'dispute' }]) + mockGetBillingPeriodUsageCost.mockRejectedValueOnce(new Error('canceling statement')) + + const result = await checkServerSideUsageLimits('user-1', null) + + expect(result).toMatchObject({ + isExceeded: true, + reason: 'billing_blocked', + message: 'Account frozen. Please contact support to resolve this issue.', + }) + }) + it('keeps blocked accounts blocked while reporting their real ledger usage', async () => { dbChainMockFns.limit.mockResolvedValueOnce([{ blocked: true, blockedReason: 'payment_failed' }]) const subscription = { diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index 811b812f10a..035551c0beb 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.ts @@ -1,7 +1,7 @@ import { db } from '@sim/db' import { userStats } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { toError } from '@sim/utils/errors' +import { getErrorMessage, toError } from '@sim/utils/errors' import { eq } from 'drizzle-orm' import { USAGE_UNAVAILABLE_MESSAGE } from '@/lib/billing/constants' import { isOrganizationBillingBlocked } from '@/lib/billing/core/access' @@ -306,6 +306,38 @@ export async function checkBillingEntityBlocked( * @param userId The ID of the user to check * @returns An object containing the exceeded status and usage details */ +/** + * A blocked account's ledger usage, for display only: `/api/users/me/usage-limits` exposes it as + * `currentPeriodCost`. The account is blocked either way, so a failed read reports zero rather + * than turning the blocked refusal into an unreadable-usage one. + */ +async function readBlockedAccountUsage( + userId: string, + preloadedSubscription: UsageLimitSubscription | null | undefined, + preloadedBillingContext: BillingContext | undefined +): Promise { + try { + const sub = + preloadedSubscription !== undefined + ? preloadedSubscription + : await getHighestPrioritySubscription(userId) + const subIsOrgScoped = isOrgScopedSubscription(sub, userId) + const billingEntity: BillingEntity = + subIsOrgScoped && sub + ? { type: 'organization', id: sub.referenceId } + : { type: 'user', id: userId } + const billingPeriod = preloadedBillingContext?.billingPeriod ?? + resolveSubscriptionUsagePeriod(sub) ?? { ...defaultBillingPeriod(), source: 'default' } + return await getBillingPeriodUsageCost(billingEntity, billingPeriod) + } catch (error) { + logger.warn('Blocked account usage could not be read', { + userId, + error: getErrorMessage(error), + }) + return 0 + } +} + export async function checkServerSideUsageLimits( userId: string, preloadedSubscription?: UsageLimitSubscription | null, @@ -331,20 +363,11 @@ export async function checkServerSideUsageLimits( const blocked = await checkBillingBlocked(userId) if (blocked.blocked) { - // Enforcement stays blocked, but surfaced usage must be the real ledger - // value — `/api/users/me/usage-limits` exposes it as `currentPeriodCost`. - const sub = - preloadedSubscription !== undefined - ? preloadedSubscription - : await getHighestPrioritySubscription(userId) - const subIsOrgScoped = isOrgScopedSubscription(sub, userId) - const billingEntity: BillingEntity = - subIsOrgScoped && sub - ? { type: 'organization', id: sub.referenceId } - : { type: 'user', id: userId } - const billingPeriod = preloadedBillingContext?.billingPeriod ?? - resolveSubscriptionUsagePeriod(sub) ?? { ...defaultBillingPeriod(), source: 'default' } - const currentUsage = await getBillingPeriodUsageCost(billingEntity, billingPeriod) + const currentUsage = await readBlockedAccountUsage( + userId, + preloadedSubscription, + preloadedBillingContext + ) return { isExceeded: true, currentUsage, diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 5ea9552674d..5104c0dfcfc 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -60,7 +60,8 @@ async function currentAttribution( ): Promise { const key = currentPeriodKey(attribution) const cached = fresh ? undefined : currentPeriodCache.get(key) - if (cached) return cached + // A cached period that has since ended is stale: the payer may already be in the next one. + if (cached && !periodHasEnded(cached)) return cached const current = await refreshAttributionPeriod(attribution) currentPeriodCache.set(key, current) return current @@ -138,7 +139,21 @@ export async function readMidRunAccountUsageVerdict( ...defaultBillingPeriod(), source: 'default' as const, } - const usage = await checkUsageStatus(decision.userId, subscription, { + // An organization payer without a subscription stays organization-scoped on the free plan, + // as `toUsageLimitSubscription` does for attributed runs, never the actor's personal ledger. + const usageSubscription = + subscription ?? + (payer.type === 'organization' + ? { + referenceId: payer.id, + plan: 'free', + status: null, + seats: null, + periodStart: billingPeriod.start, + periodEnd: billingPeriod.end, + } + : null) + const usage = await checkUsageStatus(decision.userId, usageSubscription, { billingEntity: payer, billingPeriod, }) From e1d11025ff618953539c4ac1556754b4787633fe Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 10:12:16 -0700 Subject: [PATCH 07/13] fix(billing): cache admitted direct-v1 continuation verdicts and bound the callback's standing read - A direct-v1 continuation re-read the payer's full period ledger on every resume leg. Its admitted verdict is now served for the execution gate's TTL, with concurrent misses coalesced, like the attributed path; a refusal or an unreadable ledger is always read again, and the read is skipped when billing is off. - The cost callback waits at most 1 s on the payer's standing, well inside the worker's 5 s callback timeout, and answers not exceeded past it; the abandoned read still caches its admission. - The straddling-period test now reaches the re-judge branch. --- .../app/api/billing/update-cost/route.test.ts | 24 ++++++-- apps/sim/app/api/billing/update-cost/route.ts | 25 +++++++- .../copilot/api-keys/validate/route.test.ts | 26 ++++++++- apps/sim/lib/billing/core/mid-run-usage.ts | 57 ++++++++++++++++--- .../request/lifecycle/admission.test.ts | 6 +- 5 files changed, 118 insertions(+), 20 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index db604e7273b..b25eb226b63 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -49,7 +49,7 @@ vi.mock('@/lib/billing/threshold-billing', () => ({ })) import { billingUpdateCostResponseSchema } from '@/lib/api/contracts/subscription' -import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage' +import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage' import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { BillingCallbackBody, @@ -897,7 +897,7 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { beforeEach(() => { resetUsageGateCache() - resetMidRunPeriodCache() + resetMidRunUsageCaches() setEnvFlags({ isBillingEnabled: true, isHosted: true }) mockCheckInternalApiKey.mockReturnValue({ success: true }) mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 }) @@ -1136,8 +1136,9 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { end: new Date(Date.now() + 40).toISOString(), }, } - mockRequireBillingAttributionHeader.mockReturnValue(straddling) - mockRefreshAttributionPeriod.mockResolvedValue(CURRENT_ATTRIBUTION) + mockRefreshAttributionPeriod + .mockResolvedValueOnce(straddling) + .mockResolvedValue(CURRENT_ATTRIBUTION) mockCheckAttributedUsageLimits.mockImplementation( async (attribution: typeof CURRENT_ATTRIBUTION) => { if (attribution.billingPeriod.end !== straddling.billingPeriod.end) { @@ -1151,6 +1152,21 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { const body = await (await POST(attributedCallback())).json() expect(body.usageExceeded).toBe(false) + expect(mockRefreshAttributionPeriod).toHaveBeenCalledTimes(2) + }) + + it('answers not exceeded when the standing read outlasts the callback budget', async () => { + mockCheckAttributedUsageLimits.mockImplementation(async () => { + await sleep(1500) + return { isExceeded: true, scope: 'payer' } + }) + const startedAt = Date.now() + + const res = await POST(attributedCallback()) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: false }) + expect(Date.now() - startedAt).toBeLessThan(1400) }) it('reloads a cached current period once it has ended', async () => { diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index 526e8fab89e..f36d8003728 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -22,7 +22,7 @@ import { resolveLegacyV0BillingAttribution, toBillingContext, } from '@/lib/billing/core/billing-attribution' -import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' +import { type MidRunUsageVerdict, readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' import { type CumulativeUsageContextField, CumulativeUsageContextMismatchError, @@ -35,6 +35,7 @@ import { } from '@/lib/billing/threshold-billing' import { resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' +import { withinDeadline } from '@/lib/core/utils/deadline' import { generateRequestId } from '@/lib/core/utils/request' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' import { BILLING_CALLBACK_OUTCOME } from '@/lib/mothership/generated/billing-protocol-v1' @@ -45,6 +46,14 @@ import { checkInternalApiKey } from '@/lib/mothership/request/http' import { withIncomingGoSpan } from '@/lib/mothership/request/otel' const logger = createLogger('BillingUpdateCostAPI') +/** + * How long a cost callback waits on the payer's standing. The worker gives up on the whole + * callback after 5 s, and a cold gate read can wait on the ledger far longer; past this the + * callback answers not-exceeded. The abandoned read keeps running and caches its admission, and + * the next step or re-check reads a refusal again. + */ +const USAGE_STANDING_TIMEOUT_MS = 1000 + const RETRYABLE_SETTLEMENT_RESPONSE = { code: 'BILLING_SETTLEMENT_RETRYABLE', error: 'Billing settlement temporarily unavailable', @@ -70,14 +79,24 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp * Served from the execution usage gate: an admission is cached per payer and actor for the gate * TTL and a refusal is always re-read, so steady-state steps cost no ledger read. The charge is * already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the - * refusal to the next step or re-check rather than ending a paying run on a database blip. + * refusal to the next step or re-check rather than ending a paying run on a database blip, + * and so does a read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}. */ async function readUsageStanding( userId: string, billingAttribution: BillingAttributionSnapshot | undefined ): Promise { if (!isHosted || !billingAttribution) return { usageExceeded: false } - const verdict = await readMidRunUsageVerdict(billingAttribution) + let verdict: MidRunUsageVerdict + try { + verdict = await withinDeadline( + () => readMidRunUsageVerdict(billingAttribution), + Date.now() + USAGE_STANDING_TIMEOUT_MS + ) + } catch { + logger.warn('Usage standing read outlasted the callback budget; answering not exceeded') + return { usageExceeded: false } + } // Only a spent limit pauses the run. A blocked account is refused at the run's next // continuation or re-check, with blocked-account copy rather than the upgrade card. if (verdict.status !== 'exceeded') return { usageExceeded: false } diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 859768fad1f..76c21d1f807 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -141,7 +141,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) vi.mock('@/lib/workspaces/utils', () => workspacesUtilsMock) import { validateCopilotApiKeyBodySchema } from '@/lib/api/contracts/copilot' -import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage' +import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage' import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { POST } from '@/app/api/copilot/api-keys/validate/route' @@ -524,7 +524,7 @@ describe('validation lifecycle purposes', () => { periodEnd: new Date(ATTRIBUTION.billingPeriod.end), }) resetUsageGateCache() - resetMidRunPeriodCache() + resetMidRunUsageCaches() }) it('defaults older callers to full admission and rejects unknown purposes', () => { @@ -796,6 +796,28 @@ describe('validation lifecycle purposes', () => { expect(response.status).toBe(402) }) + it('answers repeated direct-v1 continuations from the cached admission and re-reads a refusal', async () => { + for (let call = 0; call < 2; call++) queueTableRows(schemaMock.user, [{ id: 'user-1' }]) + for (let leg = 0; leg < 3; leg++) { + expect((await POST(request(body, directHeaders))).status).toBe(200) + } + expect(mockCheckUsageStatus).toHaveBeenCalledTimes(1) + + resetMidRunUsageCaches() + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) + expect((await POST(request(body, directHeaders))).status).toBe(402) + expect((await POST(request(body, directHeaders))).status).toBe(402) + expect(mockCheckUsageStatus).toHaveBeenCalledTimes(3) + }) + + it('never reads the ledger for a direct-v1 continuation when billing is off', async () => { + setEnvFlags({ isHosted: false, isBillingEnabled: false }) + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) + + expect((await POST(request(body, directHeaders))).status).toBe(200) + expect(mockCheckUsageStatus).not.toHaveBeenCalled() + }) + it('judges a direct-v1 organization payer without a subscription as that organization', async () => { mockGetOrganizationSubscription.mockResolvedValue(null) mockCheckUsageStatus.mockImplementation( diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 5104c0dfcfc..7038eae5839 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -12,7 +12,13 @@ import { import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/plan' import { resolveSubscriptionUsagePeriod } from '@/lib/billing/core/reporting-period' -import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache' +import { + checkExecutionUsageLimits, + USAGE_GATE_SETTLE_TIMEOUT_MS, + USAGE_GATE_TTL_MS, +} from '@/lib/billing/core/usage-gate-cache' +import { coalesceLocally } from '@/lib/concurrency/singleflight' +import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' const logger = createLogger('MidRunUsage') @@ -121,6 +127,17 @@ export async function readMidRunUsageVerdict( return { status: 'unknown' } } +/** + * Admitted direct-v1 verdicts, served for the execution gate's TTL like + * {@link checkExecutionUsageLimits} serves attributed ones: the worker re-validates a run on + * every resume leg, and each uncached read sums the payer's ledger for the period. Only a + * `within` verdict is stored, so a refusal or an unreadable ledger is always read again. + */ +const accountVerdictCache = new LRUCache({ + max: 10_000, + ttl: USAGE_GATE_TTL_MS, +}) + /** * The same verdict for a direct-v1 run billed to an account decision rather than an attributed * payer. The payer is the one saved in the decision at admission, never re-selected from the @@ -129,6 +146,7 @@ export async function readMidRunUsageVerdict( export async function readMidRunAccountUsageVerdict( decision: AccountBillingDecision ): Promise { + if (!isHosted || !isBillingEnabled) return { status: 'within' } try { const payer = decision.billingEntity const subscription = @@ -139,6 +157,20 @@ export async function readMidRunAccountUsageVerdict( ...defaultBillingPeriod(), source: 'default' as const, } + const key = [ + payer.type, + payer.id, + billingPeriod.start.toISOString(), + billingPeriod.end.toISOString(), + billingPeriod.source, + decision.userId, + subscription?.id ?? '', + subscription?.plan ?? '', + subscription?.status ?? '', + subscription?.seats ?? '', + ].join(':') + const cached = accountVerdictCache.get(key) + if (cached) return cached // An organization payer without a subscription stays organization-scoped on the free plan, // as `toUsageLimitSubscription` does for attributed runs, never the actor's personal ledger. const usageSubscription = @@ -153,12 +185,20 @@ export async function readMidRunAccountUsageVerdict( periodEnd: billingPeriod.end, } : null) - const usage = await checkUsageStatus(decision.userId, usageSubscription, { - billingEntity: payer, - billingPeriod, - }) + const usage = await coalesceLocally( + `mid-run-account-usage:${key}`, + () => + checkUsageStatus(decision.userId, usageSubscription, { + billingEntity: payer, + billingPeriod, + }), + USAGE_GATE_SETTLE_TIMEOUT_MS + ) if (usage.unavailable) return { status: 'unknown' } - return usage.isExceeded ? { status: 'exceeded', scope: 'payer' } : { status: 'within' } + if (usage.isExceeded) return { status: 'exceeded', scope: 'payer' } + const within: MidRunUsageVerdict = { status: 'within' } + accountVerdictCache.set(key, within) + return within } catch (error) { logger.warn('Mid-run account usage read failed; continuing the run', { error: getErrorMessage(error), @@ -167,7 +207,8 @@ export async function readMidRunAccountUsageVerdict( } } -/** Drops every cached current period. Test seam; never called in production code. */ -export function resetMidRunPeriodCache(): void { +/** Drops every cached current period and account verdict. Test seam; never called in production code. */ +export function resetMidRunUsageCaches(): void { currentPeriodCache.clear() + accountVerdictCache.clear() } diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index 471a3d9d1a1..8c1a159d231 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -6,7 +6,7 @@ import { } from '@sim/testing/mocks/billing-usage-gate-cache.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { createAttributedBillingRequestEnvelope } from '@/lib/billing/core/billing-attribution' -import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage' +import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage' import { OrchestrationError } from '@/lib/core/orchestration/types' import { BillingLimitError } from '@/lib/mothership/request/go/stream' import { authorizeLifecycleContinuation, restoreBillingAdmission } from './admission' @@ -49,7 +49,7 @@ beforeEach(() => { periodStart: new Date(attribution.billingPeriod.start), periodEnd: new Date(attribution.billingPeriod.end), }) - resetMidRunPeriodCache() + resetMidRunUsageCaches() }) afterEach(resetEnvFlagsMock) @@ -157,7 +157,7 @@ describe('continuation admission', () => { authorizeLifecycleContinuation({ ...context, billingAttribution: ended }) ).rejects.toBeInstanceOf(BillingLimitError) - resetMidRunPeriodCache() + resetMidRunUsageCaches() mockGetOrganizationSubscription.mockRejectedValue(new Error('subscription read failed')) await expect( authorizeLifecycleContinuation({ ...context, billingAttribution: ended }) From d4bd0724f81848acd515d7e8adef1b573ecfd554 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 11:25:59 -0700 Subject: [PATCH 08/13] 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. --- .../billing/update-cost/route.integration.ts | 121 ++++++++++++++++++ .../app/api/billing/update-cost/route.test.ts | 76 +++++++++++ apps/sim/app/api/billing/update-cost/route.ts | 34 +++-- .../copilot/api-keys/validate/route.test.ts | 22 +++- .../api/copilot/api-keys/validate/route.ts | 1 + .../special-tags/special-tags.test.ts | 6 +- .../billing/core/billing-attribution.test.ts | 46 +++++++ .../lib/billing/core/billing-attribution.ts | 25 ++++ apps/sim/lib/billing/core/mid-run-usage.ts | 11 +- .../lib/billing/core/usage-log.integration.ts | 24 ++++ .../authorize-chat-callback.test.ts | 44 +------ .../application/authorize-chat-callback.ts | 15 +-- .../request/lifecycle/admission.test.ts | 2 +- .../mothership/request/lifecycle/run.test.ts | 10 +- .../src/mocks/billing-attribution.mock.ts | 2 + 15 files changed, 365 insertions(+), 74 deletions(-) create mode 100644 apps/sim/app/api/billing/update-cost/route.integration.ts diff --git a/apps/sim/app/api/billing/update-cost/route.integration.ts b/apps/sim/app/api/billing/update-cost/route.integration.ts new file mode 100644 index 00000000000..082009c92dc --- /dev/null +++ b/apps/sim/app/api/billing/update-cost/route.integration.ts @@ -0,0 +1,121 @@ +/** + * Cost callbacks against real PostgreSQL: a direct-v1 run that outlives its admitted Stripe period + * records its later spend in the payer's current period, so the closed period is never topped up + * after its invoice. Only the internal-key check is stubbed. + */ +import { db } from '@sim/db' +import { subscription, usageLog, user, userStats } from '@sim/db/schema' +import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock' +import { generateId } from '@sim/utils/id' +import { eq } from 'drizzle-orm' +import { NextRequest } from 'next/server' +import { afterAll, describe, expect, it, vi } from 'vitest' + +vi.mock('@/lib/core/config/env-flags', () => ({ + ...envFlagsMock, + isHosted: true, + isBillingEnabled: true, +})) +vi.mock('@/lib/mothership/request/http', async (importOriginal) => ({ + ...(await importOriginal()), + checkInternalApiKey: () => ({ success: true }), +})) + +import { + BILLING_ACCOUNT_DECISION_HEADER, + serializeAccountBillingDecisionHeader, +} from '@/lib/billing/core/billing-attribution' +import { POST } from '@/app/api/billing/update-cost/route' + +const DAY_MS = 24 * 60 * 60 * 1000 +const userId = `update-cost-user-${generateId()}` +const subscriptionId = generateId() + +afterAll(async () => { + await db.delete(usageLog).where(eq(usageLog.userId, userId)) + await db.delete(subscription).where(eq(subscription.id, subscriptionId)) + await db.delete(userStats).where(eq(userStats.userId, userId)) + await db.delete(user).where(eq(user.id, userId)) +}) + +function callback(requestKey: string, cost: number, decision: string): NextRequest { + return new NextRequest('http://localhost:3000/api/billing/update-cost', { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-api-key': 'internal', + 'x-sim-billing-protocol': 'direct-v1', + 'x-sim-billing-request-id': requestKey, + [BILLING_ACCOUNT_DECISION_HEADER]: decision, + }, + body: JSON.stringify({ + userId, + cost, + model: 'test-model', + source: 'copilot', + idempotencyKey: requestKey, + }), + }) +} + +describe('direct-v1 cost callbacks in PostgreSQL', () => { + it('records spend after a Stripe rollover in the payer current period', async () => { + const now = Date.now() + const admitted = { start: new Date(now - 10 * DAY_MS), end: new Date(now + 20 * DAY_MS) } + const rolled = { start: new Date(now - 60 * 60 * 1000), end: new Date(now + 30 * DAY_MS) } + await db.insert(user).values({ + id: userId, + name: 'Update Cost Test', + email: `${userId}@update-cost.test`, + emailVerified: true, + createdAt: new Date(now), + updatedAt: new Date(now), + }) + await db.insert(userStats).values({ id: generateId(), userId }) + await db.insert(subscription).values({ + id: subscriptionId, + plan: 'pro', + referenceId: userId, + status: 'active', + periodStart: admitted.start, + periodEnd: admitted.end, + }) + const decision = serializeAccountBillingDecisionHeader({ + userId, + billingEntity: { type: 'user', id: userId }, + billingPeriod: { + start: admitted.start.toISOString(), + end: admitted.end.toISOString(), + source: 'stripe', + }, + payerSubscriptionId: subscriptionId, + }) + const requestKey = generateId() + + expect((await POST(callback(requestKey, 0.5, decision), {})).status).toBe(200) + await db + .update(subscription) + .set({ periodStart: rolled.start, periodEnd: rolled.end }) + .where(eq(subscription.id, subscriptionId)) + expect((await POST(callback(requestKey, 0.8, decision), {})).status).toBe(200) + + const rows = await db + .select({ + eventKey: usageLog.eventKey, + cost: usageLog.cost, + billingPeriodStart: usageLog.billingPeriodStart, + }) + .from(usageLog) + .where(eq(usageLog.userId, userId)) + const byKey = new Map(rows.map((row) => [row.eventKey, row])) + expect(rows).toHaveLength(2) + expect(Number(byKey.get(`update-cost:${requestKey}`)?.cost)).toBeCloseTo(0.5) + expect(byKey.get(`update-cost:${requestKey}`)?.billingPeriodStart?.getTime()).toBe( + admitted.start.getTime() + ) + expect(Number(byKey.get(`update-cost:${requestKey}@1`)?.cost)).toBeCloseTo(0.3) + expect(byKey.get(`update-cost:${requestKey}@1`)?.billingPeriodStart?.getTime()).toBe( + rolled.start.getTime() + ) + }) +}) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index b25eb226b63..923059a252f 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -4,10 +4,16 @@ import { billingAttributionMock, billingAttributionMockFns, } from '@sim/testing/mocks/billing-attribution.mock' +import { billingCoreMock, billingCoreMockFns } from '@sim/testing/mocks/billing-core.mock' +import { billingPlanMock, billingPlanMockFns } from '@sim/testing/mocks/billing-plan.mock' import { billingUsageLogMock, billingUsageLogMockFns, } from '@sim/testing/mocks/billing-usage-log.mock' +import { + billingUsageMonitorMock, + billingUsageMonitorMockFns, +} from '@sim/testing/mocks/billing-usage-monitor.mock' import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing/mocks/copilot-http.mock' import { mothershipOtelMock } from '@sim/testing/mocks/mothership-otel.mock' import { sleep } from '@sim/utils/helpers' @@ -42,6 +48,12 @@ vi.mock('@/lib/billing/core/usage-log', () => billingUsageLogMock) vi.mock('@/lib/billing/core/billing-attribution', () => billingAttributionMock) +vi.mock('@/lib/billing/core/billing', () => billingCoreMock) + +vi.mock('@/lib/billing/core/plan', () => billingPlanMock) + +vi.mock('@/lib/billing/calculations/usage-monitor', () => billingUsageMonitorMock) + vi.mock('@/lib/billing/threshold-billing', () => ({ checkAndBillOverageThreshold: mockCheckAndBillOverageThreshold, checkAndBillPayerOverageThreshold: mockCheckAndBillPayerOverageThreshold, @@ -977,6 +989,62 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { }) }) + describe('a direct-v1 run', () => { + function directCallback() { + callbackSequence += 1 + const billingRequestId = `0190c03f-9f7d-4b79-8b58-${String(callbackSequence).padStart(12, '0')}` + return createMockRequest( + 'POST', + { + userId: 'user-1', + cost: 0.5 * callbackSequence, + model: 'claude-opus-4.8', + source: 'workspace-chat', + idempotencyKey: billingRequestId, + }, + { + 'x-api-key': 'internal', + 'x-sim-billing-protocol': 'direct-v1', + 'x-sim-billing-request-id': billingRequestId, + 'x-sim-billing-account-decision': 'serialized-account-decision', + } + ) + } + + beforeEach(() => { + mockRequireAccountBillingDecisionHeader.mockReturnValue(ACCOUNT_BILLING_DECISION) + billingCoreMockFns.mockGetOrganizationSubscription.mockResolvedValue(null) + billingPlanMockFns.mockGetHighestPrioritySubscription.mockResolvedValue(null) + billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({ blocked: false }) + billingUsageMonitorMockFns.mockCheckUsageStatus.mockResolvedValue({ + isExceeded: true, + currentUsage: 12, + limit: 10, + }) + }) + + it('tells the worker when its admitted payer has crossed its usage limit', async () => { + const body = await (await POST(directCallback())).json() + + expect(body).toMatchObject({ + usageExceeded: true, + usageUpgrade: { reason: 'usage_limit' }, + }) + }) + + it('never pauses a blocked payer with the usage card', async () => { + billingAttributionMockFns.mockCheckAccountBillingBlocks.mockResolvedValue({ + blocked: true, + scope: 'payer', + }) + + const body = await (await POST(directCallback())).json() + + expect(body.usageExceeded).toBe(false) + expect(billingUsageMonitorMockFns.mockCheckUsageStatus).not.toHaveBeenCalled() + }) + }) + describe('a run that outlives its billing period', () => { const PAYER_SUBSCRIPTION = { id: 'sub-1', @@ -995,6 +1063,14 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { ...CURRENT_ATTRIBUTION, payerSubscription: PAYER_SUBSCRIPTION, }) + mockToBillingContext.mockReturnValue({ + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { + start: new Date('2026-07-01T00:00:00.000Z'), + end: new Date('2026-08-01T00:00:00.000Z'), + source: 'stripe', + }, + }) }) it('settles the threshold against the period the charge was stamped into', async () => { diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index f36d8003728..8f4be83be39 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -22,7 +22,11 @@ import { resolveLegacyV0BillingAttribution, toBillingContext, } from '@/lib/billing/core/billing-attribution' -import { type MidRunUsageVerdict, readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' +import { + type MidRunUsageVerdict, + readMidRunAccountUsageVerdict, + readMidRunUsageVerdict, +} from '@/lib/billing/core/mid-run-usage' import { type CumulativeUsageContextField, CumulativeUsageContextMismatchError, @@ -75,24 +79,28 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp /** * Reads the run payer's standing after a cost callback, so a long run stops at its next step * once it crosses the limit instead of at its next admission, with the card the worker writes - * to its log. A duplicate callback answers too: it is often a retry whose first answer was lost. - * Served from the execution usage gate: an admission is cached per payer and actor for the gate - * TTL and a refusal is always re-read, so steady-state steps cost no ledger read. The charge is + * to its log. The payer is the attributed run's, or the one a direct-v1 run was admitted with. + * A duplicate callback answers too: it is often a retry whose first answer was lost. An + * admission is cached per payer and actor for the gate TTL and a refusal is always re-read, so + * steady-state steps cost no ledger read. The charge is * already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the * refusal to the next step or re-check rather than ending a paying run on a database blip, * and so does a read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}. */ async function readUsageStanding( userId: string, - billingAttribution: BillingAttributionSnapshot | undefined + billingAttribution: BillingAttributionSnapshot | undefined, + accountDecision: AccountBillingDecision | undefined ): Promise { - if (!isHosted || !billingAttribution) return { usageExceeded: false } + const readVerdict = billingAttribution + ? () => readMidRunUsageVerdict(billingAttribution) + : accountDecision + ? () => readMidRunAccountUsageVerdict(accountDecision) + : null + if (!isHosted || !readVerdict) return { usageExceeded: false } let verdict: MidRunUsageVerdict try { - verdict = await withinDeadline( - () => readMidRunUsageVerdict(billingAttribution), - Date.now() + USAGE_STANDING_TIMEOUT_MS - ) + verdict = await withinDeadline(readVerdict, Date.now() + USAGE_STANDING_TIMEOUT_MS) } catch { logger.warn('Usage standing read outlasted the callback budget; answering not exceeded') return { usageExceeded: false } @@ -370,8 +378,8 @@ async function updateCostInner(req: NextRequest, span: Span): Promise { expect(mockResolveBillingAttribution).not.toHaveBeenCalled() expect(mockGetUserEntityPermissions).not.toHaveBeenCalled() expect(mockGetWorkspaceBillingSettings).not.toHaveBeenCalled() - expect(mockSerializeAccountBillingDecisionHeader).toHaveBeenCalledWith(ACCOUNT_BILLING_DECISION) + expect(mockSerializeAccountBillingDecisionHeader).toHaveBeenCalledWith({ + ...ACCOUNT_BILLING_DECISION, + payerSubscriptionId: ACCOUNT_SUBSCRIPTION.id, + }) expect(res.headers.get('x-sim-billing-account-decision')).toBe('serialized-account-decision') }) @@ -514,6 +522,8 @@ describe('validation lifecycle purposes', () => { mockCheckContinuationBilling.mockReset().mockResolvedValue({ blocked: false }) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 1, limit: 10 }) + mockCheckBillingBlocked.mockResolvedValue({ blocked: false }) + mockCheckBillingEntityBlocked.mockResolvedValue({ blocked: false }) mockIsEnterprisePlan.mockResolvedValue(false) mockGetOrganizationSubscription.mockResolvedValue({ id: 'sub-org-1', @@ -810,6 +820,14 @@ describe('validation lifecycle purposes', () => { expect(mockCheckUsageStatus).toHaveBeenCalledTimes(3) }) + it('never reads the usage gate for an attributed continuation when billing is off', async () => { + setEnvFlags({ isHosted: false, isBillingEnabled: false }) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + expect((await POST(request(body, attributedHeaders))).status).toBe(200) + expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + }) + it('never reads the ledger for a direct-v1 continuation when billing is off', async () => { setEnvFlags({ isHosted: false, isBillingEnabled: false }) mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 094a2c31df8..2166dd98972 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -319,6 +319,7 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise ? { source: billingContext.billingPeriod.source } : {}), }, + ...(subscription ? { payerSubscriptionId: subscription.id } : {}), }, } } diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts index ed3296e6f14..00894654cba 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.test.ts @@ -888,7 +888,9 @@ describe('usage card written to a worker log', () => { const { segments } = parseSpecialTags(replayed, false) - expect(segments).toContainEqual({ type: 'usage_upgrade', data: usageUpgrade }) - expect(renderedText(segments)).toBe('Finished the first report.') + expect(segments).toEqual([ + { type: 'text', content: 'Finished the first report.' }, + { type: 'usage_upgrade', data: usageUpgrade }, + ]) }) }) diff --git a/apps/sim/lib/billing/core/billing-attribution.test.ts b/apps/sim/lib/billing/core/billing-attribution.test.ts index 3bf33ee85be..99eef16619a 100644 --- a/apps/sim/lib/billing/core/billing-attribution.test.ts +++ b/apps/sim/lib/billing/core/billing-attribution.test.ts @@ -17,6 +17,7 @@ import { assertBillingAttributionOwner, assertBillingAttributionSnapshot, billingAttributionsEqual, + checkAccountBillingBlocks, checkAttributedBillingBlocks, checkAttributedUsageLimits, requireBillingAttributionHeader, @@ -333,6 +334,51 @@ describe('serialized attribution boundaries', () => { }) }) +describe('checkAccountBillingBlocks', () => { + const decision = { + userId: 'actor', + billingEntity: { type: 'organization' as const, id: 'original-payer' }, + billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + } + + beforeEach(() => { + mockCheckBillingBlocked.mockReset().mockResolvedValue({ blocked: false }) + mockCheckBillingEntityBlocked.mockReset().mockResolvedValue({ blocked: false }) + }) + + it('checks both the actor and the exact original payer', async () => { + await expect(checkAccountBillingBlocks(decision)).resolves.toMatchObject({ blocked: false }) + expect(mockCheckBillingBlocked).toHaveBeenCalledWith('actor') + expect(mockCheckBillingEntityBlocked).toHaveBeenCalledWith(decision.billingEntity) + }) + + it('refuses an actor block before reading the payer', async () => { + mockCheckBillingBlocked.mockResolvedValueOnce({ blocked: true }) + await expect(checkAccountBillingBlocks(decision)).resolves.toMatchObject({ + blocked: true, + scope: 'actor', + }) + expect(mockCheckBillingEntityBlocked).not.toHaveBeenCalled() + }) + + it('refuses a payer block independently of actor standing', async () => { + mockCheckBillingEntityBlocked.mockResolvedValueOnce({ blocked: true }) + await expect(checkAccountBillingBlocks(decision)).resolves.toMatchObject({ + blocked: true, + scope: 'payer', + }) + }) + + it('reads the same personal actor and payer only once', async () => { + await checkAccountBillingBlocks({ + ...decision, + billingEntity: { type: 'user', id: 'actor' }, + }) + expect(mockCheckBillingBlocked).toHaveBeenCalledTimes(1) + expect(mockCheckBillingEntityBlocked).not.toHaveBeenCalled() + }) +}) + describe('checkAttributedUsageLimits', () => { beforeEach(() => { resetDbChainMock() diff --git a/apps/sim/lib/billing/core/billing-attribution.ts b/apps/sim/lib/billing/core/billing-attribution.ts index 938357f5654..9e06220bf28 100644 --- a/apps/sim/lib/billing/core/billing-attribution.ts +++ b/apps/sim/lib/billing/core/billing-attribution.ts @@ -97,6 +97,12 @@ export interface AccountBillingDecision { readonly end: string readonly source?: UsagePeriodSource } + /** + * The payer's subscription at admission, so a run that outlives a Stripe period bills its + * later spend to the period it was spent in, as an attributed run's `payerSubscription` does. + * Absent for a payer without a subscription, and in decisions minted before it existed. + */ + readonly payerSubscriptionId?: string } export interface ResolveBillingAttributionParams { @@ -546,6 +552,10 @@ function assertAccountBillingDecision(value: unknown): AccountBillingDecision { ) { throw new Error('Account billing decision must contain a valid billing period source') } + const payerSubscriptionId = value.payerSubscriptionId + if (payerSubscriptionId !== undefined && !isNonEmptyString(payerSubscriptionId)) { + throw new Error('Account billing decision must contain a valid payer subscription ID') + } return Object.freeze({ userId: value.userId, @@ -558,6 +568,7 @@ function assertAccountBillingDecision(value: unknown): AccountBillingDecision { end: end.toISOString(), ...(source !== undefined ? { source } : {}), }), + ...(payerSubscriptionId !== undefined ? { payerSubscriptionId } : {}), }) } @@ -937,6 +948,20 @@ export async function checkAttributedBillingBlocks( return { blocked: false } } +/** + * The same freeze checks for a direct-v1 run: the actor's own account, then the payer saved in + * its admission decision, never one re-selected from the actor's current memberships. + */ +export async function checkAccountBillingBlocks( + decision: AccountBillingDecision +): Promise { + const actorBlock = await checkBillingBlocked(decision.userId) + if (actorBlock.blocked) return { ...actorBlock, scope: 'actor' } + const payer = decision.billingEntity + if (payer.type === 'user' && payer.id === decision.userId) return actorBlock + return { ...(await checkBillingEntityBlocked(payer)), scope: 'payer' } +} + /** * Applies hosted billing gates in canonical order: actor account, workspace * payer pool, then `(organizationId, actorUserId)` member cap. diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 7038eae5839..1c892a10853 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -7,6 +7,7 @@ import { type AccountBillingDecision, type AttributedUsageLimitsResult, type BillingAttributionSnapshot, + checkAccountBillingBlocks, refreshAttributionPeriod, } from '@/lib/billing/core/billing-attribution' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' @@ -110,6 +111,7 @@ async function readGateVerdict( export async function readMidRunUsageVerdict( attribution: BillingAttributionSnapshot ): Promise { + if (!isHosted || !isBillingEnabled) return { status: 'within' } for (let attempt = 0; attempt < 2; attempt++) { let judged: BillingAttributionSnapshot try { @@ -140,8 +142,9 @@ const accountVerdictCache = new LRUCache({ /** * The same verdict for a direct-v1 run billed to an account decision rather than an attributed - * payer. The payer is the one saved in the decision at admission, never re-selected from the - * actor's current memberships, judged against that payer's current subscription period. + * payer, in the gate's order: a blocked actor or payer first, then the payer's spend. The payer + * is the one saved in the decision at admission, never re-selected from the actor's current + * memberships, judged against that payer's current subscription period. */ export async function readMidRunAccountUsageVerdict( decision: AccountBillingDecision @@ -171,6 +174,10 @@ export async function readMidRunAccountUsageVerdict( ].join(':') const cached = accountVerdictCache.get(key) if (cached) return cached + const block = await checkAccountBillingBlocks(decision) + if (block.blocked) { + return { status: 'blocked', ...(block.message ? { message: block.message } : {}) } + } // An organization payer without a subscription stays organization-scoped on the free plan, // as `toUsageLimitSubscription` does for attributed runs, never the actor's personal ledger. const usageSubscription = diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index 83862ade045..1c42957af35 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -385,6 +385,30 @@ describe('Cumulative billing with PostgreSQL', () => { expect(invoiced).toBeCloseTo(1.5, 9) }) + it('gives each period row only the tokens spent after the rows before it', async () => { + await setSubscriptionPeriod(0) + await recordCumulativeUsage({ + ...usage(0.4), + billingPeriod: { start: periods[0], end: periods[1] }, + payerSubscriptionId: 'sub-1', + }) + await setSubscriptionPeriod(1) + await recordCumulativeUsage({ + ...usage(1), + billingPeriod: { start: periods[0], end: periods[1] }, + payerSubscriptionId: 'sub-1', + metadata: { inputTokens: 25, outputTokens: 12 }, + }) + + const rows = await connection<{ event_key: string; metadata: Record }[]>` + select event_key, metadata from usage_log order by event_key + ` + expect(rows.map((row) => [row.event_key, row.metadata])).toEqual([ + ['update-cost:shared-request', { inputTokens: 10, outputTokens: 5 }], + ['update-cost:shared-request@1', { inputTokens: 15, outputTokens: 7 }], + ]) + }) + it('never stamps a charge into a period earlier than its latest row', async () => { await setSubscriptionPeriod(0) await charge(0.4) diff --git a/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts b/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts index 798b96f37e3..83fd324f972 100644 --- a/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts +++ b/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts @@ -2,10 +2,6 @@ import { billingAttributionMock, billingAttributionMockFns, } from '@sim/testing/mocks/billing-attribution.mock' -import { - billingUsageMonitorMock, - billingUsageMonitorMockFns, -} from '@sim/testing/mocks/billing-usage-monitor.mock' import { mothershipOrganizationChatsMock, mothershipOrganizationChatsMockFns, @@ -37,16 +33,14 @@ vi.mock('@/lib/permission-groups/capability-assertions', async (importOriginal) })) vi.mock('@/lib/mothership/chat/organization-chats', () => mothershipOrganizationChatsMock) vi.mock('@/lib/billing/core/billing-attribution', () => billingAttributionMock) +const mockCheckAccountBillingBlocks = billingAttributionMockFns.mockCheckAccountBillingBlocks const mockCheckAttributedBillingBlocks = billingAttributionMockFns.mockCheckAttributedBillingBlocks const mocks = { ...hoisted, organization: mothershipOrganizationChatsMockFns.mockAuthorizeOrganizationChatDelegation, - actorBlock: billingUsageMonitorMockFns.mockCheckBillingBlocked, - payerBlock: billingUsageMonitorMockFns.mockCheckBillingEntityBlocked, loadWorkspace: workspaceContextMockFns.mockResolveActiveWorkspaceApplicationContext, permission: workspaceAuthzMockFns.mockResolveEffectiveWorkspacePermission, } -vi.mock('@/lib/billing/calculations/usage-monitor', () => billingUsageMonitorMock) const context = { userId: 'actor', @@ -80,9 +74,8 @@ beforeEach(() => { }) mocks.permission.mockResolvedValue('read') mocks.organization.mockResolvedValue(undefined) - mocks.actorBlock.mockResolvedValue({ blocked: false }) - mocks.payerBlock.mockResolvedValue({ blocked: false }) mockCheckAttributedBillingBlocks.mockResolvedValue({ blocked: false }) + mockCheckAccountBillingBlocks.mockResolvedValue({ blocked: false }) }) describe('fresh chat callback authorization', () => { @@ -196,37 +189,12 @@ describe('continuation account standing', () => { it('uses the existing attributed block policy with the original snapshot', async () => { await checkCopilotContinuationBilling({ kind: 'attributed', attribution }) expect(mockCheckAttributedBillingBlocks).toHaveBeenCalledWith(attribution) - expect(mocks.actorBlock).not.toHaveBeenCalled() - expect(mocks.payerBlock).not.toHaveBeenCalled() + expect(mockCheckAccountBillingBlocks).not.toHaveBeenCalled() }) - it('checks both actor and the exact original direct-account payer', async () => { + it('uses the account block policy with the original direct-account decision', async () => { await checkCopilotContinuationBilling({ kind: 'account', decision: account }) - expect(mocks.actorBlock).toHaveBeenCalledWith('actor') - expect(mocks.payerBlock).toHaveBeenCalledWith({ type: 'organization', id: 'original-payer' }) - }) - - it('refuses an actor block before reading the payer', async () => { - mocks.actorBlock.mockResolvedValueOnce({ blocked: true }) - await expect( - checkCopilotContinuationBilling({ kind: 'account', decision: account }) - ).resolves.toMatchObject({ blocked: true, scope: 'actor' }) - expect(mocks.payerBlock).not.toHaveBeenCalled() - }) - - it('refuses a payer block independently of actor standing', async () => { - mocks.payerBlock.mockResolvedValueOnce({ blocked: true }) - await expect( - checkCopilotContinuationBilling({ kind: 'account', decision: account }) - ).resolves.toMatchObject({ blocked: true, scope: 'payer' }) - }) - - it('reads the same personal actor/payer only once', async () => { - await checkCopilotContinuationBilling({ - kind: 'account', - decision: { ...account, billingEntity: { type: 'user', id: 'actor' } }, - }) - expect(mocks.actorBlock).toHaveBeenCalledTimes(1) - expect(mocks.payerBlock).not.toHaveBeenCalled() + expect(mockCheckAccountBillingBlocks).toHaveBeenCalledWith(account) + expect(mockCheckAttributedBillingBlocks).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/mothership/application/authorize-chat-callback.ts b/apps/sim/lib/mothership/application/authorize-chat-callback.ts index c8401eeb01f..327a1cabca2 100644 --- a/apps/sim/lib/mothership/application/authorize-chat-callback.ts +++ b/apps/sim/lib/mothership/application/authorize-chat-callback.ts @@ -1,11 +1,8 @@ import type { DelegatedPrincipal } from '@sim/auth/principal' -import { - checkBillingBlocked, - checkBillingEntityBlocked, -} from '@/lib/billing/calculations/usage-monitor' import { type AccountBillingDecision, type BillingAttributionSnapshot, + checkAccountBillingBlocks, checkAttributedBillingBlocks, } from '@/lib/billing/core/billing-attribution' import { defineAuthorizedWorkspaceUseCase } from '@/lib/core/application/authorized-workspace-use-case' @@ -98,11 +95,7 @@ export type CopilotContinuationBilling = /** Checks account standing against the original admission; never reads spend or selects a new payer. */ export async function checkCopilotContinuationBilling(billing: CopilotContinuationBilling) { - if (billing.kind === 'attributed') return checkAttributedBillingBlocks(billing.attribution) - - const actor = await checkBillingBlocked(billing.decision.userId) - if (actor.blocked) return { ...actor, scope: 'actor' } - const payer = billing.decision.billingEntity - if (payer.type === 'user' && payer.id === billing.decision.userId) return actor - return { ...(await checkBillingEntityBlocked(payer)), scope: 'payer' } + return billing.kind === 'attributed' + ? checkAttributedBillingBlocks(billing.attribution) + : checkAccountBillingBlocks(billing.decision) } diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index 8c1a159d231..84097e38bfe 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -37,7 +37,7 @@ const context = { billingAttribution: attribution, } beforeEach(() => { - setEnvFlags({ isHosted: true }) + setEnvFlags({ isHosted: true, isBillingEnabled: true }) mocks.standing.mockResolvedValue({ blocked: false }) mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: false }) mockGetOrganizationSubscription.mockResolvedValue({ diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 0bbcb68ce37..d9872ecabf2 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -2265,7 +2265,7 @@ describe('runCopilotLifecycle', () => { }, payerSubscription: null, } - setEnvFlags({ isHosted: true }) + setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockEnv.COPILOT_API_KEY = 'sim-agent-key' mockRunStreamLoop.mockImplementationOnce( async ( @@ -2349,7 +2349,7 @@ describe('runCopilotLifecycle', () => { billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, payerSubscription: null, } - setEnvFlags({ isHosted: true }) + setEnvFlags({ isHosted: true, isBillingEnabled: true }) resetUsageGateCache() const billingRequestId = generateId() const onBillingAdmission = vi.fn() @@ -2405,7 +2405,7 @@ describe('runCopilotLifecycle', () => { }, payerSubscription: null, } - setEnvFlags({ isHosted: true }) + setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, message: 'limit reached', @@ -2458,7 +2458,7 @@ describe('runCopilotLifecycle', () => { billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, payerSubscription: null, } - setEnvFlags({ isHosted: true }) + setEnvFlags({ isHosted: true, isBillingEnabled: true }) resetUsageGateCache() mockCheckAttributedUsageLimits .mockResolvedValueOnce({ isExceeded: false }) @@ -2548,7 +2548,7 @@ describe('runCopilotLifecycle', () => { }) it('rejects hosted work without immutable billing attribution before egress', async () => { - setEnvFlags({ isHosted: true }) + setEnvFlags({ isHosted: true, isBillingEnabled: true }) await expect( runCopilotLifecycle( diff --git a/packages/testing/src/mocks/billing-attribution.mock.ts b/packages/testing/src/mocks/billing-attribution.mock.ts index ed26d3c271b..a54c31a63e6 100644 --- a/packages/testing/src/mocks/billing-attribution.mock.ts +++ b/packages/testing/src/mocks/billing-attribution.mock.ts @@ -160,6 +160,7 @@ export const billingAttributionMockFns = { mockResolveLegacyV0BillingAttribution: vi.fn(), mockResolveSystemBillingAttribution: vi.fn(), mockToBillingContext: vi.fn(toBillingContext), + mockCheckAccountBillingBlocks: vi.fn(), mockCheckAttributedBillingBlocks: vi.fn(), mockCheckAttributedUsageLimits: vi.fn(), mockRefreshAttributionPeriod: vi.fn(), @@ -211,6 +212,7 @@ export const billingAttributionMock = { billingAttributionMockFns.mockResolveLegacyV0BillingAttribution, resolveSystemBillingAttribution: billingAttributionMockFns.mockResolveSystemBillingAttribution, toBillingContext: billingAttributionMockFns.mockToBillingContext, + checkAccountBillingBlocks: billingAttributionMockFns.mockCheckAccountBillingBlocks, checkAttributedBillingBlocks: billingAttributionMockFns.mockCheckAttributedBillingBlocks, checkAttributedUsageLimits: billingAttributionMockFns.mockCheckAttributedUsageLimits, refreshAttributionPeriod: billingAttributionMockFns.mockRefreshAttributionPeriod, From 19c3658b98ce95967783247b140152483764dfe3 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 11:51:31 -0700 Subject: [PATCH 09/13] test(mothership): pin a usage refusal as non-retryable under the stream retry window --- apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index 558566dbe31..42544e57c6c 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -1,5 +1,6 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { + BillingLimitError, CopilotBackendError, StreamEndedWithoutTerminalError, WorkerStreamInterruptedError, @@ -142,6 +143,7 @@ describe('stream recovery budget', () => { expect(retry.nextDelay(new DOMException('Stopped', 'AbortError'))).toBeNull() expect(retry.nextDelay(new CopilotBackendError('Forbidden', { status: 403 }))).toBeNull() expect(retry.nextDelay(new Error('Invalid operation'))).toBeNull() + expect(retry.nextDelay(new BillingLimitError('user-1'))).toBeNull() expect(retry.attempt).toBe(0) }) From 074f4b8477782e643b98f474f199e6aea7483de7 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 12:09:12 -0700 Subject: [PATCH 10/13] test(billing): pin the Stripe-only rollover gate and the decision's subscription ID parsing - A payer whose period is not a Stripe period is never rolled or share-locked. - An account decision refuses a payer subscription ID that is not a non-empty string. - Documents that a subscription replaced mid-run can only under-enforce the limit. --- .../app/api/billing/update-cost/route.test.ts | 23 +++++++++++++ .../billing/core/billing-attribution.test.ts | 33 +++++++++++++++++++ .../lib/billing/core/billing-attribution.ts | 2 ++ 3 files changed, 58 insertions(+) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index 923059a252f..a3e5923d2b6 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -1112,6 +1112,29 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(res.status).toBe(200) }) + it('keeps a payer whose period is not a Stripe period on its frozen period', async () => { + mockToBillingContext.mockReturnValue({ + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { + start: new Date('2026-07-01T00:00:00.000Z'), + end: new Date('2026-08-01T00:00:00.000Z'), + source: 'default', + }, + }) + mockRecordCumulativeUsage.mockResolvedValue({ + billed: true, + delta: 0.5, + total: 0.5, + billingPeriod: CURRENT_PERIOD, + }) + + await POST(attributedCallback()) + + expect(mockRecordCumulativeUsage).toHaveBeenCalledWith( + expect.not.objectContaining({ payerSubscriptionId: expect.anything() }) + ) + }) + it('keeps reporting-window payers on their frozen period', async () => { mockToBillingContext.mockReturnValue({ billingEntity: { type: 'organization', id: 'org-1' }, diff --git a/apps/sim/lib/billing/core/billing-attribution.test.ts b/apps/sim/lib/billing/core/billing-attribution.test.ts index 99eef16619a..5d9cf95e64f 100644 --- a/apps/sim/lib/billing/core/billing-attribution.test.ts +++ b/apps/sim/lib/billing/core/billing-attribution.test.ts @@ -20,6 +20,7 @@ import { checkAccountBillingBlocks, checkAttributedBillingBlocks, checkAttributedUsageLimits, + requireAccountBillingDecisionHeader, requireBillingAttributionHeader, requireBillingCallbackAttribution, requireBillingRequestIdHeader, @@ -196,6 +197,38 @@ describe('resolveBillingAttribution', () => { }) }) +describe('account billing decision header', () => { + const decision = { + userId: 'actor', + billingEntity: { type: 'user', id: 'actor' }, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: '2026-08-01T00:00:00.000Z', + source: 'stripe', + }, + } + const header = (value: unknown) => + new Headers({ 'x-sim-billing-account-decision': encodeURIComponent(JSON.stringify(value)) }) + + it('restores the admitted payer subscription', () => { + expect( + requireAccountBillingDecisionHeader(header({ ...decision, payerSubscriptionId: 'sub-1' })) + ).toMatchObject({ payerSubscriptionId: 'sub-1' }) + expect(requireAccountBillingDecisionHeader(header(decision))).not.toHaveProperty( + 'payerSubscriptionId' + ) + }) + + it.each([42, '', ' ', null, { id: 'sub-1' }])( + 'refuses a payer subscription of %j', + (payerSubscriptionId) => { + expect(() => + requireAccountBillingDecisionHeader(header({ ...decision, payerSubscriptionId })) + ).toThrow('Account billing decision header is malformed') + } + ) +}) + describe('serialized attribution boundaries', () => { const attribution = { actorUserId: 'actor-a', diff --git a/apps/sim/lib/billing/core/billing-attribution.ts b/apps/sim/lib/billing/core/billing-attribution.ts index 9e06220bf28..9e734dc487c 100644 --- a/apps/sim/lib/billing/core/billing-attribution.ts +++ b/apps/sim/lib/billing/core/billing-attribution.ts @@ -101,6 +101,8 @@ export interface AccountBillingDecision { * The payer's subscription at admission, so a run that outlives a Stripe period bills its * later spend to the period it was spent in, as an attributed run's `payerSubscription` does. * Absent for a payer without a subscription, and in decisions minted before it existed. + * If that subscription is replaced mid-run, spend stays in its period while the mid-run + * verdict judges the payer's current one, so the limit can only be under-enforced. */ readonly payerSubscriptionId?: string } From c544dc729758a6b449aab70d4b5322479058076a Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 12:37:33 -0700 Subject: [PATCH 11/13] fix(billing): judge a non-Stripe run against the period its charges land in A reporting-window or default-period payer's charges never roll forward, so after its admitted window ends the mid-run verdict judged an empty new window and under-enforced the limit. Both the attributed and direct-v1 verdicts now judge the current period only for a Stripe payer, matching the cost callback's rollover gate, and the admitted period otherwise, even after it ends. Stripe payers keep being judged against their current period. --- .../app/api/billing/update-cost/route.test.ts | 35 +++++++++++++++++-- .../copilot/api-keys/validate/route.test.ts | 29 +++++++++++++++ apps/sim/lib/billing/core/mid-run-usage.ts | 34 ++++++++++++------ .../request/lifecycle/admission.test.ts | 6 +++- 4 files changed, 90 insertions(+), 14 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index a3e5923d2b6..7a2f061f4d7 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -880,9 +880,18 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => { describe('POST /api/billing/update-cost — mid-run usage gate', () => { let callbackSequence = 0 + /** A Stripe-period payer admitted in a period that has since ended. */ + const STRIPE_ATTRIBUTION = { + ...ATTRIBUTION, + billingPeriod: { ...ATTRIBUTION.billingPeriod, source: 'stripe' as const }, + } const CURRENT_ATTRIBUTION = { ...ATTRIBUTION, - billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: '2099-01-01T00:00:00.000Z', + source: 'stripe' as const, + }, } function attributedCallback() { @@ -1192,7 +1201,7 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { } it('judges a run past its admitted period against the payer current period', async () => { - mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockRequireBillingAttributionHeader.mockReturnValue(STRIPE_ATTRIBUTION) mockRefreshAttributionPeriod.mockResolvedValue(CURRENT_ATTRIBUTION) refuseOnlyCurrentPeriod() @@ -1217,8 +1226,28 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(body.usageExceeded).toBe(true) }) + it('judges a reporting-window run against its admitted window after that window ends', async () => { + const admitted = { + ...ATTRIBUTION, + billingPeriod: { ...ATTRIBUTION.billingPeriod, source: 'reporting' as const }, + } + mockRequireBillingAttributionHeader.mockReturnValue(admitted) + mockRefreshAttributionPeriod.mockResolvedValue({ + ...CURRENT_ATTRIBUTION, + billingPeriod: { ...CURRENT_ATTRIBUTION.billingPeriod, source: 'reporting' as const }, + }) + mockCheckAttributedUsageLimits.mockImplementation(async (attribution: typeof admitted) => ({ + isExceeded: attribution.billingPeriod.end === admitted.billingPeriod.end, + scope: 'payer', + })) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(true) + }) + it('keeps a run going when its current period cannot be read', async () => { - mockRequireBillingAttributionHeader.mockReturnValue(ATTRIBUTION) + mockRequireBillingAttributionHeader.mockReturnValue(STRIPE_ATTRIBUTION) mockRefreshAttributionPeriod.mockRejectedValue(new Error('subscription read timed out')) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index c7ab8a9d01c..54aaded7a16 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -836,6 +836,35 @@ describe('validation lifecycle purposes', () => { expect(mockCheckUsageStatus).not.toHaveBeenCalled() }) + it('judges a direct-v1 reporting-window run against its admitted window after it ends', async () => { + const admittedWindow = { + ...ACCOUNT_BILLING_DECISION, + billingPeriod: { + start: '2026-06-01T00:00:00.000Z', + end: '2026-07-01T00:00:00.000Z', + source: 'reporting' as const, + }, + } + mockCheckUsageStatus.mockImplementation( + async ( + _userId: string, + _subscription: unknown, + context?: { billingPeriod: { start: Date } } + ) => ({ + isExceeded: + context?.billingPeriod.start.toISOString() === admittedWindow.billingPeriod.start, + currentUsage: 12, + limit: 10, + }) + ) + + const response = await POST( + request(body, { ...directHeaders, 'x-sim-billing-account-decision': encode(admittedWindow) }) + ) + + expect(response.status).toBe(402) + }) + it('judges a direct-v1 organization payer without a subscription as that organization', async () => { mockGetOrganizationSubscription.mockResolvedValue(null) mockCheckUsageStatus.mockImplementation( diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 1c892a10853..109d9eda33d 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -74,6 +74,11 @@ async function currentAttribution( return current } +/** Mirrors the cost callback's rollover gate: only a Stripe period rolls forward. */ +function rollsIntoCurrentPeriod(period: { source?: string }): boolean { + return period.source === 'stripe' +} + function periodHasEnded(attribution: BillingAttributionSnapshot): boolean { return Date.now() >= new Date(attribution.billingPeriod.end).getTime() } @@ -102,16 +107,19 @@ async function readGateVerdict( } /** - * Judges a run against its admitted payer's CURRENT subscription period, never the period it was - * admitted in: charges land in whatever period the subscription is in now (a rollover or an - * early anchor reset included), so that is the allowance they count against. A read that - * straddles the end of that period is judged again against the next one. If the current period - * cannot be read the verdict is unknown, and the run continues. + * Judges a run against the period its charges land in. A Stripe-period payer's charges roll into + * whatever period the subscription is in now (a rollover or an early anchor reset included), so + * such a run is judged against the payer's CURRENT period: a read that straddles the end of that + * period is judged again against the next one, and a current period that cannot be read makes + * the verdict unknown, so the run continues. Any other payer's charges stay in the admitted + * period (a reporting window, or the open default one), so that period is judged, even after it + * ends. */ export async function readMidRunUsageVerdict( attribution: BillingAttributionSnapshot ): Promise { if (!isHosted || !isBillingEnabled) return { status: 'within' } + if (!rollsIntoCurrentPeriod(attribution.billingPeriod)) return readGateVerdict(attribution) for (let attempt = 0; attempt < 2; attempt++) { let judged: BillingAttributionSnapshot try { @@ -144,7 +152,7 @@ const accountVerdictCache = new LRUCache({ * The same verdict for a direct-v1 run billed to an account decision rather than an attributed * payer, in the gate's order: a blocked actor or payer first, then the payer's spend. The payer * is the one saved in the decision at admission, never re-selected from the actor's current - * memberships, judged against that payer's current subscription period. + * memberships, and the period judged is the one its charges land in, as for attributed runs. */ export async function readMidRunAccountUsageVerdict( decision: AccountBillingDecision @@ -156,10 +164,16 @@ export async function readMidRunAccountUsageVerdict( payer.type === 'organization' ? await getOrganizationSubscription(payer.id, { onError: 'throw' }) : await getHighestPriorityPersonalSubscription(payer.id, { onError: 'throw' }) - const billingPeriod = resolveSubscriptionUsagePeriod(subscription) ?? { - ...defaultBillingPeriod(), - source: 'default' as const, - } + const billingPeriod = rollsIntoCurrentPeriod(decision.billingPeriod) + ? (resolveSubscriptionUsagePeriod(subscription) ?? { + ...defaultBillingPeriod(), + source: 'default' as const, + }) + : { + start: new Date(decision.billingPeriod.start), + end: new Date(decision.billingPeriod.end), + source: decision.billingPeriod.source ?? ('default' as const), + } const key = [ payer.type, payer.id, diff --git a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts index 84097e38bfe..d49e468b854 100644 --- a/apps/sim/lib/mothership/request/lifecycle/admission.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/admission.test.ts @@ -137,7 +137,11 @@ describe('continuation admission', () => { it('judges a leg past its admitted period against the payer current period', async () => { const ended = { ...attribution, - billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' }, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: '2026-08-01T00:00:00.000Z', + source: 'stripe' as const, + }, } mockGetOrganizationSubscription.mockResolvedValue({ id: 'sub-org', From 4d16e30c1be4fe094d66a0eaa39476f0a49bcc86 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 12:52:47 -0700 Subject: [PATCH 12/13] test(billing): assert observable verdicts and responses instead of mock calls - The account block, continuation delegation, cache, billing-off and rollover tests assert the verdict or HTTP response. Where behaviour depends on an input, the fake answers by that input, as the real ledger and settlement do. - Pins against real PostgreSQL that a reporting run's top-ups after its window ends are counted in that window, where its request was first charged, and a later run's charges in the next. --- .../app/api/billing/update-cost/route.test.ts | 141 +++++++----------- .../copilot/api-keys/validate/route.test.ts | 33 ++-- .../billing/core/billing-attribution.test.ts | 40 ++--- .../lib/billing/core/usage-log.integration.ts | 26 ++++ .../authorize-chat-callback.test.ts | 24 +-- 5 files changed, 137 insertions(+), 127 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index 7a2f061f4d7..19e95da7420 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -970,18 +970,18 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { it('serves a cached admission to every step and re-reads a refusal', async () => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) - for (let step = 0; step < 5; step++) { + expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(false) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + for (let step = 0; step < 4; step++) { const body = await (await POST(attributedCallback())).json() expect(body.usageExceeded).toBe(false) expect(body).not.toHaveProperty('usageUpgrade') } - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(1) resetUsageGateCache() - mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) - await POST(attributedCallback()) - await POST(attributedCallback()) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(3) + expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(true) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(false) }) it('answers a duplicate retry with the verdict its lost first answer carried', async () => { @@ -1050,7 +1050,6 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { const body = await (await POST(directCallback())).json() expect(body.usageExceeded).toBe(false) - expect(billingUsageMonitorMockFns.mockCheckUsageStatus).not.toHaveBeenCalled() }) }) @@ -1061,56 +1060,70 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { status: 'active', seats: 4, } + const ADMITTED_PERIOD = { + start: new Date('2026-07-01T00:00:00.000Z'), + end: new Date('2026-08-01T00:00:00.000Z'), + } const CURRENT_PERIOD = { start: new Date('2026-08-01T00:00:00.000Z'), end: new Date('2026-09-01T00:00:00.000Z'), } + function admittedWithSource(source: 'stripe' | 'reporting' | 'default') { + mockToBillingContext.mockReturnValue({ + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { ...ADMITTED_PERIOD, source }, + }) + } + + /** Threshold settlement for a payer whose charges belong to `period` refuses any other. */ + function settlesOnlyAgainst(period: typeof ADMITTED_PERIOD) { + mockCheckAndBillPayerOverageThreshold.mockImplementation( + async (_payer: unknown, options: { expectedBillingPeriod: typeof ADMITTED_PERIOD }) => { + if (options.expectedBillingPeriod.start.getTime() !== period.start.getTime()) { + throw new Error('Settled against a period the charge did not land in') + } + } + ) + } + beforeEach(() => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockRequireBillingAttributionHeader.mockReturnValue({ ...CURRENT_ATTRIBUTION, payerSubscription: PAYER_SUBSCRIPTION, }) - mockToBillingContext.mockReturnValue({ - billingEntity: { type: 'organization', id: 'org-1' }, - billingPeriod: { - start: new Date('2026-07-01T00:00:00.000Z'), - end: new Date('2026-08-01T00:00:00.000Z'), - source: 'stripe', - }, - }) + // As the ledger behaves: a charge given the payer's subscription lands in its current + // period, any other stays in the period it was admitted in. + mockRecordCumulativeUsage.mockImplementation( + async (params: { + payerSubscriptionId?: string + billingPeriod: typeof ADMITTED_PERIOD + }) => ({ + billed: true, + delta: 0.5, + total: 1.5, + billingPeriod: params.payerSubscriptionId + ? CURRENT_PERIOD + : { start: params.billingPeriod.start, end: params.billingPeriod.end }, + }) + ) }) - it('settles the threshold against the period the charge was stamped into', async () => { - mockRecordCumulativeUsage.mockResolvedValue({ - billed: true, - delta: 0.5, - total: 1.5, - billingPeriod: CURRENT_PERIOD, - }) - - const res = await POST(attributedCallback()) + it("records a Stripe payer's charge in its current period and settles it there", async () => { + admittedWithSource('stripe') + settlesOnlyAgainst(CURRENT_PERIOD) - expect(res.status).toBe(200) - expect(mockRecordCumulativeUsage).toHaveBeenCalledWith( - expect.objectContaining({ payerSubscriptionId: 'sub-1' }) - ) - expect(mockCheckAndBillPayerOverageThreshold).toHaveBeenCalledWith( - { type: 'organization', id: 'org-1' }, - expect.objectContaining({ expectedBillingPeriod: CURRENT_PERIOD }) - ) + expect((await POST(attributedCallback())).status).toBe(200) }) it('leaves a period that closed under a recorded charge to the cycle close', async () => { + admittedWithSource('stripe') mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 1.5, - billingPeriod: { - start: new Date('2026-07-01T00:00:00.000Z'), - end: new Date('2026-08-01T00:00:00.000Z'), - }, + billingPeriod: ADMITTED_PERIOD, }) mockCheckAndBillPayerOverageThreshold.mockRejectedValue( new MockThresholdSettlementError('billing_period_elapsed') @@ -1121,51 +1134,15 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(res.status).toBe(200) }) - it('keeps a payer whose period is not a Stripe period on its frozen period', async () => { - mockToBillingContext.mockReturnValue({ - billingEntity: { type: 'organization', id: 'org-1' }, - billingPeriod: { - start: new Date('2026-07-01T00:00:00.000Z'), - end: new Date('2026-08-01T00:00:00.000Z'), - source: 'default', - }, - }) - mockRecordCumulativeUsage.mockResolvedValue({ - billed: true, - delta: 0.5, - total: 0.5, - billingPeriod: CURRENT_PERIOD, - }) - - await POST(attributedCallback()) - - expect(mockRecordCumulativeUsage).toHaveBeenCalledWith( - expect.not.objectContaining({ payerSubscriptionId: expect.anything() }) - ) - }) - - it('keeps reporting-window payers on their frozen period', async () => { - mockToBillingContext.mockReturnValue({ - billingEntity: { type: 'organization', id: 'org-1' }, - billingPeriod: { - start: new Date('2026-07-01T00:00:00.000Z'), - end: new Date('2026-08-01T00:00:00.000Z'), - source: 'reporting', - }, - }) - mockRecordCumulativeUsage.mockResolvedValue({ - billed: true, - delta: 0.5, - total: 0.5, - billingPeriod: CURRENT_PERIOD, - }) - - await POST(attributedCallback()) + it.each(['reporting', 'default'] as const)( + 'keeps a payer with a %s period on the period it was admitted in', + async (source) => { + admittedWithSource(source) + settlesOnlyAgainst(ADMITTED_PERIOD) - expect(mockRecordCumulativeUsage).toHaveBeenCalledWith( - expect.not.objectContaining({ payerSubscriptionId: expect.anything() }) - ) - }) + expect((await POST(attributedCallback())).status).toBe(200) + } + ) }) it.each([ @@ -1280,7 +1257,6 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { const body = await (await POST(attributedCallback())).json() expect(body.usageExceeded).toBe(false) - expect(mockRefreshAttributionPeriod).toHaveBeenCalledTimes(2) }) it('answers not exceeded when the standing read outlasts the callback budget', async () => { @@ -1341,8 +1317,7 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { const res = await POST(attributedCallback()) expect(res.status).toBe(200) - await expect(res.json()).resolves.toMatchObject({ usageExceeded: false }) - expect(mockRecordCumulativeUsage).toHaveBeenCalledOnce() + await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: false }) }) it('reports no exceeded usage when billing is disabled', async () => { diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 54aaded7a16..5cde12df5ea 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -400,6 +400,9 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => { }) it('admits a direct-v1 key without Redis while ignoring a local workspace ID', async () => { + mockSerializeAccountBillingDecisionHeader.mockImplementation((decision: object) => + encodeURIComponent(JSON.stringify(decision)) + ) mockGetUserEntityPermissions.mockResolvedValueOnce(null) mockGetWorkspaceBillingSettings.mockResolvedValueOnce({ billedAccountUserId: 'different-owner', @@ -426,11 +429,9 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => { expect(mockResolveBillingAttribution).not.toHaveBeenCalled() expect(mockGetUserEntityPermissions).not.toHaveBeenCalled() expect(mockGetWorkspaceBillingSettings).not.toHaveBeenCalled() - expect(mockSerializeAccountBillingDecisionHeader).toHaveBeenCalledWith({ - ...ACCOUNT_BILLING_DECISION, - payerSubscriptionId: ACCOUNT_SUBSCRIPTION.id, - }) - expect(res.headers.get('x-sim-billing-account-decision')).toBe('serialized-account-decision') + expect( + JSON.parse(decodeURIComponent(res.headers.get('x-sim-billing-account-decision') ?? '')) + ).toEqual({ ...ACCOUNT_BILLING_DECISION, payerSubscriptionId: ACCOUNT_SUBSCRIPTION.id }) }) it('fails direct-v1 admission closed when its payer cannot be resolved', async () => { @@ -740,16 +741,16 @@ describe('validation lifecycle purposes', () => { it('answers a polled re-check from the cached admission and always re-reads a refusal', async () => { for (let call = 0; call < 2; call++) queueTableRows(schemaMock.user, [{ id: 'user-1' }]) - for (let poll = 0; poll < 3; poll++) { + expect((await POST(request(body, attributedHeaders))).status).toBe(200) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + for (let poll = 0; poll < 2; poll++) { expect((await POST(request(body, attributedHeaders))).status).toBe(200) } - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(1) resetUsageGateCache() - mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) expect((await POST(request(body, attributedHeaders))).status).toBe(402) - expect((await POST(request(body, attributedHeaders))).status).toBe(402) - expect(mockCheckAttributedUsageLimits).toHaveBeenCalledTimes(3) + mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + expect((await POST(request(body, attributedHeaders))).status).toBe(200) }) it('refuses a new turn whose usage cannot be read with neutral copy', async () => { @@ -808,16 +809,16 @@ describe('validation lifecycle purposes', () => { it('answers repeated direct-v1 continuations from the cached admission and re-reads a refusal', async () => { for (let call = 0; call < 2; call++) queueTableRows(schemaMock.user, [{ id: 'user-1' }]) - for (let leg = 0; leg < 3; leg++) { + expect((await POST(request(body, directHeaders))).status).toBe(200) + mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) + for (let leg = 0; leg < 2; leg++) { expect((await POST(request(body, directHeaders))).status).toBe(200) } - expect(mockCheckUsageStatus).toHaveBeenCalledTimes(1) resetMidRunUsageCaches() - mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) expect((await POST(request(body, directHeaders))).status).toBe(402) - expect((await POST(request(body, directHeaders))).status).toBe(402) - expect(mockCheckUsageStatus).toHaveBeenCalledTimes(3) + mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 1, limit: 10 }) + expect((await POST(request(body, directHeaders))).status).toBe(200) }) it('never reads the usage gate for an attributed continuation when billing is off', async () => { @@ -825,7 +826,6 @@ describe('validation lifecycle purposes', () => { mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) expect((await POST(request(body, attributedHeaders))).status).toBe(200) - expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() }) it('never reads the ledger for a direct-v1 continuation when billing is off', async () => { @@ -833,7 +833,6 @@ describe('validation lifecycle purposes', () => { mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 }) expect((await POST(request(body, directHeaders))).status).toBe(200) - expect(mockCheckUsageStatus).not.toHaveBeenCalled() }) it('judges a direct-v1 reporting-window run against its admitted window after it ends', async () => { diff --git a/apps/sim/lib/billing/core/billing-attribution.test.ts b/apps/sim/lib/billing/core/billing-attribution.test.ts index 5d9cf95e64f..8e40d3456da 100644 --- a/apps/sim/lib/billing/core/billing-attribution.test.ts +++ b/apps/sim/lib/billing/core/billing-attribution.test.ts @@ -379,36 +379,40 @@ describe('checkAccountBillingBlocks', () => { mockCheckBillingEntityBlocked.mockReset().mockResolvedValue({ blocked: false }) }) - it('checks both the actor and the exact original payer', async () => { - await expect(checkAccountBillingBlocks(decision)).resolves.toMatchObject({ blocked: false }) - expect(mockCheckBillingBlocked).toHaveBeenCalledWith('actor') - expect(mockCheckBillingEntityBlocked).toHaveBeenCalledWith(decision.billingEntity) - }) - - it('refuses an actor block before reading the payer', async () => { - mockCheckBillingBlocked.mockResolvedValueOnce({ blocked: true }) + it('refuses the exact actor and original payer when either is blocked', async () => { + mockCheckBillingBlocked.mockImplementation(async (userId: string) => ({ + blocked: userId === 'actor', + })) await expect(checkAccountBillingBlocks(decision)).resolves.toMatchObject({ blocked: true, scope: 'actor', }) - expect(mockCheckBillingEntityBlocked).not.toHaveBeenCalled() - }) - it('refuses a payer block independently of actor standing', async () => { - mockCheckBillingEntityBlocked.mockResolvedValueOnce({ blocked: true }) + mockCheckBillingBlocked.mockResolvedValue({ blocked: false }) + mockCheckBillingEntityBlocked.mockImplementation(async (entity: { id: string }) => ({ + blocked: entity.id === 'original-payer', + })) await expect(checkAccountBillingBlocks(decision)).resolves.toMatchObject({ blocked: true, scope: 'payer', }) }) - it('reads the same personal actor and payer only once', async () => { - await checkAccountBillingBlocks({ - ...decision, - billingEntity: { type: 'user', id: 'actor' }, + it('reports an actor block ahead of a payer block', async () => { + mockCheckBillingBlocked.mockResolvedValue({ blocked: true, message: 'Actor frozen.' }) + mockCheckBillingEntityBlocked.mockResolvedValue({ blocked: true, message: 'Payer frozen.' }) + await expect(checkAccountBillingBlocks(decision)).resolves.toEqual({ + blocked: true, + message: 'Actor frozen.', + scope: 'actor', }) - expect(mockCheckBillingBlocked).toHaveBeenCalledTimes(1) - expect(mockCheckBillingEntityBlocked).not.toHaveBeenCalled() + }) + + it('answers a personal payer from the actor standing alone', async () => { + mockCheckBillingEntityBlocked.mockResolvedValue({ blocked: true }) + await expect( + checkAccountBillingBlocks({ ...decision, billingEntity: { type: 'user', id: 'actor' } }) + ).resolves.toMatchObject({ blocked: false }) }) }) diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index 1c42957af35..4c164e6ad73 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -9,6 +9,7 @@ import type { db } from '@sim/db' import * as schema from '@sim/db/schema' import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' import { getPostgresErrorCode } from '@sim/utils/errors' +import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { sql } from 'drizzle-orm' import { drizzle } from 'drizzle-orm/postgres-js' @@ -313,6 +314,31 @@ describe('Cumulative billing with PostgreSQL', () => { } ) + it("counts a reporting run's top-ups after its window ends in that window, and a later run's charges in the next", async () => { + const payer = { type: 'organization', id: 'payer' } as const + const boundary = new Date(Date.now() + 1500) + const admitted = { + start: new Date(boundary.getTime() - 30 * 24 * 60 * 60 * 1000), + end: boundary, + source: 'reporting' as const, + } + const next = { + start: boundary, + end: new Date(boundary.getTime() + 30 * 24 * 60 * 60 * 1000), + source: 'reporting' as const, + } + const windowTotal = (period: typeof admitted) => + getBillingPeriodUsageCost(payer, period, undefined, database) + + await recordCumulativeUsage({ ...usage(0.4, 'update-cost:long-run'), billingPeriod: admitted }) + await sleep(boundary.getTime() - Date.now() + 100) + await recordCumulativeUsage({ ...usage(1, 'update-cost:long-run'), billingPeriod: admitted }) + await recordCumulativeUsage({ ...usage(0.25, 'update-cost:next-run'), billingPeriod: next }) + + expect(await windowTotal(admitted)).toBeCloseTo(1, 9) + expect(await windowTotal(next)).toBeCloseTo(0.25, 9) + }) + describe('a request that outlives its billing period', () => { // Past periods: the old period's row is written under the subscription lock only once // that period has ended. diff --git a/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts b/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts index 83fd324f972..2b84d1e20ab 100644 --- a/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts +++ b/apps/sim/lib/mothership/application/authorize-chat-callback.test.ts @@ -186,15 +186,21 @@ describe('fresh chat callback authorization', () => { }) describe('continuation account standing', () => { - it('uses the existing attributed block policy with the original snapshot', async () => { - await checkCopilotContinuationBilling({ kind: 'attributed', attribution }) - expect(mockCheckAttributedBillingBlocks).toHaveBeenCalledWith(attribution) - expect(mockCheckAccountBillingBlocks).not.toHaveBeenCalled() - }) + it('judges each run kind by its own block policy and original billing material', async () => { + mockCheckAttributedBillingBlocks.mockImplementation(async (value: unknown) => ({ + blocked: value === attribution, + scope: 'payer', + })) + mockCheckAccountBillingBlocks.mockImplementation(async (value: unknown) => ({ + blocked: value === account, + scope: 'actor', + })) - it('uses the account block policy with the original direct-account decision', async () => { - await checkCopilotContinuationBilling({ kind: 'account', decision: account }) - expect(mockCheckAccountBillingBlocks).toHaveBeenCalledWith(account) - expect(mockCheckAttributedBillingBlocks).not.toHaveBeenCalled() + await expect( + checkCopilotContinuationBilling({ kind: 'attributed', attribution }) + ).resolves.toEqual({ blocked: true, scope: 'payer' }) + await expect( + checkCopilotContinuationBilling({ kind: 'account', decision: account }) + ).resolves.toEqual({ blocked: true, scope: 'actor' }) }) }) From 1291fa8b7bc61dd12df4691e3e96c9b21e1c3a86 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 13:19:55 -0700 Subject: [PATCH 13/13] refactor(billing): drop unconsumed refusal plumbing and the straddle retry, and make the reporting-window test deterministic - A new turn's 402 is empty again and the stream no longer parses 402 bodies: the worker replaces any validation 402 body with its own message and only polls continuation, so the new-turn codes, USAGE_UNAVAILABLE and the server-side refusal reasons had no consumer. - The mid-run verdict reads its current period once; a period that ends during the read is left to the next callback. - The cumulative-usage '@' check left the ledger; the cost callback already refuses such keys. - The reporting-window integration test derives its boundary from the first row's created_at and waits on the database clock. --- .../app/api/billing/update-cost/route.test.ts | 65 ++++++------------- .../copilot/api-keys/validate/route.test.ts | 38 +---------- .../api/copilot/api-keys/validate/route.ts | 45 +------------ apps/sim/lib/api/contracts/copilot.ts | 18 +---- .../calculations/usage-monitor.test.ts | 15 +---- .../lib/billing/calculations/usage-monitor.ts | 65 +++++-------------- apps/sim/lib/billing/core/mid-run-usage.ts | 37 +++++------ .../lib/billing/core/usage-log.integration.ts | 62 +++++++++--------- apps/sim/lib/billing/core/usage-log.ts | 6 +- .../lib/mothership/request/go/stream.test.ts | 48 -------------- apps/sim/lib/mothership/request/go/stream.ts | 27 -------- 11 files changed, 93 insertions(+), 333 deletions(-) diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts index 19e95da7420..60407452e78 100644 --- a/apps/sim/app/api/billing/update-cost/route.test.ts +++ b/apps/sim/app/api/billing/update-cost/route.test.ts @@ -60,7 +60,6 @@ vi.mock('@/lib/billing/threshold-billing', () => ({ ThresholdSettlementError: MockThresholdSettlementError, })) -import { billingUpdateCostResponseSchema } from '@/lib/api/contracts/subscription' import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage' import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache' import { @@ -1233,32 +1232,6 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(body.usageExceeded).toBe(false) }) - it('rejudges a read that straddles the end of the admitted period', async () => { - const straddling = { - ...CURRENT_ATTRIBUTION, - billingPeriod: { - start: '2026-07-01T00:00:00.000Z', - end: new Date(Date.now() + 40).toISOString(), - }, - } - mockRefreshAttributionPeriod - .mockResolvedValueOnce(straddling) - .mockResolvedValue(CURRENT_ATTRIBUTION) - mockCheckAttributedUsageLimits.mockImplementation( - async (attribution: typeof CURRENT_ATTRIBUTION) => { - if (attribution.billingPeriod.end !== straddling.billingPeriod.end) { - return { isExceeded: false } - } - await sleep(80) - return { isExceeded: true, scope: 'payer' } - } - ) - - const body = await (await POST(attributedCallback())).json() - - expect(body.usageExceeded).toBe(false) - }) - it('answers not exceeded when the standing read outlasts the callback budget', async () => { mockCheckAttributedUsageLimits.mockImplementation(async () => { await sleep(1500) @@ -1273,6 +1246,26 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(Date.now() - startedAt).toBeLessThan(1400) }) + it('does not pause a run on a verdict read across the end of its period', async () => { + const straddling = { + ...CURRENT_ATTRIBUTION, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: new Date(Date.now() + 40).toISOString(), + source: 'stripe' as const, + }, + } + mockRefreshAttributionPeriod.mockResolvedValue(straddling) + mockCheckAttributedUsageLimits.mockImplementation(async () => { + await sleep(80) + return { isExceeded: true, scope: 'payer' } + }) + + const body = await (await POST(attributedCallback())).json() + + expect(body.usageExceeded).toBe(false) + }) + it('reloads a cached current period once it has ended', async () => { const ending = { ...CURRENT_ATTRIBUTION, @@ -1293,24 +1286,6 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => { expect(body.usageExceeded).toBe(true) }) - it('never answers a verdict whose card and flag disagree', () => { - const base = { - success: true, - data: { processedAt: new Date(0).toISOString(), requestId: 'request-1' }, - } - const card = { reason: 'usage_limit', action: 'upgrade_plan', message: 'Limit reached.' } - expect( - billingUpdateCostResponseSchema.safeParse({ ...base, usageExceeded: true }).success - ).toBe(false) - expect( - billingUpdateCostResponseSchema.safeParse({ - ...base, - usageExceeded: false, - usageUpgrade: card, - }).success - ).toBe(false) - }) - it('keeps a recorded charge successful when the gate read fails', async () => { mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out')) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 5cde12df5ea..f9a88e61d3f 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -753,20 +753,6 @@ describe('validation lifecycle purposes', () => { expect((await POST(request(body, attributedHeaders))).status).toBe(200) }) - it('refuses a new turn whose usage cannot be read with neutral copy', async () => { - mockCheckAttributedUsageLimits.mockResolvedValueOnce({ - isExceeded: true, - reason: 'usage_unavailable', - message: 'Usage limit exceeded: $0.00 used of $0.00 limit.', - scope: 'payer', - }) - const refused = await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders)) - expect(refused.status).toBe(402) - const refusal = await refused.json() - expect(refusal.code).toBe('USAGE_UNAVAILABLE') - expect(refusal.error).not.toMatch(/\$/) - }) - it('checks the payer saved at admission for a direct-v1 run whose actor changed orgs', async () => { const endedDecision = { ...ACCOUNT_BILLING_DECISION, @@ -877,21 +863,6 @@ describe('validation lifecycle purposes', () => { expect((await POST(request(body, directHeaders))).status).toBe(402) }) - it('refuses a blocked new turn with the blocked body the contract declares', async () => { - mockCheckAttributedUsageLimits.mockResolvedValueOnce({ - isExceeded: true, - reason: 'billing_blocked', - message: 'Billing account frozen.', - scope: 'payer', - }) - const refused = await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders)) - expect(refused.status).toBe(402) - await expect(refused.json()).resolves.toEqual({ - code: 'BILLING_BLOCKED', - error: 'Billing account frozen.', - }) - }) - it('allows cancellation without billing material or spending/standing/plan checks', async () => { const response = await POST( request({ ...body, purpose: 'cancellation' }, { 'x-sim-billing-protocol': 'attribution-v1' }) @@ -998,12 +969,9 @@ describe('validation lifecycle purposes', () => { isExceeded: true, payerUsage: { currentUsage: 120, limit: 100 }, }) - const refused = await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders)) - expect(refused.status).toBe(402) - await expect(refused.json()).resolves.toMatchObject({ - code: 'USAGE_LIMIT_EXCEEDED', - usageUpgrade: { reason: 'usage_limit' }, - }) + expect((await POST(request({ ...body, purpose: 'new-turn' }, attributedHeaders))).status).toBe( + 402 + ) expect((await POST(request({ ...body, purpose: 'new-turn' }, directHeaders))).status).toBe(400) expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() }) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 2166dd98972..8723a9d5c90 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -7,15 +7,12 @@ import { type NextRequest, NextResponse } from 'next/server' import { COPILOT_BILLING_BLOCKED_CODE, COPILOT_USAGE_LIMIT_EXCEEDED_CODE, - COPILOT_USAGE_UNAVAILABLE_CODE, type ValidateCopilotApiKeyBillingBlocked, - type ValidateCopilotApiKeyRefusal, type ValidateCopilotApiKeyUsageExceeded, validateCopilotApiKeyContract, } from '@/lib/api/contracts/copilot' import { parseRequest, validationErrorResponse } from '@/lib/api/server' import { checkServerSideUsageLimits } from '@/lib/billing/calculations/usage-monitor' -import { USAGE_UNAVAILABLE_MESSAGE } from '@/lib/billing/constants' import { type AccountBillingDecision, type BillingAttributionSnapshot, @@ -68,7 +65,6 @@ import { withIncomingGoSpan } from '@/lib/mothership/request/otel' const logger = createLogger('CopilotApiKeysValidate') const CONTINUATION_BLOCKED_MESSAGE = 'Continuation billing account is blocked' -const BILLING_BLOCKED_MESSAGE = 'Billing account is blocked' function invalidBillingProtocolResponse(): NextResponse { return NextResponse.json({ error: 'Invalid billing attribution protocol' }, { status: 400 }) @@ -249,39 +245,13 @@ function resolveContinuationBilling( } } -interface AdmissionUsage { +async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise<{ isExceeded: boolean currentUsage: number limit: number scope: string - message?: string - reason?: 'billing_blocked' | 'usage_unavailable' accountBillingDecision?: AccountBillingDecision -} - -/** The 402 body for a refused new turn: the same bodies continuation validation answers with. */ -async function admissionRefusal( - userId: string, - admission: AdmissionBillingDecision, - usage: AdmissionUsage -): Promise { - if (usage.reason === 'billing_blocked') { - return { code: COPILOT_BILLING_BLOCKED_CODE, error: usage.message ?? BILLING_BLOCKED_MESSAGE } - } - if (usage.reason === 'usage_unavailable') { - return { code: COPILOT_USAGE_UNAVAILABLE_CODE, error: USAGE_UNAVAILABLE_MESSAGE } - } - const usageUpgrade = await resolveUsageUpgradePayload( - userId, - admission.kind === 'attributed' || admission.kind === 'legacy-scoped' - ? admission.attribution - : undefined, - usage.scope === 'member' ? 'member' : undefined - ) - return { code: COPILOT_USAGE_LIMIT_EXCEEDED_CODE, error: usageUpgrade.message, usageUpgrade } -} - -async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise { +}> { if (admission.kind === 'attributed' || admission.kind === 'legacy-scoped') { const usage = await checkAttributedUsageLimits(admission.attribution) const enforcedUsage = @@ -291,8 +261,6 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: enforcedUsage?.currentUsage ?? 0, limit: enforcedUsage?.limit ?? 0, scope: usage.scope ?? 'payer', - ...(usage.message ? { message: usage.message } : {}), - ...(usage.reason ? { reason: usage.reason } : {}), } } @@ -307,8 +275,6 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: usage.currentUsage, limit: usage.limit, scope: 'account', - ...(usage.message ? { message: usage.message } : {}), - ...(usage.reason ? { reason: usage.reason } : {}), accountBillingDecision: { userId: admission.userId, billingEntity: billingContext.billingEntity, @@ -330,8 +296,6 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: usage.currentUsage, limit: usage.limit, scope: 'account', - ...(usage.message ? { message: usage.message } : {}), - ...(usage.reason ? { reason: usage.reason } : {}), } } @@ -558,10 +522,7 @@ export const POST = withRouteHandler((req: NextRequest) => }) span.setAttribute(TraceAttr.CopilotValidateOutcome, CopilotValidateOutcome.UsageExceeded) span.setAttribute(TraceAttr.HttpStatusCode, 402) - return NextResponse.json( - await admissionRefusal(userId, admission, usage), - { status: 402 } - ) + return new NextResponse(null, { status: 402 }) } const responseHeaders: Record = {} diff --git a/apps/sim/lib/api/contracts/copilot.ts b/apps/sim/lib/api/contracts/copilot.ts index 3427cfd2ed5..d3682550460 100644 --- a/apps/sim/lib/api/contracts/copilot.ts +++ b/apps/sim/lib/api/contracts/copilot.ts @@ -297,27 +297,11 @@ export type ValidateCopilotApiKeyBillingBlocked = z.output< typeof validateCopilotApiKeyBillingBlockedSchema > -export const COPILOT_USAGE_UNAVAILABLE_CODE = 'USAGE_UNAVAILABLE' - -/** - * A new turn refused because usage could not be read. Admission fails closed on it; a run - * already under way is never refused for it. - */ -export const validateCopilotApiKeyUsageUnavailableSchema = z.object({ - code: z.literal(COPILOT_USAGE_UNAVAILABLE_CODE), - error: z.string(), -}) -export type ValidateCopilotApiKeyUsageUnavailable = z.output< - typeof validateCopilotApiKeyUsageUnavailableSchema -> - -/** Every 402 from validation, new turn and continuation alike, carries one of these bodies. */ +/** A continuation 402 carries one of these bodies; a new turn's 402 is empty. */ export const validateCopilotApiKeyRefusalSchema = z.union([ validateCopilotApiKeyUsageExceededSchema, validateCopilotApiKeyBillingBlockedSchema, - validateCopilotApiKeyUsageUnavailableSchema, ]) -export type ValidateCopilotApiKeyRefusal = z.output export const listCopilotApiKeysContract = defineRouteContract({ method: 'GET', diff --git a/apps/sim/lib/billing/calculations/usage-monitor.test.ts b/apps/sim/lib/billing/calculations/usage-monitor.test.ts index c45e971f687..5b587a36342 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.test.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.test.ts @@ -226,23 +226,10 @@ describe('checkServerSideUsageLimits', () => { periodEnd: new Date('2026-07-01T00:00:00.000Z'), }) - expect(result).toMatchObject({ isExceeded: true, reason: 'usage_unavailable' }) + expect(result.isExceeded).toBe(true) expect(result.message ?? '').not.toMatch(/\$/) }) - it('keeps a blocked account blocked when its ledger cannot be read', async () => { - dbChainMockFns.limit.mockResolvedValueOnce([{ blocked: true, blockedReason: 'dispute' }]) - mockGetBillingPeriodUsageCost.mockRejectedValueOnce(new Error('canceling statement')) - - const result = await checkServerSideUsageLimits('user-1', null) - - expect(result).toMatchObject({ - isExceeded: true, - reason: 'billing_blocked', - message: 'Account frozen. Please contact support to resolve this issue.', - }) - }) - it('keeps blocked accounts blocked while reporting their real ledger usage', async () => { dbChainMockFns.limit.mockResolvedValueOnce([{ blocked: true, blockedReason: 'payment_failed' }]) const subscription = { diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index 035551c0beb..9ebb874e8b8 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.ts @@ -1,7 +1,7 @@ import { db } from '@sim/db' import { userStats } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { getErrorMessage, toError } from '@sim/utils/errors' +import { toError } from '@sim/utils/errors' import { eq } from 'drizzle-orm' import { USAGE_UNAVAILABLE_MESSAGE } from '@/lib/billing/constants' import { isOrganizationBillingBlocked } from '@/lib/billing/core/access' @@ -306,38 +306,6 @@ export async function checkBillingEntityBlocked( * @param userId The ID of the user to check * @returns An object containing the exceeded status and usage details */ -/** - * A blocked account's ledger usage, for display only: `/api/users/me/usage-limits` exposes it as - * `currentPeriodCost`. The account is blocked either way, so a failed read reports zero rather - * than turning the blocked refusal into an unreadable-usage one. - */ -async function readBlockedAccountUsage( - userId: string, - preloadedSubscription: UsageLimitSubscription | null | undefined, - preloadedBillingContext: BillingContext | undefined -): Promise { - try { - const sub = - preloadedSubscription !== undefined - ? preloadedSubscription - : await getHighestPrioritySubscription(userId) - const subIsOrgScoped = isOrgScopedSubscription(sub, userId) - const billingEntity: BillingEntity = - subIsOrgScoped && sub - ? { type: 'organization', id: sub.referenceId } - : { type: 'user', id: userId } - const billingPeriod = preloadedBillingContext?.billingPeriod ?? - resolveSubscriptionUsagePeriod(sub) ?? { ...defaultBillingPeriod(), source: 'default' } - return await getBillingPeriodUsageCost(billingEntity, billingPeriod) - } catch (error) { - logger.warn('Blocked account usage could not be read', { - userId, - error: getErrorMessage(error), - }) - return 0 - } -} - export async function checkServerSideUsageLimits( userId: string, preloadedSubscription?: UsageLimitSubscription | null, @@ -347,8 +315,6 @@ export async function checkServerSideUsageLimits( currentUsage: number limit: number message?: string - /** Why a refusal is not a spent limit: a blocked account, or usage that could not be read. */ - reason?: 'billing_blocked' | 'usage_unavailable' }> { try { if (!isBillingEnabled) { @@ -363,18 +329,21 @@ export async function checkServerSideUsageLimits( const blocked = await checkBillingBlocked(userId) if (blocked.blocked) { - const currentUsage = await readBlockedAccountUsage( - userId, - preloadedSubscription, - preloadedBillingContext - ) - return { - isExceeded: true, - currentUsage, - limit: 0, - message: blocked.message, - reason: 'billing_blocked', - } + // Enforcement stays blocked, but surfaced usage must be the real ledger + // value — `/api/users/me/usage-limits` exposes it as `currentPeriodCost`. + const sub = + preloadedSubscription !== undefined + ? preloadedSubscription + : await getHighestPrioritySubscription(userId) + const subIsOrgScoped = isOrgScopedSubscription(sub, userId) + const billingEntity: BillingEntity = + subIsOrgScoped && sub + ? { type: 'organization', id: sub.referenceId } + : { type: 'user', id: userId } + const billingPeriod = preloadedBillingContext?.billingPeriod ?? + resolveSubscriptionUsagePeriod(sub) ?? { ...defaultBillingPeriod(), source: 'default' } + const currentUsage = await getBillingPeriodUsageCost(billingEntity, billingPeriod) + return { isExceeded: true, currentUsage, limit: 0, message: blocked.message } } const usageData = await checkUsageStatus(userId, preloadedSubscription, preloadedBillingContext) @@ -395,7 +364,6 @@ export async function checkServerSideUsageLimits( : usageData.isExceeded ? exceededMessage : undefined, - ...(usageData.unavailable ? { reason: 'usage_unavailable' as const } : {}), } } catch (error) { logger.error('Error in server-side usage limit check', { @@ -412,7 +380,6 @@ export async function checkServerSideUsageLimits( isExceeded: true, currentUsage: 0, limit: 0, - reason: 'usage_unavailable', message: error instanceof Error && error.message.includes('No user stats record found') ? 'User account not properly initialized. Please contact support.' diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index 109d9eda33d..36ea034f81e 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -62,11 +62,10 @@ function currentPeriodKey(attribution: BillingAttributionSnapshot): string { /** The admitted payer's attribution for its current subscription period. */ async function currentAttribution( - attribution: BillingAttributionSnapshot, - fresh: boolean + attribution: BillingAttributionSnapshot ): Promise { const key = currentPeriodKey(attribution) - const cached = fresh ? undefined : currentPeriodCache.get(key) + const cached = currentPeriodCache.get(key) // A cached period that has since ended is stale: the payer may already be in the next one. if (cached && !periodHasEnded(cached)) return cached const current = await refreshAttributionPeriod(attribution) @@ -109,9 +108,9 @@ async function readGateVerdict( /** * Judges a run against the period its charges land in. A Stripe-period payer's charges roll into * whatever period the subscription is in now (a rollover or an early anchor reset included), so - * such a run is judged against the payer's CURRENT period: a read that straddles the end of that - * period is judged again against the next one, and a current period that cannot be read makes - * the verdict unknown, so the run continues. Any other payer's charges stay in the admitted + * such a run is judged against the payer's CURRENT period. A current period that cannot be read, + * or that ends before its verdict is read, makes the verdict unknown, so the run continues and the + * next callback judges the next period. Any other payer's charges stay in the admitted * period (a reporting window, or the open default one), so that period is judged, even after it * ends. */ @@ -120,21 +119,19 @@ export async function readMidRunUsageVerdict( ): Promise { if (!isHosted || !isBillingEnabled) return { status: 'within' } if (!rollsIntoCurrentPeriod(attribution.billingPeriod)) return readGateVerdict(attribution) - for (let attempt = 0; attempt < 2; attempt++) { - let judged: BillingAttributionSnapshot - try { - judged = await currentAttribution(attribution, attempt > 0) - if (periodHasEnded(judged)) return { status: 'unknown' } - } catch (error) { - logger.warn('Current billing period could not be read; continuing the run', { - error: getErrorMessage(error), - }) - return { status: 'unknown' } - } - const verdict = await readGateVerdict(judged) - if (!periodHasEnded(judged)) return verdict + let judged: BillingAttributionSnapshot + try { + judged = await currentAttribution(attribution) + } catch (error) { + logger.warn('Current billing period could not be read; continuing the run', { + error: getErrorMessage(error), + }) + return { status: 'unknown' } } - return { status: 'unknown' } + if (periodHasEnded(judged)) return { status: 'unknown' } + const verdict = await readGateVerdict(judged) + // A period that ended during the read is judged at the next callback, which reloads it. + return periodHasEnded(judged) ? { status: 'unknown' } : verdict } /** diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index 4c164e6ad73..e03401fb602 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -9,9 +9,8 @@ import type { db } from '@sim/db' import * as schema from '@sim/db/schema' import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' import { getPostgresErrorCode } from '@sim/utils/errors' -import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' -import { sql } from 'drizzle-orm' +import { eq, sql } from 'drizzle-orm' import { drizzle } from 'drizzle-orm/postgres-js' import postgres from 'postgres' import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' @@ -316,27 +315,40 @@ describe('Cumulative billing with PostgreSQL', () => { it("counts a reporting run's top-ups after its window ends in that window, and a later run's charges in the next", async () => { const payer = { type: 'organization', id: 'payer' } as const - const boundary = new Date(Date.now() + 1500) - const admitted = { - start: new Date(boundary.getTime() - 30 * 24 * 60 * 60 * 1000), - end: boundary, + const dayMs = 24 * 60 * 60 * 1000 + // A reporting window is summed by when each row was created; the stamped period only binds a + // request's rows to each other. + const stamp = { + start: new Date('2026-01-01'), + end: new Date('2027-01-01'), source: 'reporting' as const, } - const next = { - start: boundary, - end: new Date(boundary.getTime() + 30 * 24 * 60 * 60 * 1000), - source: 'reporting' as const, + await recordCumulativeUsage({ ...usage(0.4, 'update-cost:long-run'), billingPeriod: stamp }) + const [first] = await database + .select({ createdAt: schema.usageLog.createdAt }) + .from(schema.usageLog) + .where(eq(schema.usageLog.eventKey, 'update-cost:long-run')) + // The admitted window ends right after the run's first charge, and every later write starts + // once the database clock has passed that boundary. + const boundary = new Date(first.createdAt.getTime() + 1) + for (;;) { + const [{ passed }] = await connection<{ passed: boolean }[]>` + select clock_timestamp()::timestamp > created_at + interval '1 millisecond' as passed + from usage_log where event_key = 'update-cost:long-run' + ` + if (passed) break } - const windowTotal = (period: typeof admitted) => - getBillingPeriodUsageCost(payer, period, undefined, database) - await recordCumulativeUsage({ ...usage(0.4, 'update-cost:long-run'), billingPeriod: admitted }) - await sleep(boundary.getTime() - Date.now() + 100) - await recordCumulativeUsage({ ...usage(1, 'update-cost:long-run'), billingPeriod: admitted }) - await recordCumulativeUsage({ ...usage(0.25, 'update-cost:next-run'), billingPeriod: next }) + await recordCumulativeUsage({ ...usage(1, 'update-cost:long-run'), billingPeriod: stamp }) + await recordCumulativeUsage({ ...usage(0.25, 'update-cost:next-run'), billingPeriod: stamp }) - expect(await windowTotal(admitted)).toBeCloseTo(1, 9) - expect(await windowTotal(next)).toBeCloseTo(0.25, 9) + const windowTotal = (start: Date, end: Date) => + getBillingPeriodUsageCost(payer, { start, end, source: 'reporting' }, undefined, database) + expect(await windowTotal(new Date(boundary.getTime() - 30 * dayMs), boundary)).toBeCloseTo(1, 9) + expect(await windowTotal(boundary, new Date(boundary.getTime() + 30 * dayMs))).toBeCloseTo( + 0.25, + 9 + ) }) describe('a request that outlives its billing period', () => { @@ -512,19 +524,5 @@ describe('Cumulative billing with PostgreSQL', () => { } expect(await stampedWindowTotal(start, end)).toBeCloseTo(0.6, 9) }) - - it.each([ - ['with', { payerSubscriptionId: 'sub-1' }], - ['without', {}], - ])( - 'refuses a request key that could collide with period rows %s a payer subscription', - async (_case, extra) => { - await setSubscriptionPeriod(0) - await expect( - recordCumulativeUsage({ ...usage(0.4, 'update-cost:request@1'), ...extra }) - ).rejects.toThrow('@') - expect(await ledgerRows()).toEqual([]) - } - ) }) }) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index a780a2371d5..40276626118 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -621,7 +621,8 @@ const MAX_CUMULATIVE_PERIOD_ROWS = 12 /** * The ledger key of the `index`-th period a cumulative request rolled into; 0 is the request key. - * No cumulative request key may contain `@`, so these keys never collide with another request's. + * The cost callback refuses a request key containing `@`, so these never collide with another + * request's. */ function cumulativePeriodEventKey(eventKey: string, index: number): string { return index === 0 ? eventKey : `${eventKey}@${index}` @@ -784,9 +785,6 @@ export async function recordCumulativeUsage( } const billingContext = await resolveBillingContext(userId, billingEntity, billingPeriod) - if (eventKey.includes('@')) { - throw new Error(`Cumulative usage event "${eventKey}" must not contain "@"`) - } const startedAt = Date.now() let stage: CumulativeUsageStage = 'pool' diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index d62d441fc6c..7568ccd7626 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -79,7 +79,6 @@ vi.mock('@/lib/mothership/tools/server/files/file-preview', async () => { }) import { - BillingLimitError, buildPreviewContentUpdate, CopilotBackendError, decodeJsonStringPrefix, @@ -891,53 +890,6 @@ describe('copilot go stream helpers', () => { expect(fetch).toHaveBeenCalledTimes(1) }) - it('keeps a blocked account off the usage card when the worker refuses a leg', async () => { - vi.mocked(fetch).mockResolvedValueOnce( - new Response(JSON.stringify({ code: 'BILLING_BLOCKED', error: 'Billing account frozen.' }), { - status: 402, - }) - ) - const blocked = await runStreamLoop( - 'https://example.com/api/tools/resume', - {}, - createStreamingContext(), - turnScopedExecContext(), - { timeout: 1000 } - ).then( - () => undefined, - (error: unknown) => error - ) - expect(blocked).not.toBeInstanceOf(BillingLimitError) - expect(blocked).toMatchObject({ code: 'forbidden', message: 'Billing account frozen.' }) - - vi.mocked(fetch).mockResolvedValueOnce( - new Response( - JSON.stringify({ code: 'USAGE_UNAVAILABLE', error: 'Usage could not be verified.' }), - { status: 402 } - ) - ) - await expect( - runStreamLoop( - 'https://example.com/api/tools/resume', - {}, - createStreamingContext(), - turnScopedExecContext(), - { timeout: 1000 } - ) - ).rejects.toMatchObject({ code: 'forbidden' }) - - vi.mocked(fetch).mockResolvedValueOnce(new Response(null, { status: 402 })) - await expect( - runStreamLoop( - 'https://example.com/api/tools/resume', - {}, - createStreamingContext(), - turnScopedExecContext(), - { timeout: 1000 } - ) - ).rejects.toBeInstanceOf(BillingLimitError) - }) - it('fails closed when the shared stream ends before a terminal event', async () => { const textEvent = createEvent({ streamId: 'stream-1', diff --git a/apps/sim/lib/mothership/request/go/stream.ts b/apps/sim/lib/mothership/request/go/stream.ts index 56b76b24239..cd16ae492ad 100644 --- a/apps/sim/lib/mothership/request/go/stream.ts +++ b/apps/sim/lib/mothership/request/go/stream.ts @@ -2,11 +2,6 @@ import { type Context, SpanStatusCode } from '@opentelemetry/api' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { toRecordOrNull } from '@sim/utils/object' -import { - COPILOT_BILLING_BLOCKED_CODE, - COPILOT_USAGE_UNAVAILABLE_CODE, -} from '@/lib/api/contracts/copilot' -import { OrchestrationError } from '@/lib/core/orchestration/types' import { WORKER_STREAM_IDLE_TIMEOUT_MS } from '@/lib/mothership/constants' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { CopilotSseCloseReason } from '@/lib/mothership/generated/trace-attribute-values-v1' @@ -111,16 +106,6 @@ function userFacingRejection(value: unknown): string | undefined { return message } -const BILLING_BLOCKED_MESSAGE = 'Billing account is blocked' - -function parseJsonRecord(body: string): Record | null { - try { - return toRecordOrNull(JSON.parse(body)) - } catch { - return null - } -} - /** * What the user is told about a failed backend response. A 5xx or a gateway page * is upstream detail and stays on the error for logs; a 4xx may carry the @@ -283,18 +268,6 @@ export async function runStreamLoop( }).finally(() => clearTimeout(idleTimer)) if (response.status === 402) { - // A blocked account, or usage that could not be read, is refused as such, never with the - // usage card. - const refusal = parseJsonRecord(errorText) - if ( - refusal?.code === COPILOT_BILLING_BLOCKED_CODE || - refusal?.code === COPILOT_USAGE_UNAVAILABLE_CODE - ) { - throw new OrchestrationError( - 'forbidden', - userFacingRejection(refusal.error) ?? BILLING_BLOCKED_MESSAGE - ) - } throw new BillingLimitError(execContext.userId) }