pass logger into balance calculator

Author: charleslien

backend/src/__tests__/usage-calculation.test.ts

@@ -7,8 +7,16 @@ import {
 import { afterAll, beforeAll, describe, expect, it } from 'bun:test'
 
 import type { GrantType } from '@codebuff/common/db/schema'
+import type { Logger } from '@codebuff/types/logger'
 
 describe('Usage Calculation System', () => {
+  const logger: Logger = {
+    debug: () => {},
+    error: () => {},
+    info: () => {},
+    warn: () => {},
+  }
+
   beforeAll(() => {
     // Mock the database module before importing the function
     mockModule('@codebuff/common/db', () => ({
@@ -22,17 +30,6 @@ describe('Usage Calculation System', () => {
         }),
       },
     }))
-
-    // Mock logger
-    mockModule('@codebuff/common/util/logger', () => ({
-      logger: {
-        debug: () => {},
-        error: () => {},
-        info: () => {},
-        warn: () => {},
-      },
-      withLoggerContext: async (context: any, fn: () => Promise<any>) => fn(),
-    }))
   })
 
   afterAll(() => {
@@ -76,11 +73,12 @@ describe('Usage Calculation System', () => {
       },
     }))
 
-    const { usageThisCycle } = await calculateUsageAndBalance(
-      'test-user',
-      new Date('2024-01-01'),
-      new Date('2024-01-15'), // Pass current time when grants are active
-    )
+    const { usageThisCycle } = await calculateUsageAndBalance({
+      userId: 'test-user',
+      quotaResetDate: new Date('2024-01-01'),
+      now: new Date('2024-01-15'), // Pass current time when grants are active
+      logger,
+    })
 
     expect(usageThisCycle).toBe(400) // 200 + 200 = 400 total usage
   })
@@ -112,11 +110,12 @@ describe('Usage Calculation System', () => {
       },
     }))
 
-    const { balance, usageThisCycle } = await calculateUsageAndBalance(
-      'test-user',
-      new Date('2024-01-01'),
-      new Date('2024-01-16'), // Current time after expiry
-    )
+    const { balance, usageThisCycle } = await calculateUsageAndBalance({
+      userId: 'test-user',
+      quotaResetDate: new Date('2024-01-01'),
+      now: new Date('2024-01-16'), // Current time after expiry
+      logger,
+    })
 
     expect(balance.totalRemaining).toBe(0) // Expired grant doesn't count
     expect(balance.totalDebt).toBe(0)
@@ -158,11 +157,12 @@ describe('Usage Calculation System', () => {
       },
     }))
 
-    const { balance } = await calculateUsageAndBalance(
-      'test-user',
-      new Date('2024-01-01'),
-      new Date('2024-01-15'), // Pass current time when grants are active
-    )
+    const { balance } = await calculateUsageAndBalance({
+      userId: 'test-user',
+      quotaResetDate: new Date('2024-01-01'),
+      now: new Date('2024-01-15'), // Pass current time when grants are active
+      logger,
+    })
 
     expect(balance.totalRemaining).toBe(0)
     expect(balance.totalDebt).toBe(100)
@@ -213,11 +213,12 @@ describe('Usage Calculation System', () => {
       },
     }))
 
-    const { balance, usageThisCycle } = await calculateUsageAndBalance(
-      'test-user',
-      new Date('2024-01-01'),
-      new Date('2024-01-15'), // Pass current time when grants are active
-    )
+    const { balance, usageThisCycle } = await calculateUsageAndBalance({
+      userId: 'test-user',
+      quotaResetDate: new Date('2024-01-01'),
+      now: new Date('2024-01-15'), // Pass current time when grants are active
+      logger,
+    })
 
     // Settlement: 100 positive balance - 50 debt = 50 remaining
     expect(balance.totalRemaining).toBe(50)

backend/src/llm-apis/message-cost-tracker.ts

@@ -540,7 +540,11 @@ async function updateUserCycleUsage(
     return result
   } else {
     // Consume from personal credits
-    const result = await consumeCredits(userId, creditsUsed)
+    const result = await consumeCredits({
+      userId,
+      creditsToConsume: creditsUsed,
+      logger,
+    })
 
     if (VERBOSE) {
       logger.debug(

backend/src/websockets/middleware.ts

@@ -345,10 +345,11 @@ protec.use(async (action, clientSessionId, ws, userInfo) => {
     // Continue execution to check remaining balance
   }
 
-  const { usageThisCycle, balance } = await calculateUsageAndBalance(
+  const { usageThisCycle, balance } = await calculateUsageAndBalance({
     userId,
-    user?.next_quota_reset ?? new Date(0),
-  )
+    quotaResetDate: user?.next_quota_reset ?? new Date(0),
+    logger,
+  })
 
   // Check if we have enough remaining credits
   if (balance.totalRemaining <= 0) {

backend/src/websockets/websocket-action.ts

@@ -104,7 +104,11 @@ export async function genUsageResponse(params: {
     try {
       // Get the usage data
       const { balance: balanceDetails, usageThisCycle } =
-        await calculateUsageAndBalance(userId, new Date())
+        await calculateUsageAndBalance({
+          userId,
+          quotaResetDate: new Date(),
+          logger,
+        })
 
       return {
         type: 'usage-response' as const,
@@ -428,10 +432,7 @@ export async function requestFiles(params: {
  * @param filePath - The path of the file to request
  * @returns Promise resolving to the file contents or null if not found
  */
-export async function requestFile(params: {
-  ws: WebSocket
-  filePath: string
-}) {
+export async function requestFile(params: { ws: WebSocket; filePath: string }) {
   const { ws, filePath } = params
   const files = await requestFiles({ ws, filePaths: [filePath] })
   return files[filePath] ?? null

packages/billing/src/auto-topup.ts

@@ -233,10 +233,10 @@ export async function checkAndTriggerAutoTopup(params: {
     }
 
     // Calculate balance
-    const { balance } = await calculateUsageAndBalance(
-      userId,
-      user.next_quota_reset ?? new Date(0),
-    )
+    const { balance } = await calculateUsageAndBalance({
+      ...params,
+      quotaResetDate: user.next_quota_reset ?? new Date(0),
+    })
 
     if (
       balance.totalRemaining >= user.auto_topup_threshold &&