Unified programmatic agent #1: Remove ProgrammaticAgentTemplate. Add handeStep field.

Author: jahooma

backend/src/run-agent-step.ts

@@ -23,14 +23,14 @@ import { asyncAgentManager } from './async-agent-manager'
 import { getFileReadingUpdates } from './get-file-reading-updates'
 import { checkLiveUserInput } from './live-user-inputs'
 import { getAgentStreamFromTemplate } from './prompt-agent-stream'
-import { runProgrammaticAgent } from './run-programmatic-agent'
+import { runProgrammaticStep } from './run-programmatic-step'
 import { additionalSystemPrompts } from './system-prompt/prompts'
 import { saveAgentRequest } from './system-prompt/save-agent-request'
 import { agentTemplates } from './templates/agent-list'
 import { processAgentOverrides } from './templates/agent-overrides'
 import { agentRegistry } from './templates/agent-registry'
 import { formatPrompt, getAgentPrompt } from './templates/strings'
-import { AgentTemplateUnion } from './templates/types'
+import { AgentTemplate } from './templates/types'
 import { processStreamWithTools } from './tools/stream-parser'
 import { logger } from './util/logger'
 import {
@@ -70,7 +70,7 @@ export interface AgentOptions {
 async function getAgentTemplateWithOverrides(
   agentType: AgentTemplateType,
   fileContext: ProjectFileContext
-): Promise<AgentTemplateUnion> {
+): Promise<AgentTemplate> {
   // Initialize registry if needed
   await agentRegistry.initialize(fileContext)
 
@@ -82,11 +82,6 @@ async function getAgentTemplateWithOverrides(
     )
   }
 
-  if (baseTemplate.implementation === 'programmatic') {
-    // Programmatic agents cannot be overridden.
-    return baseTemplate
-  }
-
   return processAgentOverrides(baseTemplate, fileContext)
 }
 
@@ -106,12 +101,12 @@ export const runAgentStep = async (
     onResponseChunk,
     fileContext,
     agentType,
-    agentState,
     prompt,
     params,
     assistantMessage,
     assistantPrefix,
   } = options
+  let agentState = options.agentState
 
   const { agentContext } = agentState
 
@@ -131,16 +126,13 @@ export const runAgentStep = async (
     )
   }
 
