From 44eb95fc6bb96fbb32bb2ebda59b8c1f908e1a71 Mon Sep 17 00:00:00 2001 From: OrcaWin Date: Tue, 8 Sep 2026 19:42:18 -0700 Subject: [PATCH] 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 Co-authored-by: Neil <4138956+nwparker@users.noreply.github.com> --- ...e-chat-tool-attribution-allocation.test.ts | 50 +++++++++++ src/shared/native-chat-tool-fold.ts | 21 ++--- .../native-chat-tool-pair-limit.test.ts | 85 +++++++++++++++++++ 3 files changed, 146 insertions(+), 10 deletions(-) create mode 100644 src/shared/native-chat-tool-attribution-allocation.test.ts create mode 100644 src/shared/native-chat-tool-pair-limit.test.ts diff --git a/src/shared/native-chat-tool-attribution-allocation.test.ts b/src/shared/native-chat-tool-attribution-allocation.test.ts new file mode 100644 index 00000000000..48d28166c1f --- /dev/null +++ b/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([]) + }) +}) diff --git a/src/shared/native-chat-tool-fold.ts b/src/shared/native-chat-tool-fold.ts index c334bf0b1a0..32a2531b9ee 100644 --- a/src/shared/native-chat-tool-fold.ts +++ b/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 diff --git a/src/shared/native-chat-tool-pair-limit.test.ts b/src/shared/native-chat-tool-pair-limit.test.ts new file mode 100644 index 00000000000..f279b70c996 --- /dev/null +++ b/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 }]) + }) +})