test(agent-runtime): Update test mocks to use PromptResult type

Author: brandonkachen

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

@@ -20,6 +20,8 @@
 
 import { mock } from 'bun:test'
 
+import { promptSuccess } from '../../util/error'
+
 import type { ProjectFileContext } from '../../util/file'
 
 export const mockFileContext: ProjectFileContext = {
@@ -248,11 +250,11 @@ export function createTestAgentRuntimeParams(
           toolCallId: 'mock-id',
           input: {},
         }
-        return 'mock-message-id'
+        return promptSuccess('mock-message-id')
       }),
-    promptAiSdk: overrides.promptAiSdk ?? mock(async () => 'Mock response'),
+    promptAiSdk: overrides.promptAiSdk ?? mock(async () => promptSuccess('Mock response')),
     promptAiSdkStructured:
-      overrides.promptAiSdkStructured ?? mock(async () => ({})),
+      overrides.promptAiSdkStructured ?? mock(async () => promptSuccess({})),
     requestMcpToolData: overrides.requestMcpToolData ?? mock(async () => ({})),
     startAgentRun: overrides.startAgentRun ?? mock(async () => 'test-run-id'),
     finishAgentRun: overrides.finishAgentRun ?? mock(async () => {}),
@@ -303,10 +305,10 @@ export function createTestAgentRuntimeDeps(): Omit<
         toolCallId: 'mock-id',
         input: {},
       }
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }),
-    promptAiSdk: mock(async () => 'Mock response'),
-    promptAiSdkStructured: mock(async () => ({})),
+    promptAiSdk: mock(async () => promptSuccess('Mock response')),
+    promptAiSdkStructured: mock(async () => promptSuccess({})),
     requestMcpToolData: mock(async () => ({})),
     startAgentRun: mock(async () => 'test-run-id'),
     finishAgentRun: mock(async () => {}),

packages/agent-runtime/src/__tests__/fast-rewrite.test.ts

@@ -6,6 +6,7 @@ import {
   clearMockedModules,
   mockModule,
 } from '@codebuff/common/testing/mock-modules'
+import { promptSuccess } from '@codebuff/common/util/error'
 import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'bun:test'
 import { createPatch } from 'diff'
 
