@@ -3,8 +3,7 @@ import { eq, sql } from 'drizzle-orm'
33import { type NextRequest , NextResponse } from 'next/server'
44import { z } from 'zod'
55import { checkInternalApiKey } from '@/lib/copilot/utils'
6- import { env } from '@/lib/env'
7- import { isBillingEnabled , isProd } from '@/lib/environment'
6+ import { isBillingEnabled } from '@/lib/environment'
87import { createLogger } from '@/lib/logs/console/logger'
98import { db } from '@/db'
109import { userStats } from '@/db/schema'
@@ -17,6 +16,7 @@ const UpdateCostSchema = z.object({
1716 input : z . number ( ) . min ( 0 , 'Input tokens must be a non-negative number' ) ,
1817 output : z . number ( ) . min ( 0 , 'Output tokens must be a non-negative number' ) ,
1918 model : z . string ( ) . min ( 1 , 'Model is required' ) ,
19+ multiplier : z . number ( ) . min ( 0 ) ,
2020} )
2121
2222/**
@@ -75,27 +75,27 @@ export async function POST(req: NextRequest) {
7575 )
7676 }
7777
78- const { userId, input, output, model } = validation . data
78+ const { userId, input, output, model, multiplier } = validation . data
7979
8080 logger . info ( `[${ requestId } ] Processing cost update` , {
8181 userId,
8282 input,
8383 output,
8484 model,
85+ multiplier,
8586 } )
8687
8788 const finalPromptTokens = input
8889 const finalCompletionTokens = output
8990 const totalTokens = input + output
9091
91- // Calculate cost using COPILOT_COST_MULTIPLIER (only in production, like normal executions)
92- const copilotMultiplier = isProd ? env . COPILOT_COST_MULTIPLIER || 1 : 1
92+ // Calculate cost using provided multiplier (required)
9393 const costResult = calculateCost (
9494 model ,
9595 finalPromptTokens ,
9696 finalCompletionTokens ,
9797 false ,
98- copilotMultiplier
98+ multiplier
9999 )
100100
101101 logger . info ( `[${ requestId } ] Cost calculation result` , {
@@ -104,7 +104,7 @@ export async function POST(req: NextRequest) {
104104 promptTokens : finalPromptTokens ,
105105 completionTokens : finalCompletionTokens ,
106106 totalTokens : totalTokens ,
107- copilotMultiplier ,
107+ multiplier ,
108108 costResult,
109109 } )
110110
@@ -127,6 +127,10 @@ export async function POST(req: NextRequest) {
127127 totalTokensUsed : totalTokens ,
128128 totalCost : costToStore . toString ( ) ,
129129 currentPeriodCost : costToStore . toString ( ) ,
130+ // Copilot usage tracking
131+ totalCopilotCost : costToStore . toString ( ) ,
132+ totalCopilotTokens : totalTokens ,
133+ totalCopilotCalls : 1 ,
130134 lastActive : new Date ( ) ,
131135 } )
132136
@@ -141,6 +145,10 @@ export async function POST(req: NextRequest) {
141145 totalTokensUsed : sql `total_tokens_used + ${ totalTokens } ` ,
142146 totalCost : sql `total_cost + ${ costToStore } ` ,
143147 currentPeriodCost : sql `current_period_cost + ${ costToStore } ` ,
148+ // Copilot usage tracking increments
149+ totalCopilotCost : sql `total_copilot_cost + ${ costToStore } ` ,
150+ totalCopilotTokens : sql `total_copilot_tokens + ${ totalTokens } ` ,
151+ totalCopilotCalls : sql `total_copilot_calls + 1` ,
144152 totalApiCalls : sql `total_api_calls` ,
145153 lastActive : new Date ( ) ,
146154 }
0 commit comments