fix: bound agent communication memory growth

packages/builtin-tools/src/tools/AgentTool/__tests__/filterIncompleteToolCalls.test.ts

@@ -0,0 +1,180 @@
+import { describe, expect, test } from 'bun:test'
+import type { Message } from 'src/types/message.js'
+import { filterIncompleteToolCalls } from '../filterIncompleteToolCalls.js'
+
+describe('filterIncompleteToolCalls', () => {
+  test('drops assistant tool uses that do not have matching results', () => {
+    const messages = [
+      {
+        type: 'assistant',
+        uuid: 'a1',
+        message: {
+          role: 'assistant',
+          content: [{ type: 'tool_use', id: 'missing', name: 'Read' }],
+        },
+      },
+      {
+        type: 'user',
+        uuid: 'u1',
+        message: { role: 'user', content: 'continue' },
+      },
+    ] as unknown as Message[]
+
+    expect(
+      filterIncompleteToolCalls(messages).map(message => String(message.uuid)),
+    ).toEqual(['u1'])
+  })
+
+  test('preserves assistant text when dropping orphan tool uses', () => {
+    const messages = [
+      {
+        type: 'assistant',
+        uuid: 'a1',
+        message: {
+          role: 'assistant',
+          content: [
+            { type: 'text', text: 'I will read the file.' },
+            { type: 'tool_use', id: 'missing', name: 'Read' },
+          ],
+        },
+      },
+    ] as unknown as Message[]
+
+    const filtered = filterIncompleteToolCalls(messages)
+    expect(filtered).toHaveLength(1)
+    const first = filtered[0]!
+    const content = first.message!.content
+    expect(
+      Array.isArray(content) ? content.map(block => block.type) : [],
+    ).toEqual(['text'])
+  })
+
+  test('keeps completed parallel tool calls when dropping an orphan', () => {
+    const messages = [
+      {
+        type: 'assistant',
+        uuid: 'a1',
+        message: {
+          role: 'assistant',
+          content: [
+            { type: 'tool_use', id: 'done', name: 'Read' },
+            { type: 'tool_use', id: 'missing', name: 'Grep' },
+          ],
+        },
+      },
+      {
+        type: 'user',
+        uuid: 'u1',
+        message: {
+          role: 'user',
+          content: [{ type: 'tool_result', tool_use_id: 'done', content: 'ok' }],
+        },
+      },
+    ] as unknown as Message[]
+
+    const filtered = filterIncompleteToolCalls(messages)
+    expect(filtered.map(message => String(message.uuid))).toEqual(['a1', 'u1'])
+    const first = filtered[0]!
+    const content = first.message!.content
+    expect(
+      Array.isArray(content)
+        ? content.map(block =>
+            block.type === 'tool_use' ? block.id : block.type,
+          )
+        : [],
+    ).toEqual(['done'])
+  })
+
+  test('keeps assistant tool uses that have matching results', () => {
+    const messages = [
+      {
+        type: 'assistant',
+        uuid: 'a1',
+        message: {
+          role: 'assistant',
+          content: [{ type: 'tool_use', id: 'done', name: 'Read' }],
+        },
+      },
+      {
+        type: 'user',
+        uuid: 'u1',
+        message: {
+          role: 'user',
+          content: [{ type: 'tool_result', tool_use_id: 'done', content: 'ok' }],
+        },
+      },
+    ] as unknown as Message[]
+
+    expect(
+      filterIncompleteToolCalls(messages).map(message => String(message.uuid)),
+    ).toEqual(['a1', 'u1'])
+  })
+
+  test('drops orphan tool results when their tool use was removed', () => {
+    const messages = [
+      {
+        type: 'user',
+        uuid: 'u1',
+        message: {
+          role: 'user',
+          content: [
+            { type: 'tool_result', tool_use_id: 'missing', content: 'late' },
+          ],
+        },
+      },
+    ] as unknown as Message[]
+
+    expect(filterIncompleteToolCalls(messages)).toEqual([])
+  })
+
+  test('keeps user text while dropping orphan tool results', () => {
+    const messages = [
+      {
+        type: 'assistant',
+        uuid: 'a1',
+        message: { role: 'assistant', content: 'done' },
+      },
+      {
+        type: 'user',
+        uuid: 'u1',
+        message: {
+          role: 'user',
+          content: [
+            { type: 'text', text: 'keep this' },
+            { type: 'tool_result', tool_use_id: 'missing', content: 'late' },
+          ],
+        },
+      },
+    ] as unknown as Message[]
+
+    const filtered = filterIncompleteToolCalls(messages)
+    expect(filtered.map(message => String(message.uuid))).toEqual(['a1', 'u1'])
+    const content = filtered[1]!.message!.content
+    expect(Array.isArray(content) ? content : []).toEqual([
+      { type: 'text', text: 'keep this' },
+    ])
+  })
+
+  test('drops malformed tool blocks without ids', () => {
+    const messages = [
+      {
+        type: 'assistant',
+        uuid: 'a1',
+        message: {
+          role: 'assistant',
+          content: [{ type: 'tool_use', name: 'Read' }],
+        },
+      },
+      {
+        type: 'user',
+        uuid: 'u1',
+        message: {
+          role: 'user',
+          content: [{ type: 'tool_result', content: 'late' }],
+        },
+      },
+    ] as unknown as Message[]
+
+    expect(filterIncompleteToolCalls(messages)).toEqual([])
+  })
+})

