add promptAiSdk to AgentRuntimeDeps

Author: charleslien

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

@@ -15,16 +15,18 @@ import {
 } from 'bun:test'
 
 import * as checkTerminalCommandModule from '../check-terminal-command'
-import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { mainPrompt } from '../main-prompt'
 import * as websocketAction from '../websockets/websocket-action'
 
+import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 
 // --- Shared Mocks & Helpers ---
 
+let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
+
 class MockWebSocket {
   send(msg: string) {}
   close() {}
@@ -103,6 +105,7 @@ describe.skip('mainPrompt (Integration)', () => {
 
   afterEach(() => {
     mock.restore()
+    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
   })
 
   it('should delete a specified function while preserving other code', async () => {
@@ -337,7 +340,9 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
     )
 
     // Mock LLM calls
-    spyOn(aisdk, 'promptAiSdk').mockResolvedValue('Mocked non-stream AiSdk')
+    agentRuntimeImpl.promptAiSdk = async function () {
+      return 'Mocked non-stream AiSdk'
+    }
 
     const sessionState = getInitialSessionState(mockFileContext)
     sessionState.mainAgentState.messageHistory.push(
@@ -377,7 +382,7 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
     }
 
     const { output, sessionState: finalSessionState } = await mainPrompt({
-      ...TEST_AGENT_RUNTIME_IMPL,
+      ...agentRuntimeImpl,
       ws: new MockWebSocket() as unknown as WebSocket,
       action,
       userId: TEST_USER_ID,
@@ -419,7 +424,9 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
       ).mockResolvedValue(null)
 
       // Mock LLM calls
-      spyOn(aisdk, 'promptAiSdk').mockResolvedValue('Mocked non-stream AiSdk')
+      agentRuntimeImpl.promptAiSdk = async function () {
+        return 'Mocked non-stream AiSdk'
+      }
 
       const sessionState = getInitialSessionState(mockFileContext)
       sessionState.mainAgentState.messageHistory.push(
@@ -459,7 +466,7 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
       }
 
       await mainPrompt({
-        ...TEST_AGENT_RUNTIME_IMPL,
+        ...agentRuntimeImpl,
         ws: new MockWebSocket() as unknown as WebSocket,
         action,
         userId: TEST_USER_ID,

backend/src/__tests__/malformed-tool-call.test.ts

@@ -16,17 +16,19 @@ import {
 } from 'bun:test'
 
 import { MockWebSocket, mockFileContext } from './test-utils'
-import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { processStreamWithTools } from '../tools/stream-parser'
 import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '../templates/types'
+import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 import type {
   Message,
   ToolMessage,
 } from '@codebuff/common/types/messages/codebuff-message'
 import type { WebSocket } from 'ws'
 
+let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
+
 describe('malformed tool call error handling', () => {
   let testAgent: AgentTemplate
   let mockWs: MockWebSocket
@@ -72,16 +74,17 @@ describe('malformed tool call error handling', () => {
     }))
 
     // Mock LLM APIs
-    spyOn(aisdk, 'promptAiSdk').mockImplementation(() =>
-      Promise.resolve('Test response'),
-    )
+    agentRuntimeImpl.promptAiSdk = async function () {
+      return 'Test response'
+    }
 
     // Mock generateCompactId for consistent test results
     spyOn(stringUtils, 'generateCompactId').mockReturnValue('test-tool-call-id')
   })
 
   afterEach(() => {
     mock.restore()
+    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
   })
 
   function createMockStream(chunks: string[]) {

backend/src/__tests__/process-file-block.test.ts

@@ -1,4 +1,5 @@
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import {
   clearMockedModules,
   mockModule,
@@ -9,14 +10,9 @@ import { applyPatch } from 'diff'
 
 import { processFileBlock } from '../process-file-block'
 
-import type { Logger } from '@codebuff/common/types/contracts/logger'
+import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 
-const logger: Logger = {
-  debug: () => {},
-  info: () => {},
-  warn: () => {},
-  error: () => {},
-}
+let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
 
 describe('processFileBlockModule', () => {
   beforeAll(() => {
@@ -74,6 +70,7 @@ describe('processFileBlockModule', () => {
       const expectedContent = 'function test() {\n  return true;\n}'
 
       const result = await processFileBlock({
+        ...agentRuntimeImpl,
         path: 'test.ts',
         instructions: undefined,
         initialContentPromise: Promise.resolve(null),
@@ -85,7 +82,6 @@ describe('processFileBlockModule', () => {
         fingerprintId: 'fingerprintId',
         userInputId: 'userInputId',
         userId: TEST_USER_ID,
-        logger,
       })
 
       expect(result).not.toBeNull()
@@ -111,6 +107,7 @@ describe('processFileBlockModule', () => {
         '}\r\n'
 
       const result = await processFileBlock({
+        ...agentRuntimeImpl,
         path: 'test.ts',
         instructions: undefined,
         initialContentPromise: Promise.resolve(oldContent),
@@ -122,7 +119,6 @@ describe('processFileBlockModule', () => {
         fingerprintId: 'fingerprintId',
         userInputId: 'userInputId',
         userId: TEST_USER_ID,
-        logger,
       })
 
       expect(result).not.toBeNull()
@@ -144,6 +140,7 @@ describe('processFileBlockModule', () => {
       const newContent = 'function test() {\n  return true;\n}\n'
 
       const result = await processFileBlock({
+        ...agentRuntimeImpl,
         path: 'test.ts',
         instructions: undefined,
         initialContentPromise: Promise.resolve(oldContent),
@@ -155,7 +152,6 @@ describe('processFileBlockModule', () => {
         fingerprintId: 'fingerprintId',
         userInputId: 'userInputId',
         userId: TEST_USER_ID,
-        logger,
       })
 
       expect(result).not.toBeNull()
@@ -170,6 +166,7 @@ describe('processFileBlockModule', () => {
       const newContent = 'const x = 1;\r\nconst z = 3;\r\n'
 
       const result = await processFileBlock({
+        ...agentRuntimeImpl,
         path: 'test.ts',
         instructions: undefined,
         initialContentPromise: Promise.resolve(oldContent),
@@ -181,7 +178,6 @@ describe('processFileBlockModule', () => {
         fingerprintId: 'fingerprintId',
         userInputId: 'userInputId',
         userId: TEST_USER_ID,
-        logger,
       })
 
       expect(result).not.toBeNull()
@@ -217,6 +213,7 @@ describe('processFileBlockModule', () => {
         '// ... existing code ...\nconst x = 1;\n// ... existing code ...'
 
       const result = await processFileBlock({
+        ...agentRuntimeImpl,
         path: 'test.ts',
         instructions: undefined,
         initialContentPromise: Promise.resolve(null),
@@ -228,7 +225,6 @@ describe('processFileBlockModule', () => {
         fingerprintId: 'fingerprintId',
         userInputId: 'userInputId',
         userId: TEST_USER_ID,
-        logger,
       })
 
       expect(result).not.toBeNull()

backend/src/__tests__/run-agent-step-tools.test.ts

@@ -18,7 +18,6 @@ import {
 
 // Mock imports
 import * as liveUserInputs from '../live-user-inputs'
-import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { runAgentStep } from '../run-agent-step'
 import { clearAgentGeneratorCache } from '../run-programmatic-step'
 import { asUserMessage } from '../util/messages'
@@ -106,9 +105,9 @@ describe('runAgentStep - set_output tool', () => {
     // Don't mock requestToolCall for integration test - let real tool execution happen
 
     // Mock LLM APIs
-    spyOn(aisdk, 'promptAiSdk').mockImplementation(() =>
-      Promise.resolve('Test response'),
-    )
+    agentRuntimeImpl.promptAiSdk = async function () {
+      return 'Test response'
+    }
     clearAgentGeneratorCache(agentRuntimeImpl)
   })
 

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

@@ -23,7 +23,6 @@ import * as requestFilesPrompt from '../find-files/request-files-prompt'
 import * as liveUserInputs from '../live-user-inputs'
 import { MockWebSocket, mockFileContext } from './test-utils'
 import * as linkupApi from '../llm-apis/linkup-api'
-import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { runAgentStep } from '../run-agent-step'
 import { assembleLocalAgentTemplates } from '../templates/agent-registry'
 import * as websocketAction from '../websockets/websocket-action'
@@ -67,9 +66,9 @@ describe('web_search tool with researcher agent', () => {
     }))
 
     // Mock LLM APIs
-    spyOn(aisdk, 'promptAiSdk').mockImplementation(() =>
-      Promise.resolve('Test response'),
-    )
+    agentRuntimeImpl.promptAiSdk = async function () {
+      return 'Test response'
+    }
 
     // Mock other required modules
     spyOn(requestFilesPrompt, 'requestRelevantFiles').mockImplementation(

backend/src/impl/agent-runtime.ts

@@ -1,5 +1,8 @@
 import { addAgentStep, finishAgentRun, startAgentRun } from '../agent-run'
-import { promptAiSdkStream } from '../llm-apis/vercel-ai-sdk/ai-sdk'
+import {
+  promptAiSdk,
+  promptAiSdkStream,
+} from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from '../util/logger'
 
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
@@ -12,4 +15,5 @@ export const BACKEND_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
   addAgentStep,
 
   promptAiSdkStream,
+  promptAiSdk,
 })

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

@@ -17,12 +17,12 @@ import { openRouterLanguageModel } from '../openrouter'
 import { vertexFinetuned } from './vertex-finetuned'
 
 import type { Model, OpenAIModel } from '@codebuff/common/old-constants'
-import type { PromptAiSdkStreamFn } from '@codebuff/common/types/contracts/llm'
-import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type {
-  ParamsExcluding,
-  ParamsOf,
-} from '@codebuff/common/types/function-params'
+  PromptAiSdkFn,
+  PromptAiSdkStreamFn,
+} from '@codebuff/common/types/contracts/llm'
+import type { Logger } from '@codebuff/common/types/contracts/logger'
+import type { ParamsOf } from '@codebuff/common/types/function-params'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type {
   OpenRouterProviderOptions,
@@ -231,21 +231,8 @@ export const promptAiSdkStream = async function* (
 
 // TODO: figure out a nice way to unify stream & non-stream versions maybe?
 export const promptAiSdk = async function (
-  params: {
-    messages: Message[]
-    clientSessionId: string
-    fingerprintId: string
-    userInputId: string
-    model: Model
-    userId: string | undefined
-    chargeUser?: boolean
-    agentId?: string
-    onCostCalculated?: (credits: number) => Promise<void>
-    includeCacheControl?: boolean
-    maxRetries?: number
-    logger: Logger
-  } & ParamsExcluding<typeof generateText, 'model' | 'messages'>,
-): Promise<string> {
+  params: ParamsOf<PromptAiSdkFn>,
+): ReturnType<PromptAiSdkFn> {
   const { logger } = params
 
   if (!checkLiveUserInput(params)) {

backend/src/process-file-block.ts

@@ -8,26 +8,32 @@ import {
   parseAndGetDiffBlocksSingleFile,
   retryDiffBlocksPrompt,
 } from './generate-diffs-prompt'
-import { promptAiSdk } from './llm-apis/vercel-ai-sdk/ai-sdk'
 import { countTokens } from './util/token-counter'
 
-import type { Message } from '@codebuff/common/types/messages/codebuff-message'
+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'
+import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 
-export async function processFileBlock(params: {
-  path: string
-  instructions: string | undefined
-  initialContentPromise: Promise<string | null>
-  newContent: string
-  messages: Message[]
-  fullResponse: string
-  lastUserPrompt: string | undefined
-  clientSessionId: string
-  fingerprintId: string
-  userInputId: string
-  userId: string | undefined
-  logger: Logger
-}): Promise<
+export async function processFileBlock(
+  params: {
+    path: string
+    instructions: string | undefined
+    initialContentPromise: Promise<string | null>
+    newContent: string
+    messages: Message[]
+    fullResponse: string
+    lastUserPrompt: string | undefined
+    clientSessionId: string
+    fingerprintId: string
+    userInputId: string
+    userId: string | undefined
+    logger: Logger
+  } & ParamsExcluding<
+    typeof handleLargeFile,
+    'oldContent' | 'editSnippet' | 'filePath'
+  >,
+): Promise<
   | {
       tool: 'write_file'
       path: string
@@ -113,14 +119,10 @@ export async function processFileBlock(params: {
   )
   if (tokenCount > LARGE_FILE_TOKEN_LIMIT) {
     const largeFileContent = await handleLargeFile({
+      ...params,
       oldContent: normalizedInitialContent,
       editSnippet: normalizedEditSnippet,
-      clientSessionId,
-      fingerprintId,
-      userInputId,
-      userId,
       filePath: path,
-      logger,
     })
 
     if (!largeFileContent) {
@@ -239,6 +241,7 @@ export async function handleLargeFile(params: {
   userId: string | undefined
   filePath: string
   logger: Logger
+  promptAiSdk: PromptAiSdkFn
 }): Promise<string | null> {
   const {
     oldContent,
@@ -248,6 +251,7 @@ export async function handleLargeFile(params: {
     userInputId,
     userId,
     filePath,
+    promptAiSdk,
     logger,
   } = params
   const startTime = Date.now()