fix(billing): hold free plans to their allowance mid-run

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

@@ -249,7 +249,7 @@ Execution rate buckets can burst above the sustained per-minute rate. General AP
 
 | Plan | Concurrent runs |
 | --- | --- |
-| Free | 10 |
+| Free | 2 |
 | Pro | 50 |
 | Max | 200 |
 | Enterprise default | 1,000; customizable |

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

@@ -1,11 +1,12 @@
 /**
  * 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, and spend after the payer's terminal settlement is refused. Only the
- * internal-key check is stubbed.
+ * after its invoice, and spend after the payer's terminal settlement is refused. A free payer's
+ * callback that takes it past its limit pauses the run even while an earlier admission is
+ * cached, and a paid payer keeps the cached admission. Only the internal-key check is stubbed.
  */
 import { db } from '@sim/db'
-import { subscription, usageLog, user, userStats } from '@sim/db/schema'
+import { project, subscription, usageLog, user, userStats, workspace } from '@sim/db/schema'
 import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock'
 import { generateId } from '@sim/utils/id'
 import { eq, inArray } from 'drizzle-orm'
@@ -24,8 +25,12 @@ vi.mock('@/lib/mothership/request/http', async (importOriginal) => ({
 
 import {
   BILLING_ACCOUNT_DECISION_HEADER,
+  BILLING_ATTRIBUTION_HEADER,
+  type BillingAttributionSnapshot,
   serializeAccountBillingDecisionHeader,
+  serializeBillingAttributionHeader,
 } from '@/lib/billing/core/attribution'
+import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
 import { claimTerminalPeriod } from '@/lib/billing/cycle-close'
 import { POST } from '@/app/api/billing/update-cost/route'
 
@@ -35,6 +40,8 @@ const userIds: string[] = []
 afterAll(async () => {
   if (userIds.length === 0) return
   await db.delete(usageLog).where(inArray(usageLog.userId, userIds))
+  await db.delete(workspace).where(inArray(workspace.ownerId, userIds))
+  await db.delete(project).where(inArray(project.ownerId, userIds))
   await db.delete(subscription).where(inArray(subscription.referenceId, userIds))
   await db.delete(userStats).where(inArray(userStats.userId, userIds))
   await db.delete(user).where(inArray(user.id, userIds))
@@ -47,10 +54,9 @@ interface Payer {
   decision: string
 }
 
-/** A user on a pro subscription for `period`, whose close marker has caught up to it. */
-async function createPayer(period: { start: Date; end: Date }): Promise<Payer> {
+/** A user with the default $5 usage limit and no subscription. */
+async function createUser(): Promise<string> {
   const userId = `update-cost-user-${generateId()}`
-  const subscriptionId = generateId()
   userIds.push(userId)
   await db.insert(user).values({
     id: userId,
@@ -61,6 +67,13 @@ async function createPayer(period: { start: Date; end: Date }): Promise<Payer> {
     updatedAt: new Date(),
   })
   await db.insert(userStats).values({ id: generateId(), userId })
+  return userId
+}
+
+/** A user on a pro subscription for `period`, whose close marker has caught up to it. */
+async function createPayer(period: { start: Date; end: Date }): Promise<Payer> {
+  const userId = await createUser()
+  const subscriptionId = generateId()
   await db.insert(subscription).values({
     id: subscriptionId,
     plan: 'pro',
@@ -94,6 +107,21 @@ function requestRows(requestKey: string) {
     .where(inArray(usageLog.eventKey, [`update-cost:${requestKey}`, `update-cost:${requestKey}@1`]))
 }
 
+/** A free user's direct-v1 decision, admitted in the open default period as admission mints it. */
+async function createFreePayer(): Promise<Payer> {
+  const userId = await createUser()
+  const decision = serializeAccountBillingDecisionHeader({
+    userId,
+    billingEntity: { type: 'user', id: userId },
+    billingPeriod: {
+      start: defaultBillingPeriod().start.toISOString(),
+      end: defaultBillingPeriod().end.toISOString(),
+      source: 'default',
+    },
+  })
+  return { userId, subscriptionId: '', decision }
+}
+
 function callback(payer: Payer, requestKey: string, cost: number): NextRequest {
   return new NextRequest('http://localhost:3000/api/billing/update-cost', {
     method: 'POST',
@@ -114,6 +142,65 @@ function callback(payer: Payer, requestKey: string, cost: number): NextRequest {
   })
 }
 
+/** A free user's personal workspace and the attribution its runs are admitted with. */
+async function createFreeWorkspacePayer(): Promise<BillingAttributionSnapshot> {
+  const userId = await createUser()
+  const projectId = generateId()
+  const workspaceId = generateId()
+  await db.insert(project).values({ id: projectId, name: 'Update Cost Test', ownerId: userId })
+  await db.insert(workspace).values({
+    id: workspaceId,
+    projectId,
+    name: 'Update Cost Test',
+    ownerId: userId,
+    billedAccountUserId: userId,
+  })
+  return {
+    actorUserId: userId,
+    workspaceId,
+    organizationId: null,
+    billedAccountUserId: userId,
+    billingEntity: { type: 'user', id: userId },
+    billingPeriod: {
+      start: defaultBillingPeriod().start.toISOString(),
+      end: defaultBillingPeriod().end.toISOString(),
+      source: 'default',
+    },
+    payerSubscription: null,
+  }
+}
+
+function attributedCallback(
+  attribution: BillingAttributionSnapshot,
+  requestKey: string,
+  cost: number
+): 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': 'attribution-v1',
+      'x-sim-billing-request-id': requestKey,
+      [BILLING_ATTRIBUTION_HEADER]: serializeBillingAttributionHeader(attribution),
+    },
+    body: JSON.stringify({
+      userId: attribution.actorUserId,
+      workspaceId: attribution.workspaceId,
+      cost,
+      model: 'test-model',
+      source: 'copilot',
+      idempotencyKey: requestKey,
+    }),
+  })
+}
+
+async function usageExceeded(response: ReturnType<typeof POST>): Promise<boolean> {
+  const res = await response
+  expect(res.status).toBe(200)
+  return (await res.json()).usageExceeded
+}
+
 describe('direct-v1 cost callbacks in PostgreSQL', () => {
   it('records spend after a Stripe rollover in the payer current period', async () => {
     const now = Date.now()
@@ -157,3 +244,38 @@ describe('direct-v1 cost callbacks in PostgreSQL', () => {
     expect((await requestRows(requestKey)).map((row) => Number(row.cost))).toEqual([0.5])
   })
 })
+
+/**
+ * Each payer's first callback caches an admission ($1 of a $5 limit), and its second, inside the
+ * near-limit cache window, records spend that takes it to $6.
+ */
+describe('usage standing after a charge that crosses the limit', () => {
+  it('pauses a free direct-v1 run although its admission is cached', async () => {
+    const payer = await createFreePayer()
+
+    expect(await usageExceeded(POST(callback(payer, generateId(), 1), {}))).toBe(false)
+    expect(await usageExceeded(POST(callback(payer, generateId(), 5), {}))).toBe(true)
+  })
+
+  it('pauses a free attributed run although its admission is cached', async () => {
+    const attribution = await createFreeWorkspacePayer()
+
+    expect(await usageExceeded(POST(attributedCallback(attribution, generateId(), 1), {}))).toBe(
+      false
+    )
+    expect(await usageExceeded(POST(attributedCallback(attribution, generateId(), 5), {}))).toBe(
+      true
+    )
+  })
+
+  it("keeps serving a paid payer's cached admission", async () => {
+    const now = Date.now()
+    const payer = await createPayer({
+      start: new Date(now - 10 * DAY_MS),
+      end: new Date(now + 20 * DAY_MS),
+    })
+
+    expect(await usageExceeded(POST(callback(payer, generateId(), 1), {}))).toBe(false)
+    expect(await usageExceeded(POST(callback(payer, generateId(), 5), {}))).toBe(false)
+  })
+})

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

@@ -931,6 +931,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
       source: 'stripe' as const,
     },
   }
+  const PAID_SUBSCRIPTION = { id: 'sub-1', plan: 'team', status: 'active', seats: 4 }
+  const PAID_ATTRIBUTION = { ...CURRENT_ATTRIBUTION, payerSubscription: PAID_SUBSCRIPTION }
 
   function attributedCallback() {
     callbackSequence += 1
@@ -1005,7 +1007,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
     })
   })
 
-  it('serves a cached admission to every step and re-reads a refusal', async () => {
+  it('serves a paid payer a cached admission to every step and re-reads a refusal', async () => {
+    mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
     mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false })
 
     expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(false)
@@ -1022,6 +1025,17 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
     expect((await (await POST(attributedCallback())).json()).usageExceeded).toBe(false)
   })
 
+  it('reads a free payer standing past its cached admission after every charge', async () => {
+    mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false })
+    expect((await (await POST(attributedCallback(), {})).json()).usageExceeded).toBe(false)
+
+    mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
+    await expect((await POST(attributedCallback(), {})).json()).resolves.toMatchObject({
+      usageExceeded: true,
+      usageUpgrade: { action: 'upgrade_plan' },
+    })
+  })
+
   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' })
@@ -1209,7 +1223,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
   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) => {
+  ])('does not pause a paid run for %s', async (_case, verdict) => {
+    mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
     mockCheckAttributedUsageLimits.mockResolvedValue(verdict)
 
     const body = await (await POST(attributedCallback())).json()
@@ -1218,6 +1233,42 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
     expect(body).not.toHaveProperty('usageUpgrade')
   })
 
