opensource/sim
3
0
Fork 0
mirror of https://github.com/simstudioai/sim.git synced 2026-10-24 11:36:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

improvement(billing): migrate hot path writes away from user_stats (#4768)

Browse Source 
* improvement(billing): migrate hot path writes away from user_stats

* fix period start and end sot

* address comments

* Remove stale billing migration

Co-authored-by: Cursor <cursoragent@cursor.com>

* regen migrations

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Vikhyath Mondreti
2026-05-28 13:54:27 -07:00
committed by GitHub
co-authored by Cursor
parent e62c3ad2d3
commit 53eaa600af
25 changed files with 18308 additions and 305 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/sim/app/api/billing/update-cost/route.ts
+2 -17
View File
@@ -1,7 +1,6 @@
import type { Span } from '@opentelemetry/api'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { sql } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { billingUpdateCostContract } from '@/lib/api/contracts/subscription'
import { parseRequest } from '@/lib/api/server'
@@ -155,21 +154,6 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
source,
})
const totalTokens = inputTokens + outputTokens
const additionalStats: Record<string, ReturnType<typeof sql>> = {
totalCopilotCost: sql`total_copilot_cost + ${cost}`,
currentPeriodCopilotCost: sql`current_period_copilot_cost + ${cost}`,
totalCopilotCalls: sql`total_copilot_calls + 1`,
totalCopilotTokens: sql`total_copilot_tokens + ${totalTokens}`,
}
if (isMcp) {
additionalStats.totalMcpCopilotCost = sql`total_mcp_copilot_cost + ${cost}`
additionalStats.currentPeriodMcpCopilotCost = sql`current_period_mcp_copilot_cost + ${cost}`
additionalStats.totalMcpCopilotCalls = sql`total_mcp_copilot_calls + 1`
}
await recordUsage({
userId,
entries: [
@@ -178,10 +162,11 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
source,
description: model,
cost,
eventKey: idempotencyKey ? `update-cost:${idempotencyKey}` : undefined,
sourceReference: idempotencyKey ? `update-cost:${idempotencyKey}` : requestId,
metadata: { inputTokens, outputTokens },
},
],
additionalStats,
})
usageCommitted = true
apps/sim/app/api/mcp/copilot/route.ts
-21
View File
@@ -10,13 +10,10 @@ import {
McpError,
type RequestId,
} from '@modelcontextprotocol/sdk/types.js'
import { db } from '@sim/db'
import { userStats } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { authorizeWorkflowByWorkspacePermission } from '@sim/workflow-authz'
import { eq, sql } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { mcpRequestBodySchema, mcpToolCallParamsSchema } from '@/lib/api/contracts/mcp'
import { validateOAuthAccessToken } from '@/lib/auth/oauth-token'
@@ -299,8 +296,6 @@ function buildMcpServer(abortSignal?: AbortSignal): Server {
abortSignal
)
trackMcpCopilotCall(authResult.userId)
return result
})
@@ -391,22 +386,6 @@ export const DELETE = withRouteHandler(async (request: NextRequest) => {
return NextResponse.json(createError(0, -32000, 'Method not allowed.'), { status: 405 })
})
/**
* Increment MCP copilot call counter in userStats (fire-and-forget).
*/
function trackMcpCopilotCall(userId: string): void {
db.update(userStats)
.set({
totalMcpCopilotCalls: sql`total_mcp_copilot_calls + 1`,
lastActive: new Date(),
})
.where(eq(userStats.userId, userId))
.then(() => {})
.catch((error) => {
logger.error('Failed to track MCP copilot call', { error, userId })
})
}
async function handleToolsCall(
params: { name: string; arguments?: Record<string, unknown> },
userId: string,
apps/sim/app/api/speech/token/route.ts
+6
View File
@@ -1,3 +1,4 @@
import { createHash } from 'node:crypto'
import { db } from '@sim/db'
import { chat } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
@@ -31,6 +32,10 @@ const STT_TOKEN_RATE_LIMIT = {
refillIntervalMs: 72 * 1000,
} as const
function hashVoiceToken(token: string): string {
return createHash('sha256').update(token).digest('hex')
}
const rateLimiter = new RateLimiter()
async function validateChatAuth(
@@ -163,6 +168,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
source: 'voice-input',
description: `Voice input session (${maxMinutes} min)`,
cost: sessionCost * getCostMultiplier(),
sourceReference: `voice-input:${hashVoiceToken(data.token)}`,
},
],
}).catch((err) => {
apps/sim/app/api/wand/route.ts
+2 -5
View File
@@ -1,7 +1,7 @@
import { db } from '@sim/db'
import { workflow } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { eq, sql } from 'drizzle-orm'
import { eq } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { wandGenerateContract } from '@/lib/api/contracts'
import { parseRequest } from '@/lib/api/server'
@@ -106,7 +106,6 @@ async function updateUserStatsForWand(
}
try {
const totalTokens = usage.total_tokens || 0
const promptTokens = usage.prompt_tokens || 0
const completionTokens = usage.completion_tokens || 0
@@ -138,12 +137,10 @@ async function updateUserStatsForWand(
source: 'wand',
description: modelName,
cost: costToStore,
sourceReference: `wand:${requestId}`,
metadata: { inputTokens: promptTokens, outputTokens: completionTokens },
},
],
additionalStats: {
totalTokensUsed: sql`total_tokens_used + ${totalTokens}`,
},
})
await checkAndBillOverageThreshold(billingUserId)
apps/sim/background/workflow-column-execution.ts
+1
View File
@@ -282,6 +282,7 @@ async function runWorkflowAndWriteTerminal(
source: 'enrichment',
description: enrichment.name,
cost,
sourceReference: `enrichment:${tableId}:${rowId}:${enrichment.id}`,
metadata: { enrichmentId: enrichment.id, tableId, rowId },
},
],
apps/sim/lib/billing/calculations/usage-monitor.ts
+119 -84
View File
@@ -3,11 +3,13 @@ import { member, userStats } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { and, eq } from 'drizzle-orm'
import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
import {
getHighestPrioritySubscription,
type HighestPrioritySubscription,
} from '@/lib/billing/core/plan'
import { getPooledOrgCurrentPeriodCost, getUserUsageLimit } from '@/lib/billing/core/usage'
import { getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log'
import {
computeDailyRefreshConsumed,
getOrgMemberRefreshBounds,
@@ -37,11 +39,6 @@ interface UsageData {
organizationId: string | null
}
/**
* Sum `currentPeriodCost` across all members of an org, then subtract
* daily-refresh credits (with per-user window bounds for mid-cycle
* joiners).
*/
async function computePooledOrgUsage(
organizationId: string,
sub: {
@@ -54,25 +51,16 @@ async function computePooledOrgUsage(
const { memberIds, currentPeriodCost } = await getPooledOrgCurrentPeriodCost(organizationId)
if (memberIds.length === 0) return 0
let pooled = currentPeriodCost
const billingPeriod =
sub.periodStart && sub.periodEnd
? { start: sub.periodStart, end: sub.periodEnd }
: defaultBillingPeriod()
const ledgerUsage = await getBillingPeriodUsageCost(
{ type: 'organization', id: organizationId },
billingPeriod
)
if (isPaid(sub.plan) && sub.periodStart) {
const planDollars = getPlanTierDollars(sub.plan)
if (planDollars > 0) {
const userBounds = await getOrgMemberRefreshBounds(organizationId, sub.periodStart)
const refresh = await computeDailyRefreshConsumed({
userIds: memberIds,
periodStart: sub.periodStart,
periodEnd: sub.periodEnd ?? null,
planDollars,
seats: sub.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
})
pooled = Math.max(0, pooled - refresh)
}
}
return pooled
return applyOrgRefresh(organizationId, sub, currentPeriodCost + ledgerUsage, memberIds)
}
/**
@@ -112,78 +100,57 @@ export async function checkUsageStatus(
const scope: 'user' | 'organization' = subIsOrgScoped ? 'organization' : 'user'
const organizationId: string | null = subIsOrgScoped && sub ? sub.referenceId : null
let currentUsage = 0
if (subIsOrgScoped && sub) {
currentUsage = await computePooledOrgUsage(sub.referenceId, sub)
} else {
const statsRecords = await db
.select()
.from(userStats)
.where(eq(userStats.userId, userId))
.limit(1)
if (statsRecords.length === 0) {
logger.info('No usage stats found for user', { userId, limit })
return {
percentUsed: 0,
isWarning: false,
isExceeded: false,
currentUsage: 0,
limit,
scope: 'user',
organizationId: null,
}
}
const rawUsage = toNumber(toDecimal(statsRecords[0].currentPeriodCost))
let refresh = 0
if (sub && isPaid(sub.plan) && sub.periodStart) {
const planDollars = getPlanTierDollars(sub.plan)
if (planDollars > 0) {
refresh = await computeDailyRefreshConsumed({
userIds: [userId],
periodStart: sub.periodStart,
periodEnd: sub.periodEnd ?? null,
planDollars,
})
}
}
currentUsage = Math.max(0, rawUsage - refresh)
const currentUsage = await computePooledOrgUsage(sub.referenceId, sub)
return buildUsageData({ currentUsage, limit, scope, organizationId })
}
const percentUsed = limit > 0 ? Math.min((currentUsage / limit) * 100, 100) : 100
const isExceeded = currentUsage >= limit
const isWarning = !isExceeded && percentUsed >= WARNING_THRESHOLD
const statsRecords = await db
.select()
.from(userStats)
.where(eq(userStats.userId, userId))
.limit(1)
logger.info('Final usage statistics', {
userId,
currentUsage,
limit,
percentUsed,
isWarning,
isExceeded,
scope,
organizationId,
})
return {
percentUsed,
isWarning,
isExceeded,
currentUsage,
limit,
scope,
organizationId,
if (statsRecords.length === 0) {
logger.info('No usage stats found for user', { userId, limit })
return {
percentUsed: 0,
isWarning: false,
isExceeded: false,
currentUsage: 0,
limit,
scope: 'user',
organizationId: null,
}
}
const billingPeriod =
sub?.periodStart && sub.periodEnd
? { start: sub.periodStart, end: sub.periodEnd }
: defaultBillingPeriod()
const ledgerUsage = await getBillingPeriodUsageCost({ type: 'user', id: userId }, billingPeriod)
let currentUsage = toNumber(toDecimal(statsRecords[0].currentPeriodCost)) + ledgerUsage
if (sub && isPaid(sub.plan) && sub.periodStart) {
const planDollars = getPlanTierDollars(sub.plan)
if (planDollars > 0) {
const refresh = await computeDailyRefreshConsumed({
userIds: [userId],
periodStart: sub.periodStart,
periodEnd: sub.periodEnd ?? null,
planDollars,
billingEntity: { type: 'user', id: userId },
})
currentUsage = Math.max(0, currentUsage - refresh)
}
}
return buildUsageData({ currentUsage, limit, scope, organizationId })
} catch (error) {
logger.error('Error checking usage status', {
error: error instanceof Error ? { message: error.message, stack: error.stack } : error,
userId,
})
// Block execution if we can't determine usage status
logger.error('Cannot determine usage status - blocking execution', {
userId,
error: toError(error).message,
@@ -201,6 +168,74 @@ export async function checkUsageStatus(
}
}
async function applyOrgRefresh(
organizationId: string,
sub: {
plan: string | null
seats: number | null
periodStart: Date | null
periodEnd: Date | null
},
currentUsage: number,
preloadedMemberIds?: string[]
): Promise<number> {
if (!isPaid(sub.plan) || !sub.periodStart) {
return currentUsage
}
const memberIds =
preloadedMemberIds ?? (await getPooledOrgCurrentPeriodCost(organizationId)).memberIds
if (memberIds.length === 0) return currentUsage
const planDollars = getPlanTierDollars(sub.plan)
if (planDollars <= 0) return currentUsage
const userBounds = await getOrgMemberRefreshBounds(organizationId, sub.periodStart)
const refresh = await computeDailyRefreshConsumed({
userIds: memberIds,
periodStart: sub.periodStart,
periodEnd: sub.periodEnd ?? null,
planDollars,
seats: sub.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
billingEntity: { type: 'organization', id: organizationId },
})
return Math.max(0, currentUsage - refresh)
}
function buildUsageData(params: {
currentUsage: number
limit: number
scope: 'user' | 'organization'
organizationId: string | null
}): UsageData {
const { currentUsage, limit, scope, organizationId } = params
const percentUsed = limit > 0 ? Math.min((currentUsage / limit) * 100, 100) : 100
const isExceeded = currentUsage >= limit
const isWarning = !isExceeded && percentUsed >= WARNING_THRESHOLD
logger.info('Final usage statistics', {
currentUsage,
limit,
percentUsed,
isWarning,
isExceeded,
scope,
organizationId,
})
return {
percentUsed,
isWarning,
isExceeded,
currentUsage,
limit,
scope,
organizationId,
}
}
/**
* Displays a notification to the user when they're approaching their usage limit
* Can be called on app startup or before executing actions that might incur costs
apps/sim/lib/billing/core/billing-period.ts
+9
View File
@@ -0,0 +1,9 @@
const OPEN_BILLING_PERIOD_START = new Date(0)
const OPEN_BILLING_PERIOD_END = new Date(Date.UTC(9999, 11, 31))
export function defaultBillingPeriod(): { start: Date; end: Date } {
return {
start: OPEN_BILLING_PERIOD_START,
end: OPEN_BILLING_PERIOD_END,
}
}
apps/sim/lib/billing/core/billing.ts
+46 -8
View File
@@ -7,6 +7,7 @@ import {
type SubscriptionMetadata,
} from '@/lib/billing/core/subscription'
import { getOrgUsageLimit, getUserUsageData } from '@/lib/billing/core/usage'
import { getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log'
import { getCreditBalance } from '@/lib/billing/credits/balance'
import {
computeDailyRefreshConsumed,
@@ -185,6 +186,7 @@ export async function computeOrgOverageAmount(params: {
planDollars,
seats: params.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
billingEntity: { type: 'organization', id: params.organizationId },
})
}
@@ -223,6 +225,13 @@ export async function calculateSubscriptionOverage(sub: {
if (isOrgScoped) {
const pooled = await aggregateOrgMemberStats(sub.referenceId)
const ledgerUsage =
sub.periodStart && sub.periodEnd
? await getBillingPeriodUsageCost(
{ type: 'organization', id: sub.referenceId },
{ start: sub.periodStart, end: sub.periodEnd }
)
: 0
const orgData = await db
.select({ departedMemberUsage: organization.departedMemberUsage })
@@ -239,7 +248,7 @@ export async function calculateSubscriptionOverage(sub: {
periodStart: sub.periodStart ?? null,
periodEnd: sub.periodEnd ?? null,
organizationId: sub.referenceId,
pooledCurrentPeriodCost: pooled.currentPeriodCost,
pooledCurrentPeriodCost: pooled.currentPeriodCost + ledgerUsage,
departedMemberUsage,
memberIds: pooled.memberIds,
})
@@ -249,9 +258,10 @@ export async function calculateSubscriptionOverage(sub: {
logger.info('Calculated org-scoped overage', {
subscriptionId: sub.id,
plan: sub.plan,
currentMemberUsage: pooled.currentPeriodCost,
currentMemberUsage: pooled.currentPeriodCost + ledgerUsage,
departedMemberUsage,
totalUsage: pooled.currentPeriodCost + departedMemberUsage,
ledgerUsage,
totalUsage: pooled.currentPeriodCost + ledgerUsage + departedMemberUsage,
effectiveUsage,
baseSubscriptionAmount,
totalOverage,
@@ -275,11 +285,18 @@ export async function calculateSubscriptionOverage(sub: {
const personalCurrentUsage = statsRow ? toNumber(toDecimal(statsRow.currentPeriodCost)) : 0
const snapshotUsage = statsRow ? toNumber(toDecimal(statsRow.proPeriodCostSnapshot)) : 0
const snapshotAt = statsRow?.proPeriodCostSnapshotAt ?? null
const ledgerUsage =
sub.periodStart && sub.periodEnd
? await getBillingPeriodUsageCost(
{ type: 'user', id: sub.referenceId },
{ start: sub.periodStart, end: sub.periodEnd }
)
: 0
const joinedOrgMidCycle = snapshotAt !== null || snapshotUsage > 0
const totalProUsageDecimal = joinedOrgMidCycle
? toDecimal(snapshotUsage)
: toDecimal(personalCurrentUsage)
? toDecimal(snapshotUsage).plus(ledgerUsage)
: toDecimal(personalCurrentUsage).plus(ledgerUsage)
if (joinedOrgMidCycle) {
logger.info('Billing personal Pro only for pre-join usage (user joined org mid-cycle)', {
@@ -304,6 +321,7 @@ export async function calculateSubscriptionOverage(sub: {
periodStart: sub.periodStart,
periodEnd: refreshCap,
planDollars,
billingEntity: { type: 'user', id: sub.referenceId },
})
}
@@ -319,6 +337,7 @@ export async function calculateSubscriptionOverage(sub: {
joinedOrgMidCycle,
personalCurrentUsage,
snapshot: snapshotUsage,
ledgerUsage,
billedUsage: toNumber(totalProUsageDecimal),
dailyRefreshDeduction,
basePrice,
@@ -332,13 +351,24 @@ export async function calculateSubscriptionOverage(sub: {
.where(eq(userStats.userId, sub.referenceId))
.limit(1)
const personalCurrentUsage = statsRow ? toNumber(toDecimal(statsRow.currentPeriodCost)) : 0
const ledgerUsage =
sub.periodStart && sub.periodEnd
? await getBillingPeriodUsageCost(
{ type: 'user', id: sub.referenceId },
{ start: sub.periodStart, end: sub.periodEnd }
)
: 0
const { basePrice } = getPlanPricing(sub.plan || 'free')
totalOverageDecimal = Decimal.max(0, toDecimal(personalCurrentUsage).minus(basePrice))
totalOverageDecimal = Decimal.max(
0,
toDecimal(personalCurrentUsage).plus(ledgerUsage).minus(basePrice)
)
logger.info('Calculated overage for plan', {
subscriptionId: sub.id,
plan: sub.plan || 'free',
usage: personalCurrentUsage,
usage: personalCurrentUsage + ledgerUsage,
ledgerUsage,
basePrice,
totalOverage: toNumber(totalOverageDecimal),
})
@@ -431,6 +461,13 @@ export async function getSimplifiedBillingSummary(
// `usageData` is derived from the caller's priority subscription
// and may not match the requested org (multi-org admins, personal
// priority sub, etc.), so it cannot be reused here.
const ledgerUsage =
subscription.periodStart && subscription.periodEnd
? await getBillingPeriodUsageCost(
{ type: 'organization', id: organizationId },
{ start: subscription.periodStart, end: subscription.periodEnd }
)
: 0
let refreshDeduction = 0
if (isPaid(plan) && subscription.periodStart) {
const planDollars = getPlanTierDollars(plan)
@@ -446,10 +483,11 @@ export async function getSimplifiedBillingSummary(
planDollars,
seats: subscription.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
billingEntity: { type: 'organization', id: organizationId },
})
}
}
const effectiveCurrentUsage = Math.max(0, rawCurrentUsage - refreshDeduction)
const effectiveCurrentUsage = Math.max(0, rawCurrentUsage + ledgerUsage - refreshDeduction)
const { limit: orgUsageLimit } = await getOrgUsageLimit(
organizationId,
apps/sim/lib/billing/core/organization.ts
+9
View File
@@ -4,6 +4,7 @@ import { createLogger } from '@sim/logger'
import { and, count, eq, gt, ne } from 'drizzle-orm'
import { isOrganizationBillingBlocked } from '@/lib/billing/core/access'
import { getOrganizationSubscription, getPlanPricing } from '@/lib/billing/core/billing'
import { getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log'
import {
computeDailyRefreshConsumed,
getOrgMemberRefreshBounds,
@@ -122,6 +123,13 @@ export async function getOrganizationBillingData(
// Calculate aggregated statistics
let totalCurrentUsage = members.reduce((sum, m) => sum + m.currentUsage, 0)
if (subscription.periodStart && subscription.periodEnd) {
totalCurrentUsage += await getBillingPeriodUsageCost(
{ type: 'organization', id: subscription.referenceId },
{ start: subscription.periodStart, end: subscription.periodEnd }
)
}
if (isPaid(subscription.plan) && subscription.periodStart) {
const planDollars = getPlanTierDollars(subscription.plan)
if (planDollars > 0) {
@@ -137,6 +145,7 @@ export async function getOrganizationBillingData(
planDollars,
seats: subscription.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
billingEntity: { type: 'organization', id: subscription.referenceId },
})
totalCurrentUsage = Math.max(0, totalCurrentUsage - refreshConsumed)
}
apps/sim/lib/billing/core/usage-log.test.ts
+131
View File
@@ -0,0 +1,131 @@
/**
* @vitest-environment node
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const {
mockGetHighestPrioritySubscription,
mockInsert,
mockIsOrgScopedSubscription,
mockOnConflictDoNothing,
mockReturning,
mockValues,
} = vi.hoisted(() => ({
mockGetHighestPrioritySubscription: vi.fn(),
mockInsert: vi.fn(),
mockIsOrgScopedSubscription: vi.fn(),
mockOnConflictDoNothing: vi.fn(),
mockReturning: vi.fn(),
mockValues: vi.fn(),
}))
vi.mock('@sim/db', () => ({
db: {
insert: mockInsert,
},
}))
vi.mock('@sim/db/schema', () => ({
usageLog: {
billingEntityId: 'usageLog.billingEntityId',
billingEntityType: 'usageLog.billingEntityType',
billingPeriodEnd: 'usageLog.billingPeriodEnd',
billingPeriodStart: 'usageLog.billingPeriodStart',
category: 'usageLog.category',
cost: 'usageLog.cost',
createdAt: 'usageLog.createdAt',
description: 'usageLog.description',
eventKey: 'usageLog.eventKey',
executionId: 'usageLog.executionId',
id: 'usageLog.id',
metadata: 'usageLog.metadata',
source: 'usageLog.source',
userId: 'usageLog.userId',
workflowId: 'usageLog.workflowId',
workspaceId: 'usageLog.workspaceId',
},
}))
vi.mock('@/lib/billing/core/plan', () => ({
getHighestPrioritySubscription: mockGetHighestPrioritySubscription,
}))
vi.mock('@/lib/billing/subscriptions/utils', () => ({
isOrgScopedSubscription: mockIsOrgScopedSubscription,
}))
vi.mock('@/lib/core/config/feature-flags', () => ({
isBillingEnabled: true,
}))
import { recordUsage } from '@/lib/billing/core/usage-log'
describe('recordUsage', () => {
beforeEach(() => {
vi.clearAllMocks()
mockReturning.mockResolvedValue([{ cost: '0.10' }, { cost: '0.20' }])
mockOnConflictDoNothing.mockReturnValue({ returning: mockReturning })
mockValues.mockReturnValue({
onConflictDoNothing: mockOnConflictDoNothing,
})
mockInsert.mockReturnValue({ values: mockValues })
mockGetHighestPrioritySubscription.mockResolvedValue({
periodEnd: new Date('2026-06-01T00:00:00.000Z'),
periodStart: new Date('2026-05-01T00:00:00.000Z'),
referenceId: 'org-1',
})
mockIsOrgScopedSubscription.mockReturnValue(true)
})
it('commits canonical usage rows with deterministic event keys and billing scope', async () => {
await recordUsage({
userId: 'user-1',
workspaceId: 'workspace-1',
workflowId: 'workflow-1',
executionId: 'execution-1',
entries: [
{ category: 'fixed', source: 'workflow', description: 'execution_fee', cost: 0.1 },
{
category: 'model',
source: 'workflow',
description: 'gpt-4',
cost: 0.2,
metadata: { inputTokens: 10, outputTokens: 20 },
},
],
})
const values = mockValues.mock.calls[0][0]
expect(values).toHaveLength(2)
expect(values[0]).toMatchObject({
billingEntityId: 'org-1',
billingEntityType: 'organization',
billingPeriodEnd: new Date('2026-06-01T00:00:00.000Z'),
billingPeriodStart: new Date('2026-05-01T00:00:00.000Z'),
})
expect(values[0].eventKey).toMatch(/^[a-f0-9]{64}$/)
expect(values[1].eventKey).toMatch(/^[a-f0-9]{64}$/)
expect(values[0].eventKey).not.toBe(values[1].eventKey)
expect(mockOnConflictDoNothing).toHaveBeenCalledWith(
expect.objectContaining({ target: 'usageLog.eventKey' })
)
})
it('uses pre-resolved billing context without loading subscriptions', async () => {
await recordUsage({
userId: 'user-1',
billingEntity: { type: 'user', id: 'user-1' },
billingPeriod: {
start: new Date('2026-05-01T00:00:00.000Z'),
end: new Date('2026-06-01T00:00:00.000Z'),
},
entries: [{ category: 'fixed', source: 'workflow', description: 'execution_fee', cost: 0.1 }],
})
expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled()
expect(mockValues.mock.calls[0][0][0]).toMatchObject({
billingEntityId: 'user-1',
billingEntityType: 'user',
})
})
})
apps/sim/lib/billing/core/usage-log.ts
+164 -70
View File
@@ -1,9 +1,13 @@
import { createHash } from 'node:crypto'
import { db } from '@sim/db'
import { usageLog, userStats } from '@sim/db/schema'
import { usageLog } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { and, desc, eq, gte, lte, type SQL, sql } from 'drizzle-orm'
import { and, desc, eq, gte, lte, sql } from 'drizzle-orm'
import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
import { getHighestPrioritySubscription } from '@/lib/billing/core/plan'
import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
import { isBillingEnabled } from '@/lib/core/config/feature-flags'
const logger = createLogger('UsageLog')
@@ -41,6 +45,13 @@ export interface ModelUsageMetadata {
*/
export type UsageLogMetadata = ModelUsageMetadata | Record<string, unknown> | null
export type BillingEntityType = 'user' | 'organization'
interface BillingEntity {
type: BillingEntityType
id: string
}
/**
* A single usage entry to be recorded in the usage_log table.
*/
@@ -49,6 +60,8 @@ interface UsageEntry {
source: UsageLogSource
description: string
cost: number
eventKey?: string
sourceReference?: string
metadata?: UsageLogMetadata
}
@@ -67,99 +80,180 @@ export interface RecordUsageParams {
workflowId?: string
/** Execution context */
executionId?: string
/** Source-specific counter increments (e.g. totalCopilotCalls, totalManualExecutions) */
additionalStats?: Record<string, SQL>
/** Billing entity scope, resolved by caller when already known. */
billingEntity?: BillingEntity
/** Billing period bounds, resolved by caller when already known. */
billingPeriod?: { start: Date; end: Date }
}
function stableEventKey(parts: Record<string, unknown>): string {
const payload = Object.keys(parts)
.sort()
.map((key) => `${key}:${String(parts[key] ?? '')}`)
.join('|')
return createHash('sha256').update(payload).digest('hex')
}
async function resolveBillingContext(
userId: string,
billingEntity?: BillingEntity,
billingPeriod?: { start: Date; end: Date }
): Promise<{
billingEntity: BillingEntity
billingPeriod: { start: Date; end: Date }
}> {
if (billingEntity && billingPeriod) {
return { billingEntity, billingPeriod }
}
const subscription = await getHighestPrioritySubscription(userId)
const resolvedEntity =
billingEntity ??
(subscription && isOrgScopedSubscription(subscription, userId)
? { type: 'organization' as const, id: subscription.referenceId }
: { type: 'user' as const, id: userId })
const resolvedPeriod =
billingPeriod ??
(subscription?.periodStart && subscription.periodEnd
? { start: subscription.periodStart, end: subscription.periodEnd }
: defaultBillingPeriod())
return { billingEntity: resolvedEntity, billingPeriod: resolvedPeriod }
}
/**
* Records usage by inserting into usage_log and incrementing userStats counters.
* Returns post-cutover usage for an attributed billing entity/period.
* Legacy pre-cutover usage remains in userStats as a baseline until reset.
*/
export async function getBillingPeriodUsageCost(
billingEntity: BillingEntity,
billingPeriod: { start: Date; end: Date }
): Promise<number> {
const [row] = await db
.select({
cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
})
.from(usageLog)
.where(
and(
eq(usageLog.billingEntityType, billingEntity.type),
eq(usageLog.billingEntityId, billingEntity.id),
eq(usageLog.billingPeriodStart, billingPeriod.start),
eq(usageLog.billingPeriodEnd, billingPeriod.end)
)
)
return Number.parseFloat(row?.cost ?? '0')
}
export async function getBillingPeriodUsageCostByUser(
billingEntity: BillingEntity,
billingPeriod: { start: Date; end: Date }
): Promise<Map<string, number>> {
const rows = await db
.select({
userId: usageLog.userId,
cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
})
.from(usageLog)
.where(
and(
eq(usageLog.billingEntityType, billingEntity.type),
eq(usageLog.billingEntityId, billingEntity.id),
eq(usageLog.billingPeriodStart, billingPeriod.start),
eq(usageLog.billingPeriodEnd, billingPeriod.end)
)
)
.groupBy(usageLog.userId)
return new Map(rows.map((row) => [row.userId, Number.parseFloat(row.cost ?? '0')]))
}
/**
* Records usage as append-only billing events.
*
* The two writes are intentionally not wrapped in a transaction: under high
* concurrency for the same userId, holding BEGIN/COMMIT across the user_stats
* row-lock wait pins pgbouncer connections and exhausts the pool.
*
* usage_log is the source of truth and the INSERT propagates errors to the
* caller. The userStats UPDATE is best-effort: failures (and missing-row
* cases) are logged as warnings and swallowed. Counter drift is acceptable
* here — the long-term plan is to derive counters from usage_log directly.
* Any drift warning in logs is a signal that needs investigation.
* This intentionally avoids per-event userStats updates: userStats is retained
* as the pre-cutover period baseline and for low-frequency billing trackers,
* but usage writes no longer contend on the user_stats row.
*/
export async function recordUsage(params: RecordUsageParams): Promise<void> {
if (!isBillingEnabled) {
return
}
const { userId, entries, workspaceId, workflowId, executionId, additionalStats } = params
const { userId, entries, workspaceId, workflowId, executionId, billingEntity, billingPeriod } =
params
const validEntries = entries.filter((e) => e.cost > 0)
const totalCost = validEntries.reduce((sum, e) => sum + e.cost, 0)
if (
validEntries.length === 0 &&
(!additionalStats || Object.keys(additionalStats).length === 0)
) {
if (validEntries.length === 0) {
return
}
const RESERVED_KEYS = new Set(['totalCost', 'currentPeriodCost', 'lastActive'])
const safeStats = additionalStats
? Object.fromEntries(Object.entries(additionalStats).filter(([k]) => !RESERVED_KEYS.has(k)))
: undefined
const context = await resolveBillingContext(userId, billingEntity, billingPeriod)
if (validEntries.length > 0) {
await db.insert(usageLog).values(
validEntries.map((entry) => ({
id: generateId(),
userId,
category: entry.category,
source: entry.source,
description: entry.description,
metadata: entry.metadata ?? null,
cost: entry.cost.toString(),
workspaceId: workspaceId ?? null,
workflowId: workflowId ?? null,
executionId: executionId ?? null,
}))
)
}
const insertedRows = await db
.insert(usageLog)
.values(
validEntries.map((entry, index) => {
const sourceReference =
entry.sourceReference ??
[executionId, workflowId, workspaceId, entry.source, entry.description, index]
.filter((part) => part !== undefined && part !== null && part !== '')
.join(':')
const eventKey =
entry.eventKey ??
stableEventKey({
userId,
source: entry.source,
category: entry.category,
description: entry.description,
sourceReference,
executionId,
workflowId,
workspaceId,
index,
})
const updateFields: Record<string, SQL | Date> = {
lastActive: new Date(),
...(totalCost > 0 && {
totalCost: sql`total_cost + ${totalCost}`,
currentPeriodCost: sql`current_period_cost + ${totalCost}`,
}),
...safeStats,
}
try {
const result = await db
.update(userStats)
.set(updateFields)
.where(eq(userStats.userId, userId))
.returning({ userId: userStats.userId })
if (result.length === 0) {
logger.warn('recordUsage: userStats row not found; counter increment dropped', {
userId,
totalCost,
hadEntries: validEntries.length > 0,
additionalStatsKeys: safeStats ? Object.keys(safeStats) : [],
return {
id: generateId(),
userId,
category: entry.category,
source: entry.source,
description: entry.description,
metadata: entry.metadata ?? null,
cost: entry.cost.toString(),
eventKey,
billingEntityType: context.billingEntity.type,
billingEntityId: context.billingEntity.id,
billingPeriodStart: context.billingPeriod.start,
billingPeriodEnd: context.billingPeriod.end,
workspaceId: workspaceId ?? null,
workflowId: workflowId ?? null,
executionId: executionId ?? null,
}
})
}
} catch (error) {
logger.warn('recordUsage: userStats update failed; counter increment dropped', {
error: toError(error).message,
)
.onConflictDoNothing({
target: usageLog.eventKey,
where: sql`${usageLog.eventKey} IS NOT NULL`,
})
.returning({ cost: usageLog.cost })
const insertedCost = insertedRows.reduce((sum, row) => sum + Number.parseFloat(row.cost), 0)
if (insertedRows.length < validEntries.length) {
logger.debug('Skipped duplicate usage events', {
userId,
totalCost,
hadEntries: validEntries.length > 0,
additionalStatsKeys: safeStats ? Object.keys(safeStats) : [],
attemptedEntries: validEntries.length,
insertedEntries: insertedRows.length,
})
}
logger.debug('Recorded usage', {
userId,
totalCost,
totalCost: insertedCost,
entryCount: validEntries.length,
sources: [...new Set(validEntries.map((e) => e.source))],
})
apps/sim/lib/billing/core/usage.ts
+51 -6
View File
@@ -10,10 +10,12 @@ import {
renderUsageThresholdEmail,
} from '@/components/emails'
import { getEffectiveBillingStatus } from '@/lib/billing/core/access'
import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
import {
getHighestPrioritySubscription,
type HighestPrioritySubscription,
} from '@/lib/billing/core/plan'
import { getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log'
import {
computeDailyRefreshConsumed,
getOrgMemberRefreshBounds,
@@ -56,24 +58,31 @@ export interface OrgUsageLimitResult {
*/
export async function getPooledOrgCurrentPeriodCost(
organizationId: string
): Promise<{ memberIds: string[]; currentPeriodCost: number }> {
): Promise<{ memberIds: string[]; currentPeriodCost: number; lastPeriodCost: number }> {
const rows = await db
.select({
userId: member.userId,
currentPeriodCost: userStats.currentPeriodCost,
lastPeriodCost: userStats.lastPeriodCost,
})
.from(member)
.leftJoin(userStats, eq(member.userId, userStats.userId))
.where(eq(member.organizationId, organizationId))
let pooled = new Decimal(0)
let lastPeriodCost = new Decimal(0)
const memberIds: string[] = []
for (const row of rows) {
memberIds.push(row.userId)
pooled = pooled.plus(toDecimal(row.currentPeriodCost))
lastPeriodCost = lastPeriodCost.plus(toDecimal(row.lastPeriodCost))
}
return { memberIds, currentPeriodCost: toNumber(pooled) }
return {
memberIds,
currentPeriodCost: toNumber(pooled),
lastPeriodCost: toNumber(lastPeriodCost),
}
}
/**
@@ -186,8 +195,17 @@ export async function getUserUsageData(userId: string): Promise<UsageData> {
const stats = userStatsData[0]
const orgScoped = isOrgScopedSubscription(subscription, userId)
const billingPeriod =
subscription?.periodStart && subscription.periodEnd
? { start: subscription.periodStart, end: subscription.periodEnd }
: defaultBillingPeriod()
let currentUsageDecimal = toDecimal(stats.currentPeriodCost)
if (!orgScoped) {
currentUsageDecimal = currentUsageDecimal.plus(
await getBillingPeriodUsageCost({ type: 'user', id: userId }, billingPeriod)
)
}
// For personally-scoped Pro users, include any snapshotted usage from
// a prior org-join so the display reflects their total Pro usage.
@@ -204,6 +222,7 @@ export async function getUserUsageData(userId: string): Promise<UsageData> {
}
}
let currentUsage = toNumber(currentUsageDecimal)
let lastPeriodCost = toNumber(toDecimal(stats.lastPeriodCost))
let limit: number
// Shared between the pooled-usage and pooled-refresh blocks so we
@@ -220,7 +239,12 @@ export async function getUserUsageData(userId: string): Promise<UsageData> {
const pooled = await getPooledOrgCurrentPeriodCost(subscription.referenceId)
orgMemberIds = pooled.memberIds
currentUsage = pooled.currentPeriodCost
lastPeriodCost = pooled.lastPeriodCost
const ledgerUsage = await getBillingPeriodUsageCost(
{ type: 'organization', id: subscription.referenceId },
billingPeriod
)
currentUsage = pooled.currentPeriodCost + ledgerUsage
} else {
limit = stats.currentUsageLimit
? toNumber(toDecimal(stats.currentUsageLimit))
@@ -247,6 +271,7 @@ export async function getUserUsageData(userId: string): Promise<UsageData> {
planDollars,
seats: subscription.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
billingEntity: { type: 'organization', id: subscription.referenceId },
})
}
} else {
@@ -255,6 +280,7 @@ export async function getUserUsageData(userId: string): Promise<UsageData> {
periodStart: billingPeriodStart,
periodEnd: billingPeriodEnd,
planDollars,
billingEntity: { type: 'user', id: userId },
})
}
}
@@ -273,7 +299,7 @@ export async function getUserUsageData(userId: string): Promise<UsageData> {
isExceeded,
billingPeriodStart,
billingPeriodEnd,
lastPeriodCost: toNumber(toDecimal(stats.lastPeriodCost)),
lastPeriodCost,
}
} catch (error) {
logger.error('Failed to get user usage data', { userId, error })
@@ -642,7 +668,16 @@ export async function getEffectiveCurrentPeriodCost(userId: string): Promise<num
const pooled = await getPooledOrgCurrentPeriodCost(subscription.referenceId)
if (pooled.memberIds.length === 0) return 0
refreshUserIds = pooled.memberIds
rawCost = pooled.currentPeriodCost
const billingPeriod =
subscription.periodStart && subscription.periodEnd
? { start: subscription.periodStart, end: subscription.periodEnd }
: defaultBillingPeriod()
rawCost =
pooled.currentPeriodCost +
(await getBillingPeriodUsageCost(
{ type: 'organization', id: subscription.referenceId },
billingPeriod
))
} else {
const rows = await db
.select({ current: userStats.currentPeriodCost })
@@ -651,7 +686,13 @@ export async function getEffectiveCurrentPeriodCost(userId: string): Promise<num
.limit(1)
if (rows.length === 0) return 0
rawCost = toNumber(toDecimal(rows[0].current))
const billingPeriod =
subscription?.periodStart && subscription.periodEnd
? { start: subscription.periodStart, end: subscription.periodEnd }
: defaultBillingPeriod()
rawCost =
toNumber(toDecimal(rows[0].current)) +
(await getBillingPeriodUsageCost({ type: 'user', id: userId }, billingPeriod))
}
if (!subscription || !isPaid(subscription.plan) || !subscription.periodStart) {
@@ -673,6 +714,10 @@ export async function getEffectiveCurrentPeriodCost(userId: string): Promise<num
planDollars,
seats: subscription.seats || 1,
userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
billingEntity:
orgScoped && subscription
? { type: 'organization', id: subscription.referenceId }
: { type: 'user', id: userId },
})
return Math.max(0, rawCost - refreshConsumed)
apps/sim/lib/billing/credits/daily-refresh.ts
+19 -1
View File
@@ -48,8 +48,17 @@ export async function computeDailyRefreshConsumed(params: {
planDollars: number
seats?: number
userBounds?: Record<string, PerUserBounds>
billingEntity?: { type: 'user' | 'organization'; id: string }
}): Promise<number> {
const { userIds, periodStart, periodEnd, planDollars, seats = 1, userBounds } = params
const {
userIds,
periodStart,
periodEnd,
planDollars,
seats = 1,
userBounds,
billingEntity,
} = params
if (planDollars <= 0 || userIds.length === 0) return 0
@@ -64,6 +73,13 @@ export async function computeDailyRefreshConsumed(params: {
if (dayCount <= 0) return 0
const unboundedUsers = userBounds ? userIds.filter((id) => !(id in userBounds)) : userIds
const billingEntityFilter = billingEntity
? and(
eq(usageLog.billingEntityType, billingEntity.type),
eq(usageLog.billingEntityId, billingEntity.id),
eq(usageLog.billingPeriodStart, periodStart)
)
: undefined
const boundedClauses = userBounds
? Object.entries(userBounds).flatMap(([userId, bounds]) => {
@@ -75,6 +91,7 @@ export async function computeDailyRefreshConsumed(params: {
return [
and(
eq(usageLog.userId, userId),
billingEntityFilter,
gte(usageLog.createdAt, effectiveStart),
lt(usageLog.createdAt, effectiveEnd)
),
@@ -87,6 +104,7 @@ export async function computeDailyRefreshConsumed(params: {
? [
and(
inArray(usageLog.userId, unboundedUsers),
billingEntityFilter,
gte(usageLog.createdAt, periodStart),
lt(usageLog.createdAt, cap)
),
apps/sim/lib/billing/threshold-billing.test.ts
+7
View File
@@ -11,6 +11,7 @@ const {
mockEnqueueOutboxEvent,
mockGetEffectiveBillingStatus,
mockGetHighestPrioritySubscription,
mockGetBillingPeriodUsageCost,
mockGetOrganizationSubscriptionUsable,
mockHasUsableSubscriptionAccess,
mockIsEnterprise,
@@ -29,6 +30,7 @@ const {
mockEnqueueOutboxEvent: vi.fn(),
mockGetEffectiveBillingStatus: vi.fn(),
mockGetHighestPrioritySubscription: vi.fn(),
mockGetBillingPeriodUsageCost: vi.fn(),
mockGetOrganizationSubscriptionUsable: vi.fn(),
mockHasUsableSubscriptionAccess: vi.fn(),
mockIsEnterprise: vi.fn(),
@@ -89,6 +91,10 @@ vi.mock('@/lib/billing/core/subscription', () => ({
getOrganizationSubscriptionUsable: mockGetOrganizationSubscriptionUsable,
}))
vi.mock('@/lib/billing/core/usage-log', () => ({
getBillingPeriodUsageCost: mockGetBillingPeriodUsageCost,
}))
vi.mock('@/lib/billing/plan-helpers', () => ({
isEnterprise: mockIsEnterprise,
isFree: mockIsFree,
@@ -222,6 +228,7 @@ describe('checkAndBillOverageThreshold', () => {
mockIsFree.mockReturnValue(false)
mockIsEnterprise.mockReturnValue(false)
mockIsOrgScopedSubscription.mockReturnValue(false)
mockGetBillingPeriodUsageCost.mockResolvedValue(0)
mockDbSelect.mockImplementation(() => buildPersonalSelectChain())
mockTxSelect.mockImplementation(() => buildStatsSelectChain())
mockTxUpdate.mockImplementation(() => buildUpdateChain())
apps/sim/lib/billing/threshold-billing.ts
+16 -3
View File
@@ -9,6 +9,7 @@ import {
getHighestPrioritySubscription,
getOrganizationSubscriptionUsable,
} from '@/lib/billing/core/subscription'
import { getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log'
import { isEnterprise, isFree } from '@/lib/billing/plan-helpers'
import {
hasUsableSubscriptionAccess,
@@ -315,6 +316,14 @@ async function checkAndBillOrganizationOverageThreshold(organizationId: string):
ownerId: usageSnapshot.ownerId,
})
const ledgerUsage =
orgSubscription.periodStart && orgSubscription.periodEnd
? await getBillingPeriodUsageCost(
{ type: 'organization', id: organizationId },
{ start: orgSubscription.periodStart, end: orgSubscription.periodEnd }
)
: 0
const {
totalOverage: currentOverage,
baseSubscriptionAmount: basePrice,
@@ -325,7 +334,7 @@ async function checkAndBillOrganizationOverageThreshold(organizationId: string):
periodStart: orgSubscription.periodStart ?? null,
periodEnd: orgSubscription.periodEnd ?? null,
organizationId,
pooledCurrentPeriodCost: usageSnapshot.pooledCurrentPeriodCost,
pooledCurrentPeriodCost: usageSnapshot.pooledCurrentPeriodCost + ledgerUsage,
departedMemberUsage: usageSnapshot.departedMemberUsage,
memberIds: usageSnapshot.memberIds,
})
@@ -333,7 +342,9 @@ async function checkAndBillOrganizationOverageThreshold(organizationId: string):
if (currentOverage < threshold) {
logger.debug('Organization threshold billing check below threshold before locking', {
organizationId,
totalTeamUsage: usageSnapshot.pooledCurrentPeriodCost + usageSnapshot.departedMemberUsage,
totalTeamUsage:
usageSnapshot.pooledCurrentPeriodCost + ledgerUsage + usageSnapshot.departedMemberUsage,
ledgerUsage,
effectiveTeamUsage,
basePrice,
currentOverage,
@@ -433,7 +444,9 @@ async function checkAndBillOrganizationOverageThreshold(organizationId: string):
logger.debug('Organization threshold billing check', {
organizationId,
totalTeamUsage: usageSnapshot.pooledCurrentPeriodCost + usageSnapshot.departedMemberUsage,
totalTeamUsage:
usageSnapshot.pooledCurrentPeriodCost + ledgerUsage + usageSnapshot.departedMemberUsage,
ledgerUsage,
effectiveTeamUsage,
basePrice,
currentOverage,
apps/sim/lib/billing/webhooks/invoices.test.ts
+4
View File
@@ -34,6 +34,10 @@ vi.mock('@/lib/billing/core/billing', () => ({
isSubscriptionOrgScoped: vi.fn().mockResolvedValue(true),
}))
vi.mock('@/lib/billing/core/usage-log', () => ({
getBillingPeriodUsageCostByUser: vi.fn().mockResolvedValue(new Map()),
}))
vi.mock('@/lib/billing/credits/balance', () => ({
addCredits: vi.fn(),
getCreditBalance: vi.fn(),
apps/sim/lib/billing/webhooks/invoices.ts
+102 -12
View File
@@ -13,6 +13,7 @@ import type Stripe from 'stripe'
import { getEmailSubject, PaymentFailedEmail, renderCreditPurchaseEmail } from '@/components/emails'
import { BILLING_LOCK_TIMEOUT_MS } from '@/lib/billing/constants'
import { calculateSubscriptionOverage, isSubscriptionOrgScoped } from '@/lib/billing/core/billing'
import { getBillingPeriodUsageCostByUser } from '@/lib/billing/core/usage-log'
import { addCredits, getCreditBalanceForEntity } from '@/lib/billing/credits/balance'
import { setUsageLimitForCredits } from '@/lib/billing/credits/purchase'
import { blockOrgMembers, unblockOrgMembers } from '@/lib/billing/organizations/membership'
@@ -29,6 +30,26 @@ import { quickValidateEmail } from '@/lib/messaging/email/validation'
const logger = createLogger('StripeInvoiceWebhooks')
function getSubscriptionLinePeriod(
invoice: Stripe.Invoice,
stripeSubscriptionId: string
): { periodStart: Date; periodEnd: Date } | null {
const subscriptionLine = invoice.lines?.data?.find(
(line) =>
line.parent?.type === 'subscription_item_details' &&
line.parent.subscription_item_details?.subscription === stripeSubscriptionId
)
if (!subscriptionLine?.period?.start || !subscriptionLine.period.end) {
return null
}
return {
periodStart: new Date(subscriptionLine.period.start * 1000),
periodEnd: new Date(subscriptionLine.period.end * 1000),
}
}
const METADATA_SUBSCRIPTION_INVOICE_TYPES = new Set<string>([
'overage_billing',
'overage_threshold_billing',
@@ -387,8 +408,23 @@ export async function getBilledOverageForSubscription(sub: {
: 0
}
export async function resetUsageForSubscription(sub: { plan: string | null; referenceId: string }) {
export async function resetUsageForSubscription(sub: {
plan: string | null
referenceId: string
periodStart?: Date | null
periodEnd?: Date | null
}) {
const billingPeriod =
sub.periodStart && sub.periodEnd ? { start: sub.periodStart, end: sub.periodEnd } : null
if (await isSubscriptionOrgScoped(sub)) {
const ledgerUsageByUser = billingPeriod
? await getBillingPeriodUsageCostByUser(
{ type: 'organization', id: sub.referenceId },
billingPeriod
)
: new Map<string, number>()
await db.transaction(async (tx) => {
await tx.execute(sql.raw(`SET LOCAL lock_timeout = '${BILLING_LOCK_TIMEOUT_MS}ms'`))
@@ -441,6 +477,15 @@ export async function resetUsageForSubscription(sub: { plan: string | null; refe
return
}
const lastCostByUser = sql.join(
memberStatsRows.map((row) => {
const baseline = toNumber(toDecimal(row.current))
const ledgerUsage = ledgerUsageByUser.get(row.userId) ?? 0
const lastPeriodCost = (baseline + ledgerUsage).toString()
return sql`WHEN ${row.userId} THEN ${lastPeriodCost}`
}),
sql` `
)
const currentCostByUser = sql.join(
memberStatsRows.map((row) => sql`WHEN ${row.userId} THEN ${row.current ?? '0'}`),
sql` `
@@ -449,13 +494,14 @@ export async function resetUsageForSubscription(sub: { plan: string | null; refe
memberStatsRows.map((row) => sql`WHEN ${row.userId} THEN ${row.currentCopilot ?? '0'}`),
sql` `
)
const capturedLastCost = sql`CASE ${userStats.userId} ${lastCostByUser} ELSE '0' END`
const capturedCurrentCost = sql`CASE ${userStats.userId} ${currentCostByUser} ELSE '0' END`
const capturedCurrentCopilotCost = sql`CASE ${userStats.userId} ${currentCopilotCostByUser} ELSE '0' END`
await tx
.update(userStats)
.set({
lastPeriodCost: capturedCurrentCost,
lastPeriodCost: capturedLastCost,
lastPeriodCopilotCost: capturedCurrentCopilotCost,
currentPeriodCost: sql`GREATEST(0, ${userStats.currentPeriodCost} - (${capturedCurrentCost})::decimal)`,
currentPeriodCopilotCost: sql`GREATEST(0, ${userStats.currentPeriodCopilotCost} - (${capturedCurrentCopilotCost})::decimal)`,
@@ -483,6 +529,13 @@ export async function resetUsageForSubscription(sub: { plan: string | null; refe
const current = currentStats[0].current || '0'
const snapshot = toNumber(toDecimal(currentStats[0].snapshot))
const currentCopilot = currentStats[0].currentCopilot || '0'
const ledgerUsage = billingPeriod
? await getBillingPeriodUsageCostByUser(
{ type: 'user', id: sub.referenceId },
billingPeriod
)
: new Map<string, number>()
const userLedgerUsage = ledgerUsage.get(sub.referenceId) ?? 0
// Snapshot > 0: user joined a paid org mid-cycle. The pre-join
// portion was billed on this invoice (snapshot); `currentPeriodCost`
@@ -492,7 +545,7 @@ export async function resetUsageForSubscription(sub: { plan: string | null; refe
await db
.update(userStats)
.set({
lastPeriodCost: snapshot.toString(),
lastPeriodCost: (snapshot + userLedgerUsage).toString(),
lastPeriodCopilotCost: '0',
proPeriodCostSnapshot: '0',
proPeriodCostSnapshotAt: null,
@@ -500,7 +553,11 @@ export async function resetUsageForSubscription(sub: { plan: string | null; refe
})
.where(eq(userStats.userId, sub.referenceId))
} else {
const totalLastPeriod = toNumber(toDecimal(current).plus(snapshot)).toString()
const totalLastPeriod = (
toNumber(toDecimal(current)) +
snapshot +
userLedgerUsage
).toString()
// Delta-reset for the same reason as the org branch above.
await db
.update(userStats)
@@ -744,7 +801,16 @@ export async function handleInvoicePaymentSucceeded(event: Stripe.Event) {
}
if (wasBlocked && !isProrationInvoice) {
await resetUsageForSubscription({ plan: sub.plan, referenceId: sub.referenceId })
const invoicePeriod = getSubscriptionLinePeriod(
invoice,
resolvedInvoice.stripeSubscriptionId
)
await resetUsageForSubscription({
plan: sub.plan,
referenceId: sub.referenceId,
periodStart: invoicePeriod?.periodStart ?? null,
periodEnd: invoicePeriod?.periodEnd ?? null,
})
}
}
)
@@ -885,8 +951,25 @@ export async function handleInvoiceFinalized(event: Stripe.Event) {
if (records.length === 0) return
const sub = records[0]
const invoicePeriod = getSubscriptionLinePeriod(invoice, stripeSubscriptionId)
if (!invoicePeriod) {
logger.error('Missing subscription line period on subscription cycle invoice', {
invoiceId: invoice.id,
stripeSubscriptionId,
})
if (isEnterprise(sub.plan)) {
await resetUsageForSubscription({ plan: sub.plan, referenceId: sub.referenceId })
}
return
}
if (isEnterprise(sub.plan)) {
await resetUsageForSubscription({ plan: sub.plan, referenceId: sub.referenceId })
await resetUsageForSubscription({
plan: sub.plan,
referenceId: sub.referenceId,
periodStart: invoicePeriod.periodStart,
periodEnd: invoicePeriod.periodEnd,
})
return
}
@@ -895,13 +978,15 @@ export async function handleInvoiceFinalized(event: Stripe.Event) {
event.id,
async () => {
const stripe = requireStripeClient()
const periodEnd =
invoice.lines?.data?.[0]?.period?.end ||
invoice.period_end ||
Math.floor(Date.now() / 1000)
const periodStart = Math.floor(invoicePeriod.periodStart.getTime() / 1000)
const periodEnd = Math.floor(invoicePeriod.periodEnd.getTime() / 1000)
const billingPeriod = new Date(periodEnd * 1000).toISOString().slice(0, 7)
const totalOverage = await calculateSubscriptionOverage(sub)
const totalOverage = await calculateSubscriptionOverage({
...sub,
periodStart: new Date(periodStart * 1000),
periodEnd: new Date(periodEnd * 1000),
})
const entityType = (await isSubscriptionOrgScoped(sub)) ? 'organization' : 'user'
const entityId = sub.referenceId
@@ -1089,7 +1174,12 @@ export async function handleInvoiceFinalized(event: Stripe.Event) {
// rolls `currentPeriodCost` forward by delta. Idempotent on its
// own (delta subtraction of a value that's already been
// subtracted is a no-op).
await resetUsageForSubscription({ plan: sub.plan, referenceId: sub.referenceId })
await resetUsageForSubscription({
plan: sub.plan,
referenceId: sub.referenceId,
periodStart: invoicePeriod.periodStart,
periodEnd: invoicePeriod.periodEnd,
})
return { totalOverage, creditsApplied, amountToBillStripe }
}
apps/sim/lib/billing/webhooks/subscription.ts
+10
View File
@@ -204,6 +204,8 @@ export async function handleSubscriptionCreated(subscriptionData: {
referenceId: string
plan: string | null
status: string
periodStart?: Date | null
periodEnd?: Date | null
}) {
try {
const otherActiveSubscriptions = await db
@@ -230,6 +232,8 @@ export async function handleSubscriptionCreated(subscriptionData: {
await resetUsageForSubscription({
plan: subscriptionData.plan,
referenceId: subscriptionData.referenceId,
periodStart: subscriptionData.periodStart ?? null,
periodEnd: subscriptionData.periodEnd ?? null,
})
logger.info('Successfully reset usage for free -> paid transition', {
@@ -280,6 +284,8 @@ export async function handleSubscriptionDeleted(
referenceId: string
stripeSubscriptionId: string | null
seats?: number | null
periodStart?: Date | null
periodEnd?: Date | null
},
stripeEventId?: string
) {
@@ -308,6 +314,8 @@ export async function handleSubscriptionDeleted(
await resetUsageForSubscription({
plan: subscription.plan,
referenceId: subscription.referenceId,
periodStart: subscription.periodStart ?? null,
periodEnd: subscription.periodEnd ?? null,
})
const dormantResult = await transitionOrganizationToDormantState(
@@ -412,6 +420,8 @@ export async function handleSubscriptionDeleted(
await resetUsageForSubscription({
plan: subscription.plan,
referenceId: subscription.referenceId,
periodStart: subscription.periodStart ?? null,
periodEnd: subscription.periodEnd ?? null,
})
let restoredProCount = 0
apps/sim/lib/knowledge/documents/service.ts
+1 -3
View File
@@ -702,12 +702,10 @@ export async function processDocumentAsync(
source: 'knowledge-base',
description: embeddingModelName,
cost,
sourceReference: `knowledge-document:${documentId}:${startTime}`,
metadata: { inputTokens: totalEmbeddingTokens, outputTokens: 0 },
},
],
additionalStats: {
totalTokensUsed: sql`total_tokens_used + ${totalEmbeddingTokens}`,
},
})
await checkAndBillOverageThreshold(billingUserId)
} else {
apps/sim/lib/logs/execution/logger.ts
-21
View File
@@ -41,17 +41,6 @@ import type {
} from '@/lib/logs/types'
import type { SerializableExecutionState } from '@/executor/execution/types'
/** Maps execution trigger types to their corresponding userStats counter columns */
const TRIGGER_COUNTER_MAP: Record<string, { key: string; column: string }> = {
manual: { key: 'totalManualExecutions', column: 'total_manual_executions' },
api: { key: 'totalApiCalls', column: 'total_api_calls' },
webhook: { key: 'totalWebhookTriggers', column: 'total_webhook_triggers' },
schedule: { key: 'totalScheduledExecutions', column: 'total_scheduled_executions' },
chat: { key: 'totalChatExecutions', column: 'total_chat_executions' },
mcp: { key: 'totalMcpExecutions', column: 'total_mcp_executions' },
a2a: { key: 'totalA2aExecutions', column: 'total_a2a_executions' },
} as const
const logger = createLogger('ExecutionLogger')
const MAX_EXECUTION_DATA_BYTES = 500 * 1024
const MAX_TRACE_IO_BYTES = 8 * 1024
@@ -1069,22 +1058,12 @@ export class ExecutionLogger implements IExecutionLoggerService {
}
}
const additionalStats: Record<string, ReturnType<typeof sql>> = {
totalTokensUsed: sql`total_tokens_used + ${costSummary.totalTokens}`,
}
const triggerCounter = TRIGGER_COUNTER_MAP[trigger]
if (triggerCounter) {
additionalStats[triggerCounter.key] = sql`${sql.raw(triggerCounter.column)} + 1`
}
await recordUsage({
userId,
entries,
workspaceId: workflowRecord.workspaceId ?? undefined,
workflowId,
executionId,
additionalStats,
})
// Check if user has hit overage threshold and bill incrementally
apps/sim/lib/workflows/utils.ts
+1 -49
View File
@@ -1,5 +1,5 @@
import { db } from '@sim/db'
import { permissions, userStats, workflowFolder, workflow as workflowTable } from '@sim/db/schema'
import { permissions, workflowFolder, workflow as workflowTable } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { generateId } from '@sim/utils/id'
import { authorizeWorkflowByWorkspacePermission } from '@sim/workflow-authz'
@@ -11,7 +11,6 @@ import type { ExecutionMaterializationContext } from '@/lib/execution/payloads/m
import { getNextWorkflowColor } from '@/lib/workflows/colors'
import { buildDefaultWorkflowArtifacts } from '@/lib/workflows/defaults'
import { saveWorkflowToNormalizedTables } from '@/lib/workflows/persistence/utils'
import { getWorkspaceBilledAccountUserId } from '@/lib/workspaces/utils'
import type { ExecutionResult } from '@/executor/types'
const logger = createLogger('WorkflowUtils')
@@ -245,53 +244,6 @@ export async function updateWorkflowRunCounts(workflowId: string, runs = 1) {
})
.where(eq(workflowTable.id, workflowId))
let activityUserId: string | null = null
if (workflow.workspaceId) {
try {
activityUserId = await getWorkspaceBilledAccountUserId(workflow.workspaceId)
} catch (error) {
logger.warn(`Error resolving billed account for workspace ${workflow.workspaceId}`, {
workflowId,
error,
})
}
}
if (activityUserId) {
try {
const existing = await db
.select()
.from(userStats)
.where(eq(userStats.userId, activityUserId))
.limit(1)
if (existing.length === 0) {
logger.warn('User stats record not found - should be created during onboarding', {
userId: activityUserId,
workflowId,
})
} else {
await db
.update(userStats)
.set({
lastActive: new Date(),
})
.where(eq(userStats.userId, activityUserId))
}
} catch (error) {
logger.error(`Error updating userStats lastActive for userId ${activityUserId}:`, error)
// Don't rethrow - we want to continue even if this fails
}
} else {
logger.warn(
'Skipping userStats lastActive update: unable to resolve workspace billed account',
{
workflowId,
workspaceId: workflow.workspaceId,
}
)
}
return {
success: true,
runsAdded: runs,