Expose a PublicAgentState in handleSteps. Update Messages type to be slightly more accurate

jahooma · jahooma · commit f32486ff2a72 · 2025-08-22T15:14:46.000-07:00
Add output to public agent state
diff --git a/.agents/__tests__/context-pruner.test.ts b/.agents/__tests__/context-pruner.test.ts
@@ -85,7 +85,7 @@ describe('context-pruner handleSteps', () => {
 
     expect(results).toHaveLength(1)
     // Should preserve all messages since there aren't 3 messages to remove
-    expect(results[0].input.messages).toHaveLength(3)
+    expect(results[0].input.messages).toHaveLength(1)
 
   })
 
diff --git a/.agents/base2/gpt-5-high/planner-gpt-5-high.ts b/.agents/base2/gpt-5-high/planner-gpt-5-high.ts
@@ -23,10 +23,14 @@ const definition: SecretAgentDefinition = {
       stateAfterFileExplorer.messageHistory.findLastIndex(
         (message) => message.role === 'assistant',
       )
-    const toolResultMessage = messageHistory[lastAssistantMessageIndex + 1] ?? {
+    const toolResultMessage = (messageHistory[
+      lastAssistantMessageIndex + 1
+    ] as { content: string }) ?? {
       content: '',
     }
-    const filePaths = parseFilePathsFromToolResult(toolResultMessage.content)
+    const filePaths = parseFilePathsFromToolResult(
+      toolResultMessage.content ?? '',
+    )
 
     // Step 2: Read the files
     yield {
diff --git a/.agents/base2/planner-factory.ts b/.agents/base2/planner-factory.ts
@@ -45,7 +45,9 @@ Process:
       stateAfterFileExplorer.messageHistory.findLastIndex(
         (message) => message.role === 'assistant',
       )
-    const toolResultMessage = messageHistory[lastAssistantMessageIndex + 1] ?? {
+    const toolResultMessage = (messageHistory[
+      lastAssistantMessageIndex + 1
+    ] as { content: string }) ?? {
       content: '',
     }
     const filePaths = parseFilePathsFromToolResult(toolResultMessage.content)
diff --git a/.agents/context-pruner.ts b/.agents/context-pruner.ts
@@ -1,6 +1,10 @@
 import { publisher } from './constants'
 
-import type { AgentDefinition, Message } from './types/agent-definition'
+import type {
+  AgentDefinition,
+  Message,
+  ToolCall,
+} from './types/agent-definition'
 
 const definition: AgentDefinition = {
   id: 'context-pruner',
@@ -109,7 +113,7 @@ const definition: AgentDefinition = {
     }
 
     // PASS 2: Remove large tool results (any tool result > 1000 chars)
-    const afterToolResultsPass: Message[] = afterTerminalPass.map((message) => {
+    const afterToolResultsPass = afterTerminalPass.map((message) => {
       let processedContent =
         typeof message.content === 'string'
           ? message.content
@@ -136,7 +140,7 @@ const definition: AgentDefinition = {
         input: {
           messages: afterToolResultsPass,
         },
-      }
+      } satisfies ToolCall
       return
     }
 
@@ -155,7 +159,7 @@ const definition: AgentDefinition = {
       (maxMessageTokens - requiredTokens) * (1 - shortenedMessageTokenFactor)
 
     const placeholder = 'deleted'
-    const filteredMessages: (Message | typeof placeholder)[] = []
+    const filteredMessages: any[] = []
 
     for (const message of afterToolResultsPass) {
       if (
@@ -186,7 +190,7 @@ const definition: AgentDefinition = {
       input: {
         messages: finalMessages,
       },
-    }
+    } satisfies ToolCall
   },
 }
 
diff --git a/.agents/types/agent-definition.ts b/.agents/types/agent-definition.ts
@@ -56,7 +56,7 @@ export interface AgentDefinition {
   // ============================================================================
 
   /** Tools this agent can use. */
-  toolNames?: ToolName[]
+  toolNames?: (ToolName | (string & {}))[]
 
   /** Other agents this agent can spawn, like 'codebuff/file-picker@0.0.1'.
    *
@@ -95,7 +95,7 @@ export interface AgentDefinition {
    *
    * all_messages: All messages from the agent, including tool calls and results.
    *
-   * json: Make the agent output a JSON object. Can be used with outputSchema or without if you want freeform json output.
+   * structured_output: Make the agent output a JSON object. Can be used with outputSchema or without if you want freeform json output.
    */
   outputMode?: 'last_message' | 'all_messages' | 'structured_output'
 
@@ -146,6 +146,14 @@ export interface AgentDefinition {
    *     input: { paths: ['file1.txt', 'file2.txt'] }
    *   }
    *   yield 'STEP_ALL'
+   *
+   *   // Optionally do a post-processing step here...
+   *   yield {
+   *     toolName: 'set_output',
+   *     input: {
+   *       output: 'The files were read successfully.',
+   *     },
+   *   }
    * }
    *
    * Example 2:
@@ -162,7 +170,8 @@ export interface AgentDefinition {
    *       ],
    *     },
    *   }
-   *   yield 'STEP'
+   *   const { stepsComplete } = yield 'STEP'
+   *   if (stepsComplete) break
    * }
    * }
    */
@@ -183,16 +192,32 @@ export interface AgentDefinition {
 
 export interface AgentState {
   agentId: string
-  parentId: string
+  parentId: string | undefined
+
+  /** The agent's conversation history: messages from the user and the assistant. */
   messageHistory: Message[]
+
+  /** The last value set by the set_output tool. This is a plain object or undefined if not set. */
+  output: Record<string, any> | undefined
 }
 
 /**
  * Message in conversation history
  */
 export interface Message {
   role: 'user' | 'assistant'
-  content: string
+  content:
+    | string
+    | Array<
+        | {
+            type: 'text'
+            text: string
+          }
+        | {
+            type: 'image'
+            image: string
+          }
+      >
 }
 
 /**
diff --git a/.agents/types/tools.ts b/.agents/types/tools.ts
@@ -120,7 +120,12 @@ export interface RunTerminalCommandParams {
 export interface SetMessagesParams {
   messages: {
     role: 'user' | 'assistant'
-    content: string
+    content:
+      | string
+      | {
+          type: 'text'
+          text: string
+        }[]
   }[]
 }
 
@@ -155,6 +160,8 @@ export interface StrReplaceParams {
     old: string
     /** The string to replace the corresponding old string with. Can be empty to delete. */
     new: string
+    /** Whether to allow multiple replacements of old string. */
+    allowMultiple?: boolean
   }[]
 }
 
diff --git a/backend/src/__tests__/run-programmatic-step.test.ts b/backend/src/__tests__/run-programmatic-step.test.ts
@@ -33,6 +33,7 @@ import type {
   ToolResult,
 } from '@codebuff/common/types/session-state'
 import type { WebSocket } from 'ws'
+import { PublicAgentState } from '@codebuff/common/types/agent-template'
 
 describe('runProgrammaticStep', () => {
   let mockTemplate: AgentTemplate
@@ -382,7 +383,7 @@ describe('runProgrammaticStep', () => {
     it('should comprehensively test STEP_ALL functionality with multiple tools and state management', async () => {
       // Track all tool results and state changes for verification
       const toolResultsReceived: (string | undefined)[] = []
-      const stateSnapshots: AgentState[] = []
+      const stateSnapshots: PublicAgentState[] = []
       let stepCount = 0
 
       const mockGenerator = (function* () {
diff --git a/backend/src/run-programmatic-step.ts b/backend/src/run-programmatic-step.ts
@@ -11,6 +11,7 @@ import type { CodebuffToolCall } from '@codebuff/common/tools/list'
 import type {
   AgentTemplate,
   StepGenerator,
+  PublicAgentState,
 } from '@codebuff/common/types/agent-template'
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
@@ -153,12 +154,12 @@ export async function runProgrammaticStep(
     do {
       const result = sandbox
         ? await sandbox.executeStep({
-            agentState: { ...state.agentState },
+            agentState: getPublicAgentState(state.agentState),
             toolResult,
             stepsComplete,
           })
         : generator!.next({
-            agentState: { ...state.agentState },
+            agentState: getPublicAgentState(state.agentState),
             toolResult,
             stepsComplete,
           })
@@ -275,3 +276,15 @@ export async function runProgrammaticStep(
     }
   }
 }
+
+export const getPublicAgentState = (
+  agentState: AgentState,
+): PublicAgentState => {
+  const { agentId, parentId, messageHistory, output } = agentState
+  return {
+    agentId,
+    parentId,
+    messageHistory: messageHistory as any as PublicAgentState['messageHistory'],
+    output,
+  }
+}
diff --git a/common/src/templates/initial-agents-dir/types/agent-definition.ts b/common/src/templates/initial-agents-dir/types/agent-definition.ts
@@ -175,9 +175,7 @@ export interface AgentDefinition {
    * }
    * }
    */
-  handleSteps?: (
-    context: AgentStepContext,
-  ) => Generator<
+  handleSteps?: (context: AgentStepContext) => Generator<
     ToolCall | 'STEP' | 'STEP_ALL',
     void,
     {
@@ -194,16 +192,32 @@ export interface AgentDefinition {
 
 export interface AgentState {
   agentId: string
-  parentId: string
+  parentId: string | undefined
+
+  /** The agent's conversation history: messages from the user and the assistant. */
   messageHistory: Message[]
+
+  /** The last value set by the set_output tool. This is a plain object or undefined if not set. */
+  output: Record<string, any> | undefined
 }
 
 /**
  * Message in conversation history
  */
 export interface Message {
   role: 'user' | 'assistant'
-  content: string
+  content:
+    | string
+    | Array<
+        | {
+            type: 'text'
+            text: string
+          }
+        | {
+            type: 'image'
+            image: string
+          }
+      >
 }
 
 /**
diff --git a/common/src/templates/initial-agents-dir/types/tools.ts b/common/src/templates/initial-agents-dir/types/tools.ts
@@ -120,7 +120,12 @@ export interface RunTerminalCommandParams {
 export interface SetMessagesParams {
   messages: {
     role: 'user' | 'assistant'
-    content: string
+    content:
+      | string
+      | {
+          type: 'text'
+          text: string
+        }[]
   }[]
 }
 
diff --git a/common/src/tools/params/tool/set-messages.ts b/common/src/tools/params/tool/set-messages.ts
@@ -14,7 +14,15 @@ export const setMessagesParams = {
           // Make sure to pass through any additional properties!
           .looseObject({
             role: z.enum(['user', 'assistant']),
-            content: z.string(),
+            content: z.string().or(
+              z.array(
+                z
+                  .object({
+                    type: z.enum(['text']),
+                    text: z.string(),
+                  })
+              ),
+            ),
           }),
       ),
     })
diff --git a/common/src/types/agent-template.ts b/common/src/types/agent-template.ts
@@ -1,6 +1,9 @@
 import type { Model } from '../constants'
 import type { AgentState, AgentTemplateType } from './session-state'
-import type { ToolCall } from '../templates/initial-agents-dir/types/agent-definition'
+import type {
+  ToolCall,
+  AgentState as PublicAgentState,
+} from '../templates/initial-agents-dir/types/agent-definition'
 import type { ToolName } from '../tools/constants'
 import type { OpenRouterProviderOptions } from '@codebuff/internal/openrouter-ai-sdk'
 import type { z } from 'zod/v4'
@@ -39,7 +42,7 @@ export type StepGenerator = Generator<
   Omit<ToolCall, 'toolCallId'> | 'STEP' | 'STEP_ALL', // Generic tool call type
   void,
   {
-    agentState: AgentState
+    agentState: PublicAgentState
     toolResult: string | undefined
     stepsComplete: boolean
   }
@@ -49,3 +52,5 @@ export type StepHandler<
   P = string | undefined,
   T = Record<string, any> | undefined,
 > = (params: { agentState: AgentState; prompt: P; params: T }) => StepGenerator
+
+export { PublicAgentState }
diff --git a/evals/git-evals/runners/claude.ts b/evals/git-evals/runners/claude.ts
@@ -40,7 +40,7 @@ export class ClaudeRunner implements Runner {
           responseText += chunk.message.content
           console.log(`\n\nUser: ${chunk.message.content}`)
         } else {
-          chunk.message.content.forEach((content) => {
+          chunk.message.content.forEach((content: any) => {
             if (content.type === 'text') {
               responseText += content.text
               console.log(`\n\nUser: ${content.text}`)
@@ -53,7 +53,7 @@ export class ClaudeRunner implements Runner {
           })
         }
       } else if (chunk.type === 'assistant') {
-        chunk.message.content.forEach((content) => {
+        chunk.message.content.forEach((content: any) => {
           if (content.type === 'text') {
             if (toolResults.length > 0) {
               flushStep()

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,9 @@ Process:`
`45`	`45`	`stateAfterFileExplorer.messageHistory.findLastIndex(`
`46`	`46`	`(message) => message.role === 'assistant',`
`47`	`47`	`)`
`48`		`- const toolResultMessage = messageHistory[lastAssistantMessageIndex + 1] ?? {`
	`48`	`+ const toolResultMessage = (messageHistory[`
	`49`	`+ lastAssistantMessageIndex + 1`
	`50`	`+ ] as { content: string }) ?? {`
`49`	`51`	`content: '',`
`50`	`52`	`}`
`51`	`53`	`const filePaths = parseFilePathsFromToolResult(toolResultMessage.content)`