@@ -48,7 +49,7 @@ describe('rewriteWithOpenAI', () => {
 
     agentRuntimeImpl.promptAiSdk = async (params: any) => {
       capturedPromptText = params?.messages?.[0]?.content?.[0]?.text
-      return expectedResult.replace(/\n$/, '')
+      return promptSuccess(expectedResult.replace(/\n$/, ''))
     }
 
     const result = await rewriteWithOpenAI({

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

@@ -6,6 +6,7 @@ import {
 } from '@codebuff/common/testing/mock-modules'
 import { setupDbSpies } from '@codebuff/common/testing/mocks/database'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
+import { promptSuccess } from '@codebuff/common/util/error'
 import { assistantMessage, userMessage } from '@codebuff/common/util/messages'
 import db from '@codebuff/internal/db'
 import {
@@ -67,7 +68,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       llmCallCount++
       yield { type: 'text' as const, text: 'LLM response\n\n' }
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     })
 
     // Mock analytics
@@ -486,7 +487,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       // LLM always tries to end turn
       yield { type: 'text' as const, text: 'LLM response\n\n' }
       yield createToolCallChunk('end_turn', {})
-      return `mock-message-id-${promptCallCount}`
+      return promptSuccess(`mock-message-id-${promptCallCount}`)
     }
 
     await loopAgentSteps({
@@ -558,7 +559,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
         yield { type: 'text' as const, text: 'Ending\n\n' }
         yield createToolCallChunk('end_turn', {})
       }
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     mockAgentState.output = undefined
@@ -621,7 +622,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       yield createToolCallChunk('set_output', { result: 'success' })
       yield { type: 'text' as const, text: '\n\n' }
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     mockAgentState.output = undefined
@@ -659,13 +660,13 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
     // Mock promptAiSdk to capture the n parameter
     loopAgentStepsBaseParams.promptAiSdk = async (params: any) => {
       agentStepN = params.n
-      return JSON.stringify([
+      return promptSuccess(JSON.stringify([
         'Response 1',
         'Response 2',
         'Response 3',
         'Response 4',
         'Response 5',
-      ])
+      ]))
     }
 
     await loopAgentSteps({
@@ -705,7 +706,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       'Implementation C',
     ]
     loopAgentStepsBaseParams.promptAiSdk = async () => {
-      return JSON.stringify(expectedResponses)
+      return promptSuccess(JSON.stringify(expectedResponses))
     }
 
     await loopAgentSteps({
@@ -735,7 +736,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       llmCallNumber++
       yield { type: 'text' as const, text: 'Response without output\n\n' }
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const result = await loopAgentSteps({
@@ -788,7 +789,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
         yield { type: 'text' as const, text: '\n\n' }
         yield createToolCallChunk('end_turn', {})
       }
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     mockAgentState.output = undefined

packages/agent-runtime/src/__tests__/prompt-caching-subagents.test.ts

@@ -1,6 +1,7 @@
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
+import { promptSuccess } from '@codebuff/common/util/error'
 import { assistantMessage, userMessage } from '@codebuff/common/util/messages'
 import { beforeEach, describe, expect, it } from 'bun:test'
 
@@ -107,7 +108,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
           await options.onCostCalculated(1)
         }
 
-        return 'mock-message-id'
+        return promptSuccess('mock-message-id')
       },
       // Mock file operations
       requestFiles: async ({ filePaths }) => {

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

@@ -2,6 +2,7 @@ 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'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
+import { promptSuccess } from '@codebuff/common/util/error'
 import {
   afterEach,
 
@@ -38,7 +39,7 @@ function mockAgentStream(chunks: StreamChunk[]) {
     for (const chunk of chunks) {
       yield chunk
     }
-    return 'mock-message-id'
+    return promptSuccess('mock-message-id')
   }
   agentRuntimeImpl.promptAiSdkStream = mockPromptAiSdkStream
   runAgentStepBaseParams.promptAiSdkStream = mockPromptAiSdkStream

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

@@ -3,6 +3,7 @@ import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { setupDbSpies } from '@codebuff/common/testing/mocks/database'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
+import { promptSuccess } from '@codebuff/common/util/error'
 import { assistantMessage, userMessage } from '@codebuff/common/util/messages'
 import db from '@codebuff/internal/db'
 import {
@@ -96,7 +97,7 @@ describe('runAgentStep - set_output tool', () => {
 
     // Mock LLM APIs
     agentRuntimeImpl.promptAiSdk = async function () {
-      return 'Test response'
+      return promptSuccess('Test response')
     }
     clearAgentGeneratorCache(agentRuntimeImpl)
 
@@ -161,7 +162,7 @@ describe('runAgentStep - set_output tool', () => {
       yield createToolCallChunk('set_output', { message: 'Hi' })
       yield { type: 'text' as const, text: '\n\n' }
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)
@@ -193,7 +194,7 @@ describe('runAgentStep - set_output tool', () => {
         findings: ['Bug in auth.ts', 'Missing validation'],
       })
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)
@@ -226,7 +227,7 @@ describe('runAgentStep - set_output tool', () => {
         existingField: 'updated value',
       })
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)
@@ -259,7 +260,7 @@ describe('runAgentStep - set_output tool', () => {
     runAgentStepBaseParams.promptAiSdkStream = async function* ({}) {
       yield createToolCallChunk('set_output', {})
       yield createToolCallChunk('end_turn', {})
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)
@@ -331,7 +332,7 @@ describe('runAgentStep - set_output tool', () => {
     // Mock the LLM stream to return a response that doesn't end the turn
     runAgentStepBaseParams.promptAiSdkStream = async function* ({}) {
       yield { type: 'text' as const, text: 'Continuing with the analysis...' } // Non-empty response, no tool calls
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)
@@ -475,7 +476,7 @@ describe('runAgentStep - set_output tool', () => {
         agent_type: 'message-deleter-agent',
         prompt: 'Delete the last two assistant messages',
       })
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)

packages/agent-runtime/src/__tests__/test-utils.ts

@@ -1,6 +1,8 @@
+import { promptSuccess } from '@codebuff/common/util/error'
 import { generateCompactId } from '@codebuff/common/util/string'
 
 import type { StreamChunk } from '@codebuff/common/types/contracts/llm'
+import type { PromptResult } from '@codebuff/common/util/error'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 
 /**
@@ -26,16 +28,19 @@ export function createToolCallChunk<T extends string>(
  */
 export function createMockStreamWithToolCalls(
   chunks: (string | { toolName: string; input: Record<string, unknown> })[],
-): AsyncGenerator<StreamChunk, string | null> {
-  async function* generator(): AsyncGenerator<StreamChunk, string | null> {
+): AsyncGenerator<StreamChunk, PromptResult<string | null>> {
+  async function* generator(): AsyncGenerator<
+    StreamChunk,
+    PromptResult<string | null>
+  > {
     for (const chunk of chunks) {
       if (typeof chunk === 'string') {
         yield { type: 'text' as const, text: chunk }
       } else {
         yield createToolCallChunk(chunk.toolName, chunk.input)
       }
     }
-    return 'mock-message-id'
+    return promptSuccess('mock-message-id')
   }
   return generator()
 }

packages/agent-runtime/src/__tests__/tool-stream-parser.test.ts

@@ -1,4 +1,5 @@
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
+import { promptSuccess } from '@codebuff/common/util/error'
 import { beforeEach, describe, expect, it } from 'bun:test'
 
 import { processStreamWithTools } from '../tool-stream-parser'
@@ -13,7 +14,7 @@ describe('processStreamWithTags', () => {
       yield chunk
     }
 
-    return 'mock-message-id'
+    return promptSuccess('mock-message-id')
   }
 
   function textChunk(text: string): StreamChunk {

packages/agent-runtime/src/__tests__/tool-validation-error.test.ts

@@ -1,5 +1,6 @@
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
+import { promptSuccess } from '@codebuff/common/util/error'
 import { beforeEach, describe, expect, it } from 'bun:test'
 
 import { mockFileContext } from './test-utils'
@@ -49,9 +50,9 @@ describe('tool validation error handling', () => {
       },
     }
 
-    async function* mockStream(): AsyncGenerator<StreamChunk, string | null> {
+    async function* mockStream() {
       yield invalidToolCallChunk
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)
@@ -158,9 +159,9 @@ describe('tool validation error handling', () => {
       },
     }
 
-    async function* mockStream(): AsyncGenerator<StreamChunk, string | null> {
+    async function* mockStream() {
       yield validToolCallChunk
-      return 'mock-message-id'
+      return promptSuccess('mock-message-id')
     }
 
     const sessionState = getInitialSessionState(mockFileContext)

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

@@ -2,7 +2,7 @@ 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'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
-import { success } from '@codebuff/common/util/error'
+import { promptSuccess, success } from '@codebuff/common/util/error'
 import {
   afterEach,
 
@@ -38,7 +38,7 @@ function mockAgentStream(chunks: StreamChunk[]) {
     for (const chunk of chunks) {
       yield chunk
     }
-    return 'mock-message-id'
+    return promptSuccess('mock-message-id')
   }
 }
 
@@ -83,7 +83,7 @@ describe('web_search tool with researcher agent (via web API facade)', () => {
 
     // Mock LLM APIs
     runAgentStepBaseParams.promptAiSdk = async function () {
-      return 'Test response'
+      return promptSuccess('Test response')
     }
   })