perf(chat): order queued turns without repeated transcript scans (#34945)

Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com>

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)
+  }
+)

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<number, number[]>()
-  const moved = new Set<number>()
+  const moved = new Map<number, number>()
+  const waiting = new Set<number>()
+  const waitingByTurn = new Map<string, number[]>()
+  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<number, number[]>()
+  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) ?? [])