deepseek-harness/packages/client/ui-trajectory/tests/context-branches.spec.ts

import { describe, expect, it } from 'vitest'
import type {
  ConversationContext, ConversationNode, RequestView,
} from '@deepseek-ai/dsh-client-runtime/client'
import {
  deriveTrajectoryContextBranches,
  trajectoryBranchContainsRequest,
} from '../src/client/context-branches.ts'

const checkpoint = {
  kind: 'context',
  seq: 100,
  time: 100,
  content: [],
  source: { kind: 'plugin', plugin: 'compact' },
} as ConversationNode

const abandoned = {
  kind: 'assistant',
  seq: 20,
  time: 20,
  turn: 1,
  step: 1,
  blocks: [{ kind: 'text', text: 'abandoned' }],
} as ConversationNode

const current = {
  kind: 'user',
  seq: 110,
  time: 110,
  content: [{ type: 'text', text: 'rewound' }],
  source: { kind: 'plugin', plugin: 'rewind' },
} as ConversationNode

function request(
  purpose: RequestView['purpose'],
  startSeq: number,
  resultSeq?: number,
  replacementSeq?: number,
): RequestView {
  return {
    purpose,
    startSeq,
    turn: 1,
    step: purpose === 'assistant' ? 1 : 0,
    startedAt: startSeq,
    completedAt: startSeq + 1,
    status: 'complete',
    ...(resultSeq === undefined ? {} : { resultSeq }),
    ...(replacementSeq === undefined ? {} : { replacementSeq }),
  }
}

describe('trajectory context branches', () => {
  it('inherits nodes and requests by retained surface position rather than seq cutoff', () => {
    const contexts: ConversationContext[] = [
      { id: 0, nodes: [checkpoint, abandoned] },
      {
        id: 1,
        parentId: 0,
        origin: 'rewind',
        originSeq: 110,
        nodes: [checkpoint, current],
      },
    ]
    const branches = deriveTrajectoryContextBranches(contexts)
    const successor = branches[1]!

    expect(successor.nodes.map(node => node.seq)).toEqual([110])
    expect(trajectoryBranchContainsRequest(
      successor,
      request('assistant', 10, 20),
    )).toBe(false)
    expect(trajectoryBranchContainsRequest(
      successor,
      request('compaction', 90, 95, 100),
    )).toBe(true)
    expect(trajectoryBranchContainsRequest(
      successor,
      request('assistant', 111),
    )).toBe(true)
  })
})