pass in logger to credit delegation

Author: charleslien

backend/src/api/org.ts

@@ -42,7 +42,11 @@ async function isRepoCoveredHandler(
     }
 
     // Check if repository is covered by an organization
-    const orgLookup = await findOrganizationForRepository(userId, remoteUrl)
+    const orgLookup = await findOrganizationForRepository({
+      userId,
+      repositoryUrl: remoteUrl,
+      logger,
+    })
 
     return res.status(200).json({
       isCovered: orgLookup.found,

backend/src/tools/handlers/tool/web-search.ts

@@ -76,6 +76,7 @@ export const handleWebSearch = ((params: {
             creditsToCharge,
             repoUrl,
             context: 'web search',
+            logger,
           })
 
           if (!creditResult.success) {

backend/src/websockets/middleware.ts

@@ -181,7 +181,11 @@ protec.use(async (action, clientSessionId, ws, userInfo) => {
     const { owner, repo } = ownerRepo
 
     // Perform lookup (cache removed)
-    const orgLookup = await findOrganizationForRepository(userId, repoUrl)
+    const orgLookup = await findOrganizationForRepository({
+      userId,
+      repositoryUrl: repoUrl,
+      logger,
+    })
 
     // If an organization covers this repository, check its balance
     if (orgLookup.found && orgLookup.organizationId) {

packages/billing/src/__tests__/credit-delegation.test.ts

@@ -10,6 +10,13 @@ import {
 } from '../credit-delegation'
 
 describe('Credit Delegation', () => {
+  const logger = {
+    debug: () => {},
+    info: () => {},
+    warn: () => {},
+    error: () => {},
+  }
+
   beforeAll(() => {
     // Mock the org-billing functions that credit-delegation depends on
     mockModule('@codebuff/billing/org-billing', () => ({
@@ -47,14 +54,6 @@ describe('Credit Delegation', () => {
         is_active: 'is_active',
       },
     }))
-
-    mockModule('@codebuff/common/util/logger', () => ({
-      logger: {
-        debug: mock(() => {}),
-        info: mock(() => {}),
-        error: mock(() => {}),
-      },
-    }))
   })
 
   afterAll(() => {
@@ -66,7 +65,11 @@ describe('Credit Delegation', () => {
       const userId = 'user-123'
       const repositoryUrl = 'https://github.com/codebuffai/codebuff'
 
-      const result = await findOrganizationForRepository(userId, repositoryUrl)
+      const result = await findOrganizationForRepository({
+        userId,
+        repositoryUrl,
+        logger,
+      })
 
       expect(result.found).toBe(true)
       expect(result.organizationId).toBe('org-123')
@@ -77,7 +80,11 @@ describe('Credit Delegation', () => {
       const userId = 'user-123'
       const repositoryUrl = 'https://github.com/other/repo'
 
-      const result = await findOrganizationForRepository(userId, repositoryUrl)
+      const result = await findOrganizationForRepository({
+        userId,
+        repositoryUrl,
+        logger,
+      })
 
       expect(result.found).toBe(false)
     })
@@ -89,11 +96,12 @@ describe('Credit Delegation', () => {
       const repositoryUrl = null
       const creditsToConsume = 100
 
-      const result = await consumeCreditsWithDelegation(
+      const result = await consumeCreditsWithDelegation({
         userId,
         repositoryUrl,
         creditsToConsume,
-      )
+        logger,
+      })
 
       expect(result.success).toBe(false)
       expect(result.error).toBe('No repository URL provided')

packages/billing/src/credit-delegation.ts

@@ -1,6 +1,5 @@
 import db from '@codebuff/common/db'
 import * as schema from '@codebuff/common/db/schema'
-import { logger } from '@codebuff/common/util/logger'
 import { eq, and } from 'drizzle-orm'
 
 import { consumeCredits } from './balance-calculator'
@@ -10,6 +9,8 @@ import {
   extractOwnerAndRepo,
 } from './org-billing'
 
+import type { Logger } from '@codebuff/types/logger'
+
 export interface OrganizationLookupResult {
   found: boolean
   organizationId?: string
@@ -29,10 +30,13 @@ export interface CreditDelegationResult {
  * Finds the organization associated with a repository for a given user.
  * Uses owner/repo comparison for better matching.
  */
-export async function findOrganizationForRepository(
-  userId: string,
-  repositoryUrl: string,
-): Promise<OrganizationLookupResult> {
+export async function findOrganizationForRepository(params: {
+  userId: string
+  repositoryUrl: string
+  logger: Logger
+}): Promise<OrganizationLookupResult> {
+  const { userId, repositoryUrl, logger } = params
+
   try {
     const normalizedUrl = normalizeRepositoryUrl(repositoryUrl)
     const ownerRepo = extractOwnerAndRepo(normalizedUrl)
@@ -135,23 +139,28 @@ export async function findOrganizationForRepository(
 /**
  * Consumes credits with organization delegation if applicable.
  */
-export async function consumeCreditsWithDelegation(
-  userId: string,
-  repositoryUrl: string | null,
-  creditsToConsume: number,
-): Promise<CreditDelegationResult> {
-  try {
-    // If no repository URL, fall back to personal credits
-    if (!repositoryUrl) {
-      logger.debug(
-        { userId, creditsToConsume },
-        'No repository URL provided, falling back to personal credits',
-      )
-      return { success: false, error: 'No repository URL provided' }
-    }
+export async function consumeCreditsWithDelegation(params: {
+  userId: string
+  repositoryUrl: string | null
+  creditsToConsume: number
+  logger: Logger
+}): Promise<CreditDelegationResult> {
+  const { userId, repositoryUrl, creditsToConsume, logger } = params
 
+  // If no repository URL, fall back to personal credits
+  if (!repositoryUrl) {
+    logger.debug(
+      { userId, creditsToConsume },
+      'No repository URL provided, falling back to personal credits',
+    )
+    return { success: false, error: 'No repository URL provided' }
+  }
+
+  const withRepoUrl = { ...params, repositoryUrl }
+
+  try {
     // Find organization for this repository
-    const orgLookup = await findOrganizationForRepository(userId, repositoryUrl)
+    const orgLookup = await findOrganizationForRepository(withRepoUrl)
 
     if (!orgLookup.found || !orgLookup.organizationId) {
       logger.debug(
@@ -216,13 +225,6 @@ export async function consumeCreditsWithDelegation(
   }
 }
 
-export interface CreditConsumptionOptions {
-  userId: string
-  creditsToCharge: number
-  repoUrl?: string | null
-  context: string // Description of what the credits are for (e.g., 'web search', 'documentation lookup')
-}
-
 export interface CreditFallbackResult {
   success: boolean
   organizationId?: string
@@ -235,19 +237,23 @@ export interface CreditFallbackResult {
  * Helper function that decides whether to charge credits to an organization or user directly.
  * Tries organization delegation first if a repo URL is available, falls back to personal credits.
  */
-export async function consumeCreditsWithFallback(
-  options: CreditConsumptionOptions,
-): Promise<CreditFallbackResult> {
-  const { userId, creditsToCharge, repoUrl, context } = options
+export async function consumeCreditsWithFallback(params: {
+  userId: string
+  creditsToCharge: number
+  repoUrl?: string | null
+  context: string // Description of what the credits are for (e.g., 'web search', 'documentation lookup')
+  logger: Logger
+}): Promise<CreditFallbackResult> {
+  const { userId, creditsToCharge, repoUrl, context, logger } = params
 
   try {
     // Try organization delegation first if repo URL is available
     if (repoUrl) {
-      const delegationResult = await consumeCreditsWithDelegation(
-        userId,
-        repoUrl,
-        creditsToCharge,
-      )
+      const delegationResult = await consumeCreditsWithDelegation({
+        ...params,
+        repositoryUrl: repoUrl,
+        creditsToConsume: creditsToCharge,
+      })
 
       if (delegationResult.success) {
         logger.debug(