pass in logger to searchLibraries

Author: charleslien

backend/src/__tests__/prompt-caching-subagents.test.ts

@@ -1,15 +1,9 @@
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
-import {
-  clearMockedModules,
-  mockModule,
-} from '@codebuff/common/testing/mock-modules'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
   spyOn,
   beforeEach,
   afterEach,
-  beforeAll,
-  afterAll,
   describe,
   expect,
   it,
@@ -23,8 +17,16 @@ import * as websocketAction from '../websockets/websocket-action'
 import type { AgentTemplate } from '../templates/types'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
+import type { Logger } from '@codebuff/types/logger'
 import type { WebSocket } from 'ws'
 
+const logger: Logger = {
+  debug: () => {},
+  info: () => {},
+  warn: () => {},
+  error: () => {},
+}
+
 const mockFileContext: ProjectFileContext = {
   projectRoot: '/test',
   cwd: '/test',
@@ -62,19 +64,6 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
   let mockLocalAgentTemplates: Record<string, AgentTemplate>
   let capturedMessages: Message[] = []
 
-  beforeAll(() => {
-    // Mock logger
-    mockModule('@codebuff/backend/util/logger', () => ({
-      logger: {
-        debug: () => {},
-        error: () => {},
-        info: () => {},
-        warn: () => {},
-      },
-      withLoggerContext: async (context: any, fn: () => Promise<any>) => fn(),
-    }))
-  })
-
   beforeEach(() => {
     capturedMessages = []
 
@@ -167,10 +156,6 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
     mock.restore()
   })
 
-  afterAll(() => {
-    clearMockedModules()
-  })
-
   it('should inherit parent system prompt when inheritParentSystemPrompt is true', async () => {
     const sessionState = getInitialSessionState(mockFileContext)
     const ws = new MockWebSocket() as unknown as WebSocket
@@ -189,7 +174,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
-      logger: console as any,
+      logger,
     })
 
     // Capture parent's messages which include the system prompt
@@ -224,7 +209,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
       parentSystemPrompt: parentSystemPrompt,
-      logger: console as any,
+      logger,
     })
 
     // Verify child uses parent's system prompt
@@ -272,7 +257,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
-      logger: console as any,
+      logger,
     })
 
     const parentMessages = capturedMessages
@@ -301,7 +286,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
       parentSystemPrompt: parentSystemPrompt,
-      logger: console as any,
+      logger,
     })
 
     const childMessages = capturedMessages
@@ -350,7 +335,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
-      logger: console as any,
+      logger,
     })
 
     const parentMessages = capturedMessages
@@ -382,7 +367,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
       parentSystemPrompt: parentSystemPrompt,
-      logger: console as any,
+      logger,
     })
 
     const childMessages = capturedMessages
@@ -457,7 +442,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
-      logger: console as any,
+      logger,
     })
 
     const parentMessages = capturedMessages
@@ -486,7 +471,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
       parentSystemPrompt: parentSystemPrompt,
-      logger: console as any,
+      logger,
     })
 
     const childMessages = capturedMessages
@@ -545,7 +530,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
-      logger: console as any,
+      logger,
     })
 
     const parentMessages = capturedMessages
@@ -577,7 +562,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       clientSessionId: 'test-session',
       onResponseChunk: () => {},
       parentSystemPrompt: parentSystemPrompt,
-      logger: console as any,
+      logger,
     })
 
     const childMessages = capturedMessages

backend/src/get-documentation-for-query.ts

@@ -69,9 +69,11 @@ export async function getDocumentationForQuery(params: {
   const allRawChunks = (
     await Promise.all(
       libraries.map(({ libraryName, topic }) =>
-        fetchContext7LibraryDocumentation(libraryName, {
+        fetchContext7LibraryDocumentation({
+          query: libraryName,
           tokens,
           topic,
+          logger,
         }),
       ),
     )

backend/src/llm-apis/context7-api.ts

@@ -1,7 +1,8 @@
 import { withTimeout } from '@codebuff/common/util/promise'
 import { env } from '@codebuff/internal/env'
 
-import { logger } from '../util/logger'
+import type { ParamsOf } from '@codebuff/types/common'
+import type { Logger } from '@codebuff/types/logger'
 
 const CONTEXT7_API_BASE_URL = 'https://context7.com/api/v1'
 const DEFAULT_TYPE = 'txt'
@@ -42,9 +43,12 @@ export interface SearchResult {
  * Lists all available documentation projects from Context7
  * @returns Array of projects with their metadata, or null if the request fails
  */
-export async function searchLibraries(
-  query: string,
-): Promise<SearchResult[] | null> {
+export async function searchLibraries(params: {
+  query: string
+  logger: Logger
+}): Promise<SearchResult[] | null> {
+  const { query, logger } = params
+
   const searchStartTime = Date.now()
   const searchContext = {
     query,
@@ -127,23 +131,26 @@ export async function searchLibraries(
  * @returns The documentation text or null if the request fails
  */
 export async function fetchContext7LibraryDocumentation(
-  query: string,
-  options: {
+  params: {
+    query: string
     tokens?: number
     topic?: string
     folders?: string
-  } = {},
+    logger: Logger
+  } & ParamsOf<typeof searchLibraries>,
 ): Promise<string | null> {
+  const { query, tokens, topic, folders, logger } = params
+
   const apiStartTime = Date.now()
   const apiContext = {
     query,
-    requestedTokens: options.tokens,
-    topic: options.topic,
-    folders: options.folders,
+    requestedTokens: tokens,
+    topic,
+    folders,
   }
 
   const searchStartTime = Date.now()
-  const libraries = await searchLibraries(query)
+  const libraries = await searchLibraries(params)
   const searchDuration = Date.now() - searchStartTime
 
   if (!libraries || libraries.length === 0) {
@@ -179,10 +186,9 @@ export async function fetchContext7LibraryDocumentation(
 
   try {
     const url = new URL(`${CONTEXT7_API_BASE_URL}/${libraryId}`)
-    if (options.tokens)
-      url.searchParams.set('tokens', options.tokens.toString())
-    if (options.topic) url.searchParams.set('topic', options.topic)
-    if (options.folders) url.searchParams.set('folders', options.folders)
+    if (tokens) url.searchParams.set('tokens', tokens.toString())
+    if (topic) url.searchParams.set('topic', topic)
+    if (folders) url.searchParams.set('folders', folders)
     url.searchParams.set('type', DEFAULT_TYPE)
 
     const fetchStartTime = Date.now()

backend/src/tools/handlers/tool/read-docs.ts

@@ -1,11 +1,11 @@
 import { fetchContext7LibraryDocumentation } from '../../../llm-apis/context7-api'
 
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
-import type { Logger } from '@codebuff/types/logger'
 import type {
   CodebuffToolCall,
   CodebuffToolOutput,
 } from '@codebuff/common/tools/list'
+import type { Logger } from '@codebuff/types/logger'
 
 export const handleReadDocs = (({
   previousToolCallFinished,
@@ -60,13 +60,12 @@ export const handleReadDocs = (({
 
   const documentationPromise = (async () => {
     try {
-      const documentation = await fetchContext7LibraryDocumentation(
-        libraryTitle,
-        {
-          topic,
-          tokens: max_tokens,
-        },
-      )
+      const documentation = await fetchContext7LibraryDocumentation({
+        query: libraryTitle,
+        topic,
+        tokens: max_tokens,
+        logger,
+      })
 
       const docsDuration = Date.now() - docsStartTime
       const resultLength = documentation?.length || 0