Small refactoring suggested by codebuff

Author: jahooma

backend/src/run-agent-step.ts

@@ -15,10 +15,11 @@ import { additionalSystemPrompts } from './system-prompt/prompts'
 import { getAgentTemplate } from './templates/agent-registry'
 import { getAgentPrompt } from './templates/strings'
 import { processStreamWithTools } from './tools/stream-parser'
+import { getAgentOutput } from './util/agent-output'
 import {
   asSystemInstruction,
   asSystemMessage,
-  asUserMessage,
+  buildUserMessageContent,
   messagesWithSystem,
   expireMessages,
 } from './util/messages'
@@ -52,29 +53,7 @@ import type {
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 
-/**
- * Combines prompt, params, and content into a unified message content structure
- */
-function buildUserMessageContent(
-  prompt: string | undefined,
-  params: Record<string, any> | undefined,
-  content?: Array<TextPart | ImagePart>,
-): string | Array<TextPart | ImagePart> {
-  // If we have content, return it as-is (client should have already combined prompt + content)
-  if (content && content.length > 0) {
-    if (content.length === 1 && content[0].type === 'text') {
-      return asUserMessage(content[0].text)
-    }
-    return content
-  }
 
-  // Only prompt/params, combine and return as simple text
-  const textParts = buildArray([
-    prompt,
-    params && JSON.stringify(params, null, 2),
-  ])
-  return asUserMessage(textParts.join('\n\n'))
-}
 
 export const runAgentStep = async (
   params: {
@@ -789,42 +768,4 @@ export const loopAgentSteps = async (
   }
 }
 
-function getAgentOutput(
-  agentState: AgentState,
-  agentTemplate: AgentTemplate,
-): AgentOutput {
-  if (agentTemplate.outputMode === 'structured_output') {
-    return {
-      type: 'structuredOutput',
-      value: agentState.output ?? null,
-    }
-  }
-  if (agentTemplate.outputMode === 'last_message') {
-    const assistantMessages = agentState.messageHistory.filter(
-      (message): message is AssistantMessage => message.role === 'assistant',
-    )
-    const lastAssistantMessage = assistantMessages[assistantMessages.length - 1]
-    if (!lastAssistantMessage) {
-      return {
-        type: 'error',
-        message: 'No response from agent',
-      }
-    }
-    return {
-      type: 'lastMessage',
-      value: lastAssistantMessage.content,
-    }
-  }
-  if (agentTemplate.outputMode === 'all_messages') {
-    // Remove the first message, which includes the previous conversation history.
-    const agentMessages = agentState.messageHistory.slice(1)
-    return {
-      type: 'allMessages',
-      value: agentMessages,
-    }
-  }
-  agentTemplate.outputMode satisfies never
-  throw new Error(
-    `Unknown output mode: ${'outputMode' in agentTemplate ? agentTemplate.outputMode : 'undefined'}`,
-  )
-}
+

backend/src/util/agent-output.ts

@@ -0,0 +1,43 @@
+import type { AgentTemplate } from '@codebuff/common/types/agent-template'
+import type { AssistantMessage } from '@codebuff/common/types/messages/codebuff-message'
+import type { AgentState, AgentOutput } from '@codebuff/common/types/session-state'
+
+export function getAgentOutput(
+  agentState: AgentState,
+  agentTemplate: AgentTemplate,
+): AgentOutput {
+  if (agentTemplate.outputMode === 'structured_output') {
+    return {
+      type: 'structuredOutput',
+      value: agentState.output ?? null,
+    }
+  }
+  if (agentTemplate.outputMode === 'last_message') {
+    const assistantMessages = agentState.messageHistory.filter(
+      (message): message is AssistantMessage => message.role === 'assistant',
+    )
+    const lastAssistantMessage = assistantMessages[assistantMessages.length - 1]
+    if (!lastAssistantMessage) {
+      return {
+        type: 'error',
+        message: 'No response from agent',
+      }
+    }
+    return {
+      type: 'lastMessage',
+      value: lastAssistantMessage.content,
+    }
+  }
+  if (agentTemplate.outputMode === 'all_messages') {
+    // Remove the first message, which includes the previous conversation history.
+    const agentMessages = agentState.messageHistory.slice(1)
+    return {
+      type: 'allMessages',
+      value: agentMessages,
+    }
+  }
+  agentTemplate.outputMode satisfies never
+  throw new Error(
+    `Unknown output mode: ${'outputMode' in agentTemplate ? agentTemplate.outputMode : 'undefined'}`,
+  )
+}

backend/src/util/messages.ts

@@ -17,6 +17,10 @@ import type {
   Message,
   ToolMessage,
 } from '@codebuff/common/types/messages/codebuff-message'
+import type {
+  TextPart,
+  ImagePart,
+} from '@codebuff/common/types/messages/content-part'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 
 export function messagesWithSystem(params: {
@@ -39,6 +43,31 @@ export function messagesWithSystem(params: {
 export function asUserMessage(str: string): string {
   return `<user_message>${str}${closeXml('user_message')}`
 }
+
+/**
+ * Combines prompt, params, and content into a unified message content structure
+ */
+export function buildUserMessageContent(
+  prompt: string | undefined,
+  params: Record<string, any> | undefined,
+  content?: Array<TextPart | ImagePart>,
+): string | Array<TextPart | ImagePart> {
+  // If we have content, return it as-is (client should have already combined prompt + content)
+  if (content && content.length > 0) {
+    if (content.length === 1 && content[0].type === 'text') {
+      return asUserMessage(content[0].text)
+    }
+    return content
+  }
+
+  // Only prompt/params, combine and return as simple text
+  const textParts = buildArray([
+    prompt,
+    params && JSON.stringify(params, null, 2),
+  ])
+  return asUserMessage(textParts.join('\n\n'))
+}
+
 export function parseUserMessage(str: string): string | undefined {
   const match = str.match(/<user_message>(.*?)<\/user_message>/s)
   return match ? match[1] : undefined