-  if (agentTemplate.implementation === 'programmatic') {
-    const agentState = await runProgrammaticAgent(agentTemplate, {
+  const { handleStep } = agentTemplate
+  if (handleStep) {
+    agentState = await runProgrammaticStep(agentState, {
       ...options,
       ws,
+      template: agentTemplate,
     })
-    return {
-      agentState,
-      shouldEndTurn: true,
-      fullResponse: '',
-    }
   }
 
   const { model } = agentTemplate
@@ -569,15 +561,7 @@ export const loopAgentSteps = async (
     initialAssistantPrefix,
     stepAssistantMessage,
     stepAssistantPrefix,
-  } =
-    agentTemplate.implementation === 'llm'
-      ? agentTemplate
-      : {
-          initialAssistantMessage: undefined,
-          initialAssistantPrefix: undefined,
-          stepAssistantMessage: undefined,
-          stepAssistantPrefix: undefined,
-        }
+  } = agentTemplate
   let isFirstStep = true
   let currentPrompt = prompt
   let currentParams = params

backend/src/run-programmatic-agent.ts backend/src/run-programmatic-step.tsbackend/src/run-programmatic-agent.ts renamed to backend/src/run-programmatic-step.ts

@@ -1,30 +1,31 @@
 import {
   AgentState,
   AgentTemplateType,
+  ToolResult,
 } from '@codebuff/common/types/session-state'
 import { ProjectFileContext } from '@codebuff/common/util/file'
 import { WebSocket } from 'ws'
 import { runTool } from './run-tool'
-import {
-  ProgrammaticAgentContext,
-  ProgrammaticAgentTemplate,
-} from './templates/types'
+import { AgentTemplate, StepGenerator } from './templates/types'
 import { CodebuffToolCall } from './tools/constants'
 import { logger } from './util/logger'
 import { getRequestContext } from './websockets/request-context'
 
+// Maintains generator state for all agents. Generator state can't be serialized, so we store it in memory.
+const agentIdToGenerator: Record<string, StepGenerator | undefined> = {}
+
 // Function to handle programmatic agents
-export async function runProgrammaticAgent(
-  template: ProgrammaticAgentTemplate,
-  options: {
+export async function runProgrammaticStep(
+  agentState: AgentState,
+  params: {
+    template: AgentTemplate
     userId: string | undefined
     userInputId: string
     clientSessionId: string
     fingerprintId: string
     onResponseChunk: (chunk: string) => void
     agentType: AgentTemplateType
     fileContext: ProjectFileContext
-    agentState: AgentState
     prompt: string | undefined
     params: Record<string, any> | undefined
     assistantMessage: string | undefined
@@ -33,38 +34,55 @@ export async function runProgrammaticAgent(
   }
 ): Promise<AgentState> {
   const {
-    agentState,
+    template,
     onResponseChunk,
     ws,
     userId,
     userInputId,
     clientSessionId,
     fingerprintId,
     fileContext,
-  } = options
+  } = params
 
   logger.info(
     {
       template: template.id,
-      agentType: options.agentType,
-      prompt: options.prompt,
-      params: options.params,
+      agentType: params.agentType,
+      prompt: params.prompt,
+      params: params.params,
     },
-    'Running programmatic agent'
+    'Running programmatic step'
   )
-  // Create context for the programmatic agent
-  const context: ProgrammaticAgentContext = {
-    prompt: options.prompt || '',
-    params: options.params || {},
+
+  let generator = agentIdToGenerator[agentState.agentId]
+  if (!generator) {
+    if (!template.handleStep) {
+      throw new Error('No step handler found for agent template ' + template.id)
+    }
+    generator = template.handleStep(agentState)
+    agentIdToGenerator[agentState.agentId] = generator
   }
 
+  if (Math.random() <= 1) {
+    throw new Error('Not implemented yet!')
+  }
+
+  let toolResult: ToolResult | undefined
+
   try {
-    // Run the generator function and handle tool calls
-    const generator = template.handler(context)
-    let result = generator.next()
+    do {
+      let result = generator.next(toolResult)
+      if (result.done) {
+        break
+      }
+      if (result.value === 'STEP') {
+        break
+      }
+      if (result.value === 'STEP_ALL') {
+        break
+      }
 
-    // Process tool calls yielded by the generator
-    while (!result.done) {
+      // Process tool calls yielded by the generator
       const toolCallWithoutId = result.value
       const toolCall = {
         ...toolCallWithoutId,
@@ -79,7 +97,7 @@ export async function runProgrammaticAgent(
       const repoId = requestContext?.processedRepoId
 
       // Execute the tool call using the simplified wrapper
-      const toolResult = await runTool(toolCall, {
+      toolResult = await runTool(toolCall, {
         ws,
         userId,
         userInputId,
@@ -93,13 +111,10 @@ export async function runProgrammaticAgent(
         agentState,
       })
 
-      // Send the tool result back to the generator
-      result = generator.next(toolResult)
-
       if (result.value && result.value.toolName === 'end_turn') {
         break
       }
-    }
+    } while (true)
 
     logger.info(
       { report: agentState.report },

backend/src/run-tool.ts

@@ -23,7 +23,7 @@ import { searchWeb } from './llm-apis/linkup-api'
 import { PROFIT_MARGIN } from './llm-apis/message-cost-tracker'
 import { getSearchSystemPrompt } from './system-prompt/search-system-prompt'
 import { agentRegistry } from './templates/agent-registry'
-import { AgentTemplate, ProgrammaticAgentTemplate } from './templates/types'
+import { AgentTemplate } from './templates/types'
 import { ClientToolCall, CodebuffToolCall } from './tools/constants'
 import { logger } from './util/logger'
 import { renderReadFilesResult } from './util/parse-tool-call-xml'
@@ -43,7 +43,7 @@ export interface RunToolOptions {
   agentStepId: string
   fileContext: ProjectFileContext
   messages: CoreMessage[]
-  agentTemplate: AgentTemplate | ProgrammaticAgentTemplate
+  agentTemplate: AgentTemplate
   repoId?: string
   // Additional context for update_report
   agentState?: AgentState
@@ -681,10 +681,7 @@ export async function runToolInner(
           const agentTemplate = allTemplates[agentState.agentType!]
           let report = ''
 
-          if (
-            agentTemplate.implementation === 'programmatic' ||
-            agentTemplate.outputMode === 'report'
-          ) {
+          if (agentTemplate.outputMode === 'report') {
             report = JSON.stringify(result.value.agentState.report, null, 2)
           } else if (agentTemplate.outputMode === 'last_message') {
             const { agentState } = result.value

backend/src/templates/__tests__/agent-overrides.test.ts

@@ -23,7 +23,6 @@ describe('processAgentOverrides', () => {
     systemPrompt: 'Base system prompt',
     userInputPrompt: 'Base user input prompt',
     agentStepPrompt: 'Base agent step prompt',
-    implementation: 'llm',
   }
 
   const mockFileContext: ProjectFileContext = {