fix(billing): serve execution admission from the usage gate cache instead of re-summing the billing period per run

Author: waleedlatif1

apps/sim/app/api/v1/knowledge/search/route.test.ts

@@ -70,7 +70,9 @@ vi.mock('@/lib/billing/calculations/usage-monitor', () => ({
 vi.mock('@/lib/billing/core/billing-attribution', () => ({
   resolveBillingAttribution: mockResolveBillingAttribution,
   resolveSystemBillingAttribution: mockResolveSystemBillingAttribution,
-  checkAttributedUsageLimits: vi.fn().mockResolvedValue({ isExceeded: false }),
+}))
+vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
+  checkSearchUsageLimits: vi.fn().mockResolvedValue({ isExceeded: false }),
 }))
 
 vi.mock('@/lib/knowledge/embeddings', () => ({

apps/sim/app/api/v1/knowledge/search/route.ts

@@ -2,10 +2,10 @@ import { type NextRequest, NextResponse } from 'next/server'
 import { v1KnowledgeSearchContract } from '@/lib/api/contracts/v1/knowledge'
 import { parseRequest } from '@/lib/api/server'
 import {
-  checkAttributedUsageLimits,
   resolveBillingAttribution,
   resolveSystemBillingAttribution,
 } from '@/lib/billing/core/billing-attribution'
+import { checkSearchUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants'
 import { toKbEmbeddingDimensions } from '@/lib/knowledge/embedding-models'
@@ -82,7 +82,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
      * keys resolve their system actor and immutable payer from one workspace read.
      */
     if (billingAttribution) {
-      const usage = await checkAttributedUsageLimits(billingAttribution)
+      const usage = await checkSearchUsageLimits(billingAttribution)
       if (usage.isExceeded) {
         return NextResponse.json(
           { error: usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' },

apps/sim/background/workflow-column-execution.ts

@@ -10,9 +10,9 @@ import { and, eq, isNull, or } from 'drizzle-orm'
 import {
   assertBillingAttributionSnapshot,
   type BillingAttributionSnapshot,
-  checkAttributedUsageLimits,
   toBillingContext,
 } from '@/lib/billing/core/billing-attribution'
+import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billing'
 import { isRetryableInfrastructureError } from '@/lib/core/errors/retryable-infrastructure'
 import {
@@ -575,7 +575,7 @@ async function runWorkflowAndWriteTerminal(
          * Gate the exact workspace payer and member cap before hosted-key cost.
          * A denial clears the cell pre-stamp and surfaces the upgrade state.
          */
-        const usage = await checkAttributedUsageLimits(enrichmentBillingAttribution)
+        const usage = await checkExecutionUsageLimits(enrichmentBillingAttribution)
         if (usage.isExceeded) {
           logger.warn(
             `Usage limit reached — halting enrichment (table=${tableId} row=${rowId} group=${groupId})`

apps/sim/lib/billing/core/usage-gate-cache.test.ts

@@ -12,6 +12,7 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({
 
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import {
+  checkExecutionUsageLimits,
   checkIngestionUsageLimits,
   checkSearchUsageLimits,
   resetUsageGateCache,
@@ -158,3 +159,23 @@ describe('checkSearchUsageLimits', () => {
     expect(mockCheck).toHaveBeenCalledTimes(2)
   })
 })
+
+describe('checkExecutionUsageLimits', () => {
+  beforeEach(() => {
+    resetUsageGateCache()
+    mockCheck.mockReset().mockResolvedValue({ isExceeded: false })
+  })
+
+  it('reuses an admission across workspaces of the same payer', async () => {
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    await checkExecutionUsageLimits({ ...ATTRIBUTION, workspaceId: 'ws-2' })
+    expect(mockCheck).toHaveBeenCalledTimes(1)
+  })
+
+  it('re-reads a refusal', async () => {
+    mockCheck.mockResolvedValue({ isExceeded: true, message: 'over' })
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+})

apps/sim/lib/billing/core/usage-gate-cache.ts

@@ -99,15 +99,28 @@ export function checkIngestionUsageLimits(
  * {@link checkAttributedUsageLimits} for knowledge search. Serves only a cached
  * admission: a refusal is always re-read, so a payer who just raised their limit
  * or upgraded is never held behind a cached block while they wait on a search.
- * Every other interactive caller (uploads, execution admission, the settings
- * surfaces) keeps reading the gate fresh.
+ * Every other interactive caller (uploads, the settings surfaces) keeps reading
+ * the gate fresh.
  */
 export function checkSearchUsageLimits(
   attribution: BillingAttributionSnapshot
 ): Promise<AttributedUsageLimitsResult> {
   return checkUsageLimitsThroughCache(attribution, false)
 }
 
+/**
+ * {@link checkAttributedUsageLimits} for the execution path, the highest-volume
+ * reader of the gate. Same policy as search: only an admission is served from
+ * cache, so a raised limit applies on the next run. The admission's usage figure
+ * is up to one TTL stale, which widens the execution-slot reservation headroom
+ * by that much and no more.
+ */
+export function checkExecutionUsageLimits(
+  attribution: BillingAttributionSnapshot
+): Promise<AttributedUsageLimitsResult> {
+  return checkUsageLimitsThroughCache(attribution, false)
+}
+
 /** Drops every cached gate answer. Test seam; never called in production code. */
 export function resetUsageGateCache(): void {
   gateCache.clear()

apps/sim/lib/execution/preprocessing.test.ts

@@ -46,10 +46,12 @@ vi.mock('@/lib/billing/calculations/usage-reservation', () => ({
 }))
 vi.mock('@/lib/billing/core/billing-attribution', () => ({
   assertBillingAttributionSnapshot: vi.fn((value) => value),
-  checkAttributedUsageLimits: mockCheckAttributedUsageLimits,
   resolveBillingAttribution: mockResolveBillingAttribution,
   resolveSystemBillingAttribution: mockResolveSystemBillingAttribution,
 }))
+vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
+  checkExecutionUsageLimits: mockCheckAttributedUsageLimits,
+}))
 vi.mock('@/lib/billing/core/subscription', () => ({
   getHighestPrioritySubscription: vi.fn(),
 }))

apps/sim/lib/execution/preprocessing.ts

@@ -9,12 +9,12 @@ import {
 import {
   assertBillingAttributionSnapshot,
   type BillingAttributionSnapshot,
-  checkAttributedUsageLimits,
   resolveBillingAttribution,
   resolveSystemBillingAttribution,
 } from '@/lib/billing/core/billing-attribution'
 import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
 import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
+import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import {
   type AdmissionErrorDescriptor,
   getReservationDenialDescriptor,
@@ -574,8 +574,8 @@ export async function preprocessExecution(
     let snapshot: UsageSnapshot | null = null
     try {
       const usageCheck = await withDatabaseReadRetry(
-        () => checkAttributedUsageLimits(billingAttribution),
-        { label: 'checkAttributedUsageLimits' }
+        () => checkExecutionUsageLimits(billingAttribution),
+        { label: 'checkExecutionUsageLimits' }
       )
       snapshot = usageCheck.payerUsage
         ? {

apps/sim/lib/internal/guardrails/operations.test.ts

@@ -27,10 +27,12 @@ vi.mock('@/lib/auth/credential-access', () => ({
   authorizeCredentialUseForAuth: mocks.authorizeCredential,
 }))
 vi.mock('@/lib/billing/core/billing-attribution', () => ({
-  checkAttributedUsageLimits: mocks.checkAttributedUsageLimits,
   requireBillingAttributionHeader: mocks.requireBillingAttribution,
   toBillingContext: vi.fn(() => ({})),
 }))
+vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
+  checkExecutionUsageLimits: mocks.checkAttributedUsageLimits,
+}))
 vi.mock('@/lib/billing/threshold-billing', () => ({
   checkAndBillPayerOverageThreshold: vi.fn(),
 }))

apps/sim/lib/internal/guardrails/operations.ts

@@ -6,10 +6,10 @@ import { authorizeCredentialUseForAuth } from '@/lib/auth/credential-access'
 import { AuthType } from '@/lib/auth/hybrid'
 import {
   type BillingAttributionSnapshot,
-  checkAttributedUsageLimits,
   requireBillingAttributionHeader,
   toBillingContext,
 } from '@/lib/billing/core/billing-attribution'
+import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billing'
 import { prepareCopilotEnvironmentContext } from '@/lib/copilot/environment-context'
 import { inspectModelInputProvenanceRequest } from '@/lib/execution/model-input-provenance'
@@ -164,7 +164,7 @@ async function prepareHallucinationContext(
     throw error
   }
 
-  const usage = await checkAttributedUsageLimits(billingAttribution)
+  const usage = await checkExecutionUsageLimits(billingAttribution)
   if (usage.isExceeded) {
     fail(402, usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.')
   }

apps/sim/lib/workflows/custom-blocks/child-execution.test.ts

@@ -10,8 +10,8 @@ const { mockCheckAttributedUsageLimits, mockSubscribe, mockUnsubscribe } = vi.ho
   mockUnsubscribe: vi.fn(),
 }))
 
-vi.mock('@/lib/billing/core/billing-attribution', () => ({
-  checkAttributedUsageLimits: mockCheckAttributedUsageLimits,
+vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
+  checkExecutionUsageLimits: mockCheckAttributedUsageLimits,
 }))
 
 vi.mock('@/lib/execution/cancellation', () => ({