packages/builtin-tools/src/tools/AgentTool/filterIncompleteToolCalls.ts

@@ -0,0 +1,110 @@
+import type {
+  AssistantMessage,
+  Message,
+  UserMessage,
+} from 'src/types/message.js'
+
+/**
+ * Removes invalid or orphaned tool_use/tool_result blocks while preserving
+ * completed tool-call pairs. This is intentionally block-level, not
+ * message-level, so completed parallel tool calls stay paired with results.
+ */
+export function filterIncompleteToolCalls(messages: Message[]): Message[] {
+  const toolUseIdsWithResults = new Set<string>()
+
+  for (const message of messages) {
+    if (message?.type === 'user') {
+      const userMessage = message as UserMessage
+      const content = userMessage.message.content
+      if (Array.isArray(content)) {
+        for (const block of content) {
+          if (block.type === 'tool_result' && block.tool_use_id) {
+            toolUseIdsWithResults.add(block.tool_use_id)
+          }
+        }
+      }
+    }
+  }
+
+  const retainedToolUseIds = new Set<string>()
+  const withoutOrphanToolUses: Message[] = []
+
+  for (const message of messages) {
+    if (message?.type === 'assistant') {
+      const assistantMessage = message as AssistantMessage
+      const content = assistantMessage.message.content
+      if (Array.isArray(content)) {
+        let changed = false
+        const filteredContent = content.filter(block => {
+          if (block.type !== 'tool_use') return true
+          if (!block.id) {
+            changed = true
+            return false
+          }
+          if (toolUseIdsWithResults.has(block.id)) {
+            retainedToolUseIds.add(block.id)
+            return true
+          }
+          changed = true
+          return false
+        })
+
+        if (!changed) {
+          withoutOrphanToolUses.push(message)
+          continue
+        }
+        if (filteredContent.length > 0) {
+          withoutOrphanToolUses.push({
+            ...assistantMessage,
+            message: {
+              ...assistantMessage.message,
+              content: filteredContent,
+            },
+          })
+        }
+        continue
+      }
+    }
+    withoutOrphanToolUses.push(message)
+  }
+
+  const filteredMessages: Message[] = []
+  for (const message of withoutOrphanToolUses) {
+    if (message?.type !== 'user') {
+      filteredMessages.push(message)
+      continue
+    }
+    const userMessage = message as UserMessage
+    const content = userMessage.message.content
+    if (!Array.isArray(content)) {
+      filteredMessages.push(message)
+      continue
+    }
+    let changed = false
+    const filteredContent = content.filter(block => {
+      if (block.type !== 'tool_result') return true
+      if (!block.tool_use_id) {
+        changed = true
+        return false
+      }
+      if (retainedToolUseIds.has(block.tool_use_id)) return true
+      changed = true
+      return false
+    })
+    if (!changed) {
+      filteredMessages.push(message)
+      continue
+    }
+    if (filteredContent.length > 0) {
+      filteredMessages.push({
+        ...userMessage,
+        message: {
+          ...userMessage.message,
+          content: filteredContent,
+        },
+      })
+    }
+  }
+
+  return filteredMessages
+}

packages/builtin-tools/src/tools/AgentTool/runAgent.ts

@@ -86,8 +86,11 @@ import {
 import type { ContentReplacementState } from 'src/utils/toolResultStorage.js'
 import { createAgentId } from 'src/utils/uuid.js'
 import { resolveAgentTools } from './agentToolUtils.js'
+import { filterIncompleteToolCalls } from './filterIncompleteToolCalls.js'
 import { type AgentDefinition, isBuiltInAgent } from './loadAgentsDir.js'
 
+export { filterIncompleteToolCalls } from './filterIncompleteToolCalls.js'
+
 /**
  * Initialize agent-specific MCP servers
  * Agents can define their own MCP servers in their frontmatter that are additive
@@ -886,50 +889,6 @@ export async function* runAgent({
   }
 }
 
-/**
- * Filters out assistant messages with incomplete tool calls (tool uses without results).
- * This prevents API errors when sending messages with orphaned tool calls.
- */
-export function filterIncompleteToolCalls(messages: Message[]): Message[] {
-  // Build a set of tool use IDs that have results
-  const toolUseIdsWithResults = new Set<string>()
-
-  for (const message of messages) {
-    if (message?.type === 'user') {
-      const userMessage = message as UserMessage
-      const content = userMessage.message.content
-      if (Array.isArray(content)) {
-        for (const block of content) {
-          if (block.type === 'tool_result' && block.tool_use_id) {
-            toolUseIdsWithResults.add(block.tool_use_id)
-          }
-        }
-      }
-    }
-  }
-
-  // Filter out assistant messages that contain tool calls without results
-  return messages.filter(message => {
-    if (message?.type === 'assistant') {
-      const assistantMessage = message as AssistantMessage
-      const content = assistantMessage.message.content
-      if (Array.isArray(content)) {
-        // Check if this assistant message has any tool uses without results
-        const hasIncompleteToolCall = content.some(
-          block =>
-            block.type === 'tool_use' &&
-            block.id &&
-            !toolUseIdsWithResults.has(block.id),
-        )
-        // Exclude messages with incomplete tool calls
-        return !hasIncompleteToolCall
-      }
-    }
-    // Keep all non-assistant messages and assistant messages without tool calls
-    return true
-  })
-}
-
 async function getAgentSystemPrompt(
   agentDefinition: AgentDefinition,
   toolUseContext: Pick<ToolUseContext, 'options'>,