From fc4df896a089956594a94f4da945ae9c44659328 Mon Sep 17 00:00:00 2001 From: imccyu <276526105+imccyu@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:38:51 +0800 Subject: [PATCH] perf(ui-trajectory): index trajectory snapshot assembly --- .../src/client/trajectory-snapshot-builder.ts | 130 +++++++++---- .../tests/snapshot-builder.spec.ts | 171 +++++++++++++++++- 2 files changed, 262 insertions(+), 39 deletions(-) diff --git a/packages/client/ui-trajectory/src/client/trajectory-snapshot-builder.ts b/packages/client/ui-trajectory/src/client/trajectory-snapshot-builder.ts index 585975b0aa..2f4c697275 100644 --- a/packages/client/ui-trajectory/src/client/trajectory-snapshot-builder.ts +++ b/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 +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( - header: TrajectoryRequestHeaderState, -): { turn?: number; step?: number } { +function stepKey(turn: number, step: number): string { + return `${turn}\u0000${step}` +} + +function headerStepKey(header: TrajectoryRequestHeaderState): string | undefined { const location = header.location - if (location.kind === 'step') return { turn: location.turn.turn, step: location.step.step } - if (location.kind === 'turn') return { turn: location.turn.turn } - return {} + return location.kind === 'step' + ? stepKey(location.turn.turn, location.step.step) + : undefined } function headerFor( - request: Extract, - headers: readonly TrajectoryRequestHeaderState[], + request: AssistantRequest, + headersByStep: ReadonlyMap, + previous: TrajectoryRequestHeaderState | undefined, ): TrajectoryRequestHeaderState | undefined { - const exact = headers.findLast((header) => { - const location = coordinates(header) - return location.turn === request.turn && location.step === request.step - }) - return exact ?? headers.findLast(header => header.seq < request.startSeq) + return headersByStep.get(stepKey(request.turn, request.step)) + ?? (previous !== undefined && previous.seq < request.startSeq ? previous : undefined) } function applyHeader( - request: Extract, + request: AssistantRequest, header: TrajectoryRequestHeaderState | undefined, includeChange: boolean, -): Extract { +): AssistantRequest { return header === undefined ? request : { @@ -67,26 +69,39 @@ function withRequestConfig( function captureSchemas( block: ToolCallBlock, - tools: readonly ConversationPromptSnapshot['tools'][number][], - output: Map, + toolsByName: ReadonlyMap, + output: Map, ): void { const name = 'kind' in block ? block.call?.name : block.name - const schema = name === undefined - ? undefined - : tools.find(candidate => candidate.name === name) + const schema = name === undefined ? undefined : toolsByName.get(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 { + 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 => - request.purpose === 'compaction' - && request.startSeq < boundary.seq - && request.status === 'running') + while (nextRequest < requests.length) { + const request = requests[nextRequest] + if (request === undefined || request.startSeq >= boundary.seq) break + 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() + 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 => - request.purpose === 'assistant' && request.turn === ending.turn) + const index = lastAssistantByTurn.get(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() + private readonly positions = new Map() + 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()] - .sort((left, right) => left.anchorSeq - right.anchorSeq || left.key.localeCompare(right.key)) - const headers = contributions.flatMap(node => node.data.kind === 'request-header' - ? [node.data.header] - : []) + const headersByStep = new Map() + for (const contribution of this.contributions) { + if (contribution.data.kind !== 'request-header') continue + 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() + const callSchemas = new Map() const consumedPromptChanges = new Set() + let previousHeader: TrajectoryRequestHeaderState | undefined + let previousTools: ReadonlyMap = 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 (header !== undefined) captureSchemas(data.root, header.prompt.tools, callSchemas) + if (previousHeader !== undefined && previousHeader.seq < contribution.anchorSeq) { + 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. */ diff --git a/packages/client/ui-trajectory/tests/snapshot-builder.spec.ts b/packages/client/ui-trajectory/tests/snapshot-builder.spec.ts index 87e484a433..d3cf64ac2d 100644 --- a/packages/client/ui-trajectory/tests/snapshot-builder.spec.ts +++ b/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 { @@ -15,6 +17,47 @@ function assistantRequest(startSeq: number, step: number): Extract 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 { + 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]) + }) })