From 02db0146c620668fe59f0ff208765d6995555708 Mon Sep 17 00:00:00 2001 From: OrcaWin Date: Wed, 28 Oct 2026 20:07:52 -0700 Subject: [PATCH] perf(chat): order queued turns without repeated transcript scans (#34945) Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com> --- ...ative-chat-turn-draw-order-scaling.test.ts | 87 +++++++++++++++++++ .../transcript/native-chat-turn-grouping.ts | 60 +++++++++---- 2 files changed, 132 insertions(+), 15 deletions(-) create mode 100644 mobile/src/session/native-chat/transcript/native-chat-turn-draw-order-scaling.test.ts diff --git a/mobile/src/session/native-chat/transcript/native-chat-turn-draw-order-scaling.test.ts b/mobile/src/session/native-chat/transcript/native-chat-turn-draw-order-scaling.test.ts new file mode 100644 index 00000000000..5c7e8a08390 --- /dev/null +++ b/mobile/src/session/native-chat/transcript/native-chat-turn-draw-order-scaling.test.ts @@ -0,0 +1,87 @@ +import { expect, it } from 'vitest' +import { nativeChatTurnDrawOrder } from '../../../../../src/shared/native-chat/transcript/native-chat-turn-grouping' +import { nativeChatTurnMembership } from '../../../../../src/shared/native-chat-turn-membership' +import type { + AgentJournalItemBody, + AgentJournalRenderItem +} from '../../../../../src/shared/agent-session-journal-types' + +it.each([10, 100])( + 'reads turn ownership linearly for %i queued openers and earlier-turn rows', + (count) => { + const messages = [ + { id: 'first', role: 'user' as const }, + ...Array.from({ length: count }, (_, index) => ({ + id: `queued-${index}`, + role: 'user' as const + })), + ...Array.from({ length: count }, (_, index) => ({ + id: `tail-${index}`, + role: 'assistant' as const + })), + ...Array.from({ length: count }, (_, index) => ({ + id: `answer-${index}`, + role: 'assistant' as const + })) + ] + const keys = [ + 'first', + ...Array.from({ length: count }, (_, index) => `queued-${index}`), + ...Array.from({ length: count }, () => 'first'), + ...Array.from({ length: count }, (_, index) => `queued-${index}`) + ] + const openers = new Set(keys.slice(0, count + 1)) + let reads = 0 + keys.forEach((key, index) => { + Object.defineProperty(keys, index, { + get: () => { + reads += 1 + return key + } + }) + }) + const order = nativeChatTurnDrawOrder(messages, keys, openers) + const visits = reads + expect(order?.map((index) => messages[index]?.id)).toEqual([ + 'first', + ...Array.from({ length: count }, (_, index) => `tail-${index}`), + ...Array.from({ length: count }, (_, index) => [`queued-${index}`, `answer-${index}`]).flat() + ]) + const journalItems: AgentJournalRenderItem[] = [] + const append = (itemId: string, body: AgentJournalItemBody, turnItemId: string) => { + const sequence = journalItems.length + 1 + journalItems.push({ + itemId, + body, + revision: 0, + sequence, + observedAt: sequence, + turnScope: { kind: 'turn', turnItemId } + }) + } + append('first', { kind: 'message', role: 'user', blocks: [] }, 'turn-first') + append( + 'turn-first', + { kind: 'turn', turnId: 'first', state: 'completed', userItemId: 'first' }, + 'turn-first' + ) + for (let index = 0; index < count; index += 1) { + append(`queued-${index}`, { kind: 'message', role: 'user', blocks: [] }, `turn-${index}`) + } + for (let index = 0; index < count; index += 1) { + append(`tail-${index}`, { kind: 'message', role: 'assistant', blocks: [] }, 'turn-first') + } + for (let index = 0; index < count; index += 1) { + append( + `turn-${index}`, + { kind: 'turn', turnId: `t${index}`, state: 'completed', userItemId: `queued-${index}` }, + `turn-${index}` + ) + append(`answer-${index}`, { kind: 'message', role: 'assistant', blocks: [] }, `turn-${index}`) + } + const membership = nativeChatTurnMembership(messages, { items: journalItems, submissions: [] }) + expect(membership.turnKeys).toEqual(keys) + expect(membership.drawOrder).toEqual(order) + expect(visits).toBeLessThanOrEqual(messages.length * 4) + } +) diff --git a/src/shared/native-chat/transcript/native-chat-turn-grouping.ts b/src/shared/native-chat/transcript/native-chat-turn-grouping.ts index d326ef31810..26a6e8f0770 100644 --- a/src/shared/native-chat/transcript/native-chat-turn-grouping.ts +++ b/src/shared/native-chat/transcript/native-chat-turn-grouping.ts @@ -102,32 +102,62 @@ export function nativeChatTurnDrawOrder( firstRow.set(turnKey, index) } }) - const movedAfter = new Map() - const moved = new Set() + const moved = new Map() + const waiting = new Set() + const waitingByTurn = new Map() + let lastContentRow = -1 + const finish = (index: number): void => { + if (waiting.delete(index) && lastContentRow > index) { + moved.set(index, lastContentRow) + } + } for (const [index, message] of messages.entries()) { - if (!opensTurn(index)) { + if (opensTurn(index)) { + waiting.add(index) + const turn = waitingByTurn.get(message.id) + if (turn) { + turn.push(index) + } else { + waitingByTurn.set(message.id, [index]) + } continue } - let after = -1 - for (let row = index + 1; row < messages.length; row += 1) { - const turnKey = turnKeys[row] - if (turnKey === undefined || opensTurn(row)) { - continue - } - // Its own turn's rows, or a later turn's, end the earlier turns' rows it waited behind. - if (turnKey === message.id || (firstRow.get(turnKey) ?? row) > index) { + const turnKey = turnKeys[index] + if (turnKey === undefined) { + continue + } + const first = firstRow.get(turnKey) ?? index + for (const openerIndex of waiting) { + if (openerIndex >= first) { break } - after = row + finish(openerIndex) } - if (after !== -1) { - moved.add(index) - movedAfter.set(after, [...(movedAfter.get(after) ?? []), index]) + for (const openerIndex of waitingByTurn.get(turnKey) ?? []) { + finish(openerIndex) } + waitingByTurn.delete(turnKey) + lastContentRow = index + } + for (const index of waiting) { + finish(index) } if (moved.size === 0) { return null } + const movedAfter = new Map() + for (let index = 0; index < messages.length; index += 1) { + const after = moved.get(index) + if (after === undefined) { + continue + } + const indexes = movedAfter.get(after) + if (indexes) { + indexes.push(index) + } else { + movedAfter.set(after, [index]) + } + } return messages.flatMap((_, index) => [ ...(moved.has(index) ? [] : [index]), ...(movedAfter.get(index) ?? [])