fix(cli): avoid quadratic long-session normalization #1100

src/utils/messages.test.ts

@@ -0,0 +1,120 @@
+import { describe, expect, test } from 'bun:test'
+import type { ContentBlockParam } from '@anthropic-ai/sdk/resources/index.mjs'
+import type { AssistantMessage } from '../types/message.js'
+import {
+  createAssistantMessage,
+  createUserMessage,
+  normalizeMessagesForAPI,
+} from './messages.js'
+
+function assistant(
+  messageId: string,
+  content: AssistantMessage['message']['content'],
+): AssistantMessage {
+  const message = createAssistantMessage({ content })
+  message.message.id = messageId
+  return message
+}
+
+function toolUse(id: string): AssistantMessage['message']['content'][number] {
+  return {
+    type: 'tool_use',
+    id,
+    name: 'Read',
+    input: { file_path: `/tmp/${id}` },
+  }
+}
+
+function toolResult(id: string) {
+  return createUserMessage({
+    content: [
+      {
+        type: 'tool_result',
+        tool_use_id: id,
+        content: 'ok',
+      },
+    ] as ContentBlockParam[],
+  })
+}
+
+describe('normalizeMessagesForAPI assistant fragment indexing', () => {
+  test('preserves a 10,000-step tool-result chain', () => {
+    const messages = [createUserMessage({ content: 'start' })]
+
+    for (let i = 0; i < 10_000; i++) {
+      const toolId = `tool-${i}`
+      messages.push(
+        assistant(`response-${i}`, [toolUse(toolId)]),
+        toolResult(toolId),
+      )
+    }
+
+    const normalized = normalizeMessagesForAPI(messages)
+    const assistants = normalized.filter(
+      (message): message is AssistantMessage => message.type === 'assistant',
+    )
+    const toolResults = normalized.filter(message => message.type === 'user')
+
+    expect(normalized).toHaveLength(20_001)
+    expect(assistants).toHaveLength(10_000)
+    expect(toolResults).toHaveLength(10_001)
+    expect(assistants.at(-1)?.message.id).toBe('response-9999')
+    expect(toolResults.at(-1)?.message.content).toEqual([
+      {
+        type: 'tool_result',
+        tool_use_id: 'tool-9999',
+        content: 'ok',
+      },
+    ])
+  })
+
+  test('merges interleaved response IDs across tool-result messages', () => {
+    const normalized = normalizeMessagesForAPI([
+      assistant('response-a', [toolUse('tool-a')]),
+      toolResult('tool-a'),
+      assistant('response-b', [toolUse('tool-b')]),
+      toolResult('tool-b'),
+      assistant('response-a', [{ type: 'text', text: 'A complete' }]),
+      assistant('response-b', [{ type: 'text', text: 'B complete' }]),
+    ])
+
+    const assistants = normalized.filter(
+      (message): message is AssistantMessage => message.type === 'assistant',
+    )
+
+    expect(assistants.map(message => message.message.id)).toEqual([
+      'response-a',
+      'response-b',
+    ])
+    expect(assistants[0]!.message.content.map(block => block.type)).toEqual([
+      'tool_use',
+      'text',
+    ])
+    expect(assistants[1]!.message.content.map(block => block.type)).toEqual([
+      'tool_use',
+      'text',
+    ])
+  })
+
+  test('does not merge the same response ID across a normal user turn', () => {
+    const normalized = normalizeMessagesForAPI([
+      assistant('response-a', [{ type: 'text', text: 'before' }]),
+      createUserMessage({ content: 'next turn' }),
+      assistant('response-a', [{ type: 'text', text: 'after' }]),
+    ])
+
+    const assistants = normalized.filter(
+      (message): message is AssistantMessage => message.type === 'assistant',
+    )
+
+    expect(assistants).toHaveLength(2)
+    expect(
+      assistants.map(message =>
+        message.message.content
+          .filter(block => block.type === 'text')
+          .map(block => block.text)
+          .join(''),
+      ),
+    ).toEqual(['before', 'after'])
+  })
+})

src/utils/messages.ts

@@ -2083,6 +2083,8 @@ export function normalizeMessagesForAPI(
   }
 
   const result: (UserMessage | AssistantMessage)[] = []
+  const assistantIndexByMessageId = new Map<string, number>()
+  let indexedResultLength = 0
   reorderedMessages
     .filter(
       (
@@ -2272,24 +2274,36 @@ export function normalizeMessagesForAPI(
             },
           }
 
-          // Find a previous assistant message with the same message ID and merge.
-          // Walk backwards, skipping tool results and different-ID assistants,
-          // since concurrent agents (teammates) can interleave streaming content
-          // blocks from multiple API responses with different message IDs.
-          for (let i = result.length - 1; i >= 0; i--) {
-            const msg = result[i]!
-
-            if (msg.type !== 'assistant' && !isToolResultMessage(msg)) {
-              break
+          // Index each result message once so interleaved assistant fragments
+          // can be merged without repeatedly scanning the full tool-result chain.
+          // A normal user message is a hard turn boundary, matching the previous
+          // backward scan's stopping condition.
+          for (; indexedResultLength < result.length; indexedResultLength++) {
+            const indexedMessage = result[indexedResultLength]!
+            if (indexedMessage.type === 'assistant') {
+              assistantIndexByMessageId.set(
+                indexedMessage.message.id,
+                indexedResultLength,
+              )
+            } else if (!isToolResultMessage(indexedMessage)) {
+              assistantIndexByMessageId.clear()
             }
+          }
 
-            if (msg.type === 'assistant') {
-              if (msg.message.id === normalizedMessage.message.id) {
-                result[i] = mergeAssistantMessages(msg, normalizedMessage)
-                return
-              }
-              continue
-            }
+          const existingIndex = assistantIndexByMessageId.get(
+            normalizedMessage.message.id,
+          )
+          const existingMessage =
+            existingIndex === undefined ? undefined : result[existingIndex]
+          if (
+            existingIndex !== undefined &&
+            existingMessage?.type === 'assistant'
+          ) {
+            result[existingIndex] = mergeAssistantMessages(
+              existingMessage,
+              normalizedMessage,
+            )
+            return
           }
 
           result.push(normalizedMessage)