Refactor utility functions and middleware to pass logger as parameter

Refactored the following files to accept logger as a parameter instead of importing it directly:
- util/simplify-tool-results.ts: simplifyTerminalCommandResults now takes logger param
- util/messages.ts: All message utility functions now accept logger param
- util/check-auth.ts: checkAuth and checkAdmin refactored with logger param
- util/quickjs-sandbox.ts: dispose methods now require logger param
- Updated all callers in production code and test files
- Updated plan documentation with clearer instructions for test logger naming

🤖 Generated with Codebuff
Co-Authored-By: Codebuff <noreply@codebuff.com>

Author: charleslien

backend/src/api/usage.ts

@@ -48,6 +48,7 @@ async function usageHandler(
       fingerprintId,
       authToken,
       clientSessionId,
+      logger,
     })
     if (authResult) {
       const errorMessage =

backend/src/find-files/request-files-prompt.ts

@@ -344,16 +344,17 @@ async function getRelevantFiles(params: {
     logger,
   } = params
   const bufferTokens = 100_000
-  const messagesWithPrompt = getMessagesSubset(
-    [
+  const messagesWithPrompt = getMessagesSubset({
+    messages: [
       ...messages,
       {
         role: 'user' as const,
         content: userPrompt,
       },
     ],
-    bufferTokens,
-  )
+    otherTokens: bufferTokens,
+    logger,
+  })
   const start = performance.now()
   let codebuffMessages = messagesWithSystem({ messages: messagesWithPrompt, system })
 
@@ -435,16 +436,17 @@ async function getRelevantFilesForTraining(params: {
     logger,
   } = params
   const bufferTokens = 100_000
-  const messagesWithPrompt = getMessagesSubset(
-    [
+  const messagesWithPrompt = getMessagesSubset({
+    messages: [
       ...messages,
       {
         role: 'user' as const,
         content: userPrompt,
       },
     ],
-    bufferTokens,
-  )
+    otherTokens: bufferTokens,
+    logger,
+  })
   const start = performance.now()
   let response = await promptAiSdk({
     messages: messagesWithSystem({ messages: messagesWithPrompt, system }),

backend/src/run-programmatic-step.ts

@@ -22,6 +22,7 @@ import type {
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { AgentState } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
+import type { Logger } from '@codebuff/types/logger'
 import type { WebSocket } from 'ws'
 
 // Global sandbox manager for QuickJS contexts
@@ -38,16 +39,33 @@ export function clearAgentGeneratorCache() {
   }
   runIdToStepAll.clear()
   // Clean up QuickJS sandboxes
-  sandboxManager.dispose()
+  sandboxManager.dispose({ logger })
 }
 
 // Function to handle programmatic agents
-export async function runProgrammaticStep(
-  agentState: AgentState,
-  {
+export async function runProgrammaticStep(params: {
+  agentState: AgentState
+  template: AgentTemplate
+  prompt: string | undefined
+  toolCallParams: Record<string, any> | undefined
+  system: string | undefined
+  userId: string | undefined
+  userInputId: string
+  clientSessionId: string
+  fingerprintId: string
+  onResponseChunk: (chunk: string | PrintModeEvent) => void
+  fileContext: ProjectFileContext
+  ws: WebSocket
+  localAgentTemplates: Record<string, AgentTemplate>
+  stepsComplete: boolean
+  stepNumber: number
+  logger: Logger
+}): Promise<{ agentState: AgentState; endTurn: boolean; stepNumber: number }> {
+  const {
+    agentState,
     template,
     prompt,
-    params,
+    toolCallParams,
     system,
     userId,
     userInputId,
@@ -58,24 +76,10 @@ export async function runProgrammaticStep(
     ws,
     localAgentTemplates,
     stepsComplete,
-    stepNumber,
-  }: {
-    template: AgentTemplate
-    prompt: string | undefined
-    params: Record<string, any> | undefined
-    system: string | undefined
-    userId: string | undefined
-    userInputId: string
-    clientSessionId: string
-    fingerprintId: string
-    onResponseChunk: (chunk: string | PrintModeEvent) => void
-    fileContext: ProjectFileContext
-    ws: WebSocket
-    localAgentTemplates: Record<string, AgentTemplate>
-    stepsComplete: boolean
-    stepNumber: number
-  },
-): Promise<{ agentState: AgentState; endTurn: boolean; stepNumber: number }> {
+    logger,
+  } = params
+  let { stepNumber } = params
+
   if (!template.handleSteps) {
     throw new Error('No step handler found for agent template ' + template.id)
   }
@@ -119,11 +123,12 @@ export async function runProgrammaticStep(
         initialInput: {
           agentState,
           prompt,
-          params,
+          params: toolCallParams,
           logger: streamingLogger,
         },
         config: undefined, // config
-        logger: streamingLogger, // pass the streaming logger instance for internal use
+        sandboxLogger: streamingLogger, // pass the streaming logger instance for internal use
+        logger,
       })
     } else {
       // Initialize native generator
@@ -363,7 +368,7 @@ export async function runProgrammaticStep(
     if (endTurn) {
       if (sandbox) {
         // Clean up QuickJS sandbox if execution is complete
-        sandboxManager.removeSandbox({ runId: agentState.runId })
+        sandboxManager.removeSandbox({ runId: agentState.runId, logger })
       }
       delete runIdToGenerator[agentState.runId]
       runIdToStepAll.delete(agentState.runId)

backend/src/util/__tests__/messages.test.ts

@@ -8,7 +8,6 @@ import {
   spyOn,
 } from 'bun:test'
 
-import { logger } from '../logger'
 import {
   trimMessagesToFitTokenLimit,
   messagesWithSystem,
@@ -40,6 +39,14 @@ describe('messagesWithSystem', () => {
   })
 })
 
+// Mock logger for tests
+const logger = {
+  debug: () => {},
+  info: () => {},
+  warn: () => {},
+  error: () => {},
+}
+
 describe('trimMessagesToFitTokenLimit', () => {
   beforeEach(() => {
     // Mock countTokensJson to just count characters
@@ -211,11 +218,12 @@ describe('trimMessagesToFitTokenLimit', () => {
   it('handles all features working together correctly', () => {
     const maxTotalTokens = 3000
     const systemTokens = 0
-    const result = trimMessagesToFitTokenLimit(
-      testMessages,
+    const result = trimMessagesToFitTokenLimit({
+      messages: testMessages,
       systemTokens,
       maxTotalTokens,
-    )
+      logger,
+    })
 
     // Should have replacement message for omitted content
     expect(result.length).toBeGreaterThan(0)
@@ -236,11 +244,12 @@ describe('trimMessagesToFitTokenLimit', () => {
   it('subtracts system tokens from total tokens', () => {
     const maxTotalTokens = 10_000
     const systemTokens = 7_000
-    const result = trimMessagesToFitTokenLimit(
-      testMessages,
+    const result = trimMessagesToFitTokenLimit({
+      messages: testMessages,
       systemTokens,
       maxTotalTokens,
-    )
+      logger,
+    })
 
     // Should have replacement message for omitted content
     expect(result.length).toBeGreaterThan(0)
@@ -261,11 +270,12 @@ describe('trimMessagesToFitTokenLimit', () => {
   it('does not simplify if under token limit', () => {
     const maxTotalTokens = 10_000
     const systemTokens = 100
-    const result = trimMessagesToFitTokenLimit(
-      testMessages,
+    const result = trimMessagesToFitTokenLimit({
+      messages: testMessages,
       systemTokens,
       maxTotalTokens,
-    )
+      logger,
+    })
 
     // All messages should be unchanged
     expect(result).toHaveLength(testMessages.length)
@@ -282,7 +292,7 @@ describe('trimMessagesToFitTokenLimit', () => {
   it('handles empty messages array', () => {
     const maxTotalTokens = 200
     const systemTokens = 100
-    const result = trimMessagesToFitTokenLimit([], systemTokens, maxTotalTokens)
+    const result = trimMessagesToFitTokenLimit({ messages: [], systemTokens, maxTotalTokens, logger })
 
     expect(result).toEqual([])
   })
@@ -305,7 +315,7 @@ describe('trimMessagesToFitTokenLimit', () => {
         },
       ]
 
-      const result = trimMessagesToFitTokenLimit(messages, 0, 1000)
+      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 1000, logger })
 
       // Should contain the kept messages
       const keptMessages = result.filter(
@@ -335,7 +345,7 @@ describe('trimMessagesToFitTokenLimit', () => {
         },
       ] as Message[]
 
-      const result = trimMessagesToFitTokenLimit(messages, 0, 10000)
+      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 10000, logger })
 
       // Should be unchanged when under token limit
       expect(result).toHaveLength(2)
@@ -351,7 +361,7 @@ describe('trimMessagesToFitTokenLimit', () => {
         { role: 'user', content: 'Keep this', keepDuringTruncation: true },
       ]
 
-      const result = trimMessagesToFitTokenLimit(messages, 0, 1000)
+      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 1000, logger })
 
       // Should only have one replacement message for consecutive removals
       const replacementMessages = result.filter(
@@ -381,7 +391,7 @@ describe('trimMessagesToFitTokenLimit', () => {
         { role: 'user', content: 'C'.repeat(100) }, // Might be kept
       ]
 
-      const result = trimMessagesToFitTokenLimit(messages, 0, 2000)
+      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 2000, logger })
 
       // Should preserve the keepDuringTruncation message
       const keptMessage = result.find(
@@ -405,7 +415,7 @@ describe('trimMessagesToFitTokenLimit', () => {
         { role: 'user', content: 'C'.repeat(800) }, // Large message to force truncation
       ]
 
-      const result = trimMessagesToFitTokenLimit(messages, 0, 500)
+      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 500, logger })
 
       // Should keep both marked messages
       const keptMessages = result.filter(
@@ -428,7 +438,7 @@ describe('trimMessagesToFitTokenLimit', () => {
 
 describe('getPreviouslyReadFiles', () => {
   it('returns empty array when no messages provided', () => {
-    const result = getPreviouslyReadFiles([])
+    const result = getPreviouslyReadFiles({ messages: [], logger })
     expect(result).toEqual([])
   })
 
@@ -449,7 +459,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'write_file'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([])
   })
 
@@ -481,7 +491,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'read_files'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([
       {
         path: 'src/test.ts',
@@ -518,7 +528,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'find_files'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([
       {
         path: 'components/Button.tsx',
@@ -573,7 +583,7 @@ describe('getPreviouslyReadFiles', () => {
       },
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([
       { path: 'file1.ts', content: 'content 1' },
       { path: 'file2.ts', content: 'content 2' },
@@ -611,7 +621,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'read_files'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([
       { path: 'small-file.ts', content: 'small content' },
       { path: 'another-small-file.ts', content: 'another small content' },
@@ -633,7 +643,7 @@ describe('getPreviouslyReadFiles', () => {
       },
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([])
     expect(mockLoggerError).toHaveBeenCalled()
 
@@ -660,7 +670,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'find_files'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([])
   })
 
@@ -690,7 +700,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'read_files'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([{ path: 'test.ts', content: 'test content' }])
   })
 
@@ -712,7 +722,7 @@ describe('getPreviouslyReadFiles', () => {
       } satisfies CodebuffToolMessage<'read_files'>,
     ]
 
-    const result = getPreviouslyReadFiles(messages)
+    const result = getPreviouslyReadFiles({ messages, logger })
     expect(result).toEqual([])
   })
 })