perf(ui-trajectory): index trajectory snapshot assembly

packages/client/ui-trajectory/src/client/trajectory-snapshot-builder.ts

@@ -11,6 +11,8 @@ import type {
 
 const EMPTY_LIST: readonly never[] = []
 const EMPTY_CONTEXTS = [{ id: 0, nodes: EMPTY_LIST }]
+type AssistantRequest = Extract<RequestView, { purpose: 'assistant' }>
+type ToolSchema = ConversationPromptSnapshot['tools'][number]
 
 /** Stable empty target used until a Session has assembled Trajectory records. */
 export const EMPTY_TRAJECTORY_SNAPSHOT: TrajectorySnapshot = {
@@ -23,31 +25,31 @@ export const EMPTY_TRAJECTORY_SNAPSHOT: TrajectorySnapshot = {
   runningCalls: EMPTY_LIST,
 }
 
-function coordinates(
+function stepKey(turn: number, step: number): string {
-  header: TrajectoryRequestHeaderState,
+  return `${turn}\u0000${step}`
-): { turn?: number; step?: number } {
+}
+
+function headerStepKey(header: TrajectoryRequestHeaderState): string | undefined {
   const location = header.location
-  if (location.kind === 'step') return { turn: location.turn.turn, step: location.step.step }
+  return location.kind === 'step'
-  if (location.kind === 'turn') return { turn: location.turn.turn }
+    ? stepKey(location.turn.turn, location.step.step)
-  return {}
+    : undefined
 }
 
 function headerFor(
-  request: Extract<RequestView, { purpose: 'assistant' }>,
+  request: AssistantRequest,
-  headers: readonly TrajectoryRequestHeaderState[],
+  headersByStep: ReadonlyMap<string, TrajectoryRequestHeaderState>,
+  previous: TrajectoryRequestHeaderState | undefined,
 ): TrajectoryRequestHeaderState | undefined {
-  const exact = headers.findLast((header) => {
+  return headersByStep.get(stepKey(request.turn, request.step))
-    const location = coordinates(header)
+    ?? (previous !== undefined && previous.seq < request.startSeq ? previous : undefined)
-    return location.turn === request.turn && location.step === request.step
-  })
-  return exact ?? headers.findLast(header => header.seq < request.startSeq)
 }
 
 function applyHeader(
-  request: Extract<RequestView, { purpose: 'assistant' }>,
+  request: AssistantRequest,
   header: TrajectoryRequestHeaderState | undefined,
   includeChange: boolean,
-): Extract<RequestView, { purpose: 'assistant' }> {
+): AssistantRequest {
   return header === undefined
     ? request
     : {
@@ -67,26 +69,39 @@ function withRequestConfig(
 
 function captureSchemas(
   block: ToolCallBlock,
-  tools: readonly ConversationPromptSnapshot['tools'][number][],
+  toolsByName: ReadonlyMap<string, ToolSchema>,
-  output: Map<string, ConversationPromptSnapshot['tools'][number]>,
+  output: Map<string, ToolSchema>,
 ): void {
   const name = 'kind' in block ? block.call?.name : block.name
-  const schema = name === undefined
+  const schema = name === undefined ? undefined : toolsByName.get(name)
-    ? undefined
-    : tools.find(candidate => candidate.name === name)
   if (schema !== undefined) output.set(block.callId, schema)
-  for (const child of block.subCalls) captureSchemas(child, tools, output)
+  for (const child of block.subCalls) captureSchemas(child, toolsByName, output)
+}
+
+function indexTools(tools: readonly ToolSchema[]): ReadonlyMap<string, ToolSchema> {
+  return new Map(tools.map(tool => [tool.name, tool]))
 }
 
 function interruptCompactions(
   requests: RequestView[],
   boundaries: readonly { seq: number; time: number }[],
 ): void {
+  let nextRequest = 0
+  const runningCompactions: number[] = []
   for (const boundary of boundaries) {
-    const index = requests.findLastIndex(request =>
+    while (nextRequest < requests.length) {
-      request.purpose === 'compaction'
+      const request = requests[nextRequest]
-      && request.startSeq < boundary.seq
+      if (request === undefined || request.startSeq >= boundary.seq) break
-      && request.status === 'running')
+      if (request.purpose === 'compaction' && request.status === 'running') {
+        runningCompactions.push(nextRequest)
+      }
+      nextRequest++
+    }
+    let index = runningCompactions.pop()
+    while (index !== undefined && requests[index]?.status !== 'running') {
+      index = runningCompactions.pop()
+    }
+    if (index === undefined) continue
     const request = requests[index]
     if (request?.purpose !== 'compaction') continue
     requests[index] = {
@@ -102,10 +117,14 @@ function applyTurnErrors(
   requests: RequestView[],
   endings: readonly { turn: number; time: number; error?: string }[],
 ): void {
+  const lastAssistantByTurn = new Map<number, number>()
+  for (const [index, request] of requests.entries()) {
+    if (request.purpose === 'assistant') lastAssistantByTurn.set(request.turn, index)
+  }
   for (const ending of endings) {
     if (ending.error === undefined) continue
-    const index = requests.findLastIndex(request =>
+    const index = lastAssistantByTurn.get(ending.turn)
-      request.purpose === 'assistant' && request.turn === ending.turn)
+    if (index === undefined) continue
     const request = requests[index]
     if (request?.purpose !== 'assistant') continue
     requests[index] = {
@@ -123,6 +142,8 @@ export class TrajectorySnapshotBuilder implements ConversationViewBuilder<
   TrajectorySnapshot
 > {
   private readonly nodes = new Map<string, TrajectoryConversationViewNode>()
+  private readonly positions = new Map<string, number>()
+  private contributions: TrajectoryConversationViewNode[] = []
   readonly empty = EMPTY_TRAJECTORY_SNAPSHOT
 
   replace(input: {
@@ -130,39 +151,62 @@ export class TrajectorySnapshotBuilder implements ConversationViewBuilder<
   }): TrajectorySnapshot {
     this.nodes.clear()
     for (const node of input.nodes) this.nodes.set(node.key, node)
+    this.rebuildContributions()
     return this.snapshot()
   }
 
   apply(input: {
     readonly upserts: readonly TrajectoryConversationViewNode[]
   }): TrajectorySnapshot {
-    for (const node of input.upserts) this.nodes.set(node.key, node)
+    let structural = false
+    for (const node of input.upserts) {
+      const previous = this.nodes.get(node.key)
+      this.nodes.set(node.key, node)
+      if (previous === undefined || previous.anchorSeq !== node.anchorSeq) {
+        structural = true
+        continue
+      }
+      const position = this.positions.get(node.key)
+      if (position === undefined) structural = true
+      else this.contributions[position] = node
+    }
+    if (structural) this.rebuildContributions()
     return this.snapshot()
   }
 
   private snapshot(): TrajectorySnapshot {
-    const contributions = [...this.nodes.values()]
+    const headersByStep = new Map<string, TrajectoryRequestHeaderState>()
-      .sort((left, right) => left.anchorSeq - right.anchorSeq || left.key.localeCompare(right.key))
+    for (const contribution of this.contributions) {
-    const headers = contributions.flatMap(node => node.data.kind === 'request-header'
+      if (contribution.data.kind !== 'request-header') continue
-      ? [node.data.header]
+      const key = headerStepKey(contribution.data.header)
-      : [])
+      if (key !== undefined) headersByStep.set(key, contribution.data.header)
+    }
     const finalized: ConversationNode[] = []
     const requests: RequestView[] = []
     const boundaries: { seq: number; time: number }[] = []
     const turnEndings: { turn: number; time: number; error?: string }[] = []
-    const callSchemas = new Map<string, ConversationPromptSnapshot['tools'][number]>()
+    const callSchemas = new Map<string, ToolSchema>()
     const consumedPromptChanges = new Set<number>()
+    let previousHeader: TrajectoryRequestHeaderState | undefined
+    let previousTools: ReadonlyMap<string, ToolSchema> = new Map()
     let partial: TrajectorySnapshot['partial'] = null
     const runningCalls: TrajectorySnapshot['runningCalls'][number][] = []
 
-    for (const contribution of contributions) {
+    for (const contribution of this.contributions) {
       const data = contribution.data
+      if (data.kind === 'request-header') {
+        previousHeader = data.header
+        previousTools = indexTools(data.header.prompt.tools)
+        continue
+      }
       if (data.kind === 'node') {
         finalized.push(data.node)
         continue
       }
       if (data.kind === 'assistant') {
-        const header = data.request === undefined ? undefined : headerFor(data.request, headers)
+        const header = data.request === undefined
+          ? undefined
+          : headerFor(data.request, headersByStep, previousHeader)
         if (data.node !== undefined) finalized.push(withRequestConfig(data.node, header?.prompt))
         if (data.partial !== null) partial = data.partial
         if (data.request !== undefined) {
@@ -176,8 +220,9 @@ export class TrajectorySnapshotBuilder implements ConversationViewBuilder<
       if (data.kind === 'tool') {
         if ('kind' in data.root) finalized.push(data.root)
         else runningCalls.push(data.root)
-        const header = headers.findLast(candidate => candidate.seq < contribution.anchorSeq)
+        if (previousHeader !== undefined && previousHeader.seq < contribution.anchorSeq) {
-        if (header !== undefined) captureSchemas(data.root, header.prompt.tools, callSchemas)
+          captureSchemas(data.root, previousTools, callSchemas)
+        }
         continue
       }
       if (data.kind === 'compaction') {
@@ -212,6 +257,15 @@ export class TrajectorySnapshotBuilder implements ConversationViewBuilder<
       runningCalls,
     }
   }
+
+  private rebuildContributions(): void {
+    this.contributions = [...this.nodes.values()]
+      .sort((left, right) => left.anchorSeq - right.anchorSeq || left.key.localeCompare(right.key))
+    this.positions.clear()
+    for (const [index, contribution] of this.contributions.entries()) {
+      this.positions.set(contribution.key, index)
+    }
+  }
 }
 
 /** Trajectory target factory preserving the existing stage-oriented view model. */

packages/client/ui-trajectory/tests/snapshot-builder.spec.ts

@@ -1,6 +1,8 @@
 import { describe, expect, it } from 'vitest'
 import type { RequestView } from '@deepseek-ai/dsh-client-runtime/client'
-import type { TrajectoryConversationViewNode } from '../src/client/trajectory-contract.ts'
+import type {
+  TrajectoryContribution, TrajectoryConversationViewNode, TrajectoryRequestHeaderState,
+} from '../src/client/trajectory-contract.ts'
 import { TrajectorySnapshotBuilder } from '../src/client/trajectory-snapshot-builder.ts'
 
 function assistantRequest(startSeq: number, step: number): Extract<RequestView, { purpose: 'assistant' }> {
@@ -15,6 +17,47 @@ function assistantRequest(startSeq: number, step: number): Extract<RequestView,
   }
 }
 
+function contribution(
+  key: string,
+  anchorSeq: number,
+  data: TrajectoryContribution,
+): TrajectoryConversationViewNode {
+  return { key, kind: key, id: key, target: 'trajectory', anchorSeq, data }
+}
+
+function stepLocation(turn: number, step: number): TrajectoryRequestHeaderState['location'] {
+  const data = { get: () => undefined }
+  const stepLocation = {
+    turn,
+    step,
+    start: undefined,
+    end: undefined,
+    status: 'unknown' as const,
+    data,
+  }
+  const turnLocation = {
+    turn,
+    start: undefined,
+    end: undefined,
+    status: 'unknown' as const,
+    steps: [stepLocation],
+    data,
+  }
+  return { kind: 'step', turn: turnLocation, step: stepLocation }
+}
+
+function compactionRequest(startSeq: number): Extract<RequestView, { purpose: 'compaction' }> {
+  return {
+    purpose: 'compaction',
+    startSeq,
+    turn: null,
+    step: 0,
+    startedAt: startSeq,
+    completedAt: null,
+    status: 'running',
+  }
+}
+
 describe('TrajectorySnapshotBuilder', () => {
   it('inherits one request header across requests without repeating its prompt change', () => {
     const prompt = {
@@ -59,4 +102,130 @@ describe('TrajectorySnapshotBuilder', () => {
       ? request.promptChange?.kind
       : undefined)).toEqual(['initial', undefined])
   })
+
+  it('indexes exact step headers and the active tool schema without backward scans', () => {
+    const basePrompt = {
+      config: { provider: 'test', model: 'base' },
+      system: 'base prompt',
+      tools: [{ name: 'read', description: 'Read', parameters: { type: 'object' } }],
+    }
+    const exactPrompt = {
+      config: { provider: 'test', model: 'exact' },
+      system: 'exact prompt',
+      tools: [{ name: 'edit', description: 'Edit', parameters: { type: 'object' } }],
+    }
+    const nodes: TrajectoryConversationViewNode[] = [
+      contribution('header:base', 2, {
+        kind: 'request-header',
+        header: {
+          seq: 2,
+          time: 2,
+          prompt: basePrompt,
+          change: { seq: 2, time: 2, kind: 'initial' },
+          location: { kind: 'session' },
+        },
+      }),
+      contribution('assistant:1', 3, {
+        kind: 'assistant',
+        partial: null,
+        request: assistantRequest(3, 1),
+      }),
+      contribution('assistant:2', 5, {
+        kind: 'assistant',
+        partial: null,
+        request: assistantRequest(5, 2),
+      }),
+      contribution('header:exact', 6, {
+        kind: 'request-header',
+        header: {
+          seq: 6,
+          time: 6,
+          prompt: exactPrompt,
+          change: { seq: 6, time: 6, kind: 'system', previous: basePrompt },
+          location: stepLocation(1, 2),
+        },
+      }),
+      contribution('tool', 7, {
+        kind: 'tool',
+        root: {
+          callId: 'call-edit',
+          name: 'edit',
+          argsRaw: '{}',
+          turn: 1,
+          step: 2,
+          time: 7,
+          callView: null,
+          subCalls: [],
+        },
+      }),
+    ]
+
+    const snapshot = new TrajectorySnapshotBuilder().replace({ nodes })
+
+    expect(snapshot.requests.map(request => request.purpose === 'assistant'
+      ? request.prompt?.system
+      : undefined)).toEqual(['base prompt', 'exact prompt'])
+    expect(snapshot.callSchemas.get('call-edit')).toEqual(exactPrompt.tools[0])
+  })
+
+  it('applies session boundaries and turn errors with linear request indexes', () => {
+    const nodes: TrajectoryConversationViewNode[] = [
+      ...[assistantRequest(1, 1), assistantRequest(3, 2)].map(request => contribution(
+        `assistant:${request.step}`,
+        request.startSeq,
+        { kind: 'assistant', partial: null, request },
+      )),
+      contribution('turn-end', 5, {
+        kind: 'turn-end',
+        turn: 1,
+        time: 5,
+        error: 'turn failed',
+      }),
+      contribution('compact:10', 10, {
+        kind: 'compaction',
+        request: compactionRequest(10),
+      }),
+      contribution('compact:12', 12, {
+        kind: 'compaction',
+        request: compactionRequest(12),
+      }),
+      contribution('session-end:14', 14, { kind: 'session-end', seq: 14, time: 14 }),
+      contribution('session-end:16', 16, { kind: 'session-end', seq: 16, time: 16 }),
+    ]
+
+    const snapshot = new TrajectorySnapshotBuilder().replace({ nodes })
+
+    expect(snapshot.requests).toMatchObject([
+      { purpose: 'assistant', step: 1, status: 'complete' },
+      { purpose: 'assistant', step: 2, status: 'error', error: 'turn failed' },
+      { purpose: 'compaction', startSeq: 10, status: 'error', completedAt: 16 },
+      { purpose: 'compaction', startSeq: 12, status: 'error', completedAt: 14 },
+    ])
+  })
+
+  it('keeps cached contribution order across content updates and structural inserts', () => {
+    const builder = new TrajectorySnapshotBuilder()
+    const first = contribution('assistant:1', 1, {
+      kind: 'assistant', partial: null, request: assistantRequest(1, 1),
+    })
+    const last = contribution('assistant:3', 5, {
+      kind: 'assistant', partial: null, request: assistantRequest(5, 3),
+    })
+    expect(builder.replace({ nodes: [last, first] }).requests.map(request => request.startSeq))
+      .toEqual([1, 5])
+
+    const updatedLast = contribution('assistant:3', 5, {
+      kind: 'assistant',
+      partial: null,
+      request: { ...assistantRequest(5, 3), status: 'error', error: 'failed' },
+    })
+    expect(builder.apply({ upserts: [updatedLast] }).requests.map(request => request.startSeq))
+      .toEqual([1, 5])
+
+    const middle = contribution('assistant:2', 3, {
+      kind: 'assistant', partial: null, request: assistantRequest(3, 2),
+    })
+    expect(builder.apply({ upserts: [middle] }).requests.map(request => request.startSeq))
+      .toEqual([1, 3, 5])
+  })
 })