Implement web search and docs search in the agent runtime, using our web endpoint. Track costs

Author: jahooma

packages/agent-runtime/src/__tests__/read-docs-tool.test.ts

@@ -0,0 +1,174 @@
+import { withTimeout } from '@codebuff/common/util/promise'
+import { env } from '@codebuff/common/env'
+
+import type { Logger } from '@codebuff/common/types/contracts/logger'
+
+const FETCH_TIMEOUT_MS = 30_000
+
+export async function callWebSearchAPI(params: {
+  query: string
+  depth?: 'standard' | 'deep'
+  repoUrl?: string | null
+  fetch: typeof globalThis.fetch
+  logger: Logger
+  baseUrl?: string
+  apiKey?: string
+}): Promise<{ result?: string; error?: string; creditsUsed?: number }> {
+  const { query, depth = 'standard', repoUrl, fetch, logger } = params
+  const baseUrl = params.baseUrl ?? env.NEXT_PUBLIC_CODEBUFF_APP_URL
+  const apiKey = params.apiKey ?? process.env.CODEBUFF_API_KEY
+
+  if (!baseUrl || !apiKey) {
+    return { error: 'Missing Codebuff base URL or API key' }
+  }
+
+  const url = `${baseUrl}/api/v1/web-search`
+  const payload = { query, depth, ...(repoUrl ? { repoUrl } : {}) }
+
+  try {
+    const res = await withTimeout(
+      fetch(url, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${apiKey}`,
+          'x-codebuff-api-key': apiKey,
+        },
+        body: JSON.stringify(payload),
+      }),
+      FETCH_TIMEOUT_MS,
+    )
+
+    const text = await res.text()
+    const tryJson = () => {
+      try {
+        return JSON.parse(text)
+      } catch {
+        return null
+      }
+    }
+
+    if (!res.ok) {
+      const maybe = tryJson()
+      const err =
+        (maybe && (maybe.error || maybe.message)) || text || 'Request failed'
+      logger.warn(
+        {
+          url,
+          status: res.status,
+          statusText: res.statusText,
+          body: text?.slice(0, 500),
+        },
+        'Web API web-search request failed',
+      )
+      return { error: typeof err === 'string' ? err : 'Unknown error' }
+    }
+
+    const data = tryJson()
+    if (data && typeof data.result === 'string') {
+      return {
+        result: data.result,
+        creditsUsed: typeof data.creditsUsed === 'number' ? data.creditsUsed : undefined,
+      }
+    }
+    if (data && typeof data.error === 'string') return { error: data.error }
+    return { error: 'Invalid response format' }
+  } catch (error) {
+    logger.error(
+      {
+        error:
+          error instanceof Error
+            ? { name: error.name, message: error.message, stack: error.stack }
+            : error,
+      },
+      'Web API web-search network error',
+    )
+    return { error: error instanceof Error ? error.message : 'Network error' }
+  }
+}
+
+export async function callDocsSearchAPI(params: {
+  libraryTitle: string
+  topic?: string
+  maxTokens?: number
+  repoUrl?: string | null
+  fetch: typeof globalThis.fetch
+  logger: Logger
+  baseUrl?: string
+  apiKey?: string
+}): Promise<{ documentation?: string; error?: string; creditsUsed?: number }> {
+  const { libraryTitle, topic, maxTokens, repoUrl, fetch, logger } = params
+  const baseUrl = params.baseUrl ?? env.NEXT_PUBLIC_CODEBUFF_APP_URL
+  const apiKey = params.apiKey ?? process.env.CODEBUFF_API_KEY
+
+  if (!baseUrl || !apiKey) {
+    return { error: 'Missing Codebuff base URL or API key' }
+  }
+
+  const url = `${baseUrl}/api/v1/docs-search`
+  const payload: Record<string, any> = { libraryTitle }
+  if (topic) payload.topic = topic
+  if (typeof maxTokens === 'number') payload.maxTokens = maxTokens
+  if (repoUrl) payload.repoUrl = repoUrl
+
+  try {
+    const res = await withTimeout(
+      fetch(url, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${apiKey}`,
+          'x-codebuff-api-key': apiKey,
+        },
+        body: JSON.stringify(payload),
+      }),
+      FETCH_TIMEOUT_MS,
+    )
+
+    const text = await res.text()
+    const tryJson = () => {
+      try {
+        return JSON.parse(text) as any
+      } catch {
+        return null
+      }
+    }
+
+    if (!res.ok) {
+      const maybe = tryJson()
+      const err =
+        (maybe && (maybe.error || maybe.message)) || text || 'Request failed'
+      logger.warn(
+        {
+          url,
+          status: res.status,
+          statusText: res.statusText,
+          body: text?.slice(0, 500),
+        },
+        'Web API docs-search request failed',
+      )
+      return { error: typeof err === 'string' ? err : 'Unknown error' }
+    }
+
+    const data = tryJson()
+    if (data && typeof data.documentation === 'string') {
+      return {
+        documentation: data.documentation,
+        creditsUsed: typeof data.creditsUsed === 'number' ? data.creditsUsed : undefined,
+      }
+    }
+    if (data && typeof data.error === 'string') return { error: data.error }
+    return { error: 'Invalid response format' }
+  } catch (error) {
+    logger.error(
+      {
+        error:
+          error instanceof Error
+            ? { name: error.name, message: error.message, stack: error.stack }
+            : error,
+      },
+      'Web API docs-search network error',
+    )
+    return { error: error instanceof Error ? error.message : 'Network error' }
+  }
+}

packages/agent-runtime/src/__tests__/web-search-tool.test.ts

@@ -81,6 +81,7 @@ export const runAgentStep = async (
     | 'agentTemplate'
     | 'agentContext'
     | 'fullResponse'
+    | 'onCostCalculated'
   > &
     ParamsExcluding<
       typeof getAgentStreamFromTemplate,
@@ -245,27 +246,19 @@ export const runAgentStep = async (
 
   const { model } = agentTemplate
 
+  let stepCreditsUsed = 0
+
+  const onCostCalculated = async (credits: number) => {
+    stepCreditsUsed += credits
+    agentState.creditsUsed += credits
+    agentState.directCreditsUsed += credits
+  }
+
   const { getStream } = getAgentStreamFromTemplate({
     ...params,
     agentId: agentState.parentId ? agentState.agentId : undefined,
     template: agentTemplate,
-    onCostCalculated: async (credits: number) => {
-      try {
-        agentState.creditsUsed += credits
-        agentState.directCreditsUsed += credits
-        // Transactional cost attribution: ensure costs are actually deducted
-        // This is already handled by the saveMessage function which calls updateUserCycleUsage
-        // If that fails, the promise rejection will bubble up and halt agent execution
-      } catch (error) {
-        logger.error(
-          { agentId: agentState.agentId, credits, error },
-          'Failed to add cost to agent state',
-        )
-        throw new Error(
-          `Cost tracking failed for agent ${agentState.agentId}: ${error}`,
-        )
-      }
-    },
+    onCostCalculated,
     includeCacheControl: supportsCacheControl(agentTemplate.model),
   })
 
@@ -315,6 +308,7 @@ export const runAgentStep = async (
     agentTemplate,
     agentContext,
     fullResponse,
+    onCostCalculated,
   })
   toolResults.push(...newToolResults)
 
@@ -423,6 +417,7 @@ export const runAgentStep = async (
       toolResults,
       agentContext: newAgentContext,
       fullResponseChunks,
+      stepCreditsUsed,
     },
     `End agent ${agentType} step ${iterationNum} (${userInputId}${prompt ? ` - Prompt: ${prompt.slice(0, 20)}` : ''})`,
   )
@@ -465,6 +460,7 @@ export async function loopAgentSteps(
     | 'stepsComplete'
     | 'stepNumber'
     | 'system'
+    | 'onCostCalculated'
   > &
     ParamsExcluding<typeof getAgentTemplate, 'agentId'> &
     ParamsExcluding<
@@ -676,6 +672,10 @@ export async function loopAgentSteps(
           system,
           stepsComplete: shouldEndTurn,
           stepNumber: totalSteps,
+          onCostCalculated: async (credits: number) => {
+            agentState.creditsUsed += credits
+            agentState.directCreditsUsed += credits
+          },
         })
         const {
           agentState: programmaticAgentState,