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

* 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>

Author: icecrasher321

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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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 },
                 },
               ],