change promptAiSdk functions to function instead of const

Author: charleslien

backend/src/__tests__/cost-aggregation.integration.test.ts

@@ -99,9 +99,10 @@ class MockWebSocket {
 describe('Cost Aggregation Integration Tests', () => {
   let mockLocalAgentTemplates: Record<string, any>
   let mockWebSocket: MockWebSocket
-  let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
+  let agentRuntimeImpl: AgentRuntimeDeps
 
   beforeEach(async () => {
+    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
     mockWebSocket = new MockWebSocket()
 
     // Setup mock agent templates
@@ -230,7 +231,6 @@ describe('Cost Aggregation Integration Tests', () => {
 
   afterEach(() => {
     mock.restore()
-    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
   })
 
   it('should correctly aggregate costs across the entire main prompt flow', async () => {

backend/src/__tests__/main-prompt.test.ts

@@ -28,11 +28,10 @@ import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
-import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 
-let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
+let agentRuntimeImpl: AgentRuntimeDeps
 
 const mockAgentStream = (streamOutput: string) => {
   agentRuntimeImpl.promptAiSdkStream = async function* ({}) {
@@ -43,14 +42,10 @@ const mockAgentStream = (streamOutput: string) => {
 
 describe('mainPrompt', () => {
   let mockLocalAgentTemplates: Record<string, any>
-  const logger: Logger = {
-    debug: () => {},
-    error: () => {},
-    info: () => {},
-    warn: () => {},
-  }
 
   beforeEach(() => {
+    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
+
     // Setup common mock agent templates
     mockLocalAgentTemplates = {
       [AgentTemplateTypes.base]: {
@@ -86,12 +81,10 @@ describe('mainPrompt', () => {
         stepPrompt: '',
       } satisfies AgentTemplate,
     }
-  })
 
-  beforeEach(() => {
     // Mock analytics and tracing
     spyOn(analytics, 'initAnalytics').mockImplementation(() => {})
-    analytics.initAnalytics({ logger }) // Initialize the mock
+    analytics.initAnalytics(agentRuntimeImpl) // Initialize the mock
     spyOn(analytics, 'trackEvent').mockImplementation(() => {})
     spyOn(bigquery, 'insertTrace').mockImplementation(() =>
       Promise.resolve(true),
@@ -176,7 +169,6 @@ describe('mainPrompt', () => {
   afterEach(() => {
     // Clear all mocks after each test
     mock.restore()
-    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
   })
 
   class MockWebSocket {

backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts

@@ -53,7 +53,7 @@ const modelToAiSDKModel = (model: Model): LanguageModel => {
 // TODO: Add retries & fallbacks: likely by allowing this to instead of "model"
 // also take an array of form [{model: Model, retries: number}, {model: Model, retries: number}...]
 // eg: [{model: "gemini-2.0-flash-001"}, {model: "vertex/gemini-2.0-flash-001"}, {model: "claude-3-5-haiku", retries: 3}]
-export const promptAiSdkStream = async function* (
+export async function* promptAiSdkStream(
   params: ParamsOf<PromptAiSdkStreamFn>,
 ): ReturnType<PromptAiSdkStreamFn> {
   const { logger } = params
@@ -230,7 +230,7 @@ export const promptAiSdkStream = async function* (
 }
 
 // TODO: figure out a nice way to unify stream & non-stream versions maybe?
-export const promptAiSdk = async function (
+export async function promptAiSdk(
   params: ParamsOf<PromptAiSdkFn>,
 ): ReturnType<PromptAiSdkFn> {
   const { logger } = params
@@ -287,7 +287,7 @@ export const promptAiSdk = async function (
 }
 
 // Copied over exactly from promptAiSdk but with a schema
-export const promptAiSdkStructured = async function <T>(
+export async function promptAiSdkStructured<T>(
   params: PromptAiSdkStructuredInput<T>,
 ): PromptAiSdkStructuredOutput<T> {
   const { logger } = params