Fix context pruner to not remove user instructions prompt! And not leave extra messages

Author: jahooma

.agents/__tests__/context-pruner.test.ts

@@ -47,10 +47,12 @@ describe('context-pruner handleSteps', () => {
     })
   })
 
-  test('removes spawn_agent_inline call for context-pruner from last assistant message', () => {
+  test('removes spawn_agent_inline call for context-pruner and following messages', () => {
     const messages = [
       createMessage('user', 'Hello'),
       createMessage('assistant', 'I will spawn the context-pruner agent.\n\n<codebuff_tool_call>\n{\n  "cb_tool_name": "spawn_agent_inline",\n  "agent_type": "context-pruner"\n}\n</codebuff_tool_call>'),
+      createMessage('user', '{"params": {"maxContextLength": 100000}}'),
+      createMessage('user', 'Tools and instructions'),
     ]
 
     const results = runHandleSteps(messages)
@@ -60,16 +62,31 @@ describe('context-pruner handleSteps', () => {
     expect(results[0].input.messages[0]).toEqual(createMessage('user', 'Hello'))
   })
 
-  test('does not remove last message if it does not contain context-pruner spawn call', () => {
+  test('does not remove messages if assistant message does not contain context-pruner spawn call', () => {
     const messages = [
       createMessage('user', 'Hello'),
       createMessage('assistant', 'Regular response without spawn call'),
+      createMessage('user', 'Follow up'),
+    ]
+
+    const results = runHandleSteps(messages)
+    expect(results).toHaveLength(1)
+    expect(results[0].input.messages).toHaveLength(3)
+  })
+
+  test('handles context-pruner spawn call without enough following messages', () => {
+    const messages = [
+      createMessage('user', 'Hello'),
+      createMessage('assistant', 'I will spawn the context-pruner agent.\n\n<codebuff_tool_call>\n{\n  "cb_tool_name": "spawn_agent_inline",\n  "agent_type": "context-pruner"\n}\n</codebuff_tool_call>'),
+      createMessage('user', '{"params": {"maxContextLength": 100000}}'),
     ]
 
     const results = runHandleSteps(messages)
 
     expect(results).toHaveLength(1)
-    expect(results[0].input.messages).toHaveLength(2)
+    // Should preserve all messages since there aren't 3 messages to remove
+    expect(results[0].input.messages).toHaveLength(3)
+
   })
 
   test('removes old terminal command results while keeping recent 5', () => {
@@ -389,26 +406,41 @@ describe('context-pruner edge cases', () => {
 
   test('handles spawn_agent_inline detection with variations', () => {
     const testCases = [
-      'Regular message with spawn_agent_inline but not for context-pruner',
-      'spawn_agent_inline call for "context-pruner" with quotes',
-      'spawn_agent_inline\n  "agent_type": "context-pruner"',
-      'Multiple spawn_agent_inline calls, one for context-pruner',
+      {
+        content: 'Regular message with spawn_agent_inline but not for other-agent',
+        shouldRemove: false,
+      },
+      {
+        content: 'spawn_agent_inline call for "context-pruner" with quotes',
+        shouldRemove: true, // Has context-pruner and 3 total messages
+      },
+      {
+        content: 'spawn_agent_inline\n  "agent_type": "context-pruner"',
+        shouldRemove: true, // Has context-pruner and 3 total messages
+      },
+      {
+        content: 'Multiple spawn_agent_inline calls, one for context-pruner',
+        shouldRemove: true, // Has context-pruner and 3 total messages
+      },
     ]
 
-    testCases.forEach((content, index) => {
+    testCases.forEach(({ content, shouldRemove }, index) => {
       const messages = [
         createMessage('user', 'Hello'),
         createMessage('assistant', content),
+        createMessage('user', 'Follow up'),
+        createMessage('user', 'Tools and instructions'),
       ]
 
       const results = runHandleSteps(messages)
 
-      if (content.includes('context-pruner')) {
-        // Should remove the message containing context-pruner spawn
+      if (shouldRemove) {
+        // Should remove the assistant message and following 2 user messages
         expect(results[0].input.messages).toHaveLength(1)
+        expect(results[0].input.messages[0]).toEqual(createMessage('user', 'Hello'))
       } else {
-        // Should preserve the message
-        expect(results[0].input.messages).toHaveLength(2)
+        // Should preserve all messages
+        expect(results[0].input.messages).toHaveLength(4)
       }
     })
   })

.agents/context-pruner.ts

@@ -37,23 +37,28 @@ const definition: AgentDefinition = {
     const maxMessageTokens: number = params?.maxContextLength ?? 200_000
     const numTerminalCommandsToKeep = 5
 
-    // Remove the last assistant message if it contains the spawn call that invoked this context-pruner
     let currentMessages = [...messages]
-    if (currentMessages.length > 0) {
-      const lastMessage = currentMessages[currentMessages.length - 1]
-      if (
-        lastMessage.role === 'assistant' &&
-        typeof lastMessage.content === 'string'
-      ) {
-        // Check if this message contains a spawn_agent_inline call for context-pruner
-        if (
-          lastMessage.content.includes('spawn_agent_inline') &&
-          lastMessage.content.includes('context-pruner')
-        ) {
-          // Remove the entire message
-          currentMessages.pop()
-        }
-      }
+    const lastAssistantMessageIndex = currentMessages.findLastIndex(
+      (message) => message.role === 'assistant',
+    )
+    const lastAssistantMessage = currentMessages[lastAssistantMessageIndex]
+    const lastAssistantMessageIsToolCall =
+      typeof lastAssistantMessage?.content === 'string' &&
+      lastAssistantMessage.content.includes('spawn_agent_inline') &&
+      lastAssistantMessage.content.includes('context-pruner')
+    const inputMessage = currentMessages[lastAssistantMessageIndex + 1]
+    const userInstructionsToolsMessage =
+      currentMessages[lastAssistantMessageIndex + 2]
+    if (
+      lastAssistantMessageIsToolCall &&
+      inputMessage &&
+      userInstructionsToolsMessage
+    ) {
+      // Remove these messages:
+      // - assistant message with spawn_agent_inline call
+      // - input message with params object for this agent
+      // - user instructions message with tools description
+      currentMessages.splice(lastAssistantMessageIndex, 3)
     }
 
     // Initial check - if already under limit, return (with inline agent tool call removed)

backend/src/__tests__/run-agent-step-tools.test.ts

@@ -456,7 +456,7 @@ describe('runAgentStep - set_output tool', () => {
         // Delete the last two assistant messages by doing two iterations
         const messages = [...agentState.messageHistory]
 
-        // First iteration: find and remove the last assistant message
+        // First iteration: find and remove the last assistant message, which is the tool call to this agent
         for (let i = messages.length - 1; i >= 0; i--) {
           if (messages[i].role === 'assistant') {
             messages.splice(i, 1)
@@ -472,6 +472,14 @@ describe('runAgentStep - set_output tool', () => {
           }
         }
 
+        // Third iteration: find and remove the third assistant message
+        for (let i = messages.length - 1; i >= 0; i--) {
+          if (messages[i].role === 'assistant') {
+            messages.splice(i, 1)
+            break
+          }
+        }
+
         // Set the updated messages
         yield {
           toolName: 'set_messages',
@@ -551,7 +559,7 @@ describe('runAgentStep - set_output tool', () => {
       { role: 'user', content: 'Hello' },
       { role: 'assistant', content: 'Hi there!' },
       { role: 'user', content: 'How are you?' },
-      { role: 'assistant', content: 'I am doing well, thank you!' },
+      // { role: 'assistant', content: 'I am doing well, thank you!' },
       { role: 'user', content: 'Can you help me?' },
       {
         role: 'user',
@@ -565,6 +573,11 @@ describe('runAgentStep - set_output tool', () => {
           'Delete the last two assistant messages',
         ),
       },
+      {
+        role: 'user',
+        content: expect.stringContaining('Delete messages instructions prompt'),
+        timeToLive: 'userPrompt',
+      },
     ]
 
     expectedMessages.forEach((expected, index) => {

backend/src/tools/handlers/tool/spawn-agent-inline.ts

@@ -112,9 +112,6 @@ export const handleSpawnAgentInline = ((params: {
     // Since we share the same message array reference, this should already be updated
     let finalMessages = result.agentState?.messageHistory || state.messages
 
-    // Expire messages with timeToLive: 'userPrompt' to clean up inline agent's temporary messages
-    finalMessages = expireMessages(finalMessages, 'userPrompt')
-
     state.messages = finalMessages
 
     // Update parent agent state to reflect shared message history