perf: bound tool pairing and avoid unused attribution arrays (#19468)

* perf: bound tool pairing and avoid unused attribution arrays

* test(native-chat): cover mobile pair-limit shapes and group allocation tests

---------

Co-authored-by: m4air <m4air@m4airs-MacBook-Air.local>
Co-authored-by: Neil <4138956+nwparker@users.noreply.github.com>

src/shared/native-chat-tool-attribution-allocation.test.ts

@@ -0,0 +1,50 @@
+import { describe, expect, it } from 'vitest'
+import { foldToolMessages } from './native-chat-tool-fold'
+import type { NativeChatBlock, NativeChatMessage } from './native-chat-types'
+
+function message(id: string, blocks: NativeChatBlock[]): NativeChatMessage {
+  return {
+    id,
+    role: 'assistant',
+    blocks,
+    timestamp: null,
+    source: 'transcript'
+  }
+}
+
+describe('tool attribution allocation', () => {
+  it('does not append valid prose blocks to discarded attribution arrays', () => {
+    const prose: NativeChatBlock = { type: 'text', text: 'Ordinary prose' }
+    const messages = Array.from({ length: 1000 }, (_, i) => message(String(i), [prose]))
+    const push = Array.prototype.push
+    let appends = 0
+    Array.prototype.push = function (this: unknown[], ...items: unknown[]) {
+      if (items[0] === prose) {
+        appends += items.length
+      }
+      return push.apply(this, items)
+    }
+    let output: NativeChatMessage[]
+    try {
+      output = foldToolMessages(messages)
+    } finally {
+      Array.prototype.push = push
+    }
+    expect(appends).toBe(0)
+    output.forEach((entry, i) => expect(entry).toBe(messages[i]))
+  })
+
+  it('removes only unattributable results while preserving subsequent call/result pairs', () => {
+    const text: NativeChatBlock = { type: 'text', text: 'Prose' }
+    const call: NativeChatBlock = {
+      type: 'tool-call',
+      name: 'read',
+      input: {}
+    }
+    const result: NativeChatBlock = { type: 'tool-result', output: 'done' }
+    const input = message('one', [text, result, text, call, result, result, text])
+    expect(foldToolMessages([input])[0].blocks).toEqual([text, text, call, result, text])
+    expect(input.blocks).toHaveLength(7)
+    expect(foldToolMessages([message('two', [result])])).toEqual([])
+  })
+})

src/shared/native-chat-tool-fold.ts

@@ -52,20 +52,22 @@ function isInterruptionBoundary(message: NativeChatMessage): boolean {
 
 /** Drop tool results the renderer cannot pair within their folded message. */
 function dropUnattributableToolResults(message: NativeChatMessage): NativeChatMessage | null {
-  const blocks: NativeChatBlock[] = []
+  let blocks: NativeChatBlock[] | undefined
   let unansweredCalls = 0
-  for (const block of message.blocks) {
+  for (let index = 0; index < message.blocks.length; index++) {
+    const block = message.blocks[index]
     if (isToolCallBlock(block)) {
       unansweredCalls += 1
     } else if (isToolResultBlock(block)) {
       if (unansweredCalls === 0) {
+        blocks ??= message.blocks.slice(0, index)
         continue
       }
       unansweredCalls -= 1
     }
-    blocks.push(block)
+    blocks?.push(block)
   }
-  if (blocks.length === message.blocks.length) {
+  if (!blocks) {
     return message
   }
   return blocks.length > 0 ? { ...message, blocks } : null
@@ -141,15 +143,16 @@ export function pairToolBlocks(
   limit = Infinity
 ): NativeChatToolPair[] {
   const pairs: NativeChatToolPair[] = []
-  const callSlots: (number | null)[] = []
+  const callSlots: number[] = []
   let resultOrdinal = 0
   for (const block of blocks) {
+    if (pairs.length >= limit && resultOrdinal >= callSlots.length) {
+      break
+    }
     if (block.type === 'tool-call') {
       if (pairs.length < limit) {
         callSlots.push(pairs.length)
         pairs.push({ call: block })
-      } else {
-        callSlots.push(null)
       }
       continue
     }
@@ -163,9 +166,7 @@ export function pairToolBlocks(
       }
     } else {
       resultOrdinal += 1
-      if (slot !== null) {
-        pairs[slot]!.result = block
-      }
+      pairs[slot]!.result = block
     }
   }
   return pairs

src/shared/native-chat-tool-pair-limit.test.ts

@@ -0,0 +1,85 @@
+import { describe, expect, it } from 'vitest'
+import { pairToolBlocks, type NativeChatToolPair } from './native-chat-tool-fold'
+import type { NativeChatBlock } from './native-chat-types'
+
+function original(blocks: readonly NativeChatBlock[], limit: number): NativeChatToolPair[] {
+  const pairs: NativeChatToolPair[] = []
+  const slots: (number | null)[] = []
+  let ordinal = 0
+  for (const block of blocks) {
+    if (block.type === 'tool-call') {
+      if (pairs.length < limit) {
+        slots.push(pairs.length)
+        pairs.push({ call: block })
+      } else {
+        slots.push(null)
+      }
+    } else if (block.type === 'tool-result') {
+      const slot = slots[ordinal]
+      if (slot === undefined) {
+        if (pairs.length < limit) {
+          pairs.push({ result: block })
+        }
+      } else {
+        ordinal++
+        if (slot !== null) {
+          pairs[slot].result = block
+        }
+      }
+    }
+  }
+  return pairs
+}
+
+const call: NativeChatBlock = { type: 'tool-call', name: 'read', input: {} }
+const result: NativeChatBlock = { type: 'tool-result', output: 'done' }
+
+describe('tool pair limits', () => {
+  it('stops visiting blocks once retained pairs are fully answered', () => {
+    let reads = 0
+    const tail = Array.from({ length: 10000 }, () => ({
+      get type() {
+        reads++
+        return 'tool-call' as const
+      },
+      name: 'read',
+      input: {}
+    }))
+    expect(pairToolBlocks([call, result, ...tail], 1)).toEqual([{ call, result }])
+    expect(reads).toBe(0)
+  })
+
+  it('preserves FIFO and stray-result behavior across finite and unlimited limits', () => {
+    for (let seed = 0; seed < 128; seed++) {
+      const blocks = Array.from({ length: 30 }, (_, i) =>
+        (seed * 13 + i * 17) % 7 < 3
+          ? call
+          : (seed + i) % 3
+            ? result
+            : { type: 'text' as const, text: 'hi' }
+      )
+      for (const limit of [0, 1, 2, 5, 0.5, -1, Infinity, Number.NaN]) {
+        expect(pairToolBlocks(blocks, limit)).toEqual(original(blocks, limit))
+      }
+    }
+    expect(pairToolBlocks([call, call, result, result], 1)).toEqual(
+      original([call, call, result, result], 1)
+    )
+  })
+  it('still answers every retained call when the results trail far behind', () => {
+    // The break must not fire while a retained call is unanswered, or the mobile run
+    // would render a spinner for a tool that actually completed.
+    const blocks = [call, call, ...Array.from({ length: 5000 }, () => result)]
+    expect(pairToolBlocks(blocks, 2)).toEqual([
+      { call, result },
+      { call, result }
+    ])
+    expect(pairToolBlocks(blocks, 2)).toEqual(original(blocks, 2))
+  })
+
+  it('keeps a leading stray result and then stops at the limit', () => {
+    const blocks = [result, call, result, call, result]
+    expect(pairToolBlocks(blocks, 1)).toEqual(original(blocks, 1))
+    expect(pairToolBlocks(blocks, 1)).toEqual([{ result }])
+  })
+})