Fix for context pruner: actually pass system prompt and tools through

Author: jahooma

.agents/types/agent-definition.ts

@@ -283,11 +283,14 @@ export interface AgentState {
   /** The last value set by the set_output tool. This is a plain object or undefined if not set. */
   output: Record<string, any> | undefined
 
-  /** The system prompt for this agent. Used by context-pruner to account for system prompt size when calculating effective message budget. */
-  systemPrompt?: string
+  /** The system prompt for this agent. */
+  systemPrompt: string
 
-  /** The tool definitions for this agent. Used by context-pruner to account for tool definition tokens when calculating effective message budget. */
-  toolDefinitions?: Record<string, { name: string; description: string; parameters: JsonObjectSchema }>[]
+  /** The tool definitions for this agent. */
+  toolDefinitions: Record<
+    string,
+    { description: string | undefined; inputSchema: {} }
+  >
 }
 
 /**

common/src/templates/initial-agents-dir/types/agent-definition.ts

@@ -282,6 +282,15 @@ export interface AgentState {
 
   /** The last value set by the set_output tool. This is a plain object or undefined if not set. */
   output: Record<string, any> | undefined
+
+  /** The system prompt for this agent. */
+  systemPrompt: string
+
+  /** The tool definitions for this agent. */
+  toolDefinitions: Record<
+    string,
+    { description: string | undefined; inputSchema: {} }
+  >
 }
 
 /**

common/src/types/session-state.ts

@@ -39,8 +39,11 @@ export type AgentState = {
   directCreditsUsed: number
   output?: Record<string, any>
   parentId?: string
-  systemPrompt?: string
-  toolDefinitions?: Record<string, { description: string | undefined; inputSchema: {} }>
+  systemPrompt: string
+  toolDefinitions: Record<
+    string,
+    { description: string | undefined; inputSchema: {} }
+  >
 }
 
 export const AgentOutputSchema = z.discriminatedUnion('type', [
@@ -122,6 +125,8 @@ export function getInitialAgentState(): AgentState {
     directCreditsUsed: 0,
     output: undefined,
     parentId: undefined,
+    systemPrompt: '',
+    toolDefinitions: {},
   }
 }
 export function getInitialSessionState(

packages/agent-runtime/src/run-agent-step.ts

@@ -493,6 +493,7 @@ export async function loopAgentSteps(
     startAgentRun: StartAgentRunFn
     userId: string | undefined
     userInputId: string
+    agentTemplate?: AgentTemplate
   } & ParamsExcluding<typeof additionalToolDefinitions, 'agentTemplate'> &
     ParamsExcluding<
       typeof runProgrammaticStep,
@@ -569,10 +570,12 @@ export async function loopAgentSteps(
     userInputId,
   } = params
 
-  const agentTemplate = await getAgentTemplate({
-    ...params,
-    agentId: agentType,
-  })
+  const agentTemplate =
+    params.agentTemplate ??
+    (await getAgentTemplate({
+      ...params,
+      agentId: agentType,
+    }))
   if (!agentTemplate) {
     throw new Error(`Agent template not found for type: ${agentType}`)
   }
@@ -767,6 +770,7 @@ export async function loopAgentSteps(
           stepNumber: totalSteps,
           stepsComplete: shouldEndTurn,
           system,
+          tools,
           template: agentTemplate,
           toolCallParams: currentParams,
         })

packages/agent-runtime/src/run-programmatic-step.ts

@@ -384,13 +384,23 @@ export async function runProgrammaticStep(
 export const getPublicAgentState = (
   agentState: AgentState & Required<Pick<AgentState, 'runId'>>,
 ): PublicAgentState => {
-  const { agentId, runId, parentId, messageHistory, output } = agentState
+  const {
+    agentId,
+    runId,
+    parentId,
+    messageHistory,
+    output,
+    systemPrompt,
+    toolDefinitions,
+  } = agentState
   return {
     agentId,
     runId,
     parentId,
     messageHistory: messageHistory as any as PublicAgentState['messageHistory'],
     output,
+    systemPrompt,
+    toolDefinitions,
   }
 }
 

packages/agent-runtime/src/tools/handlers/handler-function-type.ts

@@ -50,7 +50,7 @@ export type CodebuffToolHandlerFunction<T extends ToolName = ToolName> = (
     sendSubagentChunk: SendSubagentChunkFn
     signal: AbortSignal
     system: string
-    tools?: ToolSet
+    tools: ToolSet
     trackEvent: TrackEventFn
     userId: string | undefined
     userInputId: string

packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts

@@ -18,6 +18,7 @@ import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { AgentState } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { ToolSet } from 'ai'
+import { mapValues } from 'lodash'
 
 type ToolName = 'spawn_agent_inline'
 export const handleSpawnAgentInline = (async (
@@ -33,7 +34,7 @@ export const handleSpawnAgentInline = (async (
     localAgentTemplates: Record<string, AgentTemplate>
     logger: Logger
     system: string
-    tools?: ToolSet
+    tools: ToolSet
     userId: string | undefined
     userInputId: string
     writeToClient: (chunk: string | PrintModeEvent) => void
@@ -60,9 +61,10 @@ export const handleSpawnAgentInline = (async (
     agentTemplate: parentAgentTemplate,
     fingerprintId,
     system,
-    tools: parentTools = {},
+    tools: parentTools,
     userInputId,
     writeToClient,
+    logger,
   } = params
   const {
     agent_type: agentTypeStr,
@@ -80,17 +82,31 @@ export const handleSpawnAgentInline = (async (
 
   validateAgentInput(agentTemplate, agentType, prompt, spawnParams)
 
+  // Override template for inline agent to share system prompt & message history with parent
+  const inlineTemplate = {
+    ...agentTemplate,
+    includeMessageHistory: true,
+    inheritParentSystemPrompt: true,
+  }
+
   // Create child agent state that shares message history with parent
-  const childAgentState: AgentState = createAgentState(
-    agentType,
-    agentTemplate,
-    parentAgentState,
-    parentAgentState.agentContext,
-  )
+  const childAgentState: AgentState = {
+    ...createAgentState(
+      agentType,
+      inlineTemplate,
+      parentAgentState,
+      parentAgentState.agentContext,
+    ),
+    systemPrompt: system,
+    toolDefinitions: mapValues(parentTools, (tool) => ({
+      description: tool.description,
+      inputSchema: tool.inputSchema as {},
+    })),
+  }
 
   logAgentSpawn({
     ...params,
-    agentTemplate,
+    agentTemplate: inlineTemplate,
     agentType,
     agentId: childAgentState.agentId,
     parentId: childAgentState.parentId,
@@ -104,14 +120,12 @@ export const handleSpawnAgentInline = (async (
     userInputId: `${userInputId}-inline-${agentType}${childAgentState.agentId}`,
     prompt: prompt || '',
     spawnParams,
-    agentTemplate,
+    agentTemplate: inlineTemplate,
     parentAgentState,
     agentState: childAgentState,
     fingerprintId,
     parentSystemPrompt: system,
-    parentTools: agentTemplate.inheritParentSystemPrompt
-      ? parentTools
-      : undefined,
+    parentTools,
     onResponseChunk: (chunk) => {
       // Inherits parent's onResponseChunk, except for context-pruner (TODO: add an option for it to be silent?)
       if (agentType !== 'context-pruner') {

packages/agent-runtime/src/tools/handlers/tool/spawn-agent-utils.ts

@@ -186,6 +186,8 @@ export function createAgentState(
     directCreditsUsed: 0,
     output: undefined,
     parentId: parentAgentState.agentId,
+    systemPrompt: '',
+    toolDefinitions: {},
   }
 }
 

packages/agent-runtime/src/tools/tool-executor.ts

@@ -143,7 +143,7 @@ export type ExecuteToolCallParams<T extends string = ToolName> = {
   runId: string
   signal: AbortSignal
   system: string
-  tools?: ToolSet
+  tools: ToolSet
   toolCallId: string | undefined
   toolCalls: (CodebuffToolCall | CustomToolCall)[]
   toolResults: ToolMessage[]