improvement(billing): self-heal null usage limits and debounce api-key last-used writes

Author: waleedlatif1

apps/sim/lib/api-key/service.test.ts

@@ -52,7 +52,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({
 }))
 
 import { hashApiKey } from '@/lib/api-key/crypto'
-import { authenticateApiKeyFromHeader } from '@/lib/api-key/service'
+import { authenticateApiKeyFromHeader, updateApiKeyLastUsed } from '@/lib/api-key/service'
 
 function personalKeyRecord(overrides: Partial<Record<string, unknown>> = {}) {
   return {
@@ -141,3 +141,33 @@ describe('authenticateApiKeyFromHeader', () => {
     expect(JSON.stringify(filter)).toContain(expected)
   })
 })
+
+describe('updateApiKeyLastUsed', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('only writes when the stored lastUsed is missing or stale', async () => {
+    await updateApiKeyLastUsed('key-1')
+
+    expect(dbChainMockFns.update).toHaveBeenCalledTimes(1)
+    expect(dbChainMockFns.set).toHaveBeenCalledWith({ lastUsed: expect.any(Date) })
+    const [condition] = dbChainMockFns.where.mock.calls[0]
+    expect(condition).toMatchObject({
+      type: 'and',
+      conditions: [
+        { type: 'eq', right: 'key-1' },
+        { type: 'or', conditions: [{ type: 'isNull' }, { type: 'lt' }] },
+      ],
+    })
+  })
+
+  it('swallows database errors instead of failing the request', async () => {
+    dbChainMockFns.update.mockImplementationOnce(() => {
+      throw new Error('connection lost')
+    })
+
+    await expect(updateApiKeyLastUsed('key-1')).resolves.toBeUndefined()
+    expect(serviceLogger.error).toHaveBeenCalled()
+  })
+})

apps/sim/lib/api-key/service.ts

@@ -1,7 +1,7 @@
 import { db } from '@sim/db'
 import { apiKey as apiKeyTable, user as userTable } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
-import { and, eq } from 'drizzle-orm'
+import { and, eq, isNull, lt, or } from 'drizzle-orm'
 import { hashApiKey } from '@/lib/api-key/crypto'
 import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
 import { getWorkspaceBillingSettings, type WorkspaceBillingSettings } from '@/lib/workspaces/utils'
@@ -136,12 +136,28 @@ export async function authenticateApiKeyFromHeader(
   }
 }
 
+const LAST_USED_DEBOUNCE_MS = 60_000
+
 /**
- * Update the last used timestamp for an API key
+ * Update the last used timestamp for an API key.
+ *
+ * `lastUsed` is display-only, so the write is debounced: it only fires when
+ * the stored value is older than {@link LAST_USED_DEBOUNCE_MS}. High-traffic
+ * keys otherwise rewrite the same row on every request, serializing
+ * concurrent requests behind row locks.
  */
 export async function updateApiKeyLastUsed(keyId: string): Promise<void> {
   try {
-    await db.update(apiKeyTable).set({ lastUsed: new Date() }).where(eq(apiKeyTable.id, keyId))
+    const staleBefore = new Date(Date.now() - LAST_USED_DEBOUNCE_MS)
+    await db
+      .update(apiKeyTable)
+      .set({ lastUsed: new Date() })
+      .where(
+        and(
+          eq(apiKeyTable.id, keyId),
+          or(isNull(apiKeyTable.lastUsed), lt(apiKeyTable.lastUsed, staleBefore))
+        )
+      )
   } catch (error) {
     logger.error('Error updating API key last used:', error)
   }

apps/sim/lib/billing/core/usage.test.ts

@@ -0,0 +1,136 @@
+/**
+ * Tests for getUserUsageLimit.
+ *
+ * Org-scoped members carry a null `currentUsageLimit` by design, so a user
+ * whose subscription stops being org-scoped without a resync is left null.
+ * The limit read must self-heal that state to the plan default instead of
+ * failing closed and blocking every execution.
+ *
+ * @vitest-environment node
+ */
+import { dbChainMock, dbChainMockFns } from '@sim/testing'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+vi.mock('@sim/db', () => dbChainMock)
+
+const {
+  mockGetFreeTierLimit,
+  mockGetPerUserMinimumLimit,
+  mockHasPaidSubscriptionStatus,
+  mockIsOrgScopedSubscription,
+} = vi.hoisted(() => ({
+  mockGetFreeTierLimit: vi.fn(),
+  mockGetPerUserMinimumLimit: vi.fn(),
+  mockHasPaidSubscriptionStatus: vi.fn(),
+  mockIsOrgScopedSubscription: vi.fn(),
+}))
+
+vi.mock('@/lib/billing/subscriptions/utils', () => ({
+  canEditUsageLimit: vi.fn(),
+  getFreeTierLimit: mockGetFreeTierLimit,
+  getPerUserMinimumLimit: mockGetPerUserMinimumLimit,
+  getPlanPricing: vi.fn(() => ({ basePrice: 20 })),
+  hasPaidSubscriptionStatus: mockHasPaidSubscriptionStatus,
+  hasUsableSubscriptionAccess: vi.fn(),
+  isOrgScopedSubscription: mockIsOrgScopedSubscription,
+}))
+
+vi.mock('@/lib/billing/core/plan', () => ({
+  getHighestPrioritySubscription: vi.fn(),
+}))
+
+vi.mock('@/lib/billing/core/access', () => ({
+  getEffectiveBillingStatus: vi.fn(),
+}))
+
+vi.mock('@/lib/billing/core/usage-log', () => ({
+  getBillingPeriodUsageCost: vi.fn(),
+}))
+
+vi.mock('@/lib/billing/credits/daily-refresh', () => ({
+  computeDailyRefreshConsumed: vi.fn(),
+  getOrgMemberRefreshBounds: vi.fn(),
+}))
+
+vi.mock('@/components/emails', () => ({
+  getEmailSubject: vi.fn(),
+  renderCreditsExhaustedEmail: vi.fn(),
+  renderFreeTierUpgradeEmail: vi.fn(),
+  renderUsageThresholdEmail: vi.fn(),
+}))
+
+vi.mock('@/lib/messaging/email/mailer', () => ({
+  sendEmail: vi.fn(),
+}))
+
+vi.mock('@/lib/messaging/email/unsubscribe', () => ({
+  getEmailPreferences: vi.fn(),
+}))
+
+import { getUserUsageLimit } from '@/lib/billing/core/usage'
+
+const PRO_SUBSCRIPTION = {
+  id: 'sub-1',
+  plan: 'pro',
+  status: 'active',
+  referenceId: 'user-1',
+  seats: 1,
+  periodStart: null,
+  periodEnd: null,
+} as never
+
+describe('getUserUsageLimit', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockIsOrgScopedSubscription.mockReturnValue(false)
+  })
+
+  it('returns the stored limit when set', async () => {
+    dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: '25' }])
+
+    const limit = await getUserUsageLimit('user-1', null)
+
+    expect(limit).toBe(25)
+    expect(dbChainMockFns.update).not.toHaveBeenCalled()
+  })
+
+  it('throws when no userStats row exists', async () => {
+    dbChainMockFns.limit.mockResolvedValueOnce([])
+
+    await expect(getUserUsageLimit('user-1', null)).rejects.toThrow('No user stats record found')
+  })
+
+  it('heals a null limit to the free-tier default for free users', async () => {
+    dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: null }])
+    mockGetFreeTierLimit.mockReturnValue(10)
+
+    const limit = await getUserUsageLimit('user-1', null)
+
+    expect(limit).toBe(10)
+    expect(dbChainMockFns.update).toHaveBeenCalledTimes(1)
+    expect(dbChainMockFns.set).toHaveBeenCalledWith({
+      currentUsageLimit: '10',
+      usageLimitUpdatedAt: expect.any(Date),
+    })
+    const [condition] = dbChainMockFns.where.mock.calls.at(-1) ?? []
+    expect(condition).toMatchObject({
+      type: 'and',
+      conditions: [{ type: 'eq', right: 'user-1' }, { type: 'isNull' }],
+    })
+  })
+
+  it('heals a null limit to the plan minimum for paid personal subscriptions', async () => {
+    dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: null }])
+    mockHasPaidSubscriptionStatus.mockReturnValue(true)
+    mockGetPerUserMinimumLimit.mockReturnValue(40)
+
+    const limit = await getUserUsageLimit('user-1', PRO_SUBSCRIPTION)
+
+    expect(limit).toBe(40)
+    expect(mockGetPerUserMinimumLimit).toHaveBeenCalledWith(PRO_SUBSCRIPTION)
+    expect(dbChainMockFns.set).toHaveBeenCalledWith({
+      currentUsageLimit: '40',
+      usageLimitUpdatedAt: expect.any(Date),
+    })
+  })
+})

apps/sim/lib/billing/core/usage.ts

@@ -2,7 +2,7 @@ import { db } from '@sim/db'
 import { member, organization, settings, user, userStats } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { generateId } from '@sim/utils/id'
-import { eq } from 'drizzle-orm'
+import { and, eq, isNull } from 'drizzle-orm'
 import {
   getEmailSubject,
   renderCreditsExhaustedEmail,
@@ -495,6 +495,11 @@ export async function updateUserUsageLimit(
  * Get usage limit for a user (used by checkUsageStatus for server-side
  * checks). Org-scoped subs return the organization limit;
  * personally-scoped subs return the individual user limit from userStats.
+ *
+ * Org-scoped members carry a null `currentUsageLimit` by design (see
+ * `syncUsageLimitsFromSubscription`). A user whose subscription stops being
+ * org-scoped without a resync would otherwise stay null and fail closed on
+ * every execution, so a null limit self-heals to the plan default here.
  */
 export async function getUserUsageLimit(
   userId: string,
@@ -537,9 +542,26 @@ export async function getUserUsageLimit(
   }
 
   if (!userStatsQuery[0].currentUsageLimit) {
-    throw new Error(
-      `Invalid null usage limit for ${subscription?.plan || 'free'} user: ${userId}. User stats must be properly initialized.`
-    )
+    const fallbackLimit =
+      subscription && hasPaidSubscriptionStatus(subscription.status)
+        ? getPerUserMinimumLimit(subscription)
+        : getFreeTierLimit()
+
+    await db
+      .update(userStats)
+      .set({
+        currentUsageLimit: fallbackLimit.toString(),
+        usageLimitUpdatedAt: new Date(),
+      })
+      .where(and(eq(userStats.userId, userId), isNull(userStats.currentUsageLimit)))
+
+    logger.warn('Healed null usage limit to plan default', {
+      userId,
+      plan: subscription?.plan || 'free',
+      fallbackLimit,
+    })
+
+    return fallbackLimit
   }
 
   return toNumber(toDecimal(userStatsQuery[0].currentUsageLimit))