move callMainPrompt to agent-runtime

charleslien · charleslien · commit 43cac13d9a83 · 2025-10-23T10:40:10.000-07:00
diff --git a/backend/src/__tests__/cost-aggregation.integration.test.ts b/backend/src/__tests__/cost-aggregation.integration.test.ts
@@ -1,5 +1,5 @@
 import { disableLiveUserInputCheck } from '@codebuff/agent-runtime/live-user-inputs'
-import { mainPrompt } from '@codebuff/agent-runtime/main-prompt'
+import { callMainPrompt, mainPrompt } from '@codebuff/agent-runtime/main-prompt'
 import * as agentRegistry from '@codebuff/agent-runtime/templates/agent-registry'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
@@ -15,9 +15,6 @@ import {
   beforeAll,
 } from 'bun:test'
 
-import * as messageCostTracker from '../llm-apis/message-cost-tracker'
-import * as websocketAction from '../websockets/websocket-action'
-
 import type { AgentTemplate } from '@codebuff/agent-runtime/templates/types'
 import type { ServerAction } from '@codebuff/common/actions'
 import type {
@@ -135,14 +132,6 @@ describe('Cost Aggregation Integration Tests', () => {
       } satisfies AgentTemplate,
     }
 
-    // Mock cost tracking to return 0 so only onCostCalculated contributes
-    spyOn(messageCostTracker, 'saveMessage').mockImplementation(
-      async (value) => {
-        // Return 0 so we can control costs only via onCostCalculated
-        return 0
-      },
-    )
-
     // Mock LLM streaming
     let callCount = 0
     const creditHistory: number[] = []
@@ -271,7 +260,7 @@ describe('Cost Aggregation Integration Tests', () => {
     }
 
     // Call through websocket action handler to test full integration
-    await websocketAction.callMainPrompt({
+    await callMainPrompt({
       ...agentRuntimeImpl,
       repoId: undefined,
       repoUrl: undefined,
@@ -424,17 +413,6 @@ describe('Cost Aggregation Integration Tests', () => {
   it('should not double-count costs in complex scenarios', async () => {
     // Track all saveMessage calls to ensure no duplication
     const saveMessageCalls: any[] = []
-    spyOn(messageCostTracker, 'saveMessage').mockImplementation(
-      async (value) => {
-        saveMessageCalls.push({
-          messageId: value.messageId,
-          model: value.model,
-          inputTokens: value.inputTokens,
-          outputTokens: value.outputTokens,
-        })
-        return 8 // Each LLM call costs 8 credits
-      },
-    )
 
     const sessionState = getInitialSessionState(mockFileContext)
     sessionState.mainAgentState.agentType = 'base'
@@ -490,7 +468,7 @@ describe('Cost Aggregation Integration Tests', () => {
     }
 
     // Call through websocket action to test server-side reset
-    await websocketAction.callMainPrompt({
+    await callMainPrompt({
       ...agentRuntimeImpl,
       repoId: undefined,
       repoUrl: undefined,
diff --git a/backend/src/websockets/websocket-action.ts b/backend/src/websockets/websocket-action.ts
@@ -1,10 +1,7 @@
 import {
   cancelUserInput,
-  checkLiveUserInput,
   startUserInput,
 } from '@codebuff/agent-runtime/live-user-inputs'
-import { mainPrompt } from '@codebuff/agent-runtime/main-prompt'
-import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import { calculateUsageAndBalance } from '@codebuff/billing'
 import { trackEvent } from '@codebuff/common/analytics'
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
@@ -19,13 +16,13 @@ import { getRequestContext } from './request-context'
 import { withLoggerContext } from '../util/logger'
 
 import type { ClientAction, UsageResponse } from '@codebuff/common/actions'
-import type { SendActionFn } from '@codebuff/common/types/contracts/client'
 import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type { UserInputRecord } from '@codebuff/common/types/contracts/live-user-input'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { ClientMessage } from '@codebuff/common/websockets/websocket-schema'
 import type { WebSocket } from 'ws'
+import { callMainPrompt } from '@codebuff/agent-runtime/main-prompt'
 
 /**
  * Generates a usage response object for the client
@@ -181,101 +178,6 @@ const onPrompt = async (
   )
 }
 
-export const callMainPrompt = async (
-  params: {
-    action: ClientAction<'prompt'>
-    userId: string
-    promptId: string
-    clientSessionId: string
-    sendAction: SendActionFn
-    liveUserInputRecord: UserInputRecord
-    logger: Logger
-  } & ParamsExcluding<
-    typeof mainPrompt,
-    'localAgentTemplates' | 'onResponseChunk'
-  >,
-) => {
-  const { action, userId, promptId, clientSessionId, sendAction, logger } =
-    params
-  const { fileContext } = action.sessionState
-
-  // Enforce server-side state authority: reset creditsUsed to 0
-  // The server controls cost tracking, clients cannot manipulate this value
-  action.sessionState.mainAgentState.creditsUsed = 0
-  action.sessionState.mainAgentState.directCreditsUsed = 0
-
-  // Assemble local agent templates from fileContext
-  const { agentTemplates: localAgentTemplates, validationErrors } =
-    assembleLocalAgentTemplates({ fileContext, logger })
-
-  if (validationErrors.length > 0) {
-    sendAction({
-      action: {
-        type: 'prompt-error',
-        message: `Invalid agent config: ${validationErrors.map((err) => err.message).join('\n')}`,
-        userInputId: promptId,
-      },
-    })
-  }
-
-  sendAction({
-    action: {
-      type: 'response-chunk',
-      userInputId: promptId,
-      chunk: {
-        type: 'start',
-        agentId: action.sessionState.mainAgentState.agentType ?? undefined,
-        messageHistoryLength:
-          action.sessionState.mainAgentState.messageHistory.length,
-      },
-    },
-  })
-
-  const result = await mainPrompt({
-    ...params,
-    localAgentTemplates,
-    onResponseChunk: (chunk) => {
-      if (checkLiveUserInput({ ...params, userInputId: promptId })) {
-        sendAction({
-          action: {
-            type: 'response-chunk',
-            userInputId: promptId,
-            chunk,
-          },
-        })
-      }
-    },
-  })
-
-  const { sessionState, output } = result
-
-  sendAction({
-    action: {
-      type: 'response-chunk',
-      userInputId: promptId,
-      chunk: {
-        type: 'finish',
-        agentId: sessionState.mainAgentState.agentType ?? undefined,
-        totalCost: sessionState.mainAgentState.creditsUsed,
-      },
-    },
-  })
-
-  // Send prompt data back
-  sendAction({
-    action: {
-      type: 'prompt-response',
-      promptId,
-      sessionState,
-      toolCalls: [],
-      toolResults: [],
-      output,
-    },
-  })
-
-  return result
-}
-
 /**
  * Handles initialization actions from the client
  * @param fileContext - The file context information
diff --git a/packages/agent-runtime/src/main-prompt.ts b/packages/agent-runtime/src/main-prompt.ts
@@ -3,14 +3,22 @@ import { generateCompactId } from '@codebuff/common/util/string'
 import { uniq } from 'lodash'
 
 import { checkTerminalCommand } from './check-terminal-command'
+import { checkLiveUserInput } from './live-user-inputs'
 import { loopAgentSteps } from './run-agent-step'
-import { getAgentTemplate } from './templates/agent-registry'
+import {
+  assembleLocalAgentTemplates,
+  getAgentTemplate,
+} from './templates/agent-registry'
 import { expireMessages } from './util/messages'
 
 import type { AgentTemplate } from './templates/types'
 import type { ClientAction } from '@codebuff/common/actions'
 import type { CostMode } from '@codebuff/common/old-constants'
-import type { RequestToolCallFn } from '@codebuff/common/types/contracts/client'
+import type {
+  RequestToolCallFn,
+  SendActionFn,
+} from '@codebuff/common/types/contracts/client'
+import type { UserInputRecord } from '@codebuff/common/types/contracts/live-user-input'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
@@ -20,7 +28,7 @@ import type {
   AgentOutput,
 } from '@codebuff/common/types/session-state'
 
-export const mainPrompt = async (
+export async function mainPrompt(
   params: {
     action: ClientAction<'prompt'>
 
@@ -48,7 +56,7 @@ export const mainPrompt = async (
 ): Promise<{
   sessionState: SessionState
   output: AgentOutput
-}> => {
+}> {
   const { action, localAgentTemplates, requestToolCall, logger } = params
 
   const {
@@ -239,3 +247,95 @@ export const mainPrompt = async (
     },
   }
 }
+
+export async function callMainPrompt(
+  params: {
+    action: ClientAction<'prompt'>
+    promptId: string
+    sendAction: SendActionFn
+    liveUserInputRecord: UserInputRecord
+    logger: Logger
+  } & ParamsExcluding<
+    typeof mainPrompt,
+    'localAgentTemplates' | 'onResponseChunk'
+  >,
+) {
+  const { action, promptId, sendAction, logger } = params
+  const { fileContext } = action.sessionState
+
+  // Enforce server-side state authority: reset creditsUsed to 0
+  // The server controls cost tracking, clients cannot manipulate this value
+  action.sessionState.mainAgentState.creditsUsed = 0
+  action.sessionState.mainAgentState.directCreditsUsed = 0
+
+  // Assemble local agent templates from fileContext
+  const { agentTemplates: localAgentTemplates, validationErrors } =
+    assembleLocalAgentTemplates({ fileContext, logger })
+
+  if (validationErrors.length > 0) {
+    sendAction({
+      action: {
+        type: 'prompt-error',
+        message: `Invalid agent config: ${validationErrors.map((err) => err.message).join('\n')}`,
+        userInputId: promptId,
+      },
+    })
+  }
+
+  sendAction({
+    action: {
+      type: 'response-chunk',
+      userInputId: promptId,
+      chunk: {
+        type: 'start',
+        agentId: action.sessionState.mainAgentState.agentType ?? undefined,
+        messageHistoryLength:
+          action.sessionState.mainAgentState.messageHistory.length,
+      },
+    },
+  })
+
+  const result = await mainPrompt({
+    ...params,
+    localAgentTemplates,
+    onResponseChunk: (chunk) => {
+      if (checkLiveUserInput({ ...params, userInputId: promptId })) {
+        sendAction({
+          action: {
+            type: 'response-chunk',
+            userInputId: promptId,
+            chunk,
+          },
+        })
+      }
+    },
+  })
+
+  const { sessionState, output } = result
+
+  sendAction({
+    action: {
+      type: 'response-chunk',
+      userInputId: promptId,
+      chunk: {
+        type: 'finish',
+        agentId: sessionState.mainAgentState.agentType ?? undefined,
+        totalCost: sessionState.mainAgentState.creditsUsed,
+      },
+    },
+  })
+
+  // Send prompt data back
+  sendAction({
+    action: {
+      type: 'prompt-response',
+      promptId,
+      sessionState,
+      toolCalls: [],
+      toolResults: [],
+      output,
+    },
+  })
+
+  return result
+}
