fix: cr

packages/client/ui-trajectory/src/client/layout.ts

@@ -54,6 +54,7 @@ export function deriveTrajectoryLayout(input: TrajectoryLayoutInput): readonly T
   const turns = new Map<number, { message: LaidCell[]; steps: Map<number, LaidCell[]> }>()
   let index = 0
   let prevAbsTime: number | null = null
+  let lastAssistantTurn: number | null = null
 
   const bucket = (turn: number) => {
     let entry = turns.get(turn)
@@ -74,9 +75,16 @@ export function deriveTrajectoryLayout(input: TrajectoryLayoutInput): readonly T
     steps.set(step, list)
   }
 
-  for (const node of nodes) {
+  for (let i = 0; i < nodes.length; i++) {
+    const node = nodes[i]
+    /* v8 ignore next -- dense-array guard: i stays within nodes.length, so the undefined arm needs a sparse array no caller builds. */
+    if (node === undefined) continue
     if (node.kind === 'user' || node.kind === 'steering') {
-      const turn = node.kind === 'steering' ? node.turn : 0
+      // user/message has no turn on the wire; enclose it in the next assistant
+      // (or partial) turn, else open the turn after the last assistant.
+      const turn = node.kind === 'steering'
+        ? node.turn
+        : enclosingUserTurn(nodes, i, partial, lastAssistantTurn)
       pushMessage(turn, {
         absTime: finiteTime(node.time),
         cell: {
@@ -96,6 +104,12 @@ export function deriveTrajectoryLayout(input: TrajectoryLayoutInput): readonly T
       const last = laidList[laidList.length - 1]
       if (last !== undefined) index = last.cell.index
       prevAbsTime = finiteTime(node.time) ?? prevAbsTime
+      lastAssistantTurn = node.turn
+      continue
+    }
+    if (node.kind === 'context') {
+      // No trajectory cell, but the surface still advances the duration cursor.
+      prevAbsTime = finiteTime(node.time) ?? prevAbsTime
       continue
     }
     if (node.kind === 'tool-result') {
@@ -149,6 +163,7 @@ export function deriveTrajectoryLayout(input: TrajectoryLayoutInput): readonly T
     })
   }
 
+  // Orphan turn-0 cells (orphaned tools / steering turn 0) fold into Turn 1.
   const prologue = turns.get(0)
   if (prologue !== undefined) {
     turns.delete(0)
@@ -269,11 +284,9 @@ function expandAssistant(
         index: ++index, kind: 'message', text: summarizeText(block.text),
         timeSeconds: messageDuration,
       }
-      if (!usageAttached && usage !== undefined) {
-        if (usage.inputTokens !== undefined) cell.input = usage.inputTokens
-        if (usage.outputTokens !== undefined) cell.output = usage.outputTokens
-        if (usage.reasoningTokens !== undefined) cell.think = usage.reasoningTokens
-        usageAttached = true
+      if (!usageAttached) {
+        attachUsage(cell, usage)
+        usageAttached = usage !== undefined
       }
       out.push({ absTime: nodeAbs, cell })
       continue
@@ -302,14 +315,45 @@ function expandAssistant(
   }
 
   if (out.length === 0 && !streaming) {
-    out.push({
-      absTime: nodeAbs,
-      cell: { index: ++index, kind: 'message', text: '', timeSeconds: messageDuration },
-    })
+    // Reasoning-only / empty success still owns provider usage on the Message row.
+    const cell: TrajectoryCellProps = {
+      index: ++index, kind: 'message', text: '', timeSeconds: messageDuration,
+    }
+    attachUsage(cell, usage)
+    out.push({ absTime: nodeAbs, cell })
   }
   return out
 }
 
+/**
+ * Turn that encloses a user/message: next assistant/steering turn, else the
+ * in-flight partial, else the turn after the last finalized assistant (or 1).
+ */
+function enclosingUserTurn(
+  nodes: ConversationSnapshot['nodes'],
+  userIndex: number,
+  partial: ConversationSnapshot['partial'],
+  lastAssistantTurn: number | null,
+): number {
+  for (let i = userIndex + 1; i < nodes.length; i++) {
+    const n = nodes[i]
+    /* v8 ignore next -- dense-array guard: i stays within nodes.length, so the undefined arm needs a sparse array no caller builds. */
+    if (n === undefined) continue
+    if (n.kind === 'assistant' || n.kind === 'steering') return n.turn
+  }
+  if (partial !== null) return partial.turn
+  if (lastAssistantTurn !== null) return lastAssistantTurn + 1
+  return 1
+}
+
+/** Copy provider usage onto a Message cell when present. */
+function attachUsage(cell: TrajectoryCellProps, usage: UsageLike | undefined): void {
+  if (usage === undefined) return
+  if (usage.inputTokens !== undefined) cell.input = usage.inputTokens
+  if (usage.outputTokens !== undefined) cell.output = usage.outputTokens
+  if (usage.reasoningTokens !== undefined) cell.think = usage.reasoningTokens
+}
+
 function indexResults(nodes: ConversationSnapshot['nodes']): Map<string, ToolResultNode> {
   const map = new Map<string, ToolResultNode>()
   for (const node of nodes) {

packages/client/ui-trajectory/tests/layout.spec.tsx

@@ -140,4 +140,67 @@ describe('deriveTrajectoryLayout', () => {
     const turns = deriveTrajectoryLayout({ nodes, partial: null, runningCalls: [] })
     expect(turns[0]?.groups[0]?.description).toBe('2.9s bash×2')
   })
+
+  it('assigns each user message to its enclosing turn instead of pooling into Turn 1', () => {
+    const nodes = [
+      { kind: 'user', seq: 1, time: 1_000, content: [{ type: 'text', text: 'first' }], source: null },
+      {
+        kind: 'assistant', seq: 2, time: 2_000, turn: 1, step: 0,
+        blocks: [{ kind: 'text', text: 'ok1' }],
+      },
+      { kind: 'user', seq: 3, time: 3_000, content: [{ type: 'text', text: 'second' }], source: null },
+      {
+        kind: 'assistant', seq: 4, time: 4_000, turn: 2, step: 0,
+        blocks: [{ kind: 'text', text: 'ok2' }],
+      },
+    ] as unknown as ConversationSnapshot['nodes']
+    const turns = deriveTrajectoryLayout({ nodes, partial: null, runningCalls: [] })
+    expect(turns.map((t) => t.turn)).toEqual([1, 2])
+    expect(turns[0]?.groups.flatMap((g) => g.cells.map((c) => c.text))).toEqual(['first', 'ok1'])
+    expect(turns[1]?.groups.flatMap((g) => g.cells.map((c) => c.text))).toEqual(['second', 'ok2'])
+  })
+
+  it('keeps usage on the fallback Message row when assistant has no text block', () => {
+    const nodes = [
+      {
+        kind: 'assistant', seq: 1, time: 5_000, turn: 1, step: 0,
+        blocks: [{ kind: 'reasoning', text: '…' }],
+        usage: { inputTokens: 11, outputTokens: 22, reasoningTokens: 3 },
+      },
+    ] as unknown as ConversationSnapshot['nodes']
+    const turns = deriveTrajectoryLayout({ nodes, partial: null, runningCalls: [] })
+    const message = turns[0]?.groups.flatMap((g) => g.cells).find((c) => c.kind === 'message')
+    expect(message).toMatchObject({
+      text: '', input: 11, output: 22, think: 3,
+    })
+  })
+
+  it('advances the duration cursor over context nodes', () => {
+    const nodes = [
+      { kind: 'user', seq: 1, time: 1_000, content: [{ type: 'text', text: 'hi' }], source: null },
+      {
+        kind: 'assistant', seq: 2, time: 2_000, turn: 1, step: 1,
+        blocks: [{ kind: 'tool-call', callId: 'c1', name: 'bash', argsRaw: '{}' }],
+      },
+      {
+        kind: 'tool-result', seq: 3, time: 3_000, callId: 'c1',
+        call: { name: 'bash', argsRaw: '{}' }, callTime: 2_100,
+        content: [], isError: false, callView: null, resultView: null,
+      },
+      {
+        kind: 'context', seq: 4, time: 9_000,
+        content: [{ type: 'text', text: 'extra' }], source: null,
+      },
+      {
+        kind: 'assistant', seq: 5, time: 10_000, turn: 1, step: 0,
+        blocks: [{ kind: 'text', text: 'done' }],
+      },
+    ] as unknown as ConversationSnapshot['nodes']
+    const turns = deriveTrajectoryLayout({ nodes, partial: null, runningCalls: [] })
+    const message = turns[0]?.groups
+      .flatMap((g) => g.cells)
+      .find((c) => c.kind === 'message' && c.text === 'done')
+    // From context at 9s, not from the earlier user/tool surfaces.
+    expect(message?.timeSeconds).toBe(1)
+  })
 })