fix(billing): deny the hosted key when a free payer's balance is unreadable

Author: waleedlatif1

apps/sim/lib/billing/core/chat-access.ts

@@ -7,6 +7,15 @@ import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags'
 
 const logger = createLogger('PaidAccess')
 
+interface PaidAccessOptions {
+  /**
+   * The verdict for a free payer whose credit balance cannot be read. Chat admits, because a
+   * database blip must not lock a paying customer out and every turn is bounded by the usage
+   * check that follows. A hosted key denies, because nothing after acquisition bounds it.
+   */
+  onBalanceUnavailable: 'admit' | 'deny'
+}
+
 /**
  * Whether a payer may use a feature a hosted deployment reserves for paying payers: Chat, and
  * hosted keys a tool marks `requiresPaidPlan`. A hosted deployment that bills withholds these from
@@ -15,11 +24,11 @@ const logger = createLogger('PaidAccess')
  *
  * The plan is the payer's entitled subscription (active or past due), so a canceled
  * plan reads as free; a blocked payer is refused by the usage check every caller runs next.
- *
- * A balance that cannot be read admits the payer: a database blip must not lock a paying
- * customer out, and a free payer stays bounded by its usage limit.
  */
-export async function payerHasPaidAccess(payer: UsageUpgradePayer): Promise<boolean> {
+export async function payerHasPaidAccess(
+  payer: UsageUpgradePayer,
+  options: PaidAccessOptions
+): Promise<boolean> {
   if (!isHosted || !isBillingEnabled) return true
   if (!isFree(payer.payerSubscription?.plan)) return true
   try {
@@ -29,22 +38,23 @@ export async function payerHasPaidAccess(payer: UsageUpgradePayer): Promise<bool
     )
     return balance > 0
   } catch (error) {
-    logger.warn('Credit balance unavailable; admitting the payer', {
+    logger.warn('Credit balance unavailable for a free payer', {
       billingEntity: payer.billingEntity,
+      verdict: options.onBalanceUnavailable,
       error: getErrorMessage(error),
     })
-    return true
+    return options.onBalanceUnavailable === 'admit'
   }
 }
 
 /**
  * Whether a payer may use Chat (talking to Sim): {@link payerHasPaidAccess}, since every turn
- * spends hosted model credits.
+ * spends hosted model credits, admitting a payer whose balance cannot be read.
  *
  * This is the single verdict behind every Chat plan gate: run admission, chat titles, scheduled
  * job dispatch, Copilot key admission, and the workspace host context the composer reads, so the
  * composer never offers a send that admission will refuse.
  */
 export function payerHasChatAccess(payer: UsageUpgradePayer): Promise<boolean> {
-  return payerHasPaidAccess(payer)
+  return payerHasPaidAccess(payer, { onBalanceUnavailable: 'admit' })
 }

apps/sim/tools/index.test.ts

@@ -3,6 +3,10 @@ import { createSessionPrincipal } from '@sim/testing/factories/principal.factory
 import { jsonResponse } from '@sim/testing/helpers/http'
 import { apiKeyByokMock, apiKeyByokMockFns } from '@sim/testing/mocks/api-key-byok.mock'
 import { authInternalMock, authInternalMockFns } from '@sim/testing/mocks/auth-internal.mock'
+import {
+  billingAttributionMock,
+  billingAttributionMockFns,
+} from '@sim/testing/mocks/billing-attribution.mock'
 import { billingUsageLogMock } from '@sim/testing/mocks/billing-usage-log.mock'
 import {
   executorPrincipalMock,
@@ -179,6 +183,8 @@ vi.mock('@/lib/billing/core/chat-access', () => ({
   payerHasPaidAccess: mockPayerHasPaidAccess,
 }))
 
+vi.mock('@/lib/billing/core/attribution', () => billingAttributionMock)
+
 vi.mock('@/lib/uploads/contexts/workspace/file-manager', () => workspaceFileManagerMock)
 
 vi.mock(
@@ -448,6 +454,7 @@ import { createToolConfig, getTool } from '@/tools/utils'
 import { getToolAsync } from '@/tools/utils.server'
 
 const mockGetBYOKKey = apiKeyByokMockFns.mockGetBYOKKey
+const mockResolveWorkspaceBillingPayer = billingAttributionMockFns.mockResolveWorkspaceBillingPayer
 const mockGenerateInternalToken = authInternalMockFns.mockGenerateInternalToken
 const mockCreateExecutorPrincipalFromExecutionContext =
   executorPrincipalMockFns.mockCreateExecutorPrincipalFromExecutionContext
@@ -6619,7 +6626,9 @@ describe('Hosted Key Injection', () => {
 
       expect(result.success).toBe(false)
       expect(result.error).toContain('available on paid plans')
-      expect(mockPayerHasPaidAccess).toHaveBeenCalledWith(TEST_BILLING_ATTRIBUTION)
+      expect(mockPayerHasPaidAccess).toHaveBeenCalledWith(TEST_BILLING_ATTRIBUTION, {
+        onBalanceUnavailable: 'deny',
+      })
       expect(mockRateLimiterFns.acquireKey).not.toHaveBeenCalled()
       expect(request).not.toHaveBeenCalled()
     })
@@ -6635,6 +6644,60 @@ describe('Hosted Key Injection', () => {
       expect(request).toHaveBeenCalledTimes(1)
     })
 
+    async function runPaidOnlyToolWithoutAttribution() {
+      mockGetBYOKKey.mockResolvedValue(null)
+      mockValidateUrlWithDNS.mockResolvedValue({ isValid: true, resolvedIP: '93.184.216.34' })
+      mockSecureFetchWithPinnedIP.mockResolvedValue({
+        ok: true,
+        status: 200,
+        statusText: 'OK',
+        headers: { get: () => null, toRecord: () => ({}) },
+        json: async () => ({ ok: true }),
+        text: async () => '{"ok":true}',
+      })
+      const result = await executeTool(PAID_ONLY_TOOL_ID, {
+        _context: { workspaceId: 'workspace-without-run', userId: 'user-1' },
+      })
+      return { result, request: mockSecureFetchWithPinnedIP }
+    }
+
+    it('judges the workspace payer when the run carries no attribution', async () => {
+      mockResolveWorkspaceBillingPayer.mockResolvedValue({
+        billedAccountUserId: 'payer-1',
+        organizationId: 'org-1',
+        payerSubscription: { plan: 'team' },
+      })
+      mockPayerHasPaidAccess.mockResolvedValue(true)
+
+      const { result, request } = await runPaidOnlyToolWithoutAttribution()
+
+      expect(result.success).toBe(true)
+      expect(mockResolveWorkspaceBillingPayer).toHaveBeenCalledWith('workspace-without-run', {
+        onMissing: 'return-null',
+      })
+      expect(mockPayerHasPaidAccess).toHaveBeenCalledWith(
+        {
+          billingEntity: { type: 'organization', id: 'org-1' },
+          payerSubscription: { plan: 'team' },
+        },
+        { onBalanceUnavailable: 'deny' }
+      )
+      expect(request).toHaveBeenCalledTimes(1)
+    })
+
+    it('refuses the hosted key when the workspace has no payer', async () => {
+      mockResolveWorkspaceBillingPayer.mockResolvedValue(null)
+      mockPayerHasPaidAccess.mockResolvedValue(true)
+
+      const { result, request } = await runPaidOnlyToolWithoutAttribution()
+
+      expect(result.success).toBe(false)
+      expect(result.error).toContain('available on paid plans')
+      expect(mockPayerHasPaidAccess).not.toHaveBeenCalled()
+      expect(mockRateLimiterFns.acquireKey).not.toHaveBeenCalled()
+      expect(request).not.toHaveBeenCalled()
+    })
+
     it("runs a free payer's own workspace key without consulting the plan", async () => {
       mockGetBYOKKey.mockResolvedValue({ apiKey: 'workspace-own-key' })
       mockPayerHasPaidAccess.mockResolvedValue(false)

apps/sim/tools/index.ts

@@ -557,7 +557,7 @@ async function registerAssistantCredentialSecrets(
 /**
  * Refuses a `requiresPaidPlan` hosted key to a run whose payer lacks paid access. The payer is
  * the run's admitted attribution, or the workspace's payer for a surface that carries none; a
- * workspace without a payer is refused.
+ * workspace without a payer, or a free payer whose credit balance cannot be read, is refused.
  */
 async function assertPayerMayUseHostedKey(
   tool: ToolDefinition,
@@ -581,7 +581,7 @@ async function assertPayerMayUseHostedKey(
     }
   }
 
-  if (payer && (await payerHasPaidAccess(payer))) return
+  if (payer && (await payerHasPaidAccess(payer, { onBalanceUnavailable: 'deny' }))) return
 
   throw new HostedKeyRequiresPaidPlanError(
     `${tool.name} with Sim's hosted key is available on paid plans. Upgrade your plan, or add your own API key under Settings > BYOK.`