Remove bigquery tracing from agent-runtime

Author: jahooma

packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts

@@ -37,10 +37,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
   let loopAgentStepsBaseParams: any
 
   beforeAll(async () => {
-    // Mock bigquery
-    await mockModule('@codebuff/bigquery', () => ({
-      insertTrace: () => {},
-    }))
+    // Set up mocks.
   })
 
   beforeEach(() => {

packages/agent-runtime/src/__tests__/main-prompt.test.ts

@@ -1,4 +1,3 @@
-import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { createTestAgentRuntimeParams } from '@codebuff/common/testing/fixtures/agent-runtime'
@@ -101,11 +100,8 @@ describe('mainPrompt', () => {
         }) as Response,
     }
 
-    // Mock analytics and tracing
+    // Mock analytics
     spyOn(analytics, 'trackEvent').mockImplementation(() => {})
-    spyOn(bigquery, 'insertTrace').mockImplementation(() =>
-      Promise.resolve(true),
-    ) // Return Promise<boolean>
 
     // Mock processFileBlock
     spyOn(processFileBlockModule, 'processFileBlock').mockImplementation(

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

@@ -1,4 +1,3 @@
-import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
@@ -53,9 +52,6 @@ describe('read_docs tool with researcher agent (via web API facade)', () => {
     spyOn(analytics, 'flushAnalytics').mockImplementation(() =>
       Promise.resolve(),
     )
-    spyOn(bigquery, 'insertTrace').mockImplementation(() =>
-      Promise.resolve(true),
-    )
 
     agentRuntimeImpl.requestFiles = async () => ({})
     agentRuntimeImpl.requestOptionalFile = async () => null

packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts

@@ -1,4 +1,3 @@
-import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
@@ -70,11 +69,8 @@ describe('runAgentStep - set_output tool', () => {
       })),
     } as any)
 
-    // Mock analytics and tracing
+    // Mock analytics
     spyOn(analytics, 'trackEvent').mockImplementation(() => {})
-    spyOn(bigquery, 'insertTrace').mockImplementation(() =>
-      Promise.resolve(true),
-    )
 
     agentRuntimeImpl.requestFiles = async ({ filePaths }) => {
       const results: Record<string, string | null> = {}

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

@@ -1,4 +1,3 @@
-import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
@@ -72,11 +71,8 @@ describe('web_search tool with researcher agent (via web API facade)', () => {
       userInputId: 'test-input',
     }
 
-    // Mock analytics and tracing
+    // Mock analytics
     spyOn(analytics, 'trackEvent').mockImplementation(() => {})
-    spyOn(bigquery, 'insertTrace').mockImplementation(() =>
-      Promise.resolve(true),
-    )
 
     // Mock websocket actions
     runAgentStepBaseParams.requestFiles = async () => ({})

packages/agent-runtime/src/run-agent-step.ts

@@ -1,4 +1,3 @@
-import { insertTrace } from '@codebuff/bigquery'
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
 import { supportsCacheControl } from '@codebuff/common/old-constants'
 import { TOOLS_WHICH_WONT_FORCE_NEXT_STEP } from '@codebuff/common/tools/constants'
@@ -26,7 +25,6 @@ import {
 } from './util/messages'
 import { countTokensJson } from './util/token-counter'
 
-import type { AgentResponseTrace } from '@codebuff/bigquery'
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
 import type {
@@ -352,22 +350,6 @@ export const runAgentStep = async (
 
   fullResponse = fullResponseAfterStream
 
-  const agentResponseTrace: AgentResponseTrace = {
-    type: 'agent-response',
-    created_at: new Date(),
-    agent_step_id: agentStepId,
-    user_id: userId ?? '',
-    id: crypto.randomUUID(),
-    payload: {
-      output: fullResponse,
-      user_input_id: userInputId,
-      client_session_id: clientSessionId,
-      fingerprint_id: fingerprintId,
-    },
-  }
-
-  insertTrace({ trace: agentResponseTrace, logger })
-
   agentState.messageHistory = expireMessages(
     agentState.messageHistory,
     'agentStep',

packages/agent-runtime/src/system-prompt/search-system-prompt.ts

@@ -1,5 +1,4 @@
 import { countTokens, countTokensJson } from '../util/token-counter'
-import { insertTrace } from '@codebuff/bigquery'
 import { buildArray } from '@codebuff/common/util/array'
 
 import {
@@ -63,24 +62,6 @@ export function getSearchSystemPrompt(params: {
     },
     {} as Record<number, string>,
   )
-
-  const trace = {
-    id: crypto.randomUUID(),
-    agent_step_id: options.agentStepId,
-    created_at: new Date(),
-    type: 'file-trees' as const,
-    user_id: options.userId ?? '',
-    payload: {
-      filetrees: truncatedTrees,
-      user_input_id: options.userInputId,
-      client_session_id: options.clientSessionId,
-      fingerprint_id: options.fingerprintId,
-    },
-  }
-
-  insertTrace({ trace, logger }).catch((error: Error) => {
-    logger.error({ error }, 'Failed to insert file trees trace')
-  })
   const fileTreeTokens = countTokensJson(projectFileTreePrompt)
 
   const systemInfoPrompt = getSystemInfoPrompt(fileContext)

packages/bigquery/src/client.ts

@@ -162,7 +162,7 @@ export async function insertTrace({
     return true
   } catch (error) {
     logger.warn(
-      { error, traceId: trace.id },
+      { error: getErrorObject(error), traceId: trace.id },
       'Failed to insert trace into BigQuery',
     )
     return false