+  it.each([
+    [
+      'an unreadable ledger',
+      () =>
+        mockCheckAttributedUsageLimits.mockResolvedValue({
+          isExceeded: true,
+          reason: 'usage_unavailable',
+        }),
+    ],
+    [
+      'a failed gate read',
+      () => mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out')),
+    ],
+  ])('pauses a free run for %s, keeping the recorded charge', async (_case, answer) => {
+    answer()
+
+    const res = await POST(attributedCallback(), {})
+
+    expect(res.status).toBe(200)
+    await expect(res.json()).resolves.toMatchObject({
+      success: true,
+      usageExceeded: true,
+      usageUpgrade: { action: 'upgrade_plan' },
+    })
+  })
+
+  it('does not pause a free run for a blocked account', async () => {
+    mockCheckAttributedUsageLimits.mockResolvedValue({
+      isExceeded: true,
+      reason: 'billing_blocked',
+      scope: 'payer',
+    })
+
+    expect((await (await POST(attributedCallback(), {})).json()).usageExceeded).toBe(false)
+  })
+
   it('tells a member over the cap their organization set who can raise it', async () => {
     mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'member' })
 
@@ -1284,8 +1335,11 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
     expect(body.usageExceeded).toBe(true)
   })
 
-  it('keeps a run going when its current period cannot be read', async () => {
-    mockRequireBillingAttributionHeader.mockReturnValue(STRIPE_ATTRIBUTION)
+  it('keeps a paid run going when its current period cannot be read', async () => {
+    mockRequireBillingAttributionHeader.mockReturnValue({
+      ...STRIPE_ATTRIBUTION,
+      payerSubscription: PAID_SUBSCRIPTION,
+    })
     mockRefreshAttributionPeriod.mockRejectedValue(new Error('subscription read timed out'))
     mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' })
 
@@ -1304,31 +1358,39 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
       vi.useRealTimers()
     })
 
-    it('answers not exceeded when the standing read outlasts the callback budget', async () => {
-      mockCheckAttributedUsageLimits.mockImplementation(async () => {
-        await sleep(1500)
-        return { isExceeded: true, scope: 'payer' }
-      })
+    it.each([
+      ['a paid', PAID_ATTRIBUTION, false],
+      ['a free', CURRENT_ATTRIBUTION, true],
+    ])(
+      'answers %s payer whose standing read outlasts the callback budget',
+      async (_case, attribution, exceeded) => {
+        mockRequireBillingAttributionHeader.mockReturnValue(attribution)
+        mockCheckAttributedUsageLimits.mockImplementation(async () => {
+          await sleep(1500)
+          return { isExceeded: false }
+        })
 
-      const pending = POST(attributedCallback())
-      await vi.advanceTimersByTimeAsync(1000)
-      const res = await pending
+        const pending = POST(attributedCallback(), {})
+        await vi.advanceTimersByTimeAsync(1000)
+        const res = await pending
 
-      expect(res.status).toBe(200)
-      await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: false })
-      // Settle the abandoned read so no later test coalesces onto it.
-      await vi.advanceTimersByTimeAsync(500)
-    })
+        expect(res.status).toBe(200)
+        await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: exceeded })
+        // Settle the abandoned read so no later test coalesces onto it.
+        await vi.advanceTimersByTimeAsync(500)
+      }
+    )
 
-    it('does not pause a run on a verdict read across the end of its period', async () => {
+    it('does not pause a paid run on a verdict read across the end of its period', async () => {
       const straddling = {
-        ...CURRENT_ATTRIBUTION,
+        ...PAID_ATTRIBUTION,
         billingPeriod: {
           start: '2026-07-01T00:00:00.000Z',
           end: new Date(Date.now() + 40).toISOString(),
           source: 'stripe' as const,
         },
       }
+      mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
       mockRefreshAttributionPeriod.mockResolvedValue(straddling)
       mockCheckAttributedUsageLimits.mockImplementation(async () => {
         await sleep(80)
@@ -1363,7 +1425,8 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
     })
   })
 
-  it('keeps a recorded charge successful when the gate read fails', async () => {
+  it('keeps a paid recorded charge successful and unpaused when the gate read fails', async () => {
+    mockRequireBillingAttributionHeader.mockReturnValue(PAID_ATTRIBUTION)
     mockCheckAttributedUsageLimits.mockRejectedValue(new Error('ledger read timed out'))
 
     const res = await POST(attributedCallback())