From 544169a41abc721099ef5f2ee4b5ee42d81d746c Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 9 Sep 2026 05:32:36 -0700 Subject: [PATCH] perf(desktop): stop per-event full-timeline allocations on streamed tool events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every tool.start / tool.progress / tool.complete ran three whole-timeline passes that each allocated per part: `completeOpenStreamParts` (`map` with a spread per row), the pending-row scan (`map` → `filter` → `map`, one wrapper object per part), and `generatedImageEchoSources` (`flatMap` + a throwaway array per non-image row) — then `dedupeGeneratedImageEchoesInParts` copied the array again even when it stripped nothing, handing React a fresh identity per event. Replace them with single in-place loops, and return the input array from the dedupe when there are no echoes so a no-op stays a no-op for the store. Measured on the real `useMessageStream` hook (vitest/jsdom, 2,000 starts + 2,000 completions interleaved with text deltas): 618 ms → 138 ms; the tool-only 2,000×2 case 252 ms → 67 ms. The remaining O(n) is one `findIndex` per event over a ~4k array (~8 µs), well under the cost of the copy React needs anyway, so no per-stream id index is introduced. Behaviour is unchanged: same first-match `toolCallId` resolution, same sparse/no-id contextual correlation, same ordering of parts. Co-authored-by: Xipong <217837358+Xipong@users.noreply.github.com> --- .../src/lib/chat-messages/tool-parts.ts | 35 ++++++++++++------- apps/desktop/src/lib/generated-images.test.ts | 6 ++++ apps/desktop/src/lib/generated-images.ts | 20 +++++++++-- 3 files changed, 46 insertions(+), 15 deletions(-) diff --git a/apps/desktop/src/lib/chat-messages/tool-parts.ts b/apps/desktop/src/lib/chat-messages/tool-parts.ts index 881a564c8a..392ec58b4f 100644 --- a/apps/desktop/src/lib/chat-messages/tool-parts.ts +++ b/apps/desktop/src/lib/chat-messages/tool-parts.ts @@ -179,10 +179,15 @@ function findToolPartIndex( } } - const pendingIndices = parts - .map((part, index) => ({ part, index })) - .filter(({ part }) => part.type === 'tool-call' && part.toolName === name && part.result === undefined) - .map(({ index }) => index) + const pendingIndices: number[] = [] + + for (let index = 0; index < parts.length; index += 1) { + const part = parts[index] + + if (part.type === 'tool-call' && part.toolName === name && part.result === undefined) { + pendingIndices.push(index) + } + } if (pendingIndices.length === 0) { return -1 @@ -279,11 +284,17 @@ function toolResult( } function completeOpenStreamParts(parts: ChatMessagePart[], completedAt: number): ChatMessagePart[] { - return parts.map(part => - (part.type === 'text' || part.type === 'reasoning') && part.completedAt === undefined - ? ({ ...part, completedAt } as ChatMessagePart) - : part - ) + const next = parts.slice() + + for (let index = 0; index < next.length; index += 1) { + const part = next[index] + + if ((part.type === 'text' || part.type === 'reasoning') && part.completedAt === undefined) { + next[index] = { ...part, completedAt } as ChatMessagePart + } + } + + return next } export function upsertToolPart( @@ -325,11 +336,11 @@ export function upsertToolPart( } satisfies ChatMessagePart if (index === -1) { - return [...next, base] + next.push(base) + } else { + next[index] = { ...next[index], ...base } } - next[index] = { ...next[index], ...base } - return next } diff --git a/apps/desktop/src/lib/generated-images.test.ts b/apps/desktop/src/lib/generated-images.test.ts index 2f08bf5210..7fb3210510 100644 --- a/apps/desktop/src/lib/generated-images.test.ts +++ b/apps/desktop/src/lib/generated-images.test.ts @@ -60,6 +60,12 @@ describe('generatedImageEchoSources', () => { }) describe('dedupeGeneratedImageEchoesInParts', () => { + it('preserves timeline identity when there is nothing to strip', () => { + const parts = [{ result: {}, toolName: 'read_file', type: 'tool-call' }] + + expect(dedupeGeneratedImageEchoesInParts(parts)).toBe(parts) + }) + it('keeps the agent prose while removing the duplicated image', () => { expect( dedupeGeneratedImageEchoesInParts([ diff --git a/apps/desktop/src/lib/generated-images.ts b/apps/desktop/src/lib/generated-images.ts index 9a9243796f..4ee990bd6d 100644 --- a/apps/desktop/src/lib/generated-images.ts +++ b/apps/desktop/src/lib/generated-images.ts @@ -68,7 +68,19 @@ export function generatedImageFromResult(result: unknown): string | null { /** Every path/URL a generated image might appear as in prose, for de-duping. */ export function generatedImageEchoSources(parts: readonly ToolLike[]): string[] { - return unique(parts.flatMap(part => stringFields(imageResult(part) ?? {}, ECHO_KEYS))) + // Runs on every streamed tool event over the whole timeline; skip the + // per-part array allocations for the (usual) rows that are not image results. + const sources: string[] = [] + + for (const part of parts) { + const result = imageResult(part) + + if (result) { + sources.push(...stringFields(result, ECHO_KEYS)) + } + } + + return unique(sources) } /** Strip a generated image out of prose so it only ever shows in the tool slot. @@ -94,11 +106,13 @@ export function stripGeneratedImageEchoes(text: string, sources: readonly string /** Strip generated-image echoes from text parts, dropping any part left empty. * The image lives in the tool slot; prose keeps the agent's actual words. */ -export function dedupeGeneratedImageEchoesInParts(parts: readonly T[]): T[] { +export function dedupeGeneratedImageEchoesInParts(parts: T[]): T[] { const sources = generatedImageEchoSources(parts) + // No echoes to strip: hand back the same array so React and the streaming + // reducer see a no-op instead of a fresh identity on every tool event. if (!sources.length) { - return [...parts] + return parts } return parts