perf(ui-trajectory): index trajectory snapshot assembly

This commit is contained in:
imccyu
2026-08-11 03:29:04 +08:00
parent 3469de58eb
commit fc4df896a0
2 changed files with 262 additions and 39 deletions
@@ -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(
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<RequestView, { purpose: 'assistant' }>,
headers: readonly TrajectoryRequestHeaderState[],
request: AssistantRequest,
headersByStep: ReadonlyMap<string, TrajectoryRequestHeaderState>,
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<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][],
output: Map<string, ConversationPromptSnapshot['tools'][number]>,
toolsByName: ReadonlyMap<string, ToolSchema>,
output: Map<string, ToolSchema>,
): 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<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 =>
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<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 =>
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<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()]
.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<string, TrajectoryRequestHeaderState>()
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<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 (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. */
@@ -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])
})
})