perf(ui-trajectory): index trajectory snapshot assembly
This commit is contained in:
@@ -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])
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user