perf(ui-trajectory): index trajectory snapshot assembly

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