diff --git a/apps/sim/lib/api-key/service.test.ts b/apps/sim/lib/api-key/service.test.ts index 02fec43ca54..04011bb46a8 100644 --- a/apps/sim/lib/api-key/service.test.ts +++ b/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> = {}) { 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() + }) +}) diff --git a/apps/sim/lib/api-key/service.ts b/apps/sim/lib/api-key/service.ts index aa451498c1d..4520b92f22a 100644 --- a/apps/sim/lib/api-key/service.ts +++ b/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,30 @@ export async function authenticateApiKeyFromHeader( } } +const LAST_USED_STALENESS_WINDOW_MS = 10 * 60 * 1000 + /** - * 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 uses a staleness window: it only + * fires when the stored value is older than + * {@link LAST_USED_STALENESS_WINDOW_MS}. High-traffic keys otherwise rewrite + * the same row on every request, serializing concurrent requests behind row + * locks. The 10-minute window matches GitLab's personal-access-token + * last-used tracking. */ export async function updateApiKeyLastUsed(keyId: string): Promise { try { - await db.update(apiKeyTable).set({ lastUsed: new Date() }).where(eq(apiKeyTable.id, keyId)) + const staleBefore = new Date(Date.now() - LAST_USED_STALENESS_WINDOW_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) } diff --git a/apps/sim/lib/billing/core/usage.test.ts b/apps/sim/lib/billing/core/usage.test.ts new file mode 100644 index 00000000000..a12b9ada731 --- /dev/null +++ b/apps/sim/lib/billing/core/usage.test.ts @@ -0,0 +1,157 @@ +/** + * 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 }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ currentUsageLimit: '10' }]) + 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 }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ currentUsageLimit: '40' }]) + 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), + }) + }) + + it('returns a concurrently written limit when the guarded heal matches no rows', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: null }]) + dbChainMockFns.returning.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: '30' }]) + mockGetFreeTierLimit.mockReturnValue(10) + + const limit = await getUserUsageLimit('user-1', null) + + expect(limit).toBe(30) + }) + + it('still returns the fallback when the heal write fails', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: null }]) + dbChainMockFns.returning.mockRejectedValueOnce(new Error('connection lost')) + mockGetFreeTierLimit.mockReturnValue(10) + + await expect(getUserUsageLimit('user-1', null)).resolves.toBe(10) + }) +}) diff --git a/apps/sim/lib/billing/core/usage.ts b/apps/sim/lib/billing/core/usage.ts index a39d3d1309f..3d1fcf04be2 100644 --- a/apps/sim/lib/billing/core/usage.ts +++ b/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,13 @@ 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. The + * write-back is best-effort: a limit written concurrently wins, and a failed + * write still resolves to the fallback instead of blocking execution. */ export async function getUserUsageLimit( userId: string, @@ -537,9 +544,43 @@ 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() + + try { + const healed = await db + .update(userStats) + .set({ + currentUsageLimit: fallbackLimit.toString(), + usageLimitUpdatedAt: new Date(), + }) + .where(and(eq(userStats.userId, userId), isNull(userStats.currentUsageLimit))) + .returning({ currentUsageLimit: userStats.currentUsageLimit }) + + if (healed.length === 0) { + const concurrent = await db + .select({ currentUsageLimit: userStats.currentUsageLimit }) + .from(userStats) + .where(eq(userStats.userId, userId)) + .limit(1) + + if (concurrent[0]?.currentUsageLimit) { + return toNumber(toDecimal(concurrent[0].currentUsageLimit)) + } + } + + logger.warn('Healed null usage limit to plan default', { + userId, + plan: subscription?.plan || 'free', + fallbackLimit, + }) + } catch (error) { + logger.error('Failed to heal null usage limit', { userId, fallbackLimit, error }) + } + + return fallbackLimit } return toNumber(toDecimal(userStatsQuery[0].currentUsageLimit))