feat(copilot): add billing endpoint (#855)

* Add copilot billing

* Lint

* Update logic

* Dont count as api callg

Author: Sg312

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

@@ -0,0 +1,214 @@
+import crypto from 'crypto'
+import { eq, sql } from 'drizzle-orm'
+import { type NextRequest, NextResponse } from 'next/server'
+import { z } from 'zod'
+import { env } from '@/lib/env'
+import { isProd } from '@/lib/environment'
+import { createLogger } from '@/lib/logs/console/logger'
+import { db } from '@/db'
+import { userStats } from '@/db/schema'
+import { calculateCost } from '@/providers/utils'
+
+const logger = createLogger('billing-update-cost')
+
+// Schema for the request body
+const UpdateCostSchema = z.object({
+  userId: z.string().min(1, 'User ID is required'),
+  input: z.number().min(0, 'Input tokens must be a non-negative number'),
+  output: z.number().min(0, 'Output tokens must be a non-negative number'),
+  model: z.string().min(1, 'Model is required'),
+})
+
+// Authentication function (reused from copilot/methods route)
+function checkInternalApiKey(req: NextRequest) {
+  const apiKey = req.headers.get('x-api-key')
+  const expectedApiKey = env.INTERNAL_API_SECRET
+
+  if (!expectedApiKey) {
+    return { success: false, error: 'Internal API key not configured' }
+  }
+
+  if (!apiKey) {
+    return { success: false, error: 'API key required' }
+  }
+
+  if (apiKey !== expectedApiKey) {
+    return { success: false, error: 'Invalid API key' }
+  }
+
+  return { success: true }
+}
+
+/**
+ * POST /api/billing/update-cost
+ * Update user cost based on token usage with internal API key auth
+ */
+export async function POST(req: NextRequest) {
+  const requestId = crypto.randomUUID().slice(0, 8)
+  const startTime = Date.now()
+
+  try {
+    logger.info(`[${requestId}] Update cost request started`)
+
+    // Check authentication (internal API key)
+    const authResult = checkInternalApiKey(req)
+    if (!authResult.success) {
+      logger.warn(`[${requestId}] Authentication failed: ${authResult.error}`)
+      return NextResponse.json(
+        {
+          success: false,
+          error: authResult.error || 'Authentication failed',
+        },
+        { status: 401 }
+      )
+    }
+
+    // Parse and validate request body
+    const body = await req.json()
+    const validation = UpdateCostSchema.safeParse(body)
+
+    if (!validation.success) {
+      logger.warn(`[${requestId}] Invalid request body`, {
+        errors: validation.error.issues,
+        body,
+      })
+      return NextResponse.json(
+        {
+          success: false,
+          error: 'Invalid request body',
+          details: validation.error.issues,
+        },
+        { status: 400 }
+      )
+    }
+
+    const { userId, input, output, model } = validation.data
+
+    logger.info(`[${requestId}] Processing cost update`, {
+      userId,
+      input,
+      output,
+      model,
+    })
+
+    const finalPromptTokens = input
+    const finalCompletionTokens = output
+    const totalTokens = input + output
+
+    // Calculate cost using COPILOT_COST_MULTIPLIER (only in production, like normal executions)
+    const copilotMultiplier = isProd ? env.COPILOT_COST_MULTIPLIER || 1 : 1
+    const costResult = calculateCost(
+      model,
+      finalPromptTokens,
+      finalCompletionTokens,
+      false,
+      copilotMultiplier
+    )
+
+    logger.info(`[${requestId}] Cost calculation result`, {
+      userId,
+      model,
+      promptTokens: finalPromptTokens,
+      completionTokens: finalCompletionTokens,
+      totalTokens: totalTokens,
+      copilotMultiplier,
+      costResult,
+    })
+
+    // Follow the exact same logic as ExecutionLogger.updateUserStats but with direct userId
+    const costToStore = costResult.total // No additional multiplier needed since calculateCost already applied it
+
+    // Check if user stats record exists (same as ExecutionLogger)
+    const userStatsRecords = await db.select().from(userStats).where(eq(userStats.userId, userId))
+
+    if (userStatsRecords.length === 0) {
+      // Create new user stats record (same logic as ExecutionLogger)
+      await db.insert(userStats).values({
+        id: crypto.randomUUID(),
+        userId: userId,
+        totalManualExecutions: 0,
+        totalApiCalls: 0,
+        totalWebhookTriggers: 0,
+        totalScheduledExecutions: 0,
+        totalChatExecutions: 0,
+        totalTokensUsed: totalTokens,
+        totalCost: costToStore.toString(),
+        currentPeriodCost: costToStore.toString(),
+        lastActive: new Date(),
+      })
+
+      logger.info(`[${requestId}] Created new user stats record`, {
+        userId,
+        totalCost: costToStore,
+        totalTokens,
+      })
+    } else {
+      // Update existing user stats record (same logic as ExecutionLogger)
+      const updateFields = {
+        totalTokensUsed: sql`total_tokens_used + ${totalTokens}`,
+        totalCost: sql`total_cost + ${costToStore}`,
+        currentPeriodCost: sql`current_period_cost + ${costToStore}`,
+        totalApiCalls: sql`total_api_calls`,
+        lastActive: new Date(),
+      }
+
+      await db.update(userStats).set(updateFields).where(eq(userStats.userId, userId))
+
+      logger.info(`[${requestId}] Updated user stats record`, {
+        userId,
+        addedCost: costToStore,
+        addedTokens: totalTokens,
+      })
+    }
+
+    const duration = Date.now() - startTime
+
+    logger.info(`[${requestId}] Cost update completed successfully`, {
+      userId,
+      duration,
+      cost: costResult.total,
+      totalTokens,
+    })
+
+    return NextResponse.json({
+      success: true,
+      data: {
+        userId,
+        input,
+        output,
+        totalTokens,
+        model,
+        cost: {
+          input: costResult.input,
+          output: costResult.output,
+          total: costResult.total,
+        },
+        tokenBreakdown: {
+          prompt: finalPromptTokens,
+          completion: finalCompletionTokens,
+          total: totalTokens,
+        },
+        pricing: costResult.pricing,
+        processedAt: new Date().toISOString(),
+        requestId,
+      },
+    })
+  } catch (error) {
+    const duration = Date.now() - startTime
+
+    logger.error(`[${requestId}] Cost update failed`, {
+      error: error instanceof Error ? error.message : String(error),
+      stack: error instanceof Error ? error.stack : undefined,
+      duration,
+    })
+
+    return NextResponse.json(
+      {
+        success: false,
+        error: 'Internal server error',
+        requestId,
+      },
+      { status: 500 }
+    )
+  }
+}

apps/sim/lib/env.ts

@@ -67,6 +67,7 @@ export const env = createEnv({
     // Monitoring & Analytics
     TELEMETRY_ENDPOINT:                   z.string().url().optional(),          // Custom telemetry/analytics endpoint
     COST_MULTIPLIER:                      z.number().optional(),                // Multiplier for cost calculations
+    COPILOT_COST_MULTIPLIER:              z.number().optional(),                // Multiplier for copilot cost calculations
     SENTRY_ORG:                           z.string().optional(),                // Sentry organization for error tracking
     SENTRY_PROJECT:                       z.string().optional(),                // Sentry project for error tracking
     SENTRY_AUTH_TOKEN:                    z.string().optional(),                // Sentry authentication token

apps/sim/providers/utils.ts

@@ -429,13 +429,15 @@ export async function transformBlockTool(
  * @param promptTokens Number of prompt tokens used
  * @param completionTokens Number of completion tokens used
  * @param useCachedInput Whether to use cached input pricing (default: false)
+ * @param customMultiplier Optional custom multiplier to override the default cost multiplier
  * @returns Cost calculation results with input, output and total costs
  */
 export function calculateCost(
   model: string,
   promptTokens = 0,
   completionTokens = 0,
-  useCachedInput = false
+  useCachedInput = false,
+  customMultiplier?: number
 ) {
   // First check if it's an embedding model
   let pricing = getEmbeddingModelPricing(model)
@@ -472,7 +474,7 @@ export function calculateCost(
   const outputCost = completionTokens * (pricing.output / 1_000_000)
   const totalCost = inputCost + outputCost
 
-  const costMultiplier = getCostMultiplier()
+  const costMultiplier = customMultiplier ?? getCostMultiplier()
 
   const finalInputCost = inputCost * costMultiplier
   const finalOutputCost = outputCost * costMultiplier