Skip to content

Commit 5996a5e

Browse files
committed
improvement(billing): migrate hot path writes away from user_stats
1 parent 7ddd90b commit 5996a5e

24 files changed

Lines changed: 18052 additions & 292 deletions

File tree

apps/sim/app/api/billing/update-cost/route.ts

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import type { Span } from '@opentelemetry/api'
22
import { createLogger } from '@sim/logger'
33
import { toError } from '@sim/utils/errors'
4-
import { sql } from 'drizzle-orm'
54
import { type NextRequest, NextResponse } from 'next/server'
65
import { billingUpdateCostContract } from '@/lib/api/contracts/subscription'
76
import { parseRequest } from '@/lib/api/server'
@@ -155,21 +154,6 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
155154
source,
156155
})
157156

158-
const totalTokens = inputTokens + outputTokens
159-
160-
const additionalStats: Record<string, ReturnType<typeof sql>> = {
161-
totalCopilotCost: sql`total_copilot_cost + ${cost}`,
162-
currentPeriodCopilotCost: sql`current_period_copilot_cost + ${cost}`,
163-
totalCopilotCalls: sql`total_copilot_calls + 1`,
164-
totalCopilotTokens: sql`total_copilot_tokens + ${totalTokens}`,
165-
}
166-
167-
if (isMcp) {
168-
additionalStats.totalMcpCopilotCost = sql`total_mcp_copilot_cost + ${cost}`
169-
additionalStats.currentPeriodMcpCopilotCost = sql`current_period_mcp_copilot_cost + ${cost}`
170-
additionalStats.totalMcpCopilotCalls = sql`total_mcp_copilot_calls + 1`
171-
}
172-
173157
await recordUsage({
174158
userId,
175159
entries: [
@@ -178,10 +162,11 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
178162
source,
179163
description: model,
180164
cost,
165+
eventKey: idempotencyKey ? `update-cost:${idempotencyKey}` : undefined,
166+
sourceReference: idempotencyKey ? `update-cost:${idempotencyKey}` : requestId,
181167
metadata: { inputTokens, outputTokens },
182168
},
183169
],
184-
additionalStats,
185170
})
186171
usageCommitted = true
187172

apps/sim/app/api/mcp/copilot/route.ts

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,10 @@ import {
1010
McpError,
1111
type RequestId,
1212
} from '@modelcontextprotocol/sdk/types.js'
13-
import { db } from '@sim/db'
14-
import { userStats } from '@sim/db/schema'
1513
import { createLogger } from '@sim/logger'
1614
import { toError } from '@sim/utils/errors'
1715
import { generateId } from '@sim/utils/id'
1816
import { authorizeWorkflowByWorkspacePermission } from '@sim/workflow-authz'
19-
import { eq, sql } from 'drizzle-orm'
2017
import { type NextRequest, NextResponse } from 'next/server'
2118
import { mcpRequestBodySchema, mcpToolCallParamsSchema } from '@/lib/api/contracts/mcp'
2219
import { validateOAuthAccessToken } from '@/lib/auth/oauth-token'
@@ -391,20 +388,8 @@ export const DELETE = withRouteHandler(async (request: NextRequest) => {
391388
return NextResponse.json(createError(0, -32000, 'Method not allowed.'), { status: 405 })
392389
})
393390

394-
/**
395-
* Increment MCP copilot call counter in userStats (fire-and-forget).
396-
*/
397391
function trackMcpCopilotCall(userId: string): void {
398-
db.update(userStats)
399-
.set({
400-
totalMcpCopilotCalls: sql`total_mcp_copilot_calls + 1`,
401-
lastActive: new Date(),
402-
})
403-
.where(eq(userStats.userId, userId))
404-
.then(() => {})
405-
.catch((error) => {
406-
logger.error('Failed to track MCP copilot call', { error, userId })
407-
})
392+
logger.debug('MCP copilot call tracked via request logs', { userId })
408393
}
409394

410395
async function handleToolsCall(

apps/sim/app/api/speech/token/route.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { createHash } from 'node:crypto'
12
import { db } from '@sim/db'
23
import { chat } from '@sim/db/schema'
34
import { createLogger } from '@sim/logger'
@@ -31,6 +32,10 @@ const STT_TOKEN_RATE_LIMIT = {
3132
refillIntervalMs: 72 * 1000,
3233
} as const
3334

35+
function hashVoiceToken(token: string): string {
36+
return createHash('sha256').update(token).digest('hex')
37+
}
38+
3439
const rateLimiter = new RateLimiter()
3540

3641
async function validateChatAuth(
@@ -163,6 +168,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
163168
source: 'voice-input',
164169
description: `Voice input session (${maxMinutes} min)`,
165170
cost: sessionCost * getCostMultiplier(),
171+
sourceReference: `voice-input:${hashVoiceToken(data.token)}`,
166172
},
167173
],
168174
}).catch((err) => {

apps/sim/app/api/wand/route.ts

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { db } from '@sim/db'
22
import { workflow } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
4-
import { eq, sql } from 'drizzle-orm'
4+
import { eq } from 'drizzle-orm'
55
import { type NextRequest, NextResponse } from 'next/server'
66
import { wandGenerateContract } from '@/lib/api/contracts'
77
import { parseRequest } from '@/lib/api/server'
@@ -106,7 +106,6 @@ async function updateUserStatsForWand(
106106
}
107107

108108
try {
109-
const totalTokens = usage.total_tokens || 0
110109
const promptTokens = usage.prompt_tokens || 0
111110
const completionTokens = usage.completion_tokens || 0
112111

@@ -138,12 +137,10 @@ async function updateUserStatsForWand(
138137
source: 'wand',
139138
description: modelName,
140139
cost: costToStore,
140+
sourceReference: `wand:${requestId}`,
141141
metadata: { inputTokens: promptTokens, outputTokens: completionTokens },
142142
},
143143
],
144-
additionalStats: {
145-
totalTokensUsed: sql`total_tokens_used + ${totalTokens}`,
146-
},
147144
})
148145

149146
await checkAndBillOverageThreshold(billingUserId)

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,7 @@ async function runWorkflowAndWriteTerminal(
282282
source: 'enrichment',
283283
description: enrichment.name,
284284
cost,
285+
sourceReference: `enrichment:${tableId}:${rowId}:${enrichment.id}`,
285286
metadata: { enrichmentId: enrichment.id, tableId, rowId },
286287
},
287288
],

0 commit comments

Comments
 (0)