Merge origin/main into brandon/cli-env-filter

Resolved merge conflict in common/src/types/contracts/billing.ts by accepting
upstream changes. The upstream version uses ErrorOr<CreditFallbackResult>
which is cleaner than the previous custom success/error object.

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

Author: brandonkachen

backend/src/__tests__/web-search-tool.test.ts

@@ -12,6 +12,7 @@ import {
 } from '@codebuff/common/testing/impl/agent-runtime'
 import { getToolCallString } from '@codebuff/common/tools/utils'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
+import { success } from '@codebuff/common/util/error'
 import {
   afterEach,
   beforeAll,
@@ -25,15 +26,15 @@ import {
 
 import { mockFileContext } from './test-utils'
 import researcherAgent from '../../../.agents/researcher/researcher'
-import * as linkupApi from '../llm-apis/linkup-api'
+import * as linkupApi from '@codebuff/agent-runtime/llm-apis/linkup-api'
 import { runAgentStep } from '../run-agent-step'
 
 import type {
   AgentRuntimeDeps,
   AgentRuntimeScopedDeps,
 } from '@codebuff/common/types/contracts/agent-runtime'
 
-let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
+let agentRuntimeImpl: AgentRuntimeDeps
 function mockAgentStream(content: string | string[]) {
   agentRuntimeImpl.promptAiSdkStream = async function* ({}) {
     if (typeof content === 'string') {
@@ -54,6 +55,14 @@ describe('web_search tool with researcher agent', () => {
   })
 
   beforeEach(() => {
+    agentRuntimeImpl = {
+      ...TEST_AGENT_RUNTIME_IMPL,
+      consumeCreditsWithFallback: async () => {
+        return success({
+          chargedToOrganization: false,
+        })
+      },
+    }
     agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
 
     // Mock analytics and tracing
@@ -139,11 +148,12 @@ describe('web_search tool with researcher agent', () => {
     })
 
     // Just verify that searchWeb was called
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'test query',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'test query',
+        depth: 'standard',
+      }),
+    )
   })
 
   test('should successfully perform web search with basic query', async () => {
@@ -190,11 +200,12 @@ describe('web_search tool with researcher agent', () => {
       spawnParams: undefined,
     })
 
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'Next.js 15 new features',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'Next.js 15 new features',
+        depth: 'standard',
+      }),
+    )
 
     // Check that the search results were added to the message history
     const toolResultMessages = newAgentState.messageHistory.filter(
@@ -251,11 +262,12 @@ describe('web_search tool with researcher agent', () => {
       spawnParams: undefined,
     })
 
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'React Server Components tutorial',
-      depth: 'deep',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'React Server Components tutorial',
+        depth: 'deep',
+      }),
+    )
   })
 
   test('should handle case when no search results are found', async () => {
@@ -298,11 +310,12 @@ describe('web_search tool with researcher agent', () => {
     })
 
     // Verify that searchWeb was called
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'very obscure search query that returns nothing',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'very obscure search query that returns nothing',
+        depth: 'standard',
+      }),
+    )
 
     // Check that the "no results found" message was added
     const toolResultMessages = newAgentState.messageHistory.filter(
@@ -358,11 +371,12 @@ describe('web_search tool with researcher agent', () => {
     })
 
     // Verify that searchWeb was called
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'test query',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'test query',
+        depth: 'standard',
+      }),
+    )
 
     // Check that the error message was added
     const toolResultMessages = newAgentState.messageHistory.filter(
@@ -417,11 +431,12 @@ describe('web_search tool with researcher agent', () => {
     })
 
     // Verify that searchWeb was called
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'test query',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'test query',
+        depth: 'standard',
+      }),
+    )
   })
 
   test('should handle non-Error exceptions', async () => {
@@ -466,11 +481,12 @@ describe('web_search tool with researcher agent', () => {
     })
 
     // Verify that searchWeb was called
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'test query',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'test query',
+        depth: 'standard',
+      }),
+    )
 
     // Check that the error message was added
     const toolResultMessages = newAgentState.messageHistory.filter(
@@ -527,11 +543,12 @@ describe('web_search tool with researcher agent', () => {
     })
 
     // Verify that searchWeb was called
-    expect(linkupApi.searchWeb).toHaveBeenCalledWith({
-      query: 'test formatting',
-      depth: 'standard',
-      logger: expect.anything(),
-    })
+    expect(linkupApi.searchWeb).toHaveBeenCalledWith(
+      expect.objectContaining({
+        query: 'test formatting',
+        depth: 'standard',
+      }),
+    )
 
     // Check that the search results were formatted correctly
     const toolResultMessages = newAgentState.messageHistory.filter(

backend/src/impl/agent-runtime.ts

@@ -1,3 +1,4 @@
+import { consumeCreditsWithFallback } from '@codebuff/billing'
 import { trackEvent } from '@codebuff/common/analytics'
 import { consumeCreditsWithFallback } from '@codebuff/billing'
 
@@ -22,6 +23,9 @@ export const BACKEND_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
   finishAgentRun,
   addAgentStep,
 
+  // Billing
+  consumeCreditsWithFallback,
+
   // LLM
   promptAiSdkStream,
   promptAiSdk,

backend/src/tools/handlers/list.ts

@@ -14,14 +14,14 @@ import { handleRunFileChangeHooks } from '@codebuff/agent-runtime/tools/handlers
 import { handleRunTerminalCommand } from '@codebuff/agent-runtime/tools/handlers/tool/run-terminal-command'
 import { handleSetMessages } from '@codebuff/agent-runtime/tools/handlers/tool/set-messages'
 import { handleSetOutput } from '@codebuff/agent-runtime/tools/handlers/tool/set-output'
+import { handleStrReplace } from '@codebuff/agent-runtime/tools/handlers/tool/str-replace'
 import { handleThinkDeeply } from '@codebuff/agent-runtime/tools/handlers/tool/think-deeply'
 import { handleUpdateSubgoal } from '@codebuff/agent-runtime/tools/handlers/tool/update-subgoal'
+import { handleWebSearch } from '@codebuff/agent-runtime/tools/handlers/tool/web-search'
 import { handleWriteFile } from '@codebuff/agent-runtime/tools/handlers/tool/write-file'
 
 import { handleSpawnAgentInline } from './tool/spawn-agent-inline'
 import { handleSpawnAgents } from './tool/spawn-agents'
-import { handleStrReplace } from './tool/str-replace'
-import { handleWebSearch } from './tool/web-search'
 
 import type { CodebuffToolHandlerFunction } from '@codebuff/agent-runtime/tools/handlers/handler-function-type'
 import type { ToolName } from '@codebuff/common/tools/constants'

common/src/testing/impl/agent-runtime.ts

@@ -31,6 +31,13 @@ export const TEST_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
   finishAgentRun: async () => {},
   addAgentStep: async () => 'test-agent-step-id',
 
+  // Billing
+  consumeCreditsWithFallback: async () => {
+    throw new Error(
+      'consumeCreditsWithFallback not implemented in test runtime',
+    )
+  },
+
   // LLM
   promptAiSdkStream: async function* () {
     throw new Error('promptAiSdkStream not implemented in test runtime')

common/src/types/contracts/agent-runtime.ts

@@ -33,6 +33,9 @@ export type AgentRuntimeDeps = {
   finishAgentRun: FinishAgentRunFn
   addAgentStep: AddAgentStepFn
 
+  // Billing
+  consumeCreditsWithFallback: ConsumeCreditsWithFallbackFn
+
   // LLM
   promptAiSdkStream: PromptAiSdkStreamFn
   promptAiSdk: PromptAiSdkFn

common/src/types/contracts/billing.ts

@@ -1,4 +1,5 @@
 import type { Logger } from './logger'
+import type { ErrorOr } from '../../util/error'
 
 export type GetUserUsageDataFn = (params: {
   userId: string
@@ -12,12 +13,12 @@ export type ConsumeCreditsWithFallbackFn = (params: {
   userId: string
   creditsToCharge: number
   repoUrl?: string | null
-  context: string
+  context: string // Description of what the credits are for (e.g., 'web search', 'documentation lookup')
   logger: Logger
-}) => Promise<{
-  success: boolean
+}) => Promise<ErrorOr<CreditFallbackResult>>
+
+export type CreditFallbackResult = {
   organizationId?: string
   organizationName?: string
   chargedToOrganization: boolean
-  error?: string
-}>
+}

evals/impl/agent-runtime.ts

@@ -1,3 +1,5 @@
+import { success } from '@codebuff/common/util/error'
+
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 
@@ -13,6 +15,13 @@ export const EVALS_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
   finishAgentRun: async () => {},
   addAgentStep: async () => 'test-agent-step-id',
 
+  // Backend
+  consumeCreditsWithFallback: async () => {
+    return success({
+      chargedToOrganization: false,
+    })
+  },
+
   // LLM
   promptAiSdkStream: async function* () {
     throw new Error('promptAiSdkStream not implemented in eval runtime')

packages/agent-runtime/src/find-files/request-files-prompt.ts

@@ -1,6 +1,5 @@
 import { dirname, isAbsolute, normalize } from 'path'
 
-import { insertTrace } from '@codebuff/bigquery'
 import {
   finetunedVertexModels,
   models,
@@ -17,10 +16,6 @@ import {
 } from '../util/messages'
 
 import type { TextBlock } from '../llm-api/claude'
-import type {
-  GetExpandedFileContextForTrainingTrace,
-  GetRelevantFilesTrace,
-} from '@codebuff/bigquery'
 import type { PromptAiSdkFn } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
@@ -236,29 +231,6 @@ async function getRelevantFiles(
 
   const files = validateFilePaths(response.split('\n'))
 
-  const trace: GetRelevantFilesTrace = {
-    id: crypto.randomUUID(),
-    agent_step_id: agentStepId,
-    user_id: userId ?? '',
-    created_at: new Date(),
-    type: 'get-relevant-files',
-    payload: {
-      messages: messagesWithPrompt,
-      system,
-      output: response,
-      request_type: requestType,
-      user_input_id: userInputId,
-      client_session_id: clientSessionId,
-      fingerprint_id: fingerprintId,
-      model: finetunedModel,
-      repo_name: repoId, // Use repoId parameter for trace
-    },
-  }
-
-  insertTrace({ trace, logger }).catch((error: Error) => {
-    logger.error({ error }, 'Failed to insert trace')
-  })
-
   return { files, duration, requestType, response }
 }
 
@@ -317,29 +289,6 @@ async function getRelevantFilesForTraining(
 
   const files = validateFilePaths(response.split('\n'))
 
-  const trace: GetExpandedFileContextForTrainingTrace = {
-    id: crypto.randomUUID(),
-    agent_step_id: agentStepId,
-    user_id: userId ?? '',
-    created_at: new Date(),
-    type: 'get-expanded-file-context-for-training',
-    payload: {
-      messages: messagesWithPrompt,
-      system,
-      output: response,
-      request_type: requestType,
-      user_input_id: userInputId,
-      client_session_id: clientSessionId,
-      fingerprint_id: fingerprintId,
-      model: models.ft_filepicker_005, // Use specific model for trace
-      repo_name: repoId, // Use repoId parameter for trace
-    },
-  }
-
-  insertTrace({ trace, logger }).catch((error: Error) => {
-    logger.error({ error }, 'Failed to insert trace')
-  })
-
   return { files, duration, requestType, response